fix: DailyBatchSampler yields per-day cross-sections instead of contiguous slices

qlib/contrib/model/pytorch_gats_ts.py

@@ -26,19 +26,24 @@
 class DailyBatchSampler(Sampler):
     def __init__(self, data_source):
         self.data_source = data_source
-        # calculate number of samples in each batch
-        self.daily_count = (
-            pd.Series(index=self.data_source.get_index()).groupby("datetime", group_keys=False).size().values
-        )
-        self.daily_index = np.roll(np.cumsum(self.daily_count), 1)  # calculate begin index of each batch
-        self.daily_index[0] = 0
+        # TSDataSampler rows are physically instrument-major (<instrument, datetime>).
+        # get_index() swaps the LABEL order only — the rows of one trading day are
+        # NOT contiguous in row space. Collect each day's actual row positions
+        # instead of slicing contiguous ranges (see microsoft/qlib#2319).
+        index = self.data_source.get_index()
+        positions = pd.Series(np.arange(len(index)), index=index)
+        self.daily_batches = [group.to_numpy() for _, group in positions.groupby(level=0, sort=True)]
+        # physical row numbers in iteration (day-major) order; consumers that
+        # produce one value per batch row can re-align with the original
+        # instrument-major layout via `get_index()[index_order]`
+        self.index_order = np.concatenate(self.daily_batches) if self.daily_batches else np.array([], dtype=int)
 
     def __iter__(self):
-        for idx, count in zip(self.daily_index, self.daily_count):
-            yield np.arange(idx, idx + count)
+        for batch in self.daily_batches:
+            yield batch
 
     def __len__(self):
-        return len(self.data_source)
+        return len(self.daily_batches)
 
 
 class GATs(Model):
@@ -332,7 +337,7 @@ def predict(self, dataset):
 
             preds.append(pred)
 
-        return pd.Series(np.concatenate(preds), index=dl_test.get_index())
+        return pd.Series(np.concatenate(preds), index=dl_test.get_index()[sampler_test.index_order])
 
 
 class GATModel(nn.Module):

qlib/data/dataset/__init__.py

@@ -478,6 +478,11 @@ def get_index(self):
         """
         Get the pandas index of the data, it will be useful in following scenarios
         - Special sampler will be used (e.g. user want to sample day by day)
+
+        Note: this swaps the level order of the labels only (to
+        <datetime, instrument>); the underlying rows remain physically sorted
+        <instrument, datetime>. Callers that assume rows of one datetime are
+        contiguous will get wrong results (see microsoft/qlib#2319).
         """
         return self.data_index.swaplevel()  # to align the order of multiple index of original data received by __init__
 

tests/model/test_daily_batch_sampler.py

@@ -0,0 +1,86 @@
+"""Regression tests for DailyBatchSampler (microsoft/qlib#2319).
+
+TSDataSampler stores rows instrument-major (<instrument, datetime>);
+get_index() swaps the label order only. DailyBatchSampler used to slice
+contiguous ranges by per-day counts, which produced cross-day batches of a
+single instrument. Each yielded batch must be one trading day's full
+cross-section instead.
+"""
+
+import unittest
+
+import numpy as np
+import pandas as pd
+
+from qlib.contrib.model.pytorch_gats_ts import DailyBatchSampler
+
+
+class _FakeDataSource:
+    """Mirror of TSDataSampler.get_index(): stores rows instrument-major
+    (<instrument, datetime>) and returns the SWAPPED label view
+    (<datetime, instrument>) without reordering the rows.
+    """
+
+    def __init__(self, data_index: pd.MultiIndex):
+        self._data_index = data_index
+
+    def get_index(self) -> pd.MultiIndex:
+        return self._data_index.swaplevel()
+
+
+def _instrument_major_index(instruments, dates) -> pd.MultiIndex:
+    """Rows physically sorted <instrument, datetime>, like TSDataSampler.data_index."""
+    rows = [(inst, day) for inst in instruments for day in dates]
+    return pd.MultiIndex.from_tuples(rows, names=["instrument", "datetime"])
+
+
+class TestDailyBatchSampler(unittest.TestCase):
+    def test_batches_are_per_day_cross_sections(self):
+        instruments = ["SH600000", "SH600008", "SH600009"]
+        dates = list(pd.date_range("2026-01-05", periods=4, freq="B"))
+        index = _instrument_major_index(instruments, dates)
+
+        sampler = DailyBatchSampler(_FakeDataSource(index))
+        batches = list(iter(sampler))
+
+        # one batch per trading day
+        self.assertEqual(len(batches), len(dates))
+
+        covered = []
+        for batch in batches:
+            labels = index[batch]
+            datetimes = {day for _, day in labels}
+            # every batch covers exactly one trading day ...
+            self.assertEqual(len(datetimes), 1)
+            # ... and that day's full instrument cross-section
+            self.assertEqual(len(labels), len(instruments))
+            covered.extend(batch.tolist())
+
+        # every physical row is yielded exactly once
+        self.assertEqual(sorted(covered), list(range(len(index))))
+
+    def test_batches_follow_chronological_day_order(self):
+        instruments = ["SH600000", "SH600008"]
+        dates = list(pd.date_range("2026-01-05", periods=3, freq="B"))
+        index = _instrument_major_index(instruments, dates)
+
+        sampler = DailyBatchSampler(_FakeDataSource(index))
+        batches = list(iter(sampler))
+
+        first_days = []
+        for batch in batches:
+            labels = index[batch]
+            first_days.append(min(day for _, day in labels))
+        self.assertEqual(first_days, dates)
+
+    def test_len_matches_number_of_batches(self):
+        instruments = ["SH600000", "SH600008"]
+        dates = list(pd.date_range("2026-01-05", periods=3, freq="B"))
+        index = _instrument_major_index(instruments, dates)
+
+        sampler = DailyBatchSampler(_FakeDataSource(index))
+        self.assertEqual(len(sampler), len(dates))
+
+
+if __name__ == "__main__":
+    unittest.main()