feat: add dtype_backend to result loading

packages/data-designer-config/src/data_designer/config/utils/io_helpers.py

@@ -10,15 +10,20 @@
 from decimal import Decimal
 from numbers import Number
 from pathlib import Path
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, Literal
 from urllib.parse import urlparse
 
 import requests
 import yaml
 
 import data_designer.lazy_heavy_imports as lazy
 from data_designer.config.errors import InvalidFileFormatError, InvalidFilePathError
-from data_designer.config.utils.storage_encoding import read_encoded_parquet, read_storage_plan, validate_storage_plans
+from data_designer.config.utils.storage_encoding import (
+    read_encoded_parquet,
+    read_native_parquet,
+    read_storage_plan,
+    validate_storage_plans,
+)
 
 if TYPE_CHECKING:
     import pandas as pd
@@ -110,7 +115,9 @@ def load_processor_dataset(processors_outputs_path: Path, processor_name: str) -
     raise FileNotFoundError(f"No artifacts found for processor named {processor_name!r}")
 
 
-def read_parquet_dataset(path: Path) -> pd.DataFrame:
+def read_parquet_dataset(
+    path: Path, *, dtype_backend: Literal["pyarrow", "numpy_nullable"] | None = "pyarrow"
+) -> pd.DataFrame:
     """Read a parquet dataset from a path.
 
     Directory schemas are unified permissively before reading so compatible
@@ -119,32 +126,39 @@ def read_parquet_dataset(path: Path) -> pd.DataFrame:
 
     Args:
         path: The path to the parquet dataset, can be either a file or a directory.
+        dtype_backend: The pandas dtype backend for native columns. `"pyarrow"` (the default)
+            and `"numpy_nullable"` select Arrow-backed or pandas nullable dtypes, respectively.
+            `None` uses pandas' default conversion, including persisted pandas metadata.
+            Decoded JSON storage columns remain object dtype for every backend. This option
+            does not normalize nested values or repair incompatible persisted dtype metadata.
 
     Returns:
         The parquet dataset as a pandas DataFrame.
     """
+    if dtype_backend not in ("pyarrow", "numpy_nullable", None):
+        raise ValueError(f"dtype_backend must be 'pyarrow', 'numpy_nullable' or None, got {dtype_backend!r}")
     if path.is_dir():
         dataset = lazy.pq.ParquetDataset(path)
         schemas = [fragment.physical_schema for fragment in dataset.fragments]
         plan = validate_storage_plans(schemas)
         if plan.columns:
-            return read_encoded_parquet(path, schemas)
+            return read_encoded_parquet(path, schemas, dtype_backend=dtype_backend)
         try:
             unified_schema = lazy.pa.unify_schemas([*schemas, dataset.schema], 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(dataset.files)],
+                [read_native_parquet(file, dtype_backend=dtype_backend) for file in sorted(dataset.files)],
                 ignore_index=True,
             )
-        return lazy.pd.read_parquet(
-            path, dtype_backend="pyarrow", schema=unified_schema, partitioning=dataset.partitioning
+        return read_native_parquet(
+            path, schema=unified_schema, partitioning=dataset.partitioning, dtype_backend=dtype_backend
         )
 
     schema = lazy.pq.read_schema(path)
     if read_storage_plan(schema).columns:
-        return read_encoded_parquet(path, [schema])
-    return lazy.pd.read_parquet(path, dtype_backend="pyarrow")
+        return read_encoded_parquet(path, [schema], dtype_backend=dtype_backend)
+    return read_native_parquet(path, dtype_backend=dtype_backend)
 
 
 def validate_dataset_file_path(file_path: str | Path, should_exist: bool = True) -> Path:

packages/data-designer-config/src/data_designer/config/utils/storage_encoding.py

@@ -7,7 +7,7 @@
 import math
 from copy import deepcopy
 from pathlib import Path
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, Literal
 
 from pydantic import ValidationError
 
@@ -18,6 +18,7 @@
 if TYPE_CHECKING:
     import pandas as pd
     import pyarrow as pa
+    from pyarrow.dataset import ParquetFileFragment, Partitioning
 
 
 def read_storage_plan(schema: pa.Schema) -> JsonStoragePlan:
@@ -120,35 +121,10 @@ def encode_storage_table(dataframe: pd.DataFrame, plan: JsonStoragePlan) -> pa.T
 def decode_storage_table(table: pa.Table) -> pd.DataFrame:
     """Restore logical JSON values without guessing from ordinary string contents."""
     plan = read_storage_plan(table.schema)
-    # pandas' ArrowDtype conversion casts even identical types, which can corrupt
-    # nested null children in Arrow. Wrap nested arrays directly without a cast.
-    nested_columns = [field.name for field in table.schema if lazy.pa.types.is_nested(field.type)]
-    frame = _prepare_table_for_pandas(table.drop(nested_columns)).to_pandas(types_mapper=lazy.pd.ArrowDtype)
-    column_labels: dict[str, Any] = {}
-    if nested_columns:
-        # An empty schema-only conversion restores pandas labels without touching
-        # real nested buffers. Physical Arrow names need not be pandas labels.
-        template = _prepare_table_for_pandas(table.schema.empty_table()).to_pandas(types_mapper=lazy.pd.ArrowDtype)
-        index_fields = [
-            name
-            for name in (table.schema.pandas_metadata or {}).get("index_columns", [])
-            if isinstance(name, str) and name in table.column_names
-        ]
-        data_fields = [name for name in table.column_names if name not in index_fields]
-        column_labels = dict(zip(data_fields, template.columns, strict=True))
-        if any(name in nested_columns for name in index_fields):
-            index_arrays = [lazy.pd.arrays.ArrowExtensionArray(table[name]) for name in index_fields]
-            frame.index = (
-                lazy.pd.Index(index_arrays[0], name=template.index.name)
-                if len(index_arrays) == 1
-                else lazy.pd.MultiIndex.from_arrays(index_arrays, names=template.index.names)
-            )
-        for name in nested_columns:
-            if name not in index_fields:
-                frame[column_labels[name]] = lazy.pd.Series(
-                    lazy.pd.arrays.ArrowExtensionArray(table[name]), index=frame.index
-                )
-        frame = frame.reindex(columns=template.columns)
+    frame = _arrow_table_to_pandas(table)
+    index_fields = (table.schema.pandas_metadata or {}).get("index_columns", [])
+    data_fields = [name for name in table.column_names if name not in index_fields]
+    column_labels = dict(zip(data_fields, frame.columns, strict=True))
     for name, paths in plan.columns.items():
         values = []
         for row, value in enumerate(table[name].to_pylist()):
@@ -158,7 +134,7 @@ def decode_storage_table(table: pa.Table) -> pd.DataFrame:
                 values.append(value)
             except (ValueError, TypeError, RecursionError) as exc:
                 raise InvalidFileFormatError(f"Cannot decode JSON storage column {name!r}, row {row}: {exc}") from exc
-        frame[column_labels.get(name, name)] = lazy.pd.Series(values, index=frame.index, dtype=object)
+        frame[column_labels[name]] = lazy.pd.Series(values, index=frame.index, dtype=object)
     if plan.columns:
         frame.attrs[JSON_STORAGE_RUN_KEY] = plan.model_dump()
     return frame
@@ -179,18 +155,40 @@ def decode_storage_dataframe(dataframe: pd.DataFrame, plan: JsonStoragePlan) ->
     return frame
 
 
-def read_encoded_parquet(path: Path, schemas: list[pa.Schema]) -> pd.DataFrame:
-    """Read annotated files individually to retain and validate each file's metadata."""
+def read_encoded_parquet(
+    path: Path,
+    schemas: list[pa.Schema],
+    *,
+    dtype_backend: Literal["pyarrow", "numpy_nullable"] | None = "pyarrow",
+) -> pd.DataFrame:
+    """Read native dtypes once and restore JSON values from validated fragments."""
     plan = validate_storage_plans(schemas)
     dataset = lazy.pq.ParquetDataset(path)
-    frames = []
-    for fragment in sorted(dataset.fragments, key=lambda item: item.path):
+    fragments = sorted(dataset.fragments, key=lambda item: item.path)
+    all_schemas = []
+    for fragment in fragments:
         schema = fragment.physical_schema
         for field in dataset.schema:
             if field.name not in schema.names:
                 schema = schema.append(field)
-        frames.append(decode_storage_table(fragment.to_table(schema=schema)))
-    frame = lazy.pd.concat(frames, ignore_index=True)
+        all_schemas.append(schema)
+    names = list(dict.fromkeys(name for schema in all_schemas for name in schema.names))
+    native_schemas = [
+        lazy.pa.schema([field for field in schema if field.name not in plan.columns], metadata=schema.metadata)
+        for schema in all_schemas
+    ]
+    native = _read_native_fragments(fragments, native_schemas, dataset.partitioning, dtype_backend)
+    protected = lazy.pd.concat(
+        [decode_storage_table(fragment.to_table(columns=list(plan.columns))) for fragment in fragments],
+        ignore_index=True,
+    )
+    frame = lazy.pd.concat([native.reset_index(drop=True), protected], axis=1)
+    schema = lazy.pa.schema(
+        [next(schema.field(name) for schema in all_schemas if name in schema.names) for name in names],
+        metadata=all_schemas[0].metadata,
+    )
+    columns = _arrow_table_to_pandas(schema.empty_table()).columns
+    frame = frame.reindex(columns=columns)
     frame.attrs[JSON_STORAGE_RUN_KEY] = plan.model_dump()
     return frame
 
@@ -221,6 +219,72 @@ def _prepare_table_for_pandas(table: pa.Table) -> pa.Table:
     return table.replace_schema_metadata(updated)
 
 
+def read_native_parquet(
+    path: str | Path | list[str],
+    *,
+    schema: pa.Schema | None = None,
+    partitioning: Partitioning | str | None = "hive",
+    dtype_backend: Literal["pyarrow", "numpy_nullable"] | None = "pyarrow",
+) -> pd.DataFrame:
+    """Read native columns without changing pandas' non-Arrow conversion policy.
+
+    Args:
+        path: A Parquet file, directory, or ordered fragment paths.
+        schema: The native schema to read, excluding JSON storage columns.
+        partitioning: Arrow partition discovery, or the dataset's existing partitioning.
+        dtype_backend: The requested pandas backend; None uses pandas defaults.
+
+    Returns:
+        Native columns as a DataFrame. Arrow-backed reads wrap the original
+        buffers directly; other backends use pandas' public Parquet reader.
+    """
+    kwargs: dict[str, Any] = {"partitioning": partitioning}
+    if schema is not None:
+        kwargs["schema"] = _prepare_table_for_pandas(schema.empty_table()).schema
+        kwargs["columns"] = schema.names
+    if dtype_backend == "pyarrow":
+        return _arrow_table_to_pandas(lazy.pq.read_table(path, **kwargs))
+    if dtype_backend is not None:
+        kwargs["dtype_backend"] = dtype_backend
+    return lazy.pd.read_parquet(path, **kwargs)
+
+
+def _read_native_fragments(
+    fragments: list[ParquetFileFragment],
+    schemas: list[pa.Schema],
+    partitioning: Partitioning | None,
+    dtype_backend: Literal["pyarrow", "numpy_nullable"] | None,
+) -> pd.DataFrame:
+    try:
+        unified = lazy.pa.unify_schemas(schemas, promote_options="permissive")
+    except (lazy.pa.ArrowInvalid, lazy.pa.ArrowTypeError):
+        return lazy.pd.concat(
+            [
+                read_native_parquet(
+                    [fragment.path], schema=schema, partitioning=partitioning, dtype_backend=dtype_backend
+                )
+                for fragment, schema in zip(fragments, schemas)
+            ],
+            ignore_index=True,
+        )
+    return read_native_parquet(
+        [fragment.path for fragment in fragments],
+        schema=unified,
+        partitioning=partitioning,
+        dtype_backend=dtype_backend,
+    )
+
+
+def _arrow_table_to_pandas(table: pa.Table) -> pd.DataFrame:
+    """Use pandas' Arrow extension hook without recasting the original arrays."""
+
+    class PreservingArrowDtype(lazy.pd.ArrowDtype):
+        def __from_arrow__(self, array: pa.Array | pa.ChunkedArray) -> pd.arrays.ArrowExtensionArray:
+            return lazy.pd.arrays.ArrowExtensionArray(array)
+
+    return _prepare_table_for_pandas(table).to_pandas(types_mapper=PreservingArrowDtype)
+
+
 def _transform_value(value: Any, path: list[str | None], *, encode: bool) -> Any:
     if value is None or value is lazy.pd.NA or (not encode and isinstance(value, float) and math.isnan(value)):
         return None

packages/data-designer-config/tests/config/utils/test_io_helpers.py

@@ -6,6 +6,7 @@
 import tempfile
 from datetime import date, datetime, timedelta
 from decimal import Decimal
+from pathlib import Path
 from unittest.mock import MagicMock, patch
 
 import pytest
@@ -22,7 +23,8 @@
 )
 
 
-def test_read_parquet_dataset_unifies_nested_numeric_types(tmp_path) -> None:
+@pytest.mark.parametrize("dtype_backend", ["pyarrow", "numpy_nullable", None])
+def test_read_parquet_dataset_unifies_nested_numeric_types(tmp_path: Path, dtype_backend: str | None) -> None:
     evaluations = [
         {"evaluations": [{"overall": {"score": 9}}]},
         {"evaluations": [{"overall": {"score": 9.5}}]},
@@ -36,13 +38,38 @@ def test_read_parquet_dataset_unifies_nested_numeric_types(tmp_path) -> None:
     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)
+    result = read_parquet_dataset(tmp_path, dtype_backend=dtype_backend)
 
     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)
 
 
+@pytest.mark.parametrize(
+    ("dtype_backend", "expected_dtype"),
+    [
+        ("pyarrow", lazy.pd.ArrowDtype(lazy.pa.int64())),
+        ("numpy_nullable", lazy.pd.Int64Dtype()),
+        (None, lazy.np.dtype("int64")),
+    ],
+)
+def test_read_parquet_dataset_fallback_honors_dtype_backend(
+    tmp_path: Path, dtype_backend: str | None, expected_dtype: object
+) -> None:
+    for batch in range(2):
+        lazy.pd.DataFrame({"id": [batch], "meta": [{"score": 9}]}).to_parquet(
+            tmp_path / f"batch_{batch}.parquet", index=False
+        )
+
+    with patch.object(lazy.pa, "unify_schemas", side_effect=lazy.pa.ArrowInvalid("forced")):
+        result = read_parquet_dataset(tmp_path, dtype_backend=dtype_backend)
+
+    assert result["id"].dtype == expected_dtype
+    assert result["id"].tolist() == [0, 1]
+    assert result["meta"].tolist() == [{"score": 9}, {"score": 9}]
+    assert (result["meta"].dtype == object) is (dtype_backend != "pyarrow")
+
+
 def test_smart_load_yaml():
     stub_dict = {
         "hello": "world",

packages/data-designer-config/tests/config/utils/test_storage_encoding.py

@@ -75,9 +75,10 @@ def test_null_empty_batches_have_explicit_nested_leaf_type(values: list[Any]) ->
 
 
 @pytest.mark.parametrize("counts", [(1, 3, 2), (3, 3, 3), (0, 0, 0)], ids=["variable", "uniform", "empty"])
-def test_nested_null_side_effect_survives_parquet_read(tmp_path: Path, counts: tuple[int, ...]) -> None:
+@pytest.mark.parametrize("encoded", [False, True], ids=["native", "json-encoded"])
+def test_nested_null_side_effect_survives_parquet_read(tmp_path: Path, counts: tuple[int, ...], encoded: bool) -> None:
     values = [[{"entries": [{"label": "synthetic", "optional": None} for _ in range(count)]}] for count in counts]
-    plan = JsonStoragePlan(columns={"payload": [[]]})
+    plan = JsonStoragePlan(columns={"payload": [[]]} if encoded else {})
     frame = pd.DataFrame({"side_effect": values, "payload": [{"x": 1}] * len(values), "ordinary": [1, 2, 3]})
     table = encode_storage_table(frame, plan)
     path = tmp_path / "batch.parquet"
@@ -93,7 +94,7 @@ def test_nested_null_side_effect_survives_parquet_read(tmp_path: Path, counts: t
     assert isinstance(restored["ordinary"].dtype, pd.ArrowDtype)
     assert restored["ordinary"].tolist() == [1, 2, 3]
     assert restored["payload"].tolist() == [{"x": 1}] * len(values)
-    assert restored.attrs[JSON_STORAGE_RUN_KEY] == plan.model_dump()
+    assert restored.attrs.get(JSON_STORAGE_RUN_KEY) == (plan.model_dump() if encoded else None)
     assert list(restored.columns) == list(frame.columns)
     assert original["side_effect"].to_pylist() == values
 

packages/data-designer-engine/src/data_designer/engine/storage/artifact_storage.py

@@ -11,7 +11,7 @@
 from datetime import datetime
 from functools import cached_property
 from pathlib import Path
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Literal
 from uuid import uuid4
 
 from pydantic import BaseModel, ConfigDict, Field, PrivateAttr, field_validator, model_validator
@@ -218,8 +218,13 @@ def create_batch_file_path(
             raise ArtifactStorageError("🛑 Batch number must be non-negative.")
         return self._get_stage_path(batch_stage) / BATCH_FILE_NAME_FORMAT.format(batch_number=batch_number)
 
-    def load_dataset(self, batch_stage: BatchStage = BatchStage.FINAL_RESULT) -> pd.DataFrame:
-        return read_parquet_dataset(self._get_stage_path(batch_stage))
+    def load_dataset(
+        self,
+        batch_stage: BatchStage = BatchStage.FINAL_RESULT,
+        *,
+        dtype_backend: Literal["pyarrow", "numpy_nullable"] | None = "pyarrow",
+    ) -> pd.DataFrame:
+        return read_parquet_dataset(self._get_stage_path(batch_stage), dtype_backend=dtype_backend)
 
     def load_processor_dataset(self, processor_name: str) -> pd.DataFrame:
         """Load a processor's output dataset. Raises ArtifactStorageError if not found."""

packages/data-designer/src/data_designer/interface/results.py

@@ -88,13 +88,20 @@ def load_analysis(self) -> DatasetProfilerResults:
         """
         return self._analysis
 
-    def load_dataset(self) -> pd.DataFrame:
+    def load_dataset(self, *, dtype_backend: Literal["pyarrow", "numpy_nullable"] | None = "pyarrow") -> pd.DataFrame:
         """Load the generated dataset as a pandas DataFrame.
 
+        Args:
+            dtype_backend: The pandas dtype backend for native columns. `"pyarrow"` (the default)
+                selects Arrow-backed dtypes; `"numpy_nullable"` selects pandas nullable dtypes.
+                `None` uses pandas' default conversion, including persisted pandas metadata.
+                Decoded JSON storage columns remain object dtype for every backend. This option
+                does not normalize nested values or repair incompatible persisted dtype metadata.
+
         Returns:
             A pandas DataFrame containing the full generated dataset.
         """
-        return self.artifact_storage.load_dataset()
+        return self.artifact_storage.load_dataset(dtype_backend=dtype_backend)
 
     def to_config_builder(self, columns: list[str] | None = None) -> DataDesignerConfigBuilder:
         """Create a new config builder seeded from this result dataset.