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

flaml/automl/task/time_series_task.py

@@ -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]
@@ -141,6 +143,17 @@ 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, y_train_all = self._align_y_to_X_and_promote(X_train_all, y_train_all)
+                    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)):
+                            y_train_all = y_train_all.reset_index(drop=True)
                 dataframe = TimeSeriesDataset.to_dataframe(X_train_all, y_train_all, target_names, self.time_col)
 
             elif dataframe is not None:
@@ -150,6 +163,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:
@@ -161,6 +182,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, y_val = self._align_y_to_X_and_promote(X_val, y_val)
+                    X_val, promoted_time_col_val = self._promote_datetime_index_if_needed(
+                        X_val,
+                        time_col=self.time_col,
+                        is_inferred=False,
+                    )
+                    if promoted_time_col_val is not None:
+                        if isinstance(y_val, (pd.DataFrame, pd.Series)):
+                            y_val = y_val.reset_index(drop=True)
                 val_df = TimeSeriesDataset.to_dataframe(X_val, y_val, target_names, self.time_col)
                 val_len = len(val_df)
             else:
@@ -366,9 +397,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,7 +418,98 @@ def _preprocess(self, X, transformer=None):
             X = transformer.transform(X)
         return X
 
+    @staticmethod
+    def _align_y_to_X_and_promote(X, y):
+        """Align y to X by label before index promotion or resetting.
+
+        If X has a DatetimeIndex and y has a positional RangeIndex (or non-DatetimeIndex
+        Series), preserve positional pairing by assigning X.index to y.
+        Otherwise, if y has a DatetimeIndex or both are pandas objects with indexes,
+        align y to X.index by label. If non-matching labels introduce new missing
+        values, raise ValueError.
+        """
+        if not isinstance(X, pd.DataFrame) or not isinstance(y, (pd.DataFrame, pd.Series)):
+            return X, y
+
+        if isinstance(X.index, pd.DatetimeIndex):
+            if isinstance(y, pd.Series) and not isinstance(y.index, pd.DatetimeIndex):
+                if len(y) == len(X):
+                    y = y.copy()
+                    y.index = X.index
+                    return X, y
+            elif isinstance(y, pd.DataFrame) and isinstance(y.index, pd.RangeIndex):
+                if len(y) == len(X):
+                    y = y.copy()
+                    y.index = X.index
+                    return X, y
+
+        if not y.index.equals(X.index):
+            y_aligned = y.reindex(X.index)
+            # If reindexing introduced NaNs that were not originally present, indexes do not match
+            orig_nan_count = int(y.isna().sum().sum() if isinstance(y, pd.DataFrame) else y.isna().sum())
+            new_nan_count = int(
+                y_aligned.isna().sum().sum() if isinstance(y_aligned, pd.DataFrame) else y_aligned.isna().sum()
+            )
+            if new_nan_count > orig_nan_count:
+                raise ValueError("Target index labels do not match feature index labels.")
+            y = y_aligned
+
+        return X, y
+
+    @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)

flaml/automl/time_series/ts_data.py

@@ -547,8 +547,29 @@ def normalize_ts_data(X_train_all, target_names, time_col, y_train_all=None):
                 index=X_train_all.index,
             )
         elif isinstance(y_train_all, pd.Series):
-            y_train_all = pd.DataFrame(y_train_all)
-            y_train_all.index = X_train_all.index
+            if isinstance(y_train_all.index, pd.DatetimeIndex):
+                if not y_train_all.index.equals(X_train_all.index):
+                    y_aligned = y_train_all.reindex(X_train_all.index)
+                    if y_aligned.isna().sum() > y_train_all.isna().sum():
+                        raise ValueError("Target index labels do not match feature index labels.")
+                    y_train_all = y_aligned
+                y_train_all = pd.DataFrame(y_train_all)
+            else:
+                y_train_all = pd.DataFrame(y_train_all)
+                y_train_all.index = X_train_all.index
+        elif isinstance(y_train_all, pd.DataFrame):
+            if (
+                isinstance(X_train_all.index, pd.DatetimeIndex)
+                and isinstance(y_train_all.index, pd.RangeIndex)
+                and len(y_train_all) == len(X_train_all)
+            ):
+                y_train_all = y_train_all.copy()
+                y_train_all.index = X_train_all.index
+            elif not y_train_all.index.equals(X_train_all.index):
+                y_aligned = y_train_all.reindex(X_train_all.index)
+                if y_aligned.isna().sum().sum() > y_train_all.isna().sum().sum():
+                    raise ValueError("Target index labels do not match feature index labels.")
+                y_train_all = y_aligned
 
         dataframe = pd.concat([X_train_all, y_train_all], axis=1)