chore: rename cdc_options argument to use_content_defined_chunking

Author: kszucs

src/datasets/arrow_writer.py

@@ -678,11 +678,17 @@ def finalize(self, close_stream=True):
 
 
 class ParquetWriter(ArrowWriter):
-    def __init__(self, *args, cdc_options=None, **kwargs):
+    def __init__(self, *args, use_content_defined_chunking=None, **kwargs):
         super().__init__(*args, **kwargs)
-        self.cdc_options = config.DEFAULT_CDC_OPTIONS if cdc_options is None else cdc_options
+        self.use_content_defined_chunking = (
+            config.DEFAULT_CDC_OPTIONS if use_content_defined_chunking is None else use_content_defined_chunking
+        )
 
     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)})
+        self.pa_writer = pq.ParquetWriter(
+            self.stream, self._schema, use_content_defined_chunking=self.use_content_defined_chunking
+        )
+        self.pa_writer.add_key_value_metadata(
+            {"content_defined_chunking": json.dumps(self.use_content_defined_chunking)}
+        )

src/datasets/io/parquet.py

@@ -77,33 +77,43 @@ def __init__(
         dataset: Dataset,
         path_or_buf: Union[PathLike, BinaryIO],
         batch_size: Optional[int] = None,
-        cdc_options: Optional[dict] = None,
         storage_options: Optional[dict] = None,
+        use_content_defined_chunking: Optional[dict] = None,
         **parquet_writer_kwargs,
     ):
         self.dataset = dataset
         self.path_or_buf = path_or_buf
-        self.cdc_options = cdc_options
         self.batch_size = batch_size or get_writer_batch_size(dataset.features)
         self.storage_options = storage_options or {}
         self.parquet_writer_kwargs = parquet_writer_kwargs
+        self.use_content_defined_chunking = use_content_defined_chunking
 
     def write(self) -> int:
         batch_size = self.batch_size if self.batch_size else config.DEFAULT_MAX_BATCH_SIZE
-        cdc_options = self.cdc_options if self.cdc_options else config.DEFAULT_CDC_OPTIONS
+        use_content_defined_chunking = (
+            self.use_content_defined_chunking if self.use_content_defined_chunking else config.DEFAULT_CDC_OPTIONS
+        )
 
         if isinstance(self.path_or_buf, (str, bytes, os.PathLike)):
             with fsspec.open(self.path_or_buf, "wb", **(self.storage_options or {})) as buffer:
                 written = self._write(
-                    file_obj=buffer, batch_size=batch_size, cdc_options=cdc_options, **self.parquet_writer_kwargs
+                    file_obj=buffer,
+                    batch_size=batch_size,
+                    use_content_defined_chunking=use_content_defined_chunking,
+                    **self.parquet_writer_kwargs,
                 )
         else:
             written = self._write(
-                file_obj=self.path_or_buf, batch_size=batch_size, cdc_options=cdc_options, **self.parquet_writer_kwargs
+                file_obj=self.path_or_buf,
+                batch_size=batch_size,
+                use_content_defined_chunking=use_content_defined_chunking,
+                **self.parquet_writer_kwargs,
             )
         return written
 
-    def _write(self, file_obj: BinaryIO, batch_size: int, cdc_options: dict, **parquet_writer_kwargs) -> int:
+    def _write(
+        self, file_obj: BinaryIO, batch_size: int, use_content_defined_chunking: bool | dict, **parquet_writer_kwargs
+    ) -> int:
         """Writes the pyarrow table as Parquet to a binary file handle.
 
         Caller is responsible for opening and closing the handle.
@@ -113,7 +123,7 @@ def _write(self, file_obj: BinaryIO, batch_size: int, cdc_options: dict, **parqu
         schema = self.dataset.features.arrow_schema
 
         writer = pq.ParquetWriter(
-            file_obj, schema=schema, use_content_defined_chunking=cdc_options, **parquet_writer_kwargs
+            file_obj, schema=schema, use_content_defined_chunking=use_content_defined_chunking, **parquet_writer_kwargs
         )
 
         for offset in hf_tqdm(
@@ -130,6 +140,6 @@ def _write(self, file_obj: BinaryIO, batch_size: int, cdc_options: dict, **parqu
             written += batch.nbytes
 
         # TODO(kszucs): we may want to persist multiple parameters
-        writer.add_key_value_metadata({"content_defined_chunking": json.dumps(cdc_options)})
+        writer.add_key_value_metadata({"content_defined_chunking": json.dumps(use_content_defined_chunking)})
         writer.close()
         return written

tests/io/test_parquet.py

@@ -231,7 +231,7 @@ def test_parquet_write_uses_content_defined_chunking(dataset, tmp_path):
 )
 def test_parquet_writer_persist_cdc_options_as_metadata(dataset, tmp_path, cdc_options, expected_options):
     # write the dataset to parquet with the default CDC options
-    writer = ParquetDatasetWriter(dataset, tmp_path / "foo.parquet", cdc_options=cdc_options)
+    writer = ParquetDatasetWriter(dataset, tmp_path / "foo.parquet", use_content_defined_chunking=cdc_options)
     assert writer.write() > 0
 
     # read the parquet KV metadata

tests/test_arrow_writer.py

@@ -349,7 +349,7 @@ def test_parquet_writer_write():
 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:
+        with ParquetWriter(stream=output, use_content_defined_chunking=cdc_options) as writer:
             writer.write({"col_1": "foo", "col_2": 1})
             writer.write({"col_1": "bar", "col_2": 2})
             writer.finalize()