[binaryaaron]

📋 Summary

Let callers choose pandas dtypes through DatasetCreationResults.load_dataset(dtype_backend=...), avoiding direct artifact reads or post-load casts. Preserve the Arrow default and decoded JSON values.

🔗 Related Issue

N/A — no matching dtype-selection issue found. Upstream plan review is pending.

🔄 Changes

• Preserve "pyarrow" as the default; accept "numpy_nullable" and None for pandas-default conversion.
• Forward the choice through native and JSON-encoded reads, including file/directory paths and schema-unification fallback.
• Keep decoded JSON columns as objects. Persisted pandas metadata still applies; artifact serialization and nested-container normalization are unchanged.

🧪 Testing

• make test: 5,019 passed, 3 skipped.
• Unit tests consolidated: 271 passed on each of pandas 2.3.3 and 3.0.5 with PyArrow 24.0.0.
• make check-all, commit hooks, and diff checks.
• E2E: not run; no provider-backed behavior changes.

✅ Checklist

• Follows commit message conventions.
• All commits signed off (DCO).
• External-contributor issue triage, if applicable: no issue linked yet.
• Public API docstrings updated; no architecture documentation contradicted.

---

Prepared with AI assistance.

[github-actions]

Fern preview: https://nvidia-preview-pr-1128.docs.buildwithfern.com/nemo/datadesigner

Fern previews include the docs-website version archive with PR changes synced into latest. Notebook tutorials are rendered without execution outputs in previews.

[greptile-apps]

Confidence Score: 5/5

[Medium impact] The PR appears safe to merge; no actionable defects were found.

Summary

Adds dtype_backend to DatasetCreationResults.load_dataset() and forwards it through artifact loading.

• Dataset results can load with the pandas dtype style you choose.

binaryaaron explicitly noted that persisted pandas metadata still applies and nested-container normalization is unchanged.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  A["DatasetCreationResults.load_dataset(dtype_backend)"] --> B["ArtifactStorage.load_dataset"]
  B --> C["read_parquet_dataset"]
  C --> D{"JSON storage columns?"}
  D -->|No| E["pandas.read_parquet"]
  D -->|Yes| F["read_encoded_parquet"]
  F --> G["decode_storage_table"]
  G --> H["Convert native columns; restore JSON as objects"]

Loading

_{Reviews (1) · Last reviewed commit: "test: consolidate result dtype coverage" · Reviewed by Greptile}

[svcnvidia-nemo-ci]

No qualifying findings.

I reviewed all seven changed files against the base and traced the result-loading path through engine storage and both native and JSON-encoded readers. The public API preserves the Arrow default. Explicit backend choices reach file, directory, and schema-unification fallback reads, while decoded JSON columns retain object dtype.

The inspected regressions cover public loading through real storage, native nested numeric promotion, fallback conversion, and file/directory consistency. I couldn't execute tests in this read-only environment; this approval is based on source and test inspection.

Note: formal approval is not available for this request, so this review is posted as a comment only.

Review generated by codex. This was a full review. Request a full review with /review reset=true.

[nabinchha]

Thanks for putting this together, @binaryaaron!

Summary

PR #1128 adds dtype_backend selection to DatasetCreationResults.load_dataset and forwards it through artifact loading while preserving the Arrow default. Two P2 issues affect native-column dtypes when loading JSON-encoded artifacts.

Could you share the concrete downstream use case driving this API addition? Which consumer needs None or "numpy_nullable", and why is choosing the backend during loading needed rather than converting the loaded DataFrame?

Findings

Warnings — Worth addressing

packages/data-designer-config/src/data_designer/config/utils/storage_encoding.py:127–130 — [P2] Preserve pandas conversion semantics for encoded reads

• What: The new None and numpy_nullable branches directly call table.to_pandas, bypassing pandas.read_parquet post-conversion behavior. On pandas 3.0.5 / PyArrow 24.0.0, two supported cases differ from native reads. With future.infer_string=False and dtype_backend=None, an unrelated native string column returns as str instead of object. With a fixed-offset timestamp such as +05:45, both None and numpy_nullable preserve pytz.FixedOffset instead of pandas' normalized datetime.timezone.
• Why: Native columns change behavior solely because another column uses JSON storage. Assigning an integer to the string column succeeds for native artifacts but raises TypeError for encoded artifacts. Concatenating otherwise matching native and encoded fixed-offset timestamp columns produces object dtype, and subsequent .dt.hour raises AttributeError. Valid persisted metadata does not avoid these discrepancies; this is conversion behavior rather than repairing incompatible metadata.
• Suggestion: Match pandas' public conversion semantics, including string-inference settings and timezone normalization, while restoring declared JSON columns as objects. Add regression tests for disabled string inference and fixed-offset timestamps under both new backends. Avoid importing private pandas conversion helpers as a compatibility dependency.

packages/data-designer-config/src/data_designer/config/utils/storage_encoding.py:178 — [P2] Preserve categorical dtype when combining encoded batches

• What: Converting each fragment before pandas concatenation drops native categorical dtype when batches have different category vocabularies. Write two batches whose label columns are pd.Categorical(['a']) and pd.Categorical(['b']), alongside payload=[{'x': 1}]. With payload JSON-encoded, load_dataset(dtype_backend=None) and load_dataset(dtype_backend='numpy_nullable') return label as str. Native artifacts retain category with categories ['a', 'b']; the unchanged Arrow backend retains dictionary dtype in both cases.
• Why: The fragment metadata is valid, but pandas.concat cannot preserve categoricals with different vocabularies. Downstream .cat operations therefore fail only on the encoded path. Choosing JSON storage for payload should not discard label's native dtype. The added interface test writes only one batch, so its file/directory comparison does not cover this failure.
• Suggestion: Union compatible native categorical dictionaries before pandas concatenation, or defer native conversion until the Arrow fragments are combined. Preserve fragment-level JSON-contract validation. Add a real two-batch round-trip regression under both new backends.

What Looks Good

• Keyword-only forwarding preserves existing positional batch_stage and zero-argument callers.
• Decoded JSON values remain objects, including nested containers and large integers.
• The Arrow default, serialization format, and package dependency direction remain intact.

Residual Risk

Validation: 271 targeted tests passed on pandas 3.0.5 / PyArrow 24.0.0, and lint/format checks passed for all seven changed files. The string-inference and categorical failures were reproduced through DatasetCreationResults.load_dataset; the fixed-offset timestamp concatenation failure was also reproduced. The full suite and pandas 2.3.3 matrix reported by the author were not independently rerun.

Verdict

Needs changes — align pandas conversion semantics and preserve categorical dtype across encoded batches.

---

This review was generated by an AI assistant.

[binaryaaron]

agent (review-pr): Proposed reader fix and independent Arrow upgrade; neither is applied to this PR yet.

@nabinchha, I recommend fixing the reader here and reviewing the Arrow update independently. Both findings reproduce. The timezone mismatch specifically reproduces with pandas 3 / Arrow 24 when pytz is installed. Arrow's timezone change, released in Arrow 25, removes that mismatch. Category loss matches pandas' documented concatenation behavior.

The tested options are:

Option	Result
Upgrade Arrow alone	Fixes timezone behavior; leaves the other two failures.
Fix the reader on Arrow 24	Addresses all three reproduced failures without raising the minimum Arrow version.
Reader fix plus Arrow upgrade	Same reader correction, plus upstream timezone alignment. Keep these independently reviewable.

Proposed reader fix: validate each fragment's JSON contract, read native columns together through public pd.read_parquet, and restore JSON-bearing columns from the original Arrow fragments using the existing decoder. Exclude JSON-bearing columns from native schema unification. Retain sparse columns, partitions, fallback reads, row alignment, and JSON-only row counts. Remove the new nullable-type mapper and the decoder's backend argument, leaving pandas responsible for native conversion.

The shorter alternatives failed concrete cases: reading everything through pandas before JSON decoding rounded a native nested integer from 9007199254740993 to 9007199254740992.0; combining whole Arrow tables rejected valid JSON-bearing columns whose native siblings differed across fragments. Those are the reasons for retaining per-fragment restoration.

The prototype changes three files (+91/−61 relative to this head), with 9 net production lines added and no new test functions. Existing public-load cases gain two batches with differing categories, fixed-offset timestamps, and disabled string inference, comparing native columns against pandas. These expose failures on this head and pass in all six combinations of pandas 2.3.3/3.0.5 × Arrow 24.0.0/25.0.0/26.0.1. The broader matrix passed 468 checks per combination, including 60 exploratory probes kept outside the proposed patch.

Independent dependency proposal: widen all three constraints to pyarrow>=24,<27 and lock 26.0.1; no other locked package versions change. Its full-suite run yielded 5,018 passed, 3 skipped, and one CLI help-text assertion failure that also reproduces on Arrow 24 in the same environment. Other Python/OS CI combinations remain untested locally.

On motivation: this originated in Anonymizer's result/export integration, but this API does not fix its nested-Arrow metadata export failure. The intended benefit is choosing pandas' read semantics at the artifact boundary instead of manually reading files or casting afterward. A separate consumer requiring numpy_nullable has not been established; mirroring pandas' options is the current rationale. The API-scope question remains open independently of these defects.