Snowflake 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 and publishes a regression model, inspects training diagnostics, and scores customers in Snowflake.

Sample Data

The notebook uses Snowflake’s preloaded TPC-H sample data in SNOWFLAKE_SAMPLE_DATA.TPCH_SF1. No separate dataset import is needed when the sample database is available in your account. By default, SAMPLE_CUSTOMERS = 2000 selects 2,000 customers and their related orders, line items, and referenced parts. If SNOWFLAKE_SAMPLE_DATA is missing or inaccessible, ask your administrator to make Snowflake’s sample data available and grant access.

The notebook creates CUSTOMER, PART, ORDERS, and LINEITEM views in the writable database and schema you configure. It adds a synthetic LINEITEM_ID because the source line-item key has two columns. The sample database is read-only; these views provide the source for Kumo’s staged reads. Creating them requires the appropriate privileges.

Prepare the Notebook

Use a GPU-backed Snowflake container notebook and install kumo-engine[snowflake] to begin. The notebook has a separate installation cell: run it once, restart the kernel, rerun the configuration cell, and continue at the runtime imports. See Snowflake Notebook Setup and GPU Sizing for environment and working-storage requirements.

Set the warehouse, writable database and schema, model name, and prediction output table in the notebook’s configuration cells. Confirm that your role can read the sample data, create the views, publish to Model Registry, and write predictions to the output table. Uploading the diagnostic archive requires a writable internal stage.

The main prediction cell writes predictions to the configured Snowflake table using OVERWRITE and produces local Parquet files. Review the destination before running it. Choose a unique table name to retain an earlier batch.

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 examples, suggests a model plan, fits a model, and publishes the selected model to Snowflake Model Registry. Record the returned model name and version. The notebook uses the suggested model plan without explicitly overriding its epoch settings. Inspect and adjust the plan before training to suit your task.

The optional diagnostics cells compare trials, plot available epoch metrics, and look for explanation artifacts. These files remain in local scratch unless you archive them. Run the optional archive cells to create a ZIP of the diagnostics and upload it to an internal Snowflake stage. The main prediction cells use a configurable historical anchor, write local Parquet results, and persist predictions to the configured Snowflake table. The bounded preview reads a few rows without loading the full prediction result into memory.

To score again without fitting, record the published model name and version. In a later session, rerun setup through the graph snapshot, predictive query, and task-specific prediction options. Skip the original training, prediction, and model-recording cells, set model_name and model_version to the recorded values, and run the load_from_connector() and subsequent scoring cells. Predictions from the reload section remain local. The published model does not include source tables or the graph index.