[binaryaaron]

📋 Summary

Add keyword-only dtype_backend selection to DatasetCreationResults.load_dataset() while preserving the Arrow default and decoded JSON values. Native columns now have consistent conversion across ordinary and JSON-encoded artifacts, including multi-batch datasets.

Rebased onto main after #1142 merged.

🔗 Related Issue

Builds on merged #1142. No separate dtype-selection issue is linked.

🔄 Changes

• Accept "pyarrow" (default), "numpy_nullable", and None; validate invalid choices before artifact access.
• Share native Parquet reads across both storage paths. Delegate non-Arrow conversion to public pd.read_parquet, retaining its string-inference and timezone behavior.
• Replace manual nested-column/index reconstruction with pandas' documented Arrow conversion hook. Wrap original arrays directly and retain the required metadata repair. Returned arrays keep ordinary pandas Arrow dtypes.
• Unify native fragment schemas before conversion, preserving nullable types and categorical vocabularies. Restore declared JSON columns from raw fragments so nested values and large integers remain exact.
• Extend existing integration tests for two batches, all-null shards, strings, categories, timestamps, and native/encoded nested arrays. Keep the inherited metadata/persistence regressions; no additional test functions versus the previous PR revision.

🧪 Testing

• make test: 5,256 passed, 3 skipped on Python 3.14 / pandas 3.0.5 / Arrow 24.0.0.
• 573 targeted checks passed per pandas version (2.3.3 and 3.0.5), including 60 exploratory reader probes outside the patch, with optional pytz present.
• Final focused controls: 74 passed on pandas 2.3.3 / Arrow 24 and 74 passed on pandas 3.0.5 / Arrow 26.0.1.
• Red/green checks for dtype/validation failures and native nested-array corruption.
• make check-all, complexity checks, diff checks, and normal commit hooks.
• Provider-backed E2E: not run.

Reviewer Notes

Dependencies and the serialization format are unchanged. This API selects load-time representation; it does not promise to repair incompatible exported pandas metadata or normalize nested containers. The choices follow pandas' API; a separate consumer requiring numpy_nullable has not been established.

---

Description updated 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 correctness issues were found.

Summary

Adds dtype_backend to DatasetCreationResults.load_dataset() while keeping "pyarrow" as the default.

• Dataset loading lets callers choose pandas dtypes for native columns.

No actionable new issues were found. No previous review threads were supplied.

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["read_native_parquet"]
    D -->|Yes| F["Read native columns with selected backend"]
    D -->|Yes| G["Decode JSON columns as objects"]
    F --> H["Combine columns in original order"]
    G --> H

Loading

_{Reviews (2) · Last reviewed commit: "fix: share native parquet conversion" · 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.

[andreatnvidia]

It looks like numpy_nullable can lose its nullable dtype here on pandas 3. With batches containing x=[1] and x=[None], I get an object column when a JSON-encoded sibling is present, versus Int64 without it. Would resolving native column types before converting the fragments help? An all-null shard might be a useful regression case.

[binaryaaron]

agent (review-pr): Fixed in 508dcfc6: native fragment schemas are resolved before conversion, so an all-null shard no longer turns a nullable integer into object dtype. The existing public-load test now writes two batches, including this null-shard case, and compares the native columns with pandas' reader. Checked on pandas 2.3.3 and 3.0.5.

[andreatnvidia]

I noticed a small difference from the native reader on pandas 3 with future.infer_string=False: dtype_backend=None gives object strings with None for native files, but str columns with NaN for JSON-annotated files. This seems to come from skipping pandas’ post-conversion cleanup. It might be worth aligning these so a JSON-encoded sibling doesn’t affect unrelated string columns.

[binaryaaron]

agent (review-pr): Fixed in 508dcfc6: both readers now share native conversion, with None and numpy_nullable delegated to public pd.read_parquet. The public-load regression uses future.infer_string=False and null strings and checks values and dtypes against pandas, alongside fixed-offset timestamps and categories.

[andreatnvidia]

Maybe we could validate dtype_backend at this boundary? At the moment, an invalid value with missing artifacts produces FileNotFoundError before reaching argument validation. Checking it here would give callers a clearer error, and a small public-API test could cover that.

[binaryaaron]

agent (review-pr): Implemented in 508dcfc6: read_parquet_dataset validates the backend before artifact access, so the public API raises ValueError even when artifacts are missing. The previous internal validation test was moved to load_dataset and covers both artifact-present and artifact-missing cases.

[andreatnvidia]

@binaryaaron after your comment during our sync, I looked more closely at how this PR relates to #1142. Indeed, they’re related, but the fix in #1142 looked fine to merge separately.

For this PR, I think the reader approach you proposed is the way to go: read native columns together through pandas, then restore JSON columns from the original fragments. That should also cover the null-shard and string-conversion cases I commented on yesterday. Now that #1142 is merged, maybe we can rebase this and make that reader change while keeping its regression coverage?

[binaryaaron]

agent (review-pr): Rebased onto merged #1142 and applied the reader changes in 508dcfc6.

@nabinchha @andreatnvidia, the branch now builds on updated main and retains #1142's metadata, index, and persistence regressions.

• Native columns are read together through a shared helper. None and "numpy_nullable" use public pd.read_parquet, preserving pandas' string-inference and timezone conversion. Native fragment schemas are resolved before conversion, preserving categories with different batch vocabularies and nullable integers with an all-null shard. JSON columns are restored from the original fragments.
• The Arrow path uses pandas' documented conversion hook to wrap the original arrays directly, replacing manual nested-column/index reconstruction. The shared path also fixes the nested-null-array corruption for ordinary, unannotated Parquet. The required metadata repair remains.
• Invalid backends now raise ValueError before artifact access, including when artifacts are missing.

The earlier prototype in my proposal is superseded: checking it against #1142 exposed metadata/index regressions, which this implementation covers. The final diff against updated main is 7 files, +246/−57. Existing tests gained the reported cases; the number of new test functions is unchanged from the previous PR revision.

Validation: 5,256 passed, 3 skipped in make test; make check-all and commit hooks passed. The targeted suite passed 573 checks each on pandas 2.3.3 and 3.0.5 with Arrow 24 and optional pytz installed. Final focused controls also passed on pandas 2.3.3 / Arrow 24 and pandas 3.0.5 / Arrow 26.0.1. Provider-backed E2E was not run.

Dependencies remain unchanged; the Arrow upgrade remains a separate option. The API motivation also remains as previously qualified: choosing pandas conversion at the load boundary is the benefit, but a separate consumer requiring numpy_nullable has not been established, and this does not repair Anonymizer's nested-Arrow export metadata issue.

[andreatnvidia]

@binaryaaron thanks for the update, this addresses the points from the earlier review. I ran another pass, including live generation and export workflows, and found two remaining edge cases: a unified read can fail on large nested integers without reaching the per-fragment fallback, and recombining native and JSON columns drops unrelated DataFrame attrs. I’d get the fallback case fixed before merging; the attrs issue seems less urgent.

[andreatnvidia]

There seems to be a fallback case missing here. With a JSON-encoded column alongside a native struct containing 9007199254740993 in one shard and 0.5 in another, the default load_dataset() works on the base but now raises ArrowInvalid. Schema unification succeeds; the later read fails because the integer cannot safely become a double. Could the per-fragment fallback also handle that read-time conversion failure?

[andreatnvidia]

I noticed that this concat drops unrelated DataFrame attributes. For example, attrs={'external': {'units': 'USD'}} survives an encoded round trip on the base, but disappears here. The protected frame has the JSON storage attribute added, so its attrs differ from the native frame’s. Maybe we could carry forward the attributes retained by the protected-fragment concat?