Databricks Notebook Code Example

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.

Sample Data

The notebook reads Databricks’ built-in TPC-H data in the samples.tpch catalog and schema. 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, Connect Your Data, and 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:

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.