test: cover more input parameter values for ParquetDatasetWriter

Author: kszucs

src/datasets/io/parquet.py

@@ -78,44 +78,37 @@ def __init__(
         path_or_buf: Union[PathLike, BinaryIO],
         batch_size: Optional[int] = None,
         storage_options: Optional[dict] = None,
-        use_content_defined_chunking: Optional[dict] = None,
+        use_content_defined_chunking: bool | dict = True,
         **parquet_writer_kwargs,
     ):
         self.dataset = dataset
         self.path_or_buf = path_or_buf
         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
+        if use_content_defined_chunking is True:
+            use_content_defined_chunking = config.DEFAULT_CDC_OPTIONS
         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
-        use_content_defined_chunking = (
-            config.DEFAULT_CDC_OPTIONS
-            if self.use_content_defined_chunking is None
-            else self.use_content_defined_chunking
-        )
 
         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,
-                    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,
-                use_content_defined_chunking=use_content_defined_chunking,
                 **self.parquet_writer_kwargs,
             )
         return written
 
-    def _write(
-        self, file_obj: BinaryIO, batch_size: int, use_content_defined_chunking: bool | dict, **parquet_writer_kwargs
-    ) -> int:
+    def _write(self, file_obj: BinaryIO, batch_size: int, **parquet_writer_kwargs) -> int:
         """Writes the pyarrow table as Parquet to a binary file handle.
 
         Caller is responsible for opening and closing the handle.
@@ -125,7 +118,10 @@ def _write(
         schema = self.dataset.features.arrow_schema
 
         writer = pq.ParquetWriter(
-            file_obj, schema=schema, use_content_defined_chunking=use_content_defined_chunking, **parquet_writer_kwargs
+            file_obj,
+            schema=schema,
+            use_content_defined_chunking=self.use_content_defined_chunking,
+            **parquet_writer_kwargs,
         )
 
         for offset in hf_tqdm(
@@ -142,8 +138,8 @@ def _write(
             written += batch.nbytes
 
         # TODO(kszucs): we may want to persist multiple parameters
-        if use_content_defined_chunking is not False:
-            writer.add_key_value_metadata({"content_defined_chunking": json.dumps(use_content_defined_chunking)})
+        if self.use_content_defined_chunking is not False:
+            writer.add_key_value_metadata({"content_defined_chunking": json.dumps(self.use_content_defined_chunking)})
 
         writer.close()
         return written

tests/io/test_parquet.py

@@ -219,29 +219,38 @@ def test_parquet_write_uses_content_defined_chunking(dataset, tmp_path):
         assert kwargs["use_content_defined_chunking"] == config.DEFAULT_CDC_OPTIONS
 
 
-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_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", use_content_defined_chunking=cdc_options)
-    assert writer.write() > 0
-
-    # read the parquet KV metadata
-    metadata = pq.read_metadata(tmp_path / "foo.parquet")
-    key_value_metadata = metadata.metadata
-
-    # check that the content defined chunking options are persisted
+def test_parquet_writer_persist_cdc_options_as_metadata(dataset, tmp_path):
+    def write_and_get_metadata(**kwargs):
+        # write the dataset to parquet with the default CDC options
+        writer = ParquetDatasetWriter(dataset, tmp_path / "foo.parquet", **kwargs)
+        assert writer.write() > 0
+
+        # read the parquet KV metadata
+        metadata = pq.read_metadata(tmp_path / "foo.parquet")
+        key_value_metadata = metadata.metadata
+
+        return key_value_metadata
+
+    # by default no arguments are passed, same as passing True using the default options
+    for key_value_metadata in [write_and_get_metadata(), write_and_get_metadata(use_content_defined_chunking=True)]:
+        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) == config.DEFAULT_CDC_OPTIONS
+
+    # passing False disables the content defined chunking and doesn't persist the options in metadata
+    key_value_metadata = write_and_get_metadata(use_content_defined_chunking=False)
+    assert b"content_defined_chunking" not in key_value_metadata
+
+    # passing custom options, using the custom options
+    custom_cdc_options = {
+        "min_chunk_size": 128 * 1024,  # 128 KiB
+        "max_chunk_size": 512 * 1024,  # 512 KiB
+        "norm_level": 1,
+    }
+    key_value_metadata = write_and_get_metadata(use_content_defined_chunking=custom_cdc_options)
     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
+    assert json.loads(json_encoded_options) == custom_cdc_options
 
 
 def test_dataset_to_parquet_keeps_features(shared_datadir, tmp_path):