Databricks Connector

Use the Databricks connector to create a Kumo Relational graph from tables in a Databricks SQL warehouse through Unity Catalog. The graph represents your tables, column metadata, and relationships so Kumo Relational can sample the records relevant to a prediction without loading complete source tables into application memory. Kumo Relational sends the sampled context to the configured NVIDIA Inference Microservice (NIM), the service that runs predictions.

Installation

Install the SDK with the relational and Databricks extras:

pip install "kumo-relational-client[relational,databricks]"

Before creating a graph, collect the server hostname and HTTP path for your Databricks SQL warehouse. Ensure that your Databricks access token can use the warehouse and read the tables in the selected Unity Catalog catalog and schema.

Quick Start

Pass the warehouse connection settings and Unity Catalog namespace explicitly:

import os
from kumo_relational_client import relational
graph = relational.Graph.from_databricks(
    connection={
        "server_hostname": "<workspace>.cloud.databricks.com",
        "http_path": "/sql/1.0/warehouses/<warehouse_id>",
        "access_token": os.environ["DATABRICKS_TOKEN"],
    },
    catalog="MY_CATALOG",
    schema="MY_SCHEMA",
)

This operation performs the following actions:

1. Connects to the Databricks SQL warehouse.
2. Discovers all tables in the specified catalog and schema.
3. Infers column metadata, including data types, semantic types, primary keys, and time columns.
4. Detects foreign-key relationships.
5. Prints a summary of the inferred metadata and relationships.

Review the summary to confirm that the connector detected the intended keys, time columns, and relationships before using the graph for predictions.

Specifying Tables

By default, Kumo Relational includes every table it discovers in the selected schema. Use the tables parameter to include a subset of tables or override a table’s source name, catalog, or schema:

import os
graph = relational.Graph.from_databricks(
    connection={...},
    catalog="MY_CATALOG",
    schema="MY_SCHEMA",
    tables=[
        "USERS",
        {"name": "ORDERS", "source_name": "ORDERS_SNAPSHOT"},
        {"name": "ITEMS", "schema": "OTHER_SCHEMA"},
    ],
)

The following table describes the table configuration options:

Authentication

Store the access token in an environment variable instead of a source file or notebook:

graph = relational.Graph.from_databricks(
    connection={
        "server_hostname": "<workspace>.cloud.databricks.com",
        "http_path": "/sql/1.0/warehouses/<warehouse_id>",
        "access_token": os.environ["DATABRICKS_TOKEN"],
    },
    catalog="MY_CATALOG",
    schema="MY_SCHEMA",
)

Prefer short-lived credentials or your organization’s managed identity mechanism when the deployment supports it.

When connection arguments are omitted, the connector reads the following environment variables:

With those variables configured, construct the graph without embedding connection arguments:

from kumo_relational_client import relational
graph = relational.Graph.from_databricks()

Catalog and Schema

The catalog and schema parameters set the default Unity Catalog namespace for all tables:

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

If you omit both parameters, Kumo Relational uses the current catalog and schema from the active Databricks session. Use the catalog and schema keys in a table configuration dictionary to override either default for that table.

Controlling Metadata Inference

By default, Graph.from_databricks() infers metadata and relationships during graph creation. To run metadata and relationship inference later, set infer_metadata=False, then call the inference methods when you are ready:

graph = relational.Graph.from_databricks(
    connection={...},
    catalog="MY_CATALOG",
    schema="MY_SCHEMA",
    infer_metadata=False,
    verbose=False,
)
graph.infer_metadata()
graph.infer_links()

Manual Edge Specification

An edge represents a relationship from a foreign-key column in one table to the primary key of another table. Use the edges parameter to provide relationships explicitly when you do not want to rely on automatic detection. Each tuple contains the source table, its foreign-key column, and the referenced table:

graph = relational.Graph.from_databricks(
    connection={...},
    catalog="MY_CATALOG",
    schema="MY_SCHEMA",
    edges=[
        ("ORDERS", "USER_ID", "USERS"),
        ("ORDERS", "ITEM_ID", "ITEMS"),
    ],
)

Pass edges=None to permit declared or inferred relationships. Pass edges=[] to create the graph without edges and suppress relationship inference.