Training Tables | NVIDIA Kumo

A predictive query defines the task. A training table holds the labeled rows Kumo learns from, and a prediction table holds the entities it scores. Kumo can generate both tables from your query, or you can supply your own rows. This page continues with the query, temporal_query, and snapshot from Predictive Queries.

Static and Temporal Rows

The query type decides what a row is. A static query labels each entity once, using the value already stored in its target column. A temporal query labels each entity at one or more anchor times, using what happened in the window after each anchor.

	Static query	Temporal query
Example	`PREDICT customers.CHURNED FOR EACH customers.ID`	`PREDICT COUNT(SalesOrderHeader.*, 0, 90, days) > 0 FOR EACH Customer.CustomerID`
Training rows	One row per entity whose target has a value. The label is that stored value.	One row per eligible entity at each anchor time. The label is calculated from the window after the anchor.
Prediction rows	For a direct-column target, the entities whose target is missing	The eligible entities at the prediction time
Prediction time	Not accepted	Optional. Defaults to the latest relevant source timestamp.
Default split	Random 80/10/10, or chronological when the entity table has a time column	Date-offset boundaries based on the anchor times
Generation settings	`split`	`split`, plus the anchor range and spacing below and the forecast settings

In a static training table, the same customer appears at most once. In a temporal training table, a customer can appear at several anchors, each with its own label. If every entity in a static query already has a target, Kumo has nothing to score and prediction table generation raises an error. Supply prediction rows to score those entities.

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,
)

Setting	Purpose
`split`	Choose the assignment of rows to training, validation, and test.
`train_start_offset`, `train_end_offset`	Set temporal training bounds as offsets from the latest usable timestamp. Unless the query uses a custom aggregation time unit, these values are in days.
`timeframe_step`	Set the spacing between temporal training anchors. Unless the query uses a custom aggregation time unit, this value is in days.
`forecast_length`	Set the number of future windows for forecasting.
`lag_timesteps`	Include previous target values as forecast inputs.
`year_over_year`	Include corresponding historical values from the previous year.

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:

Split string	Use
`RandomSplit([0.8, 0.1, 0.1], seed=42)`	Random row assignment; the seed controls this split.
`RandomLinkSplit([0.8, 0.1, 0.1])`	Split individual links in a static array-target task. No seed argument.
`TemporalSplit([0.8, 0.1, 0.1])`	Chronological assignment of timestamped static rows.
`ApproxDateOffsetSplit([0.8, 0.1, 0.1])`	Approximate chronological proportions for temporal targets.
`DateOffsetSplit([-60, -30], offset_unit='days')`	Explicit relative boundaries, using one or two ordered nonpositive offsets.
`TimeRangeSplit([('2024-01-01', '2024-03-01'), ('2024-03-01', '2024-04-01'), ('2024-04-01', '2024-05-01')])`	Two or three nonoverlapping chronological intervals covered by your data.
`ManualSplit(col='SPLIT', train_value=0, val_value=1, test_value=2)`	Read assignments from supplied rows. Prefer `TrainingTableSpec(split_col='SPLIT')` below.

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. See Forecasting Architecture.

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. Continue to Making Predictions. See Inspect the Training Table.

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

Kumo usually generates training and prediction rows from your query. Supply your own rows instead when you already have labels, sample weights, or split assignments, or when you want to score a specific set of entities. If Kumo generates your rows, skip this section and continue to Training.

Use Custom Training Tables

A custom training table is a file of labeled rows that Kumo trains on as is. It needs these columns:

Column	Required	Contents
`ENTITY1`	Yes	The entity key, with the same type as the graph’s primary key
`TARGET1`	Yes	The label for that entity
`WEIGHT`	No	A sample weight for the row. Enable it with `weight_col="WEIGHT"`.
`SPLIT`	No	`0` for train, `1` for validation, `2` for test. Enable it with `split_col="SPLIT"`.

The following steps use a small local example: 50 customers, of which IDs 40 through 49 have no churn label.

1. Create the example graph and query. 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",
)

2. Load your labeled rows as a training table. These 30 rows carry their own labels, weights, and splits: 20 train, 6 validation, and 4 test.

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,
)

Kumo keeps your labels as they are instead of recalculating them from the query.

3. Train and score. Pass the same spec when suggesting the model plan so Kumo accounts for the weights:

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())

Prediction returns ten rows, for IDs 40 through 49. True_PROB is the predicted probability of the positive label, and TARGET_PRED applies the threshold. These synthetic results demonstrate the workflow, not model quality.

Rules for labels, weights, and splits

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.

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. Trainer.fit receives the materialized table and its explicit snapshot.

When to include a TIMESTAMP column

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 example above does not.

Supplied-row task	Training `TIMESTAMP`	Prediction `TIMESTAMP`
Temporal task	Required	Required
Static node without entity time column	Optional	Optional
Static node with entity time column	Required	Required
Static foreign-key list ranking without graph time columns	Optional	Optional

Do not apply the static-node rule 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.

Supply list targets for ranking tasks

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 specific entities instead of the generated population, write their keys to a file and load it as a prediction table. This example scores three customers from the example above:

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. The file needs typed ENTITY1 keys and no TARGET columns. Pass the fitted trainer’s training_table_generation_plan as plan.

For a temporal query, also add a TIMESTAMP column with each row’s anchor time, such as pd.to_datetime(["2024-03-01", "2024-03-08", "2024-03-15"]), using dates covered by your graph history. An entity can appear more than once at distinct anchors. See Retain Supplied Row Anchors.

Validation notes

Both supplied-table constructors stage and validate their rows immediately. Use them to define a population rather than modifying generated files. Column validation does not establish that every identifier exists in the snapshot. Check the resulting population and avoid null, missing, or mismatched keys.

Supply Forecast Sequences

Forecasting tables need extra sequence columns for future labels and past history. Most readers should let Kumo generate them. Expand the section below only if you calculate these values yourself.

Forecast sequence columns and example

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:

Column	Training	Prediction	Values per row
`ENTITY1`, `TIMESTAMP`	Required	Required	Matching entity key and per-row anchor
`TARGET1`	Required	Omit	Scalar count label
`future_labels`	Required	Omit	Horizon values, beginning at the anchor
`autoregression_labels`	Required	Required	`lag_timesteps` values, oldest first
`year_over_year_labels`	Required when enabled	Required when enabled	`horizon + 20` padded historical values

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 or Making Predictions.