fix: unify Parquet schemas when reading

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:

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",