Model Transforms

ModelTransforms lets you enrich a source (an EventSource or JoinSource) with one or more model inference outputs — treating a model call as a first-class row-level transform in a Chronon pipeline. The transformed rows can then be:

• materialized to a table by a batch backfill (useful for training sets or as the input to a downstream GroupBy),
• aggregated into online features by a streaming GroupBy whose source is the ModelTransforms (model inference happens on the Flink hot path), or
• fetched on-demand from the serving layer via fetchModelTransforms (model inference happens at request time).

Typical use cases:

• Generating item / listing embeddings off a dimension table or Join output using an off-the-shelf model (Vertex AI, OpenAI embeddings, etc.) or a custom model you have trained.
• Scoring enriched join rows with a custom model (e.g. a CTR scorer) to produce a feature that can be served online.
• Chaining a model into a streaming pipeline so that model outputs become inputs to Chronon aggregations (e.g. LAST embedding over 7 days per user).

Behind the scenes ModelTransforms is a 1:1 row-level transform (not an aggregation), so its input dependency is day-for-day.

Core concepts

A ModelTransforms config brings together three things:

1. One Source — an EventSource or (most commonly) a JoinSource on top of a parent Join. The source provides the inputs to the model.
2. One or more Model objects — each describes a model backend (Vertex AI, SageMaker, Ray, …), how to map a source row to the model's input format (input_mapping), how to pull desired fields out of the model's response (output_mapping), and the schema of what the model returns (value_fields).
3. Passthrough + key fields — passthrough_fields names columns from the source that should flow through alongside the model outputs; key_fields is the schema of the key columns the models key off (used by online serving).

The output of a ModelTransforms is a set of rows that contains:

• the passthrough columns from the source, plus
• one column per entry in each model's output_mapping, named <sanitized_model_name>__<output_key>.

Column naming: the output column from a model is <cleaned_model_name>__<output_mapping_key>. The cleaned name already includes the version suffix (dots replaced with underscores). For example a model whose metadata name is gcp.listing.item_description_model__1 with output_mapping = {"item_embedding": "..."} produces a column named gcp_listing_item_description_model__1__item_embedding. See MetadataOps.cleanName and ModelTransformsSchemaUtil in the online module.

A minimal example

The full end-to-end example lives in the canary tree:

• Model definition: python/test/canary/models/gcp/listing.py
• ModelTransforms definition: python/test/canary/model_transforms/gcp/listing.py
• Downstream GroupBy: python/test/canary/group_bys/gcp/listing_embeddings.py

The salient bits are below.

1. Define a Model

item_description_model = Model(
    version="1",
    inference_spec=InferenceSpec(
        model_backend=ModelBackend.VERTEXAI,
        model_backend_params={
            "model_name": "gemini-embedding-001",
            "model_type": "publisher",
        },
    ),
    # One Spark SQL expression per field the model expects. Here Vertex's text
    # embedding API wants an `instance` struct with a `content` field.
    input_mapping={
        "instance": "named_struct('content', concat_ws('; ', listing_id_headline, listing_id_long_description))",
    },
    # Columns to extract from the raw model response. The RHS is Spark SQL
    # over the response columns, named `<cleaned_model_name>__<field>`.
    output_mapping={
        "item_embedding": "gcp_listing_item_description_model__1__embeddings.values",
    },
    # Schema of the raw response row the backend returns.
    value_fields=[("embeddings", embeddings)],
)

Key fields:

• inference_spec — the backend plus backend-specific params. For Vertex, model_type: publisher points at a published model; model_type: custom is used for models you've deployed yourself.
• input_mapping — one Spark SQL expression per field the backend expects. Each expression runs against a row of the source dataframe.
• output_mapping — one Spark SQL expression per column you want to extract from the model's raw response. The LHS names become the final column suffixes after the <model>__ prefix.
• value_fields — the schema of the raw response row. Required for backfill — Chronon uses it to assemble the intermediate dataframe before applying output_mapping.

2. Define a ModelTransforms

source = JoinSource(
    join=demo.v1,
    # Pre-filter rows that would produce empty content; Vertex rejects empty strings.
    query=Query(wheres=[
        "(listing_id_headline IS NOT NULL AND listing_id_headline != '') OR "
        "(listing_id_long_description IS NOT NULL AND listing_id_long_description != '')"
    ]),
)
 
v1 = ModelTransforms(
    sources=[source],
    models=[listing.item_description_model],
    passthrough_fields=["user_id", "listing_id", "listing_id_is_active"],
    version=2,
    output_namespace="data",
    # key_fields is the key schema for the online fetch path. For a JoinSource
    # it can be omitted (resolved from the join metadata).
    key_fields=[
        ("listing_id_headline", DataType.STRING),
        ("listing_id_long_description", DataType.STRING),
    ],
)

Worth noting:

• sources currently accepts exactly one source (multi-source ModelTransforms isn't supported in v1).
• passthrough_fields is additive: Chronon always keeps the partition column and ts time column even if you don't list them, so downstream GroupBys always see the fields they need.
• output_namespace is optional — if omitted, it falls back to your team's default (outputNamespace in teams.json).
• version is a free-form integer bumped whenever you change the semantic definition. Combined with the file name it becomes the config name, and in turn the output table name.

Using ModelTransforms output

Once defined, the enriched output can be consumed in three different ways. Pick the pattern that matches your serving needs.

Pattern A: Batch materialization + downstream GroupBy

Running the backfill (zipline hub backfill compiled/model_transforms/<team>/<conf> --start-ds <ds> --end-ds <ds>) populates a table at modelTransforms.table that looks exactly like a regular source table — passthrough columns plus model-output columns. You can then use the ModelTransforms as the source of a GroupBy:

from model_transforms.gcp import listing as listing_mt
 
MODEL_OUTPUT_COL = "gcp_listing_item_description_model__1__item_embedding"
 
v1 = GroupBy(
    sources=[listing_mt.v1],
    keys=["user_id"],
    aggregations=[
        Aggregation(
            input_column=MODEL_OUTPUT_COL,
            operation=Operation.LAST,
            windows=[Window(length=7, time_unit=TimeUnit.DAYS)],
        ),
    ],
    accuracy=Accuracy.TEMPORAL,
    online=True,
    version=1,
)

When Chronon compiles this GroupBy, it detects the ModelTransforms source and rewrites it to read from the backfill output table. Any query.selects / query.wheres defined inside the inner ModelTransforms are ignored at this layer — they reference pre-transform columns that have been dropped by the time the ModelTransforms writes its output. Only passthroughFields plus model outputs survive, and ts is always included.

Pattern B: Streaming (model inference on the hot path)

If your ModelTransforms sources from a streaming-capable JoinSource (one whose left side has a topic) or a streaming EventSource, a GroupBy that references it becomes a streaming job that calls the model online as events arrive:

Example — a streaming variant of the listing embeddings (from the canary file linked above):

# demo_parent.parent_join has a pubsub-backed left source and enriches with
# dim_listings.v1 (so `listing_id_headline` / `listing_id_long_description`
# are present at fetch time).
streaming_source = JoinSource(join=demo_parent.parent_join, query=Query(wheres=[...]))
 
v2 = ModelTransforms(
    sources=[streaming_source],
    models=[listing.item_description_model],
    passthrough_fields=["user_id", "listing_id"],
    version=1,
    output_namespace="data",
    key_fields=[...],
)

Pointing a GroupBy at listing.v2 and setting online=True wires up the Flink job automatically. Two async operators are chained in series — join enrichment and model inference — and each has its own tunable timeout / capacity (join_enrichment_async_timeout_ms, join_enrichment_async_capacity, model_async_timeout_ms, model_async_capacity); both fall back to async_timeout_ms / async_capacity if the specific overrides aren't set.

Supported inner-source shapes for streaming (as of this writing):

• EventSource with a topic, or
• JoinSource whose join's left is an EventSource.

Pattern C: On-demand online fetch

If you don't need aggregation on top of the model output — you just want the model result enriched with (optional) join features and returned to a caller — use Fetcher.fetchModelTransforms directly. This is useful for synchronous model scoring during request handling.

import ai.chronon.online.fetcher.Fetcher.Request
 
val request = Request(
  name = "gcp.listing.v1",            // ModelTransforms metadata name
  keys = Map("listing_id" -> 42L),    // keys the join / source expects
)
 
fetcher.fetchModelTransforms(Seq(request)).map { responses =>
  responses.foreach { r =>
    // r.values is a Try[Map[String, AnyRef]] containing passthrough + model-output columns
  }
}

Under the hood:

• If the ModelTransforms is backed by a JoinSource, the fetcher first calls fetchJoin on the parent join using the request keys, then feeds those enriched rows into the model call.
• If it's backed by a plain EventSource, the request keys are fed to the model directly (after input_mapping evaluation).
• Model failures are isolated per model: a failing model produces an exception feature (<model_name>_exception) rather than failing the whole response.

Compared to the streaming path, on-demand fetch gives up KV-cached aggregates and pays a model inference call per request — simpler to operate, but higher latency and cost per request.

Model orchestration (optional)

The Model object can also carry orchestration metadata — for custom models you train and deploy yourself — via training_conf and deployment_conf. See python/test/canary/models/gcp/click_through_rate.py for a full example that trains an XGBoost CTR model on Vertex AI and deploys it to a custom container endpoint.

When training_conf is set, Chronon's planner adds a soft dependency on a model deploy node so a ModelTransforms backfill waits for the model to be redeployed when upstream training data lands. When training_conf is absent (the embeddings case above), Chronon assumes the model is managed outside of Chronon and skips that dependency.

Published / managed models (model_type: publisher for Vertex, hosted OpenAI endpoints, etc.) don't need training_conf or deployment_conf at all.

Backend integration

Model inference goes through two interfaces defined in ai.chronon.online.Api:

• ModelPlatform — the per-backend implementation. The method you'll care about is predict(PredictRequest): Future[PredictResponse], which both the Spark backfill and the Flink / fetcher online paths call. Deployment-related methods (createEndpoint, deployModel, submitTrainingJob) are only invoked when training_conf / deployment_conf are present on the model.
• ModelPlatformProvider — a factory that resolves a ModelPlatform based on the model's ModelBackend + modelBackendParams. This is what gets handed to the Flink job, backfill, and fetcher.

Your Api subclass should override generateModelPlatformProvider to return the provider you want (GCP, AWS, etc.). The reference GCP implementation lives in cloud_gcp/.../GcpApiImpl.scala and ships a Vertex AI–backed ModelPlatform.

Compile, run, and verify

With the three files in place (models/<team>/foo.py, model_transforms/<team>/foo.py, and the consuming group_bys/<team>/foo.py if any):

# Compile (runs ruff + emits Thrift JSON under compiled/)
zipline compile
 
# Backfill the ModelTransforms output table
zipline hub backfill compiled/model_transforms/<team>/<conf> --start-ds 2026-10-20 --end-ds 2026-10-20

Once backfilled, you can query the output table with the usual tooling. The output table name is available programmatically via my_model_transforms.table.

To put it on a recurring schedule, use zipline hub schedule compiled/model_transforms/<team>/<conf> (or let a CI step run zipline hub schedule-all on merges to main). See Testing and Branching and Deployment for the full lifecycle.

Common gotchas

• Model output column names. The column after output_mapping is <cleaned_model_name>__<key>. The cleaned name is the metadata name with . → _, and it already carries the version suffix. If you reference this column in a GroupBy aggregation or a derivation, hard-code the exact name (see the MODEL_OUTPUT_COL constant in the streaming example).
• Empty / null inputs. Many managed embedding backends reject empty strings. Filter the source query with wheres to drop rows where the input fields are null / empty before they reach the model — the listing.v1 example above shows the pattern.
• value_fields is the raw response schema, not the final output schema. It describes what the backend returns per row, before output_mapping is applied. The final output columns are whatever output_mapping keys produce.
• One source per ModelTransforms (today). The thrift schema accepts a list, but the backfill, Flink wiring, and planner all currently require exactly one. This may be relaxed later.
• Streaming inner sources are restricted. For the Flink path the inner source must be an EventSource with a topic, or a JoinSource whose left is an EventSource.
• Dependency semantics are day-for-day. Because ModelTransforms is a row-level transform (no windowing), its backfill depends on the source partition for the same ds — not a window of lookback.

Related reading

• Chaining Features — the broader pattern of using a Join as a source of downstream computation; ModelTransforms is a specialization that adds a model call step.
• Join — how joins are defined; JoinSource is the usual starting point for a ModelTransforms.
• GroupBy — how aggregations downstream of a ModelTransforms are expressed.