Keep CSV row indices continuous across batches

Assign each exported DataFrame its logical dataset row offset when index output is requested. Preserve index=False behavior and cover single and multiple batches, multiprocessing, selected rows, and byte counts.

Author: j45856021-dev

src/datasets/arrow_dataset.py

@@ -5355,6 +5355,8 @@ def to_csv(
 
                 </Changed>
 
+                With `index=True`, row indices are zero-based positions in the dataset and are continuous across batches.
+
         Returns:
             `int`: The number of characters or bytes written.
 

src/datasets/io/csv.py

@@ -105,7 +105,10 @@ def _batch_csv(self, args):
             key=slice(offset, offset + self.batch_size),
             indices=self.dataset._indices,
         )
-        csv_str = batch.to_pandas(integer_object_nulls=True).to_csv(
+        batch_df = batch.to_pandas(integer_object_nulls=True)
+        if index:
+            batch_df.index = range(offset, offset + len(batch_df))
+        csv_str = batch_df.to_csv(
             path_or_buf=None, header=header if (offset == 0) else False, index=index, **to_csv_kwargs
         )
         return csv_str.encode(self.encoding)

tests/io/test_csv.py

@@ -157,6 +157,23 @@ def test_dataset_to_csv_multiproc(csv_path, tmp_path):
         assert row1 == row2
 
 
+@pytest.mark.parametrize("batch_size", [5, 20])
+@pytest.mark.parametrize("num_proc", [None, 2])
+@pytest.mark.parametrize("index", [False, True])
+@pytest.mark.parametrize("selection", [None, [11, 3, 8, 0, 10, 1, 9]])
+def test_dataset_to_csv_index_independent_of_batch_size(batch_size, num_proc, index, selection, tmp_path):
+    dataset = Dataset.from_dict({"value": list(range(12))})
+    if selection is not None:
+        dataset = dataset.select(selection)
+    expected = dataset.to_pandas().to_csv(index=index, index_label="row").encode("utf-8")
+    output_csv = tmp_path / "out.csv"
+
+    written = dataset.to_csv(output_csv, batch_size=batch_size, num_proc=num_proc, index=index, index_label="row")
+
+    assert output_csv.read_bytes() == expected
+    assert written == len(expected)
+
+
 def test_dataset_to_csv_invalidproc(csv_path, tmp_path):
     cache_dir = tmp_path / "cache"
     output_csv = os.path.join(cache_dir, "tmp.csv")