Fix batch(by_column=...) crashing after shard/shuffle/split (#8259)

MappedExamplesIterable.shuffle_data_sources, shard_data_sources and
reshard_data_sources did not forward is_batch_accumulate_arrow_table_function,
so any re-created iterable (DataLoader workers, .shard(), .shuffle(),
split_dataset_by_node) stopped injecting the tables_accumulator/length kwargs
and the accumulate function crashed:

    TypeError: _batch_accumulate_arrow_table_by_columns() missing 2 required
    positional arguments: 'tables_accumulator' and 'length'

Forward the flag in all three methods and add a regression test covering
shard() and shuffle() on a by_column-batched stream.

Co-authored-by: Pepijn <pepijn1999kooijmans@gmail.com>

Author: pkooij

src/datasets/iterable_dataset.py

@@ -1723,6 +1723,7 @@ def shuffle_data_sources(self, generator: np.random.Generator) -> "MappedExample
             formatting=self.formatting,
             features=self.features,
             max_num_running_async_map_functions_in_parallel=self.max_num_running_async_map_functions_in_parallel,
+            is_batch_accumulate_arrow_table_function=self.is_batch_accumulate_arrow_table_function,
         )
 
     def shard_data_sources(self, num_shards: int, index: int, contiguous=True) -> "MappedExamplesIterable":
@@ -1740,6 +1741,7 @@ def shard_data_sources(self, num_shards: int, index: int, contiguous=True) -> "M
             formatting=self.formatting,
             features=self.features,
             max_num_running_async_map_functions_in_parallel=self.max_num_running_async_map_functions_in_parallel,
+            is_batch_accumulate_arrow_table_function=self.is_batch_accumulate_arrow_table_function,
         )
 
     def reshard_data_sources(self) -> "MappedExamplesIterable":
@@ -1756,6 +1758,7 @@ def reshard_data_sources(self) -> "MappedExamplesIterable":
             formatting=self.formatting,
             features=self.features,
             max_num_running_async_map_functions_in_parallel=self.max_num_running_async_map_functions_in_parallel,
+            is_batch_accumulate_arrow_table_function=self.is_batch_accumulate_arrow_table_function,
         )
 
     @property

tests/test_iterable_dataset.py

@@ -2897,6 +2897,25 @@ def test_iterable_dataset_batch(num_shards: int):
         assert batch["text"] == [f"Text {3 * i}", f"Text {3 * i + 1}", f"Text {3 * i + 2}"]
 
 
+def test_iterable_dataset_batch_by_column_survives_resharding():
+    # Re-creating the iterable (shard / shuffle / split_by_node, e.g. inside torch DataLoader
+    # workers) must keep accumulating whole groups instead of crashing with a missing
+    # tables_accumulator argument (regression test).
+    data = {
+        "id": list(range(10)),
+        "category": ["A"] * 5 + ["B"] * 5,
+    }
+    ds = IterableDataset.from_dict(data, num_shards=2)
+    batched_ds = ds.batch(by_column="category")
+
+    sharded = [batch["category"][0] for i in range(2) for batch in batched_ds.shard(num_shards=2, index=i)]
+    assert sorted(sharded) == ["A", "B"]
+
+    shuffled = list(batched_ds.shuffle(seed=0, buffer_size=2))
+    assert sorted(batch["category"][0] for batch in shuffled) == ["A", "B"]
+    assert all(len(set(batch["category"])) == 1 for batch in shuffled)
+
+
 @pytest.mark.parametrize("num_shards", [1, 2, 3, 7, 10])
 def test_iterable_dataset_batch_by_column(num_shards: int):
     # Create a Dataset with a column to group by