> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/sdgm/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/sdgm/_mcp/server.

# kumoai.rfm

> KumoRFM - the pre-trained Relational Foundation Model

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

```python
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](/rfm/querying-rfm), [Prediction Types](/rfm/prediction-types), and [Filters and Operators](/rfm/filters-and-operators).

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

```sql
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.

```python
from kumoai.rfm import LocalTable

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

The DataFrame backing this table.

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.

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

The tables in the graph.

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

#### `from_data()` `classmethod`

Creates a `Graph` directly from a dictionary of DataFrames.

Mapping of table name to DataFrame.

Optional edges to add. Inferred automatically if not specified.

Whether to automatically infer column metadata.

Whether to print progress output.

**Returns** `Graph`

#### `from_sqlite()` `classmethod`

Creates a `Graph` from a SQLite database.

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

Tables to include. Includes all tables if not specified.

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

Whether to automatically infer column metadata.

**Returns** `Graph`

#### `from_snowflake()` `classmethod`

Creates a `Graph` from a Snowflake database.

The Snowflake connection object or credentials dict.

Tables to include. Includes all tables if not specified.

The Snowflake database name.

The Snowflake schema name.

Optional edges.

Whether to automatically infer column metadata.

**Returns** `Graph`

#### `from_duckdb()` `classmethod`

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

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

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

Optional edges.

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

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

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

The Unity Catalog catalog name.

The Unity Catalog schema name.

Optional edges.

Whether to automatically infer column metadata.

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.

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

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

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.

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

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

Name of the column containing per-entity observation arrays.

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

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

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

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

Number of timeframes to forecast.

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

#### `add_table()`

The table to add.

#### `remove_table()`

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

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

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

The foreign key column name in the source table.

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

#### `unlink()`

Removes a foreign key edge.

#### `infer_metadata()`

**Returns** `Graph`

#### `infer_links()`

Automatically detects foreign key relationships.

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

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.

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

The relational graph to query over.

Whether to print progress output during inference.

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.

```python
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
```

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

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

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

If `True`, includes entity embeddings in the output DataFrame.

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

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

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

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

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

Number of lag timesteps for temporal context.

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

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

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

Random seed for reproducibility.

Whether to print progress output.

**Returns** `Union[pd.DataFrame, Explanation]`

#### `evaluate()`

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

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

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

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

The evaluation anchor time.

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

The inference run mode.

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

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

Number of lag timesteps for temporal context.

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

Maximum number of sampling iterations to collect valid labeled examples.

Random seed for reproducibility.

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

**Returns** `pd.DataFrame` — Metric scores.

***

#### `predict_task()`

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

```python
result = rfm.predict_task(task)
```

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

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

If `True`, includes entity embeddings in the output DataFrame.

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

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

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

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

Whether to print progress output.

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

Whether to include the anchor timestamp as an additional feature.

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.

```python
metrics_df = rfm.evaluate_task(task)
```

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 to compute. Uses task-appropriate defaults if `None`.

The inference run mode.

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

Optional inference-time model configuration.

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

Whether to print progress output.

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

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.

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

Maximum number of retry attempts on failure.

#### `batch_mode()` `context manager`

Context manager that batches multiple predictions together for efficiency.

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

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

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.

The predictive query string or validated query object.

Maximum number of entities to generate labels for.

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

Seed for reproducibility.

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.

The task table to augment with lagged features.

The predictive query to compute lagged target features from.

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

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.

```python
from kumoai.rfm import ExplainConfig

result = rfm.predict(query, explain=ExplainConfig(skip_summary=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.

```python
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.