Repository navigation
fix: preserve nested Arrow arrays on storage reads - #1142
Conversation
Avoid the pandas ArrowDtype identity cast that corrupts nested null children. Keep native columns Arrow-backed and preserve explicit JSON decoding and metadata. Signed-off-by: Lipika Ramaswamy <lipika.ramaswamy@gmail.com>
[Medium impact] The PR appears safe to merge; no new actionable issues were found. SummaryPreserves nested Arrow arrays when reading JSON-annotated Parquet files.
|
| frame = table.to_pandas( | ||
| types_mapper=lambda dtype: None if lazy.pa.types.is_nested(dtype) else lazy.pd.ArrowDtype(dtype) | ||
| ) | ||
| for field in table.schema: | ||
| if lazy.pa.types.is_nested(field.type) and field.name in frame.columns: | ||
| frame[field.name] = lazy.pd.Series(lazy.pd.arrays.ArrowExtensionArray(table[field.name]), index=frame.index) |
There was a problem hiding this comment.
decode_storage_table converts every nested column to Python objects in table.to_pandas(), then immediately discards those objects and wraps the original Arrow arrays. Large nested columns now need extra time and memory, including JSON columns that are decoded again below. Build the frame without converting nested columns, then insert their ArrowExtensionArray values.
Prompt To Fix With AI
This is a comment left during a code review.
Path: packages/data-designer-config/src/data_designer/config/utils/storage_encoding.py
Line: 125-130
Comment:
**Nested reads do extra work**
`decode_storage_table` converts every nested column to Python objects in `table.to_pandas()`, then immediately discards those objects and wraps the original Arrow arrays. Large nested columns now need extra time and memory, including JSON columns that are decoded again below. Build the frame without converting nested columns, then insert their `ArrowExtensionArray` values.
---
For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.There was a problem hiding this comment.
Addressed in dd134a1. Nested columns are now excluded from the table sent to pandas conversion, then wrapped directly from the original Arrow arrays; there is no intermediate nested Python-object materialization. Added six edge cases covering zero-row/all-nested tables, RangeIndex/named/MultiIndex row indexes, interleaved column order, and sliced/chunked arrays. Revised patch validation: 839 DD config tests passed, 89 storage tests passed on pandas 2.3.3, and the full downstream Anonymizer suite passed (2,017 passed, 2 skipped) against the local DD source overlay. Normal commit hooks and repository lint/format scope pass. Bug report and synthetic reproducer: #1143. Self-review completed again after the change; PR CI remains separate from these local results.
|
Fern preview: https://nvidia-preview-pr-1142.docs.buildwithfern.com/nemo/datadesigner
|
svcnvidia-nemo-ci
left a comment
There was a problem hiding this comment.
No Critical or Warning findings.
I reviewed both changed files against the base and traced the annotated Parquet reader, fragment concatenation, artifact loading, and seed handoff. The decoder wraps the original nested arrays while keeping explicit JSON decoding authoritative. The regression tests cover buffer validity, values, Arrow dtypes, metadata, and column order.
This was a static review; I couldn't independently reproduce the reported test results. One delegated report was unavailable, so I inspected its public-reader, integration, and test-coverage topic directly.
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.
Signed-off-by: Lipika Ramaswamy <lipika.ramaswamy@gmail.com>
andreatnvidia
left a comment
There was a problem hiding this comment.
The nested-array fix looks good overall. I ran into two metadata edge cases where loading changes the returned columns: non-string labels can lose scalar columns, and tuple-valued indexes come back as data columns. I’d prefer to address those before merging; details are in the inline comments.
Map physical Arrow fields to pandas-restored labels before reinsertion. Restore nested persisted index fields as indexes rather than data columns, without casting the original nested arrays. Add decoder and annotated-reader regressions for non-string column labels and tuple-valued indexes. Signed-off-by: Lipika Ramaswamy <lipika.ramaswamy@gmail.com>
Filter stale field and index metadata before pandas conversion. Avoid redundant Arrow index dtype parsing while preserving nested buffers, logical labels, and all-nested row counts. Cover repeated persisted reads, metadata projections, column labels, indexes, artifact stages, and failure preservation across 85 edge cases. Signed-off-by: Lipika Ramaswamy <lipika.ramaswamy@gmail.com>
📋 Summary
Prevent valid nested Arrow arrays from becoming invalid while loading JSON-annotated Parquet artifacts. pandas'
ArrowDtypeconversion performs an identity cast; with nested null children, PyArrow 24 can produce a malformed child array even though the original Parquet table validates successfully.The synthetic regression reproduces this on pandas 2.3.3 / PyArrow 24.0.0: upstream main fails the variable-length and uniform nested-list cases; both pass with this fix.
🔗 Related Issue
Fixes #1143
🔄 Changes
ArrowExtensionArray. This avoids both the corrupting identity cast and the redundant Python-object materialization identified by Greptile.🧪 Testing
PYTHONPATHoverlays; shared environments were not modified. No provider-backed E2E/model calls were run.✅ Checklist
Description updated with AI assistance.