viirya commented on code in PR #108: URL: https://github.com/apache/spark-connect-rust/pull/108#discussion_r4409827706
########## crates/pyspark-rs/src/pyarrow_export.rs: ########## @@ -0,0 +1,106 @@ +// 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. + +//! Export decoded Spark results using arrow-rs's PyArrow reader conversion. + +use std::sync::Arc; + +use arrow::datatypes::{DataType, Field, Schema}; +use arrow::ipc::writer::FileWriter; +use arrow::pyarrow::IntoPyArrow; +use arrow::record_batch::{RecordBatch, RecordBatchIterator, RecordBatchReader}; +use pyo3::exceptions::{PyRuntimeError, PyValueError}; +use pyo3::prelude::*; +use pyo3::types::PyBytes; + +/// Share compatible value and offset buffers with PyArrow in one reader import. +/// +/// Batches have already been collected; this does not stream the network response. +/// Spark result maps are unsorted and result fields are not dictionary-encoded. +/// An empty vector retains the existing empty-schema behavior. +pub(crate) fn record_batches_to_pyarrow( + py: Python<'_>, + batches: Vec<RecordBatch>, +) -> PyResult<Py<PyAny>> { + let schema = batches + .first() + .map(|batch| batch.schema()) + .unwrap_or_else(|| Arc::new(Schema::empty())); + // RecordBatchIterator does not validate schemas. Check before either path, + // including metadata, so later batches are not relabeled with the first schema. + if batches.iter().any(|batch| batch.schema_ref() != &schema) { + return Err(PyValueError::new_err( + "Arrow record batches have inconsistent schemas", + )); + } + if schema.fields().iter().any(|field| field_needs_ipc(field)) { + // Arrow 58.3 panics when exporting NUL-containing C strings. Reuse Review Comment: Could this say where the panic happens? With the stream export, the schema is converted inside the `extern "C"` `get_schema` callback, which has no `catch_unwind`, so the panic aborts the Python process and does not surface as an exception. That is why the check has to stay even if someone later considers catching the panic instead. Something like "panics inside the C stream `get_schema` callback, which aborts the process" would help. ########## 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") Review Comment: The coverage workflow installs `protobuf-compiler` today, but if that ever changes, all of these tests would skip and the job would stay green. Could this fail instead of skipping when running in CI (e.g. when `CI` is set)? ########## 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): Review Comment: Apart from the mismatched-schema cases, these handoff tests pass unchanged against master's IPC implementation, so a later change that routed `toArrow` back through IPC would go unnoticed. The docstring on the buffer-lifetime test acknowledges it can't prove Rust-to-Python sharing either. One cheap guard is to monkeypatch `pyarrow.ipc.open_file` to raise in the non-NUL tests, or in a dedicated test. That asserts the property this PR exists for, and the NUL tests then show the fallback is the only path that uses IPC. -- 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]
