diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py index 40670caae6d1..c430ffc034dd 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py @@ -684,7 +684,9 @@ def serialize(self, engine: StreamingEngine | None) -> dict: config_options = config_options.drop_unserializable() rapidsmpf_options = engine.rapidsmpf_options.get_strings() result["config_options"] = { - "config_options": dataclasses.asdict(config_options), + "config_options": dataclasses.asdict( + config_options, dict_factory=ConfigOptions.dict_factory + ), "rapidsmpf_options": rapidsmpf_options, } # discard unserializable / unnecessary UUIDs diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index b7aeabf58140..315aa0c8f8cd 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -1002,6 +1002,27 @@ class ConfigOptions(Generic[ExecutorType]): device: int | None = None memory_resource_config: MemoryResourceConfig | None = None + @staticmethod + def dict_factory(items: list[tuple[str, Any]]) -> dict[str, Any]: + """ + ``dict_factory`` for :func:`dataclasses.asdict`. + + Converts any :data:`UNSPECIFIED` value to ``None`` + (e.g. ParquetOptions.prefetch_file_metadata) so the resulting + dict can be serialized with :func:`json.dumps`. + + Parameters + ---------- + items + The ``(key, value)`` pairs for a single dataclass level, as passed + by :func:`dataclasses.asdict`. + + Returns + ------- + A dict with :data:`UNSPECIFIED` values replaced by ``None``. + """ + return {k: (None if isinstance(v, Unspecified) else v) for k, v in items} + def drop_unserializable(self) -> ConfigOptions[ExecutorType]: """ Return a copy safe to pickle to a worker/actor. diff --git a/python/cudf_polars/tests/test_config.py b/python/cudf_polars/tests/test_config.py index 59321b07cf39..fd184cf2f2c3 100644 --- a/python/cudf_polars/tests/test_config.py +++ b/python/cudf_polars/tests/test_config.py @@ -3,6 +3,7 @@ from __future__ import annotations +import dataclasses from typing import cast import pytest @@ -553,6 +554,16 @@ def test_parquet_options_object_engine_default() -> None: assert isinstance(config.parquet_options.prefetch_file_metadata, Unspecified) +def test_parquet_options_unspecified_dict_factory() -> None: + parquet_options = ParquetOptions() + config = ConfigOptions.from_polars_engine( + pl.GPUEngine(executor="streaming", parquet_options=parquet_options) + ) + assert isinstance(config.parquet_options.prefetch_file_metadata, Unspecified) + result = dataclasses.asdict(config, dict_factory=ConfigOptions.dict_factory) + assert result["parquet_options"]["prefetch_file_metadata"] is None + + def test_validate_raise_on_fail() -> None: with pytest.raises(TypeError, match="'raise_on_fail' must be"): ConfigOptions.from_polars_engine(