From 13e35585cab99d1d86e153f3a736035908247986 Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Thu, 29 Oct 2026 01:33:49 -0700 Subject: [PATCH 1/3] [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 | 137 ++++++++++++++++++++++++++++- 1 file changed, 133 insertions(+), 4 deletions(-) diff --git a/python/tests/test_arrow_handoff.py b/python/tests/test_arrow_handoff.py index 459a65e1..2747df6b 100644 --- a/python/tests/test_arrow_handoff.py +++ b/python/tests/test_arrow_handoff.py @@ -424,7 +424,8 @@ def _arrow_input_service( } config.update(overrides) state = {"cache": {}, "uploads": 0, "relations": [], "message_sizes": [], - "config_requests": [], "upload_contexts": [], "status_contexts": []} + "config": config, "config_requests": [], "input_tables": [], + "upload_contexts": [], "status_contexts": []} server_id = "00000000-0000-0000-0000-000000000001" def attribution(request, context): @@ -444,10 +445,13 @@ 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 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() @@ -524,6 +528,7 @@ def execute(raw, context): ) rows, first, last = 0, None, None state["arrow_schemas"], state["batch_rows"] = [], [] + input_batches = [] for payload in payloads: reader = pa.ipc.open_stream(payload) state["arrow_schemas"].append(reader.schema) @@ -531,11 +536,14 @@ def execute(raw, context): state["batch_rows"].append(stream_rows) for batch in reader: stream_rows.append(batch.num_rows) + input_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["input_tables"].append(pa.Table.from_batches( + input_batches, schema=state["arrow_schemas"][0])) state["result"] = (rows, first, last) result = pa.table({"rows": [rows]}) sink = pa.BufferOutputStream() @@ -889,3 +897,124 @@ def import_module(name, *args, **kwargs): spark.createDataFrame(data) assert imported == [] assert state["config_requests"] == [] + + +_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 1f792d9d1dabd7c97eea3bbd67af82760e788a16 Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Thu, 29 Oct 2026 10:01:35 -0700 Subject: [PATCH 2/3] [SPARK-60717][CONNECT] Skip redundant canonical int64 Arrow normalization Avoid loading general converters and rebuilding canonical nullable int64 tables while preserving the general path for customized types, metadata, casts, and loaded converter hooks. Co-authored-by: Isaac --- python/pyspark/sql/_arrow_input.py | 80 +++++++- python/tests/test_arrow_handoff.py | 85 +++++++++ python/tests/test_dropin_offline.py | 278 ++++++++++++++++++++++++++++ 3 files changed, 439 insertions(+), 4 deletions(-) diff --git a/python/pyspark/sql/_arrow_input.py b/python/pyspark/sql/_arrow_input.py index 0c5b337f..948b29b5 100644 --- a/python/pyspark/sql/_arrow_input.py +++ b/python/pyspark/sql/_arrow_input.py @@ -18,15 +18,78 @@ """Schema normalization for the native Arrow Table input path.""" +import sys +from types import GetSetDescriptorType, WrapperDescriptorType + from pyspark.errors import PySparkTypeError, PySparkValueError -from pyspark.sql.pandas.types import ( - _check_arrow_table_timestamps_localize, - from_arrow_schema, - to_arrow_schema, +from pyspark._pyspark import ( + AtomicType as NativeAtomicType, + BooleanType, + ByteType, + IntegerType, + LongType, + ShortType, + StructField, + StructType as NativeStructType, ) +from pyspark.sql import types as sql_types from pyspark.sql.types import AtomicType, StructType +def _canonical_long_table(table, schema): + # A loaded converter may have customized hooks. Keep its normal behavior, + # and avoid importing it solely to normalize an already canonical table. + if ("pyspark.sql.pandas.types" in sys.modules + or sys.modules.get("pyspark.sql.types") is not sql_types + or AtomicType is not NativeAtomicType + or StructType is not NativeStructType + or type(schema) is not NativeStructType): + return False + for name, cls in (("AtomicType", NativeAtomicType), + ("BooleanType", BooleanType), ("ByteType", ByteType), + ("ShortType", ShortType), ("IntegerType", IntegerType), + ("LongType", LongType), ("StructField", StructField), + ("StructType", NativeStructType)): + if getattr(sql_types, name, None) is not cls: + return False + if LongType.__getattribute__ is not object.__getattribute__: + return False + # Native classes reject subclasses, but their class attributes are mutable. + # Reject customized access/iteration, including changes made before import. + for cls, names, descriptor_type in ( + (NativeStructType, ("fields", "names"), GetSetDescriptorType), + (NativeStructType, ("__iter__",), WrapperDescriptorType), + (StructField, ("name", "dataType", "nullable", "metadata"), GetSetDescriptorType), + ): + if cls.__getattribute__ is not object.__getattribute__: + return False + 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 False + + import pyarrow as pa + + if type(table) is not pa.Table: + return False + arrow_schema = table.schema + fields = schema.fields + if (not fields or len(fields) != len(arrow_schema) + or arrow_schema.metadata is not None): + return False + names = set() + for field, arrow_field in zip(fields, arrow_schema): + if (type(field) is not StructField or type(field.dataType) is not LongType + or field.nullable is not True or field.metadata + or field.name != arrow_field.name or field.name in names + or arrow_field.type != pa.int64() or arrow_field.nullable is not True + or arrow_field.metadata is not None): + return False + names.add(field.name) + return True + + def prepare_table(table, schema, timestamp_type, timezone, large_types): """Match Connect's Arrow casts without converting any column to Python rows.""" if schema is not None and not isinstance(schema, (AtomicType, StructType, list, tuple)): @@ -59,6 +122,15 @@ def prepare_table(table, schema, timestamp_type, timezone, large_types): }, ) + if explicit and _canonical_long_table(table, schema): + return table, schema, None + + from pyspark.sql.pandas.types import ( + _check_arrow_table_timestamps_localize, + from_arrow_schema, + to_arrow_schema, + ) + if names is not None: names.extend(f"_{i + 1}" for i in range(len(names), table.num_columns)) inferred = ( diff --git a/python/tests/test_arrow_handoff.py b/python/tests/test_arrow_handoff.py index 2747df6b..c7a9c86a 100644 --- a/python/tests/test_arrow_handoff.py +++ b/python/tests/test_arrow_handoff.py @@ -1018,3 +1018,88 @@ def test_arrow_input_empty_skips_all_configs(protocol): spark.createDataFrame(table) assert state["config_requests"] == [] assert state["uploads"] == 0 + + +@pytest.mark.parametrize("rows", [10000, 30000, 60000, 1000000]) +def test_canonical_arrow_input_matches_normalized_ipc_and_uploads(protocol, monkeypatch, rows): + import json + import sys + import numpy as np + from pyspark.sql import pandas as pandas_package + from pyspark.sql.types import LongType, StructField, StructType + + values = np.arange(rows, dtype=np.int64) + values[0], values[-1] = -(2**63), 2**63 - 1 + # Different column chunk boundaries also exercise native record batch alignment. + table = pa.table({ + "x": pa.chunked_array([values[:7], values[7:]]), + "y": pa.chunked_array([pa.array([None], type=pa.int64()), + pa.array(values[1:])]), + "z": pa.chunked_array([values[:31], values[31:]]), + }) + schema = StructType([StructField(name, LongType()) for name in table.column_names]) + monkeypatch.setattr(pandas_package, "types", getattr(pandas_package, "types", None), + raising=False) + monkeypatch.delitem(sys.modules, "pyspark.sql.pandas.types", raising=False) + with _arrow_input_service(protocol) as (spark, state): + spark.createDataFrame(table, schema).collect() + assert "pyspark.sql.pandas.types" not in sys.modules + first_uploads = state["uploads"] + first_cache = dict(state["cache"]) + # A loaded converter deliberately selects the established normalization path. + from importlib import import_module + + converters = import_module("pyspark.sql.pandas.types") + + assert converters is sys.modules["pyspark.sql.pandas.types"] + state["config"].update(dict(zip(_NORMALIZATION_KEYS, [ + "TIMESTAMP_NTZ", "America/Los_Angeles", "true"]))) + spark.createDataFrame(table, schema).collect() + assert state["config_requests"] == [("get", _NORMALIZATION_KEYS + _UPLOAD_KEYS)] * 2 + assert state["uploads"] == first_uploads + assert state["cache"] == first_cache + for actual in state["input_tables"]: + assert actual.equals(table, check_metadata=True) + assert actual.to_pydict() == table.to_pydict() + first, second = state["relations"] + assert first.WhichOneof("rel_type") == second.WhichOneof("rel_type") + if table.nbytes < 1048576: + assert first.HasField("local_relation") + assert first.local_relation.data == second.local_relation.data + assert first.local_relation.schema == second.local_relation.schema + assert first_uploads == 0 + fields = json.loads(first.local_relation.schema)["fields"] + else: + assert first.HasField("chunked_cached_local_relation") + assert first.chunked_cached_local_relation == second.chunked_cached_local_relation + assert first_uploads == 1 + fields = json.loads(state["spark_schema"])["fields"] + assert fields == [dict(name=name, type="long", nullable=True, metadata={}) + for name in table.column_names] + + +@pytest.mark.parametrize("failure,match", [ + ("crc", "CRC check failed"), ("missing", "Missing cache artifact"), + ("limit", "LOCAL_RELATION_SIZE_LIMIT_EXCEEDED"), + ("config", "configuration read failed"), +]) +def test_canonical_arrow_input_preserves_failures(protocol, monkeypatch, failure, match): + import sys + from pyspark.sql.types import LongType, StructField, StructType + + monkeypatch.delitem(sys.modules, "pyspark.sql.pandas.types", raising=False) + options = {"spark.sql.session.localRelationCacheThreshold": "0"} + if failure == "limit": + options["spark.sql.session.localRelationSizeLimit"] = "10" + with _arrow_input_service(protocol, upload_error=failure, **options) as (spark, state): + if failure == "config": + state["config_error"] = grpc.StatusCode.INVALID_ARGUMENT + with pytest.raises(Exception, match=match): + spark.createDataFrame(pa.table({"x": [1, None, 2]}), + StructType([StructField("x", LongType())])) + assert "pyspark.sql.pandas.types" not in sys.modules + assert not state["relations"] + assert state["uploads"] == (1 if failure in ("crc", "missing") else 0) + assert state["config_requests"] == ( + [("get", _NORMALIZATION_KEYS + _UPLOAD_KEYS)] + ) diff --git a/python/tests/test_dropin_offline.py b/python/tests/test_dropin_offline.py index 39e06382..28b881a6 100644 --- a/python/tests/test_dropin_offline.py +++ b/python/tests/test_dropin_offline.py @@ -4030,6 +4030,284 @@ def test_arrow_input_explicit_cast_and_nested_values(): "TIMESTAMP_LTZ", "UTC", "false") +def _arrow_input_ipc_bytes(table): + 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() + + +@pytest.mark.parametrize("layout", ["nullable", "sliced", "chunked", "empty"]) +def test_arrow_input_canonical_int64_preserves_ipc(layout): + pytest.importorskip("pyarrow") + import subprocess + import sys + + subprocess.run([sys.executable, "-c", """ +import sys +import pyarrow as pa +from pyspark.sql import types as T +from pyspark.sql._arrow_input import prepare_table + +values = [-(2**63), None, 0, 2**63 - 1] +layout = sys.argv[1] +if layout == "sliced": + column = pa.array([99, *values, 99], type=pa.int64()).slice(1, len(values)) +elif layout == "chunked": + column = pa.chunked_array([values[:1], values[1:3], values[3:]], type=pa.int64()) +else: + column = pa.array([] if layout == "empty" else values, type=pa.int64()) +table = pa.table({"x": column, "y": column}) +schema = T.StructType([T.StructField(name, T.LongType()) for name in table.column_names]) +heavy_modules = ("pyspark.errors", "pyspark.sql.pandas.types") +assert all(name not in sys.modules for name in heavy_modules) +for large_types in ("false", "true", None): + normalized, result_schema, names = prepare_table( + table, schema, "TIMESTAMP_LTZ", "UTC", large_types + ) + assert normalized is table + assert result_schema is schema and names is None +assert all(name not in sys.modules for name in heavy_modules) + +# A noncanonical input must still load the general path after warm fast calls. +fallback = pa.table({"x": pa.array([1, None], type=pa.int32()), + "y": pa.array([2, None], type=pa.int32())}) +converted, result_schema, names = prepare_table(fallback, schema, "TIMESTAMP_LTZ", "UTC", "false") +assert all(column.type == pa.int64() for column in converted.columns) +assert converted.to_pylist() == fallback.to_pylist() +assert result_schema is schema and names is None +assert all(name in sys.modules for name in heavy_modules) + +# Import the general converter only after checking the fast path and repeated +# calls. Its cast is the oracle for cells, schema metadata, and IPC bytes. +from pyspark.sql.pandas.types import to_arrow_schema +expected = table.cast(to_arrow_schema(schema, timezone="UTC")).rename_columns(schema.names) +def ipc_bytes(value): + sink = pa.BufferOutputStream() + with pa.ipc.new_stream(sink, value.schema) as writer: + writer.write_table(value) + return sink.getvalue().to_pybytes() +assert normalized.schema.equals(expected.schema, check_metadata=True) +assert normalized.to_pylist() == expected.to_pylist() +assert ipc_bytes(normalized) == ipc_bytes(expected) +""", layout], check=True) + + +@pytest.mark.parametrize("kind", [ + "schema_metadata", "schema_metadata_empty", "field_metadata", "field_metadata_empty", + "source_nonnullable", "target_metadata", + "target_nonnullable", "different_names", "cast", "int32", "inferred", "names", + "custom_type", "duplicate_names", +]) +def test_arrow_input_noncanonical_uses_general_conversion(kind, monkeypatch): + import sys + + pa = pytest.importorskip("pyarrow") + from pyspark.sql._arrow_input import prepare_table + from pyspark.sql import pandas as pandas_package + from pyspark.sql.pandas import types as pandas_types + + table = pa.table({"x": pa.array([1, None, 2], type=pa.int64())}) + schema = T.StructType([T.StructField("x", T.LongType())]) + if kind in ("schema_metadata", "schema_metadata_empty"): + metadata = {} if kind == "schema_metadata_empty" else {b"source": b"test"} + table = table.replace_schema_metadata(metadata) + elif kind in ("field_metadata", "field_metadata_empty"): + metadata = {} if kind == "field_metadata_empty" else {b"source": b"test"} + table = table.cast(pa.schema([ + pa.field("x", pa.int64(), metadata=metadata) + ])) + elif kind == "source_nonnullable": + table = pa.table({"x": [1, 2]}, schema=pa.schema([ + pa.field("x", pa.int64(), nullable=False) + ])) + elif kind == "target_metadata": + schema = T.StructType([T.StructField("x", T.LongType(), metadata={"note": "test"})]) + elif kind == "target_nonnullable": + table = pa.table({"x": [1, 2]}) + schema = T.StructType([T.StructField("x", T.LongType(), nullable=False)]) + elif kind == "different_names": + schema = T.StructType([T.StructField("renamed", T.LongType())]) + elif kind == "cast": + schema = T.StructType([T.StructField("x", T.IntegerType())]) + elif kind == "int32": + table = pa.table({"x": pa.array([1, None, 2], type=pa.int32())}) + elif kind == "inferred": + schema = None + elif kind == "names": + schema = ["renamed"] + elif kind == "custom_type": + class CustomType(T.DataType): + pass + + schema = CustomType() + elif kind == "duplicate_names": + table = pa.Table.from_arrays([table.column(0), table.column(0)], names=["x", "x"]) + schema = T.StructType([T.StructField("x", T.LongType()), T.StructField("x", T.LongType())]) + + expected_schema = ( + schema if isinstance(schema, T.StructType) else pandas_types.from_arrow_schema(table.schema) + ) + expected_error = None + try: + expected = table.cast(pandas_types.to_arrow_schema( + expected_schema, error_on_duplicated_field_names_in_struct=True, timezone="UTC" + )).rename_columns(expected_schema.names) + except Exception as error: + expected_error = (type(error), str(error)) + # A previously imported converter already disables the fast path. Remove + # its module entry so each case proves its own guard rejects the input. + # Restore both import-system references after the test. + monkeypatch.setattr(pandas_package, "types", pandas_types) + monkeypatch.delitem(sys.modules, "pyspark.sql.pandas.types") + try: + normalized, result_schema, names = prepare_table( + table, schema, "TIMESTAMP_LTZ", "UTC", "false" + ) + except Exception as error: + assert expected_error == (type(error), str(error)) + assert "pyspark.sql.pandas.types" in sys.modules + return + assert expected_error is None + assert "pyspark.sql.pandas.types" in sys.modules + assert normalized.schema.equals(expected.schema, check_metadata=True) + assert normalized.to_pylist() == expected.to_pylist() + assert _arrow_input_ipc_bytes(normalized) == _arrow_input_ipc_bytes(expected) + if kind == "names": + assert names == ["renamed"] + else: + assert names is None + if schema is not None and kind != "names": + assert result_schema is schema + + +def test_arrow_input_respects_loaded_converter_hook(monkeypatch): + pa = pytest.importorskip("pyarrow") + from pyspark.sql._arrow_input import prepare_table + from pyspark.sql.pandas import types as pandas_types + + table = pa.table({"x": [1, None, 2]}) + schema = T.StructType([T.StructField("x", T.LongType())]) + calls = [] + + def custom_arrow_schema(actual_schema, **kwargs): + calls.append(actual_schema) + return pa.schema([pa.field("x", pa.int32())]) + + monkeypatch.setattr(pandas_types, "to_arrow_schema", custom_arrow_schema) + normalized, result_schema, names = prepare_table( + table, schema, "TIMESTAMP_LTZ", "UTC", "false" + ) + assert calls == [schema] + assert normalized.column(0).type == pa.int32() + assert normalized.to_pylist() == table.to_pylist() + assert result_schema is schema and names is None + + +@pytest.mark.parametrize("alias", [ + "StructType", "StructField", "BooleanType", "ByteType", "ShortType", + "IntegerType", "LongType", +]) +def test_arrow_input_respects_rebound_type_aliases(alias): + pytest.importorskip("pyarrow") + import subprocess + import sys + + subprocess.run([sys.executable, "-c", """ +import sys +import pyarrow as pa +from pyspark.sql import types as T +from pyspark.sql._arrow_input import prepare_table + +table = pa.table({"x": [1, None, 2]}) +schema = T.StructType([T.StructField("x", T.LongType())]) +alias = sys.argv[1] +original = getattr(T, alias) +replacement = T.DataType if alias in ("StructType", "StructField") else T.LongType +if alias == "LongType": + replacement = T.IntegerType + +def capture(operation): + try: + value = operation() + except Exception as error: + return ("error", type(error), str(error)) + sink = pa.BufferOutputStream() + with pa.ipc.new_stream(sink, value.schema) as writer: + writer.write_table(value) + return ("ok", value.schema.serialize().to_pybytes(), value.to_pylist(), sink.getvalue().to_pybytes()) + +assert "pyspark.sql.pandas.types" not in sys.modules +try: + setattr(T, alias, replacement) + actual = capture(lambda: prepare_table(table, schema, "TIMESTAMP_LTZ", "UTC", "false")[0]) + assert "pyspark.sql.pandas.types" in sys.modules + from pyspark.sql.pandas.types import to_arrow_schema + expected = capture(lambda: table.cast(to_arrow_schema( + schema, error_on_duplicated_field_names_in_struct=True, timezone="UTC" + )).rename_columns(schema.names)) + assert actual == expected, (alias, actual, expected) +finally: + setattr(T, alias, original) +""", alias], check=True) + + +@pytest.mark.parametrize("hook", [ + "names", "__iter__", "field_metadata", "long_getattribute", +]) +def test_arrow_input_respects_native_descriptors_changed_before_import(hook): + pytest.importorskip("pyarrow") + import subprocess + import sys + + subprocess.run([sys.executable, "-c", """ +import sys +import pyarrow as pa +from pyspark.sql import types as T + +table = pa.table({"x": [1, None, 2]}) +schema = T.StructType([T.StructField("x", T.LongType())]) +hook = sys.argv[1] +owner = T.StructType +attribute = hook +if hook == "names": + replacement = property(lambda self: ["renamed"]) +elif hook == "__iter__": + replacement = lambda self: iter([T.StructField("x", T.IntegerType())]) +elif hook == "field_metadata": + owner, attribute = T.StructField, "metadata" + replacement = property(lambda self: {"hook": "custom"}) +else: + owner, attribute = T.LongType, "__getattribute__" + def replacement(self, name): + return T.BooleanType if name == "__class__" else object.__getattribute__(self, name) +missing = object() +original = vars(owner).get(attribute, missing) + +try: + setattr(owner, attribute, replacement) + from pyspark.sql._arrow_input import prepare_table + assert "pyspark.sql.pandas.types" not in sys.modules + normalized, result_schema, names = prepare_table(table, schema, "TIMESTAMP_LTZ", "UTC", "false") + assert "pyspark.sql.pandas.types" in sys.modules + from pyspark.sql.pandas.types import to_arrow_schema + expected = table.cast(to_arrow_schema( + schema, error_on_duplicated_field_names_in_struct=True, timezone="UTC" + )).rename_columns(schema.names) + assert normalized.schema.equals(expected.schema, check_metadata=True) + assert normalized.to_pylist() == expected.to_pylist() + assert result_schema is schema + assert names is None +finally: + if original is missing: + delattr(owner, attribute) + else: + setattr(owner, attribute, original) +""", hook], check=True) + + def test_arrow_input_roundtrip_live_server(): """Run explicitly with SPARK_REMOTE; also checks imported buffers own their lifetime.""" import gc From 7bddb7535f88a62b1acdb66d117f968021de2762 Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Thu, 5 Nov 2026 15:51:06 -0800 Subject: [PATCH 3/3] [CONNECT] Update Arrow shortcut tests for current importer contracts Co-authored-by: Isaac --- python/tests/test_dropin_offline.py | 37 +++++++++++++++++++++++++---- 1 file changed, 33 insertions(+), 4 deletions(-) diff --git a/python/tests/test_dropin_offline.py b/python/tests/test_dropin_offline.py index 28b881a6..e244b746 100644 --- a/python/tests/test_dropin_offline.py +++ b/python/tests/test_dropin_offline.py @@ -4061,15 +4061,19 @@ def test_arrow_input_canonical_int64_preserves_ipc(layout): column = pa.array([] if layout == "empty" else values, type=pa.int64()) table = pa.table({"x": column, "y": column}) schema = T.StructType([T.StructField(name, T.LongType()) for name in table.column_names]) -heavy_modules = ("pyspark.errors", "pyspark.sql.pandas.types") -assert all(name not in sys.modules for name in heavy_modules) +# Errors are already imported during current client startup. Only the general +# Arrow converter remains deferred by this shortcut. +startup_errors = sys.modules["pyspark.errors"] +converter_modules = ("pyspark.sql.pandas.types",) +assert all(name not in sys.modules for name in converter_modules) for large_types in ("false", "true", None): normalized, result_schema, names = prepare_table( table, schema, "TIMESTAMP_LTZ", "UTC", large_types ) assert normalized is table assert result_schema is schema and names is None -assert all(name not in sys.modules for name in heavy_modules) +assert all(name not in sys.modules for name in converter_modules) +assert sys.modules["pyspark.errors"] is startup_errors # A noncanonical input must still load the general path after warm fast calls. fallback = pa.table({"x": pa.array([1, None], type=pa.int32()), @@ -4078,7 +4082,8 @@ def test_arrow_input_canonical_int64_preserves_ipc(layout): assert all(column.type == pa.int64() for column in converted.columns) assert converted.to_pylist() == fallback.to_pylist() assert result_schema is schema and names is None -assert all(name in sys.modules for name in heavy_modules) +assert all(name in sys.modules for name in converter_modules) +assert sys.modules["pyspark.errors"] is startup_errors # Import the general converter only after checking the fast path and repeated # calls. Its cast is the oracle for cells, schema metadata, and IPC bytes. @@ -4147,6 +4152,24 @@ class CustomType(T.DataType): table = pa.Table.from_arrays([table.column(0), table.column(0)], names=["x", "x"]) schema = T.StructType([T.StructField("x", T.LongType()), T.StructField("x", T.LongType())]) + if kind == "custom_type": + from pyspark.errors import PySparkTypeError + + # Current master validates supported schema categories before invoking + # general Arrow conversion. Preserve its structured error and ordering. + monkeypatch.setattr(pandas_package, "types", pandas_types) + monkeypatch.delitem(sys.modules, "pyspark.sql.pandas.types") + with pytest.raises(PySparkTypeError) as error: + prepare_table(table, schema, "TIMESTAMP_LTZ", "UTC", "false") + assert type(error.value) is PySparkTypeError + assert error.value.getCondition() == "NOT_EXPECTED_TYPE" + assert error.value.getMessageParameters() == { + "arg_name": "schema", "expected_type": "list, None or StructType", + "arg_type": "CustomType", + } + assert "pyspark.sql.pandas.types" not in sys.modules + return + expected_schema = ( schema if isinstance(schema, T.StructType) else pandas_types.from_arrow_schema(table.schema) ) @@ -4177,10 +4200,16 @@ class CustomType(T.DataType): assert _arrow_input_ipc_bytes(normalized) == _arrow_input_ipc_bytes(expected) if kind == "names": assert names == ["renamed"] + elif kind == "inferred": + assert names == table.column_names else: assert names is None if schema is not None and kind != "names": assert result_schema is schema + else: + # Inferred and renamed inputs keep their types in Arrow, with names + # applied by the current importer through a toDF wrapper. + assert result_schema is None def test_arrow_input_respects_loaded_converter_hook(monkeypatch):