Skip to navigation

Databricks Connector

View as Markdown

There are three ways to use Databricks data: sample in place through a SQL warehouse, sample in place through a notebook’s own Spark session, or read a table into memory with read().

Installation

pip install "structured-data-models-client[relational,databricks]"

Sample in place through a SQL warehouse

from databricks.sdk import WorkspaceClient
from sdm_client import PredictiveQuery, Schema
from sdm_connectors import DatabricksSource
from sdm_connectors.databricks import warehouse
workspace = WorkspaceClient()
connection = warehouse(warehouse_id, workspace=workspace)
source = DatabricksSource(connection, catalog="main", schema="sales")
schema = Schema.from_source(
source,
{"users": "crm_users", "orders": "shop_orders"},
time_columns={"orders": "ordered_at"},
)
request = PredictiveQuery(
"PREDICT COUNT(orders.*, 0, 30, days) FOR EACH users.user_id", schema
).materialize(indices=[7, 12, 3])

warehouse() looks up the SQL warehouse by its id and authenticates as workspace — it reads no credential from the environment itself, so how workspace authenticates is the one place to look. The connection is yours; warehouse() does not close it. tables maps the names your query uses to the names in Databricks; leave it out to describe every table the source holds, which costs seconds per table on a large schema.

Sample in place through a notebook’s Spark session

In a Databricks notebook, pass the notebook’s own spark instead of a warehouse connection — no SQL warehouse needed:

source = DatabricksSource(spark, catalog="main", schema="sales")

This needs Databricks Runtime 17.3 LTS or serverless environment 4 or later, which have DataFrame.toArrow(). On an older runtime, name a SQL warehouse instead, as above.

Databricks metric views

A Databricks metric view already states its tables, joins, and columns. Convert it directly:

schema = Schema.from_source(source.semantic_model("sales_metrics"))

Its expressions run as SQL under your own credentials, so by default only pure, deterministic built-in functions are allowed; pass trust_expressions=True on semantic_model(...) to keep an expression that calls something else, when you trust whoever wrote the view. schema.conversion_notes lists what the conversion resolved.

Read into memory instead

import os
from sdm_connectors import read
from sdm_client import Schema
users = read(
"databricks",
server_hostname="<workspace>.cloud.databricks.com",
http_path="/sql/1.0/warehouses/<warehouse_id>",
access_token=os.environ["DATABRICKS_TOKEN"],
table="users",
)
orders = read("databricks", server_hostname="...", http_path="...", access_token="...", table="orders")
schema = Schema.from_data({"users": users, "orders": orders})

Prefer short-lived credentials or your organization’s managed identity mechanism when the deployment supports it. Do not store access tokens in source files.

Authenticate from the environment

When connection arguments are omitted, read() falls back to these environment variables:

Environment variableConnection argument
DATABRICKS_SERVER_HOSTNAMEserver_hostname
DATABRICKS_HTTP_PATHhttp_path
DATABRICKS_TOKENaccess_token
DATABRICKS_CATALOGcatalog
DATABRICKS_SCHEMAschema
from sdm_connectors import read
users = read("databricks", table="users")

Every connection or API request structured-data-models-client opens on Databricks identifies itself as nvidia_structured-data-models-client/<version>, as Databricks requires of a partner product, and a warehouse session it opens tags its statements structured_data_models_client_purpose: sampling in query history. A connection you open yourself and pass to a source is used as it is.

Review the inferred schema and relationships before predicting — see Build a Schema. For a model served on Databricks Model Serving, see Connect from Databricks.