from __future__ import annotations
import logging
from string import Template
from typing import TYPE_CHECKING, Any
from xpark.dataset.datatype import DataType
from xpark.dataset.expressions import BatchColumnClassProtocol, udf
from xpark.dataset.import_utils import lazy_import
from xpark.dataset.model import ModelSpec, cache_model
if TYPE_CHECKING:
import numpy as np
import pandas as pd
import pyarrow as pa
import torch
from chronos.chronos2 import Chronos2Pipeline
else:
np = lazy_import("numpy", rename="np")
pa = lazy_import("pyarrow", rename="pa")
pd = lazy_import("pandas", rename="pd")
torch = lazy_import("torch")
logger = logging.getLogger("ray")
ForecastModel = {
"amazon/chronos-2": {
"label": {"test", "all"},
"model_specs": {
"huggingface": {
"pytorch": {
"model_id": "amazon/chronos-2",
"model_revision": "0f8a440441931157957e2be1a9bce66627d99c76",
"quantizations": [None],
},
},
"modelscope": {
"pytorch": {
"model_id": "amazon/chronos-2",
"quantizations": [None],
},
},
},
},
}
class ForecastModelSpec(ModelSpec):
pass
ALL_FORECAST_MODELS = {k: ForecastModelSpec.model_validate(v) for k, v in ForecastModel.items()}
AVAILABLE_MODELS = [k for k, v in ALL_FORECAST_MODELS.items() if v.label & {"all"}]
def _to_float32_array(values, *, where: str) -> "np.ndarray":
"""Convert a Python list (decoded from a ``list<float>`` cell) to a 1-D ``float32`` array."""
if values is None:
raise ValueError(f"forecast input is NULL at {where}; expected a non-null list of numbers")
arr = np.asarray(values, dtype=np.float32)
if arr.ndim != 1:
raise ValueError(f"forecast input at {where} must be 1-D, got shape {arr.shape}")
return arr
_FORECAST_RETURN_DTYPE = DataType.from_arrow(
pa.struct(
[
pa.field("timestamp", pa.list_(pa.timestamp("ns"))),
pa.field("0.1", pa.list_(pa.float32())),
pa.field("0.5", pa.list_(pa.float32())),
pa.field("0.9", pa.list_(pa.float32())),
]
)
)
[docs]
@udf(return_dtype=_FORECAST_RETURN_DTYPE)
class TimeSeriesForecast(BatchColumnClassProtocol):
__doc__ = Template("""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.
Args:
_local_model: Chronos-2 model id. Available models: ${AVAILABLE_MODELS}
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:
.. code-block:: python
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:
.. code-block:: python
ds = ds.with_column(
"forecast",
TimeSeriesForecast(prediction_length=7)
.with_column(col("target"), col("ts"), col("humidity"), col("wind_speed")),
)
""").safe_substitute(AVAILABLE_MODELS=AVAILABLE_MODELS)
def __init__(
self,
_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,
):
if _local_model not in ALL_FORECAST_MODELS:
raise ValueError(f"Unsupported forecast model: {_local_model!r}. Available: {AVAILABLE_MODELS}")
if prediction_length <= 0:
raise ValueError(f"prediction_length must be positive, got {prediction_length}")
quantiles = list(quantile_levels)
if not quantiles:
raise ValueError("quantile_levels must not be empty")
for q in quantiles:
if not 0.0 < q < 1.0:
raise ValueError(f"quantile_levels entries must be in (0, 1), got {q}")
self.quantile_levels = quantiles
self.field_names = ["timestamp", *(format(q, ".10g") for q in self.quantile_levels)]
if len(self.field_names) != len(set(self.field_names)):
raise ValueError(f"field_names contains duplicates: {self.field_names}")
self.prediction_length = prediction_length
self.freq = freq
self._predict_kwargs: dict[str, Any] = {
"context_length": context_length,
"cross_learning": cross_learning,
"limit_prediction_length": limit_prediction_length,
}
from chronos.chronos2 import Chronos2Pipeline
model_spec = ALL_FORECAST_MODELS[_local_model]
model_path = cache_model("huggingface", "pytorch", model_spec, None)
logger.info(f"Loading Chronos-2 pipeline from {model_path}")
self.pipeline: Chronos2Pipeline = Chronos2Pipeline.from_pretrained(model_path, device_map="auto")
def _build_future_timestamps(
self,
timestamps: "pa.ChunkedArray",
n_rows: int,
) -> "pa.Array":
"""Derive the per-row future timestamps from each row's history.
Mirrors Chronos-2's ``predict_df`` logic: infer (or use the
user-supplied) frequency, then extend the last historical timestamp by
``[1..prediction_length] * offset``.
Each input row must contain at least 3 timestamps when ``self.freq``
is not provided, otherwise ``pandas.infer_freq`` cannot decide a
frequency. The returned ``pa.Array`` has type
``list<timestamp[ns]>`` with each row of length ``prediction_length``.
"""
if len(timestamps) != n_rows:
raise ValueError(f"timestamps column length {len(timestamps)} does not match target column length {n_rows}")
h = self.prediction_length
explicit_freq = self.freq
future_per_row: list[np.ndarray] = []
for i, history in enumerate(timestamps.to_pylist()):
if history is None:
raise ValueError(f"timestamps row {i} is NULL; expected a list of timestamps")
history_index = pd.DatetimeIndex(pd.to_datetime(list(history)))
if len(history_index) == 0:
raise ValueError(f"timestamps row {i} is empty")
freq = explicit_freq
if freq is None:
if len(history_index) < 3:
raise ValueError(
f"timestamps row {i} has only {len(history_index)} points, "
"need >= 3 to infer frequency or pass `freq=` explicitly"
)
freq = pd.infer_freq(history_index)
if not freq:
raise ValueError(
f"could not infer frequency for timestamps row {i}; "
"ensure timestamps are equally spaced or pass `freq=` explicitly"
)
offset = pd.tseries.frequencies.to_offset(freq)
last_ts = history_index[-1]
# Same recipe as ``df_utils.convert_df_input_to_list_of_dicts_input``
# in chronos-2: stack ``last_ts + step * offset`` for step=1..h.
future_index = pd.DatetimeIndex([last_ts + step * offset for step in range(1, h + 1)])
future_per_row.append(future_index.to_numpy(dtype="datetime64[ns]"))
# ``pa.array`` accepts a list of ``np.ndarray[datetime64[ns]]`` rows
# directly, producing a ``list<timestamp[ns]>`` array.
return pa.array(future_per_row, type=pa.list_(pa.timestamp("ns")))
def _build_inputs(
self,
target: "pa.ChunkedArray",
features: tuple["pa.ChunkedArray", ...],
) -> list:
"""Build the ``inputs`` argument expected by ``predict_quantiles``.
- When no covariates are supplied each row becomes a 1-D ``np.ndarray``
(the fast univariate path of Chronos-2).
- Otherwise each row becomes a dict ``{"target": ..., "past_covariates": {...}}``.
"""
n_rows = len(target)
for idx, feat in enumerate(features):
if len(feat) != n_rows:
raise ValueError(
f"feature column #{idx} length {len(feat)} does not match target column length {n_rows}"
)
targets = target.to_pylist()
if not features:
return [_to_float32_array(values, where=f"target row {i}") for i, values in enumerate(targets)]
feature_lists = [feat.to_pylist() for feat in features]
feature_names = [f"feature_{i}" for i in range(len(features))]
inputs: list[dict] = []
for i, target_values in enumerate(targets):
target_arr = _to_float32_array(target_values, where=f"target row {i}")
past_covariates: dict[str, np.ndarray] = {}
for name, feat_values in zip(feature_names, feature_lists):
cov_arr = _to_float32_array(feat_values[i], where=f"{name} row {i}")
if cov_arr.shape[0] != target_arr.shape[0]:
raise ValueError(
f"past covariate {name!r} length {cov_arr.shape[0]} does not "
f"match target length {target_arr.shape[0]} at row {i}"
)
past_covariates[name] = cov_arr
inputs.append({"target": target_arr, "past_covariates": past_covariates})
return inputs
def __call__(
self,
target: "pa.ChunkedArray",
timestamps: "pa.ChunkedArray",
*features: "pa.ChunkedArray",
) -> "pa.Array":
"""Run Chronos-2 inference for a batch of rows.
Args:
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.
"""
inputs = self._build_inputs(target, features)
n_rows = len(target)
ts_array = self._build_future_timestamps(timestamps, n_rows)
# ``predict_quantiles`` returns ``(quantiles, mean)``. Each element of
# ``quantiles`` has shape ``(n_variates, prediction_length, n_quantiles)``;
# since one row maps to a single univariate series we squeeze the
# leading variate dimension and end up with
# ``[prediction_length, n_quantiles]``.
quantiles, _ = self.pipeline.predict_quantiles(
inputs=inputs,
prediction_length=self.prediction_length,
quantile_levels=self.quantile_levels,
**self._predict_kwargs,
)
n_q = len(self.quantile_levels)
h = self.prediction_length
# Stay in torch all the way through reshaping. Each ``quantiles[i]`` has
# shape ``(n_variates, prediction_length, n_quantiles)``; with our
# one-row-per-series contract ``n_variates == 1`` so we squeeze the
# leading singleton dim and end up with ``(h, n_q)`` per row.
per_row: list[torch.Tensor] = []
for i, q in enumerate(quantiles):
t = q.detach().to("cpu", dtype=torch.float32).squeeze(0)
if t.shape != (h, n_q):
raise RuntimeError(
f"unexpected quantile tensor shape {tuple(q.shape)} for row {i}, "
f"expected leading singleton variate dim followed by ({h}, {n_q})"
)
per_row.append(t)
# Stack per-row tensors of shape ``(h, n_q)`` into ``(n_rows, h, n_q)``,
# then transpose to ``(n_q, n_rows, h)`` so axis 0 indexes the quantile
# level. This lets us slice ``stacked[k]`` to get a ``(n_rows, h)``
# matrix for the k-th quantile -- exactly what each output PyArrow
# ``list<float32>`` column needs.
stacked = torch.stack(per_row, dim=0).permute(2, 0, 1).contiguous()
list_type = pa.list_(pa.float32())
quantile_matrix = stacked.numpy()
field_arrays: list[pa.Array] = [
ts_array,
*(pa.array(list(quantile_matrix[k]), type=list_type) for k in range(n_q)),
]
return pa.StructArray.from_arrays(field_arrays, names=self.field_names)