Python SDK

kumoai.pquery

View as Markdown

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.

ValueDescription
FASTSpeeds up the search process — approximately 4× faster than NORMAL.
NORMALThe default mode.
BESTApproximately 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.

1from kumoai.pquery import PredictiveQuery
2
3pq = PredictiveQuery(
4    graph=graph,
5    query="RANK BY COUNT(orders.*, 0, 30, days) FOR EACH users.user_id",
6)
graph
GraphRequired

The Graph this predictive query is defined over.

query
strRequired

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.

verbose
boolDefaults to True

Whether to print validation output.

Returns PredictiveQuery

suggest_training_table_plan()

Generates a recommended TrainingTableGenerationPlan for this query.

run_mode
RunModeDefaults to 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.

plan
Optional[TrainingTableGenerationPlan]Defaults to None

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

non_blocking
boolDefaults to 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.

plan
Optional[PredictionTableGenerationPlan]Defaults to None

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

non_blocking
boolDefaults to False

If True, returns a PredictionTableJob immediately rather than blocking.

Returns Union[PredictionTable, PredictionTableJob]

suggest_model_plan()

Generates a recommended ModelPlan for this query.

run_mode
RunModeDefaults to RunMode.FAST

The AutoML run mode controlling training speed vs. quality.

train_table_spec
TrainingTableSpecDefaults to 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.

base_model_id
strRequired

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

run_mode
RunModeDefaults to RunMode.FAST

The AutoML run mode.

train_table_spec
Optional[TrainingTableSpec]Defaults to 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.

training_table_plan
Optional[TrainingTableGenerationPlan]Defaults to None

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

model_plan
Optional[ModelPlan]Defaults to None

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

non_blocking
boolDefaults to False

If True, returns a TrainingJob immediately rather than blocking.

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

generate_baseline()

Generates baseline metrics for comparison.

metrics
List[str]Required

The metrics to compute for the baseline.

train_table
Union[TrainingTable, TrainingTableJob]Required

The training table to use.

non_blocking
boolDefaults to False

If True, returns a BaselineJob immediately rather than blocking.

Returns Union[BaselineJob, BaselineJobResult]

save()

Saves this predictive query to Kumo.

name
strDefaults to None

Optional name for the saved query.

Returns PredictiveQueryID

load() classmethod

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

pq_id_or_template
strRequired

The predictive query ID or template name.

Returns PredictiveQuery

load_from_training_job() classmethod

Loads the predictive query associated with an existing training job.

training_job_id
strRequired

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.

1train_table = pq.generate_training_table(plan=plan)
2df = train_table.data_df()
job_id
GenerateTrainTableJobIDRequired

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.

source_table_type
strRequired

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

train_table_mod
TrainingTableModRequired

The training table modification to validate.

extensive_validation
boolDefaults to False

If True, performs more thorough validation checks.

Returns None

export()

Exports the training table to an external connector.

output_config
TrainingTableExportConfigRequired

The output destination configuration.

non_blocking
boolDefaults to 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.

source_table
SourceTableRequired

The modified source table containing the weight column.

train_table_mod
TrainingTableSpecRequired

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

validate
boolDefaults to True

Whether to validate the modified training table against the original.

extensive_validation
boolDefaults to 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.

1pred_table = pq.generate_prediction_table(plan=plan)
2df = pred_table.data_df()
job_id
GeneratePredictionTableJobIDDefaults to None

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

table_data_path
strDefaults to 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.