From dc29a19e19c06df447d0b2a42be6169864adf81e Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Wed, 28 Oct 2026 21:26:55 -0700 Subject: [PATCH 01/23] [SPARK-60676][CONNECT] Use Arrow columns for signed int64 pandas inputs --- docs/architecture.md | 7 ++ python/pyspark/sql/_pandas_input.py | 60 ++++++++++++++ python/tests/test_pandas_input.py | 117 ++++++++++++++++++++++++++++ 3 files changed, 184 insertions(+) create mode 100644 python/pyspark/sql/_pandas_input.py create mode 100644 python/tests/test_pandas_input.py diff --git a/docs/architecture.md b/docs/architecture.md index 5cce8782..75f73ad5 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -142,6 +142,13 @@ timestamp conversion helpers normalize the table, then arrow-rs imports its those batches to IPC for the local relation sent to Spark. No Python rows are materialized; casts, buffer alignment, and IPC serialization may still copy. Empty tables require an explicit schema, matching PySpark Connect. +Pandas inputs with only NumPy `int64` or pandas nullable `Int64` columns also +use this importer, preserving integer precision and null masks without Python +rows. This initial path accepts inferred schemas, complete string column names, +and explicit `StructType` schemas whose fields are all `LongType`. Other dtypes, +DDL/atomic schemas, schema conversions, and duplicate/non-string input names +retain the existing row conversion path. `verifySchema` remains ignored for +pandas Arrow input, matching Spark Connect. Inputs at or above `spark.sql.session.localRelationCacheThreshold` are serialized into independent IPC streams and uploaded as session-scoped cache artifacts. The plan carries `ChunkedCachedLocalRelation` hashes instead of embedding the diff --git a/python/pyspark/sql/_pandas_input.py b/python/pyspark/sql/_pandas_input.py new file mode 100644 index 00000000..7b5f3438 --- /dev/null +++ b/python/pyspark/sql/_pandas_input.py @@ -0,0 +1,60 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + + +"""Conservative columnar pandas input; unsupported conversions use the row path.""" + +from pyspark.sql.types import LongType, StructType + + +def prepare_int64_table(frame, schema): + """Return an Arrow table only when each column is already a signed 64-bit integer. + + In particular, never coerce a mixed pandas frame to a common NumPy dtype: + doing so can round integers above 2**53. Other dtypes and schema conversions + deliberately retain the existing row path until their semantics are covered. + ``verifySchema`` is ignored for pandas Arrow input, as in Spark Connect. + """ + import numpy as np + import pandas as pd + + width = len(frame.columns) + if not width or not frame.columns.is_unique or not all( + isinstance(name, str) for name in frame.columns + ): + return None + if not all(dtype == np.dtype("int64") or isinstance(dtype, pd.Int64Dtype) + for dtype in frame.dtypes): + return None + if isinstance(schema, StructType): + if len(schema.fields) != width or not all( + isinstance(field.dataType, LongType) for field in schema.fields + ): + return None + elif schema is not None: + if not isinstance(schema, (list, tuple)) or len(schema) != width or not all( + isinstance(name, str) for name in schema + ): + return None + + import pyarrow as pa + + # Arrow consumes each Series' integer buffer and nullable mask independently. + return pa.Table.from_arrays( + [pa.array(frame.iloc[:, i], type=pa.int64(), from_pandas=True) for i in range(width)], + names=list(frame.columns), + ) diff --git a/python/tests/test_pandas_input.py b/python/tests/test_pandas_input.py new file mode 100644 index 00000000..01adc89c --- /dev/null +++ b/python/tests/test_pandas_input.py @@ -0,0 +1,117 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""Signed-int64 pandas column handoff and conservative fallback coverage.""" + +from datetime import date, datetime, timezone +from decimal import Decimal + +import numpy as np +import pandas as pd +import pytest + +from pyspark.sql import types as T +from pyspark.sql._pandas_input import prepare_int64_table +from test_arrow_handoff import _arrow_input_service, protocol + + +@pytest.mark.parametrize("schema_kind", ["inferred", "names", "explicit"]) +@pytest.mark.parametrize("verify", [True, False]) +def test_integer_columns_preserve_values_without_python_rows(protocol, monkeypatch, schema_kind, verify): + frame = pd.DataFrame({ + "a": pd.Series([2**53 + 1, None, -(2**63)], dtype="Int64"), + "b": np.array([2**63 - 1, -1, 0], dtype=np.int64), + }) + schema = None + names = ["a", "b"] + if schema_kind == "names": + schema = names = ["renamed_a", "renamed_b"] + elif schema_kind == "explicit": + schema = T.StructType([T.StructField("a", T.LongType()), + T.StructField("b", T.LongType(), False, {"source": "test"})]) + import pyspark.sql._arrow_input as helper + original = helper.prepare_table + observed = [] + + def capture(*args): + result = original(*args) + observed.append(result) + return result + + def no_rows(*args, **kwargs): + raise AssertionError("supported pandas input must not materialize Python rows") + + monkeypatch.setattr(helper, "prepare_table", capture) + monkeypatch.setattr(pd.DataFrame, "astype", no_rows) + with _arrow_input_service(protocol) as (spark, state): + df = spark.createDataFrame(frame, schema=schema, verifySchema=verify) + assert df.isLocal() + assert df.toArrow().column(0).to_pylist() == [3] + assert state["result"] == (3, 2**53 + 1, -(2**63)) + table, resolved, renamed = observed[0] + assert table.to_pydict() == {"a": [2**53 + 1, None, -(2**63)], + "b": [2**63 - 1, -1, 0]} + assert (renamed or resolved.names) == names + if schema_kind == "explicit": + assert resolved == schema + assert not table.schema.field(1).nullable + assert resolved.fields[1].metadata == {"source": "test"} + + +@pytest.mark.parametrize("dtype", ["int64", "Int64"]) +def test_empty_integer_frame_requires_explicit_schema(protocol, dtype): + frame = pd.DataFrame({"a": pd.Series([], dtype=dtype)}) + with _arrow_input_service(protocol) as (spark, state): + with pytest.raises(ValueError): + spark.createDataFrame(frame) + df = spark.createDataFrame(frame, T.StructType([T.StructField("a", T.LongType())])) + assert df.isLocal() + + +@pytest.mark.parametrize("values", [ + [1.0, float("nan")], [True, False], ["x", None], + [datetime(2024, 1, 1), None], + [datetime(2024, 1, 1, tzinfo=timezone.utc), None], + [date(2024, 1, 1), None], [Decimal("1.20"), None], + [[1, 2], None], [{"a": 1}, None], +]) +def test_other_dtypes_keep_row_fallback(values): + frame = pd.DataFrame({"a": values}) + assert prepare_int64_table(frame, None) is None + assert prepare_int64_table(frame, T.StructType([T.StructField("a", T.LongType())])) is None + + +@pytest.mark.parametrize("schema", [ + "a long", T.LongType(), ["a", "b"], [1], + T.StructType([T.StructField("a", T.IntegerType())]), + T.StructType([T.StructField("a", T.DoubleType())]), + T.StructType([]), +]) +def test_schema_conversions_keep_row_fallback(schema): + assert prepare_int64_table(pd.DataFrame({"a": [1, 2]}), schema) is None + + +def test_duplicate_names_fallback_and_nullable_masks(): + frame = pd.DataFrame([[2**53 + 1, 2**53 + 3]], columns=["x", "x"], dtype="int64") + assert prepare_int64_table(frame, None) is None + assert prepare_int64_table(pd.DataFrame({0: [1]}), None) is None + frame = pd.DataFrame({"a": pd.Series([None, None], dtype="Int64")}) + assert prepare_int64_table(frame, None).column(0).to_pylist() == [None, None] + + +@pytest.mark.parametrize("dtype", ["int32", "uint64", "Int32", "UInt64", "int64[pyarrow]"]) +def test_other_integer_dtypes_keep_row_fallback(dtype): + assert prepare_int64_table(pd.DataFrame({"a": pd.Series([1, 2], dtype=dtype)}), None) is None From 3e79f9128ae47c5f6d32c280abcb44ad5cc6b1bb Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Wed, 28 Oct 2026 21:27:26 -0700 Subject: [PATCH 02/23] [SPARK-60676][CONNECT] Validate integer upload scales against mock RPC --- python/tests/test_pandas_input.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/python/tests/test_pandas_input.py b/python/tests/test_pandas_input.py index 01adc89c..c4e50ba2 100644 --- a/python/tests/test_pandas_input.py +++ b/python/tests/test_pandas_input.py @@ -115,3 +115,21 @@ def test_duplicate_names_fallback_and_nullable_masks(): @pytest.mark.parametrize("dtype", ["int32", "uint64", "Int32", "UInt64", "int64[pyarrow]"]) def test_other_integer_dtypes_keep_row_fallback(dtype): assert prepare_int64_table(pd.DataFrame({"a": pd.Series([1, 2], dtype=dtype)}), None) is None + + +@pytest.mark.parametrize("rows", [100, 1000, 100000, 1000000]) +@pytest.mark.parametrize("explicit", [False, True]) +def test_integer_scales_reach_mock_service(protocol, monkeypatch, rows, explicit): + values = np.arange(rows, dtype=np.int64) + (2**53 + 1) + frame = pd.DataFrame({name: values for name in ("a", "b", "c")}) + schema = T.StructType([T.StructField(name, T.LongType()) for name in frame.columns]) + + def no_rows(*args, **kwargs): + raise AssertionError("integer input must stay columnar at every scale") + + monkeypatch.setattr(pd.DataFrame, "astype", no_rows) + with _arrow_input_service(protocol) as (spark, state): + spark.createDataFrame(frame, schema if explicit else None).toArrow() + assert state["result"] == (rows, 2**53 + 1, 2**53 + rows) + if frame.memory_usage(index=False).sum() > 1048576: + assert state["relations"][0].HasField("chunked_cached_local_relation") From 8c705c08fc43f8599cda2cf6a278bd851182e7ac Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Wed, 28 Oct 2026 21:40:38 -0700 Subject: [PATCH 03/23] [SPARK-60676][CONNECT] Cover leading null inference and non-nullable schemas --- python/tests/test_pandas_input.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/python/tests/test_pandas_input.py b/python/tests/test_pandas_input.py index c4e50ba2..0b727471 100644 --- a/python/tests/test_pandas_input.py +++ b/python/tests/test_pandas_input.py @@ -133,3 +133,21 @@ def no_rows(*args, **kwargs): assert state["result"] == (rows, 2**53 + 1, 2**53 + rows) if frame.memory_usage(index=False).sum() > 1048576: assert state["relations"][0].HasField("chunked_cached_local_relation") + + +def test_nullable_inference_does_not_depend_on_first_value(protocol): + frame = pd.DataFrame({"a": pd.Series([None, 2**53 + 1], dtype="Int64")}) + with _arrow_input_service(protocol) as (spark, state): + spark.createDataFrame(frame).toArrow() + assert state["result"] == (2, None, 2**53 + 1) + + +@pytest.mark.parametrize("verify", [True, False]) +def test_non_nullable_schema_is_preserved(protocol, verify): + frame = pd.DataFrame({"a": pd.Series([1, None], dtype="Int64")}) + schema = T.StructType([T.StructField("a", T.LongType(), False)]) + with _arrow_input_service(protocol) as (spark, state): + with pytest.raises(ValueError, match="null"): + spark.createDataFrame(frame, schema, verifySchema=verify) + assert not state["relations"] + assert not state["uploads"] From fb996d982b807b8fe734ff0de66ccda983ee14ae Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Wed, 28 Oct 2026 21:48:11 -0700 Subject: [PATCH 04/23] [SPARK-60676][CONNECT] Exercise named schemas through mock rename relations --- python/tests/test_arrow_handoff.py | 3 +++ python/tests/test_pandas_input.py | 3 ++- 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/python/tests/test_arrow_handoff.py b/python/tests/test_arrow_handoff.py index 459a65e1..dbc9e2d2 100644 --- a/python/tests/test_arrow_handoff.py +++ b/python/tests/test_arrow_handoff.py @@ -509,6 +509,9 @@ def execute(raw, context): state["column_names"] = list(relation.to_df.column_names) relation = relation.to_df.input state["relations"].append(relation) + # Column-name schemas wrap the local input with a rename relation. + if relation.HasField("to_df"): + relation = relation.to_df.input if relation.HasField("chunked_cached_local_relation"): cached = relation.chunked_cached_local_relation payloads = [state["cache"]["cache/" + h] for h in cached.dataHashes] diff --git a/python/tests/test_pandas_input.py b/python/tests/test_pandas_input.py index 0b727471..0f33bbca 100644 --- a/python/tests/test_pandas_input.py +++ b/python/tests/test_pandas_input.py @@ -58,7 +58,8 @@ def no_rows(*args, **kwargs): monkeypatch.setattr(pd.DataFrame, "astype", no_rows) with _arrow_input_service(protocol) as (spark, state): df = spark.createDataFrame(frame, schema=schema, verifySchema=verify) - assert df.isLocal() + if schema_kind != "names": + assert df.isLocal() assert df.toArrow().column(0).to_pylist() == [3] assert state["result"] == (3, 2**53 + 1, -(2**63)) table, resolved, renamed = observed[0] From 79d3faf3ea710ec698d7d992a22fd1bdede3a13c Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Wed, 28 Oct 2026 21:51:06 -0700 Subject: [PATCH 05/23] [SPARK-60676][CONNECT] Avoid Arrow configuration overhead for small NumPy frames --- docs/architecture.md | 5 ++++- python/pyspark/sql/_pandas_input.py | 7 ++++++- python/tests/test_pandas_input.py | 9 ++++++++- 3 files changed, 18 insertions(+), 3 deletions(-) diff --git a/docs/architecture.md b/docs/architecture.md index 75f73ad5..45bfb2c4 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -144,7 +144,10 @@ materialized; casts, buffer alignment, and IPC serialization may still copy. Empty tables require an explicit schema, matching PySpark Connect. Pandas inputs with only NumPy `int64` or pandas nullable `Int64` columns also use this importer, preserving integer precision and null masks without Python -rows. This initial path accepts inferred schemas, complete string column names, +rows. Small frames (fewer than 1,000 rows) with only NumPy columns retain the +row path to avoid configuration RPC overhead. Nullable `Int64` columns always +use Arrow so schema inference does not depend on the first value being non-null. +This initial path accepts inferred schemas, complete string column names, and explicit `StructType` schemas whose fields are all `LongType`. Other dtypes, DDL/atomic schemas, schema conversions, and duplicate/non-string input names retain the existing row conversion path. `verifySchema` remains ignored for diff --git a/python/pyspark/sql/_pandas_input.py b/python/pyspark/sql/_pandas_input.py index 7b5f3438..1015017f 100644 --- a/python/pyspark/sql/_pandas_input.py +++ b/python/pyspark/sql/_pandas_input.py @@ -37,8 +37,13 @@ def prepare_int64_table(frame, schema): isinstance(name, str) for name in frame.columns ): return None + dtypes = frame.dtypes if not all(dtype == np.dtype("int64") or isinstance(dtype, pd.Int64Dtype) - for dtype in frame.dtypes): + for dtype in dtypes): + return None + # The importer's configuration RPCs outweigh row conversion for small NumPy + # frames. Nullable extension columns still use Arrow for reliable null typing. + if len(frame) < 1000 and not any(isinstance(dtype, pd.Int64Dtype) for dtype in dtypes): return None if isinstance(schema, StructType): if len(schema.fields) != width or not all( diff --git a/python/tests/test_pandas_input.py b/python/tests/test_pandas_input.py index 0f33bbca..4cdaee7c 100644 --- a/python/tests/test_pandas_input.py +++ b/python/tests/test_pandas_input.py @@ -128,7 +128,8 @@ def test_integer_scales_reach_mock_service(protocol, monkeypatch, rows, explicit def no_rows(*args, **kwargs): raise AssertionError("integer input must stay columnar at every scale") - monkeypatch.setattr(pd.DataFrame, "astype", no_rows) + if rows >= 1000: + monkeypatch.setattr(pd.DataFrame, "astype", no_rows) with _arrow_input_service(protocol) as (spark, state): spark.createDataFrame(frame, schema if explicit else None).toArrow() assert state["result"] == (rows, 2**53 + 1, 2**53 + rows) @@ -152,3 +153,9 @@ def test_non_nullable_schema_is_preserved(protocol, verify): spark.createDataFrame(frame, schema, verifySchema=verify) assert not state["relations"] assert not state["uploads"] + + +def test_small_numpy_frames_avoid_importer_configuration_overhead(): + assert prepare_int64_table(pd.DataFrame({"a": np.arange(999, dtype=np.int64)}), None) is None + assert prepare_int64_table(pd.DataFrame({"a": np.arange(1000, dtype=np.int64)}), None) is not None + assert prepare_int64_table(pd.DataFrame({"a": pd.Series([None], dtype="Int64")}), None) is not None From e266e8eb8b180be57d2a1951e708e77fed6e6f19 Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Wed, 28 Oct 2026 21:53:51 -0700 Subject: [PATCH 06/23] [SPARK-60676][CONNECT] Retain IPC row fallback for NUL field names --- docs/architecture.md | 2 +- python/pyspark/sql/_pandas_input.py | 7 ++++--- python/tests/test_pandas_input.py | 21 +++++++++++++++++++++ 3 files changed, 26 insertions(+), 4 deletions(-) diff --git a/docs/architecture.md b/docs/architecture.md index 45bfb2c4..089232b3 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -149,7 +149,7 @@ row path to avoid configuration RPC overhead. Nullable `Int64` columns always use Arrow so schema inference does not depend on the first value being non-null. This initial path accepts inferred schemas, complete string column names, and explicit `StructType` schemas whose fields are all `LongType`. Other dtypes, -DDL/atomic schemas, schema conversions, and duplicate/non-string input names +DDL/atomic schemas, schema conversions, NUL names, and duplicate/non-string input names retain the existing row conversion path. `verifySchema` remains ignored for pandas Arrow input, matching Spark Connect. Inputs at or above `spark.sql.session.localRelationCacheThreshold` are serialized diff --git a/python/pyspark/sql/_pandas_input.py b/python/pyspark/sql/_pandas_input.py index 1015017f..028a18c0 100644 --- a/python/pyspark/sql/_pandas_input.py +++ b/python/pyspark/sql/_pandas_input.py @@ -34,7 +34,7 @@ def prepare_int64_table(frame, schema): width = len(frame.columns) if not width or not frame.columns.is_unique or not all( - isinstance(name, str) for name in frame.columns + isinstance(name, str) and "\x00" not in name for name in frame.columns ): return None dtypes = frame.dtypes @@ -47,12 +47,13 @@ def prepare_int64_table(frame, schema): return None if isinstance(schema, StructType): if len(schema.fields) != width or not all( - isinstance(field.dataType, LongType) for field in schema.fields + isinstance(field.dataType, LongType) and "\x00" not in field.name + for field in schema.fields ): return None elif schema is not None: if not isinstance(schema, (list, tuple)) or len(schema) != width or not all( - isinstance(name, str) for name in schema + isinstance(name, str) and "\x00" not in name for name in schema ): return None diff --git a/python/tests/test_pandas_input.py b/python/tests/test_pandas_input.py index 4cdaee7c..ff19b105 100644 --- a/python/tests/test_pandas_input.py +++ b/python/tests/test_pandas_input.py @@ -159,3 +159,24 @@ def test_small_numpy_frames_avoid_importer_configuration_overhead(): assert prepare_int64_table(pd.DataFrame({"a": np.arange(999, dtype=np.int64)}), None) is None assert prepare_int64_table(pd.DataFrame({"a": np.arange(1000, dtype=np.int64)}), None) is not None assert prepare_int64_table(pd.DataFrame({"a": pd.Series([None], dtype="Int64")}), None) is not None + + +@pytest.mark.parametrize("schema_kind", ["inferred", "names", "explicit"]) +def test_nul_names_keep_lossless_ipc_row_fallback(protocol, schema_kind): + import pyarrow as pa + + name = "a\x00b" + frame = pd.DataFrame({name if schema_kind == "inferred" else "original": + np.arange(1000, dtype=np.int64)}) + schema = None + if schema_kind == "names": + schema = [name] + elif schema_kind == "explicit": + schema = T.StructType([T.StructField(name, T.LongType())]) + # C Stream field names are NUL-terminated; IPC preserves the full string. + assert prepare_int64_table(frame, schema) is None + with _arrow_input_service(protocol) as (spark, state): + spark.createDataFrame(frame, schema).toArrow() + assert state["result"] == (1000, 0, 999) + payload = state["relations"][0].local_relation.data + assert pa.ipc.open_stream(payload).schema.names == [name] From 9942f21f5b167c98aad7d6cd2a4263068c968b38 Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Wed, 28 Oct 2026 22:34:39 -0700 Subject: [PATCH 07/23] [SPARK-60676][CONNECT] Include pandas regressions in the existing Arrow CI suite --- python/tests/test_arrow_handoff.py | 162 ++++++++++++++++++++++++- python/tests/test_pandas_input.py | 182 ----------------------------- 2 files changed, 160 insertions(+), 184 deletions(-) delete mode 100644 python/tests/test_pandas_input.py diff --git a/python/tests/test_arrow_handoff.py b/python/tests/test_arrow_handoff.py index dbc9e2d2..6e81d35b 100644 --- a/python/tests/test_arrow_handoff.py +++ b/python/tests/test_arrow_handoff.py @@ -26,17 +26,19 @@ import subprocess from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager -from datetime import datetime, timezone +from datetime import date, datetime, timezone from decimal import Decimal from pathlib import Path from urllib.parse import quote import grpc +import numpy as np import pandas as pd import pyarrow as pa import pytest from google.protobuf import descriptor_pb2, descriptor_pool, message_factory -from pyspark.sql import SparkSession +from pyspark.sql import SparkSession, types as T +from pyspark.sql._pandas_input import prepare_int64_table @pytest.fixture(scope="module") @@ -892,3 +894,159 @@ def import_module(name, *args, **kwargs): spark.createDataFrame(data) assert imported == [] assert state["config_requests"] == [] + + +# Signed-int64 pandas input uses the same Arrow importer and mock service. + +@pytest.mark.parametrize("schema_kind", ["inferred", "names", "explicit"]) +@pytest.mark.parametrize("verify", [True, False]) +def test_integer_columns_preserve_values_without_python_rows(protocol, monkeypatch, schema_kind, verify): + frame = pd.DataFrame({ + "a": pd.Series([2**53 + 1, None, -(2**63)], dtype="Int64"), + "b": np.array([2**63 - 1, -1, 0], dtype=np.int64), + }) + schema = None + names = ["a", "b"] + if schema_kind == "names": + schema = names = ["renamed_a", "renamed_b"] + elif schema_kind == "explicit": + schema = T.StructType([T.StructField("a", T.LongType()), + T.StructField("b", T.LongType(), False, {"source": "test"})]) + import pyspark.sql._arrow_input as helper + original = helper.prepare_table + observed = [] + + def capture(*args): + result = original(*args) + observed.append(result) + return result + + def no_rows(*args, **kwargs): + raise AssertionError("supported pandas input must not materialize Python rows") + + monkeypatch.setattr(helper, "prepare_table", capture) + monkeypatch.setattr(pd.DataFrame, "astype", no_rows) + with _arrow_input_service(protocol) as (spark, state): + df = spark.createDataFrame(frame, schema=schema, verifySchema=verify) + if schema_kind != "names": + assert df.isLocal() + assert df.toArrow().column(0).to_pylist() == [3] + assert state["result"] == (3, 2**53 + 1, -(2**63)) + table, resolved, renamed = observed[0] + assert table.to_pydict() == {"a": [2**53 + 1, None, -(2**63)], + "b": [2**63 - 1, -1, 0]} + assert (renamed or resolved.names) == names + if schema_kind == "explicit": + assert resolved == schema + assert not table.schema.field(1).nullable + assert resolved.fields[1].metadata == {"source": "test"} + + +@pytest.mark.parametrize("dtype", ["int64", "Int64"]) +def test_empty_integer_frame_requires_explicit_schema(protocol, dtype): + frame = pd.DataFrame({"a": pd.Series([], dtype=dtype)}) + with _arrow_input_service(protocol) as (spark, state): + with pytest.raises(ValueError): + spark.createDataFrame(frame) + df = spark.createDataFrame(frame, T.StructType([T.StructField("a", T.LongType())])) + assert df.isLocal() + + +@pytest.mark.parametrize("values", [ + [1.0, float("nan")], [True, False], ["x", None], + [datetime(2024, 1, 1), None], + [datetime(2024, 1, 1, tzinfo=timezone.utc), None], + [date(2024, 1, 1), None], [Decimal("1.20"), None], + [[1, 2], None], [{"a": 1}, None], +]) +def test_other_dtypes_keep_row_fallback(values): + frame = pd.DataFrame({"a": values}) + assert prepare_int64_table(frame, None) is None + assert prepare_int64_table(frame, T.StructType([T.StructField("a", T.LongType())])) is None + + +@pytest.mark.parametrize("schema", [ + "a long", T.LongType(), ["a", "b"], [1], + T.StructType([T.StructField("a", T.IntegerType())]), + T.StructType([T.StructField("a", T.DoubleType())]), + T.StructType([]), +]) +def test_schema_conversions_keep_row_fallback(schema): + assert prepare_int64_table(pd.DataFrame({"a": [1, 2]}), schema) is None + + +def test_duplicate_names_fallback_and_nullable_masks(): + frame = pd.DataFrame([[2**53 + 1, 2**53 + 3]], columns=["x", "x"], dtype="int64") + assert prepare_int64_table(frame, None) is None + assert prepare_int64_table(pd.DataFrame({0: [1]}), None) is None + frame = pd.DataFrame({"a": pd.Series([None, None], dtype="Int64")}) + assert prepare_int64_table(frame, None).column(0).to_pylist() == [None, None] + + +@pytest.mark.parametrize("dtype", ["int32", "uint64", "Int32", "UInt64", "int64[pyarrow]"]) +def test_other_integer_dtypes_keep_row_fallback(dtype): + assert prepare_int64_table(pd.DataFrame({"a": pd.Series([1, 2], dtype=dtype)}), None) is None + + +@pytest.mark.parametrize("rows", [100, 1000, 100000, 1000000]) +@pytest.mark.parametrize("explicit", [False, True]) +def test_integer_scales_reach_mock_service(protocol, monkeypatch, rows, explicit): + values = np.arange(rows, dtype=np.int64) + (2**53 + 1) + frame = pd.DataFrame({name: values for name in ("a", "b", "c")}) + schema = T.StructType([T.StructField(name, T.LongType()) for name in frame.columns]) + + def no_rows(*args, **kwargs): + raise AssertionError("integer input must stay columnar at every scale") + + if rows >= 1000: + monkeypatch.setattr(pd.DataFrame, "astype", no_rows) + with _arrow_input_service(protocol) as (spark, state): + spark.createDataFrame(frame, schema if explicit else None).toArrow() + assert state["result"] == (rows, 2**53 + 1, 2**53 + rows) + if frame.memory_usage(index=False).sum() > 1048576: + assert state["relations"][0].HasField("chunked_cached_local_relation") + + +def test_nullable_inference_does_not_depend_on_first_value(protocol): + frame = pd.DataFrame({"a": pd.Series([None, 2**53 + 1], dtype="Int64")}) + with _arrow_input_service(protocol) as (spark, state): + spark.createDataFrame(frame).toArrow() + assert state["result"] == (2, None, 2**53 + 1) + + +@pytest.mark.parametrize("verify", [True, False]) +def test_non_nullable_schema_is_preserved(protocol, verify): + frame = pd.DataFrame({"a": pd.Series([1, None], dtype="Int64")}) + schema = T.StructType([T.StructField("a", T.LongType(), False)]) + with _arrow_input_service(protocol) as (spark, state): + with pytest.raises(ValueError, match="null"): + spark.createDataFrame(frame, schema, verifySchema=verify) + assert not state["relations"] + assert not state["uploads"] + + +def test_small_numpy_frames_avoid_importer_configuration_overhead(): + assert prepare_int64_table(pd.DataFrame({"a": np.arange(999, dtype=np.int64)}), None) is None + assert prepare_int64_table(pd.DataFrame({"a": np.arange(1000, dtype=np.int64)}), None) is not None + assert prepare_int64_table(pd.DataFrame({"a": pd.Series([None], dtype="Int64")}), None) is not None + + +@pytest.mark.parametrize("schema_kind", ["inferred", "names", "explicit"]) +def test_nul_names_keep_lossless_ipc_row_fallback(protocol, schema_kind): + import pyarrow as pa + + name = "a\x00b" + frame = pd.DataFrame({name if schema_kind == "inferred" else "original": + np.arange(1000, dtype=np.int64)}) + schema = None + if schema_kind == "names": + schema = [name] + elif schema_kind == "explicit": + schema = T.StructType([T.StructField(name, T.LongType())]) + # C Stream field names are NUL-terminated; IPC preserves the full string. + assert prepare_int64_table(frame, schema) is None + with _arrow_input_service(protocol) as (spark, state): + spark.createDataFrame(frame, schema).toArrow() + assert state["result"] == (1000, 0, 999) + payload = state["relations"][0].local_relation.data + assert pa.ipc.open_stream(payload).schema.names == [name] diff --git a/python/tests/test_pandas_input.py b/python/tests/test_pandas_input.py deleted file mode 100644 index ff19b105..00000000 --- a/python/tests/test_pandas_input.py +++ /dev/null @@ -1,182 +0,0 @@ -# -# Licensed to the Apache Software Foundation (ASF) under one or more -# contributor license agreements. See the NOTICE file distributed with -# this work for additional information regarding copyright ownership. -# The ASF licenses this file to You under the Apache License, Version 2.0 -# (the "License"); you may not use this file except in compliance with -# the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -"""Signed-int64 pandas column handoff and conservative fallback coverage.""" - -from datetime import date, datetime, timezone -from decimal import Decimal - -import numpy as np -import pandas as pd -import pytest - -from pyspark.sql import types as T -from pyspark.sql._pandas_input import prepare_int64_table -from test_arrow_handoff import _arrow_input_service, protocol - - -@pytest.mark.parametrize("schema_kind", ["inferred", "names", "explicit"]) -@pytest.mark.parametrize("verify", [True, False]) -def test_integer_columns_preserve_values_without_python_rows(protocol, monkeypatch, schema_kind, verify): - frame = pd.DataFrame({ - "a": pd.Series([2**53 + 1, None, -(2**63)], dtype="Int64"), - "b": np.array([2**63 - 1, -1, 0], dtype=np.int64), - }) - schema = None - names = ["a", "b"] - if schema_kind == "names": - schema = names = ["renamed_a", "renamed_b"] - elif schema_kind == "explicit": - schema = T.StructType([T.StructField("a", T.LongType()), - T.StructField("b", T.LongType(), False, {"source": "test"})]) - import pyspark.sql._arrow_input as helper - original = helper.prepare_table - observed = [] - - def capture(*args): - result = original(*args) - observed.append(result) - return result - - def no_rows(*args, **kwargs): - raise AssertionError("supported pandas input must not materialize Python rows") - - monkeypatch.setattr(helper, "prepare_table", capture) - monkeypatch.setattr(pd.DataFrame, "astype", no_rows) - with _arrow_input_service(protocol) as (spark, state): - df = spark.createDataFrame(frame, schema=schema, verifySchema=verify) - if schema_kind != "names": - assert df.isLocal() - assert df.toArrow().column(0).to_pylist() == [3] - assert state["result"] == (3, 2**53 + 1, -(2**63)) - table, resolved, renamed = observed[0] - assert table.to_pydict() == {"a": [2**53 + 1, None, -(2**63)], - "b": [2**63 - 1, -1, 0]} - assert (renamed or resolved.names) == names - if schema_kind == "explicit": - assert resolved == schema - assert not table.schema.field(1).nullable - assert resolved.fields[1].metadata == {"source": "test"} - - -@pytest.mark.parametrize("dtype", ["int64", "Int64"]) -def test_empty_integer_frame_requires_explicit_schema(protocol, dtype): - frame = pd.DataFrame({"a": pd.Series([], dtype=dtype)}) - with _arrow_input_service(protocol) as (spark, state): - with pytest.raises(ValueError): - spark.createDataFrame(frame) - df = spark.createDataFrame(frame, T.StructType([T.StructField("a", T.LongType())])) - assert df.isLocal() - - -@pytest.mark.parametrize("values", [ - [1.0, float("nan")], [True, False], ["x", None], - [datetime(2024, 1, 1), None], - [datetime(2024, 1, 1, tzinfo=timezone.utc), None], - [date(2024, 1, 1), None], [Decimal("1.20"), None], - [[1, 2], None], [{"a": 1}, None], -]) -def test_other_dtypes_keep_row_fallback(values): - frame = pd.DataFrame({"a": values}) - assert prepare_int64_table(frame, None) is None - assert prepare_int64_table(frame, T.StructType([T.StructField("a", T.LongType())])) is None - - -@pytest.mark.parametrize("schema", [ - "a long", T.LongType(), ["a", "b"], [1], - T.StructType([T.StructField("a", T.IntegerType())]), - T.StructType([T.StructField("a", T.DoubleType())]), - T.StructType([]), -]) -def test_schema_conversions_keep_row_fallback(schema): - assert prepare_int64_table(pd.DataFrame({"a": [1, 2]}), schema) is None - - -def test_duplicate_names_fallback_and_nullable_masks(): - frame = pd.DataFrame([[2**53 + 1, 2**53 + 3]], columns=["x", "x"], dtype="int64") - assert prepare_int64_table(frame, None) is None - assert prepare_int64_table(pd.DataFrame({0: [1]}), None) is None - frame = pd.DataFrame({"a": pd.Series([None, None], dtype="Int64")}) - assert prepare_int64_table(frame, None).column(0).to_pylist() == [None, None] - - -@pytest.mark.parametrize("dtype", ["int32", "uint64", "Int32", "UInt64", "int64[pyarrow]"]) -def test_other_integer_dtypes_keep_row_fallback(dtype): - assert prepare_int64_table(pd.DataFrame({"a": pd.Series([1, 2], dtype=dtype)}), None) is None - - -@pytest.mark.parametrize("rows", [100, 1000, 100000, 1000000]) -@pytest.mark.parametrize("explicit", [False, True]) -def test_integer_scales_reach_mock_service(protocol, monkeypatch, rows, explicit): - values = np.arange(rows, dtype=np.int64) + (2**53 + 1) - frame = pd.DataFrame({name: values for name in ("a", "b", "c")}) - schema = T.StructType([T.StructField(name, T.LongType()) for name in frame.columns]) - - def no_rows(*args, **kwargs): - raise AssertionError("integer input must stay columnar at every scale") - - if rows >= 1000: - monkeypatch.setattr(pd.DataFrame, "astype", no_rows) - with _arrow_input_service(protocol) as (spark, state): - spark.createDataFrame(frame, schema if explicit else None).toArrow() - assert state["result"] == (rows, 2**53 + 1, 2**53 + rows) - if frame.memory_usage(index=False).sum() > 1048576: - assert state["relations"][0].HasField("chunked_cached_local_relation") - - -def test_nullable_inference_does_not_depend_on_first_value(protocol): - frame = pd.DataFrame({"a": pd.Series([None, 2**53 + 1], dtype="Int64")}) - with _arrow_input_service(protocol) as (spark, state): - spark.createDataFrame(frame).toArrow() - assert state["result"] == (2, None, 2**53 + 1) - - -@pytest.mark.parametrize("verify", [True, False]) -def test_non_nullable_schema_is_preserved(protocol, verify): - frame = pd.DataFrame({"a": pd.Series([1, None], dtype="Int64")}) - schema = T.StructType([T.StructField("a", T.LongType(), False)]) - with _arrow_input_service(protocol) as (spark, state): - with pytest.raises(ValueError, match="null"): - spark.createDataFrame(frame, schema, verifySchema=verify) - assert not state["relations"] - assert not state["uploads"] - - -def test_small_numpy_frames_avoid_importer_configuration_overhead(): - assert prepare_int64_table(pd.DataFrame({"a": np.arange(999, dtype=np.int64)}), None) is None - assert prepare_int64_table(pd.DataFrame({"a": np.arange(1000, dtype=np.int64)}), None) is not None - assert prepare_int64_table(pd.DataFrame({"a": pd.Series([None], dtype="Int64")}), None) is not None - - -@pytest.mark.parametrize("schema_kind", ["inferred", "names", "explicit"]) -def test_nul_names_keep_lossless_ipc_row_fallback(protocol, schema_kind): - import pyarrow as pa - - name = "a\x00b" - frame = pd.DataFrame({name if schema_kind == "inferred" else "original": - np.arange(1000, dtype=np.int64)}) - schema = None - if schema_kind == "names": - schema = [name] - elif schema_kind == "explicit": - schema = T.StructType([T.StructField(name, T.LongType())]) - # C Stream field names are NUL-terminated; IPC preserves the full string. - assert prepare_int64_table(frame, schema) is None - with _arrow_input_service(protocol) as (spark, state): - spark.createDataFrame(frame, schema).toArrow() - assert state["result"] == (1000, 0, 999) - payload = state["relations"][0].local_relation.data - assert pa.ipc.open_stream(payload).schema.names == [name] From 5edd70017c67d5e8bb98cfc3ee6538b333722740 Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Thu, 29 Oct 2026 13:56:36 -0700 Subject: [PATCH 08/23] [SPARK-60676][CONNECT] Preserve integer pandas schema and fallback semantics Use explicit LongType schema names when constructing Arrow columns. Limit the columnar path to concrete pandas/NumPy arrays and native schema descriptors, including authentic converter aliases retained by already imported modules. Keep nullable values in non-nullable fields on the native row path so local validation still precedes configuration RPCs. Support warm Arrow importers that resolve converters locally. Add mock-RPC regressions for all emitted integer cells, schema metadata and nullability, custom conversion hooks, and repeated warm construction. Co-authored-by: Isaac --- python/pyspark/sql/_pandas_input.py | 112 ++++++++- python/tests/test_arrow_handoff.py | 356 +++++++++++++++++++++++++++- 2 files changed, 451 insertions(+), 17 deletions(-) diff --git a/python/pyspark/sql/_pandas_input.py b/python/pyspark/sql/_pandas_input.py index 028a18c0..8e5f6284 100644 --- a/python/pyspark/sql/_pandas_input.py +++ b/python/pyspark/sql/_pandas_input.py @@ -18,8 +18,6 @@ """Conservative columnar pandas input; unsupported conversions use the row path.""" -from pyspark.sql.types import LongType, StructType - def prepare_int64_table(frame, schema): """Return an Arrow table only when each column is already a signed 64-bit integer. @@ -29,38 +27,126 @@ def prepare_int64_table(frame, schema): deliberately retain the existing row path until their semantics are covered. ``verifySchema`` is ignored for pandas Arrow input, as in Spark Connect. """ + import sys + from types import GetSetDescriptorType, WrapperDescriptorType + import numpy as np import pandas as pd + try: + from pandas.core.arrays.integer import Int64Dtype, IntegerArray + from pandas.core.arrays.numpy_ import NumpyExtensionArray + from pandas.core.frame import DataFrame + from pandas.core.series import Series + except ImportError: + # Older supported pandas versions call the NumPy wrapper PandasArray. + # Keep their existing row conversion when these concrete types are absent. + return None + + from pyspark._pyspark import ( + BooleanType, ByteType, ShortType, IntegerType, LongType, StructField, StructType, + ) + from pyspark.sql import types as sql_types + + # A subclass or custom extension array can change the values exposed to Arrow. + # Check concrete identities before consulting dtype equality or conversion hooks. + if pd.DataFrame is not DataFrame or type(frame) is not DataFrame: + return None width = len(frame.columns) if not width or not frame.columns.is_unique or not all( - isinstance(name, str) and "\x00" not in name for name in frame.columns + type(name) is str and "\x00" not in name for name in frame.columns ): return None dtypes = frame.dtypes - if not all(dtype == np.dtype("int64") or isinstance(dtype, pd.Int64Dtype) + numpy_dtype = np.dtype("int64") + if not all((type(dtype) is type(numpy_dtype) and dtype == numpy_dtype) + or type(dtype) is Int64Dtype for dtype in dtypes): return None # The importer's configuration RPCs outweigh row conversion for small NumPy # frames. Nullable extension columns still use Arrow for reliable null typing. - if len(frame) < 1000 and not any(isinstance(dtype, pd.Int64Dtype) for dtype in dtypes): + if len(frame) < 1000 and not any(type(dtype) is Int64Dtype for dtype in dtypes): + return None + + # The native row path reads native schema fields directly. The Arrow converter + # instead uses Python descriptors and imported class aliases, including aliases + # captured before this helper was loaded. Customized schemas retain the row path. + if sys.modules.get("pyspark.sql.types") is not sql_types: + return None + importer = sys.modules.get("pyspark.sql._arrow_input") + if importer is not None and getattr(importer, "StructType", None) is not StructType: return None - if isinstance(schema, StructType): - if len(schema.fields) != width or not all( - isinstance(field.dataType, LongType) and "\x00" not in field.name - for field in schema.fields + converter = sys.modules.get("pyspark.sql.pandas.types") + namespaces = [vars(sql_types)] + if converter is not None: + namespaces.append(vars(converter)) + if importer is not None: + # An importer may still hold functions from a removed converter module. + for name in ("from_arrow_schema", "to_arrow_schema"): + # Later importers resolve these converters locally inside prepare_table. + if name not in vars(importer): + continue + namespace = getattr(getattr(importer, name, None), "__globals__", None) + if type(namespace) is not dict: + return None + namespaces.append(namespace) + for name, cls in ( + ("BooleanType", BooleanType), ("ByteType", ByteType), ("ShortType", ShortType), + ("IntegerType", IntegerType), ("LongType", LongType), + ("StructField", StructField), ("StructType", StructType), + ): + if any(namespace.get(name) is not cls for namespace in namespaces): + return None + for cls in (LongType, StructField, StructType): + if (cls.__getattribute__ is not object.__getattribute__ + or any("__class__" in vars(base) for base in cls.__mro__ if base is not object)): + return None + for cls, names, descriptor_type in ( + (StructType, ("fields", "names"), GetSetDescriptorType), + (StructType, ("__iter__",), WrapperDescriptorType), + (StructField, ("name", "dataType", "nullable", "metadata"), GetSetDescriptorType), + ): + for name in names: + descriptor = vars(cls).get(name) + if (type(descriptor) is not descriptor_type + or descriptor.__objclass__ is not cls or descriptor.__name__ != name): + return None + + names = list(frame.columns) + fields = None + if type(schema) is StructType: + fields = schema.fields + if len(fields) != width or not all( + type(field) is StructField and type(field.dataType) is LongType + and type(field.name) is str and "\x00" not in field.name + for field in fields ): return None + names = schema.names + if len(set(names)) != width: + return None elif schema is not None: - if not isinstance(schema, (list, tuple)) or len(schema) != width or not all( - isinstance(name, str) and "\x00" not in name for name in schema + if type(schema) not in (list, tuple) or len(schema) != width or not all( + type(name) is str and "\x00" not in name for name in schema ): return None + columns = [frame.iloc[:, i] for i in range(width)] + for i, (column, dtype) in enumerate(zip(columns, dtypes)): + if type(column) is not Series: + return None + expected = IntegerArray if type(dtype) is Int64Dtype else NumpyExtensionArray + if type(column.array) is not expected: + return None + # The native row path rejects invalid nullability before any Config RPC. + if (fields is not None and not fields[i].nullable + and type(dtype) is Int64Dtype and column.hasnans): + return None + import pyarrow as pa # Arrow consumes each Series' integer buffer and nullable mask independently. return pa.Table.from_arrays( - [pa.array(frame.iloc[:, i], type=pa.int64(), from_pandas=True) for i in range(width)], - names=list(frame.columns), + [pa.array(column, type=pa.int64(), from_pandas=True) for column in columns], + names=names, ) diff --git a/python/tests/test_arrow_handoff.py b/python/tests/test_arrow_handoff.py index 6e81d35b..dace30a4 100644 --- a/python/tests/test_arrow_handoff.py +++ b/python/tests/test_arrow_handoff.py @@ -405,6 +405,7 @@ def _arrow_input_service( protocol, *, upload_error=None, + config_error=False, user_id=None, user_agent=None, allow_file_artifacts=False, @@ -426,7 +427,8 @@ def _arrow_input_service( } config.update(overrides) state = {"cache": {}, "uploads": 0, "relations": [], "message_sizes": [], - "config_requests": [], "upload_contexts": [], "status_contexts": []} + "config_requests": [], "input_tables": [], + "upload_contexts": [], "status_contexts": []} server_id = "00000000-0000-0000-0000-000000000001" def attribution(request, context): @@ -448,6 +450,8 @@ def configure(raw, context): operation = request.operation.WhichOneof("op_type") keys = getattr(request.operation, operation).keys state["config_requests"].append((operation, list(keys))) + if config_error: + context.abort(grpc.StatusCode.UNAVAILABLE, "configuration unavailable") return protocol("ConfigResponse", session_id=request.session_id, pairs=[dict(key=k, value=config[k]) for k in keys if config.get(k) is not None] @@ -528,6 +532,7 @@ def execute(raw, context): if relation.local_relation.HasField("schema") else None ) rows, first, last = 0, None, None + batches = [] state["arrow_schemas"], state["batch_rows"] = [], [] for payload in payloads: reader = pa.ipc.open_stream(payload) @@ -536,12 +541,14 @@ def execute(raw, context): state["batch_rows"].append(stream_rows) for batch in reader: stream_rows.append(batch.num_rows) + batches.append(batch) rows += batch.num_rows if batch.num_rows: if first is None: first = batch.column(0)[0].as_py() last = batch.column(0)[-1].as_py() state["result"] = (rows, first, last) + state["input_tables"].append(pa.Table.from_batches(batches)) result = pa.table({"rows": [rows]}) sink = pa.BufferOutputStream() with pa.ipc.new_stream(sink, result.schema) as writer: @@ -1015,14 +1022,17 @@ def test_nullable_inference_does_not_depend_on_first_value(protocol): @pytest.mark.parametrize("verify", [True, False]) -def test_non_nullable_schema_is_preserved(protocol, verify): +@pytest.mark.parametrize("config_error", [False, True]) +def test_non_nullable_schema_is_preserved(protocol, verify, config_error): frame = pd.DataFrame({"a": pd.Series([1, None], dtype="Int64")}) schema = T.StructType([T.StructField("a", T.LongType(), False)]) - with _arrow_input_service(protocol) as (spark, state): - with pytest.raises(ValueError, match="null"): + assert prepare_int64_table(frame, schema) is None + with _arrow_input_service(protocol, config_error=config_error) as (spark, state): + with pytest.raises(RuntimeError, match="null"): spark.createDataFrame(frame, schema, verifySchema=verify) assert not state["relations"] assert not state["uploads"] + assert state["config_requests"] == [] def test_small_numpy_frames_avoid_importer_configuration_overhead(): @@ -1050,3 +1060,341 @@ def test_nul_names_keep_lossless_ipc_row_fallback(protocol, schema_kind): assert state["result"] == (1000, 0, 999) payload = state["relations"][0].local_relation.data assert pa.ipc.open_stream(payload).schema.names == [name] + + +@pytest.mark.parametrize("nullable", [False, True]) +@pytest.mark.parametrize("rows", [999, 1000]) +def test_integer_explicit_renamed_schema_preserves_all_ipc_cells(protocol, monkeypatch, nullable, rows): + values = ([2**53 + 1, -(2**63), 2**63 - 1, 0] * 250)[:rows] + other = [None if nullable and i % 7 == 0 else value + for i, value in enumerate(reversed(values))] + frame = pd.DataFrame({"original": np.asarray(values, dtype="int64"), + "second": pd.Series(other, dtype="Int64" if nullable else "int64")}) + schema = T.StructType([ + T.StructField("renamed", T.LongType(), False, {"source": "native", "id": 3}), + T.StructField("nullable", T.LongType(), True, {"tag": [1, 2]}), + ]) + + def no_rows(*args, **kwargs): + raise AssertionError("trusted integer columns must stay columnar") + + if rows >= 1000 or nullable: + monkeypatch.setattr(pd.DataFrame, "astype", no_rows) + with _arrow_input_service(protocol) as (spark, state): + spark.createDataFrame(frame, schema).toArrow() + table = state["input_tables"][0] + assert table.column_names == ["renamed", "nullable"] + assert [column.to_pylist() for column in table.columns] == [values, other] + assert table.schema.field(0).type == pa.int64() + assert table.schema.field(0).nullable is False + assert table.schema.field(1).nullable is True + if rows >= 1000 or nullable: + from pyspark.sql.pandas.types import from_arrow_schema + assert from_arrow_schema(table.schema).json() == schema.json(), ( + from_arrow_schema(table.schema).json(), schema.json(), table.schema) + assert state["config_requests"] + else: + assert state["config_requests"] == [] + + +def test_integer_dataframe_and_name_subclasses_keep_fallback(monkeypatch): + class Frame(pd.DataFrame): + pass + + class Name(str): + pass + + class Names(list): + pass + + frame = pd.DataFrame({"a": np.arange(1000, dtype="int64")}) + assert prepare_int64_table(Frame(frame), None) is None + assert prepare_int64_table(frame, Names(["a"])) is None + assert prepare_int64_table(frame, [Name("a")]) is None + assert prepare_int64_table(pd.DataFrame({Name("a"): frame["a"]}), None) is None + monkeypatch.setattr(pd, "DataFrame", Frame) + assert prepare_int64_table(Frame(frame), None) is None + + +def test_integer_missing_internal_pandas_array_class_keeps_fallback(): + import sys + + script = r''' +import sys +import types +import numpy as np +import pandas as pd +from pyspark.sql._pandas_input import prepare_int64_table +frame = pd.DataFrame({"a": np.arange(1000, dtype="int64")}) +nullable = pd.DataFrame({"a": pd.Series([None, 1], dtype="Int64")}) +# Older pandas exposes PandasArray rather than NumpyExtensionArray. Missing +# optional internals must disable the optimization without breaking pandas input. +sys.modules["pandas.core.arrays.numpy_"] = types.ModuleType("pandas.core.arrays.numpy_") +assert prepare_int64_table(frame, None) is None +assert prepare_int64_table(nullable, None) is None +''' + result = subprocess.run([sys.executable, "-c", script], + capture_output=True, text=True, timeout=30) + assert result.returncode == 0, result.stdout + result.stderr + + +@pytest.mark.parametrize("class_spoof", [False, True]) +def test_integer_custom_extension_dtype_cannot_select_arrow(protocol, class_spoof): + from pandas.api.extensions import ExtensionArray, ExtensionDtype, take + + arrow_calls = [] + + class Int64EqualDtype(ExtensionDtype): + name = "int64_equal" + type = int + kind = "i" + na_value = None + + @classmethod + def construct_array_type(cls): + return DifferentArrowArray + + def __eq__(self, other): + return (isinstance(other, Int64EqualDtype) + or isinstance(other, np.dtype) and other == np.dtype("int64")) + + def __hash__(self): + return hash(self.name) + + class DifferentArrowArray(ExtensionArray): + def __init__(self, values): + self._values = np.asarray(values, dtype=object) + + @property + def dtype(self): + return Int64EqualDtype() + + @property + def nbytes(self): + return self._values.nbytes + + def __len__(self): + return len(self._values) + + def __getitem__(self, key): + result = self._values[key] + return type(self)(result) if isinstance(result, np.ndarray) else result + + def isna(self): + return pd.isna(self._values) + + def take(self, indices, allow_fill=False, fill_value=None): + return type(self)(take(self._values, indices, allow_fill=allow_fill, + fill_value=fill_value)) + + def copy(self): + return type(self)(self._values.copy()) + + @classmethod + def _from_sequence(cls, scalars, *, dtype=None, copy=False): + values = np.asarray(scalars, dtype=object) + return cls(values.copy() if copy else values) + + def to_numpy(self, dtype=None, copy=False, na_value=None): + if dtype is not None: + return self._values.astype(dtype, copy=copy) + return self._values.copy() if copy else self._values + + def __arrow_array__(self, type=None): + arrow_calls.append("arrow") + return pa.array([-7] * len(self), type=type) + + values = [2**53 + 1, None, -(2**63), 0] * 250 + frame = pd.DataFrame({"a": pd.Series(DifferentArrowArray(values), copy=False)}) + if class_spoof: + Int64EqualDtype.__getattribute__ = lambda self, name: ( + type(np.dtype("int64")) if name == "__class__" + else object.__getattribute__(self, name)) + assert frame.dtypes.iloc[0] == np.dtype("int64") + assert isinstance(frame.dtypes.iloc[0], np.dtype) is class_spoof + schema = T.StructType([T.StructField("renamed", T.LongType())]) + assert prepare_int64_table(frame, schema) is None + with _arrow_input_service(protocol) as (spark, state): + spark.createDataFrame(frame, schema).toArrow() + assert state["config_requests"] == [] + assert state["input_tables"][0].column(0).to_pylist() == values + assert arrow_calls == [] + + +@pytest.mark.parametrize("custom", ["dtype", "array"]) +def test_integer_nullable_dtype_and_array_subclasses_keep_fallback(protocol, custom): + from pandas.core.arrays.integer import IntegerArray, Int64Dtype + + arrow_calls = [] + + class CustomDtype(Int64Dtype): + pass + + class CustomArray(IntegerArray): + @property + def dtype(self): + if custom == "dtype": + return CustomDtype() + return Int64Dtype() + + def __arrow_array__(self, type=None): + arrow_calls.append("arrow") + return pa.array([-7] * len(self), type=type) + + values = [2**53 + 1, 0, -(2**63), 7] * 250 + array = CustomArray(np.asarray(values, dtype="int64"), np.zeros(1000, dtype=bool)) + frame = pd.DataFrame({"a": pd.Series(array, copy=False)}) + assert type(frame["a"].array) is CustomArray + schema = T.StructType([T.StructField("a", T.LongType())]) + assert prepare_int64_table(frame, schema) is None + with _arrow_input_service(protocol) as (spark, state): + spark.createDataFrame(frame, schema).toArrow() + assert state["config_requests"] == [] + assert state["input_tables"][0].column(0).to_pylist() == values + assert arrow_calls == [] + + +def _integer_guard_process(mode, argument, before_import=False): + """Class mutation and converter imports must be isolated from other tests.""" + import sys + + if mode == "captured_alias": + import pyspark.sql._arrow_input as importer + if not hasattr(importer, "to_arrow_schema"): + pytest.skip("this importer resolves converters locally and holds no captured aliases") + + script = r''' +import importlib +import importlib.util +import sys +import tempfile +from pathlib import Path +import numpy as np +import pandas as pd +from pyspark.sql import types as T +from pyspark import _pyspark as native + +spec = importlib.util.spec_from_file_location("integer_guard_capture", sys.argv[4]) +capture = importlib.util.module_from_spec(spec) +sys.modules[spec.name] = capture +spec.loader.exec_module(capture) +values = [2**53 + 1, -(2**63), 2**63 - 1, 0] * 250 +frame = pd.DataFrame({"renamed": np.asarray(values, dtype="int64")}) +schema = T.StructType([T.StructField("renamed", T.LongType(), False, {"source": "native"})]) +mode, argument, before_import = sys.argv[1:4] +if before_import == "True": + sys.modules.pop("pyspark.sql._pandas_input", None) +if mode == "descriptor": + class_name, attribute = argument.split(":") + cls = getattr(native, class_name) + if attribute == "alias": + setattr(T, class_name, object) + elif attribute == "__getattribute__": + setattr(cls, attribute, lambda self, name: object.__getattribute__(self, name)) + elif attribute == "__class__": + setattr(cls, "__getattribute__", lambda self, name: ( + object if name == "__class__" else object.__getattribute__(self, name))) + elif attribute == "__iter__": + setattr(cls, attribute, lambda self: iter([T.StructField("changed", T.BooleanType())])) + else: + value = {"names": ["changed"], "fields": [], "name": "changed", + "dataType": T.BooleanType(), "nullable": True, + "metadata": {"changed": True}}[attribute] + setattr(cls, attribute, property(lambda self: value)) +elif mode in ("public_alias", "converter_alias", "captured_alias"): + replacement = native.BooleanType if argument == "LongType" else native.LongType + if mode == "public_alias": + sys.modules.pop("pyspark.sql.pandas.types", None) + assert "pyspark.sql.pandas.types" not in sys.modules + setattr(T, argument, replacement) + else: + converter = importlib.import_module("pyspark.sql.pandas.types") + if mode == "captured_alias": + importer = importlib.import_module("pyspark.sql._arrow_input") + setattr(converter, argument, replacement) + if mode == "captured_alias": + sys.modules.pop("pyspark.sql.pandas.types") + current = importlib.import_module("pyspark.sql.pandas.types") + assert current is not converter + assert importer.to_arrow_schema.__globals__ is vars(converter) +elif mode == "importer_alias": + importer = importlib.import_module("pyspark.sql._arrow_input") + importer.StructType = None if argument == "none" else native.LongType +else: + raise AssertionError(mode) +class Factory: + def mktemp(self, name): + return Path(tempfile.mkdtemp(prefix=name)) +with capture._arrow_input_service(capture.protocol.__wrapped__(Factory())) as (spark, state): + spark.createDataFrame(frame, schema).toArrow() + assert state["config_requests"] == [], state["config_requests"] + table = state["input_tables"][0] + assert table.schema.names == ["renamed"], table.schema + assert table.column(0).type == __import__("pyarrow").int64() + assert table.column(0).to_pylist() == values +''' + result = subprocess.run([sys.executable, "-c", script, mode, argument, + str(before_import), __file__], + capture_output=True, text=True, timeout=30) + assert result.returncode == 0, result.stdout + result.stderr + + +@pytest.mark.parametrize("class_name,attribute", [ + ("StructType", "fields"), ("StructType", "names"), ("StructType", "__iter__"), + ("StructField", "name"), ("StructField", "dataType"), ("StructField", "nullable"), + ("StructField", "metadata"), ("StructType", "__getattribute__"), + ("StructField", "__getattribute__"), ("LongType", "__getattribute__"), + ("StructType", "__class__"), ("StructField", "__class__"), ("LongType", "__class__"), + ("StructType", "alias"), ("StructField", "alias"), ("LongType", "alias"), +]) +@pytest.mark.parametrize("before_import", [False, True]) +def test_integer_custom_schema_descriptors_keep_native_row_path(class_name, attribute, before_import): + _integer_guard_process("descriptor", f"{class_name}:{attribute}", before_import) + + +@pytest.mark.parametrize("class_name", ["BooleanType", "ByteType", "ShortType", "IntegerType", "LongType"]) +@pytest.mark.parametrize("mode", ["public_alias", "converter_alias", "captured_alias"]) +def test_integer_rebound_converter_aliases_keep_native_row_path(class_name, mode): + _integer_guard_process(mode, class_name) + + +@pytest.mark.parametrize("replacement", ["none", "long"]) +def test_integer_rebound_importer_schema_alias_keeps_native_row_path(replacement): + _integer_guard_process("importer_alias", replacement) + + +def test_integer_warm_importer_without_global_converter_aliases(protocol, monkeypatch): + import pyspark.sql._arrow_input as importer + + original = importer.prepare_table + + def local_imports(table, schema, *configs): + # Newer importers import converters within prepare_table. Exercise that + # module shape even on the pinned PR118 dependency used by this branch. + from pyspark.sql.pandas.types import from_arrow_schema, to_arrow_schema + table, resolved, names = original(table, schema, *configs) + assert from_arrow_schema(table.schema).fields + assert to_arrow_schema(resolved).names == table.column_names + return table, resolved, names + + frame = pd.DataFrame({"a": np.arange(1000, dtype="int64")}) + schema = T.StructType([T.StructField("renamed", T.LongType())]) + with _arrow_input_service(protocol) as (spark, state): + spark.createDataFrame(pa.table({"warm": [1, 2]})).toArrow() + # Keep original's module globals available in the dependency importer; + # the full-stack validation separately uses the actual local-import helper. + import types + substitute = types.ModuleType(importer.__name__) + substitute.StructType = T.StructType + substitute.prepare_table = local_imports + monkeypatch.setitem(__import__("sys").modules, importer.__name__, substitute) + + def no_rows(*args, **kwargs): + raise AssertionError("warm integer input must stay columnar") + + monkeypatch.setattr(pd.DataFrame, "astype", no_rows) + for _ in range(3): + spark.createDataFrame(frame, schema).toArrow() + for table in state["input_tables"][1:]: + assert table.column_names == ["renamed"] + assert table.column(0).to_pylist() == list(range(1000)) From 080aa16133a2037963a5488867d5e902277d6ff3 Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Thu, 29 Oct 2026 17:36:46 -0700 Subject: [PATCH 09/23] [SPARK-60676][TESTS] Preserve subclass column labels on pandas 3 Build the subclass-label fallback fixture with an explicit object Index and assert its concrete label type. Pandas 3 otherwise normalizes constructor labels to builtin strings before the unchanged integer helper sees them. The original regression passes on pandas 2.3.3 and fails on 3.1.0; the fixed regression and full CI-selected Arrow/drop-in suites pass on both versions. Runtime Python and native sources remain unchanged. Co-authored-by: Isaac --- python/tests/test_arrow_handoff.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/python/tests/test_arrow_handoff.py b/python/tests/test_arrow_handoff.py index dace30a4..82a7b79a 100644 --- a/python/tests/test_arrow_handoff.py +++ b/python/tests/test_arrow_handoff.py @@ -1111,7 +1111,12 @@ class Names(list): assert prepare_int64_table(Frame(frame), None) is None assert prepare_int64_table(frame, Names(["a"])) is None assert prepare_int64_table(frame, [Name("a")]) is None - assert prepare_int64_table(pd.DataFrame({Name("a"): frame["a"]}), None) is None + # Pandas 3 normalizes constructor labels to plain strings by default. + # An object Index keeps the subclass this fallback regression exercises. + named_frame = frame.copy() + named_frame.columns = pd.Index([Name("a")], dtype=object) + assert type(named_frame.columns[0]) is Name + assert prepare_int64_table(named_frame, None) is None monkeypatch.setattr(pd, "DataFrame", Frame) assert prepare_int64_table(Frame(frame), None) is None From abbc23a7b7edac978eea9d0c3e5cd949b2833e76 Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Thu, 29 Oct 2026 23:59:29 -0700 Subject: [PATCH 10/23] [SPARK-60676][CONNECT] Preserve customized pandas columns hooks Decline noncanonical columns indexes and their direct string-storage hooks before the integer fast path inspects labels or normalizes cells. Keep native tolist errors and mutations in their original order and cover public construction on both pandas versions. Co-authored-by: Isaac --- python/pyspark/sql/_pandas_input.py | 97 +++++++++++++++ python/tests/test_arrow_handoff.py | 175 ++++++++++++++++++++++++++++ 2 files changed, 272 insertions(+) diff --git a/python/pyspark/sql/_pandas_input.py b/python/pyspark/sql/_pandas_input.py index 8e5f6284..3a7f0e3c 100644 --- a/python/pyspark/sql/_pandas_input.py +++ b/python/pyspark/sql/_pandas_input.py @@ -19,6 +19,101 @@ """Conservative columnar pandas input; unsupported conversions use the row path.""" +def _has_canonical_columns(columns): + """Decline label containers whose row-path ``tolist`` can run custom hooks.""" + from types import FunctionType + + import numpy as np + import pandas as pd + from pandas.core import base as base_module + from pandas.core.indexes import base as index_module + + Index = index_module.Index + if (pd.Index is not Index or type(columns) is not Index + or Index.__getattribute__ is not object.__getattribute__ + or "tolist" in vars(columns)): + return False + + def ordinary(method, name, module): + return (type(method) is FunctionType and method.__name__ == name + and method.__module__ == module.__name__ + and method.__globals__ is vars(module)) + + # The native row path calls columns.tolist() before normalizing any cells. + # Decline overrides without invoking them, so fallback preserves that order. + if (Index.tolist is not base_module.IndexOpsMixin.tolist + or not ordinary(Index.tolist, "tolist", base_module)): + return False + values = vars(Index).get("_values") + if (type(values) is not property + or not ordinary(values.fget, "_values", index_module)): + return False + backing = vars(columns).get("_data") + if type(backing) is np.ndarray: + return True + + try: + from pandas.core.arrays import numpy_ as numpy_module, string_ as string_module + from pandas.core.arrays import string_arrow as string_arrow_module + from pandas.core.arrays.arrow import array as arrow_module + except ImportError: + # Older pandas releases may lack these concrete string storage types. + return False + + python_strings = tuple(vars(string_module).get(name) for name in + ("StringArray", "StringArrayNumpySemantics")) + arrow_strings = tuple(vars(string_arrow_module).get(name) for name in + ("ArrowStringArray", "ArrowStringArrayNumpySemantics")) + cls = type(backing) + if not any(cls is candidate for candidate in python_strings + arrow_strings): + return False + if (cls.__getattribute__ is not object.__getattribute__ + or any(name in vars(backing) for name in + ("tolist", "to_numpy", "isna", "fillna", "_hasna")) + or getattr(cls, "tolist", None) is not getattr( + vars(string_module).get("BaseStringArray"), "tolist", None) + or not ordinary(getattr(cls, "tolist", None), "tolist", string_module)): + return False + if any(cls is candidate for candidate in python_strings): + owner, module = vars(numpy_module).get("NumpyExtensionArray"), numpy_module + string_array = vars(string_module).get("StringArray") + if owner is None or string_array is None: + return False + try: + storage = object.__getattribute__(backing, "_ndarray") + except AttributeError: + return False + if type(storage) is not np.ndarray: + return False + isna = vars(string_array).get("isna", getattr(owner, "isna", None)) + isna_module = string_module if "isna" in vars(string_array) else module + else: + import pyarrow as pa + + owner, module = vars(arrow_module).get("ArrowExtensionArray"), arrow_module + if owner is None: + return False + try: + storage = object.__getattribute__(backing, "_pa_array") + except AttributeError: + return False + if type(storage) is not pa.ChunkedArray: + return False + isna, isna_module = getattr(owner, "isna", None), module + hasna = getattr(cls, "_hasna", None) + if (type(hasna) is not property + or not ordinary(hasna.fget, "_hasna", module) + or getattr(cls, "fillna", None) is not getattr(owner, "fillna", None) + or not ordinary(getattr(cls, "fillna", None), "fillna", module)): + return False + # BaseStringArray.tolist delegates to to_numpy, which consults isna (or + # Arrow's _hasna). Authenticate those direct hooks as well as the storage. + return (getattr(cls, "to_numpy", None) is getattr(owner, "to_numpy", None) + and ordinary(getattr(cls, "to_numpy", None), "to_numpy", module) + and getattr(cls, "isna", None) is isna + and ordinary(getattr(cls, "isna", None), "isna", isna_module)) + + def prepare_int64_table(frame, schema): """Return an Arrow table only when each column is already a signed 64-bit integer. @@ -52,6 +147,8 @@ def prepare_int64_table(frame, schema): if pd.DataFrame is not DataFrame or type(frame) is not DataFrame: return None + if not _has_canonical_columns(frame.columns): + return None width = len(frame.columns) if not width or not frame.columns.is_unique or not all( type(name) is str and "\x00" not in name for name in frame.columns diff --git a/python/tests/test_arrow_handoff.py b/python/tests/test_arrow_handoff.py index 82a7b79a..eb2127f3 100644 --- a/python/tests/test_arrow_handoff.py +++ b/python/tests/test_arrow_handoff.py @@ -905,6 +905,181 @@ def import_module(name, *args, **kwargs): # Signed-int64 pandas input uses the same Arrow importer and mock service. +@pytest.mark.parametrize("storage", ["object", "default", "python", "pyarrow", "python_nan", "pyarrow_nan"]) +@pytest.mark.parametrize("dtype", ["int64", "Int64"]) +def test_integer_canonical_columns_storage_stays_columnar(protocol, monkeypatch, storage, dtype): + from pyspark.sql._pandas_input import _has_canonical_columns + + values = [2**53 + 1, -(2**63), 2**63 - 1, 0] * 250 + if dtype == "Int64": + values[1] = None + frame = pd.DataFrame({"original": pd.Series(values, dtype=dtype)}) + if storage == "object": + frame.columns = pd.Index(["original"], dtype=object) + elif storage != "default": + kind = storage.removesuffix("_nan") + label_dtype = (pd.StringDtype(storage=kind, na_value=np.nan) + if storage.endswith("_nan") else pd.StringDtype(storage=kind)) + frame.columns = pd.Index(pd.array(["original"], dtype=label_dtype)) + assert _has_canonical_columns(frame.columns) + assert prepare_int64_table(frame, None).column(0).to_pylist() == values + + def no_rows(*args, **kwargs): + raise AssertionError("ordinary columns must retain integer columnar construction") + + monkeypatch.setattr(pd.DataFrame, "astype", no_rows) + with _arrow_input_service(protocol) as (spark, state): + spark.createDataFrame(frame).toArrow() + assert state["input_tables"][0].column(0).to_pylist() == values + assert state["config_requests"] + + +@pytest.mark.parametrize("dtype", ["int64", "Int64"]) +@pytest.mark.parametrize("schema_kind", ["inferred", "names", "explicit"]) +@pytest.mark.parametrize("location", [ + "subclass", "instance", "class", "owner", "backing_subclass", "backing_instance", + "backing_to_numpy", +]) +@pytest.mark.parametrize("behavior", ["raise", "mutate"]) +def test_integer_columns_tolist_hooks_keep_native_order( + protocol, monkeypatch, dtype, schema_kind, location, behavior +): + from types import MethodType + from pandas.core.base import IndexOpsMixin + + values = [2**53 + 1, -(2**63), 2**63 - 1, 0] * 250 + if dtype == "Int64": + values[1] = None + frame = pd.DataFrame({"original": pd.Series(values, dtype=dtype)}) + schema = None + expected_name = "original" + nullable = True + metadata = {} + if schema_kind == "names": + schema = ["renamed"] + expected_name = "renamed" + elif schema_kind == "explicit": + nullable, metadata = dtype == "Int64", {"source": "native"} + schema = T.StructType([T.StructField("renamed", T.LongType(), nullable, metadata)]) + expected_name = "renamed" + original_tolist = pd.Index.tolist + original_notna, original_astype = frame.notna, frame.astype + hook_name = "to_numpy" if location == "backing_to_numpy" else "tolist" + calls = [] + + class ColumnsError(RuntimeError): + pass + + def notna(*args, **kwargs): + calls.append("notna") + return original_notna(*args, **kwargs) + + def astype(*args, **kwargs): + calls.append("astype") + return original_astype(*args, **kwargs) + + def before_tolist(): + calls.append(hook_name) + if behavior == "raise": + raise ColumnsError("columns hook precedes pandas normalization") + # Install spies only when native conversion calls the columns hook. + # Before eligibility, normalization methods are the canonical methods. + monkeypatch.setitem(vars(frame), "notna", notna) + monkeypatch.setitem(vars(frame), "astype", astype) + frame.iat[0, 0] = 2**53 + 7 + frame.iat[-1, 0] = None if dtype == "Int64" else -(2**63) + 1 + + def tolist(index): + before_tolist() + return original_tolist(index) + + if location == "subclass": + class Columns(pd.Index): + pass + monkeypatch.setattr(Columns, "tolist", tolist) + labels = pd.Index(["original"], dtype=object) + frame.columns = Columns._simple_new(labels._values, name=None) + assert type(frame.columns) is Columns + elif location == "instance": + monkeypatch.setitem(vars(frame.columns), "tolist", MethodType(tolist, frame.columns)) + assert type(frame.columns) is pd.Index + elif location == "backing_subclass": + class Names(np.ndarray): + def tolist(self): + before_tolist() + return super().tolist() + backing = np.asarray(["original"], dtype=object).view(Names) + frame.columns = pd.Index._simple_new(backing, name=None) + assert type(frame.columns) is pd.Index + assert type(frame.columns._values) is Names + elif location in ("backing_instance", "backing_to_numpy"): + frame.columns = pd.Index(pd.array(["original"], dtype="string")) + backing = frame.columns._values + original = getattr(backing, hook_name) + + def backing_tolist(*args, **kwargs): + before_tolist() + return original(*args, **kwargs) + + monkeypatch.setitem(vars(backing), hook_name, backing_tolist) + assert type(frame.columns) is pd.Index + else: + monkeypatch.setattr(pd.Index if location == "class" else IndexOpsMixin, + "tolist", tolist) + + with _arrow_input_service(protocol) as (spark, state): + if behavior == "raise": + with pytest.raises(ColumnsError, + match="columns hook precedes pandas normalization"): + spark.createDataFrame(frame, schema) + assert calls == [hook_name] + assert state["config_requests"] == [] + assert state["input_tables"] == [] + else: + spark.createDataFrame(frame, schema).toArrow() + expected = values.copy() + expected[0], expected[-1] = 2**53 + 7, (None if dtype == "Int64" else -(2**63) + 1) + table = state["input_tables"][0] + assert table.num_rows == 1000 + assert table.column_names == [expected_name] + assert table.column(0).type == pa.int64() + assert table.column(0).to_pylist() == expected + assert table.column(0).null_count == sum(value is None for value in expected) + assert table.schema.field(0).nullable is nullable + assert table.schema.field(0).metadata is None + import json + logical = json.loads(state["relations"][0].local_relation.schema) + assert logical == {"type": "struct", "fields": [{ + "name": expected_name, "type": "long", "nullable": nullable, + "metadata": metadata, + }]} + assert state["config_requests"] == [] + assert calls == [hook_name, "notna", "astype"] + + +def test_integer_rebound_index_values_keeps_native_error(protocol, monkeypatch): + frame = pd.DataFrame({"original": np.arange(1000, dtype=np.int64)}) + schema = T.StructType([T.StructField("renamed", T.LongType())]) + columns = frame.columns + original = pd.Index._values.fget + calls = [] + + def values(index): + if index is columns: + calls.append("_values") + raise RuntimeError("columns backing hook precedes normalization") + return original(index) + + monkeypatch.setattr(pd.Index, "_values", property(values)) + with _arrow_input_service(protocol) as (spark, state): + with pytest.raises(RuntimeError, + match="columns backing hook precedes normalization"): + spark.createDataFrame(frame, schema) + assert calls == ["_values"] + assert state["config_requests"] == [] + assert state["input_tables"] == [] + + @pytest.mark.parametrize("schema_kind", ["inferred", "names", "explicit"]) @pytest.mark.parametrize("verify", [True, False]) def test_integer_columns_preserve_values_without_python_rows(protocol, monkeypatch, schema_kind, verify): From 7961bdeb5098f491d1ee142b517ed5e33a3cb005 Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Fri, 30 Oct 2026 00:14:30 -0700 Subject: [PATCH 11/23] [SPARK-60676][CONNECT] Preserve rebound pandas index storage ordering Decline class-level columns storage descriptors before reading labels. Keep the single native columns storage read before pandas normalization, and cover public construction with inferred, named and explicit schemas on both pandas versions. Co-authored-by: Isaac --- python/pyspark/sql/_pandas_input.py | 2 + python/tests/test_arrow_handoff.py | 97 +++++++++++++++++++++++++++++ 2 files changed, 99 insertions(+) diff --git a/python/pyspark/sql/_pandas_input.py b/python/pyspark/sql/_pandas_input.py index 3a7f0e3c..ae4126af 100644 --- a/python/pyspark/sql/_pandas_input.py +++ b/python/pyspark/sql/_pandas_input.py @@ -48,6 +48,8 @@ def ordinary(method, name, module): if (type(values) is not property or not ordinary(values.fget, "_values", index_module)): return False + if any("_data" in vars(base) for base in Index.__mro__): + return False backing = vars(columns).get("_data") if type(backing) is np.ndarray: return True diff --git a/python/tests/test_arrow_handoff.py b/python/tests/test_arrow_handoff.py index eb2127f3..fa22952f 100644 --- a/python/tests/test_arrow_handoff.py +++ b/python/tests/test_arrow_handoff.py @@ -1578,3 +1578,100 @@ def no_rows(*args, **kwargs): for table in state["input_tables"][1:]: assert table.column_names == ["renamed"] assert table.column(0).to_pylist() == list(range(1000)) + + +@pytest.mark.parametrize("in_ci", [False, True]) +def test_missing_protoc_fails_in_ci_and_skips_locally(monkeypatch, tmp_path_factory, in_ci): + monkeypatch.setattr(shutil, "which", lambda _: None) + if in_ci: + monkeypatch.setenv("CI", "true") + else: + monkeypatch.delenv("CI", raising=False) + error = pytest.fail.Exception if in_ci else pytest.skip.Exception + with pytest.raises(error, match="require protoc"): + protocol.__wrapped__(tmp_path_factory) + + + +def test_normal_handoff_does_not_reopen_ipc(protocol, monkeypatch): + def unexpected_ipc(*args, **kwargs): + pytest.fail("Normal toArrow handoff must use the C Stream importer") + + monkeypatch.setattr(pa.ipc, "open_file", unexpected_ipc) + batch = _mixed_batch() + with _dataframe(protocol=protocol, batches=[batch]) as (frame, _): + assert frame.toArrow().equals(pa.Table.from_batches([batch]), check_metadata=True) + + +@pytest.mark.parametrize("dtype", ["int64", "Int64"]) +@pytest.mark.parametrize("schema_kind", ["inferred", "names", "explicit"]) +def test_integer_rebound_index_data_keeps_native_read_order( + protocol, monkeypatch, dtype, schema_kind +): + import json + + values = [2**53 + 1, -(2**63), 2**63 - 1, 0] * 250 + if dtype == "Int64": + values[1] = None + frame = pd.DataFrame({"original": pd.Series(values, dtype=dtype)}) + schema, expected_name, nullable, metadata = None, "original", True, {} + if schema_kind == "names": + schema, expected_name = ["renamed"], "renamed" + elif schema_kind == "explicit": + expected_name, nullable, metadata = "renamed", dtype == "Int64", {"source": "native"} + schema = T.StructType([T.StructField(expected_name, T.LongType(), nullable, metadata)]) + columns, backing = frame.columns, vars(frame.columns)["_data"] + original_notna, original_astype = frame.notna, frame.astype + calls = [] + normalizing = False + + class ColumnsDataError(RuntimeError): + pass + + def notna(*args, **kwargs): + nonlocal normalizing + normalizing = True + calls.append("notna") + return original_notna(*args, **kwargs) + + def astype(*args, **kwargs): + calls.append("astype") + return original_astype(*args, **kwargs) + + def data(index): + if index is columns: + if not normalizing: + calls.append("_data") + if calls.count("_data") > 1: + raise ColumnsDataError("columns storage read twice before normalization") + # Install spies only when the native columns storage hook runs. + # Canonical frame methods remain intact before eligibility. + monkeypatch.setitem(vars(frame), "notna", notna) + monkeypatch.setitem(vars(frame), "astype", astype) + return backing + return vars(index)["_data"] + + def set_data(index, value): + # Preserve ordinary storage initialization for indexes pandas creates + # while normalizing values; only reads of the input columns are hooked. + vars(index)["_data"] = value + + monkeypatch.setattr(pd.Index, "_data", property(data, set_data), raising=False) + with _arrow_input_service(protocol) as (spark, state): + spark.createDataFrame(frame, schema).toArrow() + assert calls == ["_data", "notna", "astype"] + assert not state["config_requests"] + assert len(state["input_tables"]) == 1 + table = state["input_tables"][0] + assert table.column(0).to_pylist() == values + assert table.column(0).null_count == sum(value is None for value in values) + assert table.schema.equals( + pa.schema([pa.field(expected_name, pa.int64(), nullable=nullable)]), + check_metadata=True, + ) + assert json.loads(state["relations"][0].local_relation.schema) == { + "type": "struct", "fields": [ + {"name": expected_name, "type": "long", "nullable": nullable, + "metadata": metadata} + ], + } From 3e3a8f5fe7507268aa6834bee74e5a25cd14feb0 Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Fri, 30 Oct 2026 20:12:25 -0700 Subject: [PATCH 12/23] Preserve integer pandas fallback with missing Arrow aliases Avoid loading optional Arrow label modules for fallback inputs, and decline integer Arrow conversion when the loaded importer AtomicType alias is absent. Co-authored-by: Isaac --- python/pyspark/sql/_pandas_input.py | 19 +++++++++++++++---- python/tests/test_arrow_handoff.py | 26 ++++++++++++++++++++++++++ 2 files changed, 41 insertions(+), 4 deletions(-) diff --git a/python/pyspark/sql/_pandas_input.py b/python/pyspark/sql/_pandas_input.py index ae4126af..13c1cd08 100644 --- a/python/pyspark/sql/_pandas_input.py +++ b/python/pyspark/sql/_pandas_input.py @@ -21,6 +21,7 @@ def _has_canonical_columns(columns): """Decline label containers whose row-path ``tolist`` can run custom hooks.""" + import sys from types import FunctionType import numpy as np @@ -54,6 +55,12 @@ def ordinary(method, name, module): if type(backing) is np.ndarray: return True + # Arrow-backed labels already hold PyArrow storage. Do not import optional + # Arrow modules just to decide whether a fallback input is eligible. + pa = sys.modules.get("pyarrow") + if pa is None: + return False + try: from pandas.core.arrays import numpy_ as numpy_module, string_ as string_module from pandas.core.arrays import string_arrow as string_arrow_module @@ -90,8 +97,6 @@ def ordinary(method, name, module): isna = vars(string_array).get("isna", getattr(owner, "isna", None)) isna_module = string_module if "isna" in vars(string_array) else module else: - import pyarrow as pa - owner, module = vars(arrow_module).get("ArrowExtensionArray"), arrow_module if owner is None: return False @@ -140,7 +145,7 @@ def prepare_int64_table(frame, schema): return None from pyspark._pyspark import ( - BooleanType, ByteType, ShortType, IntegerType, LongType, StructField, StructType, + AtomicType, BooleanType, ByteType, ShortType, IntegerType, LongType, StructField, StructType, ) from pyspark.sql import types as sql_types @@ -173,8 +178,14 @@ def prepare_int64_table(frame, schema): if sys.modules.get("pyspark.sql.types") is not sql_types: return None importer = sys.modules.get("pyspark.sql._arrow_input") - if importer is not None and getattr(importer, "StructType", None) is not StructType: + if vars(sql_types).get("AtomicType") is not AtomicType: return None + if importer is not None: + if getattr(importer, "StructType", None) is not StructType: + return None + # The importer uses this alias while validating schema widths. + if vars(importer).get("AtomicType") is not AtomicType: + return None converter = sys.modules.get("pyspark.sql.pandas.types") namespaces = [vars(sql_types)] if converter is not None: diff --git a/python/tests/test_arrow_handoff.py b/python/tests/test_arrow_handoff.py index fa22952f..861eb4d2 100644 --- a/python/tests/test_arrow_handoff.py +++ b/python/tests/test_arrow_handoff.py @@ -1497,6 +1497,21 @@ def _integer_guard_process(mode, argument, before_import=False): current = importlib.import_module("pyspark.sql.pandas.types") assert current is not converter assert importer.to_arrow_schema.__globals__ is vars(converter) +elif mode == "atomic_alias": + scope, replacement_name = argument.split(":") + replacement = None if replacement_name == "none" else native.LongType + if scope == "public_before": + sys.modules.pop("pyspark.sql._arrow_input", None) + T.AtomicType = replacement + elif scope == "public_after": + importlib.import_module("pyspark.sql._arrow_input") + T.AtomicType = replacement + else: + importer = importlib.import_module("pyspark.sql._arrow_input") + if replacement_name == "missing": + del importer.AtomicType + else: + importer.AtomicType = replacement elif mode == "importer_alias": importer = importlib.import_module("pyspark.sql._arrow_input") importer.StructType = None if argument == "none" else native.LongType @@ -1543,6 +1558,16 @@ def test_integer_rebound_importer_schema_alias_keeps_native_row_path(replacement _integer_guard_process("importer_alias", replacement) +@pytest.mark.parametrize("scope", ["public_before", "public_after", "importer"]) +@pytest.mark.parametrize("replacement", ["none", "long"]) +def test_integer_rebound_atomic_alias_keeps_native_row_path(scope, replacement): + _integer_guard_process("atomic_alias", f"{scope}:{replacement}") + + +def test_integer_deleted_importer_atomic_alias_keeps_native_row_path(): + _integer_guard_process("atomic_alias", "importer:missing") + + def test_integer_warm_importer_without_global_converter_aliases(protocol, monkeypatch): import pyspark.sql._arrow_input as importer @@ -1566,6 +1591,7 @@ def local_imports(table, schema, *configs): import types substitute = types.ModuleType(importer.__name__) substitute.StructType = T.StructType + substitute.AtomicType = T.AtomicType substitute.prepare_table = local_imports monkeypatch.setitem(__import__("sys").modules, importer.__name__, substitute) From f058ac98e81ee1783c156cb6ebf6abcaefe10da7 Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Thu, 5 Nov 2026 14:18:44 -0800 Subject: [PATCH 13/23] [SPARK-60676][TESTS] Retain current master coverage after rebase Co-authored-by: Isaac --- python/tests/test_arrow_handoff.py | 18 ------------------ 1 file changed, 18 deletions(-) diff --git a/python/tests/test_arrow_handoff.py b/python/tests/test_arrow_handoff.py index 861eb4d2..b6e1a80b 100644 --- a/python/tests/test_arrow_handoff.py +++ b/python/tests/test_arrow_handoff.py @@ -1606,27 +1606,9 @@ def no_rows(*args, **kwargs): assert table.column(0).to_pylist() == list(range(1000)) -@pytest.mark.parametrize("in_ci", [False, True]) -def test_missing_protoc_fails_in_ci_and_skips_locally(monkeypatch, tmp_path_factory, in_ci): - monkeypatch.setattr(shutil, "which", lambda _: None) - if in_ci: - monkeypatch.setenv("CI", "true") - else: - monkeypatch.delenv("CI", raising=False) - error = pytest.fail.Exception if in_ci else pytest.skip.Exception - with pytest.raises(error, match="require protoc"): - protocol.__wrapped__(tmp_path_factory) - -def test_normal_handoff_does_not_reopen_ipc(protocol, monkeypatch): - def unexpected_ipc(*args, **kwargs): - pytest.fail("Normal toArrow handoff must use the C Stream importer") - monkeypatch.setattr(pa.ipc, "open_file", unexpected_ipc) - batch = _mixed_batch() - with _dataframe(protocol=protocol, batches=[batch]) as (frame, _): - assert frame.toArrow().equals(pa.Table.from_batches([batch]), check_metadata=True) @pytest.mark.parametrize("dtype", ["int64", "Int64"]) From 121a1413e3a4b847c218fb92285241a3b8f3c1d6 Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Thu, 5 Nov 2026 14:50:04 -0800 Subject: [PATCH 14/23] Preserve current pandas label iteration and fallback order after rebase Co-authored-by: Isaac --- python/pyspark/sql/_pandas_input.py | 32 +++++--- python/tests/test_arrow_handoff.py | 118 +++++++++++++++++++--------- 2 files changed, 102 insertions(+), 48 deletions(-) diff --git a/python/pyspark/sql/_pandas_input.py b/python/pyspark/sql/_pandas_input.py index 13c1cd08..5df79db7 100644 --- a/python/pyspark/sql/_pandas_input.py +++ b/python/pyspark/sql/_pandas_input.py @@ -20,7 +20,7 @@ def _has_canonical_columns(columns): - """Decline label containers whose row-path ``tolist`` can run custom hooks.""" + """Decline label containers whose row-path iteration can run custom hooks.""" import sys from types import FunctionType @@ -32,7 +32,7 @@ def _has_canonical_columns(columns): Index = index_module.Index if (pd.Index is not Index or type(columns) is not Index or Index.__getattribute__ is not object.__getattribute__ - or "tolist" in vars(columns)): + or "__iter__" in vars(columns)): return False def ordinary(method, name, module): @@ -40,10 +40,10 @@ def ordinary(method, name, module): and method.__module__ == module.__name__ and method.__globals__ is vars(module)) - # The native row path calls columns.tolist() before normalizing any cells. + # The native row path iterates columns before normalizing any cells. # Decline overrides without invoking them, so fallback preserves that order. - if (Index.tolist is not base_module.IndexOpsMixin.tolist - or not ordinary(Index.tolist, "tolist", base_module)): + if (Index.__iter__ is not base_module.IndexOpsMixin.__iter__ + or not ordinary(Index.__iter__, "__iter__", base_module)): return False values = vars(Index).get("_values") if (type(values) is not property @@ -63,6 +63,8 @@ def ordinary(method, name, module): try: from pandas.core.arrays import numpy_ as numpy_module, string_ as string_module + from pandas.core.arrays import base as array_base_module, _mixins as mixins_module + from pandas._libs.arrays import NDArrayBacked from pandas.core.arrays import string_arrow as string_arrow_module from pandas.core.arrays.arrow import array as arrow_module except ImportError: @@ -78,7 +80,8 @@ def ordinary(method, name, module): return False if (cls.__getattribute__ is not object.__getattribute__ or any(name in vars(backing) for name in - ("tolist", "to_numpy", "isna", "fillna", "_hasna")) + ("__iter__", "__getitem__", "__len__", "_box_func", + "tolist", "to_numpy", "isna", "fillna", "_hasna")) or getattr(cls, "tolist", None) is not getattr( vars(string_module).get("BaseStringArray"), "tolist", None) or not ordinary(getattr(cls, "tolist", None), "tolist", string_module)): @@ -94,6 +97,16 @@ def ordinary(method, name, module): return False if type(storage) is not np.ndarray: return False + # ExtensionArray iteration dispatches through len and indexed boxing. + for name, module in (("__iter__", array_base_module), + ("__getitem__", mixins_module), ("_box_func", mixins_module)): + if not ordinary(getattr(cls, name, None), name, module): + return False + if (getattr(cls, "__len__", None) is not NDArrayBacked.__len__ + or getattr(cls, "ndim", None) is not NDArrayBacked.ndim + or any("_ndarray" in vars(base) for base in cls.__mro__ + if base is not NDArrayBacked)): + return False isna = vars(string_array).get("isna", getattr(owner, "isna", None)) isna_module = string_module if "isna" in vars(string_array) else module else: @@ -104,7 +117,8 @@ def ordinary(method, name, module): storage = object.__getattribute__(backing, "_pa_array") except AttributeError: return False - if type(storage) is not pa.ChunkedArray: + if (type(storage) is not pa.ChunkedArray + or not ordinary(getattr(cls, "__iter__", None), "__iter__", module)): return False isna, isna_module = getattr(owner, "isna", None), module hasna = getattr(cls, "_hasna", None) @@ -113,8 +127,8 @@ def ordinary(method, name, module): or getattr(cls, "fillna", None) is not getattr(owner, "fillna", None) or not ordinary(getattr(cls, "fillna", None), "fillna", module)): return False - # BaseStringArray.tolist delegates to to_numpy, which consults isna (or - # Arrow's _hasna). Authenticate those direct hooks as well as the storage. + # Keep the existing storage/conversion authentication as well as the + # iterator hooks: these objects also supply labels to pandas normalization. return (getattr(cls, "to_numpy", None) is getattr(owner, "to_numpy", None) and ordinary(getattr(cls, "to_numpy", None), "to_numpy", module) and getattr(cls, "isna", None) is isna diff --git a/python/tests/test_arrow_handoff.py b/python/tests/test_arrow_handoff.py index b6e1a80b..93d17f64 100644 --- a/python/tests/test_arrow_handoff.py +++ b/python/tests/test_arrow_handoff.py @@ -937,14 +937,13 @@ def no_rows(*args, **kwargs): @pytest.mark.parametrize("dtype", ["int64", "Int64"]) @pytest.mark.parametrize("schema_kind", ["inferred", "names", "explicit"]) @pytest.mark.parametrize("location", [ - "subclass", "instance", "class", "owner", "backing_subclass", "backing_instance", - "backing_to_numpy", + "subclass", "class", "owner", "backing_subclass_item", "backing_class_iter", + "backing_class_getitem", "backing_instance_box", ]) @pytest.mark.parametrize("behavior", ["raise", "mutate"]) -def test_integer_columns_tolist_hooks_keep_native_order( +def test_integer_columns_iterator_hooks_keep_native_order( protocol, monkeypatch, dtype, schema_kind, location, behavior ): - from types import MethodType from pandas.core.base import IndexOpsMixin values = [2**53 + 1, -(2**63), 2**63 - 1, 0] * 250 @@ -962,10 +961,10 @@ def test_integer_columns_tolist_hooks_keep_native_order( nullable, metadata = dtype == "Int64", {"source": "native"} schema = T.StructType([T.StructField("renamed", T.LongType(), nullable, metadata)]) expected_name = "renamed" - original_tolist = pd.Index.tolist + original_iter = pd.Index.__iter__ original_notna, original_astype = frame.notna, frame.astype - hook_name = "to_numpy" if location == "backing_to_numpy" else "tolist" calls = [] + invoked = False class ColumnsError(RuntimeError): pass @@ -978,61 +977,82 @@ def astype(*args, **kwargs): calls.append("astype") return original_astype(*args, **kwargs) - def before_tolist(): - calls.append(hook_name) + def before_hook(): + nonlocal invoked + if invoked: + return + invoked = True + calls.append("columns iterator") if behavior == "raise": raise ColumnsError("columns hook precedes pandas normalization") - # Install spies only when native conversion calls the columns hook. + # Install spies only when native iteration reaches the columns hook. # Before eligibility, normalization methods are the canonical methods. monkeypatch.setitem(vars(frame), "notna", notna) monkeypatch.setitem(vars(frame), "astype", astype) frame.iat[0, 0] = 2**53 + 7 frame.iat[-1, 0] = None if dtype == "Int64" else -(2**63) + 1 - def tolist(index): - before_tolist() - return original_tolist(index) + def iterate(index): + if index is columns: + before_hook() + return original_iter(index) if location == "subclass": class Columns(pd.Index): pass - monkeypatch.setattr(Columns, "tolist", tolist) + monkeypatch.setattr(Columns, "__iter__", iterate) labels = pd.Index(["original"], dtype=object) frame.columns = Columns._simple_new(labels._values, name=None) assert type(frame.columns) is Columns - elif location == "instance": - monkeypatch.setitem(vars(frame.columns), "tolist", MethodType(tolist, frame.columns)) - assert type(frame.columns) is pd.Index - elif location == "backing_subclass": + elif location == "backing_subclass_item": + # IndexOpsMixin iteration calls ndarray.item, rather than ndarray.__iter__. class Names(np.ndarray): - def tolist(self): - before_tolist() - return super().tolist() + def item(self, *args): + before_hook() + return super().item(*args) backing = np.asarray(["original"], dtype=object).view(Names) frame.columns = pd.Index._simple_new(backing, name=None) - assert type(frame.columns) is pd.Index assert type(frame.columns._values) is Names - elif location in ("backing_instance", "backing_to_numpy"): - frame.columns = pd.Index(pd.array(["original"], dtype="string")) + elif location.startswith("backing_"): + frame.columns = pd.Index(pd.array(["original"], dtype=pd.StringDtype(storage="python"))) backing = frame.columns._values - original = getattr(backing, hook_name) + if location == "backing_class_iter": + original = type(backing).__iter__ + + def backing_iter(array): + if array is backing: + before_hook() + return original(array) + + monkeypatch.setattr(type(backing), "__iter__", backing_iter) + elif location == "backing_class_getitem": + original = type(backing).__getitem__ + + def backing_getitem(array, key): + if array is backing: + before_hook() + return original(array, key) + + monkeypatch.setattr(type(backing), "__getitem__", backing_getitem) + else: + original = backing._box_func - def backing_tolist(*args, **kwargs): - before_tolist() - return original(*args, **kwargs) + def box(value): + before_hook() + return original(value) - monkeypatch.setitem(vars(backing), hook_name, backing_tolist) - assert type(frame.columns) is pd.Index + monkeypatch.setitem(vars(backing), "_box_func", box) else: monkeypatch.setattr(pd.Index if location == "class" else IndexOpsMixin, - "tolist", tolist) + "__iter__", iterate) + columns = frame.columns with _arrow_input_service(protocol) as (spark, state): if behavior == "raise": with pytest.raises(ColumnsError, match="columns hook precedes pandas normalization"): spark.createDataFrame(frame, schema) - assert calls == [hook_name] + assert calls == ["columns iterator"] assert state["config_requests"] == [] assert state["input_tables"] == [] else: @@ -1054,7 +1074,7 @@ def backing_tolist(*args, **kwargs): "metadata": metadata, }]} assert state["config_requests"] == [] - assert calls == [hook_name, "notna", "astype"] + assert calls == ["columns iterator", "notna", "astype"] def test_integer_rebound_index_values_keeps_native_error(protocol, monkeypatch): @@ -1110,8 +1130,8 @@ def no_rows(*args, **kwargs): monkeypatch.setattr(pd.DataFrame, "astype", no_rows) with _arrow_input_service(protocol) as (spark, state): df = spark.createDataFrame(frame, schema=schema, verifySchema=verify) - if schema_kind != "names": - assert df.isLocal() + # Inference and name lists produce a toDF wrapper on current master. + assert df.isLocal() == (schema_kind == "explicit") assert df.toArrow().column(0).to_pylist() == [3] assert state["result"] == (3, 2**53 + 1, -(2**63)) table, resolved, renamed = observed[0] @@ -1632,6 +1652,7 @@ def test_integer_rebound_index_data_keeps_native_read_order( original_notna, original_astype = frame.notna, frame.astype calls = [] normalizing = False + allowed_reads = None class ColumnsDataError(RuntimeError): pass @@ -1650,8 +1671,8 @@ def data(index): if index is columns: if not normalizing: calls.append("_data") - if calls.count("_data") > 1: - raise ColumnsDataError("columns storage read twice before normalization") + if allowed_reads is not None and calls.count("_data") > allowed_reads: + raise ColumnsDataError("eligibility added a columns storage read") # Install spies only when the native columns storage hook runs. # Canonical frame methods remain intact before eligibility. monkeypatch.setitem(vars(frame), "notna", notna) @@ -1665,12 +1686,31 @@ def set_data(index, value): vars(index)["_data"] = value monkeypatch.setattr(pd.Index, "_data", property(data, set_data), raising=False) + import pyspark.sql._pandas_input as helper + with _arrow_input_service(protocol) as (spark, state): + # Record current master's fallback rather than assuming the old tolist + # path's single read. Index iteration reads NumPy-backed _values three + # times, and extension-backed _values twice, before normalization. + with monkeypatch.context() as baseline: + baseline.setattr(helper, "prepare_int64_table", lambda *_: None) + spark.createDataFrame(frame, schema).toArrow() + baseline_calls = calls.copy() + allowed_reads = baseline_calls.index("notna") + assert allowed_reads > 0 + assert baseline_calls == ["_data"] * allowed_reads + ["notna", "astype"] + calls.clear() + normalizing = False + # Restore canonical methods before eligibility, just as for the first + # call. The storage hook installs them when native iteration starts. + monkeypatch.delitem(vars(frame), "notna") + monkeypatch.delitem(vars(frame), "astype") spark.createDataFrame(frame, schema).toArrow() - assert calls == ["_data", "notna", "astype"] + assert calls == baseline_calls assert not state["config_requests"] - assert len(state["input_tables"]) == 1 - table = state["input_tables"][0] + assert len(state["input_tables"]) == 2 + table = state["input_tables"][1] + assert table.equals(state["input_tables"][0], check_metadata=True) assert table.column(0).to_pylist() == values assert table.column(0).null_count == sum(value is None for value in values) assert table.schema.equals( From aca1d026c3f25bdc33cb5d29757312a2b15f2f30 Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Thu, 5 Nov 2026 14:54:06 -0800 Subject: [PATCH 15/23] Keep canonical string storage guard module identity Co-authored-by: Isaac --- python/pyspark/sql/_pandas_input.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/python/pyspark/sql/_pandas_input.py b/python/pyspark/sql/_pandas_input.py index 5df79db7..614287d8 100644 --- a/python/pyspark/sql/_pandas_input.py +++ b/python/pyspark/sql/_pandas_input.py @@ -98,9 +98,9 @@ def ordinary(method, name, module): if type(storage) is not np.ndarray: return False # ExtensionArray iteration dispatches through len and indexed boxing. - for name, module in (("__iter__", array_base_module), - ("__getitem__", mixins_module), ("_box_func", mixins_module)): - if not ordinary(getattr(cls, name, None), name, module): + for name, hook_module in (("__iter__", array_base_module), + ("__getitem__", mixins_module), ("_box_func", mixins_module)): + if not ordinary(getattr(cls, name, None), name, hook_module): return False if (getattr(cls, "__len__", None) is not NDArrayBacked.__len__ or getattr(cls, "ndim", None) is not NDArrayBacked.ndim From ecdd67a019516320c583f9dfbfd7c3dce2c491c5 Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Thu, 5 Nov 2026 16:01:19 -0800 Subject: [PATCH 16/23] [CONNECT] Reconcile remaining change with current inference architecture Co-authored-by: Isaac --- crates/pyspark-rs/src/session.rs | 21 +++++++++++++++++++++ python/tests/test_arrow_handoff.py | 6 ++++-- 2 files changed, 25 insertions(+), 2 deletions(-) diff --git a/crates/pyspark-rs/src/session.rs b/crates/pyspark-rs/src/session.rs index 4b9ce5c5..e17792c6 100644 --- a/crates/pyspark-rs/src/session.rs +++ b/crates/pyspark-rs/src/session.rs @@ -765,6 +765,27 @@ impl PySparkSession { } let spec = schema_spec(py, schema)?; if is_loaded_instance(py, data, "pandas", "DataFrame")? { + // Guarded integer columns share the existing Arrow importer. Schema + // validation above and unsupported pandas conversions keep master's + // normal ordering and row fallback. + let table = py + .import("pyspark.sql._pandas_input")? + .getattr("prepare_int64_table")? + .call1((data, schema))?; + if !table.is_none() { + let frame = + crate::pyarrow_import::create_dataframe(py, &self.session, &table, schema)?; + // pandas applies explicit field names with toDF even though the + // local relation already carries them. Empty input stays local. + return Ok(match &spec { + Spec::Struct(fields) if !fields.is_empty() && data.len()? > 0 => { + PyDataFrame::new(frame.dataframe.to_df( + fields.iter().map(|field| field.name.as_str()).collect(), + )) + } + _ => frame, + }); + } return self.create_from_pandas(py, data, spec); } self.create_from_local(py, data, spec) diff --git a/python/tests/test_arrow_handoff.py b/python/tests/test_arrow_handoff.py index 93d17f64..5d098808 100644 --- a/python/tests/test_arrow_handoff.py +++ b/python/tests/test_arrow_handoff.py @@ -1130,9 +1130,11 @@ def no_rows(*args, **kwargs): monkeypatch.setattr(pd.DataFrame, "astype", no_rows) with _arrow_input_service(protocol) as (spark, state): df = spark.createDataFrame(frame, schema=schema, verifySchema=verify) - # Inference and name lists produce a toDF wrapper on current master. - assert df.isLocal() == (schema_kind == "explicit") + # Every nonempty pandas input applies its names through toDF, including + # explicit struct schemas, just as the current row fallback does. + assert not df.isLocal() assert df.toArrow().column(0).to_pylist() == [3] + assert state["column_names"] == names assert state["result"] == (3, 2**53 + 1, -(2**63)) table, resolved, renamed = observed[0] assert table.to_pydict() == {"a": [2**53 + 1, None, -(2**63)], From e7ae71ae03634d1fd5bc6b8187f6278081206208 Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Thu, 5 Nov 2026 16:03:09 -0800 Subject: [PATCH 17/23] [CONNECT] Reconcile remaining change with current inference architecture Co-authored-by: Isaac --- crates/pyspark-rs/src/session.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/pyspark-rs/src/session.rs b/crates/pyspark-rs/src/session.rs index e17792c6..8fc08ff0 100644 --- a/crates/pyspark-rs/src/session.rs +++ b/crates/pyspark-rs/src/session.rs @@ -778,7 +778,7 @@ impl PySparkSession { // pandas applies explicit field names with toDF even though the // local relation already carries them. Empty input stays local. return Ok(match &spec { - Spec::Struct(fields) if !fields.is_empty() && data.len()? > 0 => { + Spec::Struct(fields) if !fields.is_empty() && table.len()? > 0 => { PyDataFrame::new(frame.dataframe.to_df( fields.iter().map(|field| field.name.as_str()).collect(), )) From 40e5266ce7d437978ed54fb8467d03cd679c0ac1 Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Thu, 5 Nov 2026 16:11:02 -0800 Subject: [PATCH 18/23] [CONNECT] Reconcile remaining change with current inference architecture Co-authored-by: Isaac --- crates/pyspark-rs/src/session.rs | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/crates/pyspark-rs/src/session.rs b/crates/pyspark-rs/src/session.rs index 8fc08ff0..4edbd2a9 100644 --- a/crates/pyspark-rs/src/session.rs +++ b/crates/pyspark-rs/src/session.rs @@ -779,9 +779,11 @@ impl PySparkSession { // local relation already carries them. Empty input stays local. return Ok(match &spec { Spec::Struct(fields) if !fields.is_empty() && table.len()? > 0 => { - PyDataFrame::new(frame.dataframe.to_df( - fields.iter().map(|field| field.name.as_str()).collect(), - )) + PyDataFrame::new( + frame + .dataframe + .to_df(fields.iter().map(|field| field.name.as_str()).collect()), + ) } _ => frame, }); From 26be6abe1cceff131004795791af651422eb5d81 Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Fri, 30 Oct 2026 02:49:42 -0700 Subject: [PATCH 19/23] [SPARK-60720][CONNECT] Encode guarded pandas String input inline Co-authored-by: Isaac --- crates/pyspark-rs/src/pandas_input.rs | 113 ++++ crates/spark-connect/src/session.rs | 17 + python/pyspark/sql/_pandas_input.py | 122 ++++ python/tests/test_arrow_handoff.py | 897 ++++++++++++++++++++++++++ 4 files changed, 1149 insertions(+) create mode 100644 crates/pyspark-rs/src/pandas_input.rs diff --git a/crates/pyspark-rs/src/pandas_input.rs b/crates/pyspark-rs/src/pandas_input.rs new file mode 100644 index 00000000..ab6eb8fc --- /dev/null +++ b/crates/pyspark-rs/src/pandas_input.rs @@ -0,0 +1,113 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Fused validation and inline encoding of guarded pandas string columns. + +use std::sync::Arc; + +use arrow::array::{ArrayRef, StringBuilder}; +use arrow::datatypes::{DataType as ArrowType, Field, Schema}; +use arrow::ipc::writer::StreamWriter; +use arrow::record_batch::RecordBatch; +use pyo3::prelude::*; +use pyo3::types::{PyBytes, PyString}; + +use crate::types::PyStructType; + +/// The Python eligibility helper authenticates pandas storage and native schema. +/// Validate each object before accepting it: extraction must never call a custom +/// numeric/string conversion hook or turn a pandas missing sentinel into null. +#[pyfunction] +pub(crate) fn _pandas_string_ipc<'py>( + py: Python<'py>, + columns: Vec>, + schema: PyRef<'py, PyStructType>, +) -> PyResult>> { + if columns.len() != schema.fields.len() { + return Ok(None); + } + // Reject unsupported objects before copying any string data. Otherwise a + // late subclass would pay for most of an Arrow build before the row path. + // Keep this pass to exact type/null checks; UTF-8 extraction stays below. + for (column, field) in columns.iter().zip(&schema.fields) { + if !matches!(&field.data_type, spark_connect::types::DataType::String { collation } + if collation == "UTF8_BINARY") + { + return Ok(None); + } + for value in column.try_iter()? { + let value = value?; + if !(value.is_none() && field.nullable || value.is_exact_instance_of::()) { + return Ok(None); + } + } + } + // A lone surrogate must fall back before allocating partial Arrow buffers. + // Drop each temporary UTF-8 bytes object immediately rather than retaining + // an additional full copy of every column. Keep the cheap type pass above + // so a late subclass does not pay for this UTF-8 traversal. + for (column, field) in columns.iter().zip(&schema.fields) { + for value in column.try_iter()? { + let value = value?; + if value.is_none() && field.nullable { + continue; + } + let Ok(value) = value.cast_exact::() else { + return Ok(None); + }; + if value.encode_utf8().is_err() { + return Ok(None); + } + } + } + let mut arrays: Vec = Vec::with_capacity(columns.len()); + let mut fields = Vec::with_capacity(columns.len()); + for (column, field) in columns.iter().zip(&schema.fields) { + let mut builder = StringBuilder::with_capacity(column.len()?, column.len()? * 8); + for value in column.try_iter()? { + let value = value?; + if value.is_none() && field.nullable { + builder.append_null(); + } else if let Ok(value) = value.cast_exact::() { + let Ok(bytes) = value.encode_utf8() else { + return Ok(None); + }; + let Ok(text) = std::str::from_utf8(bytes.as_bytes()) else { + return Ok(None); + }; + // append_value copies while the bytes object and GIL are live. + // This avoids to_cow's owned String on the abi3-py39 build. + builder.append_value(text); + } else { + return Ok(None); + } + } + arrays.push(Arc::new(builder.finish())); + // Like the row encoder, Spark metadata lives in the logical schema; + // Arrow fields carry the name, standard string type and nullability. + fields.push(Field::new(&field.name, ArrowType::Utf8, field.nullable)); + } + let bytes = py + .detach(move || -> Result, arrow::error::ArrowError> { + let batch = RecordBatch::try_new(Arc::new(Schema::new(fields)), arrays)?; + let mut writer = StreamWriter::try_new(Vec::new(), &batch.schema())?; + writer.write(&batch)?; + writer.into_inner() + }) + .map_err(|error| pyo3::exceptions::PyRuntimeError::new_err(error.to_string()))?; + Ok(Some(PyBytes::new(py, &bytes))) +} diff --git a/crates/spark-connect/src/session.rs b/crates/spark-connect/src/session.rs index 96c3b1f6..10c57aac 100644 --- a/crates/spark-connect/src/session.rs +++ b/crates/spark-connect/src/session.rs @@ -360,6 +360,23 @@ impl SparkSession { Ok(DataFrame::new(self.clone(), plan)) } + /// Internal binding bridge for Arrow IPC already validated and encoded from + /// local rows or columns. Retains the row path's inline transport policy. + #[doc(hidden)] + pub fn create_dataframe_from_inline_arrow_ipc( + &self, + data: Vec, + schema: DataType, + ) -> DataFrame { + DataFrame::new( + self.clone(), + LogicalPlan::LocalRelation { + schema, + data: Some(data), + }, + ) + } + /// Create a local DataFrame from Arrow batches without materializing rows. /// /// `schema` optionally supplies the Spark schema applied by the server. With diff --git a/python/pyspark/sql/_pandas_input.py b/python/pyspark/sql/_pandas_input.py index 614287d8..f3a5b75c 100644 --- a/python/pyspark/sql/_pandas_input.py +++ b/python/pyspark/sql/_pandas_input.py @@ -274,3 +274,125 @@ def prepare_int64_table(frame, schema): [pa.array(column, type=pa.int64(), from_pandas=True) for column in columns], names=names, ) + + +def prepare_string_input(frame, schema): + """Prepare inline IPC for explicit columns of exact strings or None. + + Inspect every object before Arrow conversion: str subclasses can override + conversion hooks, and pandas treats several other objects as missing values. + Those cases must keep the existing row conversion and its error ordering. + """ + import sys + + import numpy as np + import pandas as pd + from pandas.core import frame as frame_module, generic as generic_module + from pandas.core.frame import DataFrame + from pandas.core.series import Series + from types import FunctionType, GetSetDescriptorType, MethodDescriptorType, WrapperDescriptorType + + from pyspark._pyspark import ( + BooleanType, ByteType, ShortType, IntegerType, LongType, FloatType, + DoubleType, DecimalType, StringType, StructField, StructType as NativeStructType, + ) + from pyspark.sql import types as sql_types + + # Keep small frames on the established row path. + if (pd.DataFrame is not DataFrame or type(frame) is not DataFrame + or type(schema) is not NativeStructType + or len(frame) < 10000 + or any(name in vars(frame) for name in ("astype", "notna", "where"))): + return None + # Preserve the direct normalization hooks used by the existing row path. + for name, module in (("astype", generic_module), ("notna", frame_module), + ("where", generic_module)): + method = getattr(DataFrame, name) + if (type(method) is not FunctionType or method.__name__ != name + or method.__module__ != module.__name__ or method.__globals__ is not vars(module)): + return None + if (DataFrame.astype is not generic_module.NDFrame.astype + or DataFrame.where is not generic_module.NDFrame.where): + return None + # Earlier converter branches must not recognize a string as another type. + # Check a loaded converter too: it keeps its own imported type aliases. + if sys.modules.get("pyspark.sql.types") is not sql_types: + return None + importer = sys.modules.get("pyspark.sql._arrow_input") + if importer is not None and getattr(importer, "StructType", None) is not NativeStructType: + return None + converter = sys.modules.get("pyspark.sql.pandas.types") + for name, cls in ( + ("BooleanType", BooleanType), ("ByteType", ByteType), ("ShortType", ShortType), + ("IntegerType", IntegerType), ("LongType", LongType), ("FloatType", FloatType), + ("DoubleType", DoubleType), ("DecimalType", DecimalType), ("StringType", StringType), + ("StructField", StructField), ("StructType", NativeStructType), + ): + if (getattr(sql_types, name, None) is not cls + or converter is not None and getattr(converter, name, None) is not cls): + return None + # The row path reads native fields directly. Class attributes are mutable, + # so customized Python schema access must not select a different Arrow path. + for name, cls in (("StringType", StringType), ("StructField", StructField), + ("StructType", NativeStructType)): + if (getattr(sql_types, name, None) is not cls + or cls.__getattribute__ is not object.__getattribute__ + or any("__class__" in vars(base) for base in cls.__mro__ if base is not object)): + return None + for cls, names, descriptor_type in ( + (NativeStructType, ("fields", "names"), GetSetDescriptorType), + (NativeStructType, ("__iter__",), WrapperDescriptorType), + (StructField, ("name", "dataType", "nullable", "metadata"), GetSetDescriptorType), + (StringType, ("isUTF8BinaryCollation",), MethodDescriptorType), + ): + for name in names: + descriptor = vars(cls).get(name) + if (type(descriptor) is not descriptor_type + or descriptor.__objclass__ is not cls or descriptor.__name__ != name): + return None + if not _has_canonical_columns(frame.columns): + return None + width = len(frame.columns) + if not width or len(schema.fields) != width or not frame.columns.is_unique: + return None + if not all(type(name) is str and "\x00" not in name for name in frame.columns): + return None + if not all( + type(field) is StructField and type(field.dataType) is StringType + and field.dataType.isUTF8BinaryCollation() + and type(field.name) is str and "\x00" not in field.name + for field in schema.fields + ) or len(set(schema.names)) != width: + return None + if not all(type(dtype) is type(np.dtype("object")) and dtype == np.dtype("object") + for dtype in frame.dtypes): + return None + # A custom ExtensionArray may report a genuine NumPy dtype while its + # Arrow hook emits different values from to_numpy(). Trust only pandas' + # concrete NumPy backing array, including its name in older pandas releases. + try: + from pandas.core.arrays.numpy_ import NumpyExtensionArray + except ImportError: + try: + from pandas.core.arrays.numpy_ import PandasArray as NumpyExtensionArray + except ImportError: + return None + columns = [frame.iloc[:, i] for i in range(width)] + if any(type(column) is not Series or type(column.array) is not NumpyExtensionArray + for column in columns): + return None + + from pyspark._pyspark import _pandas_string_ipc + + # Preflight exact object types before the native Arrow buffer traversal. + # Unsupported values, non-nullable None, and surrogates return to the row + # path before serialization. Keep its inline transport and configuration + # behavior rather than selecting the public Arrow import/upload policy. + arrays = [column.array._ndarray for column in columns] + if any(type(array) is not np.ndarray or array.ndim != 1 + or type(array.dtype) is not type(np.dtype("object")) + or array.dtype != np.dtype("object") or len(array) != len(frame) + for array in arrays): + return None + return _pandas_string_ipc(arrays, schema) + diff --git a/python/tests/test_arrow_handoff.py b/python/tests/test_arrow_handoff.py index 5d098808..a5314766 100644 --- a/python/tests/test_arrow_handoff.py +++ b/python/tests/test_arrow_handoff.py @@ -1725,3 +1725,900 @@ def set_data(index, value): "metadata": metadata} ], } + + +@pytest.mark.parametrize("location", [ + "subclass", "instance", "class", "owner", "backing_subclass", "backing_instance", + "backing_to_numpy", "backing_class_tolist", "backing_class_to_numpy", + "arrow_instance", "arrow_to_numpy", "arrow_class_tolist", + "arrow_class_to_numpy", +]) +@pytest.mark.parametrize("behavior", ["raise", "mutate"]) +def test_string_columns_tolist_hooks_keep_native_order( + protocol, monkeypatch, location, behavior +): + from types import MethodType + from pandas.core.base import IndexOpsMixin + + values = ["plain"] * 10000 + frame = pd.DataFrame({"original": values}, dtype=object) + schema = T.StructType([T.StructField("renamed", T.StringType())]) + original_tolist = pd.Index.tolist + original_notna, original_astype = frame.notna, frame.astype + hook_name = next((name for name in ("to_numpy", "isna") + if location.endswith("_" + name)), "tolist") + calls = [] + + class ColumnsError(RuntimeError): + pass + + def notna(*args, **kwargs): + calls.append("notna") + return original_notna(*args, **kwargs) + + def astype(*args, **kwargs): + calls.append("astype") + return original_astype(*args, **kwargs) + + def before_tolist(): + calls.append(hook_name) + if behavior == "raise": + raise ColumnsError("columns hook precedes pandas normalization") + # Install spies only when native conversion calls the columns hook. + # Before eligibility, normalization methods are the canonical methods. + monkeypatch.setitem(vars(frame), "notna", notna) + monkeypatch.setitem(vars(frame), "astype", astype) + frame.iat[0, 0], frame.iat[-1, 0] = "changed", None + + def tolist(index): + before_tolist() + return original_tolist(index) + + if location == "subclass": + class Columns(pd.Index): + pass + monkeypatch.setattr(Columns, "tolist", tolist) + frame.columns = Columns._simple_new(np.asarray(["original"], dtype=object), name=None) + assert type(frame.columns) is Columns + elif location == "instance": + monkeypatch.setitem(vars(frame.columns), "tolist", MethodType(tolist, frame.columns)) + assert type(frame.columns) is pd.Index + elif location == "backing_subclass": + class Names(np.ndarray): + def tolist(self): + before_tolist() + return super().tolist() + backing = np.asarray(["original"], dtype=object).view(Names) + frame.columns = pd.Index._simple_new(backing, name=None) + assert type(frame.columns) is pd.Index + assert type(frame.columns._values) is Names + elif location.startswith(("backing_", "arrow_")): + storage = "pyarrow" if location.startswith("arrow_") else "python" + frame.columns = pd.Index(pd.array(["original"], dtype=pd.StringDtype(storage=storage))) + backing = frame.columns._values + original = getattr(backing, hook_name) + + def backing_tolist(*args, **kwargs): + before_tolist() + return original(*args, **kwargs) + + def class_tolist(self, *args, **kwargs): + return backing_tolist(*args, **kwargs) + + if "_class_" in location: + monkeypatch.setattr(type(backing), hook_name, class_tolist) + else: + monkeypatch.setitem(vars(backing), hook_name, backing_tolist) + assert type(frame.columns) is pd.Index + else: + monkeypatch.setattr(pd.Index if location == "class" else IndexOpsMixin, + "tolist", tolist) + + assert calls == [] + with _arrow_input_service(protocol) as (spark, state): + if behavior == "raise": + with pytest.raises(ColumnsError, + match="columns hook precedes pandas normalization"): + spark.createDataFrame(frame, schema) + assert calls == [hook_name] + assert state["config_requests"] == [] + assert state["input_tables"] == [] + assert state["uploads"] == 0 + else: + spark.createDataFrame(frame, schema).toArrow() + expected = values.copy() + expected[0], expected[-1] = "changed", None + table = state["input_tables"][0] + assert table.schema.names == ["renamed"] + assert table.column(0).type == pa.string() + assert table.column(0).to_pylist() == expected + assert table.column(0).is_null().to_pylist() == [False] * 9999 + [True] + assert state["config_requests"] == [] + assert state["uploads"] == 0 + assert calls == [hook_name, "notna", "astype"] + + + +def test_string_rebound_index_values_keeps_native_error(protocol, monkeypatch): + + frame = pd.DataFrame({"original": ["plain"] * 10000}, dtype=object) + schema = T.StructType([T.StructField("renamed", T.StringType())]) + columns = frame.columns + original = pd.Index._values.fget + calls = [] + + def values(index): + if index is columns: + calls.append("_values") + raise RuntimeError("columns backing hook precedes normalization") + return original(index) + + monkeypatch.setattr(pd.Index, "_values", property(values)) + assert calls == [] + with _arrow_input_service(protocol) as (spark, state): + with pytest.raises(RuntimeError, + match="columns backing hook precedes normalization"): + spark.createDataFrame(frame, schema) + assert calls == ["_values"] + assert state["config_requests"] == [] + assert state["input_tables"] == [] + + + +@pytest.mark.parametrize("nullable", [False, True]) +def test_string_rebound_index_data_keeps_native_read_order(protocol, monkeypatch, nullable): + import json + + values = ["plain", "é", "漢字", ""] * 2500 + if nullable: + values[1] = None + frame = pd.DataFrame({"original": values}, dtype=object) + schema = T.StructType([ + T.StructField("renamed", T.StringType(), nullable, {"source": "native"}) + ]) + columns, backing = frame.columns, vars(frame.columns)["_data"] + original_notna, original_astype = frame.notna, frame.astype + calls = [] + normalizing = False + + class ColumnsDataError(RuntimeError): + pass + + def notna(*args, **kwargs): + nonlocal normalizing + normalizing = True + calls.append("notna") + return original_notna(*args, **kwargs) + + def astype(*args, **kwargs): + calls.append("astype") + return original_astype(*args, **kwargs) + + def data(index): + if index is columns: + if not normalizing: + calls.append("_data") + if calls.count("_data") > 1: + raise ColumnsDataError("columns storage read twice before normalization") + # Canonical normalization remains intact before eligibility. + monkeypatch.setitem(vars(frame), "notna", notna) + monkeypatch.setitem(vars(frame), "astype", astype) + return backing + return vars(index)["_data"] + + def set_data(index, value): + # Keep ordinary storage initialization for indexes created by pandas. + vars(index)["_data"] = value + + monkeypatch.setattr(pd.Index, "_data", property(data, set_data), raising=False) + with _arrow_input_service(protocol) as (spark, state): + spark.createDataFrame(frame, schema).toArrow() + assert calls == ["_data", "notna", "astype"] + assert not state["config_requests"] + assert state["uploads"] == 0 + assert len(state["input_tables"]) == 1 + table = state["input_tables"][0] + assert table.column(0).to_pylist() == values + assert table.column(0).is_null().to_pylist() == [value is None for value in values] + assert table.schema.equals( + pa.schema([pa.field("renamed", pa.string(), nullable=nullable)]), + check_metadata=True, + ) + assert json.loads(state["relations"][0].local_relation.schema) == { + "type": "struct", "fields": [ + {"name": "renamed", "type": "string", "nullable": nullable, + "metadata": {"source": "native"}} + ], + } + + + +@pytest.mark.parametrize("kind", [ + "object", "default", "python", "pyarrow", "python_numpy", "pyarrow_numpy", +]) +def test_string_canonical_column_storage_stays_eligible(kind): + from pyspark.sql._pandas_input import _has_canonical_columns, prepare_string_input + + frame = pd.DataFrame({"original": ["plain"] * 10000}, dtype=object) + schema = T.StructType([T.StructField("renamed", T.StringType())]) + if kind == "object": + frame.columns = pd.Index(np.asarray(["original"], dtype=object), dtype=object) + elif kind == "default": + frame.columns = pd.Index(["original"]) + elif kind.endswith("_numpy"): + from pandas.core.arrays import string_ as strings, string_arrow as arrow_strings + + module = strings if kind.startswith("python") else arrow_strings + name = ("StringArrayNumpySemantics" if kind.startswith("python") + else "ArrowStringArrayNumpySemantics") + cls = getattr(module, name, None) + if cls is None: + pytest.skip("older infer-string storage class is absent") + frame.columns = pd.Index._simple_new(cls._from_sequence(["original"]), name=None) + else: + frame.columns = pd.Index(pd.array(["original"], dtype=pd.StringDtype(storage=kind))) + assert _has_canonical_columns(frame.columns) + assert type(prepare_string_input(frame, schema)) is bytes + + + +@pytest.mark.parametrize("missing", ["module", "owner", "base"]) +def test_string_absent_private_columns_storage_keeps_row_fallback(monkeypatch, missing): + import builtins + from pandas.core.arrays import numpy_ as numpy_module, string_ as string_module + from pyspark.sql._pandas_input import _has_canonical_columns + + columns = pd.Index(pd.array(["original"], dtype=pd.StringDtype(storage="python"))) + if missing == "module": + original = builtins.__import__ + + def without_arrow(name, *args, **kwargs): + if name == "pandas.core.arrays.arrow": + raise ImportError("private storage module is absent") + return original(name, *args, **kwargs) + + monkeypatch.setattr(builtins, "__import__", without_arrow) + elif missing == "owner": + monkeypatch.delattr(numpy_module, "NumpyExtensionArray") + else: + monkeypatch.delattr(string_module, "BaseStringArray") + assert not _has_canonical_columns(columns) + + + +@pytest.mark.parametrize("cached", [False, True]) +@pytest.mark.parametrize("verify", [False, True]) +@pytest.mark.parametrize("large", ["false", "true"]) +def test_string_columns_preserve_every_ipc_cell(protocol, monkeypatch, cached, verify, large): + import json + + from pyspark.sql.pandas.types import from_arrow_metadata + + strings = ["", "plain", "é", "e\u0301", "\x00embedded\x00", "漢字", "🙂", + "\U0010ffff", "\uffff", "\u202eright-to-left"] * 1000 + nullable = [None if i % 7 == 0 else value for i, value in enumerate(reversed(strings))] + expected = [strings, nullable, [None] * len(strings)] + frame = pd.DataFrame(dict(zip(["original", "b", "c"], expected)), dtype=object) + schema = T.StructType([ + T.StructField("renamed.`text", T.StringType(), False, {"source": "string", "id": 3}), + T.StructField("nullable", T.StringType(), True, {"tag": [1, 2]}), + T.StructField("all_null", T.StringType()), + ]) + + def no_rows(*args, **kwargs): + raise AssertionError("string fast path must not materialize Python rows") + + for method in ("to_numpy", "iterrows", "itertuples"): + monkeypatch.setattr(pd.DataFrame, method, no_rows) + monkeypatch.setattr(pd.DataFrame, "values", property(no_rows)) + monkeypatch.setattr(pd.Series, "tolist", no_rows) + overrides = { + "spark.sql.session.localRelationCacheThreshold": "1" if cached else "999999999", + "spark.sql.execution.arrow.useLargeVarTypes": large, + "spark.sql.session.localRelationChunkSizeRows": "997", + } + with _arrow_input_service(protocol, **overrides) as (spark, state): + spark.createDataFrame(frame, schema, verifySchema=verify).toArrow() + table = state["input_tables"][0] + assert table.num_rows == len(strings) + assert table.schema.names == schema.names + for column, values, field in zip(table.columns, expected, schema.fields): + assert column.type == pa.string() + assert column.to_pylist() == values + actual_field = table.schema.field(field.name) + assert actual_field.nullable == field.nullable + assert from_arrow_metadata(actual_field.metadata) is None + spark_field = json.loads(state["relations"][0].local_relation.schema)["fields"][ + schema.names.index(field.name)] + assert spark_field["metadata"] == field.metadata + assert state["config_requests"] == [] + assert state["uploads"] == 0 + + + +@pytest.mark.parametrize("reason", [ + "small", "empty", "inferred", "names", "ddl", "cast", "atomic", "mixed_schema", + "width", "duplicate_input", "duplicate_schema", "input_nul", "schema_nul", + "input_number", "datetime", "date", "nan", "pd_na", "nat", "mixed", "bytes", + "nonnullable_null", "frame_subclass", "string_subclass", "string_dtype", + "arrow_string_dtype", "late_unsupported", "surrogate", "late_surrogate", "collation", +]) +def test_string_column_guard_preserves_row_fallback(reason): + from pyspark.sql._pandas_input import prepare_string_input + + class Frame(pd.DataFrame): + pass + + class Text(str): + def __str__(self): + raise AssertionError("guard must not call conversion hooks") + + class Unsupported: + def __str__(self): + raise AssertionError("guard must not stringify unsupported objects") + + values = ["plain"] * 10000 + frame = pd.DataFrame({"text": values}, dtype=object) + schema = T.StructType([T.StructField("text", T.StringType())]) + if reason == "small": + frame = frame.iloc[:9999] + elif reason == "empty": + frame = frame.iloc[:0] + elif reason in ("inferred", "names", "ddl", "cast", "atomic"): + schema = {"inferred": None, "names": ["text"], "ddl": "text STRING", + "cast": T.StructType([T.StructField("text", T.LongType())]), + "atomic": T.StringType()}[reason] + elif reason == "mixed_schema": + frame["long"] = np.arange(10000) + schema.add("long", T.LongType()) + elif reason == "width": + schema.add("extra", T.StringType()) + elif reason in ("duplicate_input", "duplicate_schema"): + frame["other"] = values + schema.add("other", T.StringType()) + if reason == "duplicate_input": + frame.columns = ["text", "text"] + else: + schema = T.StructType([T.StructField("text", T.StringType())] * 2) + elif reason in ("input_nul", "input_number"): + frame.columns = ["text\x00" if reason == "input_nul" else 1] + elif reason == "schema_nul": + schema = T.StructType([T.StructField("text\x00", T.StringType())]) + elif reason == "frame_subclass": + frame = Frame(frame) + elif reason in ("string_dtype", "arrow_string_dtype"): + frame = pd.DataFrame({"text": pd.Series( + values, dtype=pd.StringDtype(storage="pyarrow" if reason == "arrow_string_dtype" + else "python"))}) + elif reason == "collation": + schema = T.StructType([T.StructField("text", T.StringType("UTF8_LCASE"))]) + else: + value = { + "datetime": datetime(2020, 1, 1), "date": date(2020, 1, 1), + "nan": float("nan"), "pd_na": pd.NA, "nat": pd.NaT, "mixed": 42, + "bytes": b"plain", "nonnullable_null": None, "string_subclass": Text("plain"), + "late_unsupported": Unsupported(), "surrogate": "\ud800", "late_surrogate": "\udfff", + }[reason] + if reason == "nonnullable_null": + schema = T.StructType([T.StructField("text", T.StringType(), False)]) + if reason in ("late_unsupported", "late_surrogate"): + # pandas 3 infers StringDtype for a newly assigned list column. + # Keep these unsupported-object fixtures on concrete object storage. + frame["other"] = pd.Series(values, dtype=object) + schema.add("other", T.StringType()) + backing = frame["other"].array._ndarray + assert type(backing) is np.ndarray and backing.dtype == np.dtype("object") + frame.iat[-1, -1] = value + if reason in ("late_unsupported", "late_surrogate"): + assert frame.iat[-1, -1] is value + assert prepare_string_input(frame, schema) is None + + + +@pytest.mark.parametrize("reason", [ + "string_subclass", "frame_subclass", "nan", "pd_na", "nat", "string_dtype", + "nonnullable_null", "inferred", "names", "ddl", "small", "collation", "cast", "mixed", +]) +def test_string_public_fallback_keeps_row_semantics(protocol, monkeypatch, reason): + calls = [] + + class Text(str): + def __str__(self): + calls.append("str") + return "changed" + + class Frame(pd.DataFrame): + pass + + values = ["plain"] * (9999 if reason == "small" else 10000) + if reason in ("nan", "pd_na", "nat", "nonnullable_null", "mixed"): + values[-1] = {"nan": float("nan"), "pd_na": pd.NA, "nat": pd.NaT, + "nonnullable_null": None, "mixed": 42}[reason] + elif reason == "string_subclass": + values[-1] = Text("plain") + frame_type = Frame if reason == "frame_subclass" else pd.DataFrame + frame = frame_type({"text": values}, dtype="string" if reason == "string_dtype" else object) + schema = T.StructType([T.StructField("text", T.StringType(), reason != "nonnullable_null")]) + if reason == "inferred": + schema = None + elif reason == "names": + schema = ["renamed"] + elif reason == "ddl": + schema = "text STRING" + elif reason == "collation": + schema = T.StructType([T.StructField("text", T.StringType("UTF8_LCASE"))]) + elif reason == "cast": + schema = T.StructType([T.StructField("text", T.LongType())]) + astype = pd.DataFrame.astype + + def record_rows(self, *args, **kwargs): + calls.append("astype") + return astype(self, *args, **kwargs) + + monkeypatch.setattr(pd.DataFrame, "astype", record_rows) + with _arrow_input_service(protocol) as (spark, state): + if reason in ("nonnullable_null", "cast", "mixed"): + from pyspark.errors import PySparkRuntimeError + + message = ("Column 'text' is declared as non-nullable but contains null values" + if reason == "nonnullable_null" else "Type mismatch in row data") + with pytest.raises(PySparkRuntimeError, match=message): + spark.createDataFrame(frame, schema) + assert calls == ["astype"] + assert state["config_requests"] == [] + assert state["relations"] == [] + return + spark.createDataFrame(frame, schema).toArrow() + assert calls == ["astype"] # Native extraction keeps str.__str__ overrides uncalled. + table = state["input_tables"][0] + assert table.schema.names == ["renamed" if reason == "names" else "text"] + expected = ["plain"] * len(values) + if reason in ("nan", "pd_na", "nat"): + expected[-1] = None + assert table.column(0).to_pylist() == expected + assert state["config_requests"] == [] + + + +@pytest.mark.parametrize("surrogate", ["\ud800", "\udfff"]) +def test_string_surrogate_fallback_precedes_importer_configuration(protocol, monkeypatch, surrogate): + values = ["plain"] * 10000 + values[-1] = surrogate + frame = pd.DataFrame({"text": values}, dtype=object) + schema = T.StructType([T.StructField("text", T.StringType())]) + astype = pd.DataFrame.astype + calls = [] + + def record_rows(self, *args, **kwargs): + calls.append("astype") + return astype(self, *args, **kwargs) + + monkeypatch.setattr(pd.DataFrame, "astype", record_rows) + with _arrow_input_service(protocol) as (spark, state): + state["config_error"] = grpc.StatusCode.INVALID_ARGUMENT + with pytest.raises(TypeError, match="unsupported value type for createDataFrame: str"): + spark.createDataFrame(frame, schema) + assert calls == ["astype"] + assert state["config_requests"] == [] + assert state["relations"] == [] + assert state["uploads"] == 0 + + + +@pytest.mark.parametrize("class_name,attribute", [ + ("StructType", "names"), ("StructType", "fields"), ("StructType", "__iter__"), + ("StructField", "name"), ("StructField", "dataType"), ("StructField", "nullable"), + ("StructField", "metadata"), ("StringType", "isUTF8BinaryCollation"), + ("StructType", "__getattribute__"), ("StructField", "__getattribute__"), + ("StringType", "__getattribute__"), ("StructType", "__class__"), + ("StructField", "__class__"), ("StringType", "__class__"), + ("StructType", "alias"), ("StructField", "alias"), ("StringType", "alias"), +]) +@pytest.mark.parametrize("before_import", [False, True]) +def test_string_custom_schema_descriptors_keep_native_row_path( + protocol, class_name, attribute, before_import +): + import sys + + script = r''' +import sys +import tempfile +from pathlib import Path +import pandas as pd +from pyspark.sql import types as T + +import importlib.util +spec = importlib.util.spec_from_file_location("string_schema_capture", sys.argv[4]) +capture = importlib.util.module_from_spec(spec) +sys.modules[spec.name] = capture +spec.loader.exec_module(capture) +_arrow_input_service, protocol = capture._arrow_input_service, capture.protocol +frame = pd.DataFrame({"text": ["plain"] * 10000}, dtype=object) +schema = T.StructType([T.StructField("text", T.StringType(), True, {"source": "native"})]) +class_name, attribute, before_import = sys.argv[1:4] +if before_import == "True": + # The capture module imports prepare_int64_table; force the public call to + # import a new helper after the class or public alias has been customized. + sys.modules.pop("pyspark.sql._pandas_input", None) +assert ("pyspark.sql._pandas_input" in sys.modules) == (before_import == "False") +cls = getattr(T, class_name) +if attribute == "alias": + setattr(T, class_name, object) +else: + if attribute == "__getattribute__": + replacement = lambda self, name: object.__getattribute__(self, name) + elif attribute == "__iter__": + replacement = lambda self: iter([T.StructField("changed", T.StringType())]) + elif attribute == "__class__": + attribute = "__getattribute__" + replacement = lambda self, name: object if name == "__class__" else ( + object.__getattribute__(self, name)) + elif attribute == "isUTF8BinaryCollation": + replacement = lambda self: False + else: + value = {"names": ["changed"], "fields": [], "name": "changed", + "dataType": T.StringType(), "nullable": False, "metadata": {"changed": True}}[attribute] + replacement = property(lambda self: value) + setattr(cls, attribute, replacement) +class Factory: + def mktemp(self, name): + return Path(tempfile.mkdtemp(prefix=name)) +with _arrow_input_service(protocol.__wrapped__(Factory())) as (spark, state): + spark.createDataFrame(frame, schema).toArrow() + assert state["config_requests"] == [], state["config_requests"] + table = state["input_tables"][0] + assert table.schema.names == ["text"], table.schema + assert table.column(0).to_pylist() == ["plain"] * 10000 +''' + result = subprocess.run([sys.executable, "-c", script, class_name, attribute, + str(before_import), __file__], + capture_output=True, text=True, timeout=30) + assert result.returncode == 0, result.stdout + result.stderr + + + +def test_string_rebound_pandas_alias_keeps_subclass_fallback(monkeypatch): + from pyspark.sql._pandas_input import prepare_string_input + + class Frame(pd.DataFrame): + pass + + frame = Frame({"text": ["plain"] * 10000}, dtype=object) + schema = T.StructType([T.StructField("text", T.StringType())]) + monkeypatch.setattr(pd, "DataFrame", Frame) + assert prepare_string_input(frame, schema) is None + + + +@pytest.mark.parametrize("class_name", [ + "BooleanType", "ByteType", "ShortType", "IntegerType", "LongType", + "FloatType", "DoubleType", "DecimalType", +]) +@pytest.mark.parametrize("converter_loaded", [False, True]) +def test_string_rebound_preceding_converter_alias_keeps_native_row_path( + protocol, class_name, converter_loaded +): + import sys + + script = r''' +import importlib.util +import sys +import tempfile +from pathlib import Path +import pandas as pd +from pyspark.sql import types as T + +spec = importlib.util.spec_from_file_location("string_alias_capture", sys.argv[3]) +capture = importlib.util.module_from_spec(spec) +sys.modules[spec.name] = capture +spec.loader.exec_module(capture) +_arrow_input_service, protocol = capture._arrow_input_service, capture.protocol +values = ["plain", "", "\x00", "é", "漢字", "🙂", None, "tail"] * 1250 +frame = pd.DataFrame({"text": values}, dtype=object) +schema = T.StructType([T.StructField("text", T.StringType())]) +class_name, converter_loaded = sys.argv[1:3] +if converter_loaded == "True": + import pyspark.sql.pandas.types as converter + setattr(converter, class_name, T.StringType) +else: + sys.modules.pop("pyspark.sql.pandas.types", None) + setattr(T, class_name, T.StringType) + assert "pyspark.sql.pandas.types" not in sys.modules +class Factory: + def mktemp(self, name): + return Path(tempfile.mkdtemp(prefix=name)) +with _arrow_input_service(protocol.__wrapped__(Factory())) as (spark, state): + spark.createDataFrame(frame, schema).toArrow() + assert state["config_requests"] == [], state["config_requests"] + table = state["input_tables"][0] + assert table.schema.names == ["text"], table.schema + assert table.column(0).to_pylist() == values +''' + result = subprocess.run([sys.executable, "-c", script, class_name, + str(converter_loaded), __file__], + capture_output=True, text=True, timeout=30) + assert result.returncode == 0, result.stdout + result.stderr + + + +@pytest.mark.parametrize("replacement", ["string", "none"]) +def test_string_rebound_importer_schema_alias_keeps_native_row_path(protocol, replacement): + import sys + + script = r''' +import importlib.util +import json +import sys +import tempfile +from pathlib import Path +import pandas as pd +from pyspark.sql import types as T + +spec = importlib.util.spec_from_file_location("string_importer_capture", sys.argv[2]) +capture = importlib.util.module_from_spec(spec) +sys.modules[spec.name] = capture +spec.loader.exec_module(capture) +_arrow_input_service, protocol = capture._arrow_input_service, capture.protocol +strings = ["plain", "", "\x00", "é", "漢字", "🙂", "e\u0301", "tail"] * 1250 +nullable = [None if i % 7 == 0 else value for i, value in enumerate(reversed(strings))] +expected = [strings, nullable, [None] * len(strings)] +frame = pd.DataFrame(dict(zip(["original", "b", "c"], expected)), dtype=object) +schema = T.StructType([ + T.StructField("text", T.StringType(), False, {"source": "native", "id": 3}), + T.StructField("nullable", T.StringType(), True, {"tag": [1, 2]}), + T.StructField("all_null", T.StringType()), +]) +import pyspark.sql._arrow_input as importer +importer.StructType = T.StringType if sys.argv[1] == "string" else None +class Factory: + def mktemp(self, name): + return Path(tempfile.mkdtemp(prefix=name)) +with _arrow_input_service(protocol.__wrapped__(Factory())) as (spark, state): + spark.createDataFrame(frame, schema).toArrow() + assert state["config_requests"] == [], state["config_requests"] + table = state["input_tables"][0] + assert table.schema.names == schema.names, table.schema + spark_fields = json.loads(state["relations"][0].local_relation.schema)["fields"] + for column, values, field, spark_field in zip(table.columns, expected, schema.fields, spark_fields): + assert column.to_pylist() == values + actual_field = table.schema.field(field.name) + assert actual_field.nullable == field.nullable + assert actual_field.metadata is None + assert spark_field["metadata"] == field.metadata +''' + result = subprocess.run([sys.executable, "-c", script, replacement, __file__], + capture_output=True, text=True, timeout=30) + assert result.returncode == 0, result.stdout + result.stderr + + + +@pytest.mark.parametrize("dtype_mode", ["equal", "class_spoof", "exact_numpy"]) +def test_string_custom_extension_dtype_cannot_impersonate_numpy_object(protocol, dtype_mode): + from pandas.api.extensions import ExtensionArray, ExtensionDtype, take + from pyspark.sql._pandas_input import prepare_string_input + + arrow_calls = [] + + class ObjectEqualDtype(ExtensionDtype): + name = "object_equal_string" + type = str + kind = "O" + na_value = None + + @classmethod + def construct_array_type(cls): + return DifferentArrowArray + + def __eq__(self, other): + return (isinstance(other, ObjectEqualDtype) + or isinstance(other, np.dtype) and other.kind == "O") + + def __hash__(self): + return hash(self.name) + + class DifferentArrowArray(ExtensionArray): + def __init__(self, values): + self._values = np.asarray(values, dtype=object) + + @property + def dtype(self): + return ObjectEqualDtype() + + @property + def nbytes(self): + return self._values.nbytes + + def __len__(self): + return len(self._values) + + def __getitem__(self, key): + result = self._values[key] + return type(self)(result) if isinstance(result, np.ndarray) else result + + def isna(self): + return pd.isna(self._values) + + def take(self, indices, allow_fill=False, fill_value=None): + return type(self)(take(self._values, indices, allow_fill=allow_fill, + fill_value=fill_value)) + + def copy(self): + return type(self)(self._values.copy()) + + @classmethod + def _from_sequence(cls, scalars, *, dtype=None, copy=False): + values = np.asarray(scalars, dtype=object) + return cls(values.copy() if copy else values) + + def to_numpy(self, dtype=None, copy=False, na_value=None): + if dtype is not None: + return self._values.astype(dtype, copy=copy) + return self._values.copy() if copy else self._values + + def __arrow_array__(self, type=None): + arrow_calls.append("arrow") + # A scan of to_numpy() cannot establish the values this hook emits. + return pa.array(["different"] * len(self), type=type) + + array = DifferentArrowArray(["plain"] * 10000) + frame = pd.DataFrame({"text": pd.Series(array, copy=False)}) + schema = T.StructType([T.StructField("text", T.StringType())]) + if dtype_mode == "exact_numpy": + # An ExtensionArray may change its advertised dtype after construction. + # The dtype and NumPy scan both pass while its Arrow hook changes cells. + DifferentArrowArray.dtype = property(lambda self: np.dtype("object")) + DifferentArrowArray.__array__ = lambda self, dtype=None, copy=None: self._values + elif dtype_mode == "class_spoof": + def spoof_class(self, name): + if name == "__class__": + return type(np.dtype("object")) + return object.__getattribute__(self, name) + + ObjectEqualDtype.__getattribute__ = spoof_class + dtype = frame.dtypes.iloc[0] + assert type(frame) is pd.DataFrame + assert type(frame.iloc[:, 0].array) is DifferentArrowArray + assert type(dtype) is (type(np.dtype("object")) if dtype_mode == "exact_numpy" + else ObjectEqualDtype) + assert isinstance(dtype, np.dtype) is (dtype_mode != "equal") + assert dtype == np.dtype("object") + assert prepare_string_input(frame, schema) is None + with _arrow_input_service(protocol) as (spark, state): + if dtype_mode == "exact_numpy": + # This deliberately inconsistent EA already fails pandas' row-path + # astype(object). Preserve that error before importer configuration. + with pytest.raises(TypeError, match="cannot set astype"): + spark.createDataFrame(frame, schema) + assert state["input_tables"] == [] + else: + spark.createDataFrame(frame, schema).toArrow() + assert state["input_tables"][0].column(0).to_pylist() == ["plain"] * 10000 + assert state["config_requests"] == [] + assert arrow_calls == [] + + + +@pytest.mark.parametrize("method", ["astype", "notna"]) +def test_string_instance_conversion_hooks_keep_row_fallback(protocol, method): + frame = pd.DataFrame({"text": ["plain"] * 10000}, dtype=object) + schema = T.StructType([T.StructField("text", T.StringType())]) + original = getattr(frame, method) + calls = [] + expected = ["plain"] * 10000 + if method == "astype": + expected[0] = "changed" + + def customized(*args, **kwargs): + calls.append(method) + result = original(*args, **kwargs) + result.iat[0, 0] = "changed" + return result + else: + expected[0] = None + + def customized(*args, **kwargs): + calls.append(method) + result = original(*args, **kwargs) + result.iat[0, 0] = False + return result + setattr(frame, method, customized) + with _arrow_input_service(protocol) as (spark, state): + spark.createDataFrame(frame, schema).toArrow() + assert calls == [method] + assert state["config_requests"] == [] + assert state["uploads"] == 0 + assert state["input_tables"][0].column(0).to_pylist() == expected + + + +@pytest.mark.parametrize("method", ["astype", "notna", "where"]) +@pytest.mark.parametrize("before_import", [False, True]) +def test_string_class_conversion_hooks_keep_row_fallback( + protocol, monkeypatch, method, before_import +): + import importlib + import sys + + name = "pyspark.sql._pandas_input" + if before_import: + monkeypatch.delitem(sys.modules, name, raising=False) + else: + importlib.import_module(name) + frame = pd.DataFrame({"text": ["plain"] * 10000}, dtype=object) + schema = T.StructType([T.StructField("text", T.StringType())]) + original = getattr(pd.DataFrame, method) + calls = [] + expected = ["plain"] * 10000 + expected[0] = None if method == "notna" else "changed" + + def customized(self, *args, **kwargs): + calls.append(method) + result = original(self, *args, **kwargs) + result.iat[0, 0] = False if method == "notna" else "changed" + return result + + monkeypatch.setattr(pd.DataFrame, method, customized) + with _arrow_input_service(protocol) as (spark, state): + spark.createDataFrame(frame, schema).toArrow() + assert calls == [method] + assert state["config_requests"] == [] + assert state["uploads"] == 0 + assert state["input_tables"][0].column(0).to_pylist() == expected + + + +def test_string_instance_where_override_keeps_existing_row_behavior(protocol): + from pyspark.sql._pandas_input import prepare_string_input + + frame = pd.DataFrame({"text": ["plain"] * 10000}, dtype=object) + schema = T.StructType([T.StructField("text", T.StringType())]) + + def customized(*args, **kwargs): + raise AssertionError("row normalization calls where on the astype result") + + frame.where = customized + assert prepare_string_input(frame, schema) is None + with _arrow_input_service(protocol) as (spark, state): + spark.createDataFrame(frame, schema).toArrow() + assert state["config_requests"] == [] + assert state["input_tables"][0].column(0).to_pylist() == ["plain"] * 10000 + + + +def test_string_native_backing_does_not_add_series_conversion_hooks(protocol, monkeypatch): + frame = pd.DataFrame({"text": ["plain"] * 10000}, dtype=object) + schema = T.StructType([T.StructField("text", T.StringType())]) + + def customized(*args, **kwargs): + raise AssertionError("native string input must not call Series.to_numpy") + + monkeypatch.setattr(pd.Series, "to_numpy", customized) + with _arrow_input_service(protocol) as (spark, state): + spark.createDataFrame(frame, schema).toArrow() + assert state["config_requests"] == [] + assert state["input_tables"][0].column(0).to_pylist() == ["plain"] * 10000 + + + +def test_string_large_inline_input_keeps_row_transport_and_schema(protocol): + values = ["漢字-🙂-" + str(i) for i in range(100000)] + frame = pd.DataFrame({"text": values}, dtype=object) + schema = T.StructType([T.StructField("renamed", T.StringType(), False, + {"source": "inline-string", "size": len(values)})]) + with _arrow_input_service(protocol, **{ + "spark.sql.session.localRelationCacheThreshold": "1", + "spark.sql.session.localRelationSizeLimit": "1", + "spark.sql.execution.arrow.useLargeVarTypes": "true", + }) as (spark, state): + spark.createDataFrame(frame, schema).toArrow() + assert state["config_requests"] == [] + assert state["uploads"] == 0 + table = state["input_tables"][0] + assert table.schema.names == ["renamed"] + assert table.column(0).type == pa.string() + assert table.schema.field(0).nullable is False + assert table.schema.field(0).metadata is None + assert table.column(0).to_pylist() == values + assert __import__("json").loads(state["relations"][0].local_relation.schema) == ( + __import__("json").loads(schema.json())) + From 09d0b13d38741e070252021da1623e3ec9011583 Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Fri, 30 Oct 2026 02:59:31 -0700 Subject: [PATCH 20/23] [SPARK-60720][CONNECT] Remove extra EOF blank lines Co-authored-by: Isaac --- python/pyspark/sql/_pandas_input.py | 1 - python/tests/test_arrow_handoff.py | 1 - 2 files changed, 2 deletions(-) diff --git a/python/pyspark/sql/_pandas_input.py b/python/pyspark/sql/_pandas_input.py index f3a5b75c..78d1eaf2 100644 --- a/python/pyspark/sql/_pandas_input.py +++ b/python/pyspark/sql/_pandas_input.py @@ -395,4 +395,3 @@ def prepare_string_input(frame, schema): for array in arrays): return None return _pandas_string_ipc(arrays, schema) - diff --git a/python/tests/test_arrow_handoff.py b/python/tests/test_arrow_handoff.py index a5314766..8fd50acd 100644 --- a/python/tests/test_arrow_handoff.py +++ b/python/tests/test_arrow_handoff.py @@ -2621,4 +2621,3 @@ def test_string_large_inline_input_keeps_row_transport_and_schema(protocol): assert table.column(0).to_pylist() == values assert __import__("json").loads(state["relations"][0].local_relation.schema) == ( __import__("json").loads(schema.json())) - From 4304062cf775c3d50a44ad2feb208c4d9f3c592d Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Thu, 5 Nov 2026 14:31:29 -0800 Subject: [PATCH 21/23] Reconcile current master interfaces and formatting after rebase Co-authored-by: Isaac --- crates/spark-connect/src/session.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/spark-connect/src/session.rs b/crates/spark-connect/src/session.rs index 10c57aac..e7290132 100644 --- a/crates/spark-connect/src/session.rs +++ b/crates/spark-connect/src/session.rs @@ -371,7 +371,7 @@ impl SparkSession { DataFrame::new( self.clone(), LogicalPlan::LocalRelation { - schema, + schema: Some(schema), data: Some(data), }, ) From b2274e4f6a914f85becccbbc71d5cc6801c654d8 Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Thu, 5 Nov 2026 14:54:08 -0800 Subject: [PATCH 22/23] Align string label regression tests with current native iteration Co-authored-by: Isaac --- python/tests/test_arrow_handoff.py | 126 ++++++++++++++++++----------- 1 file changed, 81 insertions(+), 45 deletions(-) diff --git a/python/tests/test_arrow_handoff.py b/python/tests/test_arrow_handoff.py index 8fd50acd..13243164 100644 --- a/python/tests/test_arrow_handoff.py +++ b/python/tests/test_arrow_handoff.py @@ -1728,26 +1728,22 @@ def set_data(index, value): @pytest.mark.parametrize("location", [ - "subclass", "instance", "class", "owner", "backing_subclass", "backing_instance", - "backing_to_numpy", "backing_class_tolist", "backing_class_to_numpy", - "arrow_instance", "arrow_to_numpy", "arrow_class_tolist", - "arrow_class_to_numpy", + "subclass", "class", "owner", "backing_subclass_item", "backing_class_iter", + "backing_class_getitem", "backing_instance_box", "arrow_class_iter", ]) @pytest.mark.parametrize("behavior", ["raise", "mutate"]) -def test_string_columns_tolist_hooks_keep_native_order( +def test_string_columns_iterator_hooks_keep_native_order( protocol, monkeypatch, location, behavior ): - from types import MethodType from pandas.core.base import IndexOpsMixin values = ["plain"] * 10000 frame = pd.DataFrame({"original": values}, dtype=object) schema = T.StructType([T.StructField("renamed", T.StringType())]) - original_tolist = pd.Index.tolist + original_iter = pd.Index.__iter__ original_notna, original_astype = frame.notna, frame.astype - hook_name = next((name for name in ("to_numpy", "isna") - if location.endswith("_" + name)), "tolist") calls = [] + invoked = False class ColumnsError(RuntimeError): pass @@ -1760,67 +1756,82 @@ def astype(*args, **kwargs): calls.append("astype") return original_astype(*args, **kwargs) - def before_tolist(): - calls.append(hook_name) + def before_hook(): + nonlocal invoked + if invoked: + return + invoked = True + calls.append("columns iterator") if behavior == "raise": raise ColumnsError("columns hook precedes pandas normalization") - # Install spies only when native conversion calls the columns hook. + # Install spies only when native iteration reaches the columns hook. # Before eligibility, normalization methods are the canonical methods. monkeypatch.setitem(vars(frame), "notna", notna) monkeypatch.setitem(vars(frame), "astype", astype) frame.iat[0, 0], frame.iat[-1, 0] = "changed", None - def tolist(index): - before_tolist() - return original_tolist(index) + def iterate(index): + if index is columns: + before_hook() + return original_iter(index) if location == "subclass": class Columns(pd.Index): pass - monkeypatch.setattr(Columns, "tolist", tolist) - frame.columns = Columns._simple_new(np.asarray(["original"], dtype=object), name=None) + monkeypatch.setattr(Columns, "__iter__", iterate) + labels = pd.Index(["original"], dtype=object) + frame.columns = Columns._simple_new(labels._values, name=None) assert type(frame.columns) is Columns - elif location == "instance": - monkeypatch.setitem(vars(frame.columns), "tolist", MethodType(tolist, frame.columns)) - assert type(frame.columns) is pd.Index - elif location == "backing_subclass": + elif location == "backing_subclass_item": + # IndexOpsMixin iteration uses ndarray.item, not ndarray.__iter__. class Names(np.ndarray): - def tolist(self): - before_tolist() - return super().tolist() + def item(self, *args): + before_hook() + return super().item(*args) backing = np.asarray(["original"], dtype=object).view(Names) frame.columns = pd.Index._simple_new(backing, name=None) - assert type(frame.columns) is pd.Index assert type(frame.columns._values) is Names elif location.startswith(("backing_", "arrow_")): storage = "pyarrow" if location.startswith("arrow_") else "python" frame.columns = pd.Index(pd.array(["original"], dtype=pd.StringDtype(storage=storage))) backing = frame.columns._values - original = getattr(backing, hook_name) + if location.endswith("class_iter"): + original = type(backing).__iter__ - def backing_tolist(*args, **kwargs): - before_tolist() - return original(*args, **kwargs) + def backing_iter(array): + if array is backing: + before_hook() + return original(array) - def class_tolist(self, *args, **kwargs): - return backing_tolist(*args, **kwargs) + monkeypatch.setattr(type(backing), "__iter__", backing_iter) + elif location.endswith("class_getitem"): + original = type(backing).__getitem__ + + def backing_getitem(array, key): + if array is backing: + before_hook() + return original(array, key) - if "_class_" in location: - monkeypatch.setattr(type(backing), hook_name, class_tolist) + monkeypatch.setattr(type(backing), "__getitem__", backing_getitem) else: - monkeypatch.setitem(vars(backing), hook_name, backing_tolist) - assert type(frame.columns) is pd.Index + original = backing._box_func + + def box(value): + before_hook() + return original(value) + + monkeypatch.setitem(vars(backing), "_box_func", box) else: monkeypatch.setattr(pd.Index if location == "class" else IndexOpsMixin, - "tolist", tolist) + "__iter__", iterate) + columns = frame.columns - assert calls == [] with _arrow_input_service(protocol) as (spark, state): if behavior == "raise": with pytest.raises(ColumnsError, match="columns hook precedes pandas normalization"): spark.createDataFrame(frame, schema) - assert calls == [hook_name] + assert calls == ["columns iterator"] assert state["config_requests"] == [] assert state["input_tables"] == [] assert state["uploads"] == 0 @@ -1828,14 +1839,19 @@ def class_tolist(self, *args, **kwargs): spark.createDataFrame(frame, schema).toArrow() expected = values.copy() expected[0], expected[-1] = "changed", None + import json + table = state["input_tables"][0] assert table.schema.names == ["renamed"] assert table.column(0).type == pa.string() assert table.column(0).to_pylist() == expected assert table.column(0).is_null().to_pylist() == [False] * 9999 + [True] + assert table.schema.equals(pa.schema([pa.field("renamed", pa.string())]), + check_metadata=True) + assert json.loads(state["relations"][0].local_relation.schema) == json.loads(schema.json()) assert state["config_requests"] == [] assert state["uploads"] == 0 - assert calls == [hook_name, "notna", "astype"] + assert calls == ["columns iterator", "notna", "astype"] @@ -1880,6 +1896,7 @@ def test_string_rebound_index_data_keeps_native_read_order(protocol, monkeypatch original_notna, original_astype = frame.notna, frame.astype calls = [] normalizing = False + allowed_reads = None class ColumnsDataError(RuntimeError): pass @@ -1898,8 +1915,8 @@ def data(index): if index is columns: if not normalizing: calls.append("_data") - if calls.count("_data") > 1: - raise ColumnsDataError("columns storage read twice before normalization") + if allowed_reads is not None and calls.count("_data") > allowed_reads: + raise ColumnsDataError("eligibility added a columns storage read") # Canonical normalization remains intact before eligibility. monkeypatch.setitem(vars(frame), "notna", notna) monkeypatch.setitem(vars(frame), "astype", astype) @@ -1911,20 +1928,39 @@ def set_data(index, value): vars(index)["_data"] = value monkeypatch.setattr(pd.Index, "_data", property(data, set_data), raising=False) + import pyspark.sql._pandas_input as helper + with _arrow_input_service(protocol) as (spark, state): + # Measure the current native fallback's storage reads, rather than the + # former tolist path's single-read contract. + with monkeypatch.context() as baseline: + baseline.setattr(helper, "prepare_string_input", lambda *_: None) + spark.createDataFrame(frame, schema).toArrow() + baseline_calls = calls.copy() + allowed_reads = baseline_calls.index("notna") + assert allowed_reads > 0 + assert baseline_calls == ["_data"] * allowed_reads + ["notna", "astype"] + calls.clear() + normalizing = False + # Restore canonical methods before the candidate eligibility check; + # native label iteration will install the spies through the storage hook. + monkeypatch.delitem(vars(frame), "notna") + monkeypatch.delitem(vars(frame), "astype") spark.createDataFrame(frame, schema).toArrow() - assert calls == ["_data", "notna", "astype"] + assert calls == baseline_calls assert not state["config_requests"] assert state["uploads"] == 0 - assert len(state["input_tables"]) == 1 - table = state["input_tables"][0] + assert len(state["input_tables"]) == 2 + assert len(state["relations"]) == 2 + table = state["input_tables"][1] + assert table.equals(state["input_tables"][0], check_metadata=True) assert table.column(0).to_pylist() == values assert table.column(0).is_null().to_pylist() == [value is None for value in values] assert table.schema.equals( pa.schema([pa.field("renamed", pa.string(), nullable=nullable)]), check_metadata=True, ) - assert json.loads(state["relations"][0].local_relation.schema) == { + assert json.loads(state["relations"][1].local_relation.schema) == { "type": "struct", "fields": [ {"name": "renamed", "type": "string", "nullable": nullable, "metadata": {"source": "native"}} From e31aad6482c1b748924991a8c01aabbea74db385 Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Thu, 5 Nov 2026 16:05:12 -0800 Subject: [PATCH 23/23] [CONNECT] Reconcile remaining change with current inference architecture Co-authored-by: Isaac --- crates/pyspark-rs/src/lib.rs | 3 +++ crates/pyspark-rs/src/session.rs | 25 +++++++++++++++++++++++++ python/tests/test_arrow_handoff.py | 10 ++++++++-- 3 files changed, 36 insertions(+), 2 deletions(-) diff --git a/crates/pyspark-rs/src/lib.rs b/crates/pyspark-rs/src/lib.rs index 15796178..807fce24 100644 --- a/crates/pyspark-rs/src/lib.rs +++ b/crates/pyspark-rs/src/lib.rs @@ -36,6 +36,7 @@ mod group; mod infer; mod ml; mod observation; +mod pandas_input; mod pipelines; mod profiler; mod pyarrow_export; @@ -130,6 +131,8 @@ fn _pyspark(_py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_function(wrap_pyfunction!(types::parse_datatype_json, m)?)?; m.add_function(wrap_pyfunction!(column::column_reference_name, m)?)?; + m.add_function(wrap_pyfunction!(pandas_input::_pandas_string_ipc, m)?)?; + // Spark Declarative Pipelines (SDP) command execution m.add_class::()?; m.add_function(wrap_pyfunction!( diff --git a/crates/pyspark-rs/src/session.rs b/crates/pyspark-rs/src/session.rs index 4edbd2a9..8345a595 100644 --- a/crates/pyspark-rs/src/session.rs +++ b/crates/pyspark-rs/src/session.rs @@ -788,6 +788,31 @@ impl PySparkSession { _ => frame, }); } + let ipc = py + .import("pyspark.sql._pandas_input")? + .getattr("prepare_string_input")? + .call1((data, schema))?; + if !ipc.is_none() { + let fields = match &spec { + Spec::Struct(fields) => fields, + _ => { + return Err(pyo3::exceptions::PyValueError::new_err( + "string input requires a struct schema", + )); + } + }; + let schema = DataType::Struct { + fields: fields.clone(), + }; + let frame = self + .session + .create_dataframe_from_inline_arrow_ipc(ipc.extract()?, schema); + // Match pandas row input: explicit nonempty field names travel + // through toDF as well as the local relation's logical schema. + return Ok(PyDataFrame::new( + frame.to_df(fields.iter().map(|field| field.name.as_str()).collect()), + )); + } return self.create_from_pandas(py, data, spec); } self.create_from_local(py, data, spec) diff --git a/python/tests/test_arrow_handoff.py b/python/tests/test_arrow_handoff.py index 13243164..f6f056f7 100644 --- a/python/tests/test_arrow_handoff.py +++ b/python/tests/test_arrow_handoff.py @@ -2054,7 +2054,10 @@ def no_rows(*args, **kwargs): "spark.sql.session.localRelationChunkSizeRows": "997", } with _arrow_input_service(protocol, **overrides) as (spark, state): - spark.createDataFrame(frame, schema, verifySchema=verify).toArrow() + result = spark.createDataFrame(frame, schema, verifySchema=verify) + assert not result.isLocal() + result.toArrow() + assert state["column_names"] == schema.names table = state["input_tables"][0] assert table.num_rows == len(strings) assert table.schema.names == schema.names @@ -2646,7 +2649,10 @@ def test_string_large_inline_input_keeps_row_transport_and_schema(protocol): "spark.sql.session.localRelationSizeLimit": "1", "spark.sql.execution.arrow.useLargeVarTypes": "true", }) as (spark, state): - spark.createDataFrame(frame, schema).toArrow() + result = spark.createDataFrame(frame, schema) + assert not result.isLocal() + result.toArrow() + assert state["column_names"] == schema.names assert state["config_requests"] == [] assert state["uploads"] == 0 table = state["input_tables"][0]