xpark.dataset.TimeSeriesForecast#

class xpark.dataset.TimeSeriesForecast(_local_model: str = 'amazon/chronos-2', /, *, prediction_length: int, quantile_levels: tuple[float, ...] | list[float] = (0.1, 0.5, 0.9), context_length: int | None = None, cross_learning: bool = False, limit_prediction_length: bool = True, freq: str | None = None)[source]#

Time-series forecasting processor backed by Chronos-2.

The processor treats one row as one complete time series. The target column is expected to be a list<float> containing the historical observations of a single series (already sorted by timestamp and equally spaced). The second positional argument is the timestamps column (list<timestamp>) carrying each row’s historical timestamps; the processor uses it to derive the future timestamps returned alongside the quantile predictions. Any number of additional feature columns can be passed as further positional arguments; each must also be a list<float> of the same length as target and is forwarded to Chronos-2 as a past-only covariate. Future covariates are intentionally not exposed here.

Parameters:
  • _local_model – Chronos-2 model id. Available models: [‘amazon/chronos-2’]

  • prediction_length – Number of future steps to predict for each series.

  • quantile_levels

    Quantile levels to return. Must not contain duplicates (otherwise the resulting field names would collide and the constructor raises ValueError). The output column is a struct containing a leading "timestamp" field of type list<timestamp[ns]> (the future timestamps inferred from the timestamps input column) followed by one list<float32> field per requested quantile level; every inner list has length prediction_length.

    Quantile field naming rule. Each quantile field name is produced by Python’s built-in format(q, ".10g"). For typical usage where q has at most 3 decimal digits (between 0.001 and 0.999), this is exactly the literal you would write yourself, so just wrap the number in quotes when indexing the output, e.g. row["forecast"]["0.1"], row["forecast"]["0.5"], row["forecast"]["0.975"].

  • context_length – Maximum history length used during inference. Defaults to the model’s built-in context length (typically 2048). Histories longer than context_length are left-truncated by Chronos-2.

  • cross_learning – Whether to enable cross-learning mode in Chronos-2. Recommended only when individual histories are very short.

  • limit_prediction_length – If True (default), raise when prediction_length is greater than the model’s default prediction length.

  • freq – Optional frequency string (e.g. "h", "D", "W") used to generate future timestamps. When None (default) the frequency is inferred per row; pass an explicit freq to skip inference (faster, and required for series shorter than 3 points).

Examples

from xpark.dataset import TimeSeriesForecast, from_items
from xpark.dataset.expressions import col

# One row per series; ``target`` is the historical observation list,
# ``ts`` is the corresponding timestamp list (equally spaced).
ds = from_items([
    {
        "id": "a",
        "target": [1.0, 2.0, 3.0, 4.0, 5.0],
        "ts": pd.date_range("2026-01-01", periods=5, freq="D").tolist(),
    },
])
ds = ds.with_column(
    "forecast",
    TimeSeriesForecast(prediction_length=3, quantile_levels=[0.1, 0.5, 0.9])
    .options(num_workers={"GPU": 1})
    .with_column(col("target"), col("ts")),
)
# row["forecast"] -> {"timestamp": [2026-01-06, 2026-01-07, 2026-01-08],
#                     "0.1": [...], "0.5": [...], "0.9": [...]}
print(ds.take_all())

With past-only covariates (each covariate is a separate list<float> column with the same length as target); covariates follow the timestamps column positionally:

ds = ds.with_column(
    "forecast",
    TimeSeriesForecast(prediction_length=7)
    .with_column(col("target"), col("ts"), col("humidity"), col("wind_speed")),
)

Methods

__call__(target, timestamps, *features)

Run Chronos-2 inference for a batch of rows.

options(**kwargs)

with_column(target, timestamps, *features)

Run Chronos-2 inference for a batch of rows.

__call__(target: pa.ChunkedArray, timestamps: pa.ChunkedArray, *features: pa.ChunkedArray) pa.Array#

Run Chronos-2 inference for a batch of rows.

Parameters:
  • target – the historical observations.

  • timestamps – the historical timestamps.

  • *features – Optional past-only covariate columns.

Returns:

StructArray with a leading "timestamp" field of type list<timestamp[ns]> followed by one list<float32> field per requested quantile level (field names are the quantile literals "0.1", "0.5", …). All inner lists have length prediction_length and follow the user-specified quantile ordering.

options(**kwargs: Unpack[ExprUDFOptions]) Self#
with_column(target: pa.ChunkedArray, timestamps: pa.ChunkedArray, *features: pa.ChunkedArray) pa.Array#

Run Chronos-2 inference for a batch of rows.

Parameters:
  • target – the historical observations.

  • timestamps – the historical timestamps.

  • *features – Optional past-only covariate columns.

Returns:

StructArray with a leading "timestamp" field of type list<timestamp[ns]> followed by one list<float32> field per requested quantile level (field names are the quantile literals "0.1", "0.5", …). All inner lists have length prediction_length and follow the user-specified quantile ordering.