From fa05502e4243b785a57c4e257a24c6b0609eebe6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=9E=97SO?=
 <142557582+Linxiushen@users.noreply.github.com>
Date: Mon, 28 Sep 2026 03:32:49 +0800
Subject: [PATCH] fix(automl): include latest observation in time-series CV

---
 flaml/automl/time_series/ts_data.py |  4 +-
 test/automl/test_ts_data.py         | 62 +++++++++++++++++++++++++++++
 2 files changed, 64 insertions(+), 2 deletions(-)

diff --git a/flaml/automl/time_series/ts_data.py b/flaml/automl/time_series/ts_data.py
index 4a650db71c..b1327fe4ad 100644
--- a/flaml/automl/time_series/ts_data.py
+++ b/flaml/automl/time_series/ts_data.py
@@ -210,10 +210,10 @@ def move_validation_boundary(self, steps: int) -> "TimeSeriesDataset":
     def cv_train_val_sets(
         self, n_splits: int, val_length: int, step_size: int
     ) -> Generator["TimeSeriesDataset", None, None]:
-        max_index = len(self.train_data) - 1
+        num_samples = len(self.train_data)
         for i in range(n_splits):
             out = copy.copy(self)
-            val_start = max_index - (n_splits - i - 1) * step_size - val_length
+            val_start = num_samples - (n_splits - i - 1) * step_size - val_length
             out.train_data = self.train_data[:val_start]
             out.test_data = self.train_data[val_start : val_start + val_length]
             yield out
diff --git a/test/automl/test_ts_data.py b/test/automl/test_ts_data.py
index 042eddedf3..5566b610d3 100644
--- a/test/automl/test_ts_data.py
+++ b/test/automl/test_ts_data.py
@@ -1,6 +1,9 @@
 import numpy as np
 import pandas as pd
+import pytest
+from sklearn.model_selection import TimeSeriesSplit
 
+from flaml import AutoML
 from flaml.automl.time_series.ts_data import TimeSeriesDataset, create_forward_frame
 
 
@@ -61,3 +64,62 @@ def test_create_forward_frame_uses_next_frequency_offset():
     pd.testing.assert_series_equal(
         quarterly_frame["ds"], pd.Series(pd.date_range("2020-06-30", periods=2, freq=quarter_end_freq), name="ds")
     )
+
+
+@pytest.mark.parametrize("period", [1, 4])
+@pytest.mark.parametrize("index_kind", ["default", "offset", "datetime"])
+def test_cv_folds_match_time_series_split(period, index_kind):
+    frame = pd.DataFrame({"ds": pd.date_range("2020-01-01", periods=24), "y": np.arange(24)})
+    if index_kind == "offset":
+        frame.index += 100
+    elif index_kind == "datetime":
+        frame.index = frame.ds
+    original = frame.copy(deep=True)
+    dataset = TimeSeriesDataset(frame, time_col="ds", target_names="y")
+    folds = list(dataset.cv_train_val_sets(n_splits=3, val_length=period, step_size=period))
+    reference = list(TimeSeriesSplit(n_splits=3, test_size=period).split(frame))
+
+    assert len(folds) == len(reference)
+    for fold, (train_indices, validation_indices) in zip(folds, reference):
+        pd.testing.assert_frame_equal(fold.train_data, frame.iloc[train_indices])
+        pd.testing.assert_frame_equal(fold.test_data, frame.iloc[validation_indices])
+    pd.testing.assert_frame_equal(dataset.train_data, original)
+    assert dataset.test_data.empty
+
+
+@pytest.mark.parametrize("step_size, expected_starts", [(2, [16, 18, 20]), (7, [6, 13, 20])])
+def test_cv_custom_step_sizes_end_at_latest_observation(step_size, expected_starts):
+    frame = pd.DataFrame({"ds": pd.date_range("2020-01-01", periods=24), "y": np.arange(24)})
+    dataset = TimeSeriesDataset(frame, time_col="ds", target_names="y")
+    folds = list(dataset.cv_train_val_sets(n_splits=3, val_length=4, step_size=step_size))
+
+    assert len(folds) == len(expected_starts)
+    for fold, start in zip(folds, expected_starts):
+        assert fold.test_data.y.tolist() == list(range(start, start + 4))
+        assert fold.train_data.y.tolist() == list(range(start))
+    assert folds[-1].test_data.ds.iloc[-1] == frame.ds.iloc[-1]
+
+
+def test_automl_cv_scores_latest_observation():
+    pytest.importorskip("statsmodels")
+    frame = pd.DataFrame({"ds": pd.date_range("2020-01-01", periods=60), "y": np.zeros(60)})
+    frame.loc[59, "y"] = 100.0
+    automl = AutoML()
+    automl.fit(
+        dataframe=frame,
+        label="y",
+        task="ts_forecast",
+        period=5,
+        estimator_list=["avg"],
+        eval_method="cv",
+        n_splits=3,
+        metric="mae",
+        max_iter=2,
+        time_budget=30,
+        retrain_full=False,
+        verbose=0,
+    )
+
+    # Every training fold contains only zeros. The final validation observation
+    # must contribute its error to the mean across all three five-row folds.
+    assert automl.best_loss == pytest.approx(100.0 / 15)