Databricks Connector | Structured Data and Graph Models

KumoRFM can connect directly to Databricks SQL warehouses using Unity Catalog, enabling predictions on enterprise-scale data without moving it out of Databricks.

Installation

The Databricks backend requires the Databricks connector:

$ pip install kumoai[databricks]

Quick Start

Connect using explicit Databricks credentials:

1 import kumoai.rfm as rfm
2 
3 graph = rfm.Graph.from_databricks(
4     connection={
5         "server_hostname": "<workspace>.cloud.databricks.com",
6         "http_path": "/sql/1.0/warehouses/<warehouse_id>",
7         "access_token": "<access_token>",
8     },
9     catalog="MY_CATALOG",
10     schema="MY_SCHEMA",
11 )

This will:

Connect to the Databricks SQL warehouse
Discover all tables in the specified catalog and schema
Infer column metadata (data types, semantic types, primary keys, time columns)
Detect foreign key relationships
Print a summary of the inferred metadata and links

Specifying Tables

Control which tables to include and customize their configuration:

1 graph = rfm.Graph.from_databricks(
2     connection={...},
3     catalog="MY_CATALOG",
4     schema="MY_SCHEMA",
5     tables=[
6         "USERS",
7         {"name": "ORDERS", "source_name": "ORDERS_SNAPSHOT"},
8         {"name": "ITEMS", "schema": "OTHER_SCHEMA"},
9     ],
10 )

Table configuration options:

Key	Description	Required
`name`	The table name used in PQL queries	Yes
`source_name`	The actual table name in Databricks (if different from `name`)	No
`catalog`	Override the default catalog for this table	No
`schema`	Override the default schema for this table	No

Authentication

Personal access token:

1 graph = rfm.Graph.from_databricks(
2     connection={
3         "server_hostname": "<workspace>.cloud.databricks.com",
4         "http_path": "/sql/1.0/warehouses/<warehouse_id>",
5         "access_token": "<personal_access_token>",
6     },
7     catalog="MY_CATALOG",
8     schema="MY_SCHEMA",
9 )

Existing connection object:

1 from kumoai.rfm.backend.databricks import connect
2 
3 conn = connect(
4     server_hostname="<workspace>.cloud.databricks.com",
5     http_path="/sql/1.0/warehouses/<warehouse_id>",
6     access_token="<access_token>",
7 )
8 graph = rfm.Graph.from_databricks(
9     connection=conn,
10     catalog="MY_CATALOG",
11     schema="MY_SCHEMA",
12 )

Catalog and Schema

The catalog and schema parameters refer to Unity Catalog naming:

catalog: The Unity Catalog catalog (top-level namespace)
schema: The schema within the catalog

If both are omitted, KumoRFM uses the current catalog and schema from the active Databricks session. Individual tables can override these defaults using catalog and schema keys in their configuration dictionary.

Controlling Metadata Inference

1 graph = rfm.Graph.from_databricks(
2     connection={...},
3     catalog="MY_CATALOG",
4     schema="MY_SCHEMA",
5     infer_metadata=False,
6     verbose=False,
7 )
8 
9 graph.infer_metadata()
10 graph.infer_links()

Manual Edge Specification

1 graph = rfm.Graph.from_databricks(
2     connection={...},
3     catalog="MY_CATALOG",
4     schema="MY_SCHEMA",
5     edges=[
6         ("ORDERS", "USER_ID", "USERS"),
7         ("ORDERS", "ITEM_ID", "ITEMS"),
8     ],
9 )