From c8a573b0a6870955b82b07bdaa58d0e849164512 Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Wed, 28 Oct 2026 22:51:55 -0700 Subject: [PATCH 1/6] [SPARK-60684][CONNECT] Fast-path exact built-in string input --- crates/pyspark-rs/src/session.rs | 8 + python/tests/test_exact_string_input.py | 235 ++++++++++++++++++++++++ 2 files changed, 243 insertions(+) create mode 100644 python/tests/test_exact_string_input.py diff --git a/crates/pyspark-rs/src/session.rs b/crates/pyspark-rs/src/session.rs index 5f6302e3..faa148e3 100644 --- a/crates/pyspark-rs/src/session.rs +++ b/crates/pyspark-rs/src/session.rs @@ -955,6 +955,14 @@ pub(crate) fn py_to_value(obj: &Bound<'_, PyAny>) -> PyResult { return Ok(Value::Null); } + // Exact built-in strings cannot implement numeric conversion protocols. + // Keep subclasses and strings that fail UTF-8 extraction on the generic path. + if let Ok(s) = obj.cast_exact::() { + if let Ok(s) = s.extract::() { + return Ok(Value::String(s)); + } + } + // 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_exact_string_input.py b/python/tests/test_exact_string_input.py new file mode 100644 index 00000000..30cb7384 --- /dev/null +++ b/python/tests/test_exact_string_input.py @@ -0,0 +1,235 @@ +# 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. + +"""Guard exact-string conversion and decode the actual outbound Arrow IPC.""" +from concurrent.futures import ThreadPoolExecutor +import multiprocessing + +import pytest +from pyspark.sql import Row, SparkSession + + +@pytest.mark.parametrize("text", ["", "1.25", "-42", "NaN", "a\x00b", "é雪😀", "e\u0301", "\uffff"]) +def test_exact_string_roundtrip(text): + actual = Row(value=text).value + assert type(actual) is str + assert actual == text + + +@pytest.mark.parametrize("text", ["\ud800", "\udfff", "a\ud800b"]) +def test_surrogate_error_preserves_generic_fallback(text): + with pytest.raises(TypeError, match="unsupported value type for createDataFrame: str"): + Row(value=text) + + +def test_string_subclass_protocol_order(): + trace = [] + + class IndexString(str): + def __index__(self): + trace.append("index") + return 17 + def __float__(self): + raise AssertionError("index must take precedence") + + class FloatString(str): + def __float__(self): + trace.append("float") + return 2.5 + def __str__(self): + raise AssertionError("numeric conversion must take precedence") + + class FallbackString(str): + def __index__(self): + trace.append("bad index") + raise ValueError("index rejected") + def __float__(self): + trace.append("bad float") + raise ValueError("float rejected") + def __str__(self): + raise AssertionError("string extraction must retain the underlying text") + + class PlainString(str): + def __str__(self): + raise AssertionError("string extraction must not call __str__") + + for value, expected in [(IndexString("x"), 17), (FloatString("x"), 2.5), + (FallbackString("original"), "original"), + (PlainString("plain"), "plain")]: + actual = Row(value=value).value + assert type(actual) is type(expected) + assert actual == expected + assert trace == ["index", "float", "bad index", "bad float"] + + +def test_nested_strings_and_nulls(): + value = Row(items=["雪", None, ""], mapping={"x": "😀", "null": None}) + actual = Row(value=value).value + assert list(actual.items) == ["雪", None, ""] + assert actual.mapping == {"x": "😀", "null": None} + + +def _varint(value): + out = bytearray() + while value > 127: + out.append((value & 127) | 128) + value >>= 7 + out.append(value) + return bytes(out) + + +def _message(field, value): + return _varint(field * 8 + 2) + _varint(len(value)) + value + + +def _fields(data): + def read(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 = read(pos) + wire = tag & 7 + if wire == 2: + size, pos = read(pos) + result[tag >> 3] = data[pos:pos + size] + pos += size + elif wire == 0: + _, pos = read(pos) + else: + raise AssertionError(f"unexpected wire type {wire}") + return result + + +def _serve(connection): + import grpc + import pyarrow as pa + + state = {} + + def analyze(raw, context): + # AnalyzePlanRequest.schema.plan.root.local_relation.data + relation = _fields(_fields(_fields(_fields(raw)[4])[1])[1]) + payload = _fields(relation[11])[1] + table = pa.ipc.open_stream(payload).read_all() + state["result"] = (table.to_pylist(), str(table.schema)) + # Return an empty StructType; the test inspects the input IPC itself. + return _message(2, _message(1, _message(21, b""))) + + with ThreadPoolExecutor(max_workers=2) as executor: + server = grpc.server(executor) + server.add_generic_rpc_handlers((grpc.method_handlers_generic_handler( + "spark.connect.SparkConnectService", { + "AnalyzePlan": grpc.unary_unary_rpc_method_handler(analyze)}),)) + 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(state["result"]) + finally: + server.stop(0).wait() + + +@pytest.fixture +def input_service(): + pytest.importorskip("grpc") + pytest.importorskip("pyarrow") + context = multiprocessing.get_context("spawn") + parent, child = context.Pipe() + process = context.Process(target=_serve, args=(child,)) + process.start() + try: + assert parent.poll(20) + spark = SparkSession.builder.remote(f"sc://127.0.0.1:{parent.recv()}").getOrCreate() + yield spark, parent + 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_mode", ["inferred", "names", "explicit"]) +def test_string_input_ipc(input_service, source, schema_mode): + spark, parent = input_service + from pyspark.sql.types import StructType, StructField, StringType + import pandas as pd + + rows = [("雪😀", "a\x00b"), ("", None), (None, "1.25"), ("e\u0301", "\uffff")] + data = pd.DataFrame(rows, columns=["a", "b"]) if source == "pandas" else rows + schema = (StructType([StructField("a", StringType()), StructField("b", StringType())]) + if schema_mode == "explicit" else ["a", "b"] if schema_mode == "names" else None) + spark.createDataFrame(data, schema).schema + parent.send("get") + assert parent.poll(10) + actual, arrow_schema = parent.recv() + names = ["col0", "col1"] if source == "tuples" and schema_mode == "inferred" else ["a", "b"] + assert actual == [dict(zip(names, row)) for row in rows] + assert all(f"{name}: string" in arrow_schema for name in names) + + +def test_nested_string_input_ipc(input_service): + from pyspark.sql.types import (StructType, StructField, StringType, ArrayType, MapType) + spark, parent = input_service + schema = StructType([ + StructField("items", ArrayType(StringType())), + StructField("mapping", MapType(StringType(), StringType())), + StructField("nested", StructType([StructField("text", StringType())])), + ]) + rows = [(["雪", None, ""], {"x": "😀", "null": None}, Row(text="a\x00b")), + (None, None, None)] + spark.createDataFrame(rows, schema).schema + parent.send("get") + assert parent.poll(10) + actual, _ = parent.recv() + # Arrow map columns decode as ordered key/value pairs. + assert actual == [ + {"items": ["雪", None, ""], "mapping": [("null", None), ("x", "😀")], + "nested": {"text": "a\x00b"}}, + {"items": None, "mapping": None, "nested": None}, + ] + + +@pytest.mark.parametrize("source", ["tuples", "pandas"]) +def test_large_string_input_ipc(input_service, source): + from pyspark.sql.types import StructType, StructField, StringType + import pandas as pd + spark, parent = input_service + names = ["a", "b", "c"] + rows = [(f"雪😀-{i}", str(i), None if i % 7 == 0 else "a\x00b") + for i in range(10000)] + data = pd.DataFrame(rows, columns=names) if source == "pandas" else rows + schema = StructType([StructField(name, StringType()) for name in names]) + spark.createDataFrame(data, schema).schema + parent.send("get") + assert parent.poll(10) + actual, _ = parent.recv() + assert actual == [dict(zip(names, row)) for row in rows] From 7ea45dd317d2c70e2461af988ee52dd5f68709c9 Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Thu, 29 Oct 2026 04:12:22 -0700 Subject: [PATCH 2/6] [SPARK-60684][CONNECT] Include string guards in CI-selected offline tests --- python/tests/test_dropin_offline.py | 221 ++++++++++++++++++++++ python/tests/test_exact_string_input.py | 235 ------------------------ 2 files changed, 221 insertions(+), 235 deletions(-) delete mode 100644 python/tests/test_exact_string_input.py diff --git a/python/tests/test_dropin_offline.py b/python/tests/test_dropin_offline.py index 9988f655..efddbd4e 100644 --- a/python/tests/test_dropin_offline.py +++ b/python/tests/test_dropin_offline.py @@ -3085,3 +3085,224 @@ def allow_response_after_request(): else: assert result == [("id", "bigint")] assert service.calls == 1 + + +# Exact built-in string conversion and full outbound Arrow IPC guards. +from concurrent.futures import ThreadPoolExecutor as _exact_string_ThreadPoolExecutor +import multiprocessing as _exact_string_multiprocessing + +from pyspark.sql import Row as _exact_string_Row, SparkSession as _exact_string_SparkSession + + +@pytest.mark.parametrize("text", ["", "1.25", "-42", "NaN", "a\x00b", "é雪😀", "e\u0301", "\uffff"]) +def test_exact_string_roundtrip(text): + actual = _exact_string_Row(value=text).value + assert type(actual) is str + assert actual == text + + +@pytest.mark.parametrize("text", ["\ud800", "\udfff", "a\ud800b"]) +def test_exact_string_surrogate_error_preserves_generic_fallback(text): + with pytest.raises(TypeError, match="unsupported value type for createDataFrame: str"): + _exact_string_Row(value=text) + + +def test_exact_string_string_subclass_protocol_order(): + trace = [] + + class IndexString(str): + def __index__(self): + trace.append("index") + return 17 + def __float__(self): + raise AssertionError("index must take precedence") + + class FloatString(str): + def __float__(self): + trace.append("float") + return 2.5 + def __str__(self): + raise AssertionError("numeric conversion must take precedence") + + class FallbackString(str): + def __index__(self): + trace.append("bad index") + raise ValueError("index rejected") + def __float__(self): + trace.append("bad float") + raise ValueError("float rejected") + def __str__(self): + raise AssertionError("string extraction must retain the underlying text") + + class PlainString(str): + def __str__(self): + raise AssertionError("string extraction must not call __str__") + + for value, expected in [(IndexString("x"), 17), (FloatString("x"), 2.5), + (FallbackString("original"), "original"), + (PlainString("plain"), "plain")]: + actual = _exact_string_Row(value=value).value + assert type(actual) is type(expected) + assert actual == expected + assert trace == ["index", "float", "bad index", "bad float"] + + +def test_exact_string_nested_strings_and_nulls(): + value = _exact_string_Row(items=["雪", None, ""], mapping={"x": "😀", "null": None}) + actual = _exact_string_Row(value=value).value + assert list(actual.items) == ["雪", None, ""] + assert actual.mapping == {"x": "😀", "null": None} + + +def _exact_string_varint(value): + out = bytearray() + while value > 127: + out.append((value & 127) | 128) + value >>= 7 + out.append(value) + return bytes(out) + + +def _exact_string_message(field, value): + return _exact_string_varint(field * 8 + 2) + _exact_string_varint(len(value)) + value + + +def _exact_string_fields(data): + def read(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 = read(pos) + wire = tag & 7 + if wire == 2: + size, pos = read(pos) + result[tag >> 3] = data[pos:pos + size] + pos += size + elif wire == 0: + _, pos = read(pos) + else: + raise AssertionError(f"unexpected wire type {wire}") + return result + + +def _exact_string_serve(connection): + import grpc + import pyarrow as pa + + state = {} + + def analyze(raw, context): + # AnalyzePlanRequest.schema.plan.root.local_relation.data + relation = _exact_string_fields(_exact_string_fields(_exact_string_fields(_exact_string_fields(raw)[4])[1])[1]) + payload = _exact_string_fields(relation[11])[1] + table = pa.ipc.open_stream(payload).read_all() + state["result"] = (table.to_pylist(), str(table.schema)) + # Return an empty StructType; the test inspects the input IPC itself. + return _exact_string_message(2, _exact_string_message(1, _exact_string_message(21, b""))) + + with _exact_string_ThreadPoolExecutor(max_workers=2) as executor: + server = grpc.server(executor) + server.add_generic_rpc_handlers((grpc.method_handlers_generic_handler( + "spark.connect.SparkConnectService", { + "AnalyzePlan": grpc.unary_unary_rpc_method_handler(analyze)}),)) + 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(state["result"]) + finally: + server.stop(0).wait() + + +@pytest.fixture +def exact_string_input_service(): + pytest.importorskip("grpc") + pytest.importorskip("pyarrow") + context = _exact_string_multiprocessing.get_context("spawn") + parent, child = context.Pipe() + process = context.Process(target=_exact_string_serve, args=(child,)) + process.start() + try: + assert parent.poll(20) + spark = _exact_string_SparkSession.builder.remote(f"sc://127.0.0.1:{parent.recv()}").getOrCreate() + yield spark, parent + 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_mode", ["inferred", "names", "explicit"]) +def test_exact_string_string_input_ipc(exact_string_input_service, source, schema_mode): + spark, parent = exact_string_input_service + from pyspark.sql.types import StructType, StructField, StringType + import pandas as pd + + rows = [("雪😀", "a\x00b"), ("", None), (None, "1.25"), ("e\u0301", "\uffff")] + data = pd.DataFrame(rows, columns=["a", "b"]) if source == "pandas" else rows + schema = (StructType([StructField("a", StringType()), StructField("b", StringType())]) + if schema_mode == "explicit" else ["a", "b"] if schema_mode == "names" else None) + spark.createDataFrame(data, schema).schema + parent.send("get") + assert parent.poll(10) + actual, arrow_schema = parent.recv() + names = ["col0", "col1"] if source == "tuples" and schema_mode == "inferred" else ["a", "b"] + assert actual == [dict(zip(names, row)) for row in rows] + assert all(f"{name}: string" in arrow_schema for name in names) + + +def test_exact_string_nested_string_input_ipc(exact_string_input_service): + from pyspark.sql.types import (StructType, StructField, StringType, ArrayType, MapType) + spark, parent = exact_string_input_service + schema = StructType([ + StructField("items", ArrayType(StringType())), + StructField("mapping", MapType(StringType(), StringType())), + StructField("nested", StructType([StructField("text", StringType())])), + ]) + rows = [(["雪", None, ""], {"x": "😀", "null": None}, _exact_string_Row(text="a\x00b")), + (None, None, None)] + spark.createDataFrame(rows, schema).schema + parent.send("get") + assert parent.poll(10) + actual, _ = parent.recv() + # Arrow map columns decode as ordered key/value pairs. + assert actual == [ + {"items": ["雪", None, ""], "mapping": [("null", None), ("x", "😀")], + "nested": {"text": "a\x00b"}}, + {"items": None, "mapping": None, "nested": None}, + ] + + +@pytest.mark.parametrize("source", ["tuples", "pandas"]) +def test_exact_string_large_string_input_ipc(exact_string_input_service, source): + from pyspark.sql.types import StructType, StructField, StringType + import pandas as pd + spark, parent = exact_string_input_service + names = ["a", "b", "c"] + rows = [(f"雪😀-{i}", str(i), None if i % 7 == 0 else "a\x00b") + for i in range(10000)] + data = pd.DataFrame(rows, columns=names) if source == "pandas" else rows + schema = StructType([StructField(name, StringType()) for name in names]) + spark.createDataFrame(data, schema).schema + parent.send("get") + assert parent.poll(10) + actual, _ = parent.recv() + assert actual == [dict(zip(names, row)) for row in rows] diff --git a/python/tests/test_exact_string_input.py b/python/tests/test_exact_string_input.py deleted file mode 100644 index 30cb7384..00000000 --- a/python/tests/test_exact_string_input.py +++ /dev/null @@ -1,235 +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. - -"""Guard exact-string conversion and decode the actual outbound Arrow IPC.""" -from concurrent.futures import ThreadPoolExecutor -import multiprocessing - -import pytest -from pyspark.sql import Row, SparkSession - - -@pytest.mark.parametrize("text", ["", "1.25", "-42", "NaN", "a\x00b", "é雪😀", "e\u0301", "\uffff"]) -def test_exact_string_roundtrip(text): - actual = Row(value=text).value - assert type(actual) is str - assert actual == text - - -@pytest.mark.parametrize("text", ["\ud800", "\udfff", "a\ud800b"]) -def test_surrogate_error_preserves_generic_fallback(text): - with pytest.raises(TypeError, match="unsupported value type for createDataFrame: str"): - Row(value=text) - - -def test_string_subclass_protocol_order(): - trace = [] - - class IndexString(str): - def __index__(self): - trace.append("index") - return 17 - def __float__(self): - raise AssertionError("index must take precedence") - - class FloatString(str): - def __float__(self): - trace.append("float") - return 2.5 - def __str__(self): - raise AssertionError("numeric conversion must take precedence") - - class FallbackString(str): - def __index__(self): - trace.append("bad index") - raise ValueError("index rejected") - def __float__(self): - trace.append("bad float") - raise ValueError("float rejected") - def __str__(self): - raise AssertionError("string extraction must retain the underlying text") - - class PlainString(str): - def __str__(self): - raise AssertionError("string extraction must not call __str__") - - for value, expected in [(IndexString("x"), 17), (FloatString("x"), 2.5), - (FallbackString("original"), "original"), - (PlainString("plain"), "plain")]: - actual = Row(value=value).value - assert type(actual) is type(expected) - assert actual == expected - assert trace == ["index", "float", "bad index", "bad float"] - - -def test_nested_strings_and_nulls(): - value = Row(items=["雪", None, ""], mapping={"x": "😀", "null": None}) - actual = Row(value=value).value - assert list(actual.items) == ["雪", None, ""] - assert actual.mapping == {"x": "😀", "null": None} - - -def _varint(value): - out = bytearray() - while value > 127: - out.append((value & 127) | 128) - value >>= 7 - out.append(value) - return bytes(out) - - -def _message(field, value): - return _varint(field * 8 + 2) + _varint(len(value)) + value - - -def _fields(data): - def read(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 = read(pos) - wire = tag & 7 - if wire == 2: - size, pos = read(pos) - result[tag >> 3] = data[pos:pos + size] - pos += size - elif wire == 0: - _, pos = read(pos) - else: - raise AssertionError(f"unexpected wire type {wire}") - return result - - -def _serve(connection): - import grpc - import pyarrow as pa - - state = {} - - def analyze(raw, context): - # AnalyzePlanRequest.schema.plan.root.local_relation.data - relation = _fields(_fields(_fields(_fields(raw)[4])[1])[1]) - payload = _fields(relation[11])[1] - table = pa.ipc.open_stream(payload).read_all() - state["result"] = (table.to_pylist(), str(table.schema)) - # Return an empty StructType; the test inspects the input IPC itself. - return _message(2, _message(1, _message(21, b""))) - - with ThreadPoolExecutor(max_workers=2) as executor: - server = grpc.server(executor) - server.add_generic_rpc_handlers((grpc.method_handlers_generic_handler( - "spark.connect.SparkConnectService", { - "AnalyzePlan": grpc.unary_unary_rpc_method_handler(analyze)}),)) - 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(state["result"]) - finally: - server.stop(0).wait() - - -@pytest.fixture -def input_service(): - pytest.importorskip("grpc") - pytest.importorskip("pyarrow") - context = multiprocessing.get_context("spawn") - parent, child = context.Pipe() - process = context.Process(target=_serve, args=(child,)) - process.start() - try: - assert parent.poll(20) - spark = SparkSession.builder.remote(f"sc://127.0.0.1:{parent.recv()}").getOrCreate() - yield spark, parent - 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_mode", ["inferred", "names", "explicit"]) -def test_string_input_ipc(input_service, source, schema_mode): - spark, parent = input_service - from pyspark.sql.types import StructType, StructField, StringType - import pandas as pd - - rows = [("雪😀", "a\x00b"), ("", None), (None, "1.25"), ("e\u0301", "\uffff")] - data = pd.DataFrame(rows, columns=["a", "b"]) if source == "pandas" else rows - schema = (StructType([StructField("a", StringType()), StructField("b", StringType())]) - if schema_mode == "explicit" else ["a", "b"] if schema_mode == "names" else None) - spark.createDataFrame(data, schema).schema - parent.send("get") - assert parent.poll(10) - actual, arrow_schema = parent.recv() - names = ["col0", "col1"] if source == "tuples" and schema_mode == "inferred" else ["a", "b"] - assert actual == [dict(zip(names, row)) for row in rows] - assert all(f"{name}: string" in arrow_schema for name in names) - - -def test_nested_string_input_ipc(input_service): - from pyspark.sql.types import (StructType, StructField, StringType, ArrayType, MapType) - spark, parent = input_service - schema = StructType([ - StructField("items", ArrayType(StringType())), - StructField("mapping", MapType(StringType(), StringType())), - StructField("nested", StructType([StructField("text", StringType())])), - ]) - rows = [(["雪", None, ""], {"x": "😀", "null": None}, Row(text="a\x00b")), - (None, None, None)] - spark.createDataFrame(rows, schema).schema - parent.send("get") - assert parent.poll(10) - actual, _ = parent.recv() - # Arrow map columns decode as ordered key/value pairs. - assert actual == [ - {"items": ["雪", None, ""], "mapping": [("null", None), ("x", "😀")], - "nested": {"text": "a\x00b"}}, - {"items": None, "mapping": None, "nested": None}, - ] - - -@pytest.mark.parametrize("source", ["tuples", "pandas"]) -def test_large_string_input_ipc(input_service, source): - from pyspark.sql.types import StructType, StructField, StringType - import pandas as pd - spark, parent = input_service - names = ["a", "b", "c"] - rows = [(f"雪😀-{i}", str(i), None if i % 7 == 0 else "a\x00b") - for i in range(10000)] - data = pd.DataFrame(rows, columns=names) if source == "pandas" else rows - schema = StructType([StructField(name, StringType()) for name in names]) - spark.createDataFrame(data, schema).schema - parent.send("get") - assert parent.poll(10) - actual, _ = parent.recv() - assert actual == [dict(zip(names, row)) for row in rows] From 45dd9bc6e7afd873f738ba3690617a559c431488 Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Fri, 30 Oct 2026 11:38:55 -0700 Subject: [PATCH 3/6] [SPARK-60684][CONNECT] Reuse shared mock server in string input tests Co-authored-by: Isaac --- python/tests/_mock_connect_server.py | 68 +++++++++++ python/tests/test_dropin_offline.py | 162 +++++---------------------- 2 files changed, 97 insertions(+), 133 deletions(-) diff --git a/python/tests/_mock_connect_server.py b/python/tests/_mock_connect_server.py index 555f829a..d26d6079 100644 --- a/python/tests/_mock_connect_server.py +++ b/python/tests/_mock_connect_server.py @@ -39,6 +39,67 @@ def message_field(field, value): return varint(field * 8 + 2) + varint(len(value)) + value +# Field numbers from crates/spark-connect-proto/proto/spark/connect/: +_ANALYZE_SCHEMA = 4 # base.proto: AnalyzePlanRequest.schema +_SCHEMA_PLAN = 1 # base.proto: AnalyzePlanRequest.Schema.plan +_PLAN_ROOT = 1 # base.proto: Plan.root +_RELATION_LOCAL_RELATION = 11 # relations.proto: Relation.local_relation +_LOCAL_RELATION_DATA = 1 # relations.proto: LocalRelation.data + + +def _message_fields(data): + """Read singular protobuf fields without generated service classes.""" + def read_varint(position): + value = 0 + for shift in range(0, 64, 7): + if position >= len(data): + raise ValueError("truncated protobuf varint") + byte = data[position] + position += 1 + if shift == 63 and byte > 1: + raise ValueError("protobuf varint exceeds 64 bits") + value |= (byte & 127) << shift + if byte < 128: + return value, position + raise ValueError("protobuf varint exceeds 64 bits") + + fields = {} + position = 0 + while position < len(data): + tag, position = read_varint(position) + field, wire_type = tag >> 3, tag & 7 + if not field: + raise ValueError("protobuf field number must be positive") + if wire_type == 0: + value, position = read_varint(position) + else: + if wire_type == 1: + size = 8 + elif wire_type == 2: + size, position = read_varint(position) + elif wire_type == 5: + size = 4 + else: + raise ValueError(f"unsupported protobuf wire type {wire_type}") + end = position + size + if end > len(data): + raise ValueError("truncated protobuf field") + value, position = data[position:end], end + fields[field] = value + return fields + + +def _local_relation_data(request): + for field in ( + _ANALYZE_SCHEMA, _SCHEMA_PLAN, _PLAN_ROOT, + _RELATION_LOCAL_RELATION, _LOCAL_RELATION_DATA, + ): + request = _message_fields(request).get(field) + if request is None: + return None + return request + + def struct_schema(name="id"): # types.proto: DataType.long=7, .struct=21; Struct.fields=1; # StructField: name=1, data_type=2, nullable=3, metadata=4. @@ -58,6 +119,7 @@ def serve(connection): # they wait for a response. class Service: calls = 0 + local_relation = None data_type = struct_schema() failure = None block_response = False @@ -65,6 +127,12 @@ class Service: def analyze(self, request, context): self.calls += 1 + payload = _local_relation_data(request) + if payload is not None: + import pyarrow as pa + + table = pa.ipc.open_stream(payload).read_all() + self.local_relation = (table.to_pylist(), str(table.schema)) if self.block_response and not self.response_gate.wait(5): context.abort( grpc.StatusCode.DEADLINE_EXCEEDED, diff --git a/python/tests/test_dropin_offline.py b/python/tests/test_dropin_offline.py index efddbd4e..079c5eb4 100644 --- a/python/tests/test_dropin_offline.py +++ b/python/tests/test_dropin_offline.py @@ -14,7 +14,7 @@ import pytest from _mock_connect_server import message_field, mock_connect_session, struct_schema -from pyspark.sql import SparkSession, types as T +from pyspark.sql import Row, SparkSession, types as T # --------------------------------------------------------------------------- types @@ -3088,26 +3088,22 @@ def allow_response_after_request(): # Exact built-in string conversion and full outbound Arrow IPC guards. -from concurrent.futures import ThreadPoolExecutor as _exact_string_ThreadPoolExecutor -import multiprocessing as _exact_string_multiprocessing - -from pyspark.sql import Row as _exact_string_Row, SparkSession as _exact_string_SparkSession @pytest.mark.parametrize("text", ["", "1.25", "-42", "NaN", "a\x00b", "é雪😀", "e\u0301", "\uffff"]) def test_exact_string_roundtrip(text): - actual = _exact_string_Row(value=text).value + actual = Row(value=text).value assert type(actual) is str assert actual == text @pytest.mark.parametrize("text", ["\ud800", "\udfff", "a\ud800b"]) def test_exact_string_surrogate_error_preserves_generic_fallback(text): - with pytest.raises(TypeError, match="unsupported value type for createDataFrame: str"): - _exact_string_Row(value=text) + with pytest.raises(TypeError): + Row(value=text) -def test_exact_string_string_subclass_protocol_order(): +def test_exact_string_subclass_protocol_order(): trace = [] class IndexString(str): @@ -3141,148 +3137,49 @@ def __str__(self): for value, expected in [(IndexString("x"), 17), (FloatString("x"), 2.5), (FallbackString("original"), "original"), (PlainString("plain"), "plain")]: - actual = _exact_string_Row(value=value).value + actual = Row(value=value).value assert type(actual) is type(expected) assert actual == expected assert trace == ["index", "float", "bad index", "bad float"] def test_exact_string_nested_strings_and_nulls(): - value = _exact_string_Row(items=["雪", None, ""], mapping={"x": "😀", "null": None}) - actual = _exact_string_Row(value=value).value + value = Row(items=["雪", None, ""], mapping={"x": "😀", "null": None}) + actual = Row(value=value).value assert list(actual.items) == ["雪", None, ""] assert actual.mapping == {"x": "😀", "null": None} -def _exact_string_varint(value): - out = bytearray() - while value > 127: - out.append((value & 127) | 128) - value >>= 7 - out.append(value) - return bytes(out) - - -def _exact_string_message(field, value): - return _exact_string_varint(field * 8 + 2) + _exact_string_varint(len(value)) + value - - -def _exact_string_fields(data): - def read(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 = read(pos) - wire = tag & 7 - if wire == 2: - size, pos = read(pos) - result[tag >> 3] = data[pos:pos + size] - pos += size - elif wire == 0: - _, pos = read(pos) - else: - raise AssertionError(f"unexpected wire type {wire}") - return result - - -def _exact_string_serve(connection): - import grpc - import pyarrow as pa - - state = {} - - def analyze(raw, context): - # AnalyzePlanRequest.schema.plan.root.local_relation.data - relation = _exact_string_fields(_exact_string_fields(_exact_string_fields(_exact_string_fields(raw)[4])[1])[1]) - payload = _exact_string_fields(relation[11])[1] - table = pa.ipc.open_stream(payload).read_all() - state["result"] = (table.to_pylist(), str(table.schema)) - # Return an empty StructType; the test inspects the input IPC itself. - return _exact_string_message(2, _exact_string_message(1, _exact_string_message(21, b""))) - - with _exact_string_ThreadPoolExecutor(max_workers=2) as executor: - server = grpc.server(executor) - server.add_generic_rpc_handlers((grpc.method_handlers_generic_handler( - "spark.connect.SparkConnectService", { - "AnalyzePlan": grpc.unary_unary_rpc_method_handler(analyze)}),)) - 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(state["result"]) - finally: - server.stop(0).wait() - - -@pytest.fixture -def exact_string_input_service(): - pytest.importorskip("grpc") - pytest.importorskip("pyarrow") - context = _exact_string_multiprocessing.get_context("spawn") - parent, child = context.Pipe() - process = context.Process(target=_exact_string_serve, args=(child,)) - process.start() - try: - assert parent.poll(20) - spark = _exact_string_SparkSession.builder.remote(f"sc://127.0.0.1:{parent.recv()}").getOrCreate() - yield spark, parent - 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_mode", ["inferred", "names", "explicit"]) -def test_exact_string_string_input_ipc(exact_string_input_service, source, schema_mode): - spark, parent = exact_string_input_service - from pyspark.sql.types import StructType, StructField, StringType +def test_exact_string_input_ipc(mock_connect, source, schema_mode): import pandas as pd + pytest.importorskip("pyarrow") + spark, service = mock_connect rows = [("雪😀", "a\x00b"), ("", None), (None, "1.25"), ("e\u0301", "\uffff")] data = pd.DataFrame(rows, columns=["a", "b"]) if source == "pandas" else rows - schema = (StructType([StructField("a", StringType()), StructField("b", StringType())]) + schema = (T.StructType([T.StructField("a", T.StringType()), T.StructField("b", T.StringType())]) if schema_mode == "explicit" else ["a", "b"] if schema_mode == "names" else None) spark.createDataFrame(data, schema).schema - parent.send("get") - assert parent.poll(10) - actual, arrow_schema = parent.recv() + actual, arrow_schema = service.local_relation names = ["col0", "col1"] if source == "tuples" and schema_mode == "inferred" else ["a", "b"] assert actual == [dict(zip(names, row)) for row in rows] assert all(f"{name}: string" in arrow_schema for name in names) -def test_exact_string_nested_string_input_ipc(exact_string_input_service): - from pyspark.sql.types import (StructType, StructField, StringType, ArrayType, MapType) - spark, parent = exact_string_input_service - schema = StructType([ - StructField("items", ArrayType(StringType())), - StructField("mapping", MapType(StringType(), StringType())), - StructField("nested", StructType([StructField("text", StringType())])), +def test_exact_string_nested_input_ipc(mock_connect): + pytest.importorskip("pyarrow") + spark, service = mock_connect + schema = T.StructType([ + T.StructField("items", T.ArrayType(T.StringType())), + T.StructField("mapping", T.MapType(T.StringType(), T.StringType())), + T.StructField("nested", T.StructType([T.StructField("text", T.StringType())])), ]) - rows = [(["雪", None, ""], {"x": "😀", "null": None}, _exact_string_Row(text="a\x00b")), + rows = [(["雪", None, ""], {"x": "😀", "null": None}, Row(text="a\x00b")), (None, None, None)] spark.createDataFrame(rows, schema).schema - parent.send("get") - assert parent.poll(10) - actual, _ = parent.recv() + actual, _ = service.local_relation # Arrow map columns decode as ordered key/value pairs. assert actual == [ {"items": ["雪", None, ""], "mapping": [("null", None), ("x", "😀")], @@ -3292,17 +3189,16 @@ def test_exact_string_nested_string_input_ipc(exact_string_input_service): @pytest.mark.parametrize("source", ["tuples", "pandas"]) -def test_exact_string_large_string_input_ipc(exact_string_input_service, source): - from pyspark.sql.types import StructType, StructField, StringType +def test_exact_string_batch_input_ipc(mock_connect, source): import pandas as pd - spark, parent = exact_string_input_service + + pytest.importorskip("pyarrow") + spark, service = mock_connect names = ["a", "b", "c"] rows = [(f"雪😀-{i}", str(i), None if i % 7 == 0 else "a\x00b") - for i in range(10000)] + for i in range(256)] data = pd.DataFrame(rows, columns=names) if source == "pandas" else rows - schema = StructType([StructField(name, StringType()) for name in names]) + schema = T.StructType([T.StructField(name, T.StringType()) for name in names]) spark.createDataFrame(data, schema).schema - parent.send("get") - assert parent.poll(10) - actual, _ = parent.recv() + actual, _ = service.local_relation assert actual == [dict(zip(names, row)) for row in rows] From aad7f70598463c0bcced192a0d516a98b3b13ac6 Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Fri, 30 Oct 2026 12:50:01 -0700 Subject: [PATCH 4/6] [SPARK-60684][CONNECT] Decode captured string IPC in tests Co-authored-by: Isaac --- python/tests/_mock_connect_server.py | 7 +------ python/tests/test_dropin_offline.py | 19 +++++++++---------- 2 files changed, 10 insertions(+), 16 deletions(-) diff --git a/python/tests/_mock_connect_server.py b/python/tests/_mock_connect_server.py index d26d6079..e39239f7 100644 --- a/python/tests/_mock_connect_server.py +++ b/python/tests/_mock_connect_server.py @@ -127,12 +127,7 @@ class Service: def analyze(self, request, context): self.calls += 1 - payload = _local_relation_data(request) - if payload is not None: - import pyarrow as pa - - table = pa.ipc.open_stream(payload).read_all() - self.local_relation = (table.to_pylist(), str(table.schema)) + self.local_relation = _local_relation_data(request) if self.block_response and not self.response_gate.wait(5): context.abort( grpc.StatusCode.DEADLINE_EXCEEDED, diff --git a/python/tests/test_dropin_offline.py b/python/tests/test_dropin_offline.py index 079c5eb4..531b3ce8 100644 --- a/python/tests/test_dropin_offline.py +++ b/python/tests/test_dropin_offline.py @@ -3153,23 +3153,23 @@ def test_exact_string_nested_strings_and_nulls(): @pytest.mark.parametrize("source", ["tuples", "pandas"]) @pytest.mark.parametrize("schema_mode", ["inferred", "names", "explicit"]) def test_exact_string_input_ipc(mock_connect, source, schema_mode): - import pandas as pd - - pytest.importorskip("pyarrow") + pd = pytest.importorskip("pandas") + pa = pytest.importorskip("pyarrow") spark, service = mock_connect rows = [("雪😀", "a\x00b"), ("", None), (None, "1.25"), ("e\u0301", "\uffff")] data = pd.DataFrame(rows, columns=["a", "b"]) if source == "pandas" else rows schema = (T.StructType([T.StructField("a", T.StringType()), T.StructField("b", T.StringType())]) if schema_mode == "explicit" else ["a", "b"] if schema_mode == "names" else None) spark.createDataFrame(data, schema).schema - actual, arrow_schema = service.local_relation + table = pa.ipc.open_stream(service.local_relation).read_all() + actual, arrow_schema = table.to_pylist(), str(table.schema) names = ["col0", "col1"] if source == "tuples" and schema_mode == "inferred" else ["a", "b"] assert actual == [dict(zip(names, row)) for row in rows] assert all(f"{name}: string" in arrow_schema for name in names) def test_exact_string_nested_input_ipc(mock_connect): - pytest.importorskip("pyarrow") + pa = pytest.importorskip("pyarrow") spark, service = mock_connect schema = T.StructType([ T.StructField("items", T.ArrayType(T.StringType())), @@ -3179,7 +3179,7 @@ def test_exact_string_nested_input_ipc(mock_connect): rows = [(["雪", None, ""], {"x": "😀", "null": None}, Row(text="a\x00b")), (None, None, None)] spark.createDataFrame(rows, schema).schema - actual, _ = service.local_relation + actual = pa.ipc.open_stream(service.local_relation).read_all().to_pylist() # Arrow map columns decode as ordered key/value pairs. assert actual == [ {"items": ["雪", None, ""], "mapping": [("null", None), ("x", "😀")], @@ -3190,9 +3190,8 @@ def test_exact_string_nested_input_ipc(mock_connect): @pytest.mark.parametrize("source", ["tuples", "pandas"]) def test_exact_string_batch_input_ipc(mock_connect, source): - import pandas as pd - - pytest.importorskip("pyarrow") + pd = pytest.importorskip("pandas") + pa = pytest.importorskip("pyarrow") spark, service = mock_connect names = ["a", "b", "c"] rows = [(f"雪😀-{i}", str(i), None if i % 7 == 0 else "a\x00b") @@ -3200,5 +3199,5 @@ def test_exact_string_batch_input_ipc(mock_connect, source): data = pd.DataFrame(rows, columns=names) if source == "pandas" else rows schema = T.StructType([T.StructField(name, T.StringType()) for name in names]) spark.createDataFrame(data, schema).schema - actual, _ = service.local_relation + actual = pa.ipc.open_stream(service.local_relation).read_all().to_pylist() assert actual == [dict(zip(names, row)) for row in rows] From 67e422b195c0b8bb4e11cc9089e66b7be1cd31da Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Fri, 30 Oct 2026 13:28:30 -0700 Subject: [PATCH 5/6] [SPARK-60684][CONNECT] Clarify compatibility guards and mock capture Co-authored-by: Isaac --- python/tests/_mock_connect_server.py | 23 ++++---- python/tests/test_dropin_offline.py | 81 +++++++++++++++++++++------- 2 files changed, 76 insertions(+), 28 deletions(-) diff --git a/python/tests/_mock_connect_server.py b/python/tests/_mock_connect_server.py index e39239f7..26a5d5d0 100644 --- a/python/tests/_mock_connect_server.py +++ b/python/tests/_mock_connect_server.py @@ -48,7 +48,7 @@ def message_field(field, value): def _message_fields(data): - """Read singular protobuf fields without generated service classes.""" + """Read length-delimited protobuf fields, skipping scalar fields.""" def read_varint(position): value = 0 for shift in range(0, 64, 7): @@ -85,18 +85,22 @@ def read_varint(position): if end > len(data): raise ValueError("truncated protobuf field") value, position = data[position:end], end - fields[field] = value + if wire_type == 2: + fields[field] = value return fields def _local_relation_data(request): - for field in ( - _ANALYZE_SCHEMA, _SCHEMA_PLAN, _PLAN_ROOT, - _RELATION_LOCAL_RELATION, _LOCAL_RELATION_DATA, - ): - request = _message_fields(request).get(field) - if request is None: - return None + try: + for field in ( + _ANALYZE_SCHEMA, _SCHEMA_PLAN, _PLAN_ROOT, + _RELATION_LOCAL_RELATION, _LOCAL_RELATION_DATA, + ): + request = _message_fields(request).get(field) + if request is None: + return None + except (ValueError, TypeError): + return None return request @@ -119,6 +123,7 @@ def serve(connection): # they wait for a response. class Service: calls = 0 + # Raw IPC from the latest AnalyzePlan request; None for non-local/unparseable requests. local_relation = None data_type = struct_schema() failure = None diff --git a/python/tests/test_dropin_offline.py b/python/tests/test_dropin_offline.py index 531b3ce8..bb418187 100644 --- a/python/tests/test_dropin_offline.py +++ b/python/tests/test_dropin_offline.py @@ -3090,7 +3090,9 @@ def allow_response_after_request(): # Exact built-in string conversion and full outbound Arrow IPC guards. -@pytest.mark.parametrize("text", ["", "1.25", "-42", "NaN", "a\x00b", "é雪😀", "e\u0301", "\uffff"]) +@pytest.mark.parametrize("text", [ + "", "1.25", "-42", "NaN", "a\x00b", "é雪😀", "e\u0301", "\uffff", +]) def test_exact_string_roundtrip(text): actual = Row(value=text).value assert type(actual) is str @@ -3104,12 +3106,15 @@ def test_exact_string_surrogate_error_preserves_generic_fallback(text): def test_exact_string_subclass_protocol_order(): + # Pin existing behavior with PyO3 0.28, not a desired contract for str + # subclasses. Revisit conversion order when upgrading PyO3. trace = [] class IndexString(str): def __index__(self): trace.append("index") return 17 + def __float__(self): raise AssertionError("index must take precedence") @@ -3117,6 +3122,7 @@ class FloatString(str): def __float__(self): trace.append("float") return 2.5 + def __str__(self): raise AssertionError("numeric conversion must take precedence") @@ -3124,9 +3130,11 @@ class FallbackString(str): def __index__(self): trace.append("bad index") raise ValueError("index rejected") + def __float__(self): trace.append("bad float") raise ValueError("float rejected") + def __str__(self): raise AssertionError("string extraction must retain the underlying text") @@ -3134,9 +3142,13 @@ class PlainString(str): def __str__(self): raise AssertionError("string extraction must not call __str__") - for value, expected in [(IndexString("x"), 17), (FloatString("x"), 2.5), - (FallbackString("original"), "original"), - (PlainString("plain"), "plain")]: + cases = [ + (IndexString("x"), 17), + (FloatString("x"), 2.5), + (FallbackString("original"), "original"), + (PlainString("plain"), "plain"), + ] + for value, expected in cases: actual = Row(value=value).value assert type(actual) is type(expected) assert actual == expected @@ -3144,7 +3156,10 @@ def __str__(self): def test_exact_string_nested_strings_and_nulls(): - value = Row(items=["雪", None, ""], mapping={"x": "😀", "null": None}) + value = Row( + items=["雪", None, ""], + mapping={"x": "😀", "null": None}, + ) actual = Row(value=value).value assert list(actual.items) == ["雪", None, ""] assert actual.mapping == {"x": "😀", "null": None} @@ -3156,14 +3171,31 @@ def test_exact_string_input_ipc(mock_connect, source, schema_mode): pd = pytest.importorskip("pandas") pa = pytest.importorskip("pyarrow") spark, service = mock_connect - rows = [("雪😀", "a\x00b"), ("", None), (None, "1.25"), ("e\u0301", "\uffff")] + rows = [ + ("雪😀", "a\x00b"), + ("", None), + (None, "1.25"), + ("e\u0301", "\uffff"), + ] data = pd.DataFrame(rows, columns=["a", "b"]) if source == "pandas" else rows - schema = (T.StructType([T.StructField("a", T.StringType()), T.StructField("b", T.StringType())]) - if schema_mode == "explicit" else ["a", "b"] if schema_mode == "names" else None) + if schema_mode == "explicit": + schema = T.StructType([ + T.StructField("a", T.StringType()), + T.StructField("b", T.StringType()), + ]) + elif schema_mode == "names": + schema = ["a", "b"] + else: + schema = None + spark.createDataFrame(data, schema).schema table = pa.ipc.open_stream(service.local_relation).read_all() - actual, arrow_schema = table.to_pylist(), str(table.schema) - names = ["col0", "col1"] if source == "tuples" and schema_mode == "inferred" else ["a", "b"] + actual = table.to_pylist() + arrow_schema = str(table.schema) + if source == "tuples" and schema_mode == "inferred": + names = ["col0", "col1"] + else: + names = ["a", "b"] assert actual == [dict(zip(names, row)) for row in rows] assert all(f"{name}: string" in arrow_schema for name in names) @@ -3174,16 +3206,24 @@ def test_exact_string_nested_input_ipc(mock_connect): schema = T.StructType([ T.StructField("items", T.ArrayType(T.StringType())), T.StructField("mapping", T.MapType(T.StringType(), T.StringType())), - T.StructField("nested", T.StructType([T.StructField("text", T.StringType())])), + T.StructField("nested", T.StructType([ + T.StructField("text", T.StringType()), + ])), ]) - rows = [(["雪", None, ""], {"x": "😀", "null": None}, Row(text="a\x00b")), - (None, None, None)] + rows = [ + (["雪", None, ""], {"x": "😀", "null": None}, Row(text="a\x00b")), + (None, None, None), + ] spark.createDataFrame(rows, schema).schema - actual = pa.ipc.open_stream(service.local_relation).read_all().to_pylist() + table = pa.ipc.open_stream(service.local_relation).read_all() + actual = table.to_pylist() # Arrow map columns decode as ordered key/value pairs. assert actual == [ - {"items": ["雪", None, ""], "mapping": [("null", None), ("x", "😀")], - "nested": {"text": "a\x00b"}}, + { + "items": ["雪", None, ""], + "mapping": [("null", None), ("x", "😀")], + "nested": {"text": "a\x00b"}, + }, {"items": None, "mapping": None, "nested": None}, ] @@ -3194,10 +3234,13 @@ def test_exact_string_batch_input_ipc(mock_connect, source): pa = pytest.importorskip("pyarrow") spark, service = mock_connect names = ["a", "b", "c"] - rows = [(f"雪😀-{i}", str(i), None if i % 7 == 0 else "a\x00b") - for i in range(256)] + rows = [ + (f"雪😀-{i}", str(i), None if i % 7 == 0 else "a\x00b") + for i in range(256) + ] data = pd.DataFrame(rows, columns=names) if source == "pandas" else rows schema = T.StructType([T.StructField(name, T.StringType()) for name in names]) spark.createDataFrame(data, schema).schema - actual = pa.ipc.open_stream(service.local_relation).read_all().to_pylist() + table = pa.ipc.open_stream(service.local_relation).read_all() + actual = table.to_pylist() assert actual == [dict(zip(names, row)) for row in rows] From 8a89df385c962e5d1deadafe75e5ec52063049b1 Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Sat, 31 Oct 2026 15:38:00 -0700 Subject: [PATCH 6/6] Trim exact-string tests to conversion guards Remove IPC-only mock parsing and tests for unchanged Arrow encoding. Reduce the subclass guard to one __index__ implementation. Co-authored-by: Isaac --- python/tests/_mock_connect_server.py | 68 ------------- python/tests/test_dropin_offline.py | 143 ++------------------------- 2 files changed, 7 insertions(+), 204 deletions(-) diff --git a/python/tests/_mock_connect_server.py b/python/tests/_mock_connect_server.py index 26a5d5d0..555f829a 100644 --- a/python/tests/_mock_connect_server.py +++ b/python/tests/_mock_connect_server.py @@ -39,71 +39,6 @@ def message_field(field, value): return varint(field * 8 + 2) + varint(len(value)) + value -# Field numbers from crates/spark-connect-proto/proto/spark/connect/: -_ANALYZE_SCHEMA = 4 # base.proto: AnalyzePlanRequest.schema -_SCHEMA_PLAN = 1 # base.proto: AnalyzePlanRequest.Schema.plan -_PLAN_ROOT = 1 # base.proto: Plan.root -_RELATION_LOCAL_RELATION = 11 # relations.proto: Relation.local_relation -_LOCAL_RELATION_DATA = 1 # relations.proto: LocalRelation.data - - -def _message_fields(data): - """Read length-delimited protobuf fields, skipping scalar fields.""" - def read_varint(position): - value = 0 - for shift in range(0, 64, 7): - if position >= len(data): - raise ValueError("truncated protobuf varint") - byte = data[position] - position += 1 - if shift == 63 and byte > 1: - raise ValueError("protobuf varint exceeds 64 bits") - value |= (byte & 127) << shift - if byte < 128: - return value, position - raise ValueError("protobuf varint exceeds 64 bits") - - fields = {} - position = 0 - while position < len(data): - tag, position = read_varint(position) - field, wire_type = tag >> 3, tag & 7 - if not field: - raise ValueError("protobuf field number must be positive") - if wire_type == 0: - value, position = read_varint(position) - else: - if wire_type == 1: - size = 8 - elif wire_type == 2: - size, position = read_varint(position) - elif wire_type == 5: - size = 4 - else: - raise ValueError(f"unsupported protobuf wire type {wire_type}") - end = position + size - if end > len(data): - raise ValueError("truncated protobuf field") - value, position = data[position:end], end - if wire_type == 2: - fields[field] = value - return fields - - -def _local_relation_data(request): - try: - for field in ( - _ANALYZE_SCHEMA, _SCHEMA_PLAN, _PLAN_ROOT, - _RELATION_LOCAL_RELATION, _LOCAL_RELATION_DATA, - ): - request = _message_fields(request).get(field) - if request is None: - return None - except (ValueError, TypeError): - return None - return request - - def struct_schema(name="id"): # types.proto: DataType.long=7, .struct=21; Struct.fields=1; # StructField: name=1, data_type=2, nullable=3, metadata=4. @@ -123,8 +58,6 @@ def serve(connection): # they wait for a response. class Service: calls = 0 - # Raw IPC from the latest AnalyzePlan request; None for non-local/unparseable requests. - local_relation = None data_type = struct_schema() failure = None block_response = False @@ -132,7 +65,6 @@ class Service: def analyze(self, request, context): self.calls += 1 - self.local_relation = _local_relation_data(request) if self.block_response and not self.response_gate.wait(5): context.abort( grpc.StatusCode.DEADLINE_EXCEEDED, diff --git a/python/tests/test_dropin_offline.py b/python/tests/test_dropin_offline.py index bb418187..578fa2f8 100644 --- a/python/tests/test_dropin_offline.py +++ b/python/tests/test_dropin_offline.py @@ -3087,7 +3087,7 @@ def allow_response_after_request(): assert service.calls == 1 -# Exact built-in string conversion and full outbound Arrow IPC guards. +# Exact built-in string conversion guards. @pytest.mark.parametrize("text", [ @@ -3105,142 +3105,13 @@ def test_exact_string_surrogate_error_preserves_generic_fallback(text): Row(value=text) -def test_exact_string_subclass_protocol_order(): - # Pin existing behavior with PyO3 0.28, not a desired contract for str - # subclasses. Revisit conversion order when upgrading PyO3. - trace = [] - +def test_exact_string_subclass_uses_generic_conversion(): + # Pin existing subclass behavior; treating subclasses as text is a separate + # behavior change. class IndexString(str): def __index__(self): - trace.append("index") return 17 - def __float__(self): - raise AssertionError("index must take precedence") - - class FloatString(str): - def __float__(self): - trace.append("float") - return 2.5 - - def __str__(self): - raise AssertionError("numeric conversion must take precedence") - - class FallbackString(str): - def __index__(self): - trace.append("bad index") - raise ValueError("index rejected") - - def __float__(self): - trace.append("bad float") - raise ValueError("float rejected") - - def __str__(self): - raise AssertionError("string extraction must retain the underlying text") - - class PlainString(str): - def __str__(self): - raise AssertionError("string extraction must not call __str__") - - cases = [ - (IndexString("x"), 17), - (FloatString("x"), 2.5), - (FallbackString("original"), "original"), - (PlainString("plain"), "plain"), - ] - for value, expected in cases: - actual = Row(value=value).value - assert type(actual) is type(expected) - assert actual == expected - assert trace == ["index", "float", "bad index", "bad float"] - - -def test_exact_string_nested_strings_and_nulls(): - value = Row( - items=["雪", None, ""], - mapping={"x": "😀", "null": None}, - ) - actual = Row(value=value).value - assert list(actual.items) == ["雪", None, ""] - assert actual.mapping == {"x": "😀", "null": None} - - -@pytest.mark.parametrize("source", ["tuples", "pandas"]) -@pytest.mark.parametrize("schema_mode", ["inferred", "names", "explicit"]) -def test_exact_string_input_ipc(mock_connect, source, schema_mode): - pd = pytest.importorskip("pandas") - pa = pytest.importorskip("pyarrow") - spark, service = mock_connect - rows = [ - ("雪😀", "a\x00b"), - ("", None), - (None, "1.25"), - ("e\u0301", "\uffff"), - ] - data = pd.DataFrame(rows, columns=["a", "b"]) if source == "pandas" else rows - if schema_mode == "explicit": - schema = T.StructType([ - T.StructField("a", T.StringType()), - T.StructField("b", T.StringType()), - ]) - elif schema_mode == "names": - schema = ["a", "b"] - else: - schema = None - - spark.createDataFrame(data, schema).schema - table = pa.ipc.open_stream(service.local_relation).read_all() - actual = table.to_pylist() - arrow_schema = str(table.schema) - if source == "tuples" and schema_mode == "inferred": - names = ["col0", "col1"] - else: - names = ["a", "b"] - assert actual == [dict(zip(names, row)) for row in rows] - assert all(f"{name}: string" in arrow_schema for name in names) - - -def test_exact_string_nested_input_ipc(mock_connect): - pa = pytest.importorskip("pyarrow") - spark, service = mock_connect - schema = T.StructType([ - T.StructField("items", T.ArrayType(T.StringType())), - T.StructField("mapping", T.MapType(T.StringType(), T.StringType())), - T.StructField("nested", T.StructType([ - T.StructField("text", T.StringType()), - ])), - ]) - rows = [ - (["雪", None, ""], {"x": "😀", "null": None}, Row(text="a\x00b")), - (None, None, None), - ] - spark.createDataFrame(rows, schema).schema - table = pa.ipc.open_stream(service.local_relation).read_all() - actual = table.to_pylist() - # Arrow map columns decode as ordered key/value pairs. - assert actual == [ - { - "items": ["雪", None, ""], - "mapping": [("null", None), ("x", "😀")], - "nested": {"text": "a\x00b"}, - }, - {"items": None, "mapping": None, "nested": None}, - ] - - -@pytest.mark.parametrize("source", ["tuples", "pandas"]) -def test_exact_string_batch_input_ipc(mock_connect, source): - pd = pytest.importorskip("pandas") - pa = pytest.importorskip("pyarrow") - spark, service = mock_connect - names = ["a", "b", "c"] - rows = [ - (f"雪😀-{i}", str(i), None if i % 7 == 0 else "a\x00b") - for i in range(256) - ] - data = pd.DataFrame(rows, columns=names) if source == "pandas" else rows - schema = T.StructType([T.StructField(name, T.StringType()) for name in names]) - spark.createDataFrame(data, schema).schema - table = pa.ipc.open_stream(service.local_relation).read_all() - actual = table.to_pylist() - assert actual == [dict(zip(names, row)) for row in rows] + actual = Row(value=IndexString("x")).value + assert type(actual) is int + assert actual == 17