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 459a65e1..c7a9c86a 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,209 @@ 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 + + +@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..e244b746 100644 --- a/python/tests/test_dropin_offline.py +++ b/python/tests/test_dropin_offline.py @@ -4030,6 +4030,313 @@ 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]) +# 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 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()), + "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 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. +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())]) + + 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) + ) + 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"] + 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): + 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