Fix silent dataset truncation in batched map with mismatched empty column (#6879)

src/datasets/arrow_writer.py

@@ -718,7 +718,7 @@ def _write_batch(
         writer_batch_size: Optional[int] = None,
         try_original_type: Optional[bool] = True,
     ):
-        if batch_examples and len(next(iter(batch_examples.values()))) == 0:
+        if batch_examples and all(len(col_values) == 0 for col_values in batch_examples.values()):
             return
         features = None if self.pa_writer is None and self.update_features else self._features
         try_features = self._features if self.pa_writer is None and self.update_features else None

tests/test_arrow_writer.py

@@ -166,6 +166,17 @@ def test_write_batch(fields, writer_batch_size):
     _check_output(output.getvalue(), expected_num_chunks=num_examples if writer_batch_size == 1 else 1)
 
 
+def test_write_batch_with_mismatched_empty_column_raises():
+    # Regression test for https://github.com/huggingface/datasets/issues/6879
+    # A batch is only skipped when *every* column is empty. A batch that mixes an
+    # empty column with a non-empty one is length-mismatched and must raise instead
+    # of being silently dropped (which used to truncate the dataset to 0 rows).
+    output = pa.BufferOutputStream()
+    with ArrowWriter(stream=output) as writer:
+        with pytest.raises(pa.lib.ArrowInvalid):
+            writer.write_batch({"col_1": [], "col_2": [1]})
+
+
 @pytest.mark.parametrize("writer_batch_size", [None, 1, 10])
 @pytest.mark.parametrize(
     "fields", [None, {"col_1": pa.string(), "col_2": pa.int64()}, {"col_1": pa.string(), "col_2": pa.int32()}]