Skip to main content
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.

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.
Graph
required
The Graph this predictive query is defined over.
str
required
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.
bool
default:"True"
Whether to print validation output.
Returns PredictiveQuery

suggest_training_table_plan()

Generates a recommended TrainingTableGenerationPlan for this query.
RunMode
default:"RunMode.FAST"
The AutoML run mode to use when generating the plan.
Returns TrainingTableGenerationPlan

generate_training_table()

Generates a training table from this predictive query.
Optional[TrainingTableGenerationPlan]
default:"None"
The plan specifying time windows, splits, and other generation parameters. If not provided, an intelligently generated default plan is used.
bool
default:"False"
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.
Optional[PredictionTableGenerationPlan]
default:"None"
The plan specifying the anchor time and other generation parameters. If not provided, an intelligently generated default plan is used.
bool
default:"False"
If True, returns a PredictionTableJob immediately rather than blocking.
Returns Union[PredictionTable, PredictionTableJob]

suggest_model_plan()

Generates a recommended ModelPlan for this query.
RunMode
default:"RunMode.FAST"
The AutoML run mode controlling training speed vs. quality.
TrainingTableSpec
default:"None"
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.
str
required
The training job ID of the base GNN model to distill from.
RunMode
default:"RunMode.FAST"
The AutoML run mode.
Optional[TrainingTableSpec]
default:"None"
Optional specification for weighted training. Obtain via TrainingTable.update().
Returns DistilledModelPlan

fit()

Trains a model on this predictive query using the auto-suggested plans.
Optional[TrainingTableGenerationPlan]
default:"None"
The plan for training table generation. If not provided, an intelligently generated default plan is used.
Optional[ModelPlan]
default:"None"
The plan for model training. If not provided, an intelligently generated default plan is used.
bool
default:"False"
If True, returns a TrainingJob immediately rather than blocking.
Returns Tuple[Trainer, Union[TrainingJobResult, TrainingJob]]

generate_baseline()

Generates baseline metrics for comparison.
List[str]
required
The metrics to compute for the baseline.
Union[TrainingTable, TrainingTableJob]
required
The training table to use.
bool
default:"False"
If True, returns a BaselineJob immediately rather than blocking.
Returns Union[BaselineJob, BaselineJobResult]

save()

Saves this predictive query to Kumo.
str
default:"None"
Optional name for the saved query.
Returns PredictiveQueryID

load() classmethod

Loads a predictive query from its ID or a named template.
str
required
The predictive query ID or template name.
Returns PredictiveQuery

load_from_training_job() classmethod

Loads the predictive query associated with an existing training job.
str
required
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().

PredictionTableGenerationPlan

Configuration for prediction table generation. Specifies the anchor time and other parameters. Obtain a recommended plan via PredictiveQuery.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.
GenerateTrainTableJobID
required
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.
str
required
The semantic type of the source table column used for the custom modification.
TrainingTableMod
required
The training table modification to validate.
bool
default:"False"
If True, performs more thorough validation checks.
Returns None

export()

Exports the training table to an external connector.
TrainingTableExportConfig
required
The output destination configuration.
bool
default:"True"
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.
SourceTable
required
The modified source table containing the weight column.
TrainingTableSpec
required
The modification spec, e.g. TrainingTableSpec(weight_col="weight").
bool
default:"True"
Whether to validate the modified training table against the original.
bool
default:"False"
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.
GeneratePredictionTableJobID
default:"None"
The ID of the completed prediction table generation job. Leave None when using table_data_path.
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.

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.