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

# kumoai.pquery

> PredictiveQuery, training tables, and prediction tables

A Kumo `PredictiveQuery` is a declarative syntax for describing a machine learning task. Predictive queries generate training and prediction tables which, together with a `Graph`, can be used to fit or predict a model.

***

## Enums

### `RunMode`

Defines the training budget for AutoML.

| Value    | Description                                                           |
| -------- | --------------------------------------------------------------------- |
| `FAST`   | Speeds up the search process — approximately 4× faster than `NORMAL`. |
| `NORMAL` | The default mode.                                                     |
| `BEST`   | Approximately 4× more thorough than `NORMAL`.                         |

***

## Predictive Query

### `PredictiveQuery`

Defines a machine learning task using PQL (Predictive Query Language), a concise SQL-like syntax. For details on writing PQL, see the [Predictive Query guide](/rfm/writing-predictive-queries).

```python
from kumoai.pquery import PredictiveQuery

pq = PredictiveQuery(
    graph=graph,
    query="RANK BY COUNT(orders.*, 0, 30, days) FOR EACH users.user_id",
)
```

The `Graph` this predictive query is defined over.

The PQL query string.

#### `id` `property`

**Returns** `str` — The unique ID for this predictive query.

#### `train_table` `property`

**Returns** `Union[TrainingTable, TrainingTableJob]` — The training table most recently generated by this query.

#### `prediction_table` `property`

**Returns** `Union[PredictionTable, PredictionTableJob]` — The prediction table most recently generated by this query.

#### `get_task_type()`

**Returns** `TaskType` — The detected task type (classification, regression, ranking, etc.).

#### `validate()`

Validates the PQL syntax of this query.

Whether to print validation output.

**Returns** `PredictiveQuery`

#### `suggest_training_table_plan()`

Generates a recommended `TrainingTableGenerationPlan` for this query.

The AutoML run mode to use when generating the plan.

**Returns** `TrainingTableGenerationPlan`

#### `generate_training_table()`

Generates a training table from this predictive query.

The plan specifying time windows, splits, and other generation parameters. If not provided, an intelligently generated default plan is used.

If `True`, returns a `TrainingTableJob` immediately rather than blocking.

**Returns** `Union[TrainingTable, TrainingTableJob]`

#### `suggest_prediction_table_plan()`

Generates a recommended `PredictionTableGenerationPlan`.

**Returns** `PredictionTableGenerationPlan`

#### `generate_prediction_table()`

Generates a prediction table from this predictive query.

The plan specifying the anchor time and other generation parameters. If not provided, an intelligently generated default plan is used.

If `True`, returns a `PredictionTableJob` immediately rather than blocking.

**Returns** `Union[PredictionTable, PredictionTableJob]`

#### `suggest_model_plan()`

Generates a recommended `ModelPlan` for this query.

The AutoML run mode controlling training speed vs. quality.

Required when the training table has been modified with a weight column via TrainingTable.update(). Import from kumoapi.train.

**Returns** `ModelPlan`

#### `suggest_distilled_model_plan()`

Generates a recommended `DistilledModelPlan` for online serving distillation.

The training job ID of the base GNN model to distill from.

The AutoML run mode.

Optional specification for weighted training. Obtain via `TrainingTable.update()`.

**Returns** `DistilledModelPlan`

#### `fit()`

Trains a model on this predictive query using the auto-suggested plans.

The plan for training table generation. If not provided, an intelligently generated default plan is used.

The plan for model training. If not provided, an intelligently generated default plan is used.

If `True`, returns a `TrainingJob` immediately rather than blocking.

**Returns** `Tuple[Trainer, Union[TrainingJobResult, TrainingJob]]`

#### `generate_baseline()`

Generates baseline metrics for comparison.

The metrics to compute for the baseline.

The training table to use.

If `True`, returns a `BaselineJob` immediately rather than blocking.

**Returns** `Union[BaselineJob, BaselineJobResult]`

#### `save()`

Saves this predictive query to Kumo.

Optional name for the saved query.

**Returns** `PredictiveQueryID`

#### `load()` `classmethod`

Loads a predictive query from its ID or a named template.

The predictive query ID or template name.

**Returns** `PredictiveQuery`

#### `load_from_training_job()` `classmethod`

Loads the predictive query associated with an existing training job.

The training job ID.

**Returns** `PredictiveQuery`

***

### `TrainingTableGenerationPlan`

Configuration for training table generation. Specifies time windows, training/validation/holdout splits, and other generation parameters. Obtain a recommended plan via [`PredictiveQuery.suggest_training_table_plan()`](#suggest-training-table-plan).

***

### `PredictionTableGenerationPlan`

Configuration for prediction table generation. Specifies the anchor time and other parameters. Obtain a recommended plan via [`PredictiveQuery.suggest_prediction_table_plan()`](#suggest-prediction-table-plan).

***

## Training Table

### `TrainingTable`

A training dataset generated from a `PredictiveQuery`. Can be initialized from the job ID of a completed training table generation job.

```python
train_table = pq.generate_training_table(plan=plan)
df = train_table.data_df()
```

The ID of the completed training table generation job.

#### `data_df()`

**Returns** `pd.DataFrame` — The generated training data.

#### `data_urls()`

**Returns** `List[str]` — Download URLs for the training table data.

#### `validate_custom_table()`

Validates a custom training table modification before applying it.

The semantic type of the source table column used for the custom modification.

The training table modification to validate.

If `True`, performs more thorough validation checks.

**Returns** `None`

#### `export()`

Exports the training table to an external connector.

The output destination configuration.

If `True`, returns an `ArtifactExportJob` immediately.

**Returns** `Union[ArtifactExportJob, ArtifactExportResult]`

#### `update()`

Modifies the training table by adding a weight column for weighted training. Import `TrainingTableSpec` from `kumoapi.train`.

The modified source table containing the weight column.

The modification spec, e.g. `TrainingTableSpec(weight_col="weight")`.

Whether to validate the modified training table against the original.

Whether to also validate row count consistency. Can be slow for large tables.

**Returns** `TrainingTable`

***

### `TrainingTableJob`

Represents an ongoing training table generation job.

#### `id` `property`

**Returns** `GenerateTrainTableJobID` - The unique job ID.

#### `result()`

Blocks until complete and returns the `TrainingTable`.

**Returns** `TrainingTable`

#### `status()`

**Returns** `JobStatusReport` — Current job status.

#### `cancel()`

Cancels the training table generation job.

#### `future()`

**Returns** `Future[TrainingTable]` — The underlying future object.

#### `load_config()`

**Returns** `GenerateTrainTableRequest` — The full configuration for this job.

***

## Prediction Table

### `PredictionTable`

A prediction dataset generated from a `PredictiveQuery`. Can be initialized from a job ID or a custom data path on supported object storage.

```python
pred_table = pq.generate_prediction_table(plan=plan)
df = pred_table.data_df()
```

The ID of the completed prediction table generation job. Leave `None` when using `table_data_path`.

Path to custom prediction table data on S3 (`s3://...`) or a Databricks UC Volume (`dbfs:/Volumes/...`). Leave `None` when using `job_id`.

#### `anchor_time` `property`

**Returns** `Optional[datetime]` — The anchor time for the generated prediction table, or `None` for custom-specified data.

#### `data_df()`

**Returns** `pd.DataFrame` — The prediction table data.

#### `data_urls()`

**Returns** `List[str]` — Download URLs for the prediction table data.

***

### `PredictionTableJob`

Represents an ongoing prediction table generation job.

#### `id` `property`

**Returns** `GeneratePredictionTableJobID` — The unique job ID.

#### `result()`

Blocks until complete and returns the `PredictionTable`.

**Returns** `PredictionTable`

#### `status()`

**Returns** `JobStatusReport`

#### `cancel()`

Cancels the prediction table generation job.

#### `future()`

**Returns** `Future[PredictionTable]` — The underlying future object.

#### `load_config()`

**Returns** `GeneratePredictionTableRequest` — The full configuration for this job.