viirya commented on code in PR #103: URL: https://github.com/apache/spark-connect-rust/pull/103#discussion_r4397189320
########## python/tests/test_arrow_handoff.py: ########## @@ -0,0 +1,521 @@ +# +# 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. +# +"""Exercise the actual Rust drop-in's Arrow handoff, without a Spark server. + +The localhost service returns deterministic Arrow IPC results. Buffer addresses +are captured after Rust decodes IPC, at its C Data export boundary: addresses +in the server's Python arrays would not prove anything about this handoff. +""" + +import ctypes +import gc +from concurrent.futures import ThreadPoolExecutor +from contextlib import contextmanager +from datetime import datetime, timezone +from decimal import Decimal + +import grpc +import pandas as pd +import pyarrow as pa +import pytest +from pyspark.sql import SparkSession + + +def _varint(value): + encoded = bytearray() + while value > 127: + encoded.append((value & 127) | 128) + value >>= 7 + encoded.append(value) + return bytes(encoded) + + +def _bytes_field(number, value): + return _varint((number << 3) | 2) + _varint(len(value)) + value + + +def _request_fields(data): + """Read only the protobuf envelope, not the plan or other nested messages.""" + position = 0 + + def varint(): + nonlocal position + value = shift = 0 + while True: + byte = data[position] + position += 1 + value |= (byte & 127) << shift + if byte < 128: + return value + shift += 7 + + fields = {} + while position < len(data): + tag = varint() + wire_type = tag & 7 + if wire_type == 2: + length = varint() + fields[tag >> 3] = data[position : position + length] + position += length + elif wire_type == 0: + varint() + elif wire_type in (1, 5): + position += 8 if wire_type == 1 else 4 + else: + raise AssertionError(f"Unexpected protobuf wire type: {wire_type}") + return fields + + +@contextmanager +def _dataframe(batches, *, payloads=None, calls=None): Review Comment: This file is 521 lines. It hand-encodes the protobuf wire format with hard-coded field numbers, mirrors `ArrowArray` in ctypes, hooks the release callbacks, and has the autouse `_forbid_c_stream_import` fixture @tdcmeehan asked about. That's a lot to maintain, and much of it depends on how the export is implemented. Buffer identity and lifetime are already covered by `conversion.rs`. Could the Python side be cut down to the public `toArrow()`/`toPandas()` behavior plus one importer-failure case, with the real-server e2e covering the rest? ########## crates/spark-python-arrow/src/bridge.rs: ########## @@ -0,0 +1,237 @@ +// 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 Arrow record batches to PyArrow through the Arrow C Data Interface. +//! +//! The caller supplies valid Rust record batches and owns the Python interpreter. +//! PyArrow's importers consume the exported descriptors and retain their buffers. + +use std::sync::Arc; + +use arrow::array::{Array, StructArray}; +use arrow::datatypes::{DataType, Field, Schema}; +use arrow::error::ArrowError; +use arrow::ffi::{FFI_ArrowArray, FFI_ArrowSchema}; +use arrow::ipc::writer::FileWriter; +use arrow::record_batch::RecordBatch; +use pyo3::exceptions::{PyRuntimeError, PyValueError}; +use pyo3::prelude::*; +use pyo3::types::{PyBytes, PyList}; + +// Called after table export has routed NUL names and timezones to the IPC fallback. +fn export_schema(py: Python<'_>, schema: &Schema) -> PyResult<Py<PyAny>> { + let data_type = DataType::Struct(schema.fields().clone()); + let mut schema = export_c_data_type(&data_type) + .and_then(|exported| exported.with_metadata(schema.metadata())) + .map_err(arrow_error)?; + let address = &mut schema as *mut FFI_ArrowSchema as usize; + Ok(py + .import("pyarrow")? + .getattr("Schema")? + .call_method1("_import_from_c", (address,))? + .unbind()) +} + +/// Build a PyArrow table, sharing data buffers for C-compatible schemas. +/// +/// Batch boundaries, schema metadata, and zero-column row counts are preserved. +/// An empty batch vector produces an empty table with an empty schema. +/// Value and offset buffers are shared; sliced validity bitmaps may be realigned. +/// Inconsistent batch schemas, including metadata and ordering flags, raise `ValueError`. +pub 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())); + if batches.iter().any(|batch| { + !Arc::ptr_eq(batch.schema_ref(), &schema) && !schemas_equal(batch.schema_ref(), &schema) + }) { + return Err(PyValueError::new_err( + "Arrow record batches have inconsistent schemas", + )); + } + let pa = py.import("pyarrow")?; + if schema.fields().iter().any(|field| field_needs_ipc(field)) { Review Comment: Instead of walking the type tree up front, could this call `export_schema` first and fall back to IPC if it returns an error? The schema is exported before any array, so nothing has been consumed when it fails. That would remove `field_needs_ipc`/`data_type_needs_ipc` (L202–225) and also cover any other schema the C exporter rejects. If you move to `arrow::pyarrow::Table`, the same approach works, but keep a clone of the batch `Vec` for the fallback (cheap, since the batches are `Arc`-backed), because `into_pyarrow` consumes the table. ########## crates/spark-python-arrow/Cargo.toml: ########## @@ -0,0 +1,33 @@ +# 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. +[package] +name = "spark-python-arrow" +version = "0.1.0" +edition = "2021" +rust-version = "1.85" Review Comment: Now that the crate is no longer built outside the workspace, `edition`, `license` and `arrow` could use `.workspace = true` like every other crate. On `rust-version`: `apache-spark-connect` already depends on arrow 58.3, which declares Rust 1.85, so the workspace's 1.82 is already wrong for the core crate too, not just this one (README L76). A separate change that raises the workspace value to 1.85 would let this crate inherit it and make that README paragraph unnecessary. ########## .github/workflows/python-arrow.yml: ########## @@ -0,0 +1,59 @@ +# +# 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. +# + +name: Python Arrow export + +on: + push: + branches: [master, "branch-*"] + pull_request: + branches: [master, "branch-*"] + +permissions: + contents: read + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +jobs: + converter: + name: Python Arrow export + runs-on: ubuntu-latest + timeout-minutes: 20 + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: "3.11" + - name: Install Rust + run: | + rustup toolchain install stable --profile minimal --component clippy Review Comment: This installs latest stable and then runs clippy with `-D warnings` (L59), so the job will start failing whenever a new stable release adds lints, even if this crate hasn't changed. Could it pin a toolchain, or match the workspace lint job's setup? -- 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]
