diff --git a/packages/data-designer-config/src/data_designer/config/utils/io_helpers.py b/packages/data-designer-config/src/data_designer/config/utils/io_helpers.py index e71ec3a10..a8c5a61ee 100644 --- a/packages/data-designer-config/src/data_designer/config/utils/io_helpers.py +++ b/packages/data-designer-config/src/data_designer/config/utils/io_helpers.py @@ -112,23 +112,29 @@ def load_processor_dataset(processors_outputs_path: Path, processor_name: str) - def read_parquet_dataset(path: Path) -> pd.DataFrame: """Read a parquet dataset from a path. + Directory schemas are unified permissively before reading so compatible + physical type drift across files, such as nested integers and floats, is + promoted to a common representation. + Args: path: The path to the parquet dataset, can be either a file or a directory. Returns: The parquet dataset as a pandas DataFrame. """ - try: - return lazy.pd.read_parquet(path, dtype_backend="pyarrow") - except Exception as e: - if path.is_dir() and "Unsupported cast" in str(e): - logger.warning("Failed to read parquets as folder, falling back to individual files") + if path.is_dir() and (parquet_files := sorted(path.glob("*.parquet"))): + schemas = [lazy.pq.read_schema(file) for file in parquet_files] + try: + unified_schema = lazy.pa.unify_schemas(schemas, promote_options="permissive") + except (lazy.pa.ArrowInvalid, lazy.pa.ArrowTypeError): + logger.warning("Failed to unify parquet schemas, falling back to individual files") return lazy.pd.concat( - [lazy.pd.read_parquet(file, dtype_backend="pyarrow") for file in sorted(path.glob("*.parquet"))], + [lazy.pd.read_parquet(file, dtype_backend="pyarrow") for file in parquet_files], ignore_index=True, ) - else: - raise e + return lazy.pd.read_parquet(path, dtype_backend="pyarrow", schema=unified_schema) + + return lazy.pd.read_parquet(path, dtype_backend="pyarrow") def validate_dataset_file_path(file_path: str | Path, should_exist: bool = True) -> Path: diff --git a/packages/data-designer-config/tests/config/utils/test_io_helpers.py b/packages/data-designer-config/tests/config/utils/test_io_helpers.py index 1d65d945d..5f6ba1244 100644 --- a/packages/data-designer-config/tests/config/utils/test_io_helpers.py +++ b/packages/data-designer-config/tests/config/utils/test_io_helpers.py @@ -16,11 +16,33 @@ from data_designer.config.utils.io_helpers import ( _maybe_rewrite_url, is_http_url, + read_parquet_dataset, serialize_data, smart_load_yaml, ) +def test_read_parquet_dataset_unifies_nested_numeric_types(tmp_path) -> None: + evaluations = [ + {"evaluations": [{"overall": {"score": 9}}]}, + {"evaluations": [{"overall": {"score": 9.5}}]}, + ] + for batch_number, evaluation in enumerate(evaluations): + lazy.pd.DataFrame({"qa_evaluations": [evaluation]}).to_parquet( + tmp_path / f"batch_{batch_number:05d}.parquet", + index=False, + ) + + schemas = [lazy.pq.read_schema(file) for file in sorted(tmp_path.glob("*.parquet"))] + assert schemas[0] != schemas[1] + + result = read_parquet_dataset(tmp_path) + + scores = [row["evaluations"][0]["overall"]["score"] for row in result["qa_evaluations"]] + assert scores == [9.0, 9.5] + assert all(isinstance(score, float) for score in scores) + + def test_smart_load_yaml(): stub_dict = { "hello": "world",