fix(automl): treat a DatetimeIndex as the ts_forecast time column

flaml/automl/automl.py

@@ -882,15 +882,25 @@ def predict(
         if estimator is None:
             logger.warning("No estimator is trained. Please run fit with enough budget.")
             return None
+        orig_index = getattr(X, "index", None)
+        if callable(orig_index):
+            orig_index = None
         X = self._state.task.preprocess(X, self._transformer)
         y_pred = estimator.predict(X, **pred_kwargs)
 
         if isinstance(y_pred, np.ndarray) and y_pred.ndim > 1 and isinstance(y_pred, np.ndarray):
             y_pred = y_pred.flatten()
         if self._label_transformer:
-            return self._label_transformer.inverse_transform(Series(y_pred.astype(int)))
-        else:
-            return y_pred
+            y_pred = self._label_transformer.inverse_transform(Series(y_pred.astype(int)))
+        if (
+            orig_index is not None
+            and not self._state.task.is_ts_forecastpanel()
+            and isinstance(y_pred, (Series, DataFrame))
+            and len(y_pred) == len(orig_index)
+        ):
+            y_pred = y_pred.copy()
+            y_pred.index = orig_index
+        return y_pred
 
     def predict_proba(self, X, **pred_kwargs):
         """Predict the probability of each class from features, only works for

flaml/automl/task/time_series_task.py

@@ -1,3 +1,4 @@
+import datetime
 import logging
 import time
 from typing import List
@@ -126,13 +127,17 @@ 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]
+                    if len(X_train_all.columns) > 0:
+                        self.time_col = X_train_all.columns[0]
                 elif dataframe is not None:
                     assert X_train_all is None, "One of dataframe and X arguments must be None"
-                    self.time_col = dataframe.columns[0]
+                    if len(dataframe.columns) > 0:
+                        self.time_col = dataframe.columns[0]
                 else:
                     self.time_col = "ds"
 
@@ -141,17 +146,68 @@ 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):
+                    needs_promotion = self._needs_datetime_index_promotion(
+                        X_train_all,
+                        time_col=self.time_col,
+                        is_inferred=time_col_inferred,
+                        allow_index_fallback=False,
+                    )
+                    X_train_all, y_train_all = self._align_y_to_X_and_promote(
+                        X_train_all,
+                        y_train_all,
+                        needs_promotion=needs_promotion,
+                    )
+                    if needs_promotion:
+                        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,
+                            allow_index_fallback=False,
+                            reserved_cols=target_names,
+                        )
+                        if promoted_time_col is not None:
+                            self.time_col = promoted_time_col
+                            if isinstance(y_train_all, (pd.DataFrame, pd.Series)):
+                                y_train_all = y_train_all.reset_index(drop=True)
+                    if not time_col_inferred and self.time_col not in X_train_all.columns:
+                        raise ValueError(
+                            f"For '{TS_FORECAST}' task, time column {self.time_col} must contain timestamp values."
+                        )
                 dataframe = TimeSeriesDataset.to_dataframe(X_train_all, y_train_all, target_names, self.time_col)
 
             elif dataframe is not None:
                 assert label is not None, "A label or list of labels must be provided."
                 assert isinstance(dataframe, pd.DataFrame), "dataframe must be a pandas DataFrame"
                 assert label in dataframe.columns, f"{label} must a column name in dataframe"
+                needs_promotion = self._needs_datetime_index_promotion(
+                    dataframe,
+                    time_col=self.time_col,
+                    is_inferred=time_col_inferred,
+                    allow_index_fallback=False,
+                )
+                if needs_promotion:
+                    dataframe, promoted_time_col = self._promote_datetime_index_if_needed(
+                        dataframe,
+                        time_col=self.time_col,
+                        is_inferred=time_col_inferred,
+                        allow_index_fallback=False,
+                        reserved_cols=target_names,
+                    )
+                    if promoted_time_col is not None:
+                        self.time_col = promoted_time_col
+                if not time_col_inferred and self.time_col not in dataframe.columns:
+                    raise ValueError(
+                        f"For '{TS_FORECAST}' task, time column {self.time_col} must contain timestamp values."
+                    )
             else:
                 raise ValueError("Must supply either X_train_all and y_train_all, or dataframe and label")
 
+            if self.time_col is None:
+                self.time_col = "ds"
+
             try:
-                dataframe.loc[:, self.time_col] = pd.to_datetime(dataframe[self.time_col])
+                dataframe[self.time_col] = pd.to_datetime(dataframe[self.time_col])
             except Exception:
                 raise ValueError(
                     f"For '{TS_FORECAST}' task, time column {self.time_col} must contain timestamp values."
@@ -161,6 +217,33 @@ 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):
+                    needs_promotion_val = self._needs_datetime_index_promotion(
+                        X_val,
+                        time_col=self.time_col,
+                        is_inferred=False,
+                        allow_index_fallback=True,
+                    )
+                    X_val, y_val = self._align_y_to_X_and_promote(
+                        X_val,
+                        y_val,
+                        needs_promotion=needs_promotion_val,
+                    )
+                    if needs_promotion_val:
+                        X_val, promoted_time_col_val = self._promote_datetime_index_if_needed(
+                            X_val,
+                            time_col=self.time_col,
+                            is_inferred=False,
+                            allow_index_fallback=True,
+                            reserved_cols=target_names,
+                        )
+                        if promoted_time_col_val is not None:
+                            if isinstance(y_val, (pd.DataFrame, pd.Series)):
+                                y_val = y_val.reset_index(drop=True)
+                    if self.time_col not in X_val.columns:
+                        raise ValueError(
+                            f"For '{TS_FORECAST}' task, time column {self.time_col} must contain timestamp values."
+                        )
                 val_df = TimeSeriesDataset.to_dataframe(X_val, y_val, target_names, self.time_col)
                 val_len = len(val_df)
             else:
@@ -366,9 +449,11 @@ def _preprocess(self, X, transformer=None):
                 X = pd.DataFrame(
                     dict(
                         [
-                            (transformer._str_columns[idx], X[idx])
-                            if isinstance(X[0], List)
-                            else (transformer._str_columns[idx], [X[idx]])
+                            (
+                                (transformer._str_columns[idx], X[idx])
+                                if isinstance(X[0], List)
+                                else (transformer._str_columns[idx], [X[idx]])
+                            )
                             for idx in range(len(X))
                         ]
                     )
@@ -385,10 +470,164 @@ def _preprocess(self, X, transformer=None):
             X = transformer.transform(X)
         return X
 
+    @staticmethod
+    def _is_valid_timestamp_col(df, col):
+        """Check if col in df is a valid datetime-like column without numeric nanosecond coercion."""
+        if col is None or not isinstance(df, pd.DataFrame) or col not in df.columns:
+            return False
+        series = df[col]
+        if pd.api.types.is_datetime64_any_dtype(series):
+            return True
+        if pd.api.types.is_numeric_dtype(series) or pd.api.types.is_bool_dtype(series):
+            return False
+        non_null = series.dropna()
+        if non_null.empty:
+            return False
+        first_val = non_null.iloc[0]
+        if not isinstance(first_val, (str, datetime.date, datetime.datetime, pd.Timestamp, np.datetime64)):
+            return False
+        try:
+            import warnings
+
+            with warnings.catch_warnings():
+                warnings.simplefilter("ignore")
+                converted = pd.to_datetime(series, errors="raise")
+            return not converted.isna().all()
+        except Exception:
+            return False
+
+    @classmethod
+    def _needs_datetime_index_promotion(cls, df, time_col=None, is_inferred=False, allow_index_fallback=False):
+        """Determine whether df's DatetimeIndex needs to be promoted to a column."""
+        if not isinstance(df, pd.DataFrame) or not isinstance(df.index, pd.DatetimeIndex):
+            return False
+        if allow_index_fallback:
+            return time_col is not None and time_col not in df.columns
+        if is_inferred:
+            if time_col is not None and cls._is_valid_timestamp_col(df, time_col):
+                return False
+            return True
+        return False
+
+    @classmethod
+    def _align_y_to_X_and_promote(cls, X, y, needs_promotion=False):
+        """Align y to X before index promotion or dataset construction.
+
+        When needs_promotion is True, X's DatetimeIndex will be promoted to the time column.
+        In that case, if y also has a DatetimeIndex, validate that index labels match and
+        align y to X by label. If y has a RangeIndex (or is a Series without DatetimeIndex),
+        pair y with X positionally.
+
+        When needs_promotion is False (e.g. X already has a valid timestamp column),
+        preserve the baseline positional contract for Series (y.index = X.index).
+        For DataFrame targets, align by label if unique labels match, or raise ValueError.
+        """
+        if not isinstance(X, pd.DataFrame) or not isinstance(y, (pd.DataFrame, pd.Series)):
+            return X, y
+
+        if len(X) != len(y):
+            raise ValueError(f"Target length ({len(y)}) does not match feature length ({len(X)}).")
+
+        # When X's DatetimeIndex is being promoted to the time column:
+        if needs_promotion:
+            if isinstance(X.index, pd.DatetimeIndex) and isinstance(y.index, pd.DatetimeIndex):
+                if not y.index.equals(X.index):
+                    if not X.index.is_unique or not y.index.is_unique:
+                        raise ValueError("Duplicate index labels are not supported.")
+                    try:
+                        if not X.index.difference(y.index).empty or not y.index.difference(X.index).empty:
+                            raise ValueError("Target index labels do not match feature index labels.")
+                    except TypeError:
+                        raise ValueError("Target index labels do not match feature index labels.")
+                    y = y.reindex(X.index)
+                return X, y
+            # Positional pairing when promotion is needed and y has RangeIndex / non-DatetimeIndex
+            y = y.copy()
+            y.index = X.index
+            return X, y
+
+        # When promotion is NOT needed (e.g. X already has a timestamp column):
+        if isinstance(y, pd.Series):
+            y = y.copy()
+            y.index = X.index
+            return X, y
+        elif isinstance(y, pd.DataFrame):
+            if not y.index.equals(X.index):
+                if not X.index.is_unique or not y.index.is_unique:
+                    raise ValueError("Duplicate index labels are not supported.")
+                try:
+                    if not X.index.difference(y.index).empty or not y.index.difference(X.index).empty:
+                        raise ValueError("Target index labels do not match feature index labels.")
+                except TypeError:
+                    raise ValueError("Target index labels do not match feature index labels.")
+                y = y.reindex(X.index)
+            return X, y
+
+        return X, y
+
+    @classmethod
+    def _promote_datetime_index_if_needed(
+        cls, df, time_col=None, is_inferred=False, allow_index_fallback=False, reserved_cols=None
+    ):
+        """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.
+            allow_index_fallback: Whether to allow falling back to df.index when time_col is absent.
+            reserved_cols: Additional column names to avoid colliding with (e.g. target names).
+
+        Returns:
+            Tuple of (df_processed, promoted_col_name or None).
+        """
+        if not cls._needs_datetime_index_promotion(
+            df,
+            time_col=time_col,
+            is_inferred=is_inferred,
+            allow_index_fallback=allow_index_fallback,
+        ):
+            return df, None
+
+        if is_inferred:
+            base_name = "ds" if df.index.name is None else df.index.name
+            promoted_col = base_name
+            # Guarantee unique column name avoiding collisions with existing columns and reserved names
+            reserved = set(reserved_cols) if reserved_cols else set()
+            target_col = promoted_col
+            i = 1
+            while target_col in df.columns or target_col in reserved:
+                target_col = f"{promoted_col}_{i}"
+                i += 1
+        else:
+            target_col = time_col
+
+        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):
+        orig_index = getattr(X, "index", None)
+        if isinstance(X, pd.DataFrame):
+            X, _ = self._promote_datetime_index_if_needed(
+                X,
+                time_col=self.time_col,
+                is_inferred=False,
+                allow_index_fallback=True,
+                reserved_cols=self.target_names,
+            )
+            if self.time_col in X.columns and not pd.api.types.is_datetime64_any_dtype(X[self.time_col]):
+                if self._is_valid_timestamp_col(X, self.time_col):
+                    X = X.copy()
+                    X[self.time_col] = pd.to_datetime(X[self.time_col])
         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)
+            X = self._preprocess(X, transformer)
+            if orig_index is not None and isinstance(X, pd.DataFrame) and len(X) == len(orig_index):
+                X.index = orig_index
+            return X
         elif isinstance(X, int):
             return X
         else:

flaml/automl/time_series/ts_data.py

@@ -539,6 +539,8 @@ def normalize_ts_data(X_train_all, target_names, time_col, y_train_all=None):
     if y_train_all is None:
         return X_train_all
     else:
+        if len(y_train_all) != len(X_train_all):
+            raise ValueError(f"Target length ({len(y_train_all)}) does not match feature length ({len(X_train_all)}).")
         if isinstance(y_train_all, np.ndarray):
             # TODO: will need to revisit this when doing multivariate y
             y_train_all = pd.DataFrame(
@@ -549,6 +551,19 @@ 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):
+            if not y_train_all.index.equals(X_train_all.index):
+                if not X_train_all.index.is_unique or not y_train_all.index.is_unique:
+                    raise ValueError("Duplicate index labels are not supported.")
+                try:
+                    if (
+                        not X_train_all.index.difference(y_train_all.index).empty
+                        or not y_train_all.index.difference(X_train_all.index).empty
+                    ):
+                        raise ValueError("Target index labels do not match feature index labels.")
+                except TypeError:
+                    raise ValueError("Target index labels do not match feature index labels.")
+                y_train_all = y_train_all.reindex(X_train_all.index)
 
         dataframe = pd.concat([X_train_all, y_train_all], axis=1)