From 5d376b3ecaca7e9752f8e7ac4ebe5d723847ce7e Mon Sep 17 00:00:00 2001
From: Aaron Gonzales <aagonzales@nvidia.com>
Date: Tue, 3 Nov 2026 00:34:53 +0000
Subject: [PATCH 1/5] feat(interface): add dtype_backend to load_dataset
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

DatasetCreationResults.load_dataset() always read parquet with
dtype_backend="pyarrow", so nested columns came back as ArrowDtype.
Add a keyword-only dtype_backend ("pyarrow" default, "numpy_nullable",
or None) threaded through ArtifactStorage.load_dataset() and
read_parquet_dataset(), including the schema-unification and encoded
JSON storage branches. None omits the pandas option, so pandas applies
its default conversion (nested columns as object).

The default is unchanged. CompositeWorkflow.load_dataset does not yet
expose the option.

🤖 Generated with Claude Code

Co-Authored-By: Claude <noreply@anthropic.com>
Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>
---
 .../data_designer/config/utils/io_helpers.py  | 24 +++--
 .../config/utils/storage_encoding.py          | 91 ++++++++++++------
 .../tests/config/utils/test_io_helpers.py     | 96 +++++++++++++++++++
 .../config/utils/test_storage_encoding.py     | 35 +++++++
 .../engine/storage/artifact_storage.py        | 11 ++-
 .../engine/storage/test_artifact_storage.py   | 17 ++++
 .../src/data_designer/interface/results.py    | 11 ++-
 .../tests/interface/test_results.py           | 27 +++++-
 8 files changed, 266 insertions(+), 46 deletions(-)

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 61402e94d..45da7488b 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
@@ -10,7 +10,7 @@
 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
@@ -110,7 +110,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 +121,36 @@ 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 the returned columns. `"pyarrow"` (the
+            default) and `"numpy_nullable"` are passed to pandas. `None` omits the option so
+            pandas applies its default conversion: nested columns are object (lists as NumPy
+            arrays, structs as dicts), while scalar columns follow pandas' defaults (text as
+            `str`, integers with nulls as `float64`).
 
     Returns:
         The parquet dataset as a pandas DataFrame.
     """
+    read_kwargs: dict[str, Any] = {} if dtype_backend is None else {"dtype_backend": dtype_backend}
     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)],
+                [lazy.pd.read_parquet(file, **read_kwargs) 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 lazy.pd.read_parquet(path, schema=unified_schema, partitioning=dataset.partitioning, **read_kwargs)
 
     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 lazy.pd.read_parquet(path, **read_kwargs)
 
 
 def validate_dataset_file_path(file_path: str | Path, should_exist: bool = True) -> Path:
diff --git a/packages/data-designer-config/src/data_designer/config/utils/storage_encoding.py b/packages/data-designer-config/src/data_designer/config/utils/storage_encoding.py
index a49df112e..dc775ce39 100644
--- a/packages/data-designer-config/src/data_designer/config/utils/storage_encoding.py
+++ b/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
 
@@ -117,38 +117,62 @@ def encode_storage_table(dataframe: pd.DataFrame, plan: JsonStoragePlan) -> pa.T
     return table
 
 
-def decode_storage_table(table: pa.Table) -> pd.DataFrame:
+def _nullable_types_mapper() -> dict[pa.DataType, Any]:
+    pa, pd = lazy.pa, lazy.pd
+    mapping: dict[pa.DataType, Any] = {
+        pa.bool_(): pd.BooleanDtype(),
+        pa.string(): pd.StringDtype(),
+        pa.large_string(): pd.StringDtype(),
+    }
+    for bits in (8, 16, 32, 64):
+        mapping[getattr(pa, f"int{bits}")()] = getattr(pd, f"Int{bits}Dtype")()
+        mapping[getattr(pa, f"uint{bits}")()] = getattr(pd, f"UInt{bits}Dtype")()
+    mapping[pa.float32()] = pd.Float32Dtype()
+    mapping[pa.float64()] = pd.Float64Dtype()
+    return mapping
+
+
+def decode_storage_table(
+    table: pa.Table, *, dtype_backend: Literal["pyarrow", "numpy_nullable"] | None = "pyarrow"
+) -> 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
+    if dtype_backend == "pyarrow":
+        # 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)
+        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)
                 )
-        frame = frame.reindex(columns=template.columns)
+            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)
+    elif dtype_backend == "numpy_nullable":
+        frame = table.to_pandas(types_mapper=_nullable_types_mapper().get)
+    elif dtype_backend is None:
+        frame = table.to_pandas()
+    else:
+        raise ValueError(f"dtype_backend must be 'pyarrow', 'numpy_nullable' or None, got {dtype_backend!r}")
     for name, paths in plan.columns.items():
         values = []
         for row, value in enumerate(table[name].to_pylist()):
@@ -179,7 +203,12 @@ def decode_storage_dataframe(dataframe: pd.DataFrame, plan: JsonStoragePlan) ->
     return frame
 
 
-def read_encoded_parquet(path: Path, schemas: list[pa.Schema]) -> pd.DataFrame:
+def read_encoded_parquet(
+    path: Path,
+    schemas: list[pa.Schema],
+    *,
+    dtype_backend: Literal["pyarrow", "numpy_nullable"] | None = "pyarrow",
+) -> pd.DataFrame:
     """Read annotated files individually to retain and validate each file's metadata."""
     plan = validate_storage_plans(schemas)
     dataset = lazy.pq.ParquetDataset(path)
@@ -189,7 +218,7 @@ def read_encoded_parquet(path: Path, schemas: list[pa.Schema]) -> pd.DataFrame:
         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)))
+        frames.append(decode_storage_table(fragment.to_table(schema=schema), dtype_backend=dtype_backend))
     frame = lazy.pd.concat(frames, ignore_index=True)
     frame.attrs[JSON_STORAGE_RUN_KEY] = plan.model_dump()
     return frame
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 5f6ba1244..6d2d5c4eb 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
@@ -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
@@ -43,6 +44,101 @@ def test_read_parquet_dataset_unifies_nested_numeric_types(tmp_path) -> None:
     assert all(isinstance(score, float) for score in scores)
 
 
+def _write_nested_batches(path: Path, evaluations: list[dict]) -> None:
+    for batch_number, evaluation in enumerate(evaluations):
+        lazy.pd.DataFrame({"id": [batch_number], "qa_evaluations": [evaluation]}).to_parquet(
+            path / f"batch_{batch_number:05d}.parquet",
+            index=False,
+        )
+
+
+def _write_simple_nested_file(path: Path) -> Path:
+    file_path = path / "batch_00000.parquet"
+    lazy.pd.DataFrame({"id": [1], "meta": [{"score": 9, "label": "x"}], "tags": [["a", "b"]]}).to_parquet(
+        file_path, index=False
+    )
+    return file_path
+
+
+@pytest.mark.parametrize("as_directory", [True, False], ids=["directory", "file"])
+def test_read_parquet_dataset_defaults_to_pyarrow_dtypes(tmp_path, as_directory: bool) -> None:
+    file_path = _write_simple_nested_file(tmp_path)
+
+    result = read_parquet_dataset(tmp_path if as_directory else file_path)
+
+    assert all(isinstance(dtype, lazy.pd.ArrowDtype) for dtype in result.dtypes)
+    assert result["meta"].iloc[0] == {"score": 9, "label": "x"}
+    assert result["tags"].iloc[0] == ["a", "b"]
+
+
+@pytest.mark.parametrize("as_directory", [True, False], ids=["directory", "file"])
+def test_read_parquet_dataset_without_dtype_backend_returns_object_columns(tmp_path, as_directory: bool) -> None:
+    file_path = _write_simple_nested_file(tmp_path)
+
+    result = read_parquet_dataset(tmp_path if as_directory else file_path, dtype_backend=None)
+
+    assert not any(isinstance(dtype, lazy.pd.ArrowDtype) for dtype in result.dtypes)
+    assert result["meta"].dtype == object
+    assert result["tags"].dtype == object
+    assert result["meta"].iloc[0] == {"score": 9, "label": "x"}
+    assert list(result["tags"].iloc[0]) == ["a", "b"]
+
+
+def test_read_parquet_dataset_without_dtype_backend_round_trips_through_pandas(tmp_path) -> None:
+    result = read_parquet_dataset(_write_simple_nested_file(tmp_path), dtype_backend=None)
+
+    result.to_parquet(tmp_path / "round_trip.parquet", index=False)
+    restored = lazy.pd.read_parquet(tmp_path / "round_trip.parquet")
+
+    assert restored["meta"].iloc[0] == {"score": 9, "label": "x"}
+    assert list(restored["tags"].iloc[0]) == ["a", "b"]
+
+
+@pytest.mark.parametrize("as_directory", [True, False], ids=["directory", "file"])
+def test_read_parquet_dataset_accepts_numpy_nullable(tmp_path, as_directory: bool) -> None:
+    file_path = _write_simple_nested_file(tmp_path)
+
+    result = read_parquet_dataset(tmp_path if as_directory else file_path, dtype_backend="numpy_nullable")
+
+    assert result["id"].dtype == lazy.pd.Int64Dtype()
+    assert result["meta"].iloc[0] == {"score": 9, "label": "x"}
+
+
+@pytest.mark.parametrize(
+    ("dtype_backend", "id_check"),
+    [
+        ("pyarrow", lambda dtype: isinstance(dtype, lazy.pd.ArrowDtype)),
+        ("numpy_nullable", lambda dtype: dtype == lazy.pd.Int64Dtype()),
+        (None, lambda dtype: dtype == "int64"),
+    ],
+)
+def test_read_parquet_dataset_schema_unification_fallback_honors_dtype_backend(
+    tmp_path, dtype_backend, id_check
+) -> None:
+    _write_nested_batches(tmp_path, [{"evaluations": [{"score": 9}]}, {"evaluations": [{"score": 10}]}])
+
+    with patch.object(lazy.pa, "unify_schemas", side_effect=lazy.pa.ArrowInvalid("forced")):
+        result = read_parquet_dataset(tmp_path, dtype_backend=dtype_backend)
+
+    assert id_check(result["id"].dtype)
+    assert result["id"].tolist() == [0, 1]
+    assert (result["qa_evaluations"].dtype == object) is (dtype_backend != "pyarrow")
+
+
+@pytest.mark.parametrize("dtype_backend", ["pyarrow", "numpy_nullable", None])
+def test_read_parquet_dataset_unifies_nested_numeric_types_for_every_backend(tmp_path, dtype_backend) -> None:
+    _write_nested_batches(
+        tmp_path,
+        [{"evaluations": [{"overall": {"score": 9}}]}, {"evaluations": [{"overall": {"score": 9.5}}]}],
+    )
+
+    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)
+
+
 def test_smart_load_yaml():
     stub_dict = {
         "hello": "world",
diff --git a/packages/data-designer-config/tests/config/utils/test_storage_encoding.py b/packages/data-designer-config/tests/config/utils/test_storage_encoding.py
index 032ab097c..8bc5f27bd 100644
--- a/packages/data-designer-config/tests/config/utils/test_storage_encoding.py
+++ b/packages/data-designer-config/tests/config/utils/test_storage_encoding.py
@@ -514,6 +514,41 @@ def test_legacy_partitioned_directory_keeps_native_columns(tmp_path: Path) -> No
     assert JSON_STORAGE_RUN_KEY not in result.attrs
 
 
+@pytest.mark.parametrize("as_directory", [True, False], ids=["directory", "file"])
+@pytest.mark.parametrize(
+    ("dtype_backend", "count_check", "label_check"),
+    [
+        ("pyarrow", lambda dtype: isinstance(dtype, pd.ArrowDtype), lambda dtype: isinstance(dtype, pd.ArrowDtype)),
+        ("numpy_nullable", lambda dtype: dtype == pd.Int64Dtype(), lambda dtype: dtype == pd.StringDtype()),
+        (None, lambda dtype: dtype == "int64", lambda dtype: pd.api.types.is_string_dtype(dtype)),
+    ],
+)
+def test_encoded_dataset_honors_dtype_backend_for_native_columns(
+    tmp_path: Path, as_directory: bool, dtype_backend: Any, count_check: Any, label_check: Any
+) -> None:
+    values = [{"alpha": 1}, {"beta": [1, "two"]}]
+    frame = pd.DataFrame({"count": [1, 2], "label": ["a", "b"], "payload": pd.Series(values, dtype=object)})
+    plan = JsonStoragePlan(columns={"payload": [[]]})
+    path = tmp_path / "batch_0.parquet"
+    pq.write_table(encode_storage_table(frame, plan), path)
+
+    result = read_parquet_dataset(tmp_path if as_directory else path, dtype_backend=dtype_backend)
+
+    assert count_check(result["count"].dtype)
+    assert label_check(result["label"].dtype)
+    assert result["payload"].dtype == object
+    assert result["payload"].tolist() == values
+    assert result.attrs[JSON_STORAGE_RUN_KEY] == plan.model_dump()
+
+
+def test_decode_storage_table_rejects_unknown_dtype_backend() -> None:
+    table = pa.Table.from_pylist([{"payload": "1"}]).replace_schema_metadata(
+        {JSON_STORAGE_METADATA_KEY: JsonStoragePlan(columns={"payload": [[]]}).model_dump_json().encode()}
+    )
+    with pytest.raises(ValueError, match="dtype_backend must be"):
+        decode_storage_table(table, dtype_backend="bogus")  # type: ignore[arg-type]
+
+
 def test_encoded_partitioned_directory_restores_json_and_partition_columns(tmp_path: Path) -> None:
     values = [{"alpha": 1}, {"beta": [1, "two"]}, None]
     frame = pd.DataFrame({"group": ["a", "b", "c"], "payload": pd.Series(values, dtype=object)})
diff --git a/packages/data-designer-engine/src/data_designer/engine/storage/artifact_storage.py b/packages/data-designer-engine/src/data_designer/engine/storage/artifact_storage.py
index 061423ae4..daee3df1d 100644
--- a/packages/data-designer-engine/src/data_designer/engine/storage/artifact_storage.py
+++ b/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."""
diff --git a/packages/data-designer-engine/tests/engine/storage/test_artifact_storage.py b/packages/data-designer-engine/tests/engine/storage/test_artifact_storage.py
index 9b9343f29..de31e27f8 100644
--- a/packages/data-designer-engine/tests/engine/storage/test_artifact_storage.py
+++ b/packages/data-designer-engine/tests/engine/storage/test_artifact_storage.py
@@ -194,6 +194,23 @@ def test_artifact_storage_read_parquet_files(stub_artifact_storage):
     lazy.pd.testing.assert_frame_equal(lazy.pd.concat([read_df1, read_df2], ignore_index=True), read_df)
 
 
+def test_artifact_storage_load_dataset_dtype_backend(stub_artifact_storage):
+    df = lazy.pd.DataFrame({"id": [1, 2], "data": [{"some_list": ["yes"]}, {"some_list": ["no"]}]})
+    stub_artifact_storage.write_batch_to_parquet_file(0, df, BatchStage.FINAL_RESULT)
+
+    default = stub_artifact_storage.load_dataset()
+    assert isinstance(default["data"].dtype, lazy.pd.ArrowDtype)
+    lazy.pd.testing.assert_frame_equal(stub_artifact_storage.load_dataset(dtype_backend="pyarrow"), default)
+
+    plain = stub_artifact_storage.load_dataset(dtype_backend=None)
+    assert plain["data"].dtype == object
+    assert not isinstance(plain["id"].dtype, lazy.pd.ArrowDtype)
+    assert plain["data"].iloc[0]["some_list"].tolist() == ["yes"]
+
+    nullable = stub_artifact_storage.load_dataset(BatchStage.FINAL_RESULT, dtype_backend="numpy_nullable")
+    assert nullable["id"].dtype == lazy.pd.Int64Dtype()
+
+
 def test_artifact_storage_path_validation(stub_artifact_storage):
     assert stub_artifact_storage.artifact_path.is_absolute()
     assert stub_artifact_storage.base_dataset_path.is_absolute()
diff --git a/packages/data-designer/src/data_designer/interface/results.py b/packages/data-designer/src/data_designer/interface/results.py
index f0c1e0361..ecb6f34ac 100644
--- a/packages/data-designer/src/data_designer/interface/results.py
+++ b/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 the returned columns. `"pyarrow"` (the
+                default) returns `ArrowDtype` columns. `"numpy_nullable"` returns pandas nullable
+                dtypes. `None` applies pandas' default conversion: nested columns are object
+                (lists as NumPy arrays, structs as dicts), and scalar columns follow pandas' defaults
+                (text as `str`, integers with nulls as `float64`).
+
         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.
diff --git a/packages/data-designer/tests/interface/test_results.py b/packages/data-designer/tests/interface/test_results.py
index 200d43810..b59007ae6 100644
--- a/packages/data-designer/tests/interface/test_results.py
+++ b/packages/data-designer/tests/interface/test_results.py
@@ -18,7 +18,7 @@
 from data_designer.config.utils.errors import DatasetSampleDisplayError
 from data_designer.config.utils.visualization import display_sample_record as display_fn
 from data_designer.engine.dataset_builders.errors import ArtifactStorageError
-from data_designer.engine.storage.artifact_storage import ArtifactStorage
+from data_designer.engine.storage.artifact_storage import ArtifactStorage, BatchStage
 from data_designer.interface.results import DatasetCreationResults
 
 
@@ -72,6 +72,31 @@ def test_load_dataset(stub_dataset_creation_results, stub_artifact_storage, stub
     lazy.pd.testing.assert_frame_equal(dataset, stub_dataframe)
 
 
+def test_load_dataset_passes_dtype_backend_to_artifact_storage(stub_dataset_creation_results, stub_artifact_storage):
+    """The default keeps the pyarrow backend; explicit values and None are passed through unchanged."""
+    stub_dataset_creation_results.load_dataset()
+    stub_artifact_storage.load_dataset.assert_called_with(dtype_backend="pyarrow")
+
+    for dtype_backend in ("numpy_nullable", None):
+        stub_dataset_creation_results.load_dataset(dtype_backend=dtype_backend)
+        stub_artifact_storage.load_dataset.assert_called_with(dtype_backend=dtype_backend)
+
+
+def test_load_dataset_without_dtype_backend_returns_object_nested_columns(tmp_path, stub_dataset_creation_results):
+    storage = ArtifactStorage(artifact_path=tmp_path)
+    frame = lazy.pd.DataFrame({"payload": [{"alpha": 1}], "tags": [["a", "b"]]})
+    storage.write_batch_to_parquet_file(0, frame, BatchStage.FINAL_RESULT)
+    stub_dataset_creation_results.artifact_storage = storage
+
+    default = stub_dataset_creation_results.load_dataset()
+    plain = stub_dataset_creation_results.load_dataset(dtype_backend=None)
+
+    assert isinstance(default["payload"].dtype, lazy.pd.ArrowDtype)
+    assert plain["payload"].dtype == object
+    assert plain["payload"].iloc[0] == {"alpha": 1}
+    assert plain["tags"].dtype == object
+
+
 def test_load_analysis(stub_dataset_creation_results, stub_dataset_profiler_results):
     """Test loading the analysis results."""
     analysis = stub_dataset_creation_results.load_analysis()

From b90c076fa404fde713c8c06f91e6da86dd5c5576 Mon Sep 17 00:00:00 2001
From: Aaron Gonzales <aagonzales@nvidia.com>
Date: Tue, 3 Nov 2026 05:16:56 +0000
Subject: [PATCH 2/5] fix: clarify result dtype conversion contract

Document that dtype_backend=None follows pandas conversion and stored
metadata. Encoded JSON columns retain object dtype for every backend;
loading does not repair incompatible persisted Arrow dtype metadata.

Extend the public result-loading regression through a Parquet write and
pandas default read for None and numpy_nullable. Preserve the existing
pyarrow default.

Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>
---
 .../src/data_designer/config/utils/io_helpers.py    | 10 +++++-----
 .../src/data_designer/interface/results.py          | 10 +++++-----
 .../data-designer/tests/interface/test_results.py   | 13 +++++++++++--
 3 files changed, 21 insertions(+), 12 deletions(-)

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 45da7488b..7027b2297 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
@@ -121,11 +121,11 @@ def read_parquet_dataset(
 
     Args:
         path: The path to the parquet dataset, can be either a file or a directory.
-        dtype_backend: The pandas dtype backend for the returned columns. `"pyarrow"` (the
-            default) and `"numpy_nullable"` are passed to pandas. `None` omits the option so
-            pandas applies its default conversion: nested columns are object (lists as NumPy
-            arrays, structs as dicts), while scalar columns follow pandas' defaults (text as
-            `str`, integers with nulls as `float64`).
+        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.
diff --git a/packages/data-designer/src/data_designer/interface/results.py b/packages/data-designer/src/data_designer/interface/results.py
index ecb6f34ac..ff70e0c46 100644
--- a/packages/data-designer/src/data_designer/interface/results.py
+++ b/packages/data-designer/src/data_designer/interface/results.py
@@ -92,11 +92,11 @@ def load_dataset(self, *, dtype_backend: Literal["pyarrow", "numpy_nullable"] |
         """Load the generated dataset as a pandas DataFrame.
 
         Args:
-            dtype_backend: The pandas dtype backend for the returned columns. `"pyarrow"` (the
-                default) returns `ArrowDtype` columns. `"numpy_nullable"` returns pandas nullable
-                dtypes. `None` applies pandas' default conversion: nested columns are object
-                (lists as NumPy arrays, structs as dicts), and scalar columns follow pandas' defaults
-                (text as `str`, integers with nulls as `float64`).
+            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.
diff --git a/packages/data-designer/tests/interface/test_results.py b/packages/data-designer/tests/interface/test_results.py
index b59007ae6..972830660 100644
--- a/packages/data-designer/tests/interface/test_results.py
+++ b/packages/data-designer/tests/interface/test_results.py
@@ -82,19 +82,28 @@ def test_load_dataset_passes_dtype_backend_to_artifact_storage(stub_dataset_crea
         stub_artifact_storage.load_dataset.assert_called_with(dtype_backend=dtype_backend)
 
 
-def test_load_dataset_without_dtype_backend_returns_object_nested_columns(tmp_path, stub_dataset_creation_results):
+@pytest.mark.parametrize("dtype_backend", [None, "numpy_nullable"], ids=["pandas-default", "nullable"])
+def test_load_dataset_nested_columns_round_trip_through_pandas(
+    tmp_path: Path, stub_dataset_creation_results: DatasetCreationResults, dtype_backend: str | None
+) -> None:
+    """A caller can opt out of Arrow dtypes before exporting ordinary nested artifacts."""
     storage = ArtifactStorage(artifact_path=tmp_path)
     frame = lazy.pd.DataFrame({"payload": [{"alpha": 1}], "tags": [["a", "b"]]})
     storage.write_batch_to_parquet_file(0, frame, BatchStage.FINAL_RESULT)
     stub_dataset_creation_results.artifact_storage = storage
 
     default = stub_dataset_creation_results.load_dataset()
-    plain = stub_dataset_creation_results.load_dataset(dtype_backend=None)
+    plain = stub_dataset_creation_results.load_dataset(dtype_backend=dtype_backend)
+    export_path = tmp_path / "consumer.parquet"
+    plain.to_parquet(export_path, index=False)
+    restored = lazy.pd.read_parquet(export_path)
 
     assert isinstance(default["payload"].dtype, lazy.pd.ArrowDtype)
     assert plain["payload"].dtype == object
     assert plain["payload"].iloc[0] == {"alpha": 1}
     assert plain["tags"].dtype == object
+    assert restored["payload"].iloc[0] == {"alpha": 1}
+    assert list(restored["tags"].iloc[0]) == ["a", "b"]
 
 
 def test_load_analysis(stub_dataset_creation_results, stub_dataset_profiler_results):

From 8e1f24150739a5917e67690ef5435274ce6ec236 Mon Sep 17 00:00:00 2001
From: Aaron Gonzales <aagonzales@nvidia.com>
Date: Tue, 3 Nov 2026 06:13:09 +0000
Subject: [PATCH 3/5] refactor: focus dtype loading on its API contract

Use an explicit typed nullable dtype mapping in encoded storage.
Consolidate public coverage around real result loading for each backend
and remove redundant pandas export/read-back coverage.

Preserve the Arrow default and decoded JSON values. Export compatibility
and nested-container normalization remain outside the loading contract.

Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>
---
 .../config/utils/storage_encoding.py          | 36 ++++++++--------
 .../tests/config/utils/test_io_helpers.py     | 10 -----
 .../tests/interface/test_results.py           | 41 ++++++++-----------
 3 files changed, 36 insertions(+), 51 deletions(-)

diff --git a/packages/data-designer-config/src/data_designer/config/utils/storage_encoding.py b/packages/data-designer-config/src/data_designer/config/utils/storage_encoding.py
index dc775ce39..cbf4c775f 100644
--- a/packages/data-designer-config/src/data_designer/config/utils/storage_encoding.py
+++ b/packages/data-designer-config/src/data_designer/config/utils/storage_encoding.py
@@ -117,21 +117,6 @@ def encode_storage_table(dataframe: pd.DataFrame, plan: JsonStoragePlan) -> pa.T
     return table
 
 
-def _nullable_types_mapper() -> dict[pa.DataType, Any]:
-    pa, pd = lazy.pa, lazy.pd
-    mapping: dict[pa.DataType, Any] = {
-        pa.bool_(): pd.BooleanDtype(),
-        pa.string(): pd.StringDtype(),
-        pa.large_string(): pd.StringDtype(),
-    }
-    for bits in (8, 16, 32, 64):
-        mapping[getattr(pa, f"int{bits}")()] = getattr(pd, f"Int{bits}Dtype")()
-        mapping[getattr(pa, f"uint{bits}")()] = getattr(pd, f"UInt{bits}Dtype")()
-    mapping[pa.float32()] = pd.Float32Dtype()
-    mapping[pa.float64()] = pd.Float64Dtype()
-    return mapping
-
-
 def decode_storage_table(
     table: pa.Table, *, dtype_backend: Literal["pyarrow", "numpy_nullable"] | None = "pyarrow"
 ) -> pd.DataFrame:
@@ -168,7 +153,7 @@ def decode_storage_table(
                     )
             frame = frame.reindex(columns=template.columns)
     elif dtype_backend == "numpy_nullable":
-        frame = table.to_pandas(types_mapper=_nullable_types_mapper().get)
+        frame = table.to_pandas(types_mapper=_get_numpy_nullable_type_mapping().get)
     elif dtype_backend is None:
         frame = table.to_pandas()
     else:
@@ -250,6 +235,25 @@ def _prepare_table_for_pandas(table: pa.Table) -> pa.Table:
     return table.replace_schema_metadata(updated)
 
 
+def _get_numpy_nullable_type_mapping() -> dict[pa.DataType, pd.api.extensions.ExtensionDtype]:
+    pa, pd = lazy.pa, lazy.pd
+    return {
+        pa.bool_(): pd.BooleanDtype(),
+        pa.int8(): pd.Int8Dtype(),
+        pa.int16(): pd.Int16Dtype(),
+        pa.int32(): pd.Int32Dtype(),
+        pa.int64(): pd.Int64Dtype(),
+        pa.uint8(): pd.UInt8Dtype(),
+        pa.uint16(): pd.UInt16Dtype(),
+        pa.uint32(): pd.UInt32Dtype(),
+        pa.uint64(): pd.UInt64Dtype(),
+        pa.float32(): pd.Float32Dtype(),
+        pa.float64(): pd.Float64Dtype(),
+        pa.string(): pd.StringDtype(),
+        pa.large_string(): pd.StringDtype(),
+    }
+
+
 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
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 6d2d5c4eb..c302ed569 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
@@ -84,16 +84,6 @@ def test_read_parquet_dataset_without_dtype_backend_returns_object_columns(tmp_p
     assert list(result["tags"].iloc[0]) == ["a", "b"]
 
 
-def test_read_parquet_dataset_without_dtype_backend_round_trips_through_pandas(tmp_path) -> None:
-    result = read_parquet_dataset(_write_simple_nested_file(tmp_path), dtype_backend=None)
-
-    result.to_parquet(tmp_path / "round_trip.parquet", index=False)
-    restored = lazy.pd.read_parquet(tmp_path / "round_trip.parquet")
-
-    assert restored["meta"].iloc[0] == {"score": 9, "label": "x"}
-    assert list(restored["tags"].iloc[0]) == ["a", "b"]
-
-
 @pytest.mark.parametrize("as_directory", [True, False], ids=["directory", "file"])
 def test_read_parquet_dataset_accepts_numpy_nullable(tmp_path, as_directory: bool) -> None:
     file_path = _write_simple_nested_file(tmp_path)
diff --git a/packages/data-designer/tests/interface/test_results.py b/packages/data-designer/tests/interface/test_results.py
index 972830660..3b58e1705 100644
--- a/packages/data-designer/tests/interface/test_results.py
+++ b/packages/data-designer/tests/interface/test_results.py
@@ -72,38 +72,29 @@ def test_load_dataset(stub_dataset_creation_results, stub_artifact_storage, stub
     lazy.pd.testing.assert_frame_equal(dataset, stub_dataframe)
 
 
-def test_load_dataset_passes_dtype_backend_to_artifact_storage(stub_dataset_creation_results, stub_artifact_storage):
-    """The default keeps the pyarrow backend; explicit values and None are passed through unchanged."""
-    stub_dataset_creation_results.load_dataset()
-    stub_artifact_storage.load_dataset.assert_called_with(dtype_backend="pyarrow")
-
-    for dtype_backend in ("numpy_nullable", None):
-        stub_dataset_creation_results.load_dataset(dtype_backend=dtype_backend)
-        stub_artifact_storage.load_dataset.assert_called_with(dtype_backend=dtype_backend)
-
-
-@pytest.mark.parametrize("dtype_backend", [None, "numpy_nullable"], ids=["pandas-default", "nullable"])
-def test_load_dataset_nested_columns_round_trip_through_pandas(
+@pytest.mark.parametrize("dtype_backend", ["pyarrow", "numpy_nullable", None])
+def test_load_dataset_selects_native_dtypes(
     tmp_path: Path, stub_dataset_creation_results: DatasetCreationResults, dtype_backend: str | None
 ) -> None:
-    """A caller can opt out of Arrow dtypes before exporting ordinary nested artifacts."""
+    """The public loader selects native dtypes without changing nested values."""
     storage = ArtifactStorage(artifact_path=tmp_path)
-    frame = lazy.pd.DataFrame({"payload": [{"alpha": 1}], "tags": [["a", "b"]]})
+    frame = lazy.pd.DataFrame({"id": [1], "payload": [{"alpha": 1}], "tags": [["a", "b"]]})
     storage.write_batch_to_parquet_file(0, frame, BatchStage.FINAL_RESULT)
     stub_dataset_creation_results.artifact_storage = storage
 
     default = stub_dataset_creation_results.load_dataset()
-    plain = stub_dataset_creation_results.load_dataset(dtype_backend=dtype_backend)
-    export_path = tmp_path / "consumer.parquet"
-    plain.to_parquet(export_path, index=False)
-    restored = lazy.pd.read_parquet(export_path)
-
-    assert isinstance(default["payload"].dtype, lazy.pd.ArrowDtype)
-    assert plain["payload"].dtype == object
-    assert plain["payload"].iloc[0] == {"alpha": 1}
-    assert plain["tags"].dtype == object
-    assert restored["payload"].iloc[0] == {"alpha": 1}
-    assert list(restored["tags"].iloc[0]) == ["a", "b"]
+    result = stub_dataset_creation_results.load_dataset(dtype_backend=dtype_backend)
+
+    assert all(isinstance(dtype, lazy.pd.ArrowDtype) for dtype in default.dtypes)
+    if dtype_backend == "pyarrow":
+        lazy.pd.testing.assert_frame_equal(result, default)
+    else:
+        assert result["id"].dtype == (lazy.pd.Int64Dtype() if dtype_backend == "numpy_nullable" else "int64")
+        assert result["payload"].dtype == object
+        assert result["tags"].dtype == object
+    assert result["id"].tolist() == [1]
+    assert result["payload"].iloc[0] == {"alpha": 1}
+    assert list(result["tags"].iloc[0]) == ["a", "b"]
 
 
 def test_load_analysis(stub_dataset_creation_results, stub_dataset_profiler_results):

From bf1697c013f5471a993bfcba9571cc23ff042a4b Mon Sep 17 00:00:00 2001
From: Aaron Gonzales <aagonzales@nvidia.com>
Date: Tue, 3 Nov 2026 06:36:42 +0000
Subject: [PATCH 4/5] test: consolidate result dtype coverage

Exercise native and JSON-encoded backend selection through real result
storage, and reuse the numeric-promotion regression across backends.
Keep focused fallback and invalid-input checks while removing duplicate
storage/reader tests and their setup helpers.

Compare dtype objects directly so the tests express conversion semantics
across supported pandas versions.

Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>
---
 .../tests/config/utils/test_io_helpers.py     | 91 ++++---------------
 .../config/utils/test_storage_encoding.py     | 27 ------
 .../engine/storage/test_artifact_storage.py   | 17 ----
 .../tests/interface/test_results.py           | 49 ++++++----
 4 files changed, 49 insertions(+), 135 deletions(-)

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 c302ed569..bd7d4d7fd 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
@@ -23,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}}]},
@@ -37,96 +38,36 @@ 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)
 
 
-def _write_nested_batches(path: Path, evaluations: list[dict]) -> None:
-    for batch_number, evaluation in enumerate(evaluations):
-        lazy.pd.DataFrame({"id": [batch_number], "qa_evaluations": [evaluation]}).to_parquet(
-            path / f"batch_{batch_number:05d}.parquet",
-            index=False,
-        )
-
-
-def _write_simple_nested_file(path: Path) -> Path:
-    file_path = path / "batch_00000.parquet"
-    lazy.pd.DataFrame({"id": [1], "meta": [{"score": 9, "label": "x"}], "tags": [["a", "b"]]}).to_parquet(
-        file_path, index=False
-    )
-    return file_path
-
-
-@pytest.mark.parametrize("as_directory", [True, False], ids=["directory", "file"])
-def test_read_parquet_dataset_defaults_to_pyarrow_dtypes(tmp_path, as_directory: bool) -> None:
-    file_path = _write_simple_nested_file(tmp_path)
-
-    result = read_parquet_dataset(tmp_path if as_directory else file_path)
-
-    assert all(isinstance(dtype, lazy.pd.ArrowDtype) for dtype in result.dtypes)
-    assert result["meta"].iloc[0] == {"score": 9, "label": "x"}
-    assert result["tags"].iloc[0] == ["a", "b"]
-
-
-@pytest.mark.parametrize("as_directory", [True, False], ids=["directory", "file"])
-def test_read_parquet_dataset_without_dtype_backend_returns_object_columns(tmp_path, as_directory: bool) -> None:
-    file_path = _write_simple_nested_file(tmp_path)
-
-    result = read_parquet_dataset(tmp_path if as_directory else file_path, dtype_backend=None)
-
-    assert not any(isinstance(dtype, lazy.pd.ArrowDtype) for dtype in result.dtypes)
-    assert result["meta"].dtype == object
-    assert result["tags"].dtype == object
-    assert result["meta"].iloc[0] == {"score": 9, "label": "x"}
-    assert list(result["tags"].iloc[0]) == ["a", "b"]
-
-
-@pytest.mark.parametrize("as_directory", [True, False], ids=["directory", "file"])
-def test_read_parquet_dataset_accepts_numpy_nullable(tmp_path, as_directory: bool) -> None:
-    file_path = _write_simple_nested_file(tmp_path)
-
-    result = read_parquet_dataset(tmp_path if as_directory else file_path, dtype_backend="numpy_nullable")
-
-    assert result["id"].dtype == lazy.pd.Int64Dtype()
-    assert result["meta"].iloc[0] == {"score": 9, "label": "x"}
-
-
 @pytest.mark.parametrize(
-    ("dtype_backend", "id_check"),
+    ("dtype_backend", "expected_dtype"),
     [
-        ("pyarrow", lambda dtype: isinstance(dtype, lazy.pd.ArrowDtype)),
-        ("numpy_nullable", lambda dtype: dtype == lazy.pd.Int64Dtype()),
-        (None, lambda dtype: dtype == "int64"),
+        ("pyarrow", lazy.pd.ArrowDtype(lazy.pa.int64())),
+        ("numpy_nullable", lazy.pd.Int64Dtype()),
+        (None, lazy.np.dtype("int64")),
     ],
 )
-def test_read_parquet_dataset_schema_unification_fallback_honors_dtype_backend(
-    tmp_path, dtype_backend, id_check
+def test_read_parquet_dataset_fallback_honors_dtype_backend(
+    tmp_path: Path, dtype_backend: str | None, expected_dtype: object
 ) -> None:
-    _write_nested_batches(tmp_path, [{"evaluations": [{"score": 9}]}, {"evaluations": [{"score": 10}]}])
+    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 id_check(result["id"].dtype)
+    assert result["id"].dtype == expected_dtype
     assert result["id"].tolist() == [0, 1]
-    assert (result["qa_evaluations"].dtype == object) is (dtype_backend != "pyarrow")
-
-
-@pytest.mark.parametrize("dtype_backend", ["pyarrow", "numpy_nullable", None])
-def test_read_parquet_dataset_unifies_nested_numeric_types_for_every_backend(tmp_path, dtype_backend) -> None:
-    _write_nested_batches(
-        tmp_path,
-        [{"evaluations": [{"overall": {"score": 9}}]}, {"evaluations": [{"overall": {"score": 9.5}}]}],
-    )
-
-    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)
+    assert result["meta"].tolist() == [{"score": 9}, {"score": 9}]
+    assert (result["meta"].dtype == object) is (dtype_backend != "pyarrow")
 
 
 def test_smart_load_yaml():
diff --git a/packages/data-designer-config/tests/config/utils/test_storage_encoding.py b/packages/data-designer-config/tests/config/utils/test_storage_encoding.py
index 8bc5f27bd..fa48b1d5f 100644
--- a/packages/data-designer-config/tests/config/utils/test_storage_encoding.py
+++ b/packages/data-designer-config/tests/config/utils/test_storage_encoding.py
@@ -514,33 +514,6 @@ def test_legacy_partitioned_directory_keeps_native_columns(tmp_path: Path) -> No
     assert JSON_STORAGE_RUN_KEY not in result.attrs
 
 
-@pytest.mark.parametrize("as_directory", [True, False], ids=["directory", "file"])
-@pytest.mark.parametrize(
-    ("dtype_backend", "count_check", "label_check"),
-    [
-        ("pyarrow", lambda dtype: isinstance(dtype, pd.ArrowDtype), lambda dtype: isinstance(dtype, pd.ArrowDtype)),
-        ("numpy_nullable", lambda dtype: dtype == pd.Int64Dtype(), lambda dtype: dtype == pd.StringDtype()),
-        (None, lambda dtype: dtype == "int64", lambda dtype: pd.api.types.is_string_dtype(dtype)),
-    ],
-)
-def test_encoded_dataset_honors_dtype_backend_for_native_columns(
-    tmp_path: Path, as_directory: bool, dtype_backend: Any, count_check: Any, label_check: Any
-) -> None:
-    values = [{"alpha": 1}, {"beta": [1, "two"]}]
-    frame = pd.DataFrame({"count": [1, 2], "label": ["a", "b"], "payload": pd.Series(values, dtype=object)})
-    plan = JsonStoragePlan(columns={"payload": [[]]})
-    path = tmp_path / "batch_0.parquet"
-    pq.write_table(encode_storage_table(frame, plan), path)
-
-    result = read_parquet_dataset(tmp_path if as_directory else path, dtype_backend=dtype_backend)
-
-    assert count_check(result["count"].dtype)
-    assert label_check(result["label"].dtype)
-    assert result["payload"].dtype == object
-    assert result["payload"].tolist() == values
-    assert result.attrs[JSON_STORAGE_RUN_KEY] == plan.model_dump()
-
-
 def test_decode_storage_table_rejects_unknown_dtype_backend() -> None:
     table = pa.Table.from_pylist([{"payload": "1"}]).replace_schema_metadata(
         {JSON_STORAGE_METADATA_KEY: JsonStoragePlan(columns={"payload": [[]]}).model_dump_json().encode()}
diff --git a/packages/data-designer-engine/tests/engine/storage/test_artifact_storage.py b/packages/data-designer-engine/tests/engine/storage/test_artifact_storage.py
index de31e27f8..9b9343f29 100644
--- a/packages/data-designer-engine/tests/engine/storage/test_artifact_storage.py
+++ b/packages/data-designer-engine/tests/engine/storage/test_artifact_storage.py
@@ -194,23 +194,6 @@ def test_artifact_storage_read_parquet_files(stub_artifact_storage):
     lazy.pd.testing.assert_frame_equal(lazy.pd.concat([read_df1, read_df2], ignore_index=True), read_df)
 
 
-def test_artifact_storage_load_dataset_dtype_backend(stub_artifact_storage):
-    df = lazy.pd.DataFrame({"id": [1, 2], "data": [{"some_list": ["yes"]}, {"some_list": ["no"]}]})
-    stub_artifact_storage.write_batch_to_parquet_file(0, df, BatchStage.FINAL_RESULT)
-
-    default = stub_artifact_storage.load_dataset()
-    assert isinstance(default["data"].dtype, lazy.pd.ArrowDtype)
-    lazy.pd.testing.assert_frame_equal(stub_artifact_storage.load_dataset(dtype_backend="pyarrow"), default)
-
-    plain = stub_artifact_storage.load_dataset(dtype_backend=None)
-    assert plain["data"].dtype == object
-    assert not isinstance(plain["id"].dtype, lazy.pd.ArrowDtype)
-    assert plain["data"].iloc[0]["some_list"].tolist() == ["yes"]
-
-    nullable = stub_artifact_storage.load_dataset(BatchStage.FINAL_RESULT, dtype_backend="numpy_nullable")
-    assert nullable["id"].dtype == lazy.pd.Int64Dtype()
-
-
 def test_artifact_storage_path_validation(stub_artifact_storage):
     assert stub_artifact_storage.artifact_path.is_absolute()
     assert stub_artifact_storage.base_dataset_path.is_absolute()
diff --git a/packages/data-designer/tests/interface/test_results.py b/packages/data-designer/tests/interface/test_results.py
index 3b58e1705..08feb0186 100644
--- a/packages/data-designer/tests/interface/test_results.py
+++ b/packages/data-designer/tests/interface/test_results.py
@@ -16,6 +16,8 @@
 from data_designer.config.errors import InvalidFileFormatError
 from data_designer.config.preview_results import PreviewResults
 from data_designer.config.utils.errors import DatasetSampleDisplayError
+from data_designer.config.utils.io_helpers import read_parquet_dataset
+from data_designer.config.utils.storage_schema import JSON_STORAGE_RUN_KEY, JsonStoragePlan
 from data_designer.config.utils.visualization import display_sample_record as display_fn
 from data_designer.engine.dataset_builders.errors import ArtifactStorageError
 from data_designer.engine.storage.artifact_storage import ArtifactStorage, BatchStage
@@ -72,29 +74,44 @@ def test_load_dataset(stub_dataset_creation_results, stub_artifact_storage, stub
     lazy.pd.testing.assert_frame_equal(dataset, stub_dataframe)
 
 
-@pytest.mark.parametrize("dtype_backend", ["pyarrow", "numpy_nullable", None])
+@pytest.mark.parametrize("encoded", [False, True], ids=["native", "json-encoded"])
+@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_load_dataset_selects_native_dtypes(
-    tmp_path: Path, stub_dataset_creation_results: DatasetCreationResults, dtype_backend: str | None
+    tmp_path: Path,
+    stub_dataset_creation_results: DatasetCreationResults,
+    encoded: bool,
+    dtype_backend: str | None,
+    expected_dtype: object,
 ) -> None:
-    """The public loader selects native dtypes without changing nested values."""
-    storage = ArtifactStorage(artifact_path=tmp_path)
-    frame = lazy.pd.DataFrame({"id": [1], "payload": [{"alpha": 1}], "tags": [["a", "b"]]})
-    storage.write_batch_to_parquet_file(0, frame, BatchStage.FINAL_RESULT)
+    plan = JsonStoragePlan(columns={"payload": [[]]} if encoded else {})
+    storage = ArtifactStorage(artifact_path=tmp_path, json_storage=plan)
+    payloads = [{"alpha": 1}, {"beta": [1, "two"]}] if encoded else [{"alpha": 1}, {"alpha": 2}]
+    frame = lazy.pd.DataFrame({"id": [1, 2], "label": ["a", "b"], "payload": payloads, "tags": [["x"], ["y"]]})
+    path = storage.write_batch_to_parquet_file(0, frame, BatchStage.FINAL_RESULT)
     stub_dataset_creation_results.artifact_storage = storage
 
     default = stub_dataset_creation_results.load_dataset()
     result = stub_dataset_creation_results.load_dataset(dtype_backend=dtype_backend)
 
-    assert all(isinstance(dtype, lazy.pd.ArrowDtype) for dtype in default.dtypes)
-    if dtype_backend == "pyarrow":
-        lazy.pd.testing.assert_frame_equal(result, default)
-    else:
-        assert result["id"].dtype == (lazy.pd.Int64Dtype() if dtype_backend == "numpy_nullable" else "int64")
-        assert result["payload"].dtype == object
-        assert result["tags"].dtype == object
-    assert result["id"].tolist() == [1]
-    assert result["payload"].iloc[0] == {"alpha": 1}
-    assert list(result["tags"].iloc[0]) == ["a", "b"]
+    assert isinstance(default["id"].dtype, lazy.pd.ArrowDtype)
+    assert result["id"].dtype == expected_dtype
+    assert result["id"].tolist() == [1, 2]
+    assert result["label"].tolist() == ["a", "b"]
+    if dtype_backend == "numpy_nullable":
+        assert result["label"].dtype == lazy.pd.StringDtype()
+    assert (result["payload"].dtype == object) is (encoded or dtype_backend != "pyarrow")
+    assert result["payload"].tolist() == payloads
+    assert (result["tags"].dtype == object) is (dtype_backend != "pyarrow")
+    assert [list(tags) for tags in result["tags"]] == [["x"], ["y"]]
+    assert result.attrs.get(JSON_STORAGE_RUN_KEY) == (plan.model_dump() if encoded else None)
+    lazy.pd.testing.assert_frame_equal(result, read_parquet_dataset(path, dtype_backend=dtype_backend))
 
 
 def test_load_analysis(stub_dataset_creation_results, stub_dataset_profiler_results):

From 508dcfc6072cdf60d5a34e6c9f5deaf619cbdd2b Mon Sep 17 00:00:00 2001
From: Aaron Gonzales <aagonzales@nvidia.com>
Date: Tue, 10 Nov 2026 23:39:20 +0000
Subject: [PATCH 5/5] fix: share native parquet conversion

Preserve Arrow buffers through the documented conversion hook and
reuse pandas public conversion for the non-Arrow backends. Resolve
native fragment schemas before conversion while restoring JSON values
from their original fragments.

Validate backend choices before artifact access and strengthen the
existing public and nested-array regressions without adding test
functions to the previous PR.

Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>
---
 .../data_designer/config/utils/io_helpers.py  |  18 +-
 .../config/utils/storage_encoding.py          | 155 +++++++++++-------
 .../config/utils/test_storage_encoding.py     |  15 +-
 .../tests/interface/test_results.py           |  71 +++++---
 4 files changed, 160 insertions(+), 99 deletions(-)

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 7027b2297..936a7c188 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
@@ -18,7 +18,12 @@
 
 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
@@ -130,7 +135,8 @@ def read_parquet_dataset(
     Returns:
         The parquet dataset as a pandas DataFrame.
     """
-    read_kwargs: dict[str, Any] = {} if dtype_backend is None else {"dtype_backend": dtype_backend}
+    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]
@@ -142,15 +148,17 @@ def read_parquet_dataset(
         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, **read_kwargs) 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, schema=unified_schema, partitioning=dataset.partitioning, **read_kwargs)
+        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], dtype_backend=dtype_backend)
-    return lazy.pd.read_parquet(path, **read_kwargs)
+    return read_native_parquet(path, dtype_backend=dtype_backend)
 
 
 def validate_dataset_file_path(file_path: str | Path, should_exist: bool = True) -> Path:
diff --git a/packages/data-designer-config/src/data_designer/config/utils/storage_encoding.py b/packages/data-designer-config/src/data_designer/config/utils/storage_encoding.py
index cbf4c775f..4cb6ed5b3 100644
--- a/packages/data-designer-config/src/data_designer/config/utils/storage_encoding.py
+++ b/packages/data-designer-config/src/data_designer/config/utils/storage_encoding.py
@@ -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:
@@ -117,47 +118,13 @@ def encode_storage_table(dataframe: pd.DataFrame, plan: JsonStoragePlan) -> pa.T
     return table
 
 
-def decode_storage_table(
-    table: pa.Table, *, dtype_backend: Literal["pyarrow", "numpy_nullable"] | None = "pyarrow"
-) -> pd.DataFrame:
+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)
-    column_labels: dict[str, Any] = {}
-    if dtype_backend == "pyarrow":
-        # 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)
-        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)
-    elif dtype_backend == "numpy_nullable":
-        frame = table.to_pandas(types_mapper=_get_numpy_nullable_type_mapping().get)
-    elif dtype_backend is None:
-        frame = table.to_pandas()
-    else:
-        raise ValueError(f"dtype_backend must be 'pyarrow', 'numpy_nullable' or None, got {dtype_backend!r}")
+    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()):
@@ -167,7 +134,7 @@ def decode_storage_table(
                 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
@@ -194,17 +161,34 @@ def read_encoded_parquet(
     *,
     dtype_backend: Literal["pyarrow", "numpy_nullable"] | None = "pyarrow",
 ) -> pd.DataFrame:
-    """Read annotated files individually to retain and validate each file's metadata."""
+    """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), dtype_backend=dtype_backend))
-    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
 
@@ -235,23 +219,70 @@ def _prepare_table_for_pandas(table: pa.Table) -> pa.Table:
     return table.replace_schema_metadata(updated)
 
 
-def _get_numpy_nullable_type_mapping() -> dict[pa.DataType, pd.api.extensions.ExtensionDtype]:
-    pa, pd = lazy.pa, lazy.pd
-    return {
-        pa.bool_(): pd.BooleanDtype(),
-        pa.int8(): pd.Int8Dtype(),
-        pa.int16(): pd.Int16Dtype(),
-        pa.int32(): pd.Int32Dtype(),
-        pa.int64(): pd.Int64Dtype(),
-        pa.uint8(): pd.UInt8Dtype(),
-        pa.uint16(): pd.UInt16Dtype(),
-        pa.uint32(): pd.UInt32Dtype(),
-        pa.uint64(): pd.UInt64Dtype(),
-        pa.float32(): pd.Float32Dtype(),
-        pa.float64(): pd.Float64Dtype(),
-        pa.string(): pd.StringDtype(),
-        pa.large_string(): pd.StringDtype(),
-    }
+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:
diff --git a/packages/data-designer-config/tests/config/utils/test_storage_encoding.py b/packages/data-designer-config/tests/config/utils/test_storage_encoding.py
index fa48b1d5f..561f4244f 100644
--- a/packages/data-designer-config/tests/config/utils/test_storage_encoding.py
+++ b/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
 
@@ -514,14 +515,6 @@ def test_legacy_partitioned_directory_keeps_native_columns(tmp_path: Path) -> No
     assert JSON_STORAGE_RUN_KEY not in result.attrs
 
 
-def test_decode_storage_table_rejects_unknown_dtype_backend() -> None:
-    table = pa.Table.from_pylist([{"payload": "1"}]).replace_schema_metadata(
-        {JSON_STORAGE_METADATA_KEY: JsonStoragePlan(columns={"payload": [[]]}).model_dump_json().encode()}
-    )
-    with pytest.raises(ValueError, match="dtype_backend must be"):
-        decode_storage_table(table, dtype_backend="bogus")  # type: ignore[arg-type]
-
-
 def test_encoded_partitioned_directory_restores_json_and_partition_columns(tmp_path: Path) -> None:
     values = [{"alpha": 1}, {"beta": [1, "two"]}, None]
     frame = pd.DataFrame({"group": ["a", "b", "c"], "payload": pd.Series(values, dtype=object)})
diff --git a/packages/data-designer/tests/interface/test_results.py b/packages/data-designer/tests/interface/test_results.py
index 08feb0186..73bdafcbc 100644
--- a/packages/data-designer/tests/interface/test_results.py
+++ b/packages/data-designer/tests/interface/test_results.py
@@ -90,28 +90,57 @@ def test_load_dataset_selects_native_dtypes(
     dtype_backend: str | None,
     expected_dtype: object,
 ) -> None:
-    plan = JsonStoragePlan(columns={"payload": [[]]} if encoded else {})
-    storage = ArtifactStorage(artifact_path=tmp_path, json_storage=plan)
-    payloads = [{"alpha": 1}, {"beta": [1, "two"]}] if encoded else [{"alpha": 1}, {"alpha": 2}]
-    frame = lazy.pd.DataFrame({"id": [1, 2], "label": ["a", "b"], "payload": payloads, "tags": [["x"], ["y"]]})
-    path = storage.write_batch_to_parquet_file(0, frame, BatchStage.FINAL_RESULT)
+    with lazy.pd.option_context("future.infer_string", False):
+        plan = JsonStoragePlan(columns={"payload": [[]]} if encoded else {})
+        storage = ArtifactStorage(artifact_path=tmp_path, json_storage=plan)
+        payloads = [{"alpha": 1}, {"beta": [1, "two"]}] if encoded else [{"alpha": 1}, {"alpha": 2}]
+        timestamps = lazy.pd.date_range("2026-01-01T12:00:00+05:45", periods=2)
+        paths = []
+        for batch_number, payload in enumerate(payloads):
+            batch = lazy.pd.DataFrame(
+                {
+                    "id": [batch_number + 1],
+                    "nullable": [1 if batch_number == 0 else None],
+                    "label": ["a" if batch_number == 0 else None],
+                    "category": lazy.pd.Categorical(["a" if batch_number == 0 else "b"]),
+                    "timestamp": timestamps[batch_number : batch_number + 1],
+                    "payload": [payload],
+                    "tags": [["x" if batch_number == 0 else "y"]],
+                }
+            )
+            paths.append(storage.write_batch_to_parquet_file(batch_number, batch, BatchStage.FINAL_RESULT))
+        stub_dataset_creation_results.artifact_storage = storage
+        result = stub_dataset_creation_results.load_dataset(dtype_backend=dtype_backend)
+
+        assert isinstance(stub_dataset_creation_results.load_dataset()["id"].dtype, lazy.pd.ArrowDtype)
+        assert result["id"].dtype == expected_dtype
+        assert (result["payload"].dtype == object) is (encoded or dtype_backend != "pyarrow")
+        assert result["payload"].tolist() == payloads
+        assert (result["tags"].dtype == object) is (dtype_backend != "pyarrow")
+        assert [list(tags) for tags in result["tags"]] == [["x"], ["y"]]
+        assert result.attrs.get(JSON_STORAGE_RUN_KEY) == (plan.model_dump() if encoded else None)
+        kwargs = {} if dtype_backend is None else {"dtype_backend": dtype_backend}
+        scalar_columns = ["id", "nullable", "label", "category", "timestamp"]
+        for source in [storage.final_dataset_path, *paths]:
+            actual = (
+                result
+                if source == storage.final_dataset_path
+                else read_parquet_dataset(source, dtype_backend=dtype_backend)
+            )
+            expected = lazy.pd.read_parquet(source, columns=scalar_columns, **kwargs)
+            lazy.pd.testing.assert_frame_equal(actual[scalar_columns], expected)
+
+
+@pytest.mark.parametrize("has_artifacts", [False, True])
+def test_load_dataset_rejects_unknown_dtype_backend(
+    tmp_path: Path, stub_dataset_creation_results: DatasetCreationResults, has_artifacts: bool
+) -> None:
+    storage = ArtifactStorage(artifact_path=tmp_path)
+    if has_artifacts:
+        storage.write_batch_to_parquet_file(0, lazy.pd.DataFrame({"id": [1]}), BatchStage.FINAL_RESULT)
     stub_dataset_creation_results.artifact_storage = storage
-
-    default = stub_dataset_creation_results.load_dataset()
-    result = stub_dataset_creation_results.load_dataset(dtype_backend=dtype_backend)
-
-    assert isinstance(default["id"].dtype, lazy.pd.ArrowDtype)
-    assert result["id"].dtype == expected_dtype
-    assert result["id"].tolist() == [1, 2]
-    assert result["label"].tolist() == ["a", "b"]
-    if dtype_backend == "numpy_nullable":
-        assert result["label"].dtype == lazy.pd.StringDtype()
-    assert (result["payload"].dtype == object) is (encoded or dtype_backend != "pyarrow")
-    assert result["payload"].tolist() == payloads
-    assert (result["tags"].dtype == object) is (dtype_backend != "pyarrow")
-    assert [list(tags) for tags in result["tags"]] == [["x"], ["y"]]
-    assert result.attrs.get(JSON_STORAGE_RUN_KEY) == (plan.model_dump() if encoded else None)
-    lazy.pd.testing.assert_frame_equal(result, read_parquet_dataset(path, dtype_backend=dtype_backend))
+    with pytest.raises(ValueError, match="dtype_backend"):
+        stub_dataset_creation_results.load_dataset(dtype_backend="bogus")  # type: ignore[arg-type]
 
 
 def test_load_analysis(stub_dataset_creation_results, stub_dataset_profiler_results):