feat: use content defined chunking in arrow_writer.ParquetWriter

Author: kszucs

src/datasets/arrow_writer.py

@@ -341,8 +341,6 @@ def __init__(
 class ArrowWriter:
     """Shuffles and writes Examples to Arrow files."""
 
-    _WRITER_CLASS = pa.RecordBatchStreamWriter
-
     def __init__(
         self,
         schema: Optional[pa.Schema] = None,
@@ -430,7 +428,7 @@ def close(self):
         if self._closable_stream and not self.stream.closed:
             self.stream.close()  # This also closes self.pa_writer if it is opened
 
-    def _build_writer(self, inferred_schema: pa.Schema):
+    def _build_schema(self, inferred_schema: pa.Schema):
         schema = self.schema
         inferred_features = Features.from_arrow_schema(inferred_schema)
         if self._features is not None:
@@ -441,19 +439,24 @@ def _build_writer(self, inferred_schema: pa.Schema):
                     if name in fields:
                         if inferred_field == fields[name]:
                             inferred_features[name] = self._features[name]
-                self._features = inferred_features
+                features = inferred_features
                 schema: pa.Schema = inferred_schema
         else:
-            self._features = inferred_features
+            features = inferred_features
             schema: pa.Schema = inferred_features.arrow_schema
+
         if self.disable_nullable:
             schema = pa.schema(pa.field(field.name, field.type, nullable=False) for field in schema)
         if self.with_metadata:
-            schema = schema.with_metadata(self._build_metadata(DatasetInfo(features=self._features), self.fingerprint))
+            schema = schema.with_metadata(self._build_metadata(DatasetInfo(features=features), self.fingerprint))
         else:
             schema = schema.with_metadata({})
-        self._schema = schema
-        self.pa_writer = self._WRITER_CLASS(self.stream, schema)
+
+        return schema, features
+
+    def _build_writer(self, inferred_schema: pa.Schema):
+        self._schema, self._features = self._build_schema(inferred_schema)
+        self.pa_writer = pa.RecordBatchStreamWriter(self.stream, self._schema)
 
     @property
     def schema(self):
@@ -674,4 +677,11 @@ def finalize(self, close_stream=True):
 
 
 class ParquetWriter(ArrowWriter):
-    _WRITER_CLASS = pq.ParquetWriter
+    def __init__(self, *args, cdc_options=None, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.cdc_options = config.DEFAULT_CDC_OPTIONS if cdc_options is None else cdc_options
+
+    def _build_writer(self, inferred_schema: pa.Schema):
+        self._schema, self._features = self._build_schema(inferred_schema)
+        self.pa_writer = pq.ParquetWriter(self.stream, self._schema, use_content_defined_chunking=self.cdc_options)
+        self.pa_writer.add_key_value_metadata({"content_defined_chunking": json.dumps(self.cdc_options)})

tests/io/test_parquet.py

@@ -213,7 +213,7 @@ def test_parquet_write_uses_content_defined_chunking(dataset, tmp_path):
         writer = ParquetDatasetWriter(dataset, tmp_path / "foo.parquet")
         writer.write()
         assert MockWriter.call_count == 1
-        args, kwargs = MockWriter.call_args
+        _, kwargs = MockWriter.call_args
         # Save or check the arguments as needed
         assert "use_content_defined_chunking" in kwargs
         assert kwargs["use_content_defined_chunking"] == config.DEFAULT_CDC_OPTIONS

tests/test_arrow_writer.py

@@ -1,4 +1,5 @@
 import copy
+import json
 import os
 import tempfile
 from unittest import TestCase
@@ -9,6 +10,7 @@
 import pyarrow.parquet as pq
 import pytest
 
+from datasets import config
 from datasets.arrow_writer import ArrowWriter, OptimizedTypedSequence, ParquetWriter, TypedSequence
 from datasets.features import Array2D, ClassLabel, Features, Image, Value
 from datasets.features.features import Array2DExtensionType, cast_to_python_objects
@@ -334,6 +336,39 @@ def test_parquet_writer_write():
     assert pa_table.to_pydict() == {"col_1": ["foo", "bar"], "col_2": [1, 2]}
 
 
+custom_cdc_options = {
+    "min_chunk_size": 128 * 1024,  # 128 KiB
+    "max_chunk_size": 512 * 1024,  # 512 KiB
+    "norm_level": 1,
+}
+
+
+@pytest.mark.parametrize(
+    ("cdc_options", "expected_options"), [(None, config.DEFAULT_CDC_OPTIONS), (custom_cdc_options, custom_cdc_options)]
+)
+def test_parquet_write_uses_content_defined_chunking(cdc_options, expected_options):
+    output = pa.BufferOutputStream()
+    with patch("pyarrow.parquet.ParquetWriter", wraps=pq.ParquetWriter) as MockWriter:
+        with ParquetWriter(stream=output, cdc_options=cdc_options) as writer:
+            writer.write({"col_1": "foo", "col_2": 1})
+            writer.write({"col_1": "bar", "col_2": 2})
+            writer.finalize()
+        assert MockWriter.call_count == 1
+        _, kwargs = MockWriter.call_args
+        assert "use_content_defined_chunking" in kwargs
+        assert kwargs["use_content_defined_chunking"] == expected_options
+
+    # read metadata from the output stream
+    with pa.input_stream(output.getvalue()) as stream:
+        metadata = pq.read_metadata(stream)
+        key_value_metadata = metadata.metadata
+
+    # check that the content defined chunking options are persisted
+    assert b"content_defined_chunking" in key_value_metadata
+    json_encoded_options = key_value_metadata[b"content_defined_chunking"].decode("utf-8")
+    assert json.loads(json_encoded_options) == expected_options
+
+
 @require_pil
 @pytest.mark.parametrize("embed_local_files", [False, True])
 def test_writer_embed_local_files(tmp_path, embed_local_files):