[SPARK-54945][PYTHON][TEST] Add tests for pa.Array.to_pandas with zero copy

[Spenserrrr]

What changes were proposed in this pull request?

This PR adds a golden-file test that monitors the zero-copy behavior of pa.Array.to_pandas, under the SPARK-54936 umbrella for detecting upstream PyArrow/pandas behavior drift.

It adds PyArrowArrayToPandasZeroCopyTests as a new subclass in the existing python/pyspark/tests/upstream/pyarrow/test_pyarrow_arrow_to_pandas_non_default.py (reusing the shared _PyArrowToPandasTestBase and the centralized golden-matrix driver), plus a _verify_zero_copy helper on the base class.

Three things are recorded per source array:

• zero_copy_only=False — the default, where PyArrow silently copies when a view is not possible; records the resulting dtype.

• zero_copy_only=True — PyArrow's own verdict: Series[dtype] when zero-copy, or ERR@ArrowInvalid when a copy is required.

• verified zero-copy — an independent check of whether the buffers were actually reused, rather than trusting the flag. It inspects whatever storage pandas returned: for Arrow-backed results it compares buffer addresses via the public __arrow_array__ protocol (keying on that protocol rather than isinstance(dtype, pd.ArrowDtype), since pandas 3's string dtype is Arrow-backed without being ArrowDtype); otherwise it uses np.shares_memory, which accounts for slice offsets.

  An Arrow array is several buffers (validity, offsets, values), so a result can borrow some while allocating others. The check therefore counts how many of the result's buffers were borrowed and reports zero-copy when all were, partial-copy when only some were, and copied when none were — rather than declaring zero-copy on the first buffer that matches, which would let a conversion reuse a cheap validity bitmap while reallocating the values and still pass. This is not hypothetical: on pandas 3 with PyArrow >= 24, pandas normalizes string to large_string, so a string source keeps its values but has its 32-bit offsets rebuilt as 64-bit (1 of 2 buffers shared) while a large_string source passes through untouched (2 of 2). The numpy branch keeps a single check, since a numpy array is one dense block and "any buffer reused" and "all buffers reused" coincide there.

The row set reuses the sibling test_pyarrow_arrow_to_pandas_default.py row set directly, so both golden files pin the same Arrow types, and appends the layout variants that only matter for zero-copy: sliced (offset) arrays and single- vs multi-chunk ChunkedArrays.

The two verification columns are recorded as data rather than asserted equal, because they legitimately disagree: a tz-aware timestamp reports zero_copy_only=True while pandas materializes it into a DatetimeTZDtype array that reuses no Arrow buffer. to_pandas is two stages — the Arrow-to-numpy step is genuinely zero-copy, but the pandas wrapper built afterwards copies — and the flag only describes the first stage. Pinning both makes that visible instead of hiding it behind an assertion.

Scope note: this PR covers the default (NumPy-backed) path. The Arrow-backed path (types_mapper=pd.ArrowDtype) is a separate argument and will be a follow-up.

Why are the changes needed?

PySpark relies on pa.Array.to_pandas throughout its conversion layer (python/pyspark/sql/pandas/conversion.py), and whether a given Arrow type converts without copying its buffers directly affects the memory and latency of toPandas and pandas UDFs. Pinning this behavior in a golden file lets CI fail loudly if it drifts across PyArrow/pandas/NumPy upgrades, instead of silently regressing performance.

Does this PR introduce any user-facing change?

No. This is a test-only change.

How was this patch tested?

This is itself a test. It runs without a Spark session (it exercises PyArrow/pandas directly). Verified both with SPARK_GENERATE_GOLDEN_FILES=1 (regenerating and reviewing every cell) and without the flag (comparing against the committed golden file).

Validated against the committed golden across 16 environments — PyArrow 18, 19, 20, 21, 22, 23, 24 and 25, each with pandas 2 and pandas 3 — all passing. The pandas 3 differences are pinned with version-guarded overrides in two tiers: pandas 3 alone only changes the string dtype, while pandas 3 together with PyArrow >= 24 is what makes the conversion Arrow-backed, so that is the only regime where buffers are reused at all and where partial-copy can arise. The golden file itself is generated on pandas 2, where storage is numpy-backed and partial reuse cannot occur.

ruff check and format pass.

Was this patch authored or co-authored using generative AI tooling?

Generated-by: Claude Code (Opus 4.8)

[Spenserrrr]

Hi @Yicong-Huang @zhengruifeng! Could you take a look when you have a moment? Thank you!

[Yicong-Huang]

1 blocking, 1 non-blocking, 0 nits.
Well-structured test-only change that reuses the golden-file family conventions and adds real value with the independent np.shares_memory column. Two points raised: the nested-type coverage looks incomplete versus the sibling default test (blocking request to add a few control rows), and a non-blocking question about whether the string/binary rows were checked across the supported pyarrow range.

Suggestions (2)

• python/pyspark/tests/upstream/pyarrow/test_pyarrow_arrow_to_pandas_non_default.py:457: Golden was generated on pyarrow 24.0.0 and the test carries no pa.version override (only a pandas>=3.0.0 one). Question for the author: were the string/binary rows' zero_copy_only verdicts checked across the supported pyarrow range (>=18), or only on 24.0.0? The sibling default test needed a pyarrow>=24.0.0 override for a pandas-3 / pyarrow-24 string-dtype change, so if this ever runs under pandas 3 or a newer pyarrow the string rows may need the same pa.version guard. -- see inline
• python/pyspark/tests/upstream/pyarrow/test_pyarrow_arrow_to_pandas_non_default.py:435: Nested-type coverage looks incomplete versus the sibling default test: only list:no-null and struct<x:int64>:no-null here, while default covers map, dictionary, large_list, fixed_size_list, nullable/empty nested, and deep nesting (list, list, list, struct, struct, struct, map<string,list>, map<string,struct>, map<string,map>). Suggest bringing the nested rows in line with the default test's set so the same shapes' zero_copy_only verdict is pinned. -- see inline

Verification

Verified the new subclass and _numpy_shares_arrow_buffer against the reused base-class API in goldenutils.py, and cross-checked representative golden cells (int8:no-null->shared, bool:no-null->not-shared, int64:multi-chunk->not-shared, int64:sliced->shared, int64:empty->n/a (empty)). Confirmed the test is registered in dev/sparktestsupport/modules.py. Ran the actual test class against the committed golden on pyarrow 18.0.0/19/20/21/22/23/24.0.0/25.0.0 x pandas 2.2.0-2.3.3 x numpy 1.26.4-2.2.6 (isolated venv, PYTHONPATH harness): all PASS, zero cell mismatches -- so the string/binary rows are stably ERR@ArrowInvalid across the supported pyarrow range and need no pa.version override. ASCII / line-length hygiene on the changed .py file is clean.

[Yicong-Huang]

This golden is generated on pyarrow 24.0.0 and there's no pa.__version__ override (only the pandas>=3 one). Did you check the string:* / binary rows across the supported pyarrow range (>=18), or just on 24.0.0?

Asking because the sibling test_pyarrow_arrow_to_pandas_default.py needed a pyarrow>=24.0.0 override for a pandas-3 / pyarrow-24 string-dtype change -- if this test is ever run under pandas 3 or a newer pyarrow, the string rows here might need the same pa.__version__ guard.

[Spenserrrr]

I re-ran against the expanded golden and added a pandas 3 axis: PyArrow 18 to 25 with pandas 2 and 3, 16/16 pass. No pa.version override needed on pandas 2. Note that pandas 3 did need a two-tier override: pandas 3 alone only changes the dtype (object to str), while pandas 3 and pyarrow >= 24 is what makes the conversion Arrow-backed so zero_copy_only succeeds.

[Spenserrrr]

Also, there is one thing worth flagging. zero_copy_only=True isn't always accurate. timestamp[us,tz=UTC] reports success but is genuinely copied, which is the only mismatch in 125 rows. to_pandas is two stages: Arrow to numpy is zero-copy, but numpy toSeries copies because DatetimeTZDtype is an extension dtype whose constructor allocates. The flag only describes the first stage. Hence, both columns are recorded as data rather than asserted equal.

[Yicong-Huang]

Confirmed -- I re-ran pyarrow 18-25 on pandas 2.3.3 and 3.0.5, zero mismatches, so the two-tier override matches the Arrow-backed boundary. Good catch on the tz-aware timestamp; recording both columns as data is the right call.

[Yicong-Huang]

The nested-type coverage looks a bit incomplete here -- just list<int64>:no-null and struct<x:int64>:no-null. The sibling default test covers the full nested set: map, dictionary, large_list, fixed_size_list, nullable/empty variants, and deep nesting (list<list>, list<struct>, list<map>, struct<struct>, struct<list>, struct<map>, map<string,list>, map<string,struct>, map<string,map>).

Could you bring the nested rows here in line with that set? Mirroring the default test's nested coverage keeps the two golden files consistent and pins the same shapes' zero_copy_only verdict.

[Spenserrrr]

Thanks for the advice @Yicong-Huang! I reuse the default test's rows directly and append a few more meaningful tests for this flag.

[Yicong-Huang]

Reusing the default test's rows directly is better than what I suggested -- coverage stays in sync now. Thanks!

[Yicong-Huang]

2 addressed, 0 remaining, 0 new.
0 blocking, 0 non-blocking, 0 nits.
Both prior findings resolved; nested coverage now reuses the default test's rows, and the version concern is closed with a correct two-tier pandas-3 / pyarrow-24 guard.

Verification

Re-ran the committed golden across pyarrow 18-25 on pandas 2.3.3 and 3.0.5: 16/16 pass, zero mismatches -- confirming no pa.version override is needed on pandas 2. Spot-checked the rewritten _verify_zero_copy (Arrow-backed address compare vs numpy shares_memory) and the tz-aware copied row.

[zhengruifeng]

Minor concern (forward-looking): _arrow_buffers includes validity, offset, and value buffers, and this loop returns zero-copy when any source buffer is present in pandas storage. A conversion could retain an offset or validity buffer while materializing the values, masking the data-copy regression this test is meant to detect. Could we compare the value buffer specifically (as described in the design doc), require all relevant buffers to match, or report a distinct partial-copy result?

[Spenserrrr]

Thanks for the advice @zhengruifeng! Actually, this problem already appears. On pandas 2 it's harmless. A numpy result is one dense block, so any/all are the same question. But on pandas 3 + pyarrow >= 24, strings become ArrowStringArray with 3 buffers of their own. string:standard borrows the values while pandas rebuilds the 32-bit offsets as 64-bit because it normalizes string to large_string. Thus, only 1 of the 2 buffers is shared. The old check can't tell this.

Thus, I count how many of the result's buffers were borrowed, then report zero-copy (all), partial-copy (some), copied (none). The edit is in the new commit. The golden file is identical, and in the override section, two rows are changed to partial-copy.

[zhengruifeng]

Merge Summary:

• merged into master 15f12eb
• merged into branch-4.x 85266b3

Posted by merge_spark_pr.py

[HyukjinKwon]

0 blocking, 0 non-blocking, 0 nits.
Clean, well-reviewed test-only monitor; the buffer-borrow classifier correctly handles partial sharing.

Verification

The buffer-borrow counting resolves the earlier any/all masking concern: on pandas 3 + pyarrow >= 24, large_string rebuilds 32-bit offsets as 64-bit so only the values buffer is shared, and the classifier reports partial-copy rather than falsely zero-copy. Nested-type coverage reuses the default test's row set (kept in sync). The reviewers ran the committed golden across pyarrow 18-25 x pandas 2-3 with zero mismatches, and the prior review converged to 0 blocking / 0 non-blocking / 0 nits.