Skip to content

Keep row validity and payloads in fixed-shape ArrayXD pandas round trips - #8889

Open
choksi2212 wants to merge 1 commit into
huggingface:mainfrom
choksi2212:fix/fixed-shape-arrayxd-pandas-nulls-8888
Open

choksi2212 wants to merge 1 commit into
huggingface:mainfrom
choksi2212:fix/fixed-shape-arrayxd-pandas-nulls-8888

Conversation

@choksi2212

Copy link
Copy Markdown

Fixes #8888

PandasArrayExtensionDtype.__from_arrow__ materialized every fixed-shape ArrayXD column through ArrayExtensionArray.to_numpy, which pads the null rows with NaN and promotes the whole column to float64. Two consequences at the pandas boundary: the Arrow row validity was gone (a null row and a legitimate all-NaN array became indistinguishable), and integer columns lost their exact payloads, with values above 2**53 rounding and Dataset.from_pandas(..., features=...) raising ArrowInvalid when converting the NaN rows back to int64.

A fixed-shape column with null rows is now materialized as one exact ndarray per valid row and None per null row, the representation the dynamic shape already uses, so Dataset.from_pandas(ds.to_pandas(), features=ds.features) keeps the validity bitmap and the exact values:

features = Features({"x": Array2D((1, 1), "float64")})
source = Dataset.from_dict({"x": [[[np.nan]], None]}, features=features)
result = Dataset.from_pandas(source.to_pandas(), features=features)
assert result.data.column("x").is_null().to_pylist() == [False, True]

The pandas dtype now also carries the fixed shape and the value type, so a column whose rows are all null still rebuilds its extension type (on main that round trip silently turned the nulls into valid all-NaN rows), and the dtype is carried through copy, slicing, take, _concat_same_type and _from_sequence. isna() used to crash on None rows (pd.isna(None).any()); it now counts scalar None/NaN rows as null without inspecting them, while ndarray rows keep the existing NaN inspection. Dynamic shapes and fixed shapes without null rows keep their previous representations, so the null-row regression from #8701 and the NumPy formatting behavior kept by #8363 are untouched.

The new tests fail on main and pass here: validity round trips for (1,1) and (2,3) columns with null rows in the middle, exact int64 payloads above 2**53, a sliced column, and an all-null column. Verified locally on the issue's pins (pandas 3.1.0, pyarrow 26.0.0, numpy 2.5.4, Python 3.13): tests/features/test_array_xd.py (111), tests/features/test_features.py, tests/test_table.py (598 passed together), and ruff check/ruff format are clean.

…pandas round trips

PandasArrayExtensionDtype.__from_arrow__ materialized every fixed-shape
ArrayXD column through ArrayExtensionArray.to_numpy, which pads the null
rows with NaN and promotes the whole column to float64. The pandas
boundary therefore lost the Arrow row validity (a null row and a
legitimate all-NaN array became indistinguishable) and the exact
payloads of integer columns: values above 2**53 lost precision, and
converting the NaN rows back to int64 raised ArrowInvalid.

A fixed-shape column with null rows is now materialized as one exact
ndarray per valid row and None per null row, like a dynamic shape
already does, so Dataset.from_pandas(ds.to_pandas(), features=...)
keeps the row validity bitmap and the payloads. The pandas dtype keeps
the fixed shape and the value type, so an all-null column can rebuild
its extension type as well, and the dtype is carried through
copy/slice/take/concat/_from_sequence. isna() no longer crashes on
None rows. Dynamic shapes and fixed shapes without null rows keep
their previous representations.

Signed-off-by: choksi2212 <choksi2212@users.noreply.github.com>
@choksi2212

Copy link
Copy Markdown
Author

CI note: this is the first PR from my fork, so the CI and Build PR Documentation workflows are sitting in action_required until someone approves the first run. Everything in the description is verified locally on the issue's pinned versions (pandas 3.1.0, pyarrow 26.0.0, numpy 2.5.4): the five new test cases fail on main and pass on this branch, and the full test_array_xd.py, test_features.py and test_table.py suites are green with ruff check/ruff format clean.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Fixed-shape ArrayXD pandas round trips lose row nulls and integer precision

1 participant