[SPARK-60717][CONNECT] Skip redundant canonical int64 Arrow normalization

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 = (

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,18 +528,22 @@ 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)
             stream_rows = []
             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)]
+        )