From dc29a19e19c06df447d0b2a42be6169864adf81e Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Wed, 28 Oct 2026 21:26:55 -0700 Subject: [PATCH 01/47] [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/47] [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/47] [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/47] [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/47] [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/47] [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/47] [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/47] [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/47] [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/47] [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/47] [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/47] 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/47] [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/47] 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/47] 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/47] [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/47] [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/47] [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 7b40f3267cc415a13787fddf91a1fa8128abeb3b Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Wed, 28 Oct 2026 22:52:15 -0700 Subject: [PATCH 19/47] [SPARK-60685][CONNECT] Fast-path exact builtin date inputs --- crates/pyspark-rs/src/session.rs | 8 ++ python/tests/test_date_input.py | 225 +++++++++++++++++++++++++++++++ 2 files changed, 233 insertions(+) create mode 100644 python/tests/test_date_input.py diff --git a/crates/pyspark-rs/src/session.rs b/crates/pyspark-rs/src/session.rs index 4edbd2a9..943115e2 100644 --- a/crates/pyspark-rs/src/session.rs +++ b/crates/pyspark-rs/src/session.rs @@ -1314,6 +1314,14 @@ pub(crate) fn py_to_value(obj: &Bound<'_, PyAny>) -> PyResult { } } + // Exact dates cannot implement numeric protocols or override toordinal. Skip + // the generic Decimal/numeric probes; subclasses (including datetime) must + // retain the fallback ordering below. PyDate also supports abi3 wheels. + if let Ok(date) = obj.cast_exact::() { + let ordinal: i64 = date.call_method0("toordinal")?.extract()?; + return Ok(Value::Date((ordinal - 719163i64) as i32)); + } + // Check for bool (before int because bool is a subclass of int in Python) if let Ok(b) = obj.extract::() { return Ok(Value::Bool(b)); diff --git a/python/tests/test_date_input.py b/python/tests/test_date_input.py new file mode 100644 index 00000000..4a3e28d6 --- /dev/null +++ b/python/tests/test_date_input.py @@ -0,0 +1,225 @@ +# 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. + +"""Date conversion preserves the existing ordinal and subclass semantics.""" +from datetime import date, datetime, timezone +from decimal import Decimal + +import pytest + +from pyspark.sql import Row + + +@pytest.mark.parametrize("value", [ + date.min, date.max, date(1, 3, 1), date(4, 2, 29), + date(1600, 2, 29), date(1900, 2, 28), date(1900, 3, 1), + date(1969, 12, 31), date(1970, 1, 1), date(1970, 1, 2), + date(2000, 2, 29), date(2100, 3, 1), +]) +def test_exact_dates_roundtrip(value): + actual = Row(value=value).value + assert type(actual) is date + assert actual == value + + +def test_date_subclasses_keep_protocol_precedence(): + trace = [] + + class Indexed(date): + def __index__(self): + trace.append("index") + return 42 + + def toordinal(self): + raise AssertionError("numeric extraction must precede date conversion") + + class Floating(date): + def __float__(self): + trace.append("float") + return 1.25 + + def toordinal(self): + raise AssertionError("numeric extraction must precede date conversion") + + class Ordinal(date): + def toordinal(self): + trace.append("ordinal") + return date(2000, 2, 29).toordinal() + + class Plain(date): + pass + + assert Row(value=Indexed(2020, 1, 1)).value == 42 + assert Row(value=Floating(2020, 1, 1)).value == 1.25 + assert Row(value=Ordinal(2020, 1, 1)).value == date(2000, 2, 29) + assert type(Row(value=Plain(2020, 1, 1)).value) is date + assert trace == ["index", "float", "ordinal"] + + +def test_datetime_and_timestamp_keep_timestamp_encoding(): + pd = pytest.importorskip("pandas") + for value in [datetime(2020, 2, 29, 12, 34, 56, 123456), + datetime(2020, 2, 29, tzinfo=timezone.utc), + pd.Timestamp("2020-02-29 12:34:56.123456"), + pd.Timestamp("2020-02-29 12:34:56.123456Z")]: + actual = Row(value=value).value + assert type(actual) is datetime + expected = value.to_pydatetime() if isinstance(value, pd.Timestamp) else value + assert actual == datetime.fromtimestamp(expected.timestamp()) + + +def test_nested_dates_and_unrelated_values(): + value = Row(items=[date.min, None, date.max], + mapping={"epoch": date(1970, 1, 1)}, + nested=Row(day=date(2000, 2, 29)), + amount=Decimal("-12.340"), number=42, flag=True, label="date") + actual = Row(value=value).value + assert actual.items == [date.min, None, date.max] + assert actual.mapping == {"epoch": date(1970, 1, 1)} + assert actual.nested.day == date(2000, 2, 29) + assert actual.amount == Decimal("-12.340") + assert (actual.number, actual.flag, actual.label) == (42, True, "date") + + +def test_named_impostor_is_not_a_date(): + impostor = type("date", (), {"toordinal": lambda self: 719163}) + with pytest.raises(TypeError, match="unsupported value type"): + Row(value=impostor()) + + +def _fields(data): + """Decode only the small protobuf envelope around a local Arrow relation.""" + def varint(pos): + value, shift = 0, 0 + while True: + byte = data[pos] + pos += 1 + value |= (byte & 127) << shift + if byte < 128: + return value, pos + shift += 7 + + result, pos = {}, 0 + while pos < len(data): + tag, pos = varint(pos) + wire = tag & 7 + if wire == 2: + length, pos = varint(pos) + result[tag >> 3] = data[pos:pos + length] + pos += length + elif wire == 0: + _, pos = varint(pos) + elif wire in (1, 5): + pos += 8 if wire == 1 else 4 + else: + raise AssertionError(f"unexpected protobuf wire type: {wire}") + return result + + +def _capture_server(connection): + from concurrent.futures import ThreadPoolExecutor + import grpc + + captured = [] + + def execute(request, context): + # ExecutePlanRequest.plan=3 -> Plan.root=1 -> + # Relation.local_relation=11 -> LocalRelation.data=1. + captured.append(_fields(_fields(_fields(_fields(request)[3])[1])[11])[1]) + context.abort(grpc.StatusCode.INVALID_ARGUMENT, "captured date input") + yield b"" # Make this a streaming handler. + + with ThreadPoolExecutor(max_workers=2) as executor: + server = grpc.server(executor) + server.add_generic_rpc_handlers((grpc.method_handlers_generic_handler( + "spark.connect.SparkConnectService", { + "ExecutePlan": grpc.unary_stream_rpc_method_handler(execute), + }),)) + port = server.add_insecure_port("127.0.0.1:0") + server.start() + connection.send(port) + try: + while True: + command = connection.recv() + if command == "stop": + break + assert command == "get" + connection.send(captured.pop(0)) + finally: + server.stop(0).wait() + + +@pytest.fixture +def capture_input(): + import multiprocessing + pytest.importorskip("grpc") + from pyspark.sql import SparkSession + + context = multiprocessing.get_context("spawn") + parent, child = context.Pipe() + process = context.Process(target=_capture_server, args=(child,)) + process.start() + try: + assert parent.poll(20), "capture service failed to start" + spark = SparkSession.builder.remote(f"sc://127.0.0.1:{parent.recv()}").getOrCreate() + + def capture(data, schema): + import pyarrow as pa + df = spark.createDataFrame(data, schema=schema) + with pytest.raises(Exception, match="captured date input"): + df.collect() + parent.send("get") + assert parent.poll(10), "capture service failed to return the IPC payload" + return pa.ipc.open_stream(parent.recv()).read_all() + + yield capture + finally: + if process.is_alive(): + parent.send("stop") + process.join(10) + if process.is_alive(): + process.terminate() + process.join(10) + parent.close() + child.close() + + +@pytest.mark.parametrize("source", ["tuples", "pandas"]) +@pytest.mark.parametrize("schema_kind", ["inferred", "names", "explicit", "ddl"]) +def test_date_input_ipc_across_calendar_range(capture_input, source, schema_kind): + import calendar + import pyarrow as pa + from pyspark.sql import types as T + + # Every year's endpoints and leap day cover the complete supported calendar + # range, including non-leap centuries, without relying on the Rust decoder. + dates = [date.min, date.max, date(1969, 12, 31), date(1970, 1, 1), + date(1970, 1, 2), None] + for year in range(1, 10000): + dates.extend([date(year, 1, 1), date(year, 12, 31)]) + if calendar.isleap(year): + dates.append(date(year, 2, 29)) + schema = {"inferred": None, "names": ["day"], "ddl": "day date", + "explicit": T.StructType([T.StructField("day", T.DateType())])}[schema_kind] + data = [(value,) for value in dates] + if source == "pandas": + pd = pytest.importorskip("pandas") + data = pd.DataFrame(data, columns=["day"]) + table = capture_input(data, schema) + assert table.column(0).type == pa.date32() + assert table.column(0).to_pylist() == dates + epoch = date(1970, 1, 1) + expected_days = [None if value is None else (value - epoch).days for value in dates] + assert table.column(0).cast(pa.int32()).to_pylist() == expected_days From 4a5005840a505baa471bfa619a90c27890750dc3 Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Thu, 29 Oct 2026 04:13:11 -0700 Subject: [PATCH 20/47] [SPARK-60685][CONNECT] Include date tests in the CI offline suite --- python/tests/test_date_input.py | 225 ------------------------- python/tests/test_dropin_offline.py | 248 ++++++++++++++++++++++++++++ 2 files changed, 248 insertions(+), 225 deletions(-) delete mode 100644 python/tests/test_date_input.py diff --git a/python/tests/test_date_input.py b/python/tests/test_date_input.py deleted file mode 100644 index 4a3e28d6..00000000 --- a/python/tests/test_date_input.py +++ /dev/null @@ -1,225 +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. - -"""Date conversion preserves the existing ordinal and subclass semantics.""" -from datetime import date, datetime, timezone -from decimal import Decimal - -import pytest - -from pyspark.sql import Row - - -@pytest.mark.parametrize("value", [ - date.min, date.max, date(1, 3, 1), date(4, 2, 29), - date(1600, 2, 29), date(1900, 2, 28), date(1900, 3, 1), - date(1969, 12, 31), date(1970, 1, 1), date(1970, 1, 2), - date(2000, 2, 29), date(2100, 3, 1), -]) -def test_exact_dates_roundtrip(value): - actual = Row(value=value).value - assert type(actual) is date - assert actual == value - - -def test_date_subclasses_keep_protocol_precedence(): - trace = [] - - class Indexed(date): - def __index__(self): - trace.append("index") - return 42 - - def toordinal(self): - raise AssertionError("numeric extraction must precede date conversion") - - class Floating(date): - def __float__(self): - trace.append("float") - return 1.25 - - def toordinal(self): - raise AssertionError("numeric extraction must precede date conversion") - - class Ordinal(date): - def toordinal(self): - trace.append("ordinal") - return date(2000, 2, 29).toordinal() - - class Plain(date): - pass - - assert Row(value=Indexed(2020, 1, 1)).value == 42 - assert Row(value=Floating(2020, 1, 1)).value == 1.25 - assert Row(value=Ordinal(2020, 1, 1)).value == date(2000, 2, 29) - assert type(Row(value=Plain(2020, 1, 1)).value) is date - assert trace == ["index", "float", "ordinal"] - - -def test_datetime_and_timestamp_keep_timestamp_encoding(): - pd = pytest.importorskip("pandas") - for value in [datetime(2020, 2, 29, 12, 34, 56, 123456), - datetime(2020, 2, 29, tzinfo=timezone.utc), - pd.Timestamp("2020-02-29 12:34:56.123456"), - pd.Timestamp("2020-02-29 12:34:56.123456Z")]: - actual = Row(value=value).value - assert type(actual) is datetime - expected = value.to_pydatetime() if isinstance(value, pd.Timestamp) else value - assert actual == datetime.fromtimestamp(expected.timestamp()) - - -def test_nested_dates_and_unrelated_values(): - value = Row(items=[date.min, None, date.max], - mapping={"epoch": date(1970, 1, 1)}, - nested=Row(day=date(2000, 2, 29)), - amount=Decimal("-12.340"), number=42, flag=True, label="date") - actual = Row(value=value).value - assert actual.items == [date.min, None, date.max] - assert actual.mapping == {"epoch": date(1970, 1, 1)} - assert actual.nested.day == date(2000, 2, 29) - assert actual.amount == Decimal("-12.340") - assert (actual.number, actual.flag, actual.label) == (42, True, "date") - - -def test_named_impostor_is_not_a_date(): - impostor = type("date", (), {"toordinal": lambda self: 719163}) - with pytest.raises(TypeError, match="unsupported value type"): - Row(value=impostor()) - - -def _fields(data): - """Decode only the small protobuf envelope around a local Arrow relation.""" - def varint(pos): - value, shift = 0, 0 - while True: - byte = data[pos] - pos += 1 - value |= (byte & 127) << shift - if byte < 128: - return value, pos - shift += 7 - - result, pos = {}, 0 - while pos < len(data): - tag, pos = varint(pos) - wire = tag & 7 - if wire == 2: - length, pos = varint(pos) - result[tag >> 3] = data[pos:pos + length] - pos += length - elif wire == 0: - _, pos = varint(pos) - elif wire in (1, 5): - pos += 8 if wire == 1 else 4 - else: - raise AssertionError(f"unexpected protobuf wire type: {wire}") - return result - - -def _capture_server(connection): - from concurrent.futures import ThreadPoolExecutor - import grpc - - captured = [] - - def execute(request, context): - # ExecutePlanRequest.plan=3 -> Plan.root=1 -> - # Relation.local_relation=11 -> LocalRelation.data=1. - captured.append(_fields(_fields(_fields(_fields(request)[3])[1])[11])[1]) - context.abort(grpc.StatusCode.INVALID_ARGUMENT, "captured date input") - yield b"" # Make this a streaming handler. - - with ThreadPoolExecutor(max_workers=2) as executor: - server = grpc.server(executor) - server.add_generic_rpc_handlers((grpc.method_handlers_generic_handler( - "spark.connect.SparkConnectService", { - "ExecutePlan": grpc.unary_stream_rpc_method_handler(execute), - }),)) - port = server.add_insecure_port("127.0.0.1:0") - server.start() - connection.send(port) - try: - while True: - command = connection.recv() - if command == "stop": - break - assert command == "get" - connection.send(captured.pop(0)) - finally: - server.stop(0).wait() - - -@pytest.fixture -def capture_input(): - import multiprocessing - pytest.importorskip("grpc") - from pyspark.sql import SparkSession - - context = multiprocessing.get_context("spawn") - parent, child = context.Pipe() - process = context.Process(target=_capture_server, args=(child,)) - process.start() - try: - assert parent.poll(20), "capture service failed to start" - spark = SparkSession.builder.remote(f"sc://127.0.0.1:{parent.recv()}").getOrCreate() - - def capture(data, schema): - import pyarrow as pa - df = spark.createDataFrame(data, schema=schema) - with pytest.raises(Exception, match="captured date input"): - df.collect() - parent.send("get") - assert parent.poll(10), "capture service failed to return the IPC payload" - return pa.ipc.open_stream(parent.recv()).read_all() - - yield capture - finally: - if process.is_alive(): - parent.send("stop") - process.join(10) - if process.is_alive(): - process.terminate() - process.join(10) - parent.close() - child.close() - - -@pytest.mark.parametrize("source", ["tuples", "pandas"]) -@pytest.mark.parametrize("schema_kind", ["inferred", "names", "explicit", "ddl"]) -def test_date_input_ipc_across_calendar_range(capture_input, source, schema_kind): - import calendar - import pyarrow as pa - from pyspark.sql import types as T - - # Every year's endpoints and leap day cover the complete supported calendar - # range, including non-leap centuries, without relying on the Rust decoder. - dates = [date.min, date.max, date(1969, 12, 31), date(1970, 1, 1), - date(1970, 1, 2), None] - for year in range(1, 10000): - dates.extend([date(year, 1, 1), date(year, 12, 31)]) - if calendar.isleap(year): - dates.append(date(year, 2, 29)) - schema = {"inferred": None, "names": ["day"], "ddl": "day date", - "explicit": T.StructType([T.StructField("day", T.DateType())])}[schema_kind] - data = [(value,) for value in dates] - if source == "pandas": - pd = pytest.importorskip("pandas") - data = pd.DataFrame(data, columns=["day"]) - table = capture_input(data, schema) - assert table.column(0).type == pa.date32() - assert table.column(0).to_pylist() == dates - epoch = date(1970, 1, 1) - expected_days = [None if value is None else (value - epoch).days for value in dates] - assert table.column(0).cast(pa.int32()).to_pylist() == expected_days diff --git a/python/tests/test_dropin_offline.py b/python/tests/test_dropin_offline.py index 39e06382..d867c040 100644 --- a/python/tests/test_dropin_offline.py +++ b/python/tests/test_dropin_offline.py @@ -4900,6 +4900,254 @@ def __index__(self): "data_type1": "StructType", "data_type2": data_type} +# --------------------------------------------------------------------------- exact date input + +from datetime import ( + date as _exact_date_date, datetime as _exact_date_datetime, + timezone as _exact_date_timezone, +) +from decimal import Decimal as _exact_date_Decimal +from pyspark.sql import Row as _exact_date_Row + + +def _exact_date_converted_value(mock_connect, value): + import pyarrow as pa + + spark, service = mock_connect + spark.createDataFrame([(value,)]).toArrow() + request = parse_message(service.execute_requests[-1]) + relation = parse_message(parse_message(request[3][0])[1][0]) + local = parse_message(relation[11][0]) + return pa.ipc.open_stream(local[1][0]).read_all().column(0)[0].as_py() + + +@pytest.mark.parametrize("value", [ + _exact_date_date.min, _exact_date_date.max, _exact_date_date(1, 3, 1), _exact_date_date(4, 2, 29), + _exact_date_date(1600, 2, 29), _exact_date_date(1900, 2, 28), _exact_date_date(1900, 3, 1), + _exact_date_date(1969, 12, 31), _exact_date_date(1970, 1, 1), _exact_date_date(1970, 1, 2), + _exact_date_date(2000, 2, 29), _exact_date_date(2100, 3, 1), +]) +def test_exact_date_builtin_roundtrip(mock_connect, value): + actual = _exact_date_converted_value(mock_connect, value) + assert type(actual) is _exact_date_date + assert actual == value + + +def test_exact_date_subclass_protocol_precedence(mock_connect): + trace = [] + + class Indexed(_exact_date_date): + def __index__(self): + trace.append("index") + return 42 + + def toordinal(self): + raise AssertionError("numeric extraction must precede date conversion") + + class Floating(_exact_date_date): + def __float__(self): + trace.append("float") + return 1.25 + + def toordinal(self): + raise AssertionError("numeric extraction must precede date conversion") + + class Ordinal(_exact_date_date): + def toordinal(self): + trace.append("ordinal") + return _exact_date_date(2000, 2, 29).toordinal() + + class Plain(_exact_date_date): + pass + + assert _exact_date_converted_value(mock_connect, Indexed(2020, 1, 1)) == 42 + assert _exact_date_converted_value(mock_connect, Floating(2020, 1, 1)) == 1.25 + assert _exact_date_converted_value(mock_connect, Ordinal(2020, 1, 1)) == _exact_date_date(2000, 2, 29) + assert type(_exact_date_converted_value(mock_connect, Plain(2020, 1, 1))) is _exact_date_date + assert trace == ["index", "float", "ordinal"] + + +def test_exact_date_datetime_timestamp_encoding(mock_connect): + pd = pytest.importorskip("pandas") + for value in [_exact_date_datetime(2020, 2, 29, 12, 34, 56, 123456), + _exact_date_datetime(2020, 2, 29, tzinfo=_exact_date_timezone.utc), + pd.Timestamp("2020-02-29 12:34:56.123456"), + pd.Timestamp("2020-02-29 12:34:56.123456Z")]: + actual = _exact_date_converted_value(mock_connect, value) + assert type(actual) is _exact_date_datetime + expected = value.to_pydatetime() if isinstance(value, pd.Timestamp) else value + assert actual == _exact_date_datetime.fromtimestamp(expected.timestamp(), _exact_date_timezone.utc) + + +def test_exact_date_nested_unrelated_values(mock_connect): + value = _exact_date_Row(items=[_exact_date_date.min, None, _exact_date_date.max], + mapping={"epoch": _exact_date_date(1970, 1, 1)}, + nested=_exact_date_Row(day=_exact_date_date(2000, 2, 29)), + amount=_exact_date_Decimal("-12.340"), number=42, flag=True, label="date") + actual = _exact_date_converted_value(mock_connect, value) + assert actual["items"] == [_exact_date_date.min, None, _exact_date_date.max] + assert dict(actual["mapping"]) == {"epoch": _exact_date_date(1970, 1, 1)} + assert actual["nested"]["day"] == _exact_date_date(2000, 2, 29) + assert actual["amount"] == _exact_date_Decimal("-12.340") + assert (actual["number"], actual["flag"], actual["label"]) == (42, True, "date") + + +def test_exact_date_named_impostor_rejected(): + impostor = type("date", (), {"toordinal": lambda self: 719163}) + with pytest.raises(TypeError, match="unsupported value type"): + _local_create_dataframe([(impostor(),)]) + + +def _exact_date_fields(data): + """Decode only the small protobuf envelope around a local Arrow relation.""" + def _exact_date_varint(pos): + value, shift = 0, 0 + while True: + byte = data[pos] + pos += 1 + value |= (byte & 127) << shift + if byte < 128: + return value, pos + shift += 7 + + result, pos = {}, 0 + while pos < len(data): + tag, pos = _exact_date_varint(pos) + wire = tag & 7 + if wire == 2: + length, pos = _exact_date_varint(pos) + result[tag >> 3] = data[pos:pos + length] + pos += length + elif wire == 0: + _, pos = _exact_date_varint(pos) + elif wire in (1, 5): + pos += 8 if wire == 1 else 4 + else: + raise AssertionError(f"unexpected protobuf wire type: {wire}") + return result + + +def _exact_date_capture_server(connection): + from concurrent.futures import ThreadPoolExecutor + import grpc + + captured = [] + + def _exact_date_execute(request, context): + # ExecutePlanRequest.plan=3 -> Plan.root=1 -> + # Relation.local_relation=11 -> LocalRelation.data=1. + relation = _exact_date_fields(_exact_date_fields(_exact_date_fields(request)[3])[1]) + while 18 in relation: # Relation.to_df -> ToDF.input. + relation = _exact_date_fields(_exact_date_fields(relation[18])[1]) + captured.append(_exact_date_fields(relation[11])[1]) + context.abort(grpc.StatusCode.INVALID_ARGUMENT, "captured date input") + yield b"" # Make this a streaming handler. + + def _exact_date_config(request, context): + from google.protobuf.internal.encoder import _VarintBytes + + def field(tag, value): + return bytes([tag]) + _VarintBytes(len(value)) + value + + values = { + "spark.sql.timestampType": "TIMESTAMP_LTZ", + "spark.sql.session.timeZone": "UTC", + "spark.sql.execution.arrow.useLargeVarTypes": "false", + "spark.sql.session.localRelationCacheThreshold": "1048576", + "spark.sql.session.localRelationSizeLimit": str(3 * 1024**3), + "spark.sql.session.localRelationChunkSizeRows": "10000", + "spark.sql.session.localRelationChunkSizeBytes": str(16 * 1024**2), + "spark.sql.session.localRelationBatchOfChunksSizeBytes": str(1024**3), + } + # ConfigResponse.pairs=2; KeyValue.key=1, value=2. + return b"".join(field(18, field(10, key.encode()) + field(18, value.encode())) + for key, value in values.items()) + + with ThreadPoolExecutor(max_workers=2) as executor: + server = grpc.server(executor) + server.add_generic_rpc_handlers((grpc.method_handlers_generic_handler( + "spark.connect.SparkConnectService", { + "ExecutePlan": grpc.unary_stream_rpc_method_handler(_exact_date_execute), + "Config": grpc.unary_unary_rpc_method_handler(_exact_date_config), + }),)) + port = server.add_insecure_port("127.0.0.1:0") + server.start() + connection.send(port) + try: + while True: + command = connection.recv() + if command == "stop": + break + assert command == "get" + connection.send(captured.pop(0)) + finally: + server.stop(0).wait() + + +@pytest.fixture +def exact_date_capture_input(): + import multiprocessing + pytest.importorskip("grpc") + from pyspark.sql import SparkSession + + context = multiprocessing.get_context("spawn") + parent, child = context.Pipe() + process = context.Process(target=_exact_date_capture_server, args=(child,)) + process.start() + try: + assert parent.poll(20), "capture service failed to start" + spark = SparkSession.builder.remote(f"sc://127.0.0.1:{parent.recv()}").getOrCreate() + + def _exact_date_capture(data, schema): + import pyarrow as pa + df = spark.createDataFrame(data, schema=schema) + with pytest.raises(Exception, match="captured date input"): + df.collect() + parent.send("get") + assert parent.poll(10), "capture service failed to return the IPC payload" + return pa.ipc.open_stream(parent.recv()).read_all() + + yield _exact_date_capture + finally: + if process.is_alive(): + parent.send("stop") + process.join(10) + if process.is_alive(): + process.terminate() + process.join(10) + parent.close() + child.close() + + +@pytest.mark.parametrize("source", ["tuples", "pandas"]) +@pytest.mark.parametrize("schema_kind", ["inferred", "names", "explicit", "ddl"]) +def test_exact_date_ipc_calendar_range(exact_date_capture_input, source, schema_kind): + import calendar + import pyarrow as pa + from pyspark.sql import types as T + + # Every year's endpoints and leap day cover the complete supported calendar + # range, including non-leap centuries, without relying on the Rust decoder. + dates = [_exact_date_date.min, _exact_date_date.max, _exact_date_date(1969, 12, 31), _exact_date_date(1970, 1, 1), + _exact_date_date(1970, 1, 2), None] + for year in range(1, 10000): + dates.extend([_exact_date_date(year, 1, 1), _exact_date_date(year, 12, 31)]) + if calendar.isleap(year): + dates.append(_exact_date_date(year, 2, 29)) + schema = {"inferred": None, "names": ["day"], "ddl": "day date", + "explicit": T.StructType([T.StructField("day", T.DateType())])}[schema_kind] + data = [(value,) for value in dates] + if source == "pandas": + pd = pytest.importorskip("pandas") + data = pd.DataFrame(data, columns=["day"]) + table = exact_date_capture_input(data, schema) + assert table.column(0).type == pa.date32() + assert table.column(0).to_pylist() == dates + epoch = _exact_date_date(1970, 1, 1) + expected_days = [None if value is None else (value - epoch).days for value in dates] + assert table.column(0).cast(pa.int32()).to_pylist() == expected_days + + # --------------------------------------------------------------------------- printSchema From fc7373c3f3791f1d4873b08f9582e2a9242ee92c Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Thu, 29 Oct 2026 01:33:49 -0700 Subject: [PATCH 21/47] [SPARK-60698][CONNECT] Batch Arrow input normalization config reads Read the three mutable Arrow normalization settings in one Get request per nonempty construction while retaining empty-input and transport behavior. Co-authored-by: Isaac --- python/tests/test_arrow_handoff.py | 135 +++++++++++++++++++++++++++-- 1 file changed, 130 insertions(+), 5 deletions(-) diff --git a/python/tests/test_arrow_handoff.py b/python/tests/test_arrow_handoff.py index 5d098808..40faf284 100644 --- a/python/tests/test_arrow_handoff.py +++ b/python/tests/test_arrow_handoff.py @@ -427,7 +427,7 @@ def _arrow_input_service( } config.update(overrides) state = {"cache": {}, "uploads": 0, "relations": [], "message_sizes": [], - "config_requests": [], "input_tables": [], + "config": config, "config_requests": [], "input_tables": [], "upload_contexts": [], "status_contexts": []} server_id = "00000000-0000-0000-0000-000000000001" @@ -448,12 +448,15 @@ def attribution(request, context): def configure(raw, context): request = protocol("ConfigRequest").FromString(raw) operation = request.operation.WhichOneof("op_type") - keys = getattr(request.operation, operation).keys - state["config_requests"].append((operation, list(keys))) + keys = list(getattr(request.operation, operation).keys) + state["config_requests"].append((operation, keys)) if config_error: context.abort(grpc.StatusCode.UNAVAILABLE, "configuration unavailable") + if error := state.pop("config_error", None): + context.abort(error, "configuration read failed") + # Deliberately reverse responses and omit missing keys to test key mapping. return protocol("ConfigResponse", session_id=request.session_id, - pairs=[dict(key=k, value=config[k]) for k in keys + pairs=[dict(key=k, value=config[k]) for k in reversed(keys) if config.get(k) is not None] ).SerializeToString() @@ -548,7 +551,8 @@ def execute(raw, context): 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)) + state["input_tables"].append(pa.Table.from_batches( + batches, schema=state["arrow_schemas"][0])) result = pa.table({"rows": [rows]}) sink = pa.BufferOutputStream() with pa.ipc.new_stream(sink, result.schema) as writer: @@ -1725,3 +1729,124 @@ def set_data(index, value): "metadata": metadata} ], } + + +_NORMALIZATION_KEYS = [ + "spark.sql.timestampType", + "spark.sql.session.timeZone", + "spark.sql.execution.arrow.useLargeVarTypes", +] +_UPLOAD_KEYS = [ + "spark.sql.session.localRelationCacheThreshold", + "spark.sql.session.localRelationSizeLimit", + "spark.sql.session.localRelationChunkSizeRows", + "spark.sql.session.localRelationChunkSizeBytes", + "spark.sql.session.localRelationBatchOfChunksSizeBytes", +] + + +@pytest.mark.parametrize("timestamp_type", ["TIMESTAMP_LTZ", "TIMESTAMP_NTZ"]) +@pytest.mark.parametrize("zone", ["UTC", "America/Los_Angeles"]) +@pytest.mark.parametrize("large", ["false", "true"]) +def test_arrow_input_batched_normalization_configs(protocol, timestamp_type, zone, large): + from zoneinfo import ZoneInfo + + value = datetime(2024, 1, 15, 12, 34, 56, 123456) + table = pa.table({"t": [value, None], "s": ["hello", None], "b": [b"xyz", None]}) + with _arrow_input_service(protocol, **dict(zip( + _NORMALIZATION_KEYS, [timestamp_type, zone, large] + ))) as (spark, state): + frame = spark.createDataFrame(table) + assert state["config_requests"] == [("get", _NORMALIZATION_KEYS + _UPLOAD_KEYS)] + frame.collect() + expected_time = (value if timestamp_type == "TIMESTAMP_NTZ" else + value.replace(tzinfo=ZoneInfo(zone)).astimezone(timezone.utc)) + expected = pa.table({ + "t": pa.array([expected_time, None], type=pa.timestamp( + "us", tz=None if timestamp_type == "TIMESTAMP_NTZ" else "UTC")), + "s": pa.array(["hello", None], type=pa.large_string() if large == "true" + else pa.string()), + "b": pa.array([b"xyz", None], type=pa.large_binary() if large == "true" + else pa.binary()), + }) + actual = state["input_tables"][0] + assert actual.schema == expected.schema + assert actual.to_pydict() == expected.to_pydict() + # Inferred input keeps its types in Arrow; names are applied via toDF. + assert not state["relations"][0].local_relation.HasField("schema") + + +@pytest.mark.parametrize("missing", _NORMALIZATION_KEYS) +@pytest.mark.parametrize("present_null", [False, True]) +def test_arrow_input_missing_normalization_config(protocol, missing, present_null): + with _arrow_input_service(protocol) as (spark, state): + if present_null: + state["config"][missing] = None + else: + del state["config"][missing] + # The current Arrow settings snapshot requires every server setting. + with pytest.raises(Exception, match="SQL_CONF_NOT_FOUND"): + spark.createDataFrame(pa.table({"x": [1, None], "s": ["hello", None]})) + assert state["config_requests"] == [("get", _NORMALIZATION_KEYS + _UPLOAD_KEYS)] + assert state["uploads"] == 0 + assert state["input_tables"] == [] + + +def test_arrow_input_normalization_config_failure_then_fresh_read(protocol): + with _arrow_input_service(protocol) as (spark, state): + state["config_error"] = grpc.StatusCode.INVALID_ARGUMENT + table = pa.table({"x": [1]}) + with pytest.raises(Exception, match="configuration read failed"): + spark.createDataFrame(table) + assert state["config_requests"] == [("get", _NORMALIZATION_KEYS + _UPLOAD_KEYS)] + assert state["uploads"] == 0 + spark.createDataFrame(table).collect() + assert state["config_requests"] == [ + ("get", _NORMALIZATION_KEYS + _UPLOAD_KEYS), + ("get", _NORMALIZATION_KEYS + _UPLOAD_KEYS)] + assert state["input_tables"][0].to_pydict() == {"x": [1]} + + +def test_arrow_input_normalization_config_transient_retry(protocol): + with _arrow_input_service(protocol) as (spark, state): + state["config_error"] = grpc.StatusCode.UNAVAILABLE + spark.createDataFrame(pa.table({"x": [1]})).collect() + assert state["config_requests"] == [ + ("get", _NORMALIZATION_KEYS + _UPLOAD_KEYS), + ("get", _NORMALIZATION_KEYS + _UPLOAD_KEYS)] + assert state["input_tables"][0].to_pydict() == {"x": [1]} + + +def test_arrow_input_normalization_configs_are_fresh(protocol): + + value = datetime(2024, 1, 15, 12) + table = pa.table({"t": [value], "s": ["hello"]}) + with _arrow_input_service(protocol) as (spark, state): + spark.createDataFrame(table).collect() + state["config"].update(dict(zip(_NORMALIZATION_KEYS, [ + "TIMESTAMP_NTZ", "America/Los_Angeles", "true"]))) + spark.createDataFrame(table).collect() + state["config"].update(dict(zip(_NORMALIZATION_KEYS, [ + "TIMESTAMP_LTZ", "America/Los_Angeles", "false"]))) + spark.createDataFrame(table).collect() + assert state["config_requests"] == [("get", _NORMALIZATION_KEYS + _UPLOAD_KEYS)] * 3 + tables = state["input_tables"] + assert [t.column("t")[0].as_py() for t in tables] == [ + value.replace(tzinfo=timezone.utc), value, + datetime(2024, 1, 15, 20, tzinfo=timezone.utc)] + assert [t.schema.field("s").type for t in tables] == [ + pa.string(), pa.large_string(), pa.string()] + assert all(not relation.local_relation.HasField("schema") + for relation in state["relations"]) + + +def test_arrow_input_empty_skips_all_configs(protocol): + from pyspark.sql.types import LongType, StructField, StructType + + with _arrow_input_service(protocol) as (spark, state): + table = pa.table({"x": pa.array([], type=pa.int64())}) + spark.createDataFrame(table, StructType([StructField("x", LongType())])) + with pytest.raises(Exception, match="CANNOT_INFER_EMPTY_SCHEMA"): + spark.createDataFrame(table) + assert state["config_requests"] == [] + assert state["uploads"] == 0 From 1fc2baba02dd9b06bcd97ea87cb0edd8840c3321 Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Thu, 29 Oct 2026 09:57:39 -0700 Subject: [PATCH 22/47] [SPARK-60716][CONNECT] Convert explicit pandas date columns through Arrow Conservatively accept ordinary object-column frames with exact built-in dates and explicit DateType schemas. Preserve native row conversion for customized schema descriptors and unsupported values. Co-authored-by: Isaac --- python/pyspark/sql/_pandas_input.py | 65 +++++++++ python/tests/test_arrow_handoff.py | 218 ++++++++++++++++++++++++++++ 2 files changed, 283 insertions(+) diff --git a/python/pyspark/sql/_pandas_input.py b/python/pyspark/sql/_pandas_input.py index 614287d8..4141f0d5 100644 --- a/python/pyspark/sql/_pandas_input.py +++ b/python/pyspark/sql/_pandas_input.py @@ -274,3 +274,68 @@ def prepare_int64_table(frame, schema): [pa.array(column, type=pa.int64(), from_pandas=True) for column in columns], names=names, ) + + +def prepare_date_table(frame, schema): + """Accept only explicit date columns whose cells are exact dates or None. + + Inspect every object before Arrow conversion: date 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. + """ + from datetime import date + + import numpy as np + import pandas as pd + from types import GetSetDescriptorType, WrapperDescriptorType + + from pyspark._pyspark import DateType, StructField, StructType as NativeStructType + from pyspark.sql import types as sql_types + + if type(frame) is not pd.DataFrame or type(schema) is not NativeStructType or len(frame) < 1000: + 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 (("DateType", DateType), ("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), + ): + 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 + 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 DateType + 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(dtype == np.dtype("object") for dtype in frame.dtypes): + return None + columns = [frame.iloc[:, i] for i in range(width)] + for column, field in zip(columns, schema.fields): + # Arrow rejects nulls in non-nullable fields; let the row path retain + # its existing behavior for that input instead. + if not all(type(value) is date or (value is None and field.nullable) + for value in column.to_numpy(copy=False)): + return None + + import pyarrow as pa + + return pa.Table.from_arrays( + [pa.array(column, type=pa.date32(), from_pandas=False) for column in columns], + names=schema.names, + ) diff --git a/python/tests/test_arrow_handoff.py b/python/tests/test_arrow_handoff.py index 40faf284..3d4eef6e 100644 --- a/python/tests/test_arrow_handoff.py +++ b/python/tests/test_arrow_handoff.py @@ -1850,3 +1850,221 @@ def test_arrow_input_empty_skips_all_configs(protocol): spark.createDataFrame(table) assert state["config_requests"] == [] assert state["uploads"] == 0 + + +# Exact built-in date columns reuse Arrow without materializing pandas rows. + +@pytest.mark.parametrize("cached", [False, True]) +@pytest.mark.parametrize("verify", [False, True]) +def test_date_columns_preserve_every_ipc_cell(protocol, monkeypatch, cached, verify): + from pyspark.sql.pandas.types import from_arrow_metadata + + dates = [date.min, date.max, date(4, 2, 29), date(1600, 2, 29), + date(1900, 2, 28), date(1900, 3, 1), date(1969, 12, 31), + date(1970, 1, 1), date(2000, 2, 29), date(2100, 3, 1)] * 100 + nullable = [None if i % 7 == 0 else value for i, value in enumerate(reversed(dates))] + expected = [dates, nullable, [None] * len(dates)] + frame = pd.DataFrame(dict(zip(["original", "b", "c"], expected)), dtype=object) + schema = T.StructType([ + T.StructField("renamed.`day", T.DateType(), False, {"source": "date", "id": 3}), + T.StructField("nullable", T.DateType(), True, {"tag": [1, 2]}), + T.StructField("all_null", T.DateType()), + ]) + + def no_rows(*args, **kwargs): + raise AssertionError("date fast path must not materialize Python rows") + + for method in ("astype", "notna", "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 "9999999"} + with _arrow_input_service(protocol, **overrides) as (spark, state): + df = spark.createDataFrame(frame, schema, verifySchema=verify) + df.toArrow() + table = state["input_tables"][0] + assert table.num_rows == len(dates) + assert table.schema.names == schema.names + epoch = date(1970, 1, 1) + for column, values, field in zip(table.columns, expected, schema.fields): + assert column.type == pa.date32() + # Independent expectations cover every date and validity bit in IPC. + assert column.cast(pa.int32()).to_pylist() == [ + None if value is None else (value - epoch).days for value in values + ] + 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) == (field.metadata or None) + assert bool(state["uploads"]) == cached + + +@pytest.mark.parametrize("reason", [ + "small", "empty", "inferred", "names", "ddl", "cast", "mixed_schema", + "width", "duplicate_input", "duplicate_schema", "input_nul", "schema_nul", + "input_number", "datetime", "timestamp", "nan", "pd_na", "nat", "mixed", + "nonnullable_null", "frame_subclass", "date_subclass", "late_unsupported", +]) +def test_date_column_guard_preserves_row_fallback(reason): + from pyspark.sql._pandas_input import prepare_date_table + + class Frame(pd.DataFrame): + pass + + class Day(date): + def toordinal(self): + raise AssertionError("guard must not call conversion hooks") + + values = [date(2020, 1, 1)] * 1000 + frame = pd.DataFrame({"day": values}, dtype=object) + schema = T.StructType([T.StructField("day", T.DateType())]) + if reason == "small": + frame = frame.iloc[:999] + elif reason == "empty": + frame = frame.iloc[:0] + elif reason in ("inferred", "names", "ddl", "cast"): + schema = {"inferred": None, "names": ["day"], "ddl": "day DATE", + "cast": T.StructType([T.StructField("day", T.StringType())])}[reason] + elif reason == "mixed_schema": + frame["long"] = np.arange(1000) + schema.add("long", T.LongType()) + elif reason == "width": + schema.add("extra", T.DateType()) + elif reason in ("duplicate_input", "duplicate_schema"): + frame["other"] = values + schema.add("other", T.DateType()) + if reason == "duplicate_input": + frame.columns = ["day", "day"] + else: + schema = T.StructType([T.StructField("day", T.DateType())] * 2) + elif reason in ("input_nul", "input_number"): + frame.columns = ["day\x00" if reason == "input_nul" else 1] + elif reason == "schema_nul": + schema = T.StructType([T.StructField("day\x00", T.DateType())]) + elif reason == "frame_subclass": + frame = Frame(frame) + else: + value = { + "datetime": datetime(2020, 1, 1), "timestamp": pd.Timestamp("2020-01-01"), + "nan": float("nan"), "pd_na": pd.NA, "nat": pd.NaT, "mixed": "2020-01-01", + "nonnullable_null": None, "date_subclass": Day(2020, 1, 1), + "late_unsupported": Day(2020, 1, 1), + }[reason] + if reason == "nonnullable_null": + schema = T.StructType([T.StructField("day", T.DateType(), False)]) + if reason == "late_unsupported": + frame["other"] = values + schema.add("other", T.DateType()) + frame.iat[-1, -1] = value + assert prepare_date_table(frame, schema) is None + + +@pytest.mark.parametrize("reason", ["date_subclass", "frame_subclass", "nan", "pd_na", "nat", + "nonnullable_null", "inferred", "ddl", "small"]) +def test_date_public_fallback_keeps_hooks(protocol, monkeypatch, reason): + calls = [] + + class Day(date): + def toordinal(self): + calls.append("toordinal") + return super().toordinal() + 7 + + class Frame(pd.DataFrame): + pass + + values = [date(2020, 1, 1)] * (999 if reason == "small" else 1000) + if reason in ("nan", "pd_na", "nat", "nonnullable_null"): + values[-1] = {"nan": float("nan"), "pd_na": pd.NA, "nat": pd.NaT, + "nonnullable_null": None}[reason] + elif reason == "date_subclass": + values[-1] = Day(2020, 1, 1) + frame_type = Frame if reason == "frame_subclass" else pd.DataFrame + frame = frame_type({"day": values}, dtype=object) + schema = T.StructType([T.StructField("day", T.DateType(), reason != "nonnullable_null")]) + if reason == "inferred": + schema = None + elif reason == "ddl": + schema = "day DATE" + 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 == "nonnullable_null": + from pyspark.errors import PySparkRuntimeError + + with pytest.raises(PySparkRuntimeError, match=( + "Column 'day' is declared as non-nullable but contains null values" + )): + spark.createDataFrame(frame, schema) + assert calls == ["astype"] + assert state["config_requests"] == [] + return + spark.createDataFrame(frame, schema).toArrow() + assert "astype" in calls + actual = state["input_tables"][0].column(0).to_pylist() + expected = [date(2020, 1, 1)] * len(values) + if reason in ("nan", "pd_na", "nat", "nonnullable_null"): + expected[-1] = None + elif reason == "date_subclass": + assert calls.count("toordinal") == 1 + expected[-1] = date(2020, 1, 8) + assert actual == expected + + +@pytest.mark.parametrize("class_name,attribute", [ + ("StructType", "names"), ("StructType", "fields"), ("StructType", "__iter__"), + ("StructField", "name"), ("StructField", "dataType"), ("StructField", "nullable"), + ("StructField", "metadata"), + ("StructType", "__getattribute__"), ("StructField", "__getattribute__"), + ("DateType", "__getattribute__"), ("StructType", "__class__"), + ("StructField", "__class__"), ("DateType", "__class__"), +]) +@pytest.mark.parametrize("before_import", [False, True]) +def test_date_custom_schema_descriptors_keep_native_row_path(class_name, attribute, before_import): + import sys + + script = r''' +import sys +import tempfile +from datetime import date +from pathlib import Path +import pandas as pd +from pyspark.sql import types as T + +frame = pd.DataFrame({"day": [date(2020, 1, 1)] * 1000}, dtype=object) +schema = T.StructType([T.StructField("day", T.DateType(), True, {"source": "native"})]) +class_name, attribute, before_import = sys.argv[1:] +if before_import == "False": + import pyspark.sql._pandas_input +cls = getattr(T, class_name) +if attribute == "__getattribute__": + replacement = lambda self, name: object.__getattribute__(self, name) +elif attribute == "__iter__": + replacement = lambda self: iter([T.StructField("changed", T.DateType())]) +elif attribute == "__class__": + attribute = "__getattribute__" + replacement = lambda self, name: object if name == "__class__" else ( + object.__getattribute__(self, name)) +else: + value = {"names": ["changed"], "fields": [], "name": "changed", + "dataType": T.DateType(), "nullable": False, "metadata": {"changed": True}}[attribute] + replacement = property(lambda self: value) +setattr(cls, attribute, replacement) +from tests.test_arrow_handoff import _arrow_input_service, protocol +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 == ["day"], table.schema + assert table.column(0).to_pylist() == [date(2020, 1, 1)] * 1000 +''' + result = subprocess.run([sys.executable, "-c", script, class_name, attribute, + str(before_import)], capture_output=True, text=True, timeout=30) + assert result.returncode == 0, result.stdout + result.stderr From 5727e1860559f6d795a790a58862e8a740dbb936 Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Thu, 29 Oct 2026 10:11:28 -0700 Subject: [PATCH 23/47] [SPARK-60716][CONNECT] Keep small date frames on the row path Use the measured 10000-row threshold to avoid config RPC overhead on smaller inputs. Cover the fast-path boundary and load fresh-process test fixtures by exact source path for CI portability. Co-authored-by: Isaac --- python/pyspark/sql/_pandas_input.py | 4 +++- python/tests/test_arrow_handoff.py | 28 +++++++++++++++++----------- 2 files changed, 20 insertions(+), 12 deletions(-) diff --git a/python/pyspark/sql/_pandas_input.py b/python/pyspark/sql/_pandas_input.py index 4141f0d5..bebd739c 100644 --- a/python/pyspark/sql/_pandas_input.py +++ b/python/pyspark/sql/_pandas_input.py @@ -292,7 +292,9 @@ def prepare_date_table(frame, schema): from pyspark._pyspark import DateType, StructField, StructType as NativeStructType from pyspark.sql import types as sql_types - if type(frame) is not pd.DataFrame or type(schema) is not NativeStructType or len(frame) < 1000: + # Keep smaller frames on the row path to avoid the importer's config RPC overhead. + if (type(frame) is not pd.DataFrame or type(schema) is not NativeStructType + or len(frame) < 10000): 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. diff --git a/python/tests/test_arrow_handoff.py b/python/tests/test_arrow_handoff.py index 3d4eef6e..57041969 100644 --- a/python/tests/test_arrow_handoff.py +++ b/python/tests/test_arrow_handoff.py @@ -1861,7 +1861,7 @@ def test_date_columns_preserve_every_ipc_cell(protocol, monkeypatch, cached, ver dates = [date.min, date.max, date(4, 2, 29), date(1600, 2, 29), date(1900, 2, 28), date(1900, 3, 1), date(1969, 12, 31), - date(1970, 1, 1), date(2000, 2, 29), date(2100, 3, 1)] * 100 + date(1970, 1, 1), date(2000, 2, 29), date(2100, 3, 1)] * 1000 nullable = [None if i % 7 == 0 else value for i, value in enumerate(reversed(dates))] expected = [dates, nullable, [None] * len(dates)] frame = pd.DataFrame(dict(zip(["original", "b", "c"], expected)), dtype=object) @@ -1878,7 +1878,7 @@ def no_rows(*args, **kwargs): 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 "9999999"} + overrides = {"spark.sql.session.localRelationCacheThreshold": "1" if cached else "999999999"} with _arrow_input_service(protocol, **overrides) as (spark, state): df = spark.createDataFrame(frame, schema, verifySchema=verify) df.toArrow() @@ -1915,18 +1915,18 @@ class Day(date): def toordinal(self): raise AssertionError("guard must not call conversion hooks") - values = [date(2020, 1, 1)] * 1000 + values = [date(2020, 1, 1)] * 10000 frame = pd.DataFrame({"day": values}, dtype=object) schema = T.StructType([T.StructField("day", T.DateType())]) if reason == "small": - frame = frame.iloc[:999] + frame = frame.iloc[:9999] elif reason == "empty": frame = frame.iloc[:0] elif reason in ("inferred", "names", "ddl", "cast"): schema = {"inferred": None, "names": ["day"], "ddl": "day DATE", "cast": T.StructType([T.StructField("day", T.StringType())])}[reason] elif reason == "mixed_schema": - frame["long"] = np.arange(1000) + frame["long"] = np.arange(10000) schema.add("long", T.LongType()) elif reason == "width": schema.add("extra", T.DateType()) @@ -1972,7 +1972,7 @@ def toordinal(self): class Frame(pd.DataFrame): pass - values = [date(2020, 1, 1)] * (999 if reason == "small" else 1000) + values = [date(2020, 1, 1)] * (9999 if reason == "small" else 10000) if reason in ("nan", "pd_na", "nat", "nonnullable_null"): values[-1] = {"nan": float("nan"), "pd_na": pd.NA, "nat": pd.NaT, "nonnullable_null": None}[reason] @@ -2035,9 +2035,9 @@ def test_date_custom_schema_descriptors_keep_native_row_path(class_name, attribu import pandas as pd from pyspark.sql import types as T -frame = pd.DataFrame({"day": [date(2020, 1, 1)] * 1000}, dtype=object) +frame = pd.DataFrame({"day": [date(2020, 1, 1)] * 10000}, dtype=object) schema = T.StructType([T.StructField("day", T.DateType(), True, {"source": "native"})]) -class_name, attribute, before_import = sys.argv[1:] +class_name, attribute, before_import = sys.argv[1:4] if before_import == "False": import pyspark.sql._pandas_input cls = getattr(T, class_name) @@ -2054,7 +2054,12 @@ def test_date_custom_schema_descriptors_keep_native_row_path(class_name, attribu "dataType": T.DateType(), "nullable": False, "metadata": {"changed": True}}[attribute] replacement = property(lambda self: value) setattr(cls, attribute, replacement) -from tests.test_arrow_handoff import _arrow_input_service, protocol +import importlib.util +spec = importlib.util.spec_from_file_location("date_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 class Factory: def mktemp(self, name): return Path(tempfile.mkdtemp(prefix=name)) @@ -2063,8 +2068,9 @@ def mktemp(self, name): assert state["config_requests"] == [], state["config_requests"] table = state["input_tables"][0] assert table.schema.names == ["day"], table.schema - assert table.column(0).to_pylist() == [date(2020, 1, 1)] * 1000 + assert table.column(0).to_pylist() == [date(2020, 1, 1)] * 10000 ''' result = subprocess.run([sys.executable, "-c", script, class_name, attribute, - str(before_import)], capture_output=True, text=True, timeout=30) + str(before_import), __file__], + capture_output=True, text=True, timeout=30) assert result.returncode == 0, result.stdout + result.stderr From bd725f938f99b1fbacb30fcfd338f995ed50707e Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Thu, 29 Oct 2026 10:15:23 -0700 Subject: [PATCH 24/47] [SPARK-60716][CONNECT] Respect the optional protoc test fixture Use the shared protocol fixture so fresh-process descriptor regressions skip consistently with the Arrow suite when protoc is unavailable. Co-authored-by: Isaac --- python/tests/test_arrow_handoff.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/python/tests/test_arrow_handoff.py b/python/tests/test_arrow_handoff.py index 57041969..985d8b4f 100644 --- a/python/tests/test_arrow_handoff.py +++ b/python/tests/test_arrow_handoff.py @@ -2024,7 +2024,9 @@ def record_rows(self, *args, **kwargs): ("StructField", "__class__"), ("DateType", "__class__"), ]) @pytest.mark.parametrize("before_import", [False, True]) -def test_date_custom_schema_descriptors_keep_native_row_path(class_name, attribute, before_import): +def test_date_custom_schema_descriptors_keep_native_row_path( + protocol, class_name, attribute, before_import +): import sys script = r''' From 21ca8f1dc7db2791b5865442b6e0b2ea2d5539b8 Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Thu, 29 Oct 2026 14:09:38 -0700 Subject: [PATCH 25/47] [SPARK-60716][CONNECT] Preserve custom pandas date fallback Require concrete pandas NumPy storage and unchanged converter aliases before selecting the date Arrow path. Preserve custom array values and native fallback error ordering. Co-authored-by: Isaac --- python/pyspark/sql/_pandas_input.py | 55 ++++++- python/tests/test_arrow_handoff.py | 246 ++++++++++++++++++++++++++++ 2 files changed, 298 insertions(+), 3 deletions(-) diff --git a/python/pyspark/sql/_pandas_input.py b/python/pyspark/sql/_pandas_input.py index bebd739c..0b789aa3 100644 --- a/python/pyspark/sql/_pandas_input.py +++ b/python/pyspark/sql/_pandas_input.py @@ -284,18 +284,55 @@ def prepare_date_table(frame, schema): Those cases must keep the existing row conversion and its error ordering. """ from datetime import date + import sys import numpy as np import pandas as pd + from pandas.core.frame import DataFrame + from pandas.core.series import Series from types import GetSetDescriptorType, WrapperDescriptorType - from pyspark._pyspark import DateType, StructField, StructType as NativeStructType + from pyspark._pyspark import ( + BooleanType, ByteType, ShortType, IntegerType, LongType, FloatType, + DoubleType, DecimalType, StringType, BinaryType, DateType, StructField, + StructType as NativeStructType, + ) from pyspark.sql import types as sql_types # Keep smaller frames on the row path to avoid the importer's config RPC overhead. - if (type(frame) is not pd.DataFrame or type(schema) is not NativeStructType + if (pd.DataFrame is not DataFrame or type(frame) is not DataFrame + or type(schema) is not NativeStructType or len(frame) < 10000): return None + # Earlier converter branches must not recognize a date as another type. + # Check loaded converters and globals captured by older importer functions. + 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") + namespaces = [vars(sql_types)] + if converter is not None: + namespaces.append(vars(converter)) + if importer is not None: + for name in ("from_arrow_schema", "to_arrow_schema"): + # Newer importers resolve these functions locally in 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), ("FloatType", FloatType), + ("DoubleType", DoubleType), ("DecimalType", DecimalType), ("StringType", StringType), + ("BinaryType", BinaryType), ("DateType", DateType), ("StructField", StructField), + ("StructType", NativeStructType), + ): + if any(namespace.get(name) is not cls for namespace in namespaces): + 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 (("DateType", DateType), ("StructField", StructField), @@ -325,10 +362,22 @@ def prepare_date_table(frame, schema): for field in schema.fields ) or len(set(schema.names)) != width: return None - if not all(dtype == np.dtype("object") for dtype in frame.dtypes): + 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 can report an object dtype while its Arrow hook + # emits different cells. Trust only pandas' concrete NumPy backing array. + 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)] for column, field in zip(columns, schema.fields): + if type(column) is not Series or type(column.array) is not NumpyExtensionArray: + return None # Arrow rejects nulls in non-nullable fields; let the row path retain # its existing behavior for that input instead. if not all(type(value) is date or (value is None and field.nullable) diff --git a/python/tests/test_arrow_handoff.py b/python/tests/test_arrow_handoff.py index 985d8b4f..7dbfee29 100644 --- a/python/tests/test_arrow_handoff.py +++ b/python/tests/test_arrow_handoff.py @@ -2076,3 +2076,249 @@ def mktemp(self, name): str(before_import), __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_date_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_date_table + + arrow_calls = [] + + class ObjectEqualDtype(ExtensionDtype): + name = "object_equal_date" + type = date + 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([date(2020, 1, 2)] * len(self), type=type) + + array = DifferentArrowArray([date(2020, 1, 1)] * 10000) + frame = pd.DataFrame({"day": pd.Series(array, copy=False)}) + schema = T.StructType([T.StructField("day", T.DateType())]) + 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_date_table(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() == [date(2020, 1, 1)] * 10000 + assert state["config_requests"] == [] + assert arrow_calls == [] + + +def test_date_rebound_dataframe_alias_keeps_subclass_on_row_path(monkeypatch): + from pyspark.sql._pandas_input import prepare_date_table + + class Frame(pd.DataFrame): + pass + + frame = Frame({"day": [date(2020, 1, 1)] * 10000}, dtype=object) + schema = T.StructType([T.StructField("day", T.DateType())]) + monkeypatch.setattr(pd, "DataFrame", Frame) + assert prepare_date_table(frame, schema) is None + + +def _date_alias_guard_process(mode, argument): + """Keep mutations of converter aliases 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 json +import sys +import tempfile +from datetime import date +from pathlib import Path +import pandas as pd +import pyarrow as pa +from pyspark.sql import types as T +from pyspark import _pyspark as native + +spec = importlib.util.spec_from_file_location("date_alias_capture", sys.argv[3]) +capture = importlib.util.module_from_spec(spec) +sys.modules[spec.name] = capture +spec.loader.exec_module(capture) +values = [date.min, date.max, date(1969, 12, 31), None] * 2500 +frame = pd.DataFrame({"original": values}, dtype=object) +schema = T.StructType([T.StructField("renamed", T.DateType(), True, {"source": "native"})]) +mode, argument = sys.argv[1:3] +if mode in ("public_alias", "converter_alias", "captured_alias"): + replacement = native.BooleanType if argument == "DateType" else native.DateType + if mode == "public_alias": + sys.modules.pop("pyspark.sql.pandas.types", None) + 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.DateType +elif mode == "importer_function": + importer = importlib.import_module("pyspark.sql._arrow_input") + setattr(importer, argument, None) +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 == pa.date32() + assert table.column(0).to_pylist() == values + assert table.schema.field(0).nullable + assert table.schema.field(0).metadata is None + relation_schema = json.loads(state["relations"][0].local_relation.schema) + assert relation_schema["fields"][0]["metadata"] == {"source": "native"} +''' + result = subprocess.run([sys.executable, "-c", script, mode, argument, __file__], + capture_output=True, text=True, timeout=30) + assert result.returncode == 0, result.stdout + result.stderr + + +@pytest.mark.parametrize("class_name", [ + "BooleanType", "ByteType", "ShortType", "IntegerType", "LongType", "FloatType", + "DoubleType", "DecimalType", "StringType", "BinaryType", "DateType", "StructField", "StructType", +]) +@pytest.mark.parametrize("mode", ["public_alias", "converter_alias", "captured_alias"]) +def test_date_rebound_converter_aliases_keep_native_row_path(class_name, mode): + _date_alias_guard_process(mode, class_name) + + +@pytest.mark.parametrize("replacement", ["none", "date"]) +def test_date_rebound_importer_schema_alias_keeps_native_row_path(replacement): + _date_alias_guard_process("importer_alias", replacement) + + +@pytest.mark.parametrize("name", ["from_arrow_schema", "to_arrow_schema"]) +def test_date_present_invalid_importer_function_keeps_native_row_path(name): + _date_alias_guard_process("importer_function", name) + + +def test_date_warm_importer_without_global_converter_aliases(protocol, monkeypatch): + import pyspark.sql._arrow_input as importer + import sys + import types + + original = importer.prepare_table + + def local_imports(table, schema, *configs): + 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 + + values = [date.min, date.max, date(1970, 1, 1), None] * 2500 + frame = pd.DataFrame({"day": values}, dtype=object) + schema = T.StructType([T.StructField("renamed", T.DateType())]) + with _arrow_input_service(protocol) as (spark, state): + spark.createDataFrame(pa.table({"warm": [1, 2]})).toArrow() + # Exercise the local-import shape on this branch's pinned dependency. + substitute = types.ModuleType(importer.__name__) + substitute.StructType = T.StructType + substitute.prepare_table = local_imports + monkeypatch.setitem(sys.modules, importer.__name__, substitute) + + def no_rows(*args, **kwargs): + raise AssertionError("warm date 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() == values From bbd3c15b017c8d5aa35077fc7330552efb9073d8 Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Thu, 29 Oct 2026 14:20:20 -0700 Subject: [PATCH 26/47] [SPARK-60716][CONNECT] Retain direct pandas date normalization hooks Keep astype, notna and where class replacements and instance hooks on the native row path. Authenticate ordinary pandas normalization methods before columnar date conversion. Co-authored-by: Isaac --- python/pyspark/sql/_pandas_input.py | 21 ++++++++++-- python/tests/test_arrow_handoff.py | 51 +++++++++++++++++++++++++++-- 2 files changed, 68 insertions(+), 4 deletions(-) diff --git a/python/pyspark/sql/_pandas_input.py b/python/pyspark/sql/_pandas_input.py index 0b789aa3..982dc31c 100644 --- a/python/pyspark/sql/_pandas_input.py +++ b/python/pyspark/sql/_pandas_input.py @@ -288,9 +288,10 @@ def prepare_date_table(frame, schema): 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 GetSetDescriptorType, WrapperDescriptorType + from types import FunctionType, GetSetDescriptorType, WrapperDescriptorType from pyspark._pyspark import ( BooleanType, ByteType, ShortType, IntegerType, LongType, FloatType, @@ -302,7 +303,23 @@ def prepare_date_table(frame, schema): # Keep smaller frames on the row path to avoid the importer's config RPC overhead. if (pd.DataFrame is not DataFrame or type(frame) is not DataFrame or type(schema) is not NativeStructType - or len(frame) < 10000): + or len(frame) < 10000 + or any(name in vars(frame) for name in ("astype", "notna", "where"))): + return None + # The row path calls these normalization methods. Custom class or instance + # hooks must retain those calls and the values they produce before Arrow RPCs. + for name, module, qualname in ( + ("astype", generic_module, "NDFrame.astype"), + ("notna", frame_module, "DataFrame.notna"), + ("where", generic_module, "NDFrame.where"), + ): + method = getattr(DataFrame, name) + if (type(method) is not FunctionType or method.__name__ != name + or method.__qualname__ != qualname 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 date as another type. # Check loaded converters and globals captured by older importer functions. diff --git a/python/tests/test_arrow_handoff.py b/python/tests/test_arrow_handoff.py index 7dbfee29..4d4045c8 100644 --- a/python/tests/test_arrow_handoff.py +++ b/python/tests/test_arrow_handoff.py @@ -1874,7 +1874,7 @@ def test_date_columns_preserve_every_ipc_cell(protocol, monkeypatch, cached, ver def no_rows(*args, **kwargs): raise AssertionError("date fast path must not materialize Python rows") - for method in ("astype", "notna", "to_numpy", "iterrows", "itertuples"): + 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) @@ -2316,9 +2316,56 @@ def local_imports(table, schema, *configs): def no_rows(*args, **kwargs): raise AssertionError("warm date input must stay columnar") - monkeypatch.setattr(pd.DataFrame, "astype", no_rows) + monkeypatch.setattr(pd.DataFrame, "values", property(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() == values + + +@pytest.mark.parametrize("location", ["class", "instance"]) +@pytest.mark.parametrize("method", ["astype", "notna", "where"]) +def test_date_custom_normalization_hooks_keep_native_row_path( + protocol, monkeypatch, location, method +): + from types import MethodType + from pyspark.sql._pandas_input import prepare_date_table + + values = [date(2020, 1, 1)] * 10000 + frame = pd.DataFrame({"original": values}, dtype=object) + schema = T.StructType([T.StructField("renamed", T.DateType())]) + original = getattr(pd.DataFrame, method) + calls = [] + + def replacement(self, *args, **kwargs): + calls.append(method) + result = original(self, *args, **kwargs) + if method == "notna": + result.iat[-1, 0] = False + else: + result.iat[0, 0] = date(2020, 1, 2 if method == "astype" else 3) + return result + + if location == "class": + monkeypatch.setattr(pd.DataFrame, method, replacement) + else: + # Avoid pandas' attribute-to-column routing when installing an instance hook. + monkeypatch.setitem(vars(frame), method, MethodType(replacement, frame)) + assert prepare_date_table(frame, schema) is None + expected = values.copy() + if method == "notna": + expected[-1] = None + elif method == "astype": + expected[0] = date(2020, 1, 2) + elif location == "class": + expected[0] = date(2020, 1, 3) + with _arrow_input_service(protocol) as (spark, state): + spark.createDataFrame(frame, schema).toArrow() + assert state["config_requests"] == [] + table = state["input_tables"][0] + assert table.column_names == ["renamed"] + assert table.column(0).to_pylist() == expected + # astype returns a new frame, so an original instance's where hook is not + # used by the native normalization sequence. Preserve that ordering too. + assert calls == ([] if method == "where" and location == "instance" else [method]) From b9d1a5de6a1f2263c9fe91b3ea9e1f44cee3ff19 Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Fri, 30 Oct 2026 01:01:46 -0700 Subject: [PATCH 27/47] [SPARK-60716][CONNECT] Preserve date column-label conversion hooks Decline customized column label storage before date column inspection so native row conversion retains hook values, exceptions, and normalization order. Co-authored-by: Isaac --- python/pyspark/sql/_pandas_input.py | 2 + python/tests/test_arrow_handoff.py | 240 ++++++++++++++++++++++++++++ 2 files changed, 242 insertions(+) diff --git a/python/pyspark/sql/_pandas_input.py b/python/pyspark/sql/_pandas_input.py index 982dc31c..edd33b16 100644 --- a/python/pyspark/sql/_pandas_input.py +++ b/python/pyspark/sql/_pandas_input.py @@ -306,6 +306,8 @@ def prepare_date_table(frame, schema): or len(frame) < 10000 or any(name in vars(frame) for name in ("astype", "notna", "where"))): return None + if not _has_canonical_columns(frame.columns): + return None # The row path calls these normalization methods. Custom class or instance # hooks must retain those calls and the values they produce before Arrow RPCs. for name, module, qualname in ( diff --git a/python/tests/test_arrow_handoff.py b/python/tests/test_arrow_handoff.py index 4d4045c8..7e9a65b1 100644 --- a/python/tests/test_arrow_handoff.py +++ b/python/tests/test_arrow_handoff.py @@ -2369,3 +2369,243 @@ def replacement(self, *args, **kwargs): # astype returns a new frame, so an original instance's where hook is not # used by the native normalization sequence. Preserve that ordering too. assert calls == ([] if method == "where" and location == "instance" else [method]) + + +def _date_columns_guard_input(native_rows=False): + class NativeRows(pd.DataFrame): + pass + + days = [date.min, date.max, date(4, 2, 29), date(1600, 2, 29), + date(1900, 2, 28), date(1900, 3, 1), date(1969, 12, 31), + date(1970, 1, 1), date(2000, 2, 29), date(2100, 3, 1)] * 1000 + nullable = [None if i % 7 == 0 else value for i, value in enumerate(reversed(days))] + expected = [days, nullable, [None] * len(days)] + frame_type = NativeRows if native_rows else pd.DataFrame + frame = frame_type(dict(zip(["original", "b", "c"], expected)), dtype=object) + schema = T.StructType([ + T.StructField("renamed.`day", T.DateType(), False, {"source": "date", "id": 3}), + T.StructField("nullable", T.DateType(), True, {"tag": [1, 2]}), + T.StructField("all_null", T.DateType(), True, {"source": "all-null"}), + ]) + return frame, schema, expected + + +def _assert_date_columns_guard_payload(state, schema, expected, *, row_path): + import json + from pyspark.sql.pandas.types import from_arrow_metadata + + assert len(state["input_tables"]) == 1 + table = state["input_tables"][0] + assert table.num_rows == 10000 + assert table.column_names == ["renamed.`day", "nullable", "all_null"] + assert not state["uploads"] + epoch = date(1970, 1, 1) + for column, values, field in zip(table.columns, expected, schema.fields): + assert column.type == pa.date32() + assert column.to_pylist() == values + assert column.cast(pa.int32()).to_pylist() == [ + None if value is None else (value - epoch).days for value in values + ] + assert column.null_count == sum(value is None for value in values) + physical = table.schema.field(field.name) + assert physical.nullable is field.nullable + if row_path: + assert physical.metadata is None + else: + assert from_arrow_metadata(physical.metadata) == field.metadata + assert json.loads(state["relations"][0].local_relation.schema) == { + "type": "struct", "fields": [ + {"name": "renamed.`day", "type": "date", "nullable": False, + "metadata": {"source": "date", "id": 3}}, + {"name": "nullable", "type": "date", "nullable": True, + "metadata": {"tag": [1, 2]}}, + {"name": "all_null", "type": "date", "nullable": True, + "metadata": {"source": "all-null"}}, + ], + } + assert (state["config_requests"] == []) is row_path + + +@pytest.mark.parametrize("native_rows", [False, True]) +@pytest.mark.parametrize("location", [ + "subclass", "instance", "class", "owner", "backing_subclass", "backing_instance", + "backing_to_numpy", +]) +@pytest.mark.parametrize("behavior", ["raise", "mutate"]) +def test_date_columns_tolist_hooks_keep_native_order( + protocol, monkeypatch, native_rows, location, behavior +): + from types import MethodType + from pandas.core.base import IndexOpsMixin + + frame, schema, values = _date_columns_guard_input(native_rows) + 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("date columns hook precedes pandas normalization") + # Install normalization spies only when native conversion invokes the + # columns hook. Eligibility sees the authentic normalization methods. + monkeypatch.setitem(vars(frame), "notna", notna) + monkeypatch.setitem(vars(frame), "astype", astype) + frame.iat[0, 0] = date(2004, 2, 29) + frame.iat[-1, 1] = None + frame.iat[17, 2] = date.min + + 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", "b", "c"], 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", "b", "c"], 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", "b", "c"], 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="date 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 = [column.copy() for column in values] + expected[0][0], expected[1][-1], expected[2][17] = date(2004, 2, 29), None, date.min + _assert_date_columns_guard_payload(state, schema, expected, row_path=True) + assert calls == [hook_name, "notna", "astype"] + + +@pytest.mark.parametrize("native_rows", [False, True]) +def test_date_rebound_index_values_keeps_native_error(protocol, monkeypatch, native_rows): + frame, schema, _ = _date_columns_guard_input(native_rows) + columns = frame.columns + original = pd.Index._values.fget + calls = [] + + def values(index): + if index is columns: + calls.append("_values") + raise RuntimeError("date 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="date columns backing hook precedes normalization"): + spark.createDataFrame(frame, schema) + assert calls == ["_values"] + assert state["config_requests"] == [] + assert state["input_tables"] == [] + assert state["uploads"] == 0 + + +@pytest.mark.parametrize("native_rows", [False, True]) +def test_date_rebound_index_data_keeps_native_read_order(protocol, monkeypatch, native_rows): + frame, schema, values = _date_columns_guard_input(native_rows) + 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("date columns storage read twice before normalization") + # Preserve the native first columns read and only then attach + # observers; no private optimization call determines fallback. + monkeypatch.setitem(vars(frame), "notna", notna) + monkeypatch.setitem(vars(frame), "astype", astype) + return backing + return vars(index)["_data"] + + def set_data(index, value): + 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_date_columns_guard_payload(state, schema, values, row_path=True) + + +@pytest.mark.parametrize("storage", ["object", "default", "python", "pyarrow", "python_nan", "pyarrow_nan"]) +def test_date_canonical_columns_storage_stays_columnar(protocol, monkeypatch, storage): + frame, schema, values = _date_columns_guard_input() + labels = ["original", "b", "c"] + if storage == "object": + frame.columns = pd.Index(labels, 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(labels, dtype=label_dtype)) + + def no_rows(*args, **kwargs): + raise AssertionError("ordinary date columns must retain columnar construction") + + monkeypatch.setattr(pd.DataFrame, "values", property(no_rows)) + with _arrow_input_service(protocol) as (spark, state): + spark.createDataFrame(frame, schema).toArrow() + _assert_date_columns_guard_payload(state, schema, values, row_path=False) From 1d9869ccc6fc77ff734b716378f41d90a659c8a6 Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Thu, 5 Nov 2026 14:21:02 -0800 Subject: [PATCH 28/47] [SPARK-60716][CONNECT] Preserve current Arrow alias fallback semantics Co-authored-by: Isaac --- python/pyspark/sql/_pandas_input.py | 9 +++++++-- python/tests/test_arrow_handoff.py | 26 ++++++++++++++++++++++++++ 2 files changed, 33 insertions(+), 2 deletions(-) diff --git a/python/pyspark/sql/_pandas_input.py b/python/pyspark/sql/_pandas_input.py index edd33b16..10438470 100644 --- a/python/pyspark/sql/_pandas_input.py +++ b/python/pyspark/sql/_pandas_input.py @@ -294,7 +294,7 @@ def prepare_date_table(frame, schema): from types import FunctionType, GetSetDescriptorType, WrapperDescriptorType from pyspark._pyspark import ( - BooleanType, ByteType, ShortType, IntegerType, LongType, FloatType, + AtomicType, BooleanType, ByteType, ShortType, IntegerType, LongType, FloatType, DoubleType, DecimalType, StringType, BinaryType, DateType, StructField, StructType as NativeStructType, ) @@ -328,8 +328,13 @@ def prepare_date_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 NativeStructType: + if vars(sql_types).get("AtomicType") is not AtomicType: return None + if importer is not None: + if getattr(importer, "StructType", None) is not NativeStructType: + return None + 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 7e9a65b1..0e20a2fb 100644 --- a/python/tests/test_arrow_handoff.py +++ b/python/tests/test_arrow_handoff.py @@ -2240,6 +2240,21 @@ def _date_alias_guard_process(mode, argument): 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.DateType + 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.DateType @@ -2282,6 +2297,16 @@ def test_date_rebound_importer_schema_alias_keeps_native_row_path(replacement): _date_alias_guard_process("importer_alias", replacement) +@pytest.mark.parametrize("scope", ["public_before", "public_after", "importer"]) +@pytest.mark.parametrize("replacement", ["none", "date"]) +def test_date_rebound_atomic_alias_keeps_native_row_path(scope, replacement): + _date_alias_guard_process("atomic_alias", f"{scope}:{replacement}") + + +def test_date_deleted_importer_atomic_alias_keeps_native_row_path(): + _date_alias_guard_process("atomic_alias", "importer:missing") + + @pytest.mark.parametrize("name", ["from_arrow_schema", "to_arrow_schema"]) def test_date_present_invalid_importer_function_keeps_native_row_path(name): _date_alias_guard_process("importer_function", name) @@ -2310,6 +2335,7 @@ def local_imports(table, schema, *configs): # Exercise the local-import shape on this branch's pinned dependency. substitute = types.ModuleType(importer.__name__) substitute.StructType = T.StructType + substitute.AtomicType = T.AtomicType substitute.prepare_table = local_imports monkeypatch.setitem(sys.modules, importer.__name__, substitute) From 53808965f98d656e107542dd2f2b4dd56aa114ee Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Thu, 5 Nov 2026 14:54:07 -0800 Subject: [PATCH 29/47] Align date label regression tests with current native iteration Co-authored-by: Isaac --- python/tests/test_arrow_handoff.py | 121 +++++++++++++++++++---------- 1 file changed, 82 insertions(+), 39 deletions(-) diff --git a/python/tests/test_arrow_handoff.py b/python/tests/test_arrow_handoff.py index 0e20a2fb..058aa095 100644 --- a/python/tests/test_arrow_handoff.py +++ b/python/tests/test_arrow_handoff.py @@ -2454,21 +2454,20 @@ def _assert_date_columns_guard_payload(state, schema, expected, *, row_path): @pytest.mark.parametrize("native_rows", [False, True]) @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_date_columns_tolist_hooks_keep_native_order( +def test_date_columns_iterator_hooks_keep_native_order( protocol, monkeypatch, native_rows, location, behavior ): - from types import MethodType from pandas.core.base import IndexOpsMixin frame, schema, values = _date_columns_guard_input(native_rows) - 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 @@ -2481,71 +2480,93 @@ 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("date columns hook precedes pandas normalization") - # Install normalization spies only when native conversion invokes the - # columns hook. Eligibility sees the authentic normalization methods. + # 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] = date(2004, 2, 29) frame.iat[-1, 1] = None frame.iat[17, 2] = date.min - 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", "b", "c"], 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", "b", "c"], 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", "b", "c"], dtype="string")) + elif location.startswith("backing_"): + frame.columns = pd.Index(pd.array(["original", "b", "c"], 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_tolist(*args, **kwargs): - before_tolist() - return original(*args, **kwargs) + def backing_iter(array): + if array is backing: + before_hook() + return original(array) - monkeypatch.setitem(vars(backing), hook_name, backing_tolist) - assert type(frame.columns) is pd.Index + 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 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 with _arrow_input_service(protocol) as (spark, state): if behavior == "raise": with pytest.raises(ColumnsError, match="date 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 + assert calls == ["columns iterator"] + assert not state["config_requests"] + assert not state["uploads"] + assert not state["relations"] + assert not state["input_tables"] else: spark.createDataFrame(frame, schema).toArrow() expected = [column.copy() for column in values] expected[0][0], expected[1][-1], expected[2][17] = date(2004, 2, 29), None, date.min _assert_date_columns_guard_payload(state, schema, expected, row_path=True) - assert calls == [hook_name, "notna", "astype"] + assert calls == ["columns iterator", "notna", "astype"] @pytest.mark.parametrize("native_rows", [False, True]) @@ -2579,6 +2600,7 @@ def test_date_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 @@ -2597,8 +2619,8 @@ def data(index): if index is columns: if not normalizing: calls.append("_data") - if calls.count("_data") > 1: - raise ColumnsDataError("date 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") # Preserve the native first columns read and only then attach # observers; no private optimization call determines fallback. monkeypatch.setitem(vars(frame), "notna", notna) @@ -2610,10 +2632,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 + + # Measure current native fallback dispatch without either eligibility probe. + with monkeypatch.context() as baseline: + baseline.setattr(helper, "prepare_int64_table", lambda *_: None) + baseline.setattr(helper, "prepare_date_table", lambda *_: None) + with _arrow_input_service(protocol) as (spark, state): + spark.createDataFrame(frame, schema).toArrow() + _assert_date_columns_guard_payload(state, schema, values, row_path=True) + control_table = state["input_tables"][0] + 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 ordinary methods before candidate eligibility; the storage hook + # installs observers only when the native columns iterator starts. + monkeypatch.delitem(vars(frame), "notna") + monkeypatch.delitem(vars(frame), "astype") with _arrow_input_service(protocol) as (spark, state): spark.createDataFrame(frame, schema).toArrow() - assert calls == ["_data", "notna", "astype"] + assert calls == baseline_calls _assert_date_columns_guard_payload(state, schema, values, row_path=True) + assert state["input_tables"][0].equals(control_table, check_metadata=True) @pytest.mark.parametrize("storage", ["object", "default", "python", "pyarrow", "python_nan", "pyarrow_nan"]) From 280bdee0255c0cc4a752ab50f495023a2b9335ca Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Fri, 30 Oct 2026 02:55:07 -0700 Subject: [PATCH 30/47] [SPARK-60727][CONNECT] Keep canonical double input on inline transport Build on the canonical pandas Index and integer guards from PR121. Serialize eligible native float64 columns to inline Arrow IPC without row materialization, preserving explicit schema and pandas missing values. Include the shared internal inline IPC bridge and public Float coverage. Co-authored-by: Isaac --- crates/spark-connect/src/session.rs | 17 + python/pyspark/sql/_pandas_input.py | 135 ++++ python/tests/test_arrow_handoff.py | 917 ++++++++++++++++++++++++++++ 3 files changed, 1069 insertions(+) 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..9f028dbf 100644 --- a/python/pyspark/sql/_pandas_input.py +++ b/python/pyspark/sql/_pandas_input.py @@ -274,3 +274,138 @@ def prepare_int64_table(frame, schema): [pa.array(column, type=pa.int64(), from_pandas=True) for column in columns], names=names, ) + +def prepare_float64_table(frame, schema): + """Keep explicit native double columns columnar without changing pandas nulls.""" + 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, WrapperDescriptorType + + from pyspark import _pyspark as native + from pyspark.sql import types as sql_types + + # First-use Arrow imports outweigh the conversion saving for smaller frames. + # Keep those on the existing row path. + if (pd.DataFrame is not DataFrame or type(frame) is not DataFrame + or type(schema) is not native.StructType or len(frame) < 100000 + or any(name in vars(frame) for name in ("astype", "notna"))): + return None + if not _has_canonical_columns(frame.columns): + return None + # Class or instance conversion hooks must retain the row path's calls. + 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 + 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 + ): + return None + double_dtype = np.dtype("float64") + if not all(type(dtype) is type(double_dtype) and dtype == double_dtype + for dtype in frame.dtypes): + return None + # A custom ExtensionArray can advertise the literal NumPy dtype while its + # Arrow hook emits other values. Accept only pandas' concrete NumPy storage. + 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 + # Native row conversion reads Rust fields, bypassing mutable Python accessors. + for name in ("DoubleType", "StructField", "StructType"): + cls = getattr(native, name) + 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 ( + (native.StructType, ("fields", "names"), GetSetDescriptorType), + (native.StructType, ("__iter__",), WrapperDescriptorType), + (native.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 + fields = schema.fields + if len(fields) != width or not all( + type(field) is native.StructField and type(field.dataType) is native.DoubleType + and type(field.name) is str and "\x00" not in field.name for field in fields + ) or len(set(schema.names)) != width: + return None + + from pyspark.sql import _arrow_input + from pyspark.sql.pandas import types as pandas_types + + for name in ("to_arrow_schema", "to_arrow_type", "to_arrow_metadata", + "_check_arrow_table_timestamps_localize"): + converter = getattr(pandas_types, name) + if (type(converter) is not FunctionType or converter.__name__ != name + or converter.__module__ != pandas_types.__name__ + or converter.__globals__ is not vars(pandas_types)): + return None + # Earlier converter branches can reinterpret DoubleType after alias rebinding. + for name in ("BooleanType", "ByteType", "ShortType", "IntegerType", "LongType", + "FloatType", "DoubleType", "StructField", "StructType"): + cls = getattr(native, name) + if (getattr(sql_types, name, None) is not cls + or getattr(pandas_types, name, None) is not cls): + return None + if (_arrow_input.StructType is not native.StructType + or getattr(_arrow_input, "to_arrow_schema", pandas_types.to_arrow_schema) + is not pandas_types.to_arrow_schema + or getattr(_arrow_input, "_check_arrow_table_timestamps_localize", + pandas_types._check_arrow_table_timestamps_localize) + is not pandas_types._check_arrow_table_timestamps_localize): + return None + + # Use the same pandas missing-value semantics as the row path, including the + # legacy use_inf_as_na option. Infinities and signed zero otherwise stay intact. + valid = frame.notna() + if any(not field.nullable and not valid.iloc[:, i].all() + for i, field in enumerate(fields)): + return None + + import pyarrow as pa + + return pa.Table.from_arrays( + [pa.array(column, type=pa.float64(), + mask=~valid.iloc[:, i], from_pandas=False) + for i, column in enumerate(columns)], + # Match the row serializer's physical fields. The complete Spark + # metadata remains in the explicit logical schema, not Arrow metadata. + schema=pa.schema([pa.field(field.name, pa.float64(), nullable=field.nullable) + for field in fields]), + ) + + +def prepare_float64_input(frame, schema): + """Serialize supported double columns with the existing row transport policy.""" + table = prepare_float64_table(frame, schema) + if table is None: + return None + + import pyarrow as pa + + sink = pa.BufferOutputStream() + with pa.ipc.new_stream(sink, table.schema) as writer: + writer.write_table(table) + return sink.getvalue().to_pybytes() + diff --git a/python/tests/test_arrow_handoff.py b/python/tests/test_arrow_handoff.py index 5d098808..6d606b60 100644 --- a/python/tests/test_arrow_handoff.py +++ b/python/tests/test_arrow_handoff.py @@ -452,6 +452,8 @@ def configure(raw, context): state["config_requests"].append((operation, list(keys))) if config_error: context.abort(grpc.StatusCode.UNAVAILABLE, "configuration unavailable") + if error := state.pop("config_error", None): + context.abort(error, "configuration read failed") 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] @@ -1725,3 +1727,918 @@ def set_data(index, value): "metadata": metadata} ], } + +def _assert_float64_ipc(table, columns, schema, *, missing=None): + assert table.num_rows == len(columns[0]) + assert table.schema.names == schema.names + assert table.num_columns == len(columns) + for index, (column, expected, field) in enumerate(zip(table.columns, columns, schema.fields)): + mask = np.isnan(expected) if missing is None else missing[index] + assert column.type == pa.float64() + # Check every IPC validity bit and the exact bits of every non-null value, + # including negative zero, subnormals, and both infinities. + assert column.is_null().to_pylist() == mask.tolist() + actual = column.drop_null().combine_chunks().to_numpy(zero_copy_only=False) + np.testing.assert_array_equal(actual.view(np.uint64), expected[~mask].view(np.uint64)) + actual_field = table.schema.field(index) + assert actual_field.nullable == field.nullable + assert not actual_field.metadata + + + +@pytest.mark.parametrize("cached", [False, True]) +@pytest.mark.parametrize("verify", [False, True]) +def test_float64_columns_preserve_every_ipc_cell(protocol, monkeypatch, cached, verify): + finite = np.array([ + 0.0, -0.0, 1.5, -2.25, np.finfo(np.float64).max, -np.finfo(np.float64).max, + np.finfo(np.float64).tiny, -np.finfo(np.float64).tiny, + np.nextafter(0.0, 1.0), -np.nextafter(0.0, 1.0), + ], dtype=np.float64) + first = np.tile(finite, 10000) + second = np.tile(np.array([np.nan, np.inf, -np.inf, -0.0, 1.5], dtype=np.float64), 20000) + columns = [first, second, np.full(100000, np.nan, dtype=np.float64)] + frame = pd.DataFrame(dict(zip(["original", "b", "c"], columns))) + schema = T.StructType([ + T.StructField("renamed.`double", T.DoubleType(), False, {"source": "float64", "id": 3}), + T.StructField("nullable", T.DoubleType(), True, {"tag": [1, 2], "label": "é"}), + T.StructField("all_null", T.DoubleType()), + ]) + + def no_rows(*args, **kwargs): + raise AssertionError("float64 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"} + with _arrow_input_service(protocol, config_error=True, **overrides) as (spark, state): + spark.createDataFrame(frame, schema, verifySchema=verify).toArrow() + _assert_float64_ipc(state["input_tables"][0], columns, schema) + assert state["uploads"] == 0 + assert state["config_requests"] == [] + + + +def test_float64_columnar_ipc_matches_existing_row_conversion(protocol, monkeypatch): + import json + + import pyspark.sql._pandas_input as helper + + values = np.tile(np.array([np.nan, np.inf, -np.inf, -0.0, 0.0, 1.25]), 20000) + frame = pd.DataFrame({"original": values}) + schema = T.StructType([T.StructField("renamed", T.DoubleType(), True, {"source": "parity"})]) + with _arrow_input_service(protocol) as (spark, state): + spark.createDataFrame(frame, schema).toArrow() + monkeypatch.setattr(helper, "prepare_float64_table", lambda *args: None) + spark.createDataFrame(frame, schema).toArrow() + assert len(state["input_tables"]) == 2 + for table in state["input_tables"]: + _assert_float64_ipc(table, [values], schema) + assert state["input_tables"][0].equals(state["input_tables"][1], check_metadata=True) + assert state["config_requests"] == [] + assert state["uploads"] == 0 + assert [json.loads(relation.local_relation.schema) for relation in state["relations"]] == [ + json.loads(schema.json()), json.loads(schema.json()) + ] + + + +@pytest.mark.parametrize("nullable", [False, True]) +def test_float64_pandas_infinity_missing_option_keeps_row_semantics(protocol, nullable): + from pyspark.errors import PySparkRuntimeError + + try: + pd.get_option("mode.use_inf_as_na") + except (KeyError, AttributeError): + pytest.skip("pandas does not expose mode.use_inf_as_na") + values = np.tile(np.array([np.inf, -np.inf, np.nan, -0.0, 1.25]), 20000) + frame = pd.DataFrame({"x": values}) + schema = T.StructType([T.StructField("x", T.DoubleType(), nullable)]) + with pd.option_context("mode.use_inf_as_na", True): + with _arrow_input_service(protocol) as (spark, state): + if not nullable: + with pytest.raises(PySparkRuntimeError, match=( + "Column 'x' is declared as non-nullable but contains null values" + )): + spark.createDataFrame(frame, schema) + assert state["config_requests"] == [] + assert state["uploads"] == 0 + return + spark.createDataFrame(frame, schema).toArrow() + _assert_float64_ipc(state["input_tables"][0], [values], schema, + missing=[~np.isfinite(values)]) + + + +@pytest.mark.parametrize("reason", [ + "small", "empty", "zero_columns", "inferred", "names", "ddl", "float_schema", + "long_schema", "mixed_schema", "width", "duplicate_input", "duplicate_schema", + "input_nul", "schema_nul", "input_number", "input_string_subclass", "frame_subclass", + "schema_proxy", "extension", "object", "float32", "float16", "non_native", + "integer", "mixed", "nonnullable_nan", "late_unsupported", +]) +def test_float64_column_guard_preserves_row_fallback(reason): + from pyspark.sql._pandas_input import prepare_float64_table + + class Frame(pd.DataFrame): + pass + + class Name(str): + pass + + class SchemaProxy: + @property + def __class__(self): + return T.StructType + + values = np.arange(100000, dtype=np.float64) + frame = pd.DataFrame({"x": values}) + schema = T.StructType([T.StructField("x", T.DoubleType())]) + if reason == "small": + frame = frame.iloc[:99999] + elif reason == "empty": + frame = frame.iloc[:0] + elif reason == "zero_columns": + frame = frame.iloc[:, :0] + schema = T.StructType([]) + elif reason in ("inferred", "names", "ddl", "float_schema", "long_schema"): + schema = { + "inferred": None, "names": ["x"], "ddl": "x DOUBLE", + "float_schema": T.StructType([T.StructField("x", T.FloatType())]), + "long_schema": T.StructType([T.StructField("x", T.LongType())]), + }[reason] + elif reason == "mixed_schema": + frame["y"] = values + schema.add("y", T.FloatType()) + elif reason == "width": + schema.add("extra", T.DoubleType()) + elif reason in ("duplicate_input", "duplicate_schema"): + frame["y"] = values + schema.add("y", T.DoubleType()) + if reason == "duplicate_input": + frame.columns = ["x", "x"] + else: + schema = T.StructType([T.StructField("x", T.DoubleType())] * 2) + elif reason == "input_string_subclass": + frame.columns = pd.Index([Name("x")], dtype=object) + assert type(frame.columns[0]) is Name + elif reason in ("input_nul", "input_number"): + frame.columns = [{"input_nul": "x\x00", "input_number": 1}[reason]] + elif reason == "schema_nul": + schema = T.StructType([T.StructField("x\x00", T.DoubleType())]) + elif reason == "frame_subclass": + frame = Frame(frame) + elif reason == "schema_proxy": + schema = SchemaProxy() + elif reason == "nonnullable_nan": + frame.iat[-1, 0] = np.nan + schema = T.StructType([T.StructField("x", T.DoubleType(), False)]) + elif reason in ("mixed", "late_unsupported"): + frame["y"] = "text" if reason == "mixed" else values.astype(object) + schema.add("y", T.StringType() if reason == "mixed" else T.DoubleType()) + else: + dtype = {"extension": "Float64", "object": object, "float32": "float32", + "float16": "float16", "non_native": ">f8", "integer": "int64"}[reason] + frame = pd.DataFrame({"x": pd.Series(values, dtype=dtype)}) + assert prepare_float64_table(frame, schema) is None + + + +@pytest.mark.parametrize("reason", [ + "small", "inferred", "names", "ddl", "float32", "extension", "object", "frame_subclass", + "duplicate_input", "duplicate_schema", "input_nul", "schema_nul", +]) +def test_float64_public_fallback_preserves_every_ipc_cell(protocol, reason): + + class Frame(pd.DataFrame): + pass + + values = np.tile(np.array([1.5, -0.0, np.inf, -np.inf, np.nan]), 20000) + if reason == "small": + values = values[:99999] + dtype = {"float32": "float32", "extension": "Float64", "object": object}.get(reason) + frame_type = Frame if reason == "frame_subclass" else pd.DataFrame + frame = frame_type({"x": pd.Series(values, dtype=dtype)}) + columns = [values] + expected_schema = T.StructType([T.StructField("x", T.DoubleType())]) + schema = expected_schema + if reason in ("inferred", "names", "ddl"): + schema = {"inferred": None, "names": ["x"], "ddl": "x DOUBLE"}[reason] + elif reason in ("duplicate_input", "duplicate_schema"): + frame["y"] = values + columns.append(values) + expected_schema = T.StructType([ + T.StructField("x", T.DoubleType()), + T.StructField("x" if reason == "duplicate_schema" else "y", T.DoubleType()), + ]) + schema = expected_schema + if reason == "duplicate_input": + frame.columns = ["x", "x"] + elif reason == "input_nul": + frame.columns = ["x\x00original"] + elif reason == "schema_nul": + schema = expected_schema = T.StructType([T.StructField("x\x00renamed", T.DoubleType())]) + with _arrow_input_service(protocol) as (spark, state): + spark.createDataFrame(frame, schema).toArrow() + assert state["config_requests"] == [] + _assert_float64_ipc(state["input_tables"][0], columns, expected_schema) + + + +@pytest.mark.parametrize("verify", [False, True]) +def test_float64_nonnullable_nan_keeps_existing_error_before_configs(protocol, verify): + from pyspark.errors import PySparkRuntimeError + + frame = pd.DataFrame({"x": np.full(100000, np.nan, dtype=np.float64)}) + schema = T.StructType([T.StructField("x", T.DoubleType(), False)]) + with _arrow_input_service(protocol) as (spark, state): + state["config_error"] = grpc.StatusCode.INVALID_ARGUMENT + with pytest.raises(PySparkRuntimeError, match=( + "Column 'x' is declared as non-nullable but contains null values" + )): + spark.createDataFrame(frame, schema, verifySchema=verify) + 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"), + ("StructType", "__getattribute__"), ("StructField", "__getattribute__"), + ("DoubleType", "__getattribute__"), ("StructType", "__class__"), + ("StructField", "__class__"), ("DoubleType", "__class__"), + ("StructType", "alias"), ("StructField", "alias"), ("DoubleType", "alias"), +]) +@pytest.mark.parametrize("before_import", [False, True]) +def test_float64_custom_schema_descriptors_keep_native_row_path( + protocol, class_name, attribute, before_import +): + import sys + + script = r''' +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 + +spec = importlib.util.spec_from_file_location("float64_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 +values = np.tile(np.array([1.5, -0.0, np.inf, -np.inf, np.nan]), 20000) +frame = pd.DataFrame({"x": values}) +schema = T.StructType([T.StructField("x", T.DoubleType(), True, {"source": "native"})]) +class_name, attribute, before_import = sys.argv[1:4] +if before_import == "True": + 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.DoubleType())]) + elif attribute == "__class__": + attribute = "__getattribute__" + replacement = lambda self, name: object if name == "__class__" else ( + object.__getattribute__(self, name)) + else: + value = {"names": ["changed"], "fields": [], "name": "changed", + "dataType": T.DoubleType(), "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 == ["x"], table.schema + assert table.schema.field(0).nullable is True + assert table.schema.field(0).type == capture.pa.float64() + assert capture.T is T + mask = np.isnan(values) + assert table.column(0).is_null().to_pylist() == mask.tolist() + actual = table.column(0).drop_null().combine_chunks().to_numpy(zero_copy_only=False) + np.testing.assert_array_equal(actual.view(np.uint64), values[~mask].view(np.uint64)) + from pyspark.sql.pandas.types import from_arrow_metadata + assert from_arrow_metadata(table.schema.field(0).metadata) is None + import json + assert json.loads(state["relations"][0].local_relation.schema)["fields"][0]["metadata"] == { + "source": "native" + } +''' + 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 + + + +@pytest.mark.parametrize("class_name", [ + "BooleanType", "ByteType", "ShortType", "IntegerType", "LongType", "FloatType", + "DoubleType", "StructType", "StructField", +]) +@pytest.mark.parametrize("converter_loaded", [False, True]) +def test_float64_rebound_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 numpy as np +import pandas as pd +from pyspark.sql import types as T + +spec = importlib.util.spec_from_file_location("float64_alias_capture", sys.argv[3]) +capture = importlib.util.module_from_spec(spec) +sys.modules[spec.name] = capture +spec.loader.exec_module(capture) +values = np.tile(np.array([1.5, -0.0, np.inf, -np.inf, np.nan]), 20000) +frame = pd.DataFrame({"x": values}) +schema = T.StructType([T.StructField("x", T.DoubleType())]) +class_name, converter_loaded = sys.argv[1:3] +replacement = object if class_name in ("DoubleType", "StructType", "StructField") else T.DoubleType +if converter_loaded == "True": + import pyspark.sql.pandas.types as converter + setattr(converter, class_name, replacement) +else: + sys.modules.pop("pyspark.sql.pandas.types", None) + setattr(T, class_name, replacement) + assert "pyspark.sql.pandas.types" not in sys.modules +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 == ["x"], table.schema + assert table.schema.field(0).type == capture.pa.float64() + mask = np.isnan(values) + assert table.column(0).is_null().to_pylist() == mask.tolist() + actual = table.column(0).drop_null().combine_chunks().to_numpy(zero_copy_only=False) + np.testing.assert_array_equal(actual.view(np.uint64), values[~mask].view(np.uint64)) +''' + 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 + + + +def test_float64_rebound_arrow_schema_alias_keeps_native_row_path(protocol, monkeypatch): + import pyspark.sql._arrow_input as helper + + frame = pd.DataFrame({"x": np.arange(100000, dtype=np.float64)}) + schema = T.StructType([T.StructField("x", T.DoubleType())]) + monkeypatch.setattr(helper, "StructType", object) + with _arrow_input_service(protocol) as (spark, state): + spark.createDataFrame(frame, schema).toArrow() + assert state["config_requests"] == [] + _assert_float64_ipc(state["input_tables"][0], [frame["x"].to_numpy()], schema) + + + +def test_float64_rebound_pandas_alias_keeps_subclass_fallback(monkeypatch): + from pyspark.sql._pandas_input import prepare_float64_table + + class Frame(pd.DataFrame): + pass + + frame = Frame({"x": np.arange(100000, dtype=np.float64)}) + schema = T.StructType([T.StructField("x", T.DoubleType())]) + monkeypatch.setattr(pd, "DataFrame", Frame) + assert prepare_float64_table(frame, schema) is None + + + +@pytest.mark.parametrize("class_spoof", [False, True]) +def test_float64_custom_extension_dtype_cannot_impersonate_numpy_float64(class_spoof): + from pandas.api.extensions import ExtensionArray, ExtensionDtype, take + from pyspark.sql._pandas_input import prepare_float64_table + + arrow_calls = [] + + class FloatEqualDtype(ExtensionDtype): + name = "float64_equal" + type = np.float64 + kind = "f" + na_value = np.nan + + @classmethod + def construct_array_type(cls): + return DifferentArrowArray + + def __eq__(self, other): + return (isinstance(other, FloatEqualDtype) + or isinstance(other, np.dtype) and other == np.dtype("float64")) + + def __hash__(self): + return hash(self.name) + + class DifferentArrowArray(ExtensionArray): + def __init__(self, values): + self._values = np.asarray(values, dtype=np.float64) + + @property + def dtype(self): + return FloatEqualDtype() + + @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 np.isnan(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=np.float64) + 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([999.0] * len(self), type=type) + + frame = pd.DataFrame({"x": pd.Series(DifferentArrowArray(np.ones(100000)), copy=False)}) + schema = T.StructType([T.StructField("x", T.DoubleType())]) + if class_spoof: + def spoof_class(self, name): + if name == "__class__": + return type(np.dtype("float64")) + return object.__getattribute__(self, name) + + FloatEqualDtype.__getattribute__ = spoof_class + dtype = frame.dtypes.iloc[0] + assert type(frame) is pd.DataFrame + assert type(dtype) is FloatEqualDtype + assert isinstance(dtype, np.dtype) is class_spoof + assert dtype == np.dtype("float64") + assert prepare_float64_table(frame, schema) is None + assert arrow_calls == [] + + + +@pytest.mark.parametrize("method", ["astype", "notna"]) +def test_float64_instance_conversion_hooks_keep_row_fallback(protocol, method): + values = np.arange(100000, dtype=np.float64) + frame = pd.DataFrame({"x": values}) + schema = T.StructType([T.StructField("x", T.DoubleType())]) + original = getattr(frame, method) + calls = [] + missing = np.zeros(len(values), dtype=bool) + expected = values + + if method == "astype": + expected = values + 7.0 + + def customized(*args, **kwargs): + calls.append("astype") + converted = original(*args, **kwargs) + converted.iloc[:, 0] = converted.iloc[:, 0] + 7.0 + return converted + else: + missing[::3] = True + + def customized(): + calls.append("notna") + mask = original() + mask.iloc[::3, 0] = False + return mask + + 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_float64_ipc(state["input_tables"][0], [expected], schema, missing=[missing]) + + + +@pytest.mark.parametrize("module_name", ["converter", "arrow_input"]) +@pytest.mark.parametrize("attribute", ["to_arrow_schema", "_check_arrow_table_timestamps_localize"]) +def test_float64_rebound_converter_function_keeps_row_fallback( + protocol, monkeypatch, module_name, attribute +): + import pyspark.sql._arrow_input as arrow_input + import pyspark.sql.pandas.types as converter + + values = np.tile(np.array([np.nan, -0.0, np.inf, -np.inf, 1.5]), 20000) + frame = pd.DataFrame({"x": values}) + schema = T.StructType([T.StructField("x", T.DoubleType())]) + + def customized(*args, **kwargs): + raise AssertionError("changed Arrow converter must keep float64 input on the row path") + + monkeypatch.setattr({"converter": converter, "arrow_input": arrow_input}[module_name], + attribute, customized, raising=False) + with _arrow_input_service(protocol) as (spark, state): + spark.createDataFrame(frame, schema).toArrow() + assert state["config_requests"] == [] + _assert_float64_ipc(state["input_tables"][0], [values], schema) + + + +@pytest.mark.parametrize("method", ["astype", "notna", "where"]) +@pytest.mark.parametrize("before_import", [False, True]) +def test_float64_class_conversion_hooks_keep_row_fallback( + protocol, monkeypatch, method, before_import +): + import importlib + import sys + + helper_name = "pyspark.sql._pandas_input" + if before_import: + monkeypatch.delitem(sys.modules, helper_name, raising=False) + else: + importlib.import_module(helper_name) + values = np.arange(100000, dtype=np.float64) + frame = pd.DataFrame({"x": values}) + schema = T.StructType([T.StructField("x", T.DoubleType())]) + original = getattr(pd.DataFrame, method) + calls = [] + missing = np.zeros(len(values), dtype=bool) + expected = values + + if method == "astype": + expected = values + 7.0 + + def customized(self, *args, **kwargs): + calls.append("astype") + converted = original(self, *args, **kwargs) + converted.iloc[:, 0] = converted.iloc[:, 0] + 7.0 + return converted + elif method == "notna": + missing[::3] = True + + def customized(self): + calls.append("notna") + mask = original(self) + mask.iloc[::3, 0] = False + return mask + else: + expected = values + 11.0 + + def customized(self, *args, **kwargs): + calls.append("where") + converted = original(self, *args, **kwargs) + converted.iloc[:, 0] = converted.iloc[:, 0] + 11.0 + return converted + + monkeypatch.setattr(pd.DataFrame, method, customized) + if before_import: + importlib.import_module(helper_name) + with _arrow_input_service(protocol) as (spark, state): + spark.createDataFrame(frame, schema).toArrow() + assert calls == [method] + assert state["config_requests"] == [] + _assert_float64_ipc(state["input_tables"][0], [expected], schema, missing=[missing]) + + + +@pytest.mark.parametrize("attribute", ["to_arrow_type", "to_arrow_metadata"]) +def test_float64_custom_type_and_metadata_converters_keep_row_fallback( + protocol, monkeypatch, attribute +): + import pyspark.sql.pandas.types as converter + + values = np.tile(np.array([np.nan, -0.0, np.inf, -np.inf, 1.5]), 20000) + frame = pd.DataFrame({"x": values}) + schema = T.StructType([T.StructField("x", T.DoubleType())]) + + def customized(*args, **kwargs): + raise AssertionError("changed type or metadata converter must retain the row path") + + monkeypatch.setattr(converter, attribute, customized) + with _arrow_input_service(protocol) as (spark, state): + spark.createDataFrame(frame, schema).toArrow() + assert state["config_requests"] == [] + _assert_float64_ipc(state["input_tables"][0], [values], schema) + + + +def test_float64_literal_numpy_dtype_custom_backing_keeps_row_error(protocol): + from pyspark.sql._pandas_input import prepare_float64_table + + calls = [] + + class DifferentArrowArray(pd.arrays.FloatingArray): + def __arrow_array__(self, type=None): + calls.append("arrow") + return pa.array(np.full(len(self), 999.0), type=type) + + values = np.tile(np.array([1.5, -0.0, np.inf, -np.inf, np.nan]), 20000) + array = DifferentArrowArray(values.copy(), np.isnan(values)) + frame = pd.DataFrame({"x": pd.Series(array, copy=False)}) + # Construct extension storage before advertising a native dtype; otherwise + # pandas legitimately normalizes the constructor input into a NumPy block. + DifferentArrowArray.dtype = property(lambda self: np.dtype("float64")) + schema = T.StructType([T.StructField("x", T.DoubleType())]) + assert type(frame) is pd.DataFrame + assert type(frame["x"]) is pd.Series + assert type(frame["x"].array) is DifferentArrowArray + assert type(frame.dtypes.iloc[0]) is type(np.dtype("float64")) + assert frame.dtypes.iloc[0] == np.dtype("float64") + assert prepare_float64_table(frame, schema) is None + assert calls == [] + with _arrow_input_service(protocol) as (spark, state): + # This inconsistent ExtensionArray already fails pandas object conversion. + # Preserve that baseline error before Arrow hooks or configuration RPCs. + with pytest.raises(AttributeError, match="na_value"): + spark.createDataFrame(frame, schema) + assert state["config_requests"] == [] + assert state["input_tables"] == [] + assert calls == [] + + + +def test_float64_missing_concrete_numpy_backing_class_keeps_fallback(monkeypatch): + import pandas.core.arrays.numpy_ as numpy_arrays + from pyspark.sql._pandas_input import prepare_float64_table + + frame = pd.DataFrame({"x": np.arange(100000, dtype=np.float64)}) + schema = T.StructType([T.StructField("x", T.DoubleType())]) + monkeypatch.delattr(numpy_arrays, "NumpyExtensionArray", raising=False) + monkeypatch.delattr(numpy_arrays, "PandasArray", raising=False) + assert prepare_float64_table(frame, schema) is None + + + +@pytest.mark.parametrize("rows", [10000, 99999, 100000]) +def test_float64_public_threshold_preserves_row_path_below_100k(protocol, rows): + import json + + values = np.resize(np.array([np.nan, -0.0, 0.0, np.inf, -np.inf, 1.5]), rows) + frame = pd.DataFrame({"original": values}) + schema = T.StructType([ + T.StructField("renamed", T.DoubleType(), True, {"boundary": {"rows": rows}}), + ]) + with _arrow_input_service( + protocol, **{"spark.sql.session.localRelationCacheThreshold": "999999999"} + ) as (spark, state): + state["config_error"] = grpc.StatusCode.INVALID_ARGUMENT + spark.createDataFrame(frame, schema).toArrow() + _assert_float64_ipc(state["input_tables"][0], [values], schema) + assert json.loads(state["relations"][0].local_relation.schema) == json.loads(schema.json()) + assert state["uploads"] == 0 + assert state["config_requests"] == [] + + + +def _float64_columns_guard_input(rows=100000): + finite = np.array([ + 0.0, -0.0, 1.5, -2.25, np.finfo(np.float64).max, -np.finfo(np.float64).max, + np.finfo(np.float64).tiny, -np.finfo(np.float64).tiny, + np.nextafter(0.0, 1.0), -np.nextafter(0.0, 1.0), + ], dtype=np.float64) + values = [ + np.resize(finite, rows), + np.resize(np.array([np.nan, np.inf, -np.inf, -0.0, 1.5]), rows), + np.full(rows, np.nan, dtype=np.float64), + ] + frame = pd.DataFrame(dict(zip(["original", "b", "c"], values))) + schema = T.StructType([ + T.StructField("renamed.`double", T.DoubleType(), False, {"source": "float64", "id": 3}), + T.StructField("nullable", T.DoubleType(), True, {"tag": [1, 2]}), + T.StructField("all_null", T.DoubleType(), True, {"source": "all-null"}), + ]) + return frame, schema, values + + +def _assert_float64_columns_guard_payload(state, schema, values): + import json + + assert len(state["input_tables"]) == 1 + assert len(state["relations"]) == 1 + _assert_float64_ipc(state["input_tables"][0], values, schema) + assert json.loads(state["relations"][0].local_relation.schema) == json.loads(schema.json()) + assert state["config_requests"] == [] + assert state["uploads"] == 0 + + + +@pytest.mark.parametrize("rows", [10000, 100000]) +@pytest.mark.parametrize("location", [ + "subclass", "instance", "class", "owner", "backing_subclass", "backing_instance", + "backing_to_numpy", +]) +@pytest.mark.parametrize("behavior", ["raise", "mutate"]) +def test_float64_columns_tolist_hooks_keep_native_order( + protocol, monkeypatch, rows, location, behavior +): + from types import MethodType + from pandas.core.base import IndexOpsMixin + + values = np.arange(rows, dtype=np.float64) + frame = pd.DataFrame({"original": values}) + schema = T.StructType([T.StructField("renamed", T.DoubleType())]) + 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], frame.iat[-1, 0] = -0.0, np.nan + + 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(frame.columns._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"] == [] + assert state["uploads"] == 0 + else: + spark.createDataFrame(frame, schema).toArrow() + expected = values.copy() + expected[0], expected[-1] = -0.0, np.nan + _assert_float64_columns_guard_payload(state, schema, [expected]) + assert calls == [hook_name, "notna", "astype"] + + + +@pytest.mark.parametrize("rows", [10000, 100000]) +def test_float64_rebound_index_values_keeps_native_error(protocol, monkeypatch, rows): + + frame = pd.DataFrame({"original": np.arange(rows, dtype=np.float64)}) + schema = T.StructType([T.StructField("renamed", T.DoubleType())]) + 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"] == [] + assert state["uploads"] == 0 + + + +@pytest.mark.parametrize("rows", [10000, 100000]) +def test_float64_rebound_index_data_keeps_native_read_order(protocol, monkeypatch, rows): + frame, schema, values = _float64_columns_guard_input(rows) + 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("float columns storage read twice before normalization") + # Preserve the native first columns read and only then attach + # observers; no private optimization call determines fallback. + monkeypatch.setitem(vars(frame), "notna", notna) + monkeypatch.setitem(vars(frame), "astype", astype) + return backing + return vars(index)["_data"] + + def set_data(index, value): + 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_float64_columns_guard_payload(state, schema, values) + + + +@pytest.mark.parametrize("storage", [ + "object", "default", "python", "pyarrow", "python_nan", "pyarrow_nan", +]) +def test_float64_canonical_columns_storage_stays_columnar(protocol, monkeypatch, storage): + frame, schema, values = _float64_columns_guard_input() + labels = ["original", "b", "c"] + if storage == "object": + frame.columns = pd.Index(labels, 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(labels, dtype=label_dtype)) + + def no_rows(*args, **kwargs): + raise AssertionError("ordinary float64 columns must retain columnar construction") + + monkeypatch.setattr(pd.DataFrame, "values", property(no_rows)) + with _arrow_input_service(protocol) as (spark, state): + spark.createDataFrame(frame, schema).toArrow() + _assert_float64_columns_guard_payload(state, schema, values) From 11ff3b70eb4a07fe4d8828bd5472e8fe836413b4 Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Fri, 30 Oct 2026 02:55:26 -0700 Subject: [PATCH 31/47] [SPARK-60727][CONNECT] Normalize Float helper terminal newline Remove the extra blank line at EOF copied from the donor suffix. All helper function bodies remain byte-identical. Co-authored-by: Isaac --- python/pyspark/sql/_pandas_input.py | 1 - 1 file changed, 1 deletion(-) diff --git a/python/pyspark/sql/_pandas_input.py b/python/pyspark/sql/_pandas_input.py index 9f028dbf..0d72e195 100644 --- a/python/pyspark/sql/_pandas_input.py +++ b/python/pyspark/sql/_pandas_input.py @@ -408,4 +408,3 @@ def prepare_float64_input(frame, schema): with pa.ipc.new_stream(sink, table.schema) as writer: writer.write_table(table) return sink.getvalue().to_pybytes() - From b9498d1007386c3646242c62e540a8cf7b0c40bf Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Thu, 5 Nov 2026 14:31:28 -0800 Subject: [PATCH 32/47] 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 2bd46383c2b4e3abe9f63d8d6ca8e04b79ceb2be Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Thu, 5 Nov 2026 14:54:07 -0800 Subject: [PATCH 33/47] Align float label regression tests with current native iteration Co-authored-by: Isaac --- python/tests/test_arrow_handoff.py | 123 ++++++++++++++++++++--------- 1 file changed, 86 insertions(+), 37 deletions(-) diff --git a/python/tests/test_arrow_handoff.py b/python/tests/test_arrow_handoff.py index 6d606b60..3b66e19f 100644 --- a/python/tests/test_arrow_handoff.py +++ b/python/tests/test_arrow_handoff.py @@ -2455,23 +2455,22 @@ def _assert_float64_columns_guard_payload(state, schema, values): @pytest.mark.parametrize("rows", [10000, 100000]) @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", "arrow_class_iter", ]) @pytest.mark.parametrize("behavior", ["raise", "mutate"]) -def test_float64_columns_tolist_hooks_keep_native_order( +def test_float64_columns_iterator_hooks_keep_native_order( protocol, monkeypatch, rows, location, behavior ): - from types import MethodType from pandas.core.base import IndexOpsMixin values = np.arange(rows, dtype=np.float64) frame = pd.DataFrame({"original": values}) schema = T.StructType([T.StructField("renamed", T.DoubleType())]) - 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 @@ -2484,59 +2483,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] = -0.0, np.nan - 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(frame.columns._values, 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 in ("backing_instance", "backing_to_numpy"): - frame.columns = pd.Index(pd.array(["original"], dtype="string")) + 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_iter(array): + if array is backing: + before_hook() + return original(array) + + 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) - def backing_tolist(*args, **kwargs): - before_tolist() - return original(*args, **kwargs) + monkeypatch.setattr(type(backing), "__getitem__", backing_getitem) + else: + original = backing._box_func - monkeypatch.setitem(vars(backing), hook_name, backing_tolist) - assert type(frame.columns) is pd.Index + 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 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 @@ -2545,7 +2567,7 @@ def backing_tolist(*args, **kwargs): expected = values.copy() expected[0], expected[-1] = -0.0, np.nan _assert_float64_columns_guard_payload(state, schema, [expected]) - assert calls == [hook_name, "notna", "astype"] + assert calls == ["columns iterator", "notna", "astype"] @@ -2583,6 +2605,7 @@ def test_float64_rebound_index_data_keeps_native_read_order(protocol, monkeypatc original_notna, original_astype = frame.notna, frame.astype calls = [] normalizing = False + allowed_reads = None class ColumnsDataError(RuntimeError): pass @@ -2601,8 +2624,8 @@ def data(index): if index is columns: if not normalizing: calls.append("_data") - if calls.count("_data") > 1: - raise ColumnsDataError("float 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") # Preserve the native first columns read and only then attach # observers; no private optimization call determines fallback. monkeypatch.setitem(vars(frame), "notna", notna) @@ -2614,10 +2637,36 @@ 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_float64_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_float64_columns_guard_payload(state, schema, values) + assert calls == baseline_calls + import json + + assert state["config_requests"] == [] + assert state["uploads"] == 0 + assert len(state["input_tables"]) == 2 + assert len(state["relations"]) == 2 + assert state["input_tables"][1].equals(state["input_tables"][0], check_metadata=True) + for table, relation in zip(state["input_tables"], state["relations"]): + _assert_float64_ipc(table, values, schema) + assert json.loads(relation.local_relation.schema) == json.loads(schema.json()) From ccc1b481ad89f78ea3861464fb5dd0dd61416d75 Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Thu, 5 Nov 2026 16:11:06 -0800 Subject: [PATCH 34/47] [CONNECT] Reconcile remaining change with current inference architecture Co-authored-by: Isaac --- crates/pyspark-rs/src/session.rs | 18 ++++++++++++++++++ python/tests/test_arrow_handoff.py | 3 +++ 2 files changed, 21 insertions(+) diff --git a/crates/pyspark-rs/src/session.rs b/crates/pyspark-rs/src/session.rs index 943115e2..3aee7aee 100644 --- a/crates/pyspark-rs/src/session.rs +++ b/crates/pyspark-rs/src/session.rs @@ -788,6 +788,24 @@ impl PySparkSession { _ => frame, }); } + let table = py + .import("pyspark.sql._pandas_input")? + .getattr("prepare_date_table")? + .call1((data, schema))?; + if !table.is_none() { + let frame = + crate::pyarrow_import::create_dataframe(py, &self.session, &table, schema)?; + 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()), + ) + } + _ => 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 058aa095..622a19b1 100644 --- a/python/tests/test_arrow_handoff.py +++ b/python/tests/test_arrow_handoff.py @@ -1881,7 +1881,9 @@ def no_rows(*args, **kwargs): overrides = {"spark.sql.session.localRelationCacheThreshold": "1" if cached else "999999999"} with _arrow_input_service(protocol, **overrides) as (spark, state): df = spark.createDataFrame(frame, schema, verifySchema=verify) + assert not df.isLocal() df.toArrow() + assert state["column_names"] == schema.names table = state["input_tables"][0] assert table.num_rows == len(dates) assert table.schema.names == schema.names @@ -2421,6 +2423,7 @@ def _assert_date_columns_guard_payload(state, schema, expected, *, row_path): from pyspark.sql.pandas.types import from_arrow_metadata assert len(state["input_tables"]) == 1 + assert state["column_names"] == schema.names table = state["input_tables"][0] assert table.num_rows == 10000 assert table.column_names == ["renamed.`day", "nullable", "all_null"] From 26be6abe1cceff131004795791af651422eb5d81 Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Fri, 30 Oct 2026 02:49:42 -0700 Subject: [PATCH 35/47] [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 36/47] [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 37/47] 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 38/47] 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 cfde2b6349181ab720978aabbe195f81ff2800a2 Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Thu, 5 Nov 2026 16:05:11 -0800 Subject: [PATCH 39/47] [CONNECT] Reconcile remaining change with current inference architecture Co-authored-by: Isaac --- crates/pyspark-rs/src/session.rs | 25 +++++++++++++++++++++++++ python/tests/test_arrow_handoff.py | 16 +++++++++++++--- 2 files changed, 38 insertions(+), 3 deletions(-) diff --git a/crates/pyspark-rs/src/session.rs b/crates/pyspark-rs/src/session.rs index 4edbd2a9..bfe3c624 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_float64_input")? + .call1((data, schema))?; + if !ipc.is_none() { + let fields = match &spec { + Spec::Struct(fields) => fields, + _ => { + return Err(pyo3::exceptions::PyValueError::new_err( + "float64 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 3b66e19f..3c61a058 100644 --- a/python/tests/test_arrow_handoff.py +++ b/python/tests/test_arrow_handoff.py @@ -1773,7 +1773,10 @@ def no_rows(*args, **kwargs): monkeypatch.setattr(pd.Series, "tolist", no_rows) overrides = {"spark.sql.session.localRelationCacheThreshold": "1" if cached else "999999999"} with _arrow_input_service(protocol, config_error=True, **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 _assert_float64_ipc(state["input_tables"][0], columns, schema) assert state["uploads"] == 0 assert state["config_requests"] == [] @@ -1789,9 +1792,15 @@ def test_float64_columnar_ipc_matches_existing_row_conversion(protocol, monkeypa frame = pd.DataFrame({"original": values}) schema = T.StructType([T.StructField("renamed", T.DoubleType(), True, {"source": "parity"})]) with _arrow_input_service(protocol) 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 monkeypatch.setattr(helper, "prepare_float64_table", lambda *args: None) - spark.createDataFrame(frame, schema).toArrow() + result = spark.createDataFrame(frame, schema) + assert not result.isLocal() + result.toArrow() + assert state["column_names"] == schema.names assert len(state["input_tables"]) == 2 for table in state["input_tables"]: _assert_float64_ipc(table, [values], schema) @@ -2446,6 +2455,7 @@ def _assert_float64_columns_guard_payload(state, schema, values): assert len(state["input_tables"]) == 1 assert len(state["relations"]) == 1 + assert state["column_names"] == schema.names _assert_float64_ipc(state["input_tables"][0], values, schema) assert json.loads(state["relations"][0].local_relation.schema) == json.loads(schema.json()) assert state["config_requests"] == [] From e31aad6482c1b748924991a8c01aabbea74db385 Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Thu, 5 Nov 2026 16:05:12 -0800 Subject: [PATCH 40/47] [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] From 053bbcf10e9a52843a9559e098bcce8576b5e8d4 Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Thu, 5 Nov 2026 16:21:41 -0800 Subject: [PATCH 41/47] [CONNECT] Reconcile remaining change with current inference architecture Co-authored-by: Isaac --- python/tests/test_dropin_offline.py | 27 ++++++++++++++++++++++++--- 1 file changed, 24 insertions(+), 3 deletions(-) diff --git a/python/tests/test_dropin_offline.py b/python/tests/test_dropin_offline.py index d867c040..b52aa552 100644 --- a/python/tests/test_dropin_offline.py +++ b/python/tests/test_dropin_offline.py @@ -4913,12 +4913,33 @@ def __index__(self): def _exact_date_converted_value(mock_connect, value): import pyarrow as pa + pd = pytest.importorskip("pandas") spark, service = mock_connect - spark.createDataFrame([(value,)]).toArrow() - request = parse_message(service.execute_requests[-1]) + # py_to_value remains the generic pandas converter. Local rows now use + # infer.convert, which intentionally has different subclass/type semantics. + data = pd.DataFrame({"value": [value]}, dtype=object) + before = len(service.execute_requests) + spark.createDataFrame(data).toArrow() + requests = service.execute_requests + assert len(requests) == before + 1 + request = parse_message(requests[-1]) relation = parse_message(parse_message(request[3][0])[1][0]) + renames = [] + while 18 in relation: # Relation.to_df -> ToDF.input; column_names=2. + rename = parse_message(relation[18][0]) + renames.append([name.decode() for name in rename.get(2, [])]) + relation = parse_message(rename[1][0]) local = parse_message(relation[11][0]) - return pa.ipc.open_stream(local[1][0]).read_all().column(0)[0].as_py() + table = pa.ipc.open_stream(local[1][0]).read_all() + table.validate(full=True) + assert (table.num_rows, table.num_columns) == (1, 1) + assert table.column_names == ["value"] + assert [field["name"] for field in json.loads(local[2][0])["fields"]] == table.column_names + for names in renames: + assert len(names) == table.num_columns + # The outermost toDF supplies observable names without rewriting raw IPC. + assert (renames[0] if renames else table.column_names) == ["value"] + return table.column(0)[0].as_py() @pytest.mark.parametrize("value", [ From 9135b9b8b51745aa9e370c66e1ca917a6d578ede Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Thu, 5 Nov 2026 16:28:59 -0800 Subject: [PATCH 42/47] [CONNECT] Reconcile remaining change with current inference architecture Co-authored-by: Isaac --- python/tests/test_dropin_offline.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/python/tests/test_dropin_offline.py b/python/tests/test_dropin_offline.py index b52aa552..08a14164 100644 --- a/python/tests/test_dropin_offline.py +++ b/python/tests/test_dropin_offline.py @@ -5014,9 +5014,10 @@ def test_exact_date_nested_unrelated_values(mock_connect): def test_exact_date_named_impostor_rejected(): + pd = pytest.importorskip("pandas") impostor = type("date", (), {"toordinal": lambda self: 719163}) with pytest.raises(TypeError, match="unsupported value type"): - _local_create_dataframe([(impostor(),)]) + _local_create_dataframe(pd.DataFrame({"value": [impostor()]}, dtype=object)) def _exact_date_fields(data): From 489a663d8476fad7fc3e9b5eba43ee3c8fed7194 Mon Sep 17 00:00:00 2001 From: Hyukjin Kwon Date: Fri, 6 Nov 2026 20:00:22 +0900 Subject: [PATCH 43/47] [SPARK-60727][CONNECT] Drop the converter guards that the float64 path never reaches, and document the path The float64 path builds its Arrow table and IPC itself, so the checks of the converters of pyspark.sql.pandas.types and of the Arrow importer guarded code it does not call, and imported pyspark.sql.pandas.types for nothing. Drop them, with the three tests that only exercised them (they pass with either path), and describe the path in the architecture docs. Co-authored-by: Isaac --- docs/architecture.md | 4 ++ python/pyspark/sql/_pandas_input.py | 25 ------------- python/tests/test_arrow_handoff.py | 58 ----------------------------- 3 files changed, 4 insertions(+), 83 deletions(-) diff --git a/docs/architecture.md b/docs/architecture.md index d7f0ca5c..ee734f0d 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -155,6 +155,10 @@ and explicit `StructType` schemas whose fields are all `LongType`. Other dtypes, 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. +Pandas frames of at least 100,000 rows whose columns are all NumPy `float64`, +with an explicit `StructType` of `DoubleType` fields, are encoded to Arrow IPC +from their columns, with pandas' missing values as nulls, and sent inline as the +row path sends them, without the importer's configuration RPC. 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 index 5a07973c..2d041e3e 100644 --- a/python/pyspark/sql/_pandas_input.py +++ b/python/pyspark/sql/_pandas_input.py @@ -353,31 +353,6 @@ def prepare_float64_table(frame, schema): ) or len(set(schema.names)) != width: return None - from pyspark.sql import _arrow_input - from pyspark.sql.pandas import types as pandas_types - - for name in ("to_arrow_schema", "to_arrow_type", "to_arrow_metadata", - "_check_arrow_table_timestamps_localize"): - converter = getattr(pandas_types, name) - if (type(converter) is not FunctionType or converter.__name__ != name - or converter.__module__ != pandas_types.__name__ - or converter.__globals__ is not vars(pandas_types)): - return None - # Earlier converter branches can reinterpret DoubleType after alias rebinding. - for name in ("BooleanType", "ByteType", "ShortType", "IntegerType", "LongType", - "FloatType", "DoubleType", "StructField", "StructType"): - cls = getattr(native, name) - if (getattr(sql_types, name, None) is not cls - or getattr(pandas_types, name, None) is not cls): - return None - if (_arrow_input.StructType is not native.StructType - or getattr(_arrow_input, "to_arrow_schema", pandas_types.to_arrow_schema) - is not pandas_types.to_arrow_schema - or getattr(_arrow_input, "_check_arrow_table_timestamps_localize", - pandas_types._check_arrow_table_timestamps_localize) - is not pandas_types._check_arrow_table_timestamps_localize): - return None - # Use the same pandas missing-value semantics as the row path, including the # legacy use_inf_as_na option. Infinities and signed zero otherwise stay intact. valid = frame.notna() diff --git a/python/tests/test_arrow_handoff.py b/python/tests/test_arrow_handoff.py index 3be64d98..bbfb4b4f 100644 --- a/python/tests/test_arrow_handoff.py +++ b/python/tests/test_arrow_handoff.py @@ -2338,19 +2338,6 @@ def mktemp(self, name): -def test_float64_rebound_arrow_schema_alias_keeps_native_row_path(protocol, monkeypatch): - import pyspark.sql._arrow_input as helper - - frame = pd.DataFrame({"x": np.arange(100000, dtype=np.float64)}) - schema = T.StructType([T.StructField("x", T.DoubleType())]) - monkeypatch.setattr(helper, "StructType", object) - with _arrow_input_service(protocol) as (spark, state): - spark.createDataFrame(frame, schema).toArrow() - assert state["config_requests"] == [] - _assert_float64_ipc(state["input_tables"][0], [frame["x"].to_numpy()], schema) - - - def test_float64_rebound_pandas_alias_keeps_subclass_fallback(monkeypatch): from pyspark.sql._pandas_input import prepare_float64_table @@ -2486,30 +2473,6 @@ def customized(): -@pytest.mark.parametrize("module_name", ["converter", "arrow_input"]) -@pytest.mark.parametrize("attribute", ["to_arrow_schema", "_check_arrow_table_timestamps_localize"]) -def test_float64_rebound_converter_function_keeps_row_fallback( - protocol, monkeypatch, module_name, attribute -): - import pyspark.sql._arrow_input as arrow_input - import pyspark.sql.pandas.types as converter - - values = np.tile(np.array([np.nan, -0.0, np.inf, -np.inf, 1.5]), 20000) - frame = pd.DataFrame({"x": values}) - schema = T.StructType([T.StructField("x", T.DoubleType())]) - - def customized(*args, **kwargs): - raise AssertionError("changed Arrow converter must keep float64 input on the row path") - - monkeypatch.setattr({"converter": converter, "arrow_input": arrow_input}[module_name], - attribute, customized, raising=False) - with _arrow_input_service(protocol) as (spark, state): - spark.createDataFrame(frame, schema).toArrow() - assert state["config_requests"] == [] - _assert_float64_ipc(state["input_tables"][0], [values], schema) - - - @pytest.mark.parametrize("method", ["astype", "notna", "where"]) @pytest.mark.parametrize("before_import", [False, True]) def test_float64_class_conversion_hooks_keep_row_fallback( @@ -2567,27 +2530,6 @@ def customized(self, *args, **kwargs): -@pytest.mark.parametrize("attribute", ["to_arrow_type", "to_arrow_metadata"]) -def test_float64_custom_type_and_metadata_converters_keep_row_fallback( - protocol, monkeypatch, attribute -): - import pyspark.sql.pandas.types as converter - - values = np.tile(np.array([np.nan, -0.0, np.inf, -np.inf, 1.5]), 20000) - frame = pd.DataFrame({"x": values}) - schema = T.StructType([T.StructField("x", T.DoubleType())]) - - def customized(*args, **kwargs): - raise AssertionError("changed type or metadata converter must retain the row path") - - monkeypatch.setattr(converter, attribute, customized) - with _arrow_input_service(protocol) as (spark, state): - spark.createDataFrame(frame, schema).toArrow() - assert state["config_requests"] == [] - _assert_float64_ipc(state["input_tables"][0], [values], schema) - - - def test_float64_literal_numpy_dtype_custom_backing_keeps_row_error(protocol): from pyspark.sql._pandas_input import prepare_float64_table From 9fbcd15ad4e9d10c4d01e15f9e78dd63747d1dca Mon Sep 17 00:00:00 2001 From: Hyukjin Kwon Date: Fri, 6 Nov 2026 20:07:18 +0900 Subject: [PATCH 44/47] [SPARK-60727][CONNECT] Do not pin the reads of master's pandas row path in the float64 test As in the int64 and timestamp tests: a later row conversion of the same frame reads Index._data fewer times than the first, as pandas caches some reads. Co-authored-by: Isaac --- python/tests/test_arrow_handoff.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/python/tests/test_arrow_handoff.py b/python/tests/test_arrow_handoff.py index bbfb4b4f..a23042b2 100644 --- a/python/tests/test_arrow_handoff.py +++ b/python/tests/test_arrow_handoff.py @@ -2824,9 +2824,8 @@ def set_data(index, value): baseline.setattr(helper, "prepare_float64_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"] + allowed_reads = baseline_calls.count("_data") + assert allowed_reads > 0 and "notna" in baseline_calls calls.clear() normalizing = False # Restore canonical methods before the candidate eligibility check; @@ -2834,7 +2833,11 @@ def set_data(index, value): monkeypatch.delitem(vars(frame), "notna") monkeypatch.delitem(vars(frame), "astype") spark.createDataFrame(frame, schema).toArrow() - assert calls == baseline_calls + # The rebound storage declines the columnar path without reading the columns + # more than the row path does. pandas caches some of these reads, so a later + # row conversion of the same frame can read them fewer times than the first. + assert "notna" in calls + assert calls.count("_data") <= allowed_reads import json assert state["config_requests"] == [] From 6eca4a4a2eb252d0e3df9c8c60a9e1dda64bfee3 Mon Sep 17 00:00:00 2001 From: Hyukjin Kwon Date: Fri, 6 Nov 2026 20:18:38 +0900 Subject: [PATCH 45/47] [SPARK-60720][CONNECT] Encode each pandas string once, and drop the converter checks the string path never reaches _pandas_string_ipc encoded every string to UTF-8 in a validation pass and again while building; build in one pass and fall back on the first string that cannot be encoded. prepare_string_input calls the Rust encoder with the native schema, so the checks of the aliases of the Arrow importer and of pyspark.sql.pandas.types guarded code it does not call; drop them, with the two tests that only exercised them. Co-authored-by: Isaac --- crates/pyspark-rs/src/pandas_input.rs | 21 +----- python/pyspark/sql/_pandas_input.py | 24 +----- python/tests/test_arrow_handoff.py | 102 -------------------------- 3 files changed, 4 insertions(+), 143 deletions(-) diff --git a/crates/pyspark-rs/src/pandas_input.rs b/crates/pyspark-rs/src/pandas_input.rs index ab6eb8fc..4d02adba 100644 --- a/crates/pyspark-rs/src/pandas_input.rs +++ b/crates/pyspark-rs/src/pandas_input.rs @@ -56,24 +56,9 @@ pub(crate) fn _pandas_string_ipc<'py>( } } } - // 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); - } - } - } + // Encode each string once, while building. A lone surrogate cannot be encoded: it + // falls back to the row path, which drops the partial Arrow buffers. The type pass + // above keeps a late subclass from paying for the build. 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) { diff --git a/python/pyspark/sql/_pandas_input.py b/python/pyspark/sql/_pandas_input.py index fed54358..c5c7466c 100644 --- a/python/pyspark/sql/_pandas_input.py +++ b/python/pyspark/sql/_pandas_input.py @@ -394,8 +394,6 @@ def prepare_string_input(frame, schema): 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 @@ -403,10 +401,7 @@ def prepare_string_input(frame, schema): 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._pyspark import StringType, StructField, StructType as NativeStructType from pyspark.sql import types as sql_types # Keep small frames on the established row path. @@ -425,23 +420,6 @@ def prepare_string_input(frame, schema): 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), diff --git a/python/tests/test_arrow_handoff.py b/python/tests/test_arrow_handoff.py index 2a5a4045..67213093 100644 --- a/python/tests/test_arrow_handoff.py +++ b/python/tests/test_arrow_handoff.py @@ -3476,108 +3476,6 @@ class Frame(pd.DataFrame): -@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 6104f543cd0dd35c19a6645e00c1c8f702bb368b Mon Sep 17 00:00:00 2001 From: Hyukjin Kwon Date: Fri, 6 Nov 2026 20:25:22 +0900 Subject: [PATCH 46/47] [SPARK-60720][CONNECT] Do not pin the reads of master's pandas row path in the string test As in the int64, timestamp and float64 tests: a later row conversion of the same frame reads Index._data fewer times than the first, as pandas caches some reads. Co-authored-by: Isaac --- python/tests/test_arrow_handoff.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/python/tests/test_arrow_handoff.py b/python/tests/test_arrow_handoff.py index 67213093..c7965e32 100644 --- a/python/tests/test_arrow_handoff.py +++ b/python/tests/test_arrow_handoff.py @@ -3084,9 +3084,8 @@ def set_data(index, value): 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"] + allowed_reads = baseline_calls.count("_data") + assert allowed_reads > 0 and "notna" in baseline_calls calls.clear() normalizing = False # Restore canonical methods before the candidate eligibility check; @@ -3094,7 +3093,11 @@ def set_data(index, value): monkeypatch.delitem(vars(frame), "notna") monkeypatch.delitem(vars(frame), "astype") spark.createDataFrame(frame, schema).toArrow() - assert calls == baseline_calls + # The rebound storage declines the string path without reading the columns + # more than the row path does. pandas caches some of these reads, so a later + # row conversion of the same frame can read them fewer times than the first. + assert "notna" in calls + assert calls.count("_data") <= allowed_reads assert not state["config_requests"] assert state["uploads"] == 0 assert len(state["input_tables"]) == 2 From adc009b3daff68864a00dcdc231bc549f78f7bfb Mon Sep 17 00:00:00 2001 From: Hyukjin Kwon Date: Fri, 6 Nov 2026 21:17:18 +0900 Subject: [PATCH 47/47] [SPARK-60716][CONNECT] Serve the Arrow importer's settings in the exact date capture mock, and do not pin the reads of master's pandas row path The explicit DateType pandas case of test_exact_date_ipc_calendar_range has more than 10,000 rows, so it now goes through the Arrow importer, which reads its settings with a Config RPC that the capture mock did not serve. As in the other pandas tests, a later row conversion of the same frame reads Index._data fewer times than the first. Co-authored-by: Isaac --- python/tests/test_arrow_handoff.py | 11 +++++++---- python/tests/test_dropin_offline.py | 24 ++++++++++++++++++++++++ 2 files changed, 31 insertions(+), 4 deletions(-) diff --git a/python/tests/test_arrow_handoff.py b/python/tests/test_arrow_handoff.py index 19e0f577..a41a792f 100644 --- a/python/tests/test_arrow_handoff.py +++ b/python/tests/test_arrow_handoff.py @@ -5129,9 +5129,8 @@ def set_data(index, value): _assert_date_columns_guard_payload(state, schema, values, row_path=True) control_table = state["input_tables"][0] baseline_calls = calls.copy() - allowed_reads = baseline_calls.index("notna") - assert allowed_reads > 0 - assert baseline_calls == ["_data"] * allowed_reads + ["notna", "astype"] + allowed_reads = baseline_calls.count("_data") + assert allowed_reads > 0 and "notna" in baseline_calls calls.clear() normalizing = False # Restore ordinary methods before candidate eligibility; the storage hook @@ -5140,7 +5139,11 @@ def set_data(index, value): monkeypatch.delitem(vars(frame), "astype") with _arrow_input_service(protocol) as (spark, state): spark.createDataFrame(frame, schema).toArrow() - assert calls == baseline_calls + # The rebound storage declines the Arrow path without reading the columns more + # than the row path does. pandas caches some of these reads, so a later row + # conversion of the same frame can read them fewer times than the first. + assert "notna" in calls + assert calls.count("_data") <= allowed_reads _assert_date_columns_guard_payload(state, schema, values, row_path=True) assert state["input_tables"][0].equals(control_table, check_metadata=True) diff --git a/python/tests/test_dropin_offline.py b/python/tests/test_dropin_offline.py index 4228baa8..f74c8b04 100644 --- a/python/tests/test_dropin_offline.py +++ b/python/tests/test_dropin_offline.py @@ -4856,11 +4856,35 @@ def _exact_date_execute(request, context): context.abort(grpc.StatusCode.INVALID_ARGUMENT, "captured date input") yield b"" # Make this a streaming handler. + # Explicit DateType pandas input of 10,000 rows or more goes through the Arrow + # importer, which reads these settings first; keep its relation inline. + confs = { + "spark.sql.timestampType": "TIMESTAMP_LTZ", + "spark.sql.session.timeZone": "UTC", + "spark.sql.execution.arrow.useLargeVarTypes": "false", + "spark.sql.session.localRelationCacheThreshold": str(1 << 60), + "spark.sql.session.localRelationSizeLimit": str(1 << 60), + "spark.sql.session.localRelationChunkSizeRows": str(1 << 30), + "spark.sql.session.localRelationChunkSizeBytes": str(1 << 40), + "spark.sql.session.localRelationBatchOfChunksSizeBytes": str(1 << 40), + } + + def _exact_date_config(request, context): + # ConfigRequest.session_id=1, operation=3, Operation.get=2, Get.keys=1; + # ConfigResponse.session_id=1, pairs=2, KeyValue.key=1, value=2. + request = parse_message(request) + keys = [key.decode() for key in + parse_message(parse_message(request[3][0])[2][0]).get(1, [])] + return message_field(1, request[1][0]) + b"".join( + message_field(2, message_field(1, key) + message_field(2, confs[key])) + for key in keys) + with ThreadPoolExecutor(max_workers=2) as executor: server = grpc.server(executor) server.add_generic_rpc_handlers((grpc.method_handlers_generic_handler( "spark.connect.SparkConnectService", { "ExecutePlan": grpc.unary_stream_rpc_method_handler(_exact_date_execute), + "Config": grpc.unary_unary_rpc_method_handler(_exact_date_config), }),)) port = server.add_insecure_port("127.0.0.1:0") server.start()