Skip to content
85 changes: 85 additions & 0 deletions flaml/automl/task/time_series_task.py
Original file line number Diff line number Diff line change
Expand Up @@ -126,7 +126,9 @@ def validate_data(
else:
target_names = label

time_col_inferred = False
if self.time_col is None:
time_col_inferred = True
if isinstance(X_train_all, pd.DataFrame):
assert dataframe is None, "One of dataframe and X arguments must be None"
self.time_col = X_train_all.columns[0]
Expand All @@ -141,6 +143,18 @@ def validate_data(
if X_train_all is not None:
assert y_train_all is not None, "If X_train_all is not None, y_train_all must also be"
assert dataframe is None, "If X_train_all is provided, dataframe must be None"
if isinstance(X_train_all, pd.DataFrame):
X_train_all, promoted_time_col = self._promote_datetime_index_if_needed(
X_train_all,
time_col=self.time_col,
is_inferred=time_col_inferred,
)
if promoted_time_col is not None:
self.time_col = promoted_time_col
if isinstance(y_train_all, (pd.DataFrame, pd.Series)) and isinstance(X_train_all, pd.DataFrame):
if not y_train_all.index.equals(X_train_all.index):
y_train_all = y_train_all.copy()
y_train_all.index = X_train_all.index
dataframe = TimeSeriesDataset.to_dataframe(X_train_all, y_train_all, target_names, self.time_col)

elif dataframe is not None:
Expand All @@ -150,6 +164,14 @@ def validate_data(
else:
raise ValueError("Must supply either X_train_all and y_train_all, or dataframe and label")

dataframe, promoted_time_col = self._promote_datetime_index_if_needed(
dataframe,
time_col=self.time_col,
is_inferred=time_col_inferred,
)
if promoted_time_col is not None:
self.time_col = promoted_time_col

try:
dataframe.loc[:, self.time_col] = pd.to_datetime(dataframe[self.time_col])
except Exception:
Expand All @@ -161,6 +183,16 @@ def validate_data(

if X_val is not None:
assert y_val is not None, "If X_val is not None, y_val must also be"
if isinstance(X_val, pd.DataFrame):
X_val, _ = self._promote_datetime_index_if_needed(
X_val,
time_col=self.time_col,
is_inferred=False,
)
if isinstance(y_val, (pd.DataFrame, pd.Series)) and isinstance(X_val, pd.DataFrame):
if not y_val.index.equals(X_val.index):
y_val = y_val.copy()
y_val.index = X_val.index
val_df = TimeSeriesDataset.to_dataframe(X_val, y_val, target_names, self.time_col)
val_len = len(val_df)
else:
Expand Down Expand Up @@ -385,7 +417,60 @@ def _preprocess(self, X, transformer=None):
X = transformer.transform(X)
return X

@staticmethod
def _promote_datetime_index_if_needed(df, time_col=None, is_inferred=False):
"""Promote a DatetimeIndex to a guaranteed-unique column if needed.

Args:
df: A pandas DataFrame.
time_col: The current timestamp column name (if known).
is_inferred: Whether time_col was inferred rather than explicitly specified.

Returns:
Tuple of (df_processed, promoted_col_name or None).
"""
if not isinstance(df, pd.DataFrame) or not isinstance(df.index, pd.DatetimeIndex):
return df, None

# Precondition check:
# 1. If time_col was inferred: promote if df lacks a datetime-typed time_col.
# 2. If time_col was explicit: promote if time_col is missing from df.columns.
needs_promotion = False
if is_inferred:
if time_col is None or time_col not in df.columns or not pd.api.types.is_datetime64_any_dtype(df[time_col]):
needs_promotion = True
else:
if time_col is not None and time_col not in df.columns:
needs_promotion = True

if not needs_promotion:
return df, None

base_name = df.index.name or "ds"
promoted_col = base_name if not is_inferred and time_col else base_name
if not is_inferred and time_col:
promoted_col = time_col

# Guarantee unique column name avoiding collisions with existing columns
target_col = promoted_col
i = 1
while target_col in df.columns:
target_col = f"{promoted_col}_{i}"
i += 1

df_out = df.copy()
df_out.insert(0, target_col, df.index)
df_out = df_out.reset_index(drop=True)
logger.info(f"Timestamp column not found, promoted DataFrame index as '{target_col}'.")
return df_out, target_col

def preprocess(self, X, transformer=None):
if isinstance(X, pd.DataFrame):
X, _ = self._promote_datetime_index_if_needed(
X,
time_col=self.time_col,
is_inferred=False,
)
if isinstance(X, (pd.DataFrame, np.ndarray, pd.Series)):
X = normalize_ts_data(X.copy(), self.target_names, self.time_col)
return self._preprocess(X, transformer)
Expand Down
3 changes: 3 additions & 0 deletions flaml/automl/time_series/ts_data.py
Original file line number Diff line number Diff line change
Expand Up @@ -549,6 +549,9 @@ def normalize_ts_data(X_train_all, target_names, time_col, y_train_all=None):
elif isinstance(y_train_all, pd.Series):
y_train_all = pd.DataFrame(y_train_all)
y_train_all.index = X_train_all.index
elif isinstance(y_train_all, pd.DataFrame):
y_train_all = y_train_all.copy()
y_train_all.index = X_train_all.index

dataframe = pd.concat([X_train_all, y_train_all], axis=1)

Expand Down
163 changes: 163 additions & 0 deletions test/automl/test_ts_forecast_datetime_index.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,163 @@
"""Regression tests for ts_forecast input where the timestamp is the DataFrame index.

Flaml expects the timestamp to live in a *column* (`time_col`). When a user passes the
idiomatic pandas shape instead -- timestamps as a `DatetimeIndex` -- `validate_data`
defaults `time_col` to `dataframe.columns[0]`, which is a value column (often the label
itself), so the real timestamps are never looked at.
"""

import numpy as np
import pandas as pd
import pytest

from flaml import AutoML


def _make_train_df(periods=120):
return pd.DataFrame(
{"y": np.sin(np.arange(periods) / 6) * 10 + 50},
index=pd.date_range("2018-01-01", periods=periods, freq="MS"),
)


def _fit(automl, **kwargs):
settings = {
"task": "ts_forecast",
"label": "y",
"period": 12,
"time_budget": 5,
"estimator_list": ["lgbm"],
"metric": "mape",
"log_file_name": False,
"verbose": 0,
}
settings.update(kwargs)
automl.fit(**settings)


def test_unnamed_datetime_index_becomes_the_time_column():
automl = AutoML()
_fit(automl, dataframe=_make_train_df())
assert automl._state.task.time_col == "ds"


def test_named_datetime_index_keeps_its_name():
df = _make_train_df()
df.index.name = "month"
automl = AutoML()
_fit(automl, dataframe=df)
assert automl._state.task.time_col == "month"


def test_promoted_index_keeps_every_training_row():
"""Pre-promotion the label column was read as epoch nanoseconds, collapsing the 120 real
timestamps onto a handful of 1970-01-01 values and dropping the duplicates."""
df = _make_train_df()
automl = AutoML()
_fit(automl, dataframe=df)
assert automl._state.data_size[0] == len(df)


def test_explicit_time_col_still_wins_over_the_index():
df = _make_train_df()
df["ts"] = df.index
automl = AutoML()
_fit(automl, dataframe=df, time_col="ts")
assert automl._state.task.time_col == "ts"


def test_genuine_missing_timestamp_still_raises():
df = pd.DataFrame({"note": ["not-a-date"] * 120, "y": np.arange(120, dtype=float)})
automl = AutoML()
with pytest.raises(ValueError) as exc:
_fit(automl, dataframe=df)
assert "must contain timestamp values" in str(exc.value)


def test_timestamp_column_input_is_unchanged():
df = _make_train_df().reset_index().rename(columns={"index": "ds"})
automl = AutoML()
_fit(automl, dataframe=df, time_col="ds")
assert automl._state.task.time_col == "ds"
assert len(automl.predict(df[["ds"]].tail(12))) == 12


def test_index_name_collision_resolves_to_unique_column():
"""When dataframe already has a column named 'ds' (or 'index') that is not datetime,
promoting an unnamed DatetimeIndex must not overwrite or conflict with the existing column."""
df = _make_train_df()
df["ds"] = np.arange(len(df), dtype=float) # Collision candidate
df["ds_1"] = np.arange(len(df), dtype=float) # Double collision candidate
automl = AutoML()
_fit(automl, dataframe=df)
# Target column should be uniquely chosen as ds_2 avoiding collision
assert automl._state.task.time_col == "ds_2"
assert "ds" in automl._feature_names_in_
assert "ds_1" in automl._feature_names_in_
assert "ds_2" in automl._feature_names_in_
assert automl.data_size_full == len(df)


def test_validation_data_with_datetime_index():
"""Verify that validation data passed as DataFrame with DatetimeIndex is properly promoted."""
df_train = _make_train_df(periods=100)
val_index = pd.date_range("2026-05-01", periods=20, freq="MS")
df_val = pd.DataFrame({"y": np.sin(np.arange(100, 120) / 6) * 10 + 50}, index=val_index)
X_val = pd.DataFrame(index=val_index)
y_val = df_val["y"]
automl = AutoML()
_fit(
automl,
dataframe=df_train,
X_val=X_val,
y_val=y_val,
)
assert automl._state.task.time_col == "ds"
assert automl._state.eval_method == "holdout"
assert len(automl.predict(X_val)) == len(val_index)


def test_validation_data_with_dataframe_target_and_datetime_index():
"""Verify that validation data with DataFrame y_val and DatetimeIndex preserves aligned indexes."""
df_train = _make_train_df(periods=100)
val_index = pd.date_range("2026-05-01", periods=20, freq="MS")
df_val = pd.DataFrame({"y": np.sin(np.arange(100, 120) / 6) * 10 + 50}, index=val_index)
X_val = pd.DataFrame(index=val_index)
y_val = df_val[["y"]] # DataFrame target with DatetimeIndex
automl = AutoML()
_fit(
automl,
dataframe=df_train,
X_val=X_val,
y_val=y_val,
)
assert automl._state.task.time_col == "ds"
assert automl._state.eval_method == "holdout"
assert len(automl.predict(X_val)) == len(val_index)


def test_xtrain_ytrain_with_datetime_index():
"""Verify that (X_train, y_train) inputs with DatetimeIndex are promoted and work seamlessly."""
train_idx = pd.date_range("2018-01-01", periods=100, freq="MS")
val_idx = pd.date_range("2026-05-01", periods=20, freq="MS")
X_train = pd.DataFrame({"feat": np.arange(100, dtype=float)}, index=train_idx)
y_train = pd.Series(np.sin(np.arange(100) / 6) * 10 + 50, index=train_idx, name="y")
X_val = pd.DataFrame({"feat": np.arange(100, 120, dtype=float)}, index=val_idx)
y_val = pd.DataFrame({"y": np.sin(np.arange(100, 120) / 6) * 10 + 50}, index=val_idx)
automl = AutoML()
_fit(automl, X_train=X_train, y_train=y_train, X_val=X_val, y_val=y_val)
assert automl._state.task.time_col == "ds"
assert len(automl.predict(X_val)) == 20


def test_predict_with_datetime_index():
"""Verify that future prediction data passed with DatetimeIndex works without explicit time_col."""
df = _make_train_df(periods=120)
automl = AutoML()
_fit(automl, dataframe=df)

# Future prediction input using DatetimeIndex
future_index = pd.date_range("2028-01-01", periods=12, freq="MS")
future_df = pd.DataFrame(index=future_index)
preds = automl.predict(future_df)
assert len(preds) == 12
Loading