Predictive Queries

A predictive query states what you want to predict and for which entities. Temporal queries also specify a future time window. The query API is the same on Linux and in Snowflake or Databricks notebooks. This page continues with a graph and snapshot. See Predictive Query Examples for common tasks and the PQL structure guide for query syntax and supported expressions.

Define and Validate a Query

For the AdventureWorks graph, this query predicts whether an active customer orders again in the next 90 days:

from kumo.engine import PredictiveQuery
query = PredictiveQuery(
    graph,
    "PREDICT COUNT(SalesOrderHeader.*, 0, 90, days) > 0 "
    "FOR EACH Customer.CustomerID "
    "WHERE COUNT(SalesOrderHeader.*, -365, 0, days) > 0",
)
query.validate()
print(query.get_task_type())

Validation checks the query’s syntax, referenced tables and columns, and target against the graph definition. Confirm that the returned task type matches your intended problem before generating rows. This example returns binary_classification. The WHERE clause selects customers with an order in the preceding 365 days.

A temporal regression query can instead predict the number of customer orders during the next 90 days:

temporal_query = PredictiveQuery(
    graph,
    """
    PREDICT COUNT(SalesOrderHeader.*, 0, 90, days)
    FOR EACH Customer.CustomerID
    WHERE COUNT(SalesOrderHeader.*, -365, 0, days) > 0
    """,
)
temporal_query.validate()
print(temporal_query.get_task_type())

This example returns regression. The FOR EACH clause identifies the entity table, and WHERE filters the eligible population at each anchor. A window excludes its start boundary and includes its end boundary. At anchor t, (-365, 0, days) looks back before and through t; (0, 90, days) labels events after t through t + 90 days. Prediction windows use nonnegative offsets; lookback filters use past offsets. Keep the feature history and label window separate to avoid leaking future information into population selection. Generated static direct-column predictions select entities whose target is missing. Supply explicit prediction rows when you need a different population.

Generate a Training Table

Generate labeled rows from your query and snapshot:

training_table = query.generate_training_table(snapshot)
print("Training rows:", training_table.train_table_path)

This synchronous call creates the training, validation, and test artifacts used by Training. Classification training tables need at least two observed target values and must stay within the configured class limit. A one-class table raises Too few classes to predict; adjust the query or supplied labels. The returned table retains its rows, labels, split assignments, and generation settings. Training consumes those artifacts without regenerating the rows. Keep their files in scratch until the operations using them have finished.

To customize generation, import and construct a TrainingTableGenerationPlan, then pass it to the generation call. For a temporal task, offsets count backward from the latest usable timestamp in the snapshot:

from kumoapi.model_plan import TrainingTableGenerationPlan
generation_plan = TrainingTableGenerationPlan(
    train_start_offset=700,
    train_end_offset=100,
    timeframe_step=90,
)
temporal_training_table = temporal_query.generate_training_table(
    snapshot,
    plan=generation_plan,
)

These are table-generation settings, separate from the model and optimization settings in ModelPlan. For a temporal query, choose a range that has enough history for features and enough future data to calculate training labels. When both bounds are supplied, train_end_offset must be smaller than train_start_offset; equal or reversed bounds are rejected. Integer offsets and spacing use the query’s time unit, including calendar offsets for a query expressed in months.

Choose Split Assignments

Pass the split constructor as a string. For example, assign static node rows reproducibly:

random_plan = TrainingTableGenerationPlan(
    split="RandomSplit([0.8, 0.1, 0.1], seed=42)",
)

Use a split that matches the task and available timestamps:

Ratios have two or three nonnegative entries. Temporal ratio splits must cover the full dataset; random splits can leave rows unused. An explicit TimeRangeSplit still uses the temporal query’s label-window width and the generation plan’s anchor spacing. It does not make unavailable labels appear. For static array-link tasks, a supplied RandomSplit or TemporalSplit is normalized to RandomLinkSplit; its ratio survives, but a random seed does not.

If you omit split, Kumo chooses by task:

• Static node rows without time use random 80/10/10; timestamped static nodes use chronological 80/10/10.
• Static array links use random per-link 80/10/10. The public foreign-key LIST_DISTINCT ranking path uses this per-link assignment.
• Temporal non-link targets starting at offset zero use approximate date-offset 80/10/10. Other temporal targets use date-offset boundaries based on the query’s target end offset.
• Forecasting reserves chronological windows based on horizon and window width. For horizon H and window width W, the default split offsets are -2*H*W + W and -H*W + W in the query’s time unit.

Review the resulting stage counts before fitting. A split configuration does not guarantee every class or entity appears in every stage.

Generate Forecast Inputs

For the existing temporal order-count query, request three future windows, two historical target values, and year-over-year inputs:

forecast_plan = TrainingTableGenerationPlan(
    forecast_length=3,
    lag_timesteps=2,
    year_over_year=True,
)
forecast_table = temporal_query.generate_training_table(
    snapshot, plan=forecast_plan,
)

forecast_length overrides a horizon written in PQL. Each row contains three future_labels values and two autoregression_labels values, ordered from oldest to most recent. Forecasting defaults to one lag and requires a positive integer; lag_timesteps=0 is rejected. Year-over-year inputs use a 364-day offset with ten steps of padding on each side, giving horizon + 20 values. They are not a calendar-year join. Use source history covering those lookbacks and inspect missing values before fitting.

Fit forecast_table through the normal Training workflow. For prediction, reuse trainer.training_table_generation_plan, including after loading a bundle, as shown in Making Predictions. For architecture options, see Forecasting Architecture.

Filter the Entity Population

Use a WHERE clause to restrict the entities included in generated rows. Create the filtered query before generating its training or prediction table. For the AdventureWorks graph, this query predicts repeat purchases for customers in territory 1:

filtered_query = PredictiveQuery(
    graph,
    "PREDICT COUNT(SalesOrderHeader.*, 0, 90, days) > 0 "
    "FOR EACH Customer.CustomerID "
    "WHERE Customer.TerritoryID = 1",
)
filtered_query.validate()

The entity property filter applies without a temporal lookback. For activity relative to each anchor, use the earlier WHERE COUNT(SalesOrderHeader.*, -365, 0, days) > 0 pattern. Filtering can leave no prediction rows or only one training class. Inspect the selected rows before changing the training budget. Supplied prediction tables explicitly select rows; do not assume their construction reruns your generated population filter.

Inspect and Retain Rows

Read a bounded batch from the generated training rows before fitting:

import pyarrow.dataset as ds
dataset = ds.dataset(str(training_table.train_table_path), format="parquet")
print(dataset.schema)
print(dataset.head(5).to_pandas())

Training rows include entity keys, labels, and applicable timestamps or forecast sequences. For prediction populations and their selected rows, continue to Making Predictions. See Inspect the Training Table for inspecting rows in the classification example.

To retain a row artifact outside disposable scratch, copy its file or partitioned directory to a new destination:

from pathlib import Path
import shutil
source = Path(training_table.train_table_path)
destination = Path("/path/to/retained-data/training-rows")
destination.parent.mkdir(parents=True, exist_ok=True)
if source.is_dir():
    shutil.copytree(source, destination)
else:
    if destination.exists():
        raise FileExistsError(destination)
    shutil.copy2(source, destination)

A row copy is useful for inspection or a separately prepared supplied-table workflow. It is not a portable TrainingTable: native manifests, split artifacts, and graph indexes are separate. Keep original scratch artifacts for ongoing operations. Retaining row files does not save the model or its graph snapshot.

Advanced Configurations

Use these options when you already have labeled training rows, sample weights, split assignments, or a specific population to score. If Kumo generates your rows, skip these options and continue to Training.

Use Custom Training Tables

Use your own labels, sample weights, and split assignments when they already exist. The following small local fixture demonstrates the complete scalar binary-classification path. It creates 50 customers with integer IDs; IDs 40 through 49 have missing graph labels and form the default prediction population. Choose an unused source directory on local storage that is separate from, and not nested inside, your configured Kumo scratch directory:

from pathlib import Path
import pandas as pd
import kumo.engine as kumo
from kumoapi.train import TrainingTableSpec
fixture_dir = Path("/mnt/localssd/custom-training-example")
fixture_dir.mkdir(parents=True, exist_ok=False)
customers_frame = pd.DataFrame({
    "ID": range(50),
    "CHURNED": pd.array(
        [bool(i % 2) for i in range(40)] + [None] * 10,
        dtype="boolean",
    ),
    "VISITS": [float(i % 7) for i in range(50)],
})
customers_frame.to_parquet(fixture_dir / "customers.parquet", index=False)
connector = kumo.FileConnector(root_dir=str(fixture_dir))
customers = kumo.Table.from_source_table(
    connector["customers.parquet"], primary_key="ID",
).infer_metadata()
custom_graph = kumo.Graph(tables={"customers": customers}, edges=[])
custom_snapshot = custom_graph.snapshot()
custom_query = kumo.PredictiveQuery(
    custom_graph,
    "PREDICT customers.CHURNED FOR EACH customers.ID",
)
labeled_rows = pd.DataFrame({
    "ENTITY1": range(30),
    "TARGET1": [bool(i % 2) for i in range(30)],
    "WEIGHT": [1.0 + i % 3 for i in range(30)],
    "SPLIT": [0] * 20 + [1] * 6 + [2] * 4,
})
labeled_rows.to_parquet(fixture_dir / "training.parquet", index=False)
spec = TrainingTableSpec(weight_col="WEIGHT", split_col="SPLIT")
custom_training_table = kumo.TrainingTable.from_source_table(
    connector["training.parquet"], spec,
    pq=custom_query, snapshot=custom_snapshot,
)
custom_model_plan = custom_query.suggest_model_plan(
    custom_snapshot, train_table_spec=spec,
)
custom_model_plan.optimization.weight_mode = ["weighted_loss"]
# Small demonstration budget; tune it for real training.
custom_model_plan.optimization.max_epochs = 1
custom_model_plan.optimization.max_steps_per_epoch = 2
custom_model_plan.training_job.disable_explain = True
custom_trainer = kumo.Trainer(model_plan=custom_model_plan)
fit_result = custom_trainer.fit(
    custom_training_table, snapshot=custom_snapshot,
)
custom_prediction_table = custom_query.generate_prediction_table(
    custom_snapshot,
)
custom_result = custom_trainer.predict(
    custom_prediction_table, snapshot=custom_snapshot,
    binary_classification_threshold=0.5,
)
print(custom_result.frame[["ENTITY", "TARGET_PRED", "True_PROB"]].head())

The supplied table has 30 labeled rows: 20 train, 6 validation, and 4 test. Prediction returns ten rows for IDs 40 through 49. True_PROB is the predicted probability of the supplied positive label; TARGET_PRED applies the threshold. These synthetic results demonstrate the workflow, not model quality. Use the same spec in construction and model-plan suggestion. Trainer.fit receives the materialized table and its explicit snapshot.

Name ENTITY1 and TARGET1 exactly. Entity keys must match the graph primary-key type; use real keys from that snapshot, not row positions. Scalar tasks require scalar labels and scalar weights. Weights must be finite, non-null, and not all zero; negative weights are rejected for this task. WEIGHT changes each row’s contribution under weighted_loss; sample uses weights for training-row sampling, while None ignores them. See Weighting and Class Imbalance for supported combinations.

SPLIT uses only 0 (train), 1 (validation), and 2 (test). split_col overrides the generation plan’s split with this mapping. Omit weight_col or split_col if the source lacks that column; without manual assignments, the selected or inferred split applies. This constructor supports those two spec fields, not custom feature or candidate-column specifications. It preserves supplied labels instead of recalculating them from PQL.

Supply Task-Appropriate Timestamps and Targets

For temporal node, link, and forecast tasks, include a real per-row TIMESTAMP using timezone-naive UTC Arrow timestamps. Static node tasks also need TIMESTAMP when their entity table has a time_column; the untimestamped scalar fixture above does not. Do not apply that static-node rule universally to links. A static foreign-key LIST_DISTINCT ranking query rejects a graph containing relevant time columns before it checks supplied rows. Use a temporal aggregation, or a separate graph definition without those time columns, rather than trying to fix that error by adding a supplied timestamp:

Link labels use a variable-length list of target entity keys in each TARGET1 row. Per-element weight lists must have the same lengths as their corresponding target lists; scalar weights do not stand in for per-element lists. Negative weights have a separate temporal-link contract and must not be copied from that task into static or scalar examples. For example, suppose an untimestamped graph link_graph has integer-key users and products, with orders.UID and orders.PID relationships to them. After creating link_snapshot, prepare list rows whose user and product IDs exist in that snapshot:

import pyarrow as pa
import pyarrow.parquet as pq
link_query = kumo.PredictiveQuery(
    link_graph,
    "PREDICT LIST_DISTINCT(orders.PID) RANK TOP 3 FOR EACH users.ID",
)
link_rows = pa.table({
    "ENTITY1": pa.array(range(20), type=pa.int64()),
    "TARGET1": pa.array([list(range(10))] * 20, type=pa.list_(pa.int64())),
    "WEIGHT": pa.array([[1.0] * 10] * 20, type=pa.list_(pa.float64())),
    "SPLIT": pa.array([0] * 12 + [1] * 4 + [2] * 4, type=pa.int64()),
})
pq.write_table(link_rows, fixture_dir / "link-training.parquet")
link_spec = TrainingTableSpec(weight_col="WEIGHT", split_col="SPLIT")
link_training = kumo.TrainingTable.from_source_table(
    connector["link-training.parquet"], link_spec,
    pq=link_query, snapshot=link_snapshot,
)

Use link_spec in link_query.suggest_model_plan(link_snapshot, train_table_spec=link_spec) and pass link_training to Trainer.fit with that snapshot. Choose optimization.weight_mode=["weighted_loss"] for these supplied link weights; weighted sampling is not supported for this task. For scoring, supply only the user keys and pass num_classes_to_return=3 to Trainer.predict. Expect one result row per returned user/product pair.

Use Supplied Prediction Rows

To score a specific population, supply entity identifiers and, for temporal queries, each row’s anchor time:

from kumo.engine import PredictionTable
prediction_table = PredictionTable.from_source_table(
    connector["scoring_population"],
    pq=query,
    snapshot=snapshot,
    plan=trainer.training_table_generation_plan,
)

This source needs typed ENTITY1 keys and the TIMESTAMP required by the task table above, with no recognized TARGET columns. For example, use the scalar fixture’s query to score three specified customers:

scoring_rows = pd.DataFrame({"ENTITY1": [40, 42, 49]})
scoring_rows.to_parquet(fixture_dir / "scoring.parquet", index=False)
selected_table = kumo.PredictionTable.from_source_table(
    connector["scoring.parquet"], pq=custom_query,
    snapshot=custom_snapshot,
    plan=custom_trainer.training_table_generation_plan,
)
selected_result = custom_trainer.predict(
    selected_table, snapshot=custom_snapshot,
    binary_classification_threshold=0.5,
)
print(selected_result.frame[["ENTITY", "True_PROB"]])

Expect IDs 40, 42, and 49, each once. For a temporal query, add a TIMESTAMP column such as pd.to_datetime(["2024-03-01", "2024-03-08", "2024-03-15"]), using dates supported by your graph history, and construct against that temporal query’s snapshot and fitted trainer. An entity can appear more than once at distinct anchors. See Retain Supplied Row Anchors for a complete temporal example and metadata output. Column validation does not establish that every identifier exists in the snapshot. Check the resulting population and avoid null, missing, or mismatched keys. Both supplied-table constructors stage and validate their rows immediately. Use them to define a population rather than modifying generated files.

Supply Forecast Sequences

Use a generated table with the intended forecast plan as a schema template before replacing values with independently calculated labels and history. For the count forecast above, the sequence columns are variable-length Arrow lists of integers:

This round trip checks the representation without changing values. Keep it small: to_table() materializes the row data in memory.

import pyarrow.dataset as ds
import pyarrow.parquet as pq
forecast_rows = ds.dataset(
    str(forecast_table.train_table_path), format="parquet",
).to_table()
# Replace labels/history only with arrays of the same types and lengths.
pq.write_table(forecast_rows, fixture_dir / "forecast-training.parquet")
forecast_spec = TrainingTableSpec()
supplied_forecast = kumo.TrainingTable.from_source_table(
    connector["forecast-training.parquet"], forecast_spec,
    pq=temporal_query, snapshot=snapshot, plan=forecast_plan,
)

Fit the supplied rows, then use the fitted plan to prepare a scoring template:

from kumoapi.model_plan import ModelPlan
forecast_model_plan = ModelPlan()
forecast_model_plan.optimization.max_epochs = 1
forecast_model_plan.optimization.max_steps_per_epoch = 2
forecast_model_plan.training_job.disable_explain = True
forecast_trainer = kumo.Trainer(model_plan=forecast_model_plan)
forecast_trainer.fit(supplied_forecast, snapshot=snapshot)
forecast_population = temporal_query.generate_prediction_table(
    snapshot, plan=forecast_trainer.training_table_generation_plan,
)
forecast_scoring_rows = ds.dataset(
    str(forecast_population.table_path), format="parquet",
).to_table()
pq.write_table(forecast_scoring_rows, fixture_dir / "forecast-scoring.parquet")
supplied_forecast_population = kumo.PredictionTable.from_source_table(
    connector["forecast-scoring.parquet"], pq=temporal_query,
    snapshot=snapshot, plan=forecast_trainer.training_table_generation_plan,
)
forecast_result = forecast_trainer.predict(
    supplied_forecast_population, snapshot=snapshot,
)
print("Forecast result rows:", forecast_result.num_rows)

Expect three horizon rows for each selected entity and anchor in this example. Replace keys, anchors, and historical arrays together when adapting the scoring rows. Do not add TARGET1 or future_labels to scoring rows. The constructors preserve supplied arrays; they do not derive them from keys and timestamps. Keep the fitted generation plan after saving and loading the model too.

For horizon 3 and lag 2, lengths are 3, 2, and 23. Both generated and supplied forecasting reject zero lag. Do not use a null list, a differently sized sequence, or an invented fill value to represent unavailable history; retain the generated task’s aggregation semantics when calculating your own inputs. This count example does not define missing-value handling for every aggregation or certify arbitrary Arrow list representations.

Continue to Training to fit a model, or Making Predictions to score with an existing model.