> ## Documentation Index
> Fetch the complete documentation index at: https://kumo.ai/docs/llms.txt
> Use this file to discover all available pages before exploring further.

# 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 theme={null}
from kumoai.pquery import PredictiveQuery

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

<ParamField body="graph" type="Graph" required>
  The `Graph` this predictive query is defined over.
</ParamField>

<ParamField body="query" type="str" required>
  The PQL query string.
</ParamField>

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

<ParamField body="verbose" type="bool" default="True">
  Whether to print validation output.
</ParamField>

**Returns** `PredictiveQuery`

#### `suggest_training_table_plan()`

Generates a recommended `TrainingTableGenerationPlan` for this query.

<ParamField body="run_mode" type="RunMode" default="RunMode.FAST">
  The AutoML run mode to use when generating the plan.
</ParamField>

**Returns** `TrainingTableGenerationPlan`

#### `generate_training_table()`

Generates a training table from this predictive query.

<ParamField body="plan" type="Optional[TrainingTableGenerationPlan]" default="None">
  The plan specifying time windows, splits, and other generation parameters. If not provided, an intelligently generated default plan is used.
</ParamField>

<ParamField body="non_blocking" type="bool" default="False">
  If `True`, returns a `TrainingTableJob` immediately rather than blocking.
</ParamField>

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

<ParamField body="plan" type="Optional[PredictionTableGenerationPlan]" default="None">
  The plan specifying the anchor time and other generation parameters. If not provided, an intelligently generated default plan is used.
</ParamField>

<ParamField body="non_blocking" type="bool" default="False">
  If `True`, returns a `PredictionTableJob` immediately rather than blocking.
</ParamField>

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

#### `suggest_model_plan()`

Generates a recommended `ModelPlan` for this query.

<ParamField body="run_mode" type="RunMode" default="RunMode.FAST">
  The AutoML run mode controlling training speed vs. quality.
</ParamField>

<ParamField body="train_table_spec" type="TrainingTableSpec" default="None">
  Required when the training table has been modified with a weight column via TrainingTable.update(). Import from kumoapi.train.
</ParamField>

**Returns** `ModelPlan`

#### `suggest_distilled_model_plan()`

Generates a recommended `DistilledModelPlan` for online serving distillation.

<ParamField body="base_model_id" type="str" required>
  The training job ID of the base GNN model to distill from.
</ParamField>

<ParamField body="run_mode" type="RunMode" default="RunMode.FAST">
  The AutoML run mode.
</ParamField>

<ParamField body="train_table_spec" type="Optional[TrainingTableSpec]" default="None">
  Optional specification for weighted training. Obtain via `TrainingTable.update()`.
</ParamField>

**Returns** `DistilledModelPlan`

#### `fit()`

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

<ParamField body="training_table_plan" type="Optional[TrainingTableGenerationPlan]" default="None">
  The plan for training table generation. If not provided, an intelligently generated default plan is used.
</ParamField>

<ParamField body="model_plan" type="Optional[ModelPlan]" default="None">
  The plan for model training. If not provided, an intelligently generated default plan is used.
</ParamField>

<ParamField body="non_blocking" type="bool" default="False">
  If `True`, returns a `TrainingJob` immediately rather than blocking.
</ParamField>

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

#### `generate_baseline()`

Generates baseline metrics for comparison.

<ParamField body="metrics" type="List[str]" required>
  The metrics to compute for the baseline.
</ParamField>

<ParamField body="train_table" type="Union[TrainingTable, TrainingTableJob]" required>
  The training table to use.
</ParamField>

<ParamField body="non_blocking" type="bool" default="False">
  If `True`, returns a `BaselineJob` immediately rather than blocking.
</ParamField>

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

#### `save()`

Saves this predictive query to Kumo.

<ParamField body="name" type="str" default="None">
  Optional name for the saved query.
</ParamField>

**Returns** `PredictiveQueryID`

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

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

<ParamField body="pq_id_or_template" type="str" required>
  The predictive query ID or template name.
</ParamField>

**Returns** `PredictiveQuery`

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

Loads the predictive query associated with an existing training job.

<ParamField body="training_job_id" type="str" required>
  The training job ID.
</ParamField>

**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 theme={null}
train_table = pq.generate_training_table(plan=plan)
df = train_table.data_df()
```

<ParamField body="job_id" type="GenerateTrainTableJobID" required>
  The ID of the completed training table generation job.
</ParamField>

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

<ParamField body="source_table_type" type="str" required>
  The semantic type of the source table column used for the custom modification.
</ParamField>

<ParamField body="train_table_mod" type="TrainingTableMod" required>
  The training table modification to validate.
</ParamField>

<ParamField body="extensive_validation" type="bool" default="False">
  If `True`, performs more thorough validation checks.
</ParamField>

**Returns** `None`

#### `export()`

Exports the training table to an external connector.

<ParamField body="output_config" type="TrainingTableExportConfig" required>
  The output destination configuration.
</ParamField>

<ParamField body="non_blocking" type="bool" default="True">
  If `True`, returns an `ArtifactExportJob` immediately.
</ParamField>

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

#### `update()`

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

<ParamField body="source_table" type="SourceTable" required>
  The modified source table containing the weight column.
</ParamField>

<ParamField body="train_table_mod" type="TrainingTableSpec" required>
  The modification spec, e.g. `TrainingTableSpec(weight_col="weight")`.
</ParamField>

<ParamField body="validate" type="bool" default="True">
  Whether to validate the modified training table against the original.
</ParamField>

<ParamField body="extensive_validation" type="bool" default="False">
  Whether to also validate row count consistency. Can be slow for large tables.
</ParamField>

**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 theme={null}
pred_table = pq.generate_prediction_table(plan=plan)
df = pred_table.data_df()
```

<ParamField body="job_id" type="GeneratePredictionTableJobID" default="None">
  The ID of the completed prediction table generation job. Leave `None` when using `table_data_path`.
</ParamField>

<ParamField body="table_data_path" type="str" default="None">
  Path to custom prediction table data on S3 (`s3://...`) or a Databricks UC Volume (`dbfs:/Volumes/...`). Leave `None` when using `job_id`.
</ParamField>

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