Spenserrrr commented on code in PR #57624:
URL: https://github.com/apache/spark/pull/57624#discussion_r3708134335
##########
python/pyspark/tests/upstream/pyarrow/test_pyarrow_arrow_to_pandas_non_default.py:
##########
@@ -258,6 +309,182 @@ def compute_cell(row_name, col_name):
)
[email protected](
+ not have_pyarrow or not have_pandas or not have_numpy,
+ pyarrow_requirement_message or pandas_requirement_message or
numpy_requirement_message,
+)
+class PyArrowArrayToPandasZeroCopyTests(_PyArrowToPandasTestBase):
+ """
+ Tests pa.Array.to_pandas(zero_copy_only=True) via golden file comparison.
+
+ PySpark converts Arrow data to (numpy-backed) pandas objects throughout its
+ conversion layer (``python/pyspark/sql/pandas/conversion.py``); whether a
+ given Arrow type can make that conversion WITHOUT copying its buffers
+ directly affects the memory and latency of ``toPandas`` and pandas UDFs.
+ ``zero_copy_only=True`` makes PyArrow raise ``ArrowInvalid`` instead of
+ silently copying, so it is the natural probe for "is this conversion
+ zero-copy?". These tests record, per Arrow type, whether the conversion is
+ zero-copy so CI fails loudly if that changes across pandas/PyArrow/NumPy
+ upgrades.
+
+ Two output columns are recorded for each source array:
+
+ - ``zero_copy_only=True``: PyArrow's own verdict -- ``Series[dtype]`` when
the
+ conversion is zero-copy, or ``ERR@ArrowInvalid`` when a copy is required.
+ - ``numpy shares arrow buffer``: an INDEPENDENT verification (via
+ ``np.shares_memory``) of whether the default ``to_pandas()`` result
+ actually shares memory with the Arrow buffer. This does not simply trust
+ PyArrow's flag; it confirms the physical sharing. The two columns are
+ expected to agree except in vacuous or backend-specific cases (e.g. empty
+ arrays, or pandas 3 arrow-backed string where the flag reports zero-copy
+ but ``to_numpy()`` materializes) -- such differences are pinned as data.
+
+ The row set targets the layout properties that determine zero-copy, rather
+ than re-enumerating every Arrow type (that is covered by
+ ``test_pyarrow_arrow_to_pandas_default.py``): fixed-width numerics with and
+ without nulls, bool, string/binary, temporal types across units, sliced
+ (offset) arrays, and single- vs multi-chunk ChunkedArrays.
+ """
+
+ def _build_source_arrays(self):
+ """Build an ordered dict of named source PyArrow arrays for testing."""
+ sources = {}
+
+ # =====================================================================
+ # Fixed-width numerics WITHOUT nulls: the canonical zero-copy case --
+ # the Arrow value buffer already matches numpy's layout exactly.
+ # =====================================================================
+ for bits, pa_type in [
+ (8, pa.int8()),
+ (16, pa.int16()),
+ (32, pa.int32()),
+ (64, pa.int64()),
+ ]:
+ sources[f"int{bits}:no-null"] = pa.array([0, 1, -1], pa_type)
+ sources["uint32:no-null"] = pa.array([0, 1, 2], pa.uint32())
+ sources["float32:no-null"] = pa.array([0.0, 1.5, -1.5], pa.float32())
+ sources["float64:no-null"] = pa.array([0.0, 1.5, -1.5], pa.float64())
+
+ # =====================================================================
+ # Same numerics WITH a null: Arrow tracks nulls in a separate validity
+ # bitmap that numpy has no equivalent for, so pandas must
promote/rebuild
+ # (e.g. int -> float64 with NaN) -- a copy.
+ # =====================================================================
+ sources["int32:with-null"] = pa.array([0, 1, None], pa.int32())
+ sources["int64:with-null"] = pa.array([0, 1, None], pa.int64())
+ sources["float64:with-null"] = pa.array([0.0, 1.5, None], pa.float64())
+
+ # =====================================================================
+ # Boolean: Arrow stores bits (1 bit/value); numpy uses 1 byte/value, so
+ # the buffer must be unpacked -- always a copy, even without nulls.
+ # =====================================================================
+ sources["bool:no-null"] = pa.array([True, False, True], pa.bool_())
+ sources["bool:with-null"] = pa.array([True, None, False], pa.bool_())
+
+ # =====================================================================
+ # String / binary: variable-length data materialized into an object
+ # array of Python str/bytes -- always a copy.
+ # =====================================================================
+ sources["string:no-null"] = pa.array(["hello", "world", ""],
pa.string())
+ sources["binary:no-null"] = pa.array([b"hello", b"world"], pa.binary())
+
+ # =====================================================================
+ # Temporal types. Fixed-width timestamp/duration convert zero-copy
+ # without nulls (they map to numpy datetime64/timedelta64); date and
+ # time map to object dtype and therefore always copy. A nullable
+ # timestamp requires a copy (bitmap reconciliation, like the numerics).
+ # =====================================================================
+ dt = datetime.datetime(2024, 1, 1, 12, 0, 0)
+ for unit in ["s", "ms", "us", "ns"]:
+ sources[f"timestamp[{unit}]:no-null"] = pa.array([dt, dt],
pa.timestamp(unit))
+ sources["timestamp[us]:with-null"] = pa.array([dt, None],
pa.timestamp("us"))
+ td = datetime.timedelta(days=1)
+ sources["duration[us]:no-null"] = pa.array([td, td], pa.duration("us"))
+ sources["date32:no-null"] = pa.array(
+ [datetime.date(2024, 1, 1), datetime.date(2024, 6, 15)],
pa.date32()
+ )
+ sources["time64[us]:no-null"] = pa.array(
+ [datetime.time(12, 30), datetime.time(18, 45)], pa.time64("us")
+ )
+
+ # =====================================================================
+ # Sliced (offset) arrays: still zero-copy when the slice views a
+ # contiguous no-null primitive region, but the numpy data starts
partway
+ # into the parent buffer -- the case that defeats naive address
equality.
+ # =====================================================================
+ sources["int64:sliced"] = pa.array(list(range(10)),
pa.int64()).slice(2, 3)
+ sources["int64:sliced-with-null"] = pa.array([1, 2, None, 4, 5],
pa.int64()).slice(1, 3)
+
+ # =====================================================================
+ # ChunkedArray: a single chunk is zero-copy, but multiple chunks must
be
+ # concatenated into one contiguous numpy buffer -- a copy. This is the
+ # common shape in real PySpark, where each partition is its own chunk.
+ # =====================================================================
+ sources["int64:single-chunk"] = pa.chunked_array([pa.array([1, 2, 3],
pa.int64())])
+ sources["int64:multi-chunk"] = pa.chunked_array(
+ [pa.array([1, 2], pa.int64()), pa.array([3, 4], pa.int64())]
+ )
+
+ # =====================================================================
+ # Empty and nested (controls). Empty reports zero-copy (nothing to
copy)
+ # but shares no buffer; nested types materialize to object -- a copy.
+ # =====================================================================
+ sources["int64:empty"] = pa.array([], pa.int64())
+ sources["string:empty"] = pa.array([], pa.string())
+ sources["list<int64>:no-null"] = pa.array([[1, 2], [3]],
pa.list_(pa.int64()))
+ sources["struct:no-null"] = pa.array([{"x": 1}, {"x": 2}],
pa.struct([("x", pa.int64())]))
+
+ return sources
+
+ # Output column recording PyArrow's own zero-copy verdict.
+ COL_ZERO_COPY_ONLY = "zero_copy_only=True"
+
+ # Output column independently verifying physical memory sharing.
+ COL_SHARES_BUFFER = "numpy shares arrow buffer"
+
+ def test_to_pandas_zero_copy_only(self):
+ """Test pa.Array.to_pandas(zero_copy_only=True) against golden file."""
+ sources = self._build_source_arrays()
+ row_names = list(sources.keys())
+ col_names = [
+ "pyarrow array",
+ self.COL_ZERO_COPY_ONLY,
+ self.COL_SHARES_BUFFER,
+ ]
+
+ # Version-specific expected values go here, keyed by (row, col), when a
+ # newer pandas/PyArrow/NumPy legitimately changes a cell's output.
+ overrides: dict[tuple[str, str], str] = {}
Review Comment:
Thanks for the advice @Yicong-Huang! I reuse the default test's rows
directly and append a few more meaningful tests for this flag.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
