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

# Databricks Notebook Code Example

> Train and score a TPC-H customer order model in a Databricks notebook

Predict how many orders a customer who ordered in the previous year will place in the next 90 days. The downloadable notebook builds a four-table TPC-H graph, trains a regression model, inspects training diagnostics, and scores customers on Databricks GPU compute.

Download the Databricks TPC-H notebook

## Sample Data

The notebook reads Databricks' built-in TPC-H data in the [`samples.tpch` catalog and schema](https://docs.databricks.com/aws/en/discover/databricks-datasets). This sample dataset is supplied with Databricks workspaces, so no separate dataset import is needed. Confirm that your workspace identity can read the sample tables.

The notebook creates `customer`, `part`, `orders`, and `lineitem` views in your configured Unity Catalog schema. Its `lineitem` view adds a synthetic `lineitem_id` because the sample's line-item key has two columns. A `DatabricksConnector` reads those views through the configured SQL warehouse while training runs on the notebook's GPU compute.

## Prepare the Notebook

Use Databricks GPU compute and install `kumo-engine[databricks]`. Set the catalog, schema, SQL warehouse ID, and local scratch directory before connecting to the data. Configure the Unity Catalog volume destination before running the optional archive or model-retention sections. See [Databricks Notebook Setup](/kumo-engine/1.0.0/get-started-with-kumo-engine/installation/databricks-notebook), [Connect Your Data](/kumo-engine/1.0.0/user-code-guide/connect-your-data-databricks#databricks-unity-catalog), and [GPU Sizing](/kumo-engine/1.0.0/get-started-with-kumo-engine/gpu-sizing).

Run the installation cell once, restart Python, rerun configuration, and continue at the runtime imports. Check your permissions to read `samples.tpch`, create views in the destination schema, use the SQL warehouse, and write an optional prediction table. The optional retention sections also require write access to a Unity Catalog volume.

Prediction writeback is enabled in the supplied notebook and uses `OVERWRITE` on its dedicated output table. Review that name before running the scoring cell, or set `WRITE_PREDICTIONS = False` to keep results only in local scratch.

## Run the Example

The notebook defines keys, time columns, and feature types for the four tables; validates their relationships; and materializes a local graph snapshot. Its predictive query is:

```text
PREDICT COUNT(orders.*, 0, 90, days)
FOR EACH customer.c_custkey
WHERE COUNT(orders.*, -365, 0, days) > 0
```

It generates historical training examples and fits a model. The complete model bundle initially stays in temporary local scratch. Run the optional model-retention section to copy it to a Unity Catalog volume. The notebook uses the suggested model plan without explicitly overriding its epoch settings. Inspect and adjust the plan before training to suit your task.

Optional cells compare trials, plot available epoch metrics, and read explanation artifacts. Diagnostic archiving and model retention are separate optional actions. Set the Volume destination before running either action. The prediction cells use a configurable historical anchor, write local Parquet results, and optionally persist the output to a Unity Catalog table. The bounded preview reads a few rows without collecting every result file into memory.

The downloadable copy has no saved execution outputs, so running it uses your own configuration and data access. Retain the complete model bundle and a compatible graph snapshot or the source data needed to rebuild it for later scoring.