Python SDK

kumoai.rfm

View as Markdown

KumoRFM (Kumo Relational Foundation Model) provides a powerful interface for querying relational data using a pre-trained foundation model. Unlike traditional ML approaches that require feature engineering and model training, KumoRFM generates predictions directly from raw relational data using PQL queries.

Overview

KumoRFM consists of three main components:

  1. LocalTable — A pandas.DataFrame wrapper that manages metadata including semantic types, primary keys, and time columns.
  2. Graph — A collection of LocalTable objects with edges defining relationships between tables.
  3. KumoRFM — The main interface for querying the foundation model.

Workflow

  1. Load relational data into pandas.DataFrame objects.
  2. Create LocalTable objects (or use Graph.from_data() directly).
  3. Build a Graph defining the relationships between tables.
  4. Initialize KumoRFM with your graph.
  5. Execute predictive queries to get predictions, explanations, or evaluations.
1import pandas as pd
2from kumoai.rfm import Graph, KumoRFM
3
4graph = Graph.from_data({
5    "users": users_df,
6    "orders": orders_df,
7})
8graph.link("orders", "user_id", "users")
9
10rfm = KumoRFM(graph)
11result = rfm.predict("PREDICT COUNT(orders.*, 0, 30, days)>0 FOR users.user_id IN (1, 2, 3)")

Query Language

KumoRFM uses Predictive Query Language (PQL). For a full introduction see the Querying guide, Prediction Types, and Filters and Operators.

The KumoRFM PQL syntax requires specifying the entity to predict for:

1PREDICT <aggregation_expression> FOR <entity_specification>

Entities can be specified as:

  • A single entity: users.user_id=1
  • A tuple of entities: users.user_id IN (1, 2, 3)

Table

Abstract base class for tables in a KumoRFM graph. Implemented by LocalTable.

LocalTable

A single in-memory table backed by a pandas.DataFrame, with metadata support for primary keys, time columns, and semantic types.

1from kumoai.rfm import LocalTable
2
3table = LocalTable(df=users_df, name="users")
4table.infer_metadata()
5table.primary_key = "user_id"
df
pd.DataFrameRequired

The DataFrame backing this table.

name
strRequired

A unique name for this table within the graph.

primary_key property

Returns Optional[str] — The primary key column name.

Set via table.primary_key = "column_name".

time_column property

Returns Optional[str] — The time column name.

Set via table.time_column = "column_name".

infer_metadata()

Automatically infers dtype and stype for all columns.

Returns LocalTable

metadata property

Returns Dict — Full column metadata dictionary.

Graph

A collection of LocalTable objects with edges defining foreign key relationships — analogous to a relational database schema.

1from kumoai.rfm import Graph
2
3# From DataFrames directly:
4graph = Graph.from_data({
5    "users": users_df,
6    "orders": orders_df,
7})
8
9# Manual construction:
10graph = Graph(tables=[users_table, orders_table])
11graph.link("orders", "user_id", "users")
12graph.validate()
tables
Sequence[Table]Required

The tables in the graph.

edges
Sequence[EdgeLike]Defaults to None

Foreign key relationships as (src_table, fkey, dst_table) tuples.

from_data() classmethod

Creates a Graph directly from a dictionary of DataFrames.

df_dict
Dict[str, pd.DataFrame]Required

Mapping of table name to DataFrame.

edges
Sequence[EdgeLike]Defaults to None

Optional edges to add. Inferred automatically if not specified.

infer_metadata
boolDefaults to True

Whether to automatically infer column metadata.

verbose
boolDefaults to True

Whether to print progress output.

Returns Graph

from_sqlite() classmethod

Creates a Graph from a SQLite database.

connection
Union[AdbcSqliteConnection, SqliteConnectionConfig, str, Path, dict]Required

The SQLite connection — a path string, Path, connection config dict, or ADBC connection object.

tables
Sequence[Union[str, dict]]Defaults to None

Tables to include. Includes all tables if not specified.

edges
Sequence[EdgeLike]Defaults to None

Optional edges. Inferred from foreign key constraints if not specified.

infer_metadata
boolDefaults to True

Whether to automatically infer column metadata.

Returns Graph

from_snowflake() classmethod

Creates a Graph from a Snowflake database.

connection
Union[SnowflakeConnection, dict, None]Defaults to None

The Snowflake connection object or credentials dict.

tables
Sequence[Union[str, dict]]Defaults to None

Tables to include. Includes all tables if not specified.

database
strDefaults to None

The Snowflake database name.

schema
strDefaults to None

The Snowflake schema name.

edges
Sequence[EdgeLike]Defaults to None

Optional edges.

infer_metadata
boolDefaults to True

Whether to automatically infer column metadata.

Returns Graph

from_duckdb() classmethod

Creates a Graph from a DuckDB database. Requires pip install kumoai[duckdb].

connection
Union[AdbcDuckDBConnection, str, Path, dict, None]Defaults to None

The DuckDB connection, path to a database file, or None for an in-memory database.

tables
Sequence[Union[str, dict]]Defaults to None

Tables to include. Includes all non-temporary tables if not specified.

edges
Sequence[EdgeLike]Defaults to None

Optional edges.

infer_metadata
boolDefaults to True

Whether to automatically infer column metadata.

Returns Graph

from_databricks() classmethod

Creates a Graph from a Databricks SQL warehouse (Unity Catalog). Requires pip install kumoai[databricks].

connection
Union[DatabricksConnection, dict, None]Defaults to None

A Databricks connection object or credentials dict (e.g., server_hostname, http_path, access_token). If None, opens a connection from environment variables.

tables
Sequence[Union[str, dict]]Defaults to None

Tables to include. Includes all tables in the catalog and schema if not specified.

catalog
strDefaults to None

The Unity Catalog catalog name.

schema
strDefaults to None

The Unity Catalog schema name.

edges
Sequence[EdgeLike]Defaults to None

Optional edges.

infer_metadata
boolDefaults to True

Whether to automatically infer column metadata.

verbose
boolDefaults to True

Whether to log progress information during graph construction.

Returns Graph

from_snowflake_semantic_view() classmethod

Creates a Graph from a Snowflake Semantic View. Reads the semantic view schema via SYSTEM$READ_YAML_FROM_SEMANTIC_VIEW and reconstructs tables, columns, primary keys, time columns, and relationships automatically.

semantic_view_name
strRequired

The fully-qualified name of the Snowflake Semantic View, e.g. CRM.CRM_SEMANTIC_VIEW.

connection
Union[SnowflakeConnection, dict, None]Defaults to None

A Snowflake connection object or credentials dict. If None, uses the active Snowpark session (available in Snowflake Notebooks).

verbose
boolDefaults to True

Whether to print graph metadata after construction.

Returns Graph

graph_and_pquery_from_timeseries() classmethod

Creates a Graph and a predictive query string from a time-series dataset stored as a single flat table. Each row represents one entity; the timeseries_col column holds an array of historical observations. The method splits the input into an entity table and a target table, links them, and returns a ready-to-use predictive query.

1import pandas as pd
2import kumoai.rfm as rfm
3
4df = pd.DataFrame({
5    "customer_id": [1, 2, 3],
6    "sales": [[10, 20, 15, 30], [5, 8, 12], [100, 95, 80]],
7})
8anchor = pd.Timestamp("2024-01-10")
9graph, pquery = rfm.Graph.graph_and_pquery_from_timeseries(
10    df, timeseries_col="sales", entity_col="customer_id",
11    time_delta=pd.Timedelta("1D"), anchor_time=anchor, num_timeframes=4,
12)
13model = rfm.KumoRFM(graph)
14result = model.predict(pquery, anchor_time=anchor)
df
pd.DataFrameRequired

Input DataFrame. Each row is one entity; timeseries_col holds a list of scalar observations.

timeseries_col
strRequired

Name of the column containing per-entity observation arrays.

timestamps_col
Optional[str]Defaults to None

Column holding per-entity timestamp arrays. When None, synthetic timestamps are generated from anchor_time and time_delta.

time_delta
Optional[pd.Timedelta]Defaults to None

Step size between consecutive observations. Required when timestamps_col is None. Also sets the prediction-window size in the generated query.

anchor_time
Optional[pd.Timestamp]Defaults to None

Forecast cutoff timestamp. Required when timestamps_col is None. Pass the same value to KumoRFM.predict().

entity_col
Optional[str]Defaults to None

Existing column to use as the entity primary key. When None, integer IDs are generated in a new entity_id column.

num_timeframes
intDefaults to 1

Number of timeframes to forecast.

Returns tuple[Graph, str] - The constructed graph and the predictive query string.

add_table()

table
TableRequired

The table to add.

remove_table()

Removes a table and all its connected edges from the graph.

name
strRequired

Name of the table to remove.

Returns Graph - The updated graph (supports method chaining).

Raises KeyError if no table with the given name exists.

has_table()

name
strRequired

Name of the table to check.

Returns bool - True if the graph contains a table with the given name.

table()

Returns the table object for a given name.

name
strRequired

Name of the table to retrieve.

Returns Table

Raises KeyError if no table with the given name exists.

tables property

Returns dict[str, Table] - Dictionary mapping table names to their Table objects.

edges property

Returns list[Edge] - All foreign key edges in the graph.

metadata property

Returns pd.DataFrame - DataFrame summarizing all tables with columns Name, Primary Key, Time Column, and End Time Column.

backend property

Returns DataBackend | None - The shared database backend for all tables in the graph, or None if the graph has no tables.

Adds a foreign key edge.

src_table
strRequired

The source table name (the one with the foreign key).

fkey
strRequired

The foreign key column name in the source table.

dst_table
strRequired

The destination table name (the one with the primary key).

Removes a foreign key edge.

src_table
strRequired
fkey
strRequired
dst_table
strRequired

infer_metadata()

verbose
boolDefaults to True

Returns Graph

Automatically detects foreign key relationships.

verbose
boolDefaults to True

Returns Graph

validate()

Validates the graph before use with KumoRFM.

Returns Graph

Prints metadata for all tables in the graph.

Prints all edges in the graph.

visualize()

Renders an interactive visualization of the graph schema.

update_connection()

Swaps the active database connection for all tables in the graph. Useful when reconnecting after a session timeout or switching to a new database instance without rebuilding the graph.

connection
AdbcSqliteConnection | AdbcDuckDBConnection | SnowflakeConnection | DatabricksConnectionRequired

The new connection object. Must match the backend type of the graph (SQLite, DuckDB, Snowflake, or Databricks).

KumoRFM

The main interface to the Kumo Relational Foundation Model. Generates predictions for any relational dataset without training.

1from kumoai.rfm import KumoRFM
2
3rfm = KumoRFM(graph)
4result = rfm.predict("PREDICT COUNT(orders.*, 0, 30, days)>0 FOR users.user_id IN (1, 2, 3)")
graph
GraphRequired

The relational graph to query over.

verbose
boolDefaults to True

Whether to print progress output during inference.

optimize
boolDefaults to False

If True, optimizes the underlying data backend for repeated querying (e.g. creates missing indices on transactional databases). Requires write access to the data backend.

predict()

Returns predictions for a PQL query.

1result = rfm.predict(
2    "PREDICT COUNT(orders.*, 0, 30, days)>0 FOR users.user_id IN (1, 2, 3)"
3)
4# Returns a DataFrame with columns: entity_id, prediction_score
5
6result_with_explain = rfm.predict(query, explain=True)
7prediction_df, summary_text = result_with_explain
query
strRequired

A PQL query string specifying the prediction task and target entities.

indices
Sequence[Union[str, float, int]]Defaults to None

Specific entity indices to predict for. Predicts for all entities if None.

explain
Union[bool, ExplainConfig, dict]Defaults to False

If True or an ExplainConfig, returns an Explanation object instead of a plain DataFrame.

return_embeddings
boolDefaults to False

If True, includes entity embeddings in the output DataFrame.

anchor_time
Union[pd.Timestamp, Literal['entity']]Defaults to None

The prediction anchor time. Uses the most recent available time if None. Pass 'entity' to use each entity’s own timestamp.

context_anchor_time
Optional[pd.Timestamp]Defaults to None

The maximum anchor time for context examples. If None, anchor_time determines the context anchor time.

run_mode
Union[RunMode, str]Defaults to RunMode.FAST

The inference run mode controlling speed vs. accuracy trade-off.

num_neighbors
List[int]Defaults to None

Per-hop neighbor counts for subgraph sampling. Uses defaults if None.

num_hops
intDefaults to 2

Number of hops for subgraph sampling. Deprecated in favor of num_neighbors.

lag_timesteps
intDefaults to 0

Number of lag timesteps for temporal context.

use_prediction_time
boolDefaults to False

Whether to use the anchor timestamp as an additional feature during prediction.

inference_config
Optional[Union[InferenceConfig, dict]]Defaults to None

Optional inference-time model configuration controlling ensembling. Supports num_estimators (1-4), column_shuffle, category_shuffle, hop_shuffle. Classification adds class_shuffle; regression/forecasting add target_transforms and output_type.

max_pq_iterations
intDefaults to 10

Maximum number of sampling iterations to collect valid labeled examples. Increase when the query has strict entity filters.

random_seed
Optional[int]Defaults to 42

Random seed for reproducibility.

verbose
boolDefaults to True

Whether to print progress output.

Returns Union[pd.DataFrame, Explanation]

evaluate()

Evaluates a PQL query against labeled data and returns metric scores.

1metrics = rfm.evaluate("PREDICT COUNT(orders.*, 0, 30, days)>0 FOR users.user_id IN (1, 2)")
query
strRequired

The PQL query string. The target entities must have ground-truth labels.

metrics
List[str]Defaults to None

Metrics to compute. Uses task-appropriate defaults if None.

anchor_time
Union[pd.Timestamp, Literal['entity']]Defaults to None

The evaluation anchor time.

context_anchor_time
Optional[pd.Timestamp]Defaults to None

The maximum anchor time for context examples. If None, anchor_time determines the context anchor time.

run_mode
Union[RunMode, str]Defaults to RunMode.FAST

The inference run mode.

num_neighbors
Optional[List[int]]Defaults to None

Per-hop neighbor counts for subgraph sampling. Uses defaults if None. Takes precedence over num_hops when provided.

use_prediction_time
boolDefaults to False

Whether to use the anchor timestamp as an additional feature during evaluation.

lag_timesteps
intDefaults to 0

Number of lag timesteps for temporal context.

inference_config
Optional[Union[InferenceConfig, dict]]Defaults to None

Optional inference-time model configuration. See predict() for supported options.

max_pq_iterations
intDefaults to 10

Maximum number of sampling iterations to collect valid labeled examples.

random_seed
Optional[int]Defaults to 42

Random seed for reproducibility.

num_hops
intDefaults to 2

Number of hops for subgraph sampling. Deprecated in favor of num_neighbors.

verbose
boolDefaults to True

Returns pd.DataFrame — Metric scores.

predict_task()

Returns predictions for a custom task specification using a TaskTable object.

1result = rfm.predict_task(task)
task
TaskTableRequired

The custom task specification, including entity, target, and context split.

explain
Union[bool, ExplainConfig, dict]Defaults to False

If True or an ExplainConfig, returns an Explanation object instead of a plain DataFrame.

return_embeddings
boolDefaults to False

If True, includes entity embeddings in the output DataFrame.

run_mode
Union[RunMode, str]Defaults to RunMode.FAST

The inference run mode controlling speed vs. accuracy trade-off.

num_neighbors
List[int]Defaults to None

Per-hop neighbor counts for subgraph sampling. Overrides num_hops when provided.

inference_config
Optional[Union[InferenceConfig, dict]]Defaults to None

Optional inference-time model configuration for ensembling or output format.

num_hops
intDefaults to 2

Number of hops for subgraph sampling. Ignored when num_neighbors is set.

verbose
boolDefaults to True

Whether to print progress output.

exclude_cols_dict
Optional[Dict[str, List[str]]]Defaults to None

Columns to exclude from model input, keyed by table name.

use_prediction_time
boolDefaults to False

Whether to include the anchor timestamp as an additional feature.

top_k
Optional[int]Defaults to None

The number of top predictions to return per entity.

Returns Union[pd.DataFrame, Explanation]

evaluate_task()

Evaluates a custom task specification against labeled data and returns metric scores.

1metrics_df = rfm.evaluate_task(task)
task
TaskTableRequired

The custom task specification. For evaluation, the prediction examples (pred_df) provided to the TaskTable must include the target column with ground-truth labels.

metrics
List[str]Defaults to None

Metrics to compute. Uses task-appropriate defaults if None.

run_mode
Union[RunMode, str]Defaults to RunMode.FAST

The inference run mode.

num_neighbors
List[int]Defaults to None

Per-hop neighbor counts for subgraph sampling. Overrides num_hops when provided.

inference_config
Optional[Union[InferenceConfig, dict]]Defaults to None

Optional inference-time model configuration.

num_hops
intDefaults to 2

Number of hops for subgraph sampling. Ignored when num_neighbors is set.

verbose
boolDefaults to True

Whether to print progress output.

exclude_cols_dict
Optional[Dict[str, List[str]]]Defaults to None

Columns to exclude from model input, keyed by table name.

use_prediction_time
boolDefaults to False

Whether to include the anchor timestamp as an additional feature.

Returns pd.DataFrame — Metric scores.

retry() context manager

Context manager that retries failed queries up to num_retries times.

1with rfm.retry(num_retries=3):
2    result = rfm.predict(query)
num_retries
intDefaults to 1

Maximum number of retry attempts on failure.

batch_mode() context manager

Context manager that batches multiple predictions together for efficiency.

1with rfm.batch_mode(batch_size=32):
2    result = rfm.predict(query)
batch_size
Union[int, Literal['max']]Defaults to "max"

Number of entities per batch. 'max' uses the largest batch size supported by the model.

num_retries
intDefaults to 1

Number of retry attempts per batch on failure.

get_train_table()

Returns the labels (training targets) of a predictive query for a given anchor time as a DataFrame. Useful for inspecting what the model will train on before launching a full training job.

query
Union[str, ValidatedPredictiveQuery]Required

The predictive query string or validated query object.

size
intRequired

Maximum number of entities to generate labels for.

anchor_time
Optional[Union[pd.Timestamp, Literal[entity]]]Defaults to None

The anchor timestamp for the query. None uses the maximum timestamp in the data. "entity" uses the timestamp of each entity.

random_seed
Optional[int]Defaults to 42

Seed for reproducibility.

max_iterations
intDefaults to 10

Maximum sampling steps before aborting.

Returns pd.DataFrame - The labels for the specified entities and anchor time.

add_lagged_target()

Adds lagged target values as input features to a TaskTable. Only supported for temporal predictive queries. Requires the task table to have a time column.

task
TaskTableRequired

The task table to augment with lagged features.

query
Union[str, ValidatedPredictiveQuery]Required

The predictive query to compute lagged target features from.

lag_timesteps
intRequired

Number of previous timesteps to use as lagged target features. Must be a positive integer.

Returns TaskTable - The task table with lagged target columns added.

update_connection()

Swaps the active database connection on the sampler. Useful when reconnecting after a session timeout without reinitializing KumoRFM.

connection
AdbcSqliteConnection | AdbcDuckDBConnection | SnowflakeConnection | DatabricksConnectionRequired

The new connection object. Must match the backend type used when the graph was created (SQLite, DuckDB, Snowflake, or Databricks).

ExplainConfig

Configuration for explainability output.

1from kumoai.rfm import ExplainConfig
2
3result = rfm.predict(query, explain=ExplainConfig(skip_summary=False))
skip_summary
boolDefaults to False

If True, skips generating a human-readable natural language summary of the explanation.

Explanation

The result of a predict() call with explain=True. Contains both the prediction scores and a natural language explanation.

1explanation = rfm.predict(query, explain=True)
2
3prediction_df = explanation.prediction  # pd.DataFrame
4summary_text = explanation.summary      # str
5
6# Supports unpacking:
7prediction_df, summary_text = explanation
8
9# Renders nicely in Jupyter:
10explanation.print()

prediction

Type pd.DataFrame — Prediction scores, one row per entity.

summary

Type str — Human-readable explanation of the most important features.

print()

Prints the prediction DataFrame and explanation summary to stdout.