Fix filter indices when batched (#5113)

* Test filter indices

* Fix filter indices when batched

* Rename test

Author: albertvillanova

src/datasets/arrow_dataset.py

@@ -2961,7 +2961,7 @@ def init_buffer_and_writer():
                             else:
                                 writer.write(example)
                 else:
-                    for i, batch in enumerate(pbar):
+                    for i, batch in zip(range(0, num_rows, batch_size), pbar):
                         indices = list(
                             range(*(slice(i, i + batch_size).indices(input_dataset.num_rows)))
                         )  # Something simpler?

tests/test_arrow_dataset.py

@@ -3081,6 +3081,12 @@ def test_dataset_add_item_introduce_feature_type():
     assert dataset[:] == {"col_1": [None, None, None, "a"]}
 
 
+def test_dataset_filter_batched_indices():
+    ds = Dataset.from_dict({"num": [0, 1, 2, 3]})
+    ds = ds.filter(lambda num: num % 2 == 0, input_columns="num", batch_size=2)
+    assert all(item["num"] % 2 == 0 for item in ds)
+
+
 @pytest.mark.parametrize("in_memory", [False, True])
 def test_dataset_from_file(in_memory, dataset, arrow_file):
     filename = arrow_file