kumoai.rfm - Kumo.ai

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:

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

Workflow

Load relational data into pandas.DataFrame objects.
Create LocalTable objects (or use Graph.from_data() directly).
Build a Graph defining the relationships between tables.
Initialize KumoRFM with your graph.
Execute predictive queries to get predictions, explanations, or evaluations.

import pandas as pd
from kumoai.rfm import Graph, KumoRFM

graph = Graph.from_data({
    "users": users_df,
    "orders": orders_df,
})
graph.link("orders", "user_id", "users")

rfm = KumoRFM(graph)
result = 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:

PREDICT <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.

from kumoai.rfm import LocalTable

table = LocalTable(df=users_df, name="users")
table.infer_metadata()
table.primary_key = "user_id"

pd.DataFrame

required

The DataFrame backing this table.

str

required

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.

from kumoai.rfm import Graph

# From DataFrames directly:
graph = Graph.from_data({
    "users": users_df,
    "orders": orders_df,
})

# Manual construction:
graph = Graph(tables=[users_table, orders_table])
graph.link("orders", "user_id", "users")
graph.validate()

Sequence[Table]

required

The tables in the graph.

Sequence[EdgeLike]

default:"None"

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

`from_data()` `classmethod`

Creates a Graph directly from a dictionary of DataFrames.

Dict[str, pd.DataFrame]

required

Mapping of table name to DataFrame.

Sequence[EdgeLike]

default:"None"

Optional edges to add. Inferred automatically if not specified.

bool

default:"True"

Whether to automatically infer column metadata.

bool

default:"True"

Whether to print progress output.

Returns Graph

`from_sqlite()` `classmethod`

Creates a Graph from a SQLite database.

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

required

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

Sequence[Union[str, dict]]

default:"None"

Tables to include. Includes all tables if not specified.

Sequence[EdgeLike]

default:"None"

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

bool

default:"True"

Whether to automatically infer column metadata.

Returns Graph

`from_snowflake()` `classmethod`

Creates a Graph from a Snowflake database.

Union[SnowflakeConnection, dict, None]

default:"None"

The Snowflake connection object or credentials dict.

Sequence[Union[str, dict]]

default:"None"

Tables to include. Includes all tables if not specified.

str

default:"None"

The Snowflake database name.

str

default:"None"

The Snowflake schema name.

Sequence[EdgeLike]

default:"None"

Optional edges.

bool

default:"True"

Whether to automatically infer column metadata.

Returns Graph

`from_duckdb()` `classmethod`

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

Union[AdbcDuckDBConnection, str, Path, dict, None]

default:"None"

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

Sequence[Union[str, dict]]

default:"None"

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

Sequence[EdgeLike]

default:"None"

Optional edges.

bool

default:"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].

Union[DatabricksConnection, dict, None]

default:"None"

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

Sequence[Union[str, dict]]

default:"None"

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

str

default:"None"

The Unity Catalog catalog name.

str

default:"None"

The Unity Catalog schema name.

Sequence[EdgeLike]

default:"None"

Optional edges.

bool

default:"True"

Whether to automatically infer column metadata.

bool

default:"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.

str

required

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

Union[SnowflakeConnection, dict, None]

default:"None"

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

bool

default:"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.

import pandas as pd
import kumoai.rfm as rfm

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

pd.DataFrame

required

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

str

required

Name of the column containing per-entity observation arrays.

Optional[str]

default:"None"

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

Optional[pd.Timedelta]

default:"None"

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

Optional[pd.Timestamp]

default:"None"

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

Optional[str]

default:"None"

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

int

default:"1"

Number of timeframes to forecast.

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

`add_table()`

Table

required

The table to add.

`remove_table()`

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

str

required

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()`

str

required

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.

str

required

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.

`link()`

Adds a foreign key edge.

str

required

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

str

required

The foreign key column name in the source table.

str

required

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

`unlink()`

Removes a foreign key edge.

str

required

str

required

str

required

`infer_metadata()`

bool

default:"True"

Returns Graph

`infer_links()`

Automatically detects foreign key relationships.

bool

default:"True"

Returns Graph

`validate()`

Validates the graph before use with KumoRFM. Returns Graph

`print_metadata()`

Prints metadata for all tables in the graph.

`print_links()`

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.

AdbcSqliteConnection | AdbcDuckDBConnection | SnowflakeConnection | DatabricksConnection

required

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.

from kumoai.rfm import KumoRFM

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

Graph

required

The relational graph to query over.

bool

default:"True"

Whether to print progress output during inference.

bool

default:"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.

result = rfm.predict(
    "PREDICT COUNT(orders.*, 0, 30, days)>0 FOR users.user_id IN (1, 2, 3)"
)
# Returns a DataFrame with columns: entity_id, prediction_score

result_with_explain = rfm.predict(query, explain=True)
prediction_df, summary_text = result_with_explain

str

required

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

Sequence[Union[str, float, int]]

default:"None"

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

Union[bool, ExplainConfig, dict]

default:"False"

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

bool

default:"False"

If True, includes entity embeddings in the output DataFrame.

Union[pd.Timestamp, Literal['entity']]

default:"None"

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

Optional[pd.Timestamp]

default:"None"

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

Union[RunMode, str]

default:"RunMode.FAST"

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

List[int]

default:"None"

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

int

default:"2"

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

int

default:"0"

Number of lag timesteps for temporal context.

bool

default:"False"

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

Optional[Union[InferenceConfig, dict]]

default:"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.

int

default:"10"

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

Optional[int]

default:"42"

Random seed for reproducibility.

bool

default:"True"

Whether to print progress output.

Returns Union[pd.DataFrame, Explanation]

`evaluate()`

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

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

str

required

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

List[str]

default:"None"

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

Union[pd.Timestamp, Literal['entity']]

default:"None"

The evaluation anchor time.

Optional[pd.Timestamp]

default:"None"

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

Union[RunMode, str]

default:"RunMode.FAST"

The inference run mode.

Optional[List[int]]

default:"None"

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

bool

default:"False"

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

int

default:"0"

Number of lag timesteps for temporal context.

Optional[Union[InferenceConfig, dict]]

default:"None"

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

int

default:"10"

Maximum number of sampling iterations to collect valid labeled examples.

Optional[int]

default:"42"

Random seed for reproducibility.

int

default:"2"

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

bool

default:"True"

Returns pd.DataFrame — Metric scores.

`predict_task()`

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

result = rfm.predict_task(task)

TaskTable

required

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

Union[bool, ExplainConfig, dict]

default:"False"

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

bool

default:"False"

If True, includes entity embeddings in the output DataFrame.

Union[RunMode, str]

default:"RunMode.FAST"

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

List[int]

default:"None"

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

Optional[Union[InferenceConfig, dict]]

default:"None"

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

int

default:"2"

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

bool

default:"True"

Whether to print progress output.

Optional[Dict[str, List[str]]]

default:"None"

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

bool

default:"False"

Whether to include the anchor timestamp as an additional feature.

Optional[int]

default:"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.

metrics_df = rfm.evaluate_task(task)

TaskTable

required

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

List[str]

default:"None"

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

Union[RunMode, str]

default:"RunMode.FAST"

The inference run mode.

List[int]

default:"None"

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

Optional[Union[InferenceConfig, dict]]

default:"None"

Optional inference-time model configuration.

int

default:"2"

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

bool

default:"True"

Whether to print progress output.

Optional[Dict[str, List[str]]]

default:"None"

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

bool

default:"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.

with rfm.retry(num_retries=3):
    result = rfm.predict(query)

int

default:"1"

Maximum number of retry attempts on failure.

`batch_mode()` `context manager`

Context manager that batches multiple predictions together for efficiency.

with rfm.batch_mode(batch_size=32):
    result = rfm.predict(query)

Union[int, Literal['max']]

default:"\"max\""

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

int

default:"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.

Union[str, ValidatedPredictiveQuery]

required

The predictive query string or validated query object.

int

required

Maximum number of entities to generate labels for.

Optional[Union[pd.Timestamp, Literal[entity]]]

default:"None"

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

Optional[int]

default:"42"

Seed for reproducibility.

int

default:"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.

TaskTable

required

The task table to augment with lagged features.

Union[str, ValidatedPredictiveQuery]

required

The predictive query to compute lagged target features from.

int

required

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.

AdbcSqliteConnection | AdbcDuckDBConnection | SnowflakeConnection | DatabricksConnection

required

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.

from kumoai.rfm import ExplainConfig

result = rfm.predict(query, explain=ExplainConfig(skip_summary=False))

bool

default:"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.

explanation = rfm.predict(query, explain=True)

prediction_df = explanation.prediction  # pd.DataFrame
summary_text = explanation.summary      # str

# Supports unpacking:
prediction_df, summary_text = explanation

# Renders nicely in Jupyter:
explanation.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.

​Overview

​Workflow

​Query Language

​Table

​LocalTable

​primary_key property

​time_column property

​infer_metadata()

​metadata property

​Graph

​from_data() classmethod

​from_sqlite() classmethod

​from_snowflake() classmethod

​from_duckdb() classmethod

​from_databricks() classmethod

​from_snowflake_semantic_view() classmethod

​graph_and_pquery_from_timeseries() classmethod

​add_table()

​remove_table()

​has_table()

​table()

​tables property

​edges property

​metadata property

​backend property

​link()

​unlink()

​infer_metadata()

​infer_links()

​validate()

​print_metadata()

​print_links()

​visualize()

​update_connection()

​KumoRFM

​predict()

​evaluate()

​predict_task()

​evaluate_task()

​retry() context manager

​batch_mode() context manager

​get_train_table()

​add_lagged_target()

​update_connection()

​ExplainConfig

​Explanation

​prediction

​summary

​print()

Overview

Workflow

Query Language

`Table`

`LocalTable`

`primary_key` `property`

`time_column` `property`

`infer_metadata()`

`metadata` `property`

`Graph`

`from_data()` `classmethod`

`from_sqlite()` `classmethod`

`from_snowflake()` `classmethod`

`from_duckdb()` `classmethod`

`from_databricks()` `classmethod`

`from_snowflake_semantic_view()` `classmethod`

`graph_and_pquery_from_timeseries()` `classmethod`

`add_table()`

`remove_table()`

`has_table()`

`table()`

`tables` `property`

`edges` `property`

`metadata` `property`

`backend` `property`

`link()`

`unlink()`

`infer_metadata()`

`infer_links()`

`validate()`

`print_metadata()`

`print_links()`

`visualize()`

`update_connection()`

`KumoRFM`

`predict()`

`evaluate()`

`predict_task()`

`evaluate_task()`

`retry()` `context manager`

`batch_mode()` `context manager`

`get_train_table()`

`add_lagged_target()`

`update_connection()`

`ExplainConfig`

`Explanation`

`prediction`

`summary`

`print()`