tdcmeehan commented on code in PR #108: URL: https://github.com/apache/spark-connect-rust/pull/108#discussion_r4413210456
########## python/tests/test_arrow_handoff.py: ########## @@ -0,0 +1,349 @@ +# +# 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. +# +"""Check public Arrow handoff behavior against a deterministic mock RPC service. + +Cases cover Spark-output values, schemas, chunks, and lifetimes through public +APIs. Server-side Python addresses cannot prove Rust-to-Python buffer sharing. +""" + +import gc +import os +import shutil +import subprocess +from concurrent.futures import ThreadPoolExecutor +from contextlib import contextmanager +from datetime import datetime, timezone +from decimal import Decimal +from pathlib import Path + +import grpc +import pandas as pd +import pyarrow as pa +import pytest +from google.protobuf import descriptor_pb2, descriptor_pool, message_factory +from pyspark.sql import SparkSession + + [email protected](scope="module") +def protocol(tmp_path_factory): + """Compile the repository protocol instead of hard-coding protobuf field numbers.""" + protoc = shutil.which(os.environ.get("PROTOC", "protoc")) + if protoc is None: + pytest.skip("Arrow handoff tests require protoc; install it or set PROTOC") + proto_root = Path(__file__).resolve().parents[2] / "crates/spark-connect-proto/proto" + descriptor = tmp_path_factory.mktemp("connect-proto") / "connect.pb" + command = [protoc, f"-I{proto_root}"] + if include := os.environ.get("PROTOC_INCLUDE"): + command.append(f"-I{include}") + subprocess.run( + command + + [ + "--include_imports", + f"--descriptor_set_out={descriptor}", + str(proto_root / "spark/connect/base.proto"), + ], + check=True, + ) + pool = descriptor_pool.DescriptorPool() + for file in descriptor_pb2.FileDescriptorSet.FromString(descriptor.read_bytes()).file: + pool.Add(file) + + def message(name, **fields): + return message_factory.GetMessageClass(pool.FindMessageTypeByName(f"spark.connect.{name}"))( + **fields + ) + + return message + + +@contextmanager +def _dataframe(batches, *, protocol, payloads=None, calls=None): + """Serve Arrow results using the repository's protobuf message definitions.""" + if payloads is None: + payloads = [] + for batch in batches: + sink = pa.BufferOutputStream() + with pa.ipc.new_stream(sink, batch.schema) as writer: + writer.write_batch(batch) + payloads.append((batch.num_rows, sink.getvalue().to_pybytes())) + server_id = "00000000-0000-0000-0000-000000000001" + + def execute(request, context): + if calls is not None: + calls["execute"] = calls.get("execute", 0) + 1 + request = protocol("ExecutePlanRequest").FromString(request) + header = dict( + session_id=request.session_id, + operation_id=request.operation_id, + server_side_session_id=server_id, + ) + for index, (rows, payload) in enumerate(payloads): + yield protocol( + "ExecutePlanResponse", + **header, + response_id=str(index), + arrow_batch=dict(row_count=rows, data=payload), + ).SerializeToString() + yield protocol( + "ExecutePlanResponse", **header, response_id="complete", result_complete={} + ).SerializeToString() + + def release_execute(request, context): + request = protocol("ReleaseExecuteRequest").FromString(request) + return protocol( + "ReleaseExecuteResponse", + session_id=request.session_id, + operation_id=request.operation_id, + server_side_session_id=server_id, + ).SerializeToString() + + def release_session(request, context): + request = protocol("ReleaseSessionRequest").FromString(request) + return protocol( + "ReleaseSessionResponse", + session_id=request.session_id, + server_side_session_id=server_id, + ).SerializeToString() + + with ThreadPoolExecutor(max_workers=2) as executor: + server = grpc.server(executor) + server.add_generic_rpc_handlers( + [ + grpc.method_handlers_generic_handler( + "spark.connect.SparkConnectService", + { + "ExecutePlan": grpc.unary_stream_rpc_method_handler(execute), + "ReleaseExecute": grpc.unary_unary_rpc_method_handler(release_execute), + "ReleaseSession": grpc.unary_unary_rpc_method_handler(release_session), + }, + ) + ] + ) + port = server.add_insecure_port("127.0.0.1:0") + assert port != 0 + server.start() + session = None + try: + session = SparkSession.builder.remote(f"sc://127.0.0.1:{port}").getOrCreate() + yield session.range(1), session + finally: + try: + if session is not None: + session.stop() + finally: + server.stop(0).wait(timeout=5) + + +def _mixed_batch(): + nested_type = pa.struct( + [ + pa.field("items", pa.list_(pa.int64()), metadata={b"nested": b"metadata"}), + pa.field("labels", pa.map_(pa.string(), pa.int64())), + ] + ) + arrays = [ + pa.array([1, None, -3, 4], type=pa.int64()), + pa.array([True, None, False, True]), + pa.array(["alpha", None, "", "omega"]), + pa.array([b"a", None, b"", b"xyz"]), + pa.array( + [Decimal("1.23"), None, Decimal("-4.56"), Decimal("0.00")], type=pa.decimal128(12, 2) + ), + pa.array( + [ + datetime(2020, 1, 1, tzinfo=timezone.utc), + None, + datetime(2020, 1, 2, tzinfo=timezone.utc), + datetime(2020, 1, 3, tzinfo=timezone.utc), + ], + type=pa.timestamp("us", tz="UTC"), + ), + pa.array( + [ + {"items": [1, None], "labels": [("a", 1)]}, + None, + {"items": [], "labels": []}, + {"items": None, "labels": [("b", None)]}, + ], + type=nested_type, + ), + pa.nulls(4), + ] + fields = [ + pa.field(f"c{index}", array.type, metadata={b"field": str(index).encode()}) + for index, array in enumerate(arrays) + ] + schema = pa.schema(fields, metadata={b"schema": b"preserved"}) + return pa.RecordBatch.from_arrays(arrays, schema=schema) + + +def test_handoff_preserves_values_metadata_and_chunks(protocol): + batch = _mixed_batch() + batches = [batch.slice(1, 2), batch.slice(0, 1), batch.slice(3, 1)] + expected = pa.Table.from_batches(batches) + calls = {} + with _dataframe(protocol=protocol, batches=batches, calls=calls) as (frame, _): + actual = frame.toArrow() + assert actual.equals(expected, check_metadata=True) + for column in actual.columns: + assert [len(chunk) for chunk in column.chunks] == [2, 1, 1] + assert calls == {"execute": 1} + + +def test_handoff_buffers_outlive_dataframe_session_and_table(protocol): + batch = _mixed_batch() + expected = batch.to_pydict() + calls = {} + + def collect(): + with _dataframe(protocol=protocol, batches=[batch], calls=calls) as (frame, session): + table = frame.toArrow() + session.stop() + return table + + table = collect() + gc.collect() + assert table.to_pydict() == expected + numeric = table.column(0).chunk(0).slice(1, 3) + strings = table.column(2).chunk(0).slice(1, 3) + addresses = [buffer.address for buffer in (numeric.buffers()[1], strings.buffers()[2])] + del table + gc.collect() + assert numeric.to_pylist() == [None, -3, 4] + assert strings.to_pylist() == [None, "", "omega"] + assert numeric.buffers()[1].address == addresses[0] + assert strings.buffers()[2].address == addresses[1] + assert calls == {"execute": 1} + + [email protected]("kind", ["no-batches", "empty-typed", "zero-columns"]) +def test_handoff_empty_and_zero_column_results(protocol, kind): + if kind == "no-batches": + batches = [] + expected = pa.table({}) + elif kind == "empty-typed": + batches = [_mixed_batch().slice(0, 0)] + expected = pa.Table.from_batches(batches) + else: + structs = pa.array([{}] * 5, type=pa.struct([])) + batch = pa.RecordBatch.from_struct_array(structs) + batch = batch.replace_schema_metadata({b"zero-column": b"preserved"}) + batches = [batch.slice(0, 3), batch.slice(3, 2)] + expected = pa.Table.from_batches(batches) + assert expected.num_rows == 5 + calls = {} + with _dataframe(protocol=protocol, batches=batches, calls=calls) as (frame, _): + actual = frame.toArrow() + assert actual.num_rows == expected.num_rows + assert actual.num_columns == expected.num_columns + assert actual.equals(expected, check_metadata=True) + assert calls == {"execute": 1} + + +def test_to_pandas_uses_the_same_handoff(protocol): + batch = _mixed_batch() + indices = [0, 2, 4, 5] + schema = pa.schema( + [batch.schema.field(index) for index in indices], metadata=batch.schema.metadata + ) + arrays = [batch.column(index) for index in indices] + selected = pa.RecordBatch.from_arrays(arrays, schema=schema) + batches = [selected.slice(0, 2), selected.slice(2, 2)] + expected = pa.Table.from_batches(batches).to_pandas() + calls = {} + with _dataframe(protocol=protocol, batches=batches, calls=calls) as (frame, _): + actual = frame.toPandas() + pd.testing.assert_frame_equal(actual, expected) + assert calls == {"execute": 1} + + [email protected]("kind", ["field", "nested-field", "timezone"]) Review Comment: These cases reach the fallback through a top-level name, a struct child and a top-level timezone, but none needs the `List`/`Map` arm of `data_type_needs_ipc` (`pyarrow_export.rs` L90-95) to find the NUL. If that arm regressed, a list of structs, or a map whose value struct has a NUL field name, would reach `CString::new(name).unwrap()` ([arrow-schema ffi.rs#L174](https://github.com/apache/arrow-rs/blob/58.3.0/arrow-schema/src/ffi.rs#L174)) inside the stream's `extern "C"` `get_schema` callback ([ffi_stream.rs#L127](https://github.com/apache/arrow-rs/blob/58.3.0/arrow-array/src/ffi_stream.rs#L127), [#L220](https://github.com/apache/arrow-rs/blob/58.3.0/arrow-array/src/ffi_stream.rs#L220)), and the process would abort instead of raising. Could `kind` add a list-of-struct case and a map-value-struct case? ########## .github/workflows/coverage.yml: ########## @@ -226,7 +226,9 @@ jobs: - name: Measure Python coverage with offline test suite run: | - cd python && PYTHONPATH=$(pwd) python -m coverage run --rcfile=.coveragerc -m pytest tests/test_dropin_offline.py + cd python + PYTHONPATH=$(pwd) python -m coverage run --rcfile=.coveragerc -m pytest \ + tests/test_dropin_offline.py tests/test_arrow_handoff.py Review Comment: Carrying over from #103: `pyproject.toml` still declares `pyarrow>=10.0.0` (L30), but this step is the only CI run of the handoff tests, and it gets the latest PyArrow from L93 (26.0.0 in [this head's run](https://github.com/apache/spark-connect-rust/actions/runs/40796536658/job/123194254142)). The preflight for older PyArrow is gone, and `toArrow()` now goes through PyArrow's private `RecordBatchReader._import_from_c` ([arrow-pyarrow lib.rs#L460-L466](https://github.com/apache/arrow-rs/blob/58.3.0/arrow-pyarrow/src/lib.rs#L460-L466)) instead of the public `pa.ipc.open_file`, so the local PyArrow 10 run is the only evidence for the floor. Could CI add a leg pinned to `pyarrow==10.0.0`? If that isn't practical, could the docs state which PyArrow versions are tested? ########## python/tests/test_arrow_handoff.py: ########## @@ -0,0 +1,349 @@ +# +# 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. +# +"""Check public Arrow handoff behavior against a deterministic mock RPC service. + +Cases cover Spark-output values, schemas, chunks, and lifetimes through public +APIs. Server-side Python addresses cannot prove Rust-to-Python buffer sharing. +""" + +import gc +import os +import shutil +import subprocess +from concurrent.futures import ThreadPoolExecutor +from contextlib import contextmanager +from datetime import datetime, timezone +from decimal import Decimal +from pathlib import Path + +import grpc +import pandas as pd +import pyarrow as pa +import pytest +from google.protobuf import descriptor_pb2, descriptor_pool, message_factory +from pyspark.sql import SparkSession + + [email protected](scope="module") +def protocol(tmp_path_factory): + """Compile the repository protocol instead of hard-coding protobuf field numbers.""" + protoc = shutil.which(os.environ.get("PROTOC", "protoc")) + if protoc is None: + pytest.skip("Arrow handoff tests require protoc; install it or set PROTOC") + proto_root = Path(__file__).resolve().parents[2] / "crates/spark-connect-proto/proto" + descriptor = tmp_path_factory.mktemp("connect-proto") / "connect.pb" + command = [protoc, f"-I{proto_root}"] + if include := os.environ.get("PROTOC_INCLUDE"): + command.append(f"-I{include}") + subprocess.run( + command + + [ + "--include_imports", + f"--descriptor_set_out={descriptor}", + str(proto_root / "spark/connect/base.proto"), + ], + check=True, + ) + pool = descriptor_pool.DescriptorPool() + for file in descriptor_pb2.FileDescriptorSet.FromString(descriptor.read_bytes()).file: + pool.Add(file) + + def message(name, **fields): + return message_factory.GetMessageClass(pool.FindMessageTypeByName(f"spark.connect.{name}"))( + **fields + ) + + return message + + +@contextmanager +def _dataframe(batches, *, protocol, payloads=None, calls=None): + """Serve Arrow results using the repository's protobuf message definitions.""" + if payloads is None: + payloads = [] + for batch in batches: + sink = pa.BufferOutputStream() + with pa.ipc.new_stream(sink, batch.schema) as writer: + writer.write_batch(batch) + payloads.append((batch.num_rows, sink.getvalue().to_pybytes())) + server_id = "00000000-0000-0000-0000-000000000001" + + def execute(request, context): + if calls is not None: + calls["execute"] = calls.get("execute", 0) + 1 + request = protocol("ExecutePlanRequest").FromString(request) + header = dict( + session_id=request.session_id, + operation_id=request.operation_id, + server_side_session_id=server_id, + ) + for index, (rows, payload) in enumerate(payloads): + yield protocol( + "ExecutePlanResponse", + **header, + response_id=str(index), + arrow_batch=dict(row_count=rows, data=payload), + ).SerializeToString() + yield protocol( + "ExecutePlanResponse", **header, response_id="complete", result_complete={} + ).SerializeToString() + + def release_execute(request, context): + request = protocol("ReleaseExecuteRequest").FromString(request) + return protocol( + "ReleaseExecuteResponse", + session_id=request.session_id, + operation_id=request.operation_id, + server_side_session_id=server_id, + ).SerializeToString() + + def release_session(request, context): + request = protocol("ReleaseSessionRequest").FromString(request) + return protocol( + "ReleaseSessionResponse", + session_id=request.session_id, + server_side_session_id=server_id, + ).SerializeToString() + + with ThreadPoolExecutor(max_workers=2) as executor: + server = grpc.server(executor) + server.add_generic_rpc_handlers( + [ + grpc.method_handlers_generic_handler( + "spark.connect.SparkConnectService", + { + "ExecutePlan": grpc.unary_stream_rpc_method_handler(execute), + "ReleaseExecute": grpc.unary_unary_rpc_method_handler(release_execute), + "ReleaseSession": grpc.unary_unary_rpc_method_handler(release_session), + }, + ) + ] + ) + port = server.add_insecure_port("127.0.0.1:0") + assert port != 0 + server.start() + session = None + try: + session = SparkSession.builder.remote(f"sc://127.0.0.1:{port}").getOrCreate() + yield session.range(1), session + finally: + try: + if session is not None: + session.stop() + finally: + server.stop(0).wait(timeout=5) + + +def _mixed_batch(): + nested_type = pa.struct( + [ + pa.field("items", pa.list_(pa.int64()), metadata={b"nested": b"metadata"}), + pa.field("labels", pa.map_(pa.string(), pa.int64())), + ] + ) + arrays = [ + pa.array([1, None, -3, 4], type=pa.int64()), + pa.array([True, None, False, True]), + pa.array(["alpha", None, "", "omega"]), + pa.array([b"a", None, b"", b"xyz"]), + pa.array( + [Decimal("1.23"), None, Decimal("-4.56"), Decimal("0.00")], type=pa.decimal128(12, 2) + ), + pa.array( + [ + datetime(2020, 1, 1, tzinfo=timezone.utc), + None, + datetime(2020, 1, 2, tzinfo=timezone.utc), + datetime(2020, 1, 3, tzinfo=timezone.utc), + ], + type=pa.timestamp("us", tz="UTC"), + ), + pa.array( + [ + {"items": [1, None], "labels": [("a", 1)]}, + None, + {"items": [], "labels": []}, + {"items": None, "labels": [("b", None)]}, + ], + type=nested_type, + ), + pa.nulls(4), + ] + fields = [ + pa.field(f"c{index}", array.type, metadata={b"field": str(index).encode()}) + for index, array in enumerate(arrays) + ] + schema = pa.schema(fields, metadata={b"schema": b"preserved"}) + return pa.RecordBatch.from_arrays(arrays, schema=schema) + + +def test_handoff_preserves_values_metadata_and_chunks(protocol): + batch = _mixed_batch() + batches = [batch.slice(1, 2), batch.slice(0, 1), batch.slice(3, 1)] + expected = pa.Table.from_batches(batches) + calls = {} + with _dataframe(protocol=protocol, batches=batches, calls=calls) as (frame, _): + actual = frame.toArrow() + assert actual.equals(expected, check_metadata=True) + for column in actual.columns: + assert [len(chunk) for chunk in column.chunks] == [2, 1, 1] + assert calls == {"execute": 1} + + +def test_handoff_buffers_outlive_dataframe_session_and_table(protocol): + batch = _mixed_batch() + expected = batch.to_pydict() + calls = {} + + def collect(): + with _dataframe(protocol=protocol, batches=[batch], calls=calls) as (frame, session): + table = frame.toArrow() + session.stop() + return table + + table = collect() + gc.collect() + assert table.to_pydict() == expected + numeric = table.column(0).chunk(0).slice(1, 3) + strings = table.column(2).chunk(0).slice(1, 3) + addresses = [buffer.address for buffer in (numeric.buffers()[1], strings.buffers()[2])] + del table + gc.collect() + assert numeric.to_pylist() == [None, -3, 4] + assert strings.to_pylist() == [None, "", "omega"] + assert numeric.buffers()[1].address == addresses[0] + assert strings.buffers()[2].address == addresses[1] + assert calls == {"execute": 1} + + [email protected]("kind", ["no-batches", "empty-typed", "zero-columns"]) +def test_handoff_empty_and_zero_column_results(protocol, kind): + if kind == "no-batches": + batches = [] + expected = pa.table({}) + elif kind == "empty-typed": + batches = [_mixed_batch().slice(0, 0)] + expected = pa.Table.from_batches(batches) + else: + structs = pa.array([{}] * 5, type=pa.struct([])) + batch = pa.RecordBatch.from_struct_array(structs) + batch = batch.replace_schema_metadata({b"zero-column": b"preserved"}) + batches = [batch.slice(0, 3), batch.slice(3, 2)] + expected = pa.Table.from_batches(batches) + assert expected.num_rows == 5 + calls = {} + with _dataframe(protocol=protocol, batches=batches, calls=calls) as (frame, _): + actual = frame.toArrow() + assert actual.num_rows == expected.num_rows + assert actual.num_columns == expected.num_columns + assert actual.equals(expected, check_metadata=True) + assert calls == {"execute": 1} + + +def test_to_pandas_uses_the_same_handoff(protocol): + batch = _mixed_batch() + indices = [0, 2, 4, 5] + schema = pa.schema( + [batch.schema.field(index) for index in indices], metadata=batch.schema.metadata + ) + arrays = [batch.column(index) for index in indices] + selected = pa.RecordBatch.from_arrays(arrays, schema=schema) + batches = [selected.slice(0, 2), selected.slice(2, 2)] + expected = pa.Table.from_batches(batches).to_pandas() + calls = {} + with _dataframe(protocol=protocol, batches=batches, calls=calls) as (frame, _): + actual = frame.toPandas() + pd.testing.assert_frame_equal(actual, expected) + assert calls == {"execute": 1} + + [email protected]("kind", ["field", "nested-field", "timezone"]) +def test_nul_schema_strings_fall_back_without_executing_again(protocol, kind): + if kind == "field": + name = "bad\x00name" + values = pa.array([1, None, -3, 4]) + elif kind == "nested-field": + name = "nested" + data_type = pa.struct([pa.field("bad\x00name", pa.int64())]) + values = pa.array( + [{"bad\x00name": 1}, None, {"bad\x00name": None}, {"bad\x00name": 4}], + type=data_type, + ) + else: + name = "timestamp" + values = pa.array([0, None, 2, 3], type=pa.timestamp("us", tz="UTC\x00suffix")) + schema = pa.schema( + [pa.field(name, values.type, metadata={b"field": b"preserved"})], + metadata={b"schema": b"preserved"}, + ) + batch = pa.RecordBatch.from_arrays([values], schema=schema) + batches = [batch.slice(0, 2), batch.slice(2, 2)] + expected = pa.Table.from_batches(batches) + calls = {} + with _dataframe(protocol=protocol, batches=batches, calls=calls) as (frame, _): + actual = frame.toArrow() + # Compare raw Arrow data, not Python datetimes: the deliberately unusual + # timezone string is preserved but is not a Python timezone identifier. + assert actual.equals(expected, check_metadata=True) + assert [len(chunk) for chunk in actual.column(0).chunks] == [2, 2] + assert calls == {"execute": 1} + + +def test_nul_schema_and_field_metadata_are_preserved(protocol): + field = pa.field("value", pa.int64(), metadata={b"field\x00key": b"field\x00value"}) + schema = pa.schema([field], metadata={b"schema\x00key": b"schema\x00value"}) + batch = pa.RecordBatch.from_arrays([pa.array([1, None, 3])], schema=schema) + calls = {} + with _dataframe(protocol=protocol, batches=[batch], calls=calls) as (frame, _): + actual = frame.toArrow() + assert actual.equals(pa.Table.from_batches([batch]), check_metadata=True) + assert calls == {"execute": 1} + + [email protected]("difference", ["type", "name", "schema-metadata", "field-metadata"]) Review Comment: The description says nullability differences are rejected too, and `Field::eq` does compare `nullable` ([field.rs#L108-L115](https://github.com/apache/arrow-rs/blob/58.3.0/arrow-schema/src/field.rs#L108-L115)), but none of these cases changes it, so a comparison that ignored nullability would still pass all eight. Could `difference` add a case where the second batch's field is `nullable=False`? ########## python/tests/test_arrow_handoff.py: ########## @@ -0,0 +1,349 @@ +# +# 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. +# +"""Check public Arrow handoff behavior against a deterministic mock RPC service. + +Cases cover Spark-output values, schemas, chunks, and lifetimes through public +APIs. Server-side Python addresses cannot prove Rust-to-Python buffer sharing. +""" + +import gc +import os +import shutil +import subprocess +from concurrent.futures import ThreadPoolExecutor +from contextlib import contextmanager +from datetime import datetime, timezone +from decimal import Decimal +from pathlib import Path + +import grpc +import pandas as pd +import pyarrow as pa +import pytest +from google.protobuf import descriptor_pb2, descriptor_pool, message_factory +from pyspark.sql import SparkSession + + [email protected](scope="module") +def protocol(tmp_path_factory): + """Compile the repository protocol instead of hard-coding protobuf field numbers.""" + protoc = shutil.which(os.environ.get("PROTOC", "protoc")) + if protoc is None: + pytest.skip("Arrow handoff tests require protoc; install it or set PROTOC") + proto_root = Path(__file__).resolve().parents[2] / "crates/spark-connect-proto/proto" + descriptor = tmp_path_factory.mktemp("connect-proto") / "connect.pb" + command = [protoc, f"-I{proto_root}"] + if include := os.environ.get("PROTOC_INCLUDE"): + command.append(f"-I{include}") + subprocess.run( + command + + [ + "--include_imports", + f"--descriptor_set_out={descriptor}", + str(proto_root / "spark/connect/base.proto"), + ], + check=True, + ) + pool = descriptor_pool.DescriptorPool() + for file in descriptor_pb2.FileDescriptorSet.FromString(descriptor.read_bytes()).file: + pool.Add(file) + + def message(name, **fields): + return message_factory.GetMessageClass(pool.FindMessageTypeByName(f"spark.connect.{name}"))( + **fields + ) + + return message + + +@contextmanager +def _dataframe(batches, *, protocol, payloads=None, calls=None): + """Serve Arrow results using the repository's protobuf message definitions.""" + if payloads is None: + payloads = [] + for batch in batches: + sink = pa.BufferOutputStream() + with pa.ipc.new_stream(sink, batch.schema) as writer: + writer.write_batch(batch) + payloads.append((batch.num_rows, sink.getvalue().to_pybytes())) + server_id = "00000000-0000-0000-0000-000000000001" + + def execute(request, context): + if calls is not None: + calls["execute"] = calls.get("execute", 0) + 1 + request = protocol("ExecutePlanRequest").FromString(request) + header = dict( + session_id=request.session_id, + operation_id=request.operation_id, + server_side_session_id=server_id, + ) + for index, (rows, payload) in enumerate(payloads): + yield protocol( + "ExecutePlanResponse", + **header, + response_id=str(index), + arrow_batch=dict(row_count=rows, data=payload), + ).SerializeToString() + yield protocol( + "ExecutePlanResponse", **header, response_id="complete", result_complete={} + ).SerializeToString() + + def release_execute(request, context): + request = protocol("ReleaseExecuteRequest").FromString(request) + return protocol( + "ReleaseExecuteResponse", + session_id=request.session_id, + operation_id=request.operation_id, + server_side_session_id=server_id, + ).SerializeToString() + + def release_session(request, context): + request = protocol("ReleaseSessionRequest").FromString(request) + return protocol( + "ReleaseSessionResponse", + session_id=request.session_id, + server_side_session_id=server_id, + ).SerializeToString() + + with ThreadPoolExecutor(max_workers=2) as executor: + server = grpc.server(executor) + server.add_generic_rpc_handlers( + [ + grpc.method_handlers_generic_handler( + "spark.connect.SparkConnectService", + { + "ExecutePlan": grpc.unary_stream_rpc_method_handler(execute), + "ReleaseExecute": grpc.unary_unary_rpc_method_handler(release_execute), + "ReleaseSession": grpc.unary_unary_rpc_method_handler(release_session), + }, + ) + ] + ) + port = server.add_insecure_port("127.0.0.1:0") + assert port != 0 + server.start() + session = None + try: + session = SparkSession.builder.remote(f"sc://127.0.0.1:{port}").getOrCreate() + yield session.range(1), session + finally: + try: + if session is not None: + session.stop() + finally: + server.stop(0).wait(timeout=5) + + +def _mixed_batch(): + nested_type = pa.struct( + [ + pa.field("items", pa.list_(pa.int64()), metadata={b"nested": b"metadata"}), + pa.field("labels", pa.map_(pa.string(), pa.int64())), + ] + ) + arrays = [ + pa.array([1, None, -3, 4], type=pa.int64()), + pa.array([True, None, False, True]), + pa.array(["alpha", None, "", "omega"]), + pa.array([b"a", None, b"", b"xyz"]), + pa.array( + [Decimal("1.23"), None, Decimal("-4.56"), Decimal("0.00")], type=pa.decimal128(12, 2) + ), + pa.array( + [ + datetime(2020, 1, 1, tzinfo=timezone.utc), + None, + datetime(2020, 1, 2, tzinfo=timezone.utc), + datetime(2020, 1, 3, tzinfo=timezone.utc), + ], + type=pa.timestamp("us", tz="UTC"), + ), + pa.array( + [ + {"items": [1, None], "labels": [("a", 1)]}, + None, + {"items": [], "labels": []}, + {"items": None, "labels": [("b", None)]}, + ], + type=nested_type, + ), + pa.nulls(4), + ] + fields = [ + pa.field(f"c{index}", array.type, metadata={b"field": str(index).encode()}) + for index, array in enumerate(arrays) + ] + schema = pa.schema(fields, metadata={b"schema": b"preserved"}) + return pa.RecordBatch.from_arrays(arrays, schema=schema) + + +def test_handoff_preserves_values_metadata_and_chunks(protocol): + batch = _mixed_batch() + batches = [batch.slice(1, 2), batch.slice(0, 1), batch.slice(3, 1)] + expected = pa.Table.from_batches(batches) + calls = {} + with _dataframe(protocol=protocol, batches=batches, calls=calls) as (frame, _): + actual = frame.toArrow() + assert actual.equals(expected, check_metadata=True) + for column in actual.columns: + assert [len(chunk) for chunk in column.chunks] == [2, 1, 1] + assert calls == {"execute": 1} + + +def test_handoff_buffers_outlive_dataframe_session_and_table(protocol): + batch = _mixed_batch() + expected = batch.to_pydict() + calls = {} + + def collect(): + with _dataframe(protocol=protocol, batches=[batch], calls=calls) as (frame, session): + table = frame.toArrow() + session.stop() + return table + + table = collect() + gc.collect() + assert table.to_pydict() == expected + numeric = table.column(0).chunk(0).slice(1, 3) + strings = table.column(2).chunk(0).slice(1, 3) + addresses = [buffer.address for buffer in (numeric.buffers()[1], strings.buffers()[2])] + del table + gc.collect() + assert numeric.to_pylist() == [None, -3, 4] + assert strings.to_pylist() == [None, "", "omega"] + assert numeric.buffers()[1].address == addresses[0] + assert strings.buffers()[2].address == addresses[1] Review Comment: Nit: `numeric` and `strings` keep their buffers alive for the whole test, so these addresses can't change and the two asserts can't fail. The `to_pylist()` checks after `del table; gc.collect()` (L226-227) are what exercise the lifetime. Could these lines and `addresses` (L223) be dropped? -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
