[miladh]

Stacked on #2875.

Description

DataFrame variants of forecast and impute, following the insert_dataframe / retrieve_dataframe convention:

• Each numeric column is one time series, labelled by its column name.
• NaN values are marked missing: hidden from the model, and for impute_dataframe, predicted.
• cohort= takes one cohort for every column, or a {column: cohort} mapping. Columns not in the mapping are processed independently.
• quantile_levels selects the quantile levels in the result: the median (0.5) by default, a list of levels, or None for every level the API returns.
• The result has a DatetimeIndex and a (column, quantile level) column for every column and included quantile level. It keeps the input's column names and timezone; a naive index is taken to be in UTC, as in insert_dataframe.
• A DataFrame without a DatetimeIndex, with repeated or non-numeric columns or infinite values, or a cohort mapping that names a column not in the DataFrame, is rejected before anything is sent.

Usage

import numpy as np
import pandas as pd
from cognite.client import CogniteClient

client = CogniteClient()
df = pd.DataFrame(
    {
        "23-PT-1101": [42.1, 42.4, 42.9],
        "23-PT-1201": [39.8, 40.1, 40.0],
        "24-TT-3001": [31.0, np.nan, 30.9],   # NaN is marked missing
    },
    index=pd.date_range("2026-10-01", periods=3, freq="1min"),
)

# Every column independently
forecast = client.ai.time_series.forecast_dataframe(df)

# The two compressor pressures jointly; 24-TT-3001 stays independent
forecast = client.ai.time_series.forecast_dataframe(
    df, cohort={"23-PT-1101": "compression-train-a", "23-PT-1201": "compression-train-a"}
)
forecast["23-PT-1101"][0.5]   # the median forecast, the default
all_levels = client.ai.time_series.forecast_dataframe(df, quantile_levels=None)

# Predict the NaNs, and fill them with the predicted median
imputed = client.ai.time_series.impute_dataframe(df)
filled = df.fillna(imputed.xs(0.5, axis=1, level=1))

pandas is an optional extra, so the new tests are marked dsl. The docstrings note that the endpoints are in beta. The ai.time_series docstring test moves to its own dsl-marked test, like datapoints and raw.

Checklist:

• Tests added/updated: TestDataFrames in test_api/test_ai_time_series.py.
• Documentation updated: docstrings.
• The PR title follows the Conventional Commit spec.

🤖 Generated with Claude Code

[gemini-code-assist]

Code Review

This pull request introduces forecast_dataframe and impute_dataframe methods to both the async and sync AI Time Series APIs, allowing users to forecast and reconstruct missing values directly from pandas DataFrames. Feedback on the implementation of the _dataframe_to_input helper function highlights the need for improved robustness, specifically regarding DatetimeIndex validation, timezone-aware index handling, and validating that all columns are numeric to prevent cryptic runtime errors.