feat(python): add reshuffle_blocks_per_epoch option for 2-phase reads

Some users want the per-epoch row-level shuffle without blocks being
reassigned to different splits every epoch, so a block's on-disk
location (and thus read locality) stays stable for the whole run.
Defaults to True (matches existing behavior: epoch feeds into both the
block shuffle and the row-level shuffle).

python/python/lancedb/streaming.py

@@ -531,6 +531,15 @@ class StreamingDataset(IterableDataset):
         ``window_blocks * block_size``, so the default window actually
         holds ``window_blocks`` blocks in RAM at once; pass this explicitly
         to decouple the shuffle bound from ``window_blocks``.
+    reshuffle_blocks_per_epoch:
+        Whether ``epoch`` feeds into which blocks are assigned to which
+        split, in addition to the row-level shuffle within each split's
+        blocks (which always varies by epoch).  Requires ``block_size``.
+        ``True`` (the default) reassigns blocks every epoch, matching
+        1-phase's per-epoch row reassignment.  ``False`` keeps each
+        split's blocks fixed across epochs -- only the row order within
+        them changes -- which keeps a block's on-disk location, and thus
+        its read locality, stable across an entire training run.
     filter:
         Optional SQL filter expression (e.g. ``"label = 'dog'"``).  Only rows
         that satisfy the predicate are included in the permutation.  The filter
@@ -656,6 +665,7 @@ class StreamingDataset(IterableDataset):
         block_size: Optional[int] = None,
         window_blocks: Optional[int] = None,
         max_shuffle_distance: Optional[int] = None,
+        reshuffle_blocks_per_epoch: Optional[bool] = None,
         filter: Optional[str] = None,
         transform: Optional[Callable] = None,
         transform_parallelism: Optional[int] = None,
@@ -697,6 +707,10 @@ class StreamingDataset(IterableDataset):
                 raise ValueError(
                     "max_shuffle_distance requires block_size to be set"
                 )
+            if reshuffle_blocks_per_epoch is not None:
+                raise ValueError(
+                    "reshuffle_blocks_per_epoch requires block_size to be set"
+                )
         else:
             if window_blocks is not None and window_blocks < 1:
                 raise ValueError("window_blocks must be at least 1")
@@ -928,9 +942,23 @@ class StreamingDataset(IterableDataset):
                     f"block(s) from {num_rows} row(s), fewer than num_splits "
                     f"({num_splits}); use a smaller block_size or fewer splits"
                 )
+            # Whether the epoch feeds into the block shuffle (so which
+            # blocks a split owns changes every epoch, same as 1-phase's
+            # row assignment) or only into the row-level shuffle within
+            # each split's fixed blocks (so blocks stay pinned to the same
+            # split/location across epochs, and only their row order
+            # varies).  Defaults to True: matches the pre-existing
+            # behavior of using epoch for both.
+            self._reshuffle_blocks_per_epoch = (
+                reshuffle_blocks_per_epoch
+                if reshuffle_blocks_per_epoch is not None
+                else True
+            )
             block_order = list(range(self._num_blocks))
             if shuffle:
-                block_seed = shuffle_seed + epoch * _EPOCH_PRIME
+                block_seed = shuffle_seed
+                if self._reshuffle_blocks_per_epoch:
+                    block_seed += epoch * _EPOCH_PRIME
                 random.Random(block_seed).shuffle(block_order)
             # Kept as a plain Python list for simplicity, unlike the
             # 1-phase row mapping below, which is Arrow-backed because it

python/python/tests/test_elastic_dataloader.py

@@ -3366,6 +3366,30 @@ def test_two_phase_resume_checkpoint_round_trip(lance_table):
     )
 
 
+def test_two_phase_reshuffle_blocks_per_epoch_toggle(lance_table):
+    """reshuffle_blocks_per_epoch=False keeps each split's blocks fixed
+    across epochs; the default (True) reassigns them every epoch."""
+    block_size = 10  # NUM_ROWS=120 -> 12 blocks
+
+    def block_perm_for(epoch, reshuffle):
+        ds = StreamingDataset(
+            lance_table,
+            num_splits=1,
+            shuffle_seed=SHUFFLE_SEED,
+            block_size=block_size,
+            epoch=epoch,
+            reshuffle_blocks_per_epoch=reshuffle,
+        )
+        return list(ds._block_perm)
+
+    assert block_perm_for(0, False) == block_perm_for(1, False), (
+        "blocks must stay pinned to the same split across epochs when disabled"
+    )
+    assert block_perm_for(0, True) != block_perm_for(1, True), (
+        "blocks should be reassigned across epochs by default"
+    )
+
+
 def test_two_phase_too_few_blocks_for_num_splits_raises(tmp_path):
     """Fewer blocks than num_splits would leave some splits with no data."""
     db = lancedb.connect(tmp_path)
@@ -3386,6 +3410,10 @@ def test_two_phase_too_few_blocks_for_num_splits_raises(tmp_path):
         pytest.param(
             {"max_shuffle_distance": 5}, id="max_shuffle_distance_without_block_size"
         ),
+        pytest.param(
+            {"reshuffle_blocks_per_epoch": True},
+            id="reshuffle_blocks_per_epoch_without_block_size",
+        ),
     ],
 )
 def test_two_phase_constructor_validation_raises(lance_table, extra_kwargs):