Prediction Guide

Generate predictions with your existing model and save them to a retained destination. Continue with trainer, query, and snapshot from Training, or retrieve a published model. Complete your data connection for Linux, Snowflake, or Databricks in the same Python session. The prediction calls below are shared across these environments; output setup depends on the destination.

Select the Population

Generate the rows to score and carry the model’s training-table generation settings forward:

prediction_table = query.generate_prediction_table(
    snapshot,
    plan=trainer.training_table_generation_plan,
)
print("Selected rows:", prediction_table.row_count)
print("Prediction rows:", prediction_table.table_path)

For temporal queries, set a timezone-naive UTC prediction_time for repeatable batches, or omit it to use the latest relevant source timestamp. Reuse the trained generation plan, particularly for forecasts. Static queries do not accept a prediction time. See population and anchor options.

Choose an Export Location

On Linux, save results in the same service as your source, under a dedicated prediction prefix or schema. Snowflake notebooks use a Snowflake output table. Databricks notebooks use a Unity Catalog table, including when reading source data from S3. Reuse your configured source identity where possible; source-read access does not grant output-write access.

Choose your destination: S3, GCS, BigQuery, Snowflake, Databricks, or local files. Each example defines output_connector and output_name for writing. Continue to Generate Predictions after configuring one destination.

For files, S3, and GCS, reuse the source connector when its configured root includes a suitable, separate output prefix: set output_connector = connector and output_name = "kumo-results/batch-001/predictions.parquet". Otherwise create an output connector with the same identity as shown below. Choose a new batch path to retain previous results. Do not write into source table directories.

Amazon S3

The existing AWS identity needs bucket listing and object read/write access to the output prefix, including permission to inspect the export markers. Reuse the installed S3 integration and configured credentials.

from kumo.engine import S3Connector
output_connector = S3Connector(
    identity=connector.identity,
    root_dir="s3://my-bucket/kumo-results/batch-001/",
)
output_name = "predictions.parquet"

Google Cloud Storage

The existing Google identity needs object listing, read, create, and overwrite/delete permissions on the destination. Reuse the installed GCS integration and credentials.

from kumo.engine import GCSConnector
output_connector = GCSConnector(
    identity=connector.identity,
    root_dir="gs://my-bucket/kumo-results/batch-001/",
)
output_name = "predictions.parquet"

BigQuery Output

Use an existing output dataset in the configured location. The identity needs the BigQuery output permissions, including table creation and data writes. Reuse the source identity and billing project.

from kumo.engine import BigQueryConnector
output_connector = BigQueryConnector(
    identity=connector.identity,
    project_id=connector.project_id,
    billing_project_id=connector.billing_project_id,
    dataset_id="predictions",
    location=connector.location,
)
output_name = "customer_predictions_batch_001"

Use an unqualified table ID without a dataset prefix or file extension. OVERWRITE creates the table or replaces its rows and schema.

Snowflake Output

First prepare the Snowflake output schema and grants. The same requirements apply when Python runs on Linux. Reuse the source connector’s identity and query warehouse:

from kumo.engine import SnowflakeConnector
output_connector = SnowflakeConnector(
    identity=connector.identity,
    warehouse=connector.warehouse,
    database=connector.database,
    role=connector.role,
    schema_name="PREDICTIONS",
)
output_name = "CUSTOMER_PREDICTIONS_BATCH_001"

Replace PREDICTIONS with your output schema. Use an unqualified table name. OVERWRITE retains grants on an existing compatible destination; its column names must match the generated output. If the source connector already targets your output schema, reuse it with output_connector = connector. Table retention follows Snowflake policies.

Databricks Output

Prepare a dedicated output schema and staging volume using Databricks output setup. Use the output_connector defined there. With a Unity Catalog source, it reuses the source identity and compute target. With an S3 source in a Databricks notebook, create the separate Databricks connector using notebook authentication and an explicit SQL warehouse, catalog, and schema. S3 access does not grant Unity Catalog access. The Unity Catalog connector settings also apply when Python runs on Linux.

readback_name = "customer_predictions_batch_001"
output_name = f"{output_connector.schema_name}__kumo__{readback_name}"

The encoded name is for writes; readback uses the bare table name. Table retention follows Databricks policies.

Local Files

For a file source, select a writable retained directory separately from the source table directories and execution scratch. Retention follows the lifetime of that filesystem.

from kumo.engine import FileConnector
output_connector = FileConnector(root_dir="/data/kumo-results/batch-001")
output_name = "predictions.parquet"

Generate Predictions

Score and save the selected population in one call. Binary classification requires a finite threshold from 0 to 1; link prediction requires an explicit candidate count from 1 to 500. The example sets these only for the applicable task. Other tasks use their query’s defaults unless you customize the output.

from kumo.engine import PredictionOutputConfig
from kumoapi.jobs import PredictionArtifactType, WriteMode
from kumoapi.task import TaskType
prediction_options = {}
task_type = query.get_task_type()
if task_type == TaskType.BINARY_CLASSIFICATION:
    prediction_options["binary_classification_threshold"] = 0.5
elif task_type.is_link_pred:
    prediction_options["num_classes_to_return"] = 10
output_config = PredictionOutputConfig(
    connector=output_connector,
    artifact_type=PredictionArtifactType.PREDICTIONS,
    table_name=output_name,
    write_mode=WriteMode.OVERWRITE,
)
result = trainer.predict(
    prediction_table,
    snapshot=snapshot,
    output_config=output_config,
    **prediction_options,
)
print("Output rows:", result.num_rows)

Prediction runs synchronously. The snapshot must support the trained target, entity, relationships, and features, but it can contain newer data. The returned result still references local Parquet files; the saved destination has its own lifetime.

Continue to Advanced Configurations for population filters, task output, timestamps, parallel scoring, exports, and embeddings.