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.

name

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

tables

Sequence[Table]

required

The tables in the graph.

edges

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.

df_dict

Dict[str, pd.DataFrame]

required

Mapping of table name to DataFrame.

edges

Sequence[EdgeLike]

default:"None"

Optional edges to add. Inferred automatically if not specified.

infer_metadata

bool

default:"True"

Whether to automatically infer column metadata.

verbose

bool

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

default:"None"

Tables to include. Includes all tables if not specified.

edges

Sequence[EdgeLike]

default:"None"

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

infer_metadata

bool

default:"True"

Whether to automatically infer column metadata.

Returns Graph

`from_snowflake()` `classmethod`

Creates a Graph from a Snowflake database.

connection

Union[SnowflakeConnection, dict, None]

default:"None"

The Snowflake connection object or credentials dict.

tables

Sequence[Union[str, dict]]

default:"None"

Tables to include. Includes all tables if not specified.

database

str

default:"None"

The Snowflake database name.

schema

str

default:"None"

The Snowflake schema name.

edges

Sequence[EdgeLike]

default:"None"

Optional edges.

infer_metadata

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

connection

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

default:"None"

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

tables

Sequence[Union[str, dict]]

default:"None"

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

edges

Sequence[EdgeLike]

default:"None"

Optional edges.

infer_metadata

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

connection

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.

tables

Sequence[Union[str, dict]]

default:"None"

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

catalog

str

default:"None"

The Unity Catalog catalog name.

schema

str

default:"None"

The Unity Catalog schema name.

edges

Sequence[EdgeLike]

default:"None"

Optional edges.

infer_metadata

bool

default:"True"

Whether to automatically infer column metadata.

verbose

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.

semantic_view_name

str

required

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

connection

Union[SnowflakeConnection, dict, None]

default:"None"

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

verbose

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.

timeseries_col

str

required

Name of the column containing per-entity observation arrays.

timestamps_col

Optional[str]

default:"None"

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

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.

anchor_time

Optional[pd.Timestamp]

default:"None"

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

entity_col

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.

num_timeframes

int

default:"1"

Number of timeframes to forecast.

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

`add_table()`

table

Table

required

The table to add.

`remove_table()`

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

name

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

name

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.

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.

src_table

str

required

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

fkey

str

required

The foreign key column name in the source table.

dst_table

str

required

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

`unlink()`

Removes a foreign key edge.

src_table

str

required

fkey

str

required

dst_table

str

required

`infer_metadata()`

verbose

bool

default:"True"

Returns Graph

`infer_links()`

Automatically detects foreign key relationships.

verbose

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.

connection

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

Graph

required

The relational graph to query over.

verbose

bool

default:"True"

Whether to print progress output during inference.

optimize

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

query

str

required

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

indices

Sequence[Union[str, float, int]]

default:"None"

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

explain

Union[bool, ExplainConfig, dict]

default:"False"

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

return_embeddings

bool

default:"False"

If True, includes entity embeddings in the output DataFrame.

anchor_time

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.

context_anchor_time

Optional[pd.Timestamp]

default:"None"

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

run_mode

Union[RunMode, str]

default:"RunMode.FAST"

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

num_neighbors

List[int]

default:"None"

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

num_hops

int

default:"2"

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

lag_timesteps

int

default:"0"

Number of lag timesteps for temporal context.

use_prediction_time

bool

default:"False"

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

inference_config

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.

max_pq_iterations

int

default:"10"

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

random_seed

Optional[int]

default:"42"

Random seed for reproducibility.

verbose

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)")

query

str

required

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

metrics

List[str]

default:"None"

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

anchor_time

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

default:"None"

The evaluation anchor time.

context_anchor_time

Optional[pd.Timestamp]

default:"None"

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

run_mode

Union[RunMode, str]

default:"RunMode.FAST"

The inference run mode.

num_neighbors

Optional[List[int]]

default:"None"

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

use_prediction_time

bool

default:"False"

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

lag_timesteps

int

default:"0"

Number of lag timesteps for temporal context.

inference_config

Optional[Union[InferenceConfig, dict]]

default:"None"

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

max_pq_iterations

int

default:"10"

Maximum number of sampling iterations to collect valid labeled examples.

random_seed

Optional[int]

default:"42"

Random seed for reproducibility.

num_hops

int

default:"2"

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

verbose

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)

task

TaskTable

required

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

explain

Union[bool, ExplainConfig, dict]

default:"False"

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

return_embeddings

bool

default:"False"

If True, includes entity embeddings in the output DataFrame.

run_mode

Union[RunMode, str]

default:"RunMode.FAST"

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

num_neighbors

List[int]

default:"None"

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

inference_config

Optional[Union[InferenceConfig, dict]]

default:"None"

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

num_hops

int

default:"2"

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

verbose

bool

default:"True"

Whether to print progress output.

exclude_cols_dict

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

default:"None"

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

use_prediction_time

bool

default:"False"

Whether to include the anchor timestamp as an additional feature.

top_k

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)

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.

metrics

List[str]

default:"None"

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

run_mode

Union[RunMode, str]

default:"RunMode.FAST"

The inference run mode.

num_neighbors

List[int]

default:"None"

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

inference_config

Optional[Union[InferenceConfig, dict]]

default:"None"

Optional inference-time model configuration.

num_hops

int

default:"2"

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

verbose

bool

default:"True"

Whether to print progress output.

exclude_cols_dict

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

default:"None"

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

use_prediction_time

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)

num_retries

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)

batch_size

Union[int, Literal['max']]

default:"\"max\""

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

num_retries

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.

query

Union[str, ValidatedPredictiveQuery]

required

The predictive query string or validated query object.

size

int

required

Maximum number of entities to generate labels for.

anchor_time

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.

random_seed

Optional[int]

default:"42"

Seed for reproducibility.

max_iterations

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.

task

TaskTable

required

The task table to augment with lagged features.

query

Union[str, ValidatedPredictiveQuery]

required

The predictive query to compute lagged target features from.

lag_timesteps

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.

connection

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

skip_summary

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