From d7cebc3a215f0d4dfe0643abb3c4bca62cb7946b Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Tue, 27 Oct 2026 16:20:35 -0700 Subject: [PATCH 1/7] Export Arrow tables through the C Stream interface Provide an alternative to the per-batch C Data handoff in PR #103 while retaining PyO3 0.28.3. Add ownership, buffer-sharing, compatibility, and public client handoff tests. Generated-by: OpenAI Codex --- .github/workflows/coverage.yml | 4 +- .github/workflows/python-arrow.yml | 59 ++ Cargo.lock | 10 + Cargo.toml | 3 + crates/pyspark-rs/Cargo.toml | 1 + crates/pyspark-rs/src/dataframe.rs | 18 +- crates/spark-python-arrow/Cargo.toml | 33 + crates/spark-python-arrow/README.md | 113 ++++ crates/spark-python-arrow/src/bridge.rs | 416 ++++++++++++ crates/spark-python-arrow/src/lib.rs | 23 + .../tests/bridge_regressions.rs | 218 +++++++ crates/spark-python-arrow/tests/conversion.rs | 281 ++++++++ dev/design/architecture.md | 24 +- docs/architecture.md | 25 +- python/tests/test_arrow_handoff.py | 616 ++++++++++++++++++ 15 files changed, 1828 insertions(+), 16 deletions(-) create mode 100644 .github/workflows/python-arrow.yml create mode 100644 crates/spark-python-arrow/Cargo.toml create mode 100644 crates/spark-python-arrow/README.md create mode 100644 crates/spark-python-arrow/src/bridge.rs create mode 100644 crates/spark-python-arrow/src/lib.rs create mode 100644 crates/spark-python-arrow/tests/bridge_regressions.rs create mode 100644 crates/spark-python-arrow/tests/conversion.rs create mode 100644 python/tests/test_arrow_handoff.py diff --git a/.github/workflows/coverage.yml b/.github/workflows/coverage.yml index 7f017ad7..5afc7c0d 100644 --- a/.github/workflows/coverage.yml +++ b/.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 PYTHONPATH=$(pwd) python -m coverage json --rcfile=.coveragerc -o coverage-python.json - name: Parse Python coverage and create badge JSON diff --git a/.github/workflows/python-arrow.yml b/.github/workflows/python-arrow.yml new file mode 100644 index 00000000..c4bde403 --- /dev/null +++ b/.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 + rustup default stable + - name: Install PyArrow + run: python -m pip install "pyarrow==26.0.0" + - name: Test Arrow table export + # Test separately from pyspark-rs: extension-module feature unification + # would disable libpython linking for the embedded Python tests. + run: | + cargo test -p spark-python-arrow --locked --features pyo3/abi3-py39 + - name: Check converter Clippy + run: | + cargo clippy -p spark-python-arrow \ + --locked --all-targets --features pyo3/abi3-py39 -- -D warnings diff --git a/Cargo.lock b/Cargo.lock index 39d506da..92e73626 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -341,6 +341,7 @@ version = "58.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f633dbfdf39c039ada1bf9e34c694816eb71fbb7dc78f613993b7245e078a1ed" dependencies = [ + "bitflags", "serde_core", "serde_json", ] @@ -3395,6 +3396,7 @@ dependencies = [ "prost", "pyo3", "serde_json", + "spark-python-arrow", "tokio", "tonic", ] @@ -3996,6 +3998,14 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "spark-python-arrow" +version = "0.1.0" +dependencies = [ + "arrow", + "pyo3", +] + [[package]] name = "sqlparser" version = "0.60.0" diff --git a/Cargo.toml b/Cargo.toml index 5fe753a8..31fe8090 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -24,6 +24,7 @@ members = [ "crates/spark-connect-macros", "crates/spark-connect-build", "crates/pyspark-rs", + "crates/spark-python-arrow", "examples", ] exclude = [ @@ -34,6 +35,8 @@ exclude = [ # Rust-native `examples` crate both link the pure-Rust core directly; both are # workspace members but excluded from `default-members` so `cargo test`/`cargo build` # stay clean. Build them explicitly (`cargo build -p pyspark-rs` / `-p examples`). +# `spark-python-arrow` is also explicit: its tests embed Python and require PyArrow. +# Test it separately to avoid extension-module feature unification disabling libpython linking. # # The `examples/wasm-udfs` and `examples/wasm-udf-inline` crates are excluded from # the workspace entirely because they require the `wasm32-unknown-unknown` target, diff --git a/crates/pyspark-rs/Cargo.toml b/crates/pyspark-rs/Cargo.toml index 16148d41..05278bad 100644 --- a/crates/pyspark-rs/Cargo.toml +++ b/crates/pyspark-rs/Cargo.toml @@ -22,6 +22,7 @@ non_snake_case = "allow" spark-connect = { package = "apache-spark-connect", path = "../spark-connect" } spark-connect-core = { package = "apache-spark-connect-core", path = "../spark-connect-core" } spark-connect-proto = { package = "apache-spark-connect-proto", path = "../spark-connect-proto" } +spark-python-arrow = { path = "../spark-python-arrow" } pyo3.workspace = true arrow.workspace = true tokio.workspace = true diff --git a/crates/pyspark-rs/src/dataframe.rs b/crates/pyspark-rs/src/dataframe.rs index 7101aadf..25529563 100644 --- a/crates/pyspark-rs/src/dataframe.rs +++ b/crates/pyspark-rs/src/dataframe.rs @@ -1166,19 +1166,15 @@ impl PyDataFrame { /// Collect the DataFrame into a `pyarrow.Table`. /// /// Mirrors `pyspark.sql.DataFrame.toArrow()`: the result is decoded from the - /// server's Arrow IPC stream and handed to `pyarrow` (no row-by-row copy). + /// server's Arrow IPC stream and handed to `pyarrow` through the shared + /// converter, sharing compatible decoded Arrow value buffers. Bitmap alignment + /// and schemas requiring the IPC compatibility fallback can still copy. #[pyo3(name = "toArrow")] fn to_arrow(&self, py: Python<'_>) -> PyResult> { - use pyo3::types::PyBytes; - // `DataFrame::to_arrow` produces Arrow IPC *file* bytes (FileWriter). - let bytes = py.detach(|| self.dataframe.to_arrow()).to_pyerr()?; - let pa = py.import("pyarrow")?; - let py_bytes = PyBytes::new(py, &bytes); - let table = pa - .getattr("ipc")? - .call_method1("open_file", (py_bytes,))? - .call_method0("read_all")?; - Ok(table.unbind()) + let batches = py + .detach(|| self.dataframe.collect_record_batches()) + .to_pyerr()?; + spark_python_arrow::bridge::record_batches_to_pyarrow(py, batches) } /// Collect the DataFrame into a pandas DataFrame. diff --git a/crates/spark-python-arrow/Cargo.toml b/crates/spark-python-arrow/Cargo.toml new file mode 100644 index 00000000..ada92c39 --- /dev/null +++ b/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" +license = "Apache-2.0" +publish = false +description = "Arrow record batch export to PyArrow for Spark clients" +repository = "https://github.com/apache/spark-connect-rust" + +[dependencies] +arrow = { version = "58", default-features = false, features = ["ffi", "ipc"] } +# The consumer chooses the Python ABI floor and embedding/extension-module linking. +pyo3 = { version = "0.28" } + +[lints.rust] +unsafe_op_in_unsafe_fn = "deny" diff --git a/crates/spark-python-arrow/README.md b/crates/spark-python-arrow/README.md new file mode 100644 index 00000000..2caa87f7 --- /dev/null +++ b/crates/spark-python-arrow/README.md @@ -0,0 +1,113 @@ + + +# Arrow table export to PyArrow + +`spark-python-arrow` exports decoded Rust Arrow record batches to PyArrow for +`pyspark-rs`. It has one public entry point, +`bridge::record_batches_to_pyarrow`, used by `DataFrame.toArrow()`. It does not +depend on Spark Connect transport, generated protobufs, or an async runtime. + +Python-to-Rust Arrow imports, standalone array/type exports, and worker-side UDF +conversion are outside this change. Scalar input conversion and row +materialization remain in `pyspark-rs` until they have another concrete consumer. + +## Export contract + +The Arrow C Stream Interface path preserves batch boundaries, schema and field +metadata, and zero-column row counts. It shares compatible Arrow value and offset +buffers with PyArrow. The caller supplies valid Rust record batches and owns its +Python interpreter; the library does not initialize Python or retain session state. + +`toArrow()` first collects all decoded batches into a Rust `Vec`. +The bridge exports a native Arrow reader over those batches, imports it once as +a PyArrow `RecordBatchReader`, and calls `read_all()`. PyArrow pulls each batch +through native callbacks, without a Python import call per batch or a temporary +Python list of batches. The callbacks do not call Python. C Stream is an +in-process ownership interface, not an IPC byte stream; it does not make result +collection network-streaming or bounded-memory. + +PyArrow consumes the exported stream and retains each batch's release callback, +so the table's buffers outlive the reader and Rust's temporary wrappers. If an +importer raises before consuming the stream, Rust releases it and its remaining +batches. PyArrow and Python code in the same interpreter are trusted, including +the public aliases used by the export path. + +Field names or timezone strings containing NUL use an IPC compatibility fallback +over the already collected batches, without another RPC. Length-prefixed schema +metadata can contain NUL without this fallback. The stream's `get_schema` callback +uses a schema descriptor that retains sorted-map and dictionary-order flags when +present; the other callbacks forward to Arrow's native stream implementation. + +Table export raises `ValueError` if batches have inconsistent schemas, including +field metadata or dictionary ordering. This is stricter than the previous IPC +export path, which wrote batches under the first batch's schema. An empty batch +vector still produces an empty table with an empty schema; preserving schemas +from schema-only IPC streams requires a separate change to the decoding path. + +Zero-copy describes compatible Arrow **value and offset buffers**, not schema +descriptors, Python scalar/row objects, pandas materialization, IPC decoding, or +the IPC fallback. The C Stream callbacks export batches using C Data descriptors; +sliced arrays whose value and null offsets differ can require copying the validity +bitmap during that export. + +## Integration + +This is an unpublished workspace implementation crate (`publish = false`), not +a separately released package. `pyspark-rs` uses a path dependency. The current +package name and `0.1.0` version do not identify a crates.io release; independent +publication would require Apache package naming and release-flow integration. +Source-based consumers can pin a repository revision and validate it in their +own dependency graph. The repository's root `LICENSE` covers this crate. + +The library does not enable PyO3's `extension-module` feature or select a minimum Python ABI; +`pyspark-rs` selects these features for the Python extension. + +PyO3 must resolve to one compatible version in each consuming dependency graph. +The manifest uses the same PyO3 0.28 line as the Python extension, with 0.28.3 +pinned in the workspace lockfile. This change does not upgrade PyO3. +This crate declares Rust 1.85, matching Arrow 58.3.0's declared minimum. +The existing workspace default of 1.82 is unchanged; it is not a claim that +these Python/Arrow packages build on 1.82. The minimum compiler has not been +independently tested locally. + +The Python Arrow export CI job tests this crate with the workspace lockfile. +These tests run separately from the extension build and are not repeated in coverage. + +## Tests + +With PyArrow installed in the Python environment used by PyO3: + +```sh +CARGO_NET_OFFLINE=true CARGO_REGISTRIES_CRATES_IO_PROTOCOL=git \ + cargo test -p spark-python-arrow --locked --features pyo3/abi3-py39 +``` + +Run this package separately from the extension's test/build command: Cargo +feature unification with an extension-module consumer can disable libpython +linking required by the embedded test executable. For a virtual environment, +set `PYO3_PYTHON` to its interpreter and expose its site-packages to the embedded +interpreter with `PYTHONPATH` when necessary. + +The standalone tests exercise the table export entry point: buffer identity and +lifetime, schema metadata, slices, chunks, empty shapes, schema mismatch errors, +and NUL fallback. The Python client's `test_arrow_handoff.py` exercises the public +`toArrow()` and `toPandas()` paths, including importer failures. Scalar regressions +live with the client tests. These are not substitutes for Spark's complete +Connect parity suite. diff --git a/crates/spark-python-arrow/src/bridge.rs b/crates/spark-python-arrow/src/bridge.rs new file mode 100644 index 00000000..a791ddb6 --- /dev/null +++ b/crates/spark-python-arrow/src/bridge.rs @@ -0,0 +1,416 @@ +// 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 Stream 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::ffi::{c_char, c_int, CString}; +use std::sync::Arc; + +use arrow::datatypes::{DataType, Field, Schema}; +use arrow::error::ArrowError; +use arrow::ffi::{FFI_ArrowArray, FFI_ArrowSchema}; +use arrow::ffi_stream::FFI_ArrowArrayStream; +use arrow::ipc::writer::FileWriter; +use arrow::record_batch::{RecordBatch, RecordBatchIterator}; +use pyo3::exceptions::{PyRuntimeError, PyValueError}; +use pyo3::prelude::*; +use pyo3::types::PyBytes; + +// Called after table export has routed NUL names and timezones to the IPC fallback. +fn export_c_schema(schema: &Schema) -> Result { + let data_type = DataType::Struct(schema.fields().clone()); + export_c_data_type(&data_type)?.with_metadata(schema.metadata()) +} + +fn validate_pyarrow_schema(py: Python<'_>, schema: &Schema) -> PyResult<()> { + let mut schema = export_c_schema(schema).map_err(arrow_error)?; + let address = &mut schema as *mut FFI_ArrowSchema as usize; + py.import("pyarrow")? + .getattr("Schema")? + .call_method1("_import_from_c", (address,))?; + Ok(()) +} + +/// 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) -> PyResult> { + 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)) { + // C schema strings cannot contain NUL. Use the already collected + // batches for this rare compatibility path without repeating any RPC. + let bytes = py + .detach(|| -> arrow::error::Result> { + let mut bytes = Vec::new(); + let mut writer = FileWriter::try_new(&mut bytes, &schema)?; + for batch in &batches { + writer.write(batch)?; + } + writer.finish()?; + drop(writer); + Ok(bytes) + }) + .map_err(arrow_error)?; + return Ok(pa + .getattr("ipc")? + .call_method1("open_file", (PyBytes::new(py, &bytes),))? + .call_method0("read_all")? + .unbind()); + } + + // Older PyArrow versions can abort on unsupported stream schemas. Validate + // the same schema through the fallible standalone importer before streaming. + validate_pyarrow_schema(py, &schema)?; + let batches = RecordBatchIterator::new(batches.into_iter().map(Ok), schema.clone()); + let mut stream = with_exported_schema(FFI_ArrowArrayStream::new(Box::new(batches)), schema); + let address = &mut stream as *mut FFI_ArrowArrayStream as usize; + // The importer moves the stream and clears its release callback. Rust drops + // any unconsumed stream on failure; imported arrays retain their own buffers. + // The pointer importer is available at our minimum supported PyArrow (10). + let reader = pa + .getattr("RecordBatchReader")? + .call_method1("_import_from_c", (address,))?; + Ok(reader.call_method0("read_all")?.unbind()) +} + +fn arrow_error(error: ArrowError) -> PyErr { + PyRuntimeError::new_err(error.to_string()) +} + +// Arrow 58's stream exporter needs the same sorted-map schema repair as the +// standalone exporter below. Own the native stream and override only its schema +// conversion; array export, iteration, and remaining-batch cleanup stay in Arrow. +struct ExportStream { + inner: FFI_ArrowArrayStream, + schema: Arc, + schema_error: Option, +} + +fn with_exported_schema(inner: FFI_ArrowArrayStream, schema: Arc) -> FFI_ArrowArrayStream { + let state = Box::new(ExportStream { + inner, + schema, + schema_error: None, + }); + FFI_ArrowArrayStream { + get_schema: Some(stream_schema), + get_next: Some(stream_next), + get_last_error: Some(stream_error), + release: Some(stream_release), + private_data: Box::into_raw(state).cast(), + } +} + +// These callbacks run without the GIL and must not call Python. The C Stream +// contract gives each callback exclusive access to a live stream and, where +// applicable, writable output storage. private_data is our Box until release. +unsafe extern "C" fn stream_schema( + stream: *mut FFI_ArrowArrayStream, + out: *mut FFI_ArrowSchema, +) -> c_int { + // SAFETY: the consumer supplies the live stream constructed above. + let state = unsafe { &mut *(*stream).private_data.cast::() }; + state.schema_error = None; + match export_c_schema(&state.schema) { + Ok(schema) => { + // SAFETY: out is writable. Each call transfers a fresh descriptor, + // which the consumer releases independently of the stream. + unsafe { std::ptr::write(out, schema) }; + 0 + } + Err(error) => { + // Error strings may contain NUL; never panic across the C boundary. + state.schema_error = + Some(CString::new(error.to_string().replace('\0', "\\0")).unwrap_or_default()); + 22 // EINVAL, specified by the C Stream interface. + } + } +} + +unsafe extern "C" fn stream_next( + stream: *mut FFI_ArrowArrayStream, + out: *mut FFI_ArrowArray, +) -> c_int { + // SAFETY: the consumer supplies the live stream constructed above. + let state = unsafe { &mut *(*stream).private_data.cast::() }; + state.schema_error = None; + match state.inner.get_next { + // SAFETY: delegate with the inner stream's own private_data, not ours. + Some(next) => unsafe { next(&mut state.inner, out) }, + None => { + state.schema_error = Some(c"Arrow stream has no get_next callback".into()); + 22 + } + } +} + +unsafe extern "C" fn stream_error(stream: *mut FFI_ArrowArrayStream) -> *const c_char { + // SAFETY: the consumer supplies the live stream constructed above. + let state = unsafe { &mut *(*stream).private_data.cast::() }; + if let Some(error) = &state.schema_error { + return error.as_ptr(); + } + match state.inner.get_last_error { + // SAFETY: the inner error string stays owned by the inner stream. + Some(error) => unsafe { error(&mut state.inner) }, + None => std::ptr::null(), + } +} + +unsafe extern "C" fn stream_release(stream: *mut FFI_ArrowArrayStream) { + if stream.is_null() { + return; + } + // SAFETY: the consumer supplies writable storage for this stream. + let stream = unsafe { &mut *stream }; + if stream.release.is_none() { + return; + } + let state = stream.private_data.cast::(); + stream.release = None; + stream.private_data = std::ptr::null_mut(); + stream.get_schema = None; + stream.get_next = None; + stream.get_last_error = None; + // SAFETY: this is the Box installed above, consumed exactly once. Dropping + // the inner stream releases unread batches, not independently exported arrays. + drop(unsafe { Box::from_raw(state) }); +} + +// Arrow 58's Field conversion overwrites a map's type flags with its field flags. +// Rebuild the descriptor tree using public constructors, retaining both. Only +// descriptors are rebuilt: array buffers are not involved in schema conversion. +fn export_c_data_type(data_type: &DataType) -> Result { + fn has_sorted_map(data_type: &DataType) -> bool { + matches!(data_type, DataType::Map(_, true)) + || child_fields(data_type) + .iter() + .any(|field| has_sorted_map(field.data_type())) + || matches!(data_type, DataType::Dictionary(_, values) if has_sorted_map(values)) + } + + fn rebuild( + source: &FFI_ArrowSchema, + data_type: &DataType, + ) -> Result { + let children = child_fields(data_type) + .into_iter() + .enumerate() + .map(|(index, field)| rebuild(source.child(index), field.data_type())) + .collect::, _>>()?; + let dictionary = match (source.dictionary(), data_type) { + (Some(schema), DataType::Dictionary(_, values)) => Some(rebuild(schema, values)?), + _ => None, + }; + let mut schema = FFI_ArrowSchema::try_new(source.format(), children, dictionary)? + .with_metadata(source.metadata()?)?; + if let Some(name) = source.name() { + schema = schema.with_name(name)?; + } + schema.flags = source.flags; + if matches!(data_type, DataType::Map(_, true)) { + // ARROW_FLAG_MAP_KEYS_SORTED in the public Arrow C Data Interface. + schema.flags |= 0b100; + } + Ok(schema) + } + + let schema = FFI_ArrowSchema::try_from(data_type)?; + if has_sorted_map(data_type) { + rebuild(&schema, data_type) + } else { + Ok(schema) + } +} + +fn child_fields(data_type: &DataType) -> Vec<&Field> { + match data_type { + DataType::List(field) + | DataType::LargeList(field) + | DataType::ListView(field) + | DataType::LargeListView(field) + | DataType::FixedSizeList(field, _) + | DataType::Map(field, _) => vec![field], + DataType::Struct(fields) => fields.iter().map(AsRef::as_ref).collect(), + DataType::Union(fields, _) => fields.iter().map(|(_, field)| field.as_ref()).collect(), + DataType::RunEndEncoded(run_ends, values) => vec![run_ends, values], + _ => vec![], + } +} + +fn schemas_equal(left: &Schema, right: &Schema) -> bool { + fn same_order(left: &DataType, right: &DataType) -> bool { + let fields_match = + child_fields(left) + .into_iter() + .zip(child_fields(right)) + .all(|(left, right)| { + left.dict_is_ordered() == right.dict_is_ordered() + && same_order(left.data_type(), right.data_type()) + }); + fields_match + && match (left, right) { + (DataType::Dictionary(_, left), DataType::Dictionary(_, right)) => { + same_order(left, right) + } + _ => true, + } + } + + // Arrow's structural equality intentionally excludes dictionary-order flags. + // PyArrow includes them, so a table must not relabel later batches' ordering. + left == right + && same_order( + &DataType::Struct(left.fields().clone()), + &DataType::Struct(right.fields().clone()), + ) +} + +fn field_needs_ipc(field: &Field) -> bool { + field.name().contains('\0') || data_type_needs_ipc(field.data_type()) +} + +fn data_type_needs_ipc(data_type: &DataType) -> bool { + match data_type { + DataType::Timestamp(_, Some(timezone)) => timezone.contains('\0'), + DataType::List(field) + | DataType::LargeList(field) + | DataType::ListView(field) + | DataType::LargeListView(field) + | DataType::FixedSizeList(field, _) + | DataType::Map(field, _) => field_needs_ipc(field), + DataType::Struct(fields) => fields.iter().any(|field| field_needs_ipc(field)), + DataType::Union(fields, _) => fields.iter().any(|(_, field)| field_needs_ipc(field)), + DataType::RunEndEncoded(run_ends, values) => { + field_needs_ipc(run_ends) || field_needs_ipc(values) + } + DataType::Dictionary(keys, values) => { + data_type_needs_ipc(keys) || data_type_needs_ipc(values) + } + _ => false, + } +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + use std::ffi::CStr; + + use arrow::datatypes::TimeUnit; + + use super::*; + + #[test] + fn invalid_time_type_is_rejected_without_panicking() { + assert!(export_c_data_type(&DataType::Time32(TimeUnit::Nanosecond)).is_err()); + } + + #[test] + fn stream_schema_can_be_requested_and_released_repeatedly() { + let entries = Field::new( + "entries", + DataType::Struct( + vec![ + Field::new("key", DataType::Utf8, false), + Field::new("value", DataType::Int64, true), + ] + .into(), + ), + false, + ); + let schema = Arc::new(Schema::new_with_metadata( + vec![Field::new( + "maps", + DataType::Map(Arc::new(entries), true), + true, + )], + HashMap::from([("key".into(), "value\0metadata".into())]), + )); + let batches = RecordBatchIterator::new(std::iter::empty(), schema.clone()); + let mut stream = with_exported_schema(FFI_ArrowArrayStream::new(Box::new(batches)), schema); + let mut exported = Vec::new(); + for _ in 0..3 { + let mut out = FFI_ArrowSchema::empty(); + // SAFETY: both descriptors are live and out is writable. + assert_eq!(unsafe { stream_schema(&mut stream, &mut out) }, 0); + assert_ne!(out.child(0).flags & 0b100, 0); + assert_eq!(out.metadata().unwrap()["key"], "value\0metadata"); + exported.push(out); + } + drop(stream); + // Schema descriptors have their own owners after the stream is gone. + for schema in exported { + assert_eq!(schema.child(0).name(), Some("maps")); + } + } + + #[test] + fn stream_schema_errors_have_stable_storage_until_next_operation() { + let schema = Arc::new(Schema::new(vec![Field::new( + "invalid", + DataType::Time32(TimeUnit::Nanosecond), + true, + )])); + let batches = RecordBatchIterator::new(std::iter::empty(), schema.clone()); + let mut stream = with_exported_schema(FFI_ArrowArrayStream::new(Box::new(batches)), schema); + let mut out = FFI_ArrowSchema::empty(); + // SAFETY: the stream and output are live, writable descriptors. + assert_eq!(unsafe { stream_schema(&mut stream, &mut out) }, 22); + // SAFETY: the error pointer is owned by the stream until the next operation. + let error = unsafe { stream_error(&mut stream) }; + assert!(!error.is_null()); + assert_eq!(error, unsafe { stream_error(&mut stream) }); + assert!(!unsafe { CStr::from_ptr(error) }.to_bytes().is_empty()); + let mut array = FFI_ArrowArray::empty(); + assert_eq!(unsafe { stream_next(&mut stream, &mut array) }, 0); + assert!(unsafe { stream_error(&mut stream) }.is_null()); + } + + #[test] + fn stream_forwards_inner_iteration_errors() { + let schema = Arc::new(Schema::empty()); + let batches = RecordBatchIterator::new( + std::iter::once(Err(ArrowError::InvalidArgumentError( + "test read failure".into(), + ))), + schema.clone(), + ); + let mut stream = with_exported_schema(FFI_ArrowArrayStream::new(Box::new(batches)), schema); + let mut array = FFI_ArrowArray::empty(); + // SAFETY: the live stream owns the error string, and array is writable. + assert_eq!(unsafe { stream_next(&mut stream, &mut array) }, 22); + let message = unsafe { CStr::from_ptr(stream_error(&mut stream)) }; + assert!(message.to_str().unwrap().contains("test read failure")); + } +} diff --git a/crates/spark-python-arrow/src/lib.rs b/crates/spark-python-arrow/src/lib.rs new file mode 100644 index 00000000..60c5d084 --- /dev/null +++ b/crates/spark-python-arrow/src/lib.rs @@ -0,0 +1,23 @@ +// 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. + +//! Arrow record batch export to PyArrow without a dependency on Spark transport. +//! +//! The bridge shares compatible Arrow value buffers. Scalar conversion and row materialization +//! remain in the client binding. The caller owns its Python interpreter. + +pub mod bridge; diff --git a/crates/spark-python-arrow/tests/bridge_regressions.rs b/crates/spark-python-arrow/tests/bridge_regressions.rs new file mode 100644 index 00000000..4939b81d --- /dev/null +++ b/crates/spark-python-arrow/tests/bridge_regressions.rs @@ -0,0 +1,218 @@ +// 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. + +use std::ffi::CStr; +use std::sync::Arc; + +use arrow::array::{ + Array, ArrayRef, DictionaryArray, Int64Array, Int8Array, ListArray, MapArray, StringArray, + StructArray, +}; +use arrow::buffer::{NullBuffer, OffsetBuffer}; +use arrow::datatypes::{DataType, Field, Int8Type, Schema}; +use arrow::record_batch::RecordBatch; +use pyo3::exceptions::PyValueError; +use pyo3::prelude::*; +use pyo3::types::PyDict; +use spark_python_arrow::bridge::record_batches_to_pyarrow; + +fn assert_export_equal(py: Python<'_>, batch: RecordBatch, expected_code: &CStr) -> PyResult<()> { + let locals = PyDict::new(py); + locals.set_item("pa", py.import("pyarrow")?)?; + py.run(expected_code, Some(&locals), Some(&locals))?; + let expected = locals.get_item("expected")?.expect("expected table"); + let actual = record_batches_to_pyarrow(py, vec![batch])?; + assert!(actual + .bind(py) + .call_method1("equals", (&expected,))? + .is_truthy()?); + assert!(actual + .bind(py) + .getattr("schema")? + .call_method1("equals", (expected.getattr("schema")?, true))? + .is_truthy()?); + Ok(()) +} + +#[test] +fn table_export_preserves_sorted_map_flags_in_nested_fields() { + Python::initialize(); + Python::attach(|py| -> PyResult<()> { + let entries = StructArray::new( + vec![ + Field::new("key", DataType::Utf8, false), + Field::new("value", DataType::Int64, true), + ] + .into(), + vec![ + Arc::new(StringArray::from(vec!["a", "b"])), + Arc::new(Int64Array::from(vec![1, 2])), + ], + None, + ); + let maps: ArrayRef = Arc::new(MapArray::new( + Arc::new(Field::new("entries", entries.data_type().clone(), false)), + OffsetBuffer::new(vec![0, 2, 2, 2].into()), + entries, + Some(NullBuffer::from(vec![true, false, true])), + true, + )); + let lists: ArrayRef = Arc::new(ListArray::new( + Arc::new(Field::new("item", maps.data_type().clone(), true)), + OffsetBuffer::new(vec![0, 1, 2, 3].into()), + maps.clone(), + None, + )); + let nested: ArrayRef = Arc::new(StructArray::new( + vec![Field::new("maps", lists.data_type().clone(), true)].into(), + vec![lists], + None, + )); + let schema = Arc::new(Schema::new(vec![ + Field::new("maps", maps.data_type().clone(), true), + Field::new("nested", nested.data_type().clone(), true), + ])); + let batch = RecordBatch::try_new(schema, vec![maps, nested]).unwrap(); + assert_export_equal( + py, + batch, + c"map_type = pa.map_(pa.string(), pa.int64(), keys_sorted=True) +maps = pa.array([[('a', 1), ('b', 2)], None, []], type=map_type) +nested_type = pa.struct([pa.field('maps', pa.list_(map_type))]) +nested = pa.array([ + {'maps': [[('a', 1), ('b', 2)]]}, {'maps': [None]}, {'maps': [[]]} +], type=nested_type) +expected = pa.Table.from_arrays([maps, nested], names=['maps', 'nested'])", + ) + }) + .unwrap(); +} + +#[test] +fn table_export_preserves_ordered_dictionaries_in_nested_fields() { + Python::initialize(); + Python::attach(|py| -> PyResult<()> { + let dictionary: ArrayRef = Arc::new(DictionaryArray::::new( + Int8Array::from(vec![Some(1), None, Some(0)]), + Arc::new(StringArray::from(vec!["a", "b"])), + )); + let field = Field::new("dictionary", dictionary.data_type().clone(), true) + .with_dict_is_ordered(true); + let nested: ArrayRef = Arc::new(StructArray::new( + vec![field.clone()].into(), + vec![dictionary.clone()], + None, + )); + let lists: ArrayRef = Arc::new(ListArray::new( + Arc::new(field.clone().with_name("item")), + OffsetBuffer::new(vec![0, 1, 2, 3].into()), + dictionary.clone(), + None, + )); + let schema = Arc::new(Schema::new(vec![ + field, + Field::new("nested", nested.data_type().clone(), true), + Field::new("lists", lists.data_type().clone(), true), + ])); + let batch = RecordBatch::try_new(schema, vec![dictionary, nested, lists]).unwrap(); + assert_export_equal( + py, + batch, + c"dictionary = pa.DictionaryArray.from_arrays( + pa.array([1, None, 0], type=pa.int8()), pa.array(['a', 'b']), ordered=True +) +nested = pa.StructArray.from_arrays([dictionary], names=['dictionary']) +lists = pa.ListArray.from_arrays(pa.array([0, 1, 2, 3], type=pa.int32()), dictionary) +expected = pa.Table.from_arrays([dictionary, nested, lists], + names=['dictionary', 'nested', 'lists'])", + ) + }) + .unwrap(); +} + +#[test] +fn table_export_rejects_inconsistent_nested_dictionary_order() { + Python::initialize(); + Python::attach(|py| { + let fields = |ordered| { + let dictionary = Field::new( + "dictionary", + DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)), + true, + ) + .with_dict_is_ordered(ordered); + [ + dictionary.clone(), + Field::new( + "nested", + DataType::Struct(vec![dictionary.clone()].into()), + true, + ), + Field::new("lists", DataType::List(Arc::new(dictionary)), true), + ] + }; + for (ordered, unordered) in fields(true).into_iter().zip(fields(false)) { + let batches = vec![ + RecordBatch::new_empty(Arc::new(Schema::new(vec![ordered]))), + RecordBatch::new_empty(Arc::new(Schema::new(vec![unordered]))), + ]; + let error = record_batches_to_pyarrow(py, batches).unwrap_err(); + assert!(error.is_instance_of::(py)); + assert!(error.to_string().contains("inconsistent schemas")); + } + }); +} + +#[test] +fn sliced_validity_shares_values_without_requiring_bitmap_identity() { + Python::initialize(); + Python::attach(|py| -> PyResult<()> { + let original = Int64Array::from(vec![Some(1), Some(2), None, Some(4), None]); + let slice = original.slice(1, 3); + assert_eq!(slice.nulls().unwrap().offset(), 1); + assert_eq!(slice.to_data().offset(), 0); + let address = slice.values().as_ptr() as usize; + let schema = Arc::new(Schema::new(vec![Field::new("x", DataType::Int64, true)])); + let batch = RecordBatch::try_new(schema, vec![Arc::new(slice)]).unwrap(); + let table = record_batches_to_pyarrow(py, vec![batch])?; + drop(original); + py.import("gc")?.call_method0("collect")?; + let array = table + .bind(py) + .call_method1("column", (0,))? + .call_method1("chunk", (0,))?; + assert_eq!( + array + .call_method0("buffers")? + .get_item(1)? + .getattr("address")? + .extract::()?, + address + ); + // Arrow may realign validity to the value buffer's offset. The bitmap's + // identity is not part of the value-buffer sharing promise. + assert_eq!(array.getattr("offset")?.extract::()?, 0); + assert_eq!( + array + .call_method0("to_pylist")? + .extract::>>()?, + vec![Some(2), None, Some(4)] + ); + Ok(()) + }) + .unwrap(); +} diff --git a/crates/spark-python-arrow/tests/conversion.rs b/crates/spark-python-arrow/tests/conversion.rs new file mode 100644 index 00000000..b5773e41 --- /dev/null +++ b/crates/spark-python-arrow/tests/conversion.rs @@ -0,0 +1,281 @@ +// 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. + +//! Table export checks against the installed PyArrow. No Spark server is needed. + +use std::collections::HashMap; +use std::ffi::CStr; +use std::sync::Arc; + +use arrow::array::{Array, Int64Array, StringArray, StructArray}; +use arrow::datatypes::{DataType, Field, Schema, TimeUnit}; +use arrow::record_batch::{RecordBatch, RecordBatchOptions}; +use pyo3::exceptions::PyValueError; +use pyo3::prelude::*; +use pyo3::types::PyDict; +use spark_python_arrow::bridge::record_batches_to_pyarrow; + +fn evaluate<'py>(py: Python<'py>, expression: &CStr) -> PyResult> { + let locals = PyDict::new(py); + locals.set_item("pa", py.import("pyarrow")?)?; + py.eval(expression, Some(&locals), Some(&locals)) +} + +fn assert_schema_equal(actual: &Bound<'_, PyAny>, expected: &Bound<'_, PyAny>) -> PyResult<()> { + assert!(actual + .call_method1("equals", (expected, true))? + .is_truthy()?); + Ok(()) +} + +#[test] +fn table_export_shares_value_and_offset_buffers_after_rust_wrappers_drop() { + Python::initialize(); + Python::attach(|py| -> PyResult<()> { + let numbers = Int64Array::from(vec![Some(12), None, Some(34)]); + let strings = StringArray::from(vec![Some("first"), None, Some("last")]); + let addresses = [ + (0, 1, numbers.values().as_ptr() as usize), + (1, 1, strings.value_offsets().as_ptr() as usize), + (1, 2, strings.value_data().as_ptr() as usize), + ]; + let schema = Arc::new(Schema::new(vec![ + Field::new("number", DataType::Int64, true), + Field::new("text", DataType::Utf8, true), + ])); + let batch = + RecordBatch::try_new(schema, vec![Arc::new(numbers), Arc::new(strings)]).unwrap(); + // Export consumes the batch; only PyArrow owns the buffers afterward. + let table = record_batches_to_pyarrow(py, vec![batch])?; + py.import("gc")?.call_method0("collect")?; + for (column, buffer, address) in addresses { + let chunk = table + .bind(py) + .call_method1("column", (column,))? + .call_method1("chunk", (0,))?; + assert_eq!( + chunk + .call_method0("buffers")? + .get_item(buffer)? + .getattr("address")? + .extract::()?, + address + ); + } + let expected = evaluate( + py, + c"pa.table({'number': pa.array([12, None, 34], type=pa.int64()), + 'text': ['first', None, 'last']})", + )?; + assert!(table + .bind(py) + .call_method1("equals", (expected,))? + .is_truthy()?); + Ok(()) + }) + .unwrap(); +} + +#[test] +fn table_export_preserves_chunks_and_shares_each_buffer() { + Python::initialize(); + Python::attach(|py| -> PyResult<()> { + let schema = Arc::new(Schema::new(vec![Field::new("x", DataType::Int64, true)])); + let arrays = [ + Arc::new(Int64Array::from(vec![Some(1), None])), + Arc::new(Int64Array::from(vec![Some(3)])), + ]; + let addresses: Vec<_> = arrays + .iter() + .map(|array| array.values().as_ptr() as usize) + .collect(); + let batches = arrays + .into_iter() + .map(|array| RecordBatch::try_new(schema.clone(), vec![array]).unwrap()) + .collect(); + let table = record_batches_to_pyarrow(py, batches)?; + let column = table.bind(py).call_method1("column", (0,))?; + assert_eq!(column.getattr("num_chunks")?.extract::()?, 2); + for (index, address) in addresses.iter().enumerate() { + let chunk = column.call_method1("chunk", (index,))?; + assert_eq!( + chunk + .call_method0("buffers")? + .get_item(1)? + .getattr("address")? + .extract::()?, + *address + ); + } + assert_eq!( + column + .call_method0("to_pylist")? + .extract::>>()?, + vec![Some(1), None, Some(3)] + ); + Ok(()) + }) + .unwrap(); +} + +#[test] +fn table_export_preserves_nested_values_and_nul_metadata() { + Python::initialize(); + Python::attach(|py| -> PyResult<()> { + let field = Arc::new( + Field::new("value", DataType::Int64, true) + .with_metadata(HashMap::from([("field\0key".into(), "a\0b".into())])), + ); + let nested = StructArray::new( + vec![field].into(), + vec![Arc::new(Int64Array::from(vec![Some(1), None, Some(3)]))], + None, + ); + let schema = Arc::new(Schema::new_with_metadata( + vec![Field::new("nested", nested.data_type().clone(), true)], + HashMap::from([("schema_key".into(), "c\0d".into())]), + )); + let batch = RecordBatch::try_new(schema, vec![Arc::new(nested)]).unwrap(); + let table = record_batches_to_pyarrow(py, vec![batch])?; + let expected_schema = evaluate( + py, + c"pa.schema([pa.field('nested', pa.struct([ + pa.field('value', pa.int64(), metadata={b'field\\x00key': b'a\\x00b'}) + ]))], metadata={b'schema_key': b'c\\x00d'})", + )?; + assert_schema_equal(&table.bind(py).getattr("schema")?, &expected_schema)?; + let actual = table.bind(py).call_method0("to_pylist")?; + let expected = evaluate( + py, + c"[{'nested': {'value': 1}}, {'nested': {'value': None}}, + {'nested': {'value': 3}}]", + )?; + assert!(actual.eq(expected)?); + Ok(()) + }) + .unwrap(); +} + +#[test] +fn table_export_preserves_empty_schemas_and_zero_column_rows() { + Python::initialize(); + Python::attach(|py| -> PyResult<()> { + let empty = record_batches_to_pyarrow(py, vec![])?; + assert_eq!(empty.bind(py).getattr("num_rows")?.extract::()?, 0); + assert_eq!( + empty.bind(py).getattr("num_columns")?.extract::()?, + 0 + ); + assert_schema_equal( + &empty.bind(py).getattr("schema")?, + &evaluate(py, c"pa.schema([])")?, + )?; + + let schema = Arc::new(Schema::new(vec![Field::new("x", DataType::Int64, true)])); + let empty = record_batches_to_pyarrow(py, vec![RecordBatch::new_empty(schema)])?; + assert_eq!(empty.bind(py).getattr("num_rows")?.extract::()?, 0); + assert_schema_equal( + &empty.bind(py).getattr("schema")?, + &evaluate(py, c"pa.schema([('x', pa.int64())])")?, + )?; + + for rows in [0, 4] { + let schema = + Schema::empty().with_metadata(HashMap::from([("empty".into(), "metadata".into())])); + let options = RecordBatchOptions::new().with_row_count(Some(rows)); + let batch = + RecordBatch::try_new_with_options(Arc::new(schema), vec![], &options).unwrap(); + let table = record_batches_to_pyarrow(py, vec![batch])?; + assert_eq!( + table.bind(py).getattr("num_rows")?.extract::()?, + rows + ); + assert_eq!( + table.bind(py).getattr("num_columns")?.extract::()?, + 0 + ); + assert_schema_equal( + &table.bind(py).getattr("schema")?, + &evaluate(py, c"pa.schema([], metadata={b'empty': b'metadata'})")?, + )?; + } + Ok(()) + }) + .unwrap(); +} + +#[test] +fn table_export_rejects_inconsistent_fields_and_metadata() { + Python::initialize(); + Python::attach(|py| { + let field = Field::new("x", DataType::Int64, true); + let schema = Schema::new(vec![field.clone()]); + for other in [ + Schema::new(vec![Field::new("other", DataType::Int64, true)]), + schema + .clone() + .with_metadata(HashMap::from([("schema".into(), "other".into())])), + Schema::new(vec![field + .clone() + .with_metadata(HashMap::from([("field".into(), "other".into())]))]), + ] { + let batches = vec![ + RecordBatch::new_empty(Arc::new(schema.clone())), + RecordBatch::new_empty(Arc::new(other)), + ]; + let error = record_batches_to_pyarrow(py, batches).unwrap_err(); + assert!(error.is_instance_of::(py)); + assert!(error.to_string().contains("inconsistent schemas")); + } + }); +} + +#[test] +fn table_export_falls_back_to_ipc_for_nul_names_and_timezones() { + Python::initialize(); + Python::attach(|py| -> PyResult<()> { + let cases: [(Field, &CStr); 3] = [ + ( + Field::new("bad\0name", DataType::Int64, true), + c"pa.schema([('bad\\x00name', pa.int64())])", + ), + ( + Field::new( + "nested", + DataType::Struct(vec![Field::new("bad\0name", DataType::Int64, true)].into()), + true, + ), + c"pa.schema([('nested', pa.struct([('bad\\x00name', pa.int64())]))])", + ), + ( + Field::new( + "time", + DataType::Timestamp(TimeUnit::Microsecond, Some("UTC\0suffix".into())), + true, + ), + c"pa.schema([('time', pa.timestamp('us', tz='UTC\\x00suffix'))])", + ), + ]; + for (field, expected) in cases { + let batch = RecordBatch::new_empty(Arc::new(Schema::new(vec![field]))); + let table = record_batches_to_pyarrow(py, vec![batch])?; + assert_schema_equal(&table.bind(py).getattr("schema")?, &evaluate(py, expected)?)?; + } + Ok(()) + }) + .unwrap(); +} diff --git a/dev/design/architecture.md b/dev/design/architecture.md index 87303768..09cf031d 100644 --- a/dev/design/architecture.md +++ b/dev/design/architecture.md @@ -47,6 +47,7 @@ crates/ column, plan, dataframe, functions, session, readwriter, catalog, group, window, udf/udtf, streaming, observation pyspark-rs PyO3 bindings -> native module `_pyspark` + spark-python-arrow Arrow C Stream table export to PyArrow; no transport python/ pyspark/… thin drop-in package; public modules re-export from `_pyspark` so `from pyspark.sql import SparkSession` works @@ -67,9 +68,26 @@ Server-streaming RPCs (`ExecutePlan`) are consumed into Arrow record batches. ## Data path -Results arrive as Arrow IPC streams. `arrow-rs` decodes them; hand-off to Python -uses the **Arrow C Data Interface** (zero-copy to `pyarrow`), which is also how -we build `toPandas()` / `createDataFrame(pandas)` without re-serializing. +Results arrive as Arrow IPC streams. `arrow-rs` decodes them; `toArrow()` collects +all batches into a Rust vector and passes it to `spark-python-arrow`. The bridge +exports one native Arrow C Stream reader, imports it as a PyArrow +`RecordBatchReader`, and calls `read_all()`. PyArrow pulls batches through native +callbacks that do not call Python, avoiding a per-batch Python import loop. +This is an in-process buffer handoff, not network streaming or bounded-memory +result collection. + +Compatible Arrow value and offset buffers are shared with PyArrow; sliced +validity bitmaps may need realignment. The bridge validates batch schema +consistency, including metadata and dictionary ordering. Schemas with embedded +NUL field names or timezones use an IPC compatibility fallback, without another +RPC. Schema metadata and ordering flags are preserved. `toPandas()` then calls +PyArrow's pandas conversion, which may copy. + +Scalar input conversion and row materialization stay in `pyspark-rs`. The export +crate does not own schema inference, transport, or worker execution. `collect()` +still uses the native client's owned rows; `createDataFrame(pandas)` still +materializes Python rows. Neither path is claimed to be zero-copy. See the +[table export contract](../../crates/spark-python-arrow/README.md). ## UDF strategy (full compatibility) diff --git a/docs/architecture.md b/docs/architecture.md index f7fdb920..83e33c8c 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -63,6 +63,7 @@ A Cargo workspace of library crates, plus a Python skin: | `apache-spark-connect-core` | `crates/spark-connect-core` | Transport: channel, retries, reattach, artifacts, errors | | `apache-spark-connect` | `crates/spark-connect` | The DataFrame API: session, dataframe, column, functions, plan, group, catalog, window, readwriter, streaming, types | | `pyspark-rs` | `crates/pyspark-rs` | PyO3 bindings - builds the `_pyspark` extension module | +| `spark-python-arrow` (unpublished) | `crates/spark-python-arrow` | Arrow C Stream table export to PyArrow | | Python skin | `python/pyspark` | The drop-in `pyspark` package (+ vendored `pyspark.pandas`, `cloudpickle`, `pyspark.testing`) | ```mermaid @@ -70,6 +71,7 @@ flowchart LR PROTO["spark-connect-proto"] --> CORE["spark-connect-core"] CORE --> SC["spark-connect (DataFrame API)"] SC --> PYO3["pyspark-rs (PyO3)"] + PYARROW["spark-python-arrow (unpublished)"] --> PYO3 PYO3 --> SKIN["python/pyspark skin"] SC --> RUSTAPP["your Rust app"] ``` @@ -102,6 +104,27 @@ triggers `ExecutePlan`. If the response stream breaks mid-flight, the core's reattachable-execute iterator resumes from the last response it saw, so long results survive transient disconnects. +### Python Arrow table handoff + +`DataFrame.toArrow()` collects the decoded Arrow batches in Rust, then exports +one reader through the Arrow C Stream Interface. PyArrow imports that reader +once and calls `read_all()` to build the table. Native callbacks transfer batch +descriptors and ownership without calling Python for each batch. C Stream is +not an IPC byte stream, and this path still holds all result batches in memory +before the handoff; it does not stream results from the network into Python. + +Compatible Arrow value and offset buffers are shared rather than serialized +again. Schema descriptors still need allocation, sliced validity bitmaps may +need realignment, and `toPandas()` may copy during pandas conversion. The server's +IPC responses must still be decoded. Row materialization in `collect()` is +unchanged and is not zero-copy. + +The bridge preserves batch boundaries, schema metadata, and ordering flags. It +rejects batches with inconsistent schemas. Field names or timezones containing +NUL use an IPC compatibility fallback over the collected batches, without +another RPC; length-prefixed schema metadata can contain NUL without fallback. +An empty batch vector still produces an empty table with an empty schema. + ## Two front-ends, one core === "Rust (native)" @@ -152,7 +175,7 @@ See [Rust UDFs via WebAssembly](udfs.md) for the full story. 3. **Correctness is gated in CI** - golden-proto tests guard plan building; the official Apache Spark Connect test suite guards transport and Arrow paths (see the CI notes under `dev/design/`). -4. **Version = Spark version** - the crate/wheel version tracks the Spark release +4. **Version = Spark version** - the published crate/wheel version tracks the Spark release it targets, so the number tells you which Spark it speaks. For the deeper design write-ups (design notes, CI, the official test suite, and diff --git a/python/tests/test_arrow_handoff.py b/python/tests/test_arrow_handoff.py new file mode 100644 index 00000000..ce92faca --- /dev/null +++ b/python/tests/test_arrow_handoff.py @@ -0,0 +1,616 @@ +# +# 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 Stream 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): + """Serve just ExecutePlan and release RPCs; unexpected RPCs fail closed. + + Field numbers are from spark/connect/base.proto. Using the wire envelope + avoids depending on official pyspark's generated protobuf package, which + cannot be imported alongside the Rust-backed drop-in under the same name. + """ + 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 = b"00000000-0000-0000-0000-000000000001" + + def execute(request, context): + if calls is not None: + calls["execute"] = calls.get("execute", 0) + 1 + fields = _request_fields(request) + header = ( + _bytes_field(1, fields[1]) + _bytes_field(12, fields[6]) + _bytes_field(15, server_id) + ) + for index, (rows, payload) in enumerate(payloads): + arrow_batch = _varint(8) + _varint(rows) + _bytes_field(2, payload) + yield (header + _bytes_field(13, str(index).encode()) + _bytes_field(2, arrow_batch)) + yield header + _bytes_field(13, b"complete") + _bytes_field(14, b"") + + def release_execute(request, context): + fields = _request_fields(request) + return _bytes_field(1, fields[1]) + _bytes_field(2, fields[3]) + _bytes_field(3, server_id) + + def release_session(request, context): + return _bytes_field(1, _request_fields(request)[1]) + _bytes_field(2, server_id) + + 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) + + +class _ArrowArray(ctypes.Structure): + pass + + +_ArrowArray._fields_ = [ + ("length", ctypes.c_int64), + ("null_count", ctypes.c_int64), + ("offset", ctypes.c_int64), + ("n_buffers", ctypes.c_int64), + ("n_children", ctypes.c_int64), + ("buffers", ctypes.POINTER(ctypes.c_void_p)), + ("children", ctypes.POINTER(ctypes.POINTER(_ArrowArray))), + ("dictionary", ctypes.POINTER(_ArrowArray)), + ("release", ctypes.c_void_p), + ("private_data", ctypes.c_void_p), +] + + +class _ArrowArrayStream(ctypes.Structure): + _fields_ = [ + (name, ctypes.c_void_p) + for name in ( + "get_schema", + "get_next", + "get_last_error", + "release", + "private_data", + ) + ] + + +_ArrayRelease = ctypes.CFUNCTYPE(None, ctypes.POINTER(_ArrowArray)) +_GetNext = ctypes.CFUNCTYPE( + ctypes.c_int, + ctypes.POINTER(_ArrowArrayStream), + ctypes.POINTER(_ArrowArray), +) +_GetLastError = ctypes.CFUNCTYPE(ctypes.c_void_p, ctypes.POINTER(_ArrowArrayStream)) +_StreamRelease = ctypes.CFUNCTYPE(None, ctypes.POINTER(_ArrowArrayStream)) + +# ctypes owns the executable callback trampolines, but Arrow only retains raw +# function pointers. Keep these test callbacks alive even if a failed assertion +# leaves an imported stream or array in a pytest traceback. No Arrow buffers +# are retained by these callbacks. +_FFI_CALLBACK_KEEPALIVE = [] + + +@pytest.fixture(autouse=True) +def _forbid_per_batch_c_data_import(monkeypatch): + real_batch = pa.RecordBatch + + class ProxyMeta(type): + def __getattr__(cls, name): + return getattr(real_batch, name) + + def __instancecheck__(cls, instance): + return isinstance(instance, real_batch) + + class RecordBatchProxy(metaclass=ProxyMeta): + @staticmethod + def _import_from_c(*args): + raise AssertionError("The C Stream handoff must not import batches in Python") + + _import_from_c_capsule = _import_from_c + + monkeypatch.setattr(pa, "RecordBatch", RecordBatchProxy) + + +def _snapshot(array): + """Copy only addresses/shape while the producer's exported array is live.""" + return { + "length": array.length, + "offset": array.offset, + "buffers": [array.buffers[i] for i in range(array.n_buffers)], + "children": [_snapshot(array.children[i].contents) for i in range(array.n_children)], + "dictionary": _snapshot(array.dictionary.contents) if array.dictionary else None, + } + + +class _StreamSpy: + def __init__( + self, monkeypatch, *, reject=False, consume_before_reject=False, fail_next_at=None + ): + self.batches = [] + self.errors = [] + self.release_counts = [] + self.import_count = 0 + self.next_count = 0 + self.stream_releases = 0 + self.callbacks = [] + real_reader = pa.RecordBatchReader + spy = self + + class ProxyMeta(type): + def __getattr__(cls, name): + return getattr(real_reader, name) + + def __instancecheck__(cls, instance): + return isinstance(instance, real_reader) + + class ReaderProxy(metaclass=ProxyMeta): + @staticmethod + def _import_from_c(address): + stream = ctypes.cast(address, ctypes.POINTER(_ArrowArrayStream)) + original_next = _GetNext(stream.contents.get_next) + original_release = _StreamRelease(stream.contents.release) + + @_GetNext + def get_next(current, output): + spy.next_count += 1 + if spy.next_count == fail_next_at: + return 22 # EINVAL: PyArrow must propagate a failed stream read. + result = original_next(current, output) + if result == 0 and output.contents.release: + try: + spy._observe_array(output) + except Exception as error: + # Exceptions cannot cross a ctypes callback boundary. + spy.errors.append(error) + return result + + @_StreamRelease + def release(current): + spy.stream_releases += 1 + original_release(current) + + spy._keep_alive(original_next, original_release, get_next, release) + stream.contents.get_next = ctypes.cast(get_next, ctypes.c_void_p).value + stream.contents.release = ctypes.cast(release, ctypes.c_void_p).value + if fail_next_at is not None: + original_error = _GetLastError(stream.contents.get_last_error) + message = ctypes.create_string_buffer(b"test stream read failure") + + @_GetLastError + def get_last_error(current): + if spy.next_count == fail_next_at: + return ctypes.addressof(message) + return original_error(current) + + spy._keep_alive(original_error, get_last_error, message) + stream.contents.get_last_error = ctypes.cast( + get_last_error, ctypes.c_void_p + ).value + if reject and not consume_before_reject: + raise RuntimeError("test importer rejected the stream") + imported = real_reader._import_from_c(address) + spy.import_count += 1 + if reject: + # PyArrow owns the stream now. The exception traceback may + # retain this local until gc.collect(), just like user code. + raise RuntimeError("test importer rejected the consumed stream") + return imported + + monkeypatch.setattr(pa, "RecordBatchReader", ReaderProxy) + + def _keep_alive(self, *callbacks): + self.callbacks.extend(callbacks) + _FFI_CALLBACK_KEEPALIVE.extend(callbacks) + + def _observe_array(self, array): + batch_index = len(self.batches) + self.batches.append(_snapshot(array.contents)) + self.release_counts.append(0) + original_release = _ArrayRelease(array.contents.release) + + @_ArrayRelease + def release(current): + self.release_counts[batch_index] += 1 + original_release(current) + + self._keep_alive(original_release, release) + array.contents.release = ctypes.cast(release, ctypes.c_void_p).value + + +def _assert_buffers(array, exported): + assert len(array) == exported["length"] + assert array.offset == exported["offset"] + own_buffers = array.buffers()[: len(exported["buffers"])] + for buffer, address in zip(own_buffers, exported["buffers"]): + if buffer is None: + assert address is None + elif buffer.size: + assert buffer.address == address + if pa.types.is_dictionary(array.type): + _assert_buffers(array.dictionary, exported["dictionary"]) + elif pa.types.is_struct(array.type): + assert array.type.num_fields == len(exported["children"]) + for index, child in enumerate(exported["children"]): + _assert_buffers(array.field(index), child) + elif pa.types.is_list(array.type) or pa.types.is_map(array.type): + assert len(exported["children"]) == 1 + _assert_buffers(array.values, exported["children"][0]) + else: + assert not exported["children"] + + +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(), keys_sorted=True)), + ] + ) + 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.DictionaryArray.from_arrays( + pa.array([0, None, 1, 0], type=pa.int8()), + pa.array(["red", "blue"]), + ordered=True, + ), + 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_chunks_and_buffer_addresses(monkeypatch): + batch = _mixed_batch() + # Sliced fixtures exercise nontrivial IPC input. The measured zero-copy + # boundary starts after Rust decodes it, not before IPC serialization. + batches = [batch.slice(1, 2), batch.slice(0, 1), batch.slice(3, 1)] + expected = pa.Table.from_batches(batches) + spy = _StreamSpy(monkeypatch) + with _dataframe(batches) as (frame, _): + actual = frame.toArrow() + assert actual.equals(expected, check_metadata=True) + assert len(spy.batches) == len(batches) + assert spy.import_count == 1 + assert spy.next_count == len(batches) + 1 + assert spy.stream_releases == 1 + assert not spy.errors + assert spy.release_counts == [0] * len(batches) + for batch_index, exported in enumerate(spy.batches): + assert exported["length"] == batches[batch_index].num_rows + assert len(exported["children"]) == actual.num_columns + for column_index, child in enumerate(exported["children"]): + column = actual.column(column_index) + assert column.num_chunks == len(batches) + _assert_buffers(column.chunk(batch_index), child) + del column, actual + gc.collect() + assert spy.release_counts == [1] * len(batches) + + +def test_handoff_buffers_outlive_dataframe_session_and_table(monkeypatch): + batch = _mixed_batch() + expected = batch.to_pydict() + spy = _StreamSpy(monkeypatch) + + def collect(): + with _dataframe([batch]) as (frame, session): + table = frame.toArrow() + session.stop() + return table + + table = collect() + gc.collect() + assert spy.import_count == 1 + assert spy.stream_releases == 1 + assert not spy.errors + assert spy.release_counts == [0] + assert table.to_pydict() == expected + chunk = table.column(0).chunk(0).slice(1, 3) + address = chunk.buffers()[1].address + del table + gc.collect() + assert chunk.to_pylist() == [None, -3, 4] + assert chunk.buffers()[1].address == address + assert spy.release_counts == [0] + del chunk + gc.collect() + assert spy.release_counts == [1] + + +@pytest.mark.parametrize("kind", ["no-batches", "empty-typed", "zero-columns"]) +def test_handoff_empty_and_zero_column_results(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([{}, {}, {}], type=pa.struct([])) + batch = pa.RecordBatch.from_struct_array(structs) + batches = [batch.replace_schema_metadata({b"zero-column": b"preserved"})] + expected = pa.Table.from_batches(batches) + with _dataframe(batches) 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) + + +def test_to_pandas_uses_the_same_handoff(monkeypatch): + batch = _mixed_batch() + indices = [0, 2, 4, 5, 7] + schema = pa.schema( + [batch.schema.field(index) for index in indices], metadata=batch.schema.metadata + ) + arrays = [batch.column(index) for index in indices] + batches = [pa.RecordBatch.from_arrays(arrays, schema=schema)] + expected = pa.Table.from_batches(batches).to_pandas() + spy = _StreamSpy(monkeypatch) + with _dataframe(batches) as (frame, _): + actual = frame.toPandas() + pd.testing.assert_frame_equal(actual, expected) + assert len(spy.batches) == 1 + assert spy.import_count == 1 + assert spy.stream_releases == 1 + assert not spy.errors + del actual + gc.collect() + assert spy.release_counts == [1] + + +@pytest.mark.parametrize("consume_before_reject", [False, True]) +def test_import_rejection_releases_stream_once(monkeypatch, consume_before_reject): + batch = _mixed_batch() + spy = _StreamSpy(monkeypatch, reject=True, consume_before_reject=consume_before_reject) + with _dataframe([batch]) as (frame, _): + with pytest.raises(RuntimeError, match="test importer rejected"): + frame.toArrow() + gc.collect() + assert spy.import_count == int(consume_before_reject) + assert spy.next_count == 0 + assert spy.stream_releases == 1 + assert not spy.batches + assert not spy.release_counts + assert not spy.errors + + +@pytest.mark.parametrize("fail_next_at", [1, 2]) +def test_stream_read_failure_releases_stream_and_arrays_once(monkeypatch, fail_next_at): + batch = _mixed_batch() + spy = _StreamSpy(monkeypatch, fail_next_at=fail_next_at) + with _dataframe([batch, batch, batch]) as (frame, _): + with pytest.raises((pa.ArrowInvalid, OSError), match="test stream read failure"): + frame.toArrow() + gc.collect() + assert len(spy.batches) == fail_next_at - 1 + assert spy.import_count == 1 + assert spy.next_count == fail_next_at + assert spy.stream_releases == 1 + assert not spy.errors + # A batch already yielded to PyArrow has its own release callback. Releasing + # the failed stream also drops the Rust batches that were never exported. + assert spy.release_counts == [1] * (fail_next_at - 1) + + +def test_schema_import_rejection_prevents_stream_import(monkeypatch): + batch = _mixed_batch() + spy = _StreamSpy(monkeypatch) + + class SchemaProxy: + @staticmethod + def _import_from_c(address): + raise RuntimeError("test schema importer rejected the schema") + + monkeypatch.setattr(pa, "Schema", SchemaProxy) + with _dataframe([batch]) as (frame, _): + with pytest.raises(RuntimeError, match="test schema importer rejected"): + frame.toArrow() + assert not spy.callbacks + + +@pytest.mark.parametrize("kind", ["field", "nested-field", "timezone", "dictionary-timezone"]) +def test_nul_schema_strings_fall_back_without_executing_again(monkeypatch, kind): + if kind == "field": + batch = pa.record_batch([pa.array([1, None])], names=["bad\x00name"]) + elif kind == "nested-field": + data_type = pa.struct([pa.field("bad\x00name", pa.int64())]) + values = pa.array([{"bad\x00name": 1}, None], type=data_type) + batch = pa.record_batch([values], names=["nested"]) + else: + values = pa.array([0, 1], type=pa.timestamp("us", tz="UTC\x00suffix")) + if kind == "dictionary-timezone": + indices = pa.array([0, None, 1], type=pa.int8()) + values = pa.DictionaryArray.from_arrays(indices, values) + batch = pa.record_batch([values], names=["timestamp"]) + expected = pa.Table.from_batches([batch]) + spy = _StreamSpy(monkeypatch) + calls = {} + with _dataframe([batch], 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 calls == {"execute": 1} + assert not spy.callbacks + + +def test_nul_metadata_still_uses_zero_copy_handoff(monkeypatch): + 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) + spy = _StreamSpy(monkeypatch) + with _dataframe([batch]) as (frame, _): + actual = frame.toArrow() + assert actual.equals(pa.Table.from_batches([batch]), check_metadata=True) + assert len(spy.batches) == 1 + assert spy.import_count == 1 + assert spy.stream_releases == 1 + assert not spy.errors + _assert_buffers(actual.column(0).chunk(0), spy.batches[0]["children"][0]) + del actual + gc.collect() + assert spy.release_counts == [1] + + +def test_mismatched_batch_schemas_are_rejected_before_import(monkeypatch): + batches = [ + pa.record_batch([pa.array([1])], names=["value"]), + pa.record_batch([pa.array(["text"])], names=["value"]), + ] + spy = _StreamSpy(monkeypatch) + with _dataframe(batches) as (frame, _): + with pytest.raises(ValueError, match="inconsistent schemas"): + frame.toArrow() + assert not spy.callbacks + + +def test_mismatched_dictionary_order_is_rejected_before_import(monkeypatch): + ordered = pa.DictionaryArray.from_arrays( + pa.array([0, 1], type=pa.int8()), pa.array(["a", "b"]), ordered=True + ) + unordered = pa.DictionaryArray.from_arrays(ordered.indices, ordered.dictionary) + batches = [pa.record_batch([array], names=["value"]) for array in (ordered, unordered)] + spy = _StreamSpy(monkeypatch) + with _dataframe(batches) as (frame, _): + with pytest.raises(ValueError, match="inconsistent schemas"): + frame.toArrow() + assert not spy.callbacks + + +def test_malformed_ipc_is_reported_before_import(monkeypatch): + spy = _StreamSpy(monkeypatch) + with _dataframe([], payloads=[(1, b"not arrow ipc")]) as (frame, _): + with pytest.raises(RuntimeError, match="Arrow"): + frame.toArrow() + assert not spy.callbacks From afb69bcac6bcdd3d772c05b10ac38f2fbe82c95c Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Wed, 28 Oct 2026 16:49:38 -0700 Subject: [PATCH 2/7] Use arrow-rs reader conversion for the C Stream handoff Apply LC's suggestion to export one boxed RecordBatchReader through arrow-pyarrow and materialize it with read_all. Remove the custom stream callbacks, unused converter crate, and dedicated converter workflow. Keep schema consistency checks, standalone schema preflight, and the NUL schema IPC fallback. Retain handoff tests in the coverage workflow and keep PyO3 at 0.28.3. Generated-by: OpenAI Codex (application version unavailable). --- .github/workflows/python-arrow.yml | 59 -- Cargo.lock | 22 +- Cargo.toml | 3 - crates/pyspark-rs/Cargo.toml | 5 +- crates/pyspark-rs/src/arrow.rs | 106 ++++ crates/pyspark-rs/src/dataframe.rs | 6 +- crates/pyspark-rs/src/lib.rs | 1 + crates/spark-python-arrow/Cargo.toml | 33 -- crates/spark-python-arrow/README.md | 113 ---- crates/spark-python-arrow/src/bridge.rs | 416 -------------- crates/spark-python-arrow/src/lib.rs | 23 - .../tests/bridge_regressions.rs | 218 -------- crates/spark-python-arrow/tests/conversion.rs | 281 ---------- dev/design/architecture.md | 16 +- docs/architecture.md | 18 +- python/tests/test_arrow_handoff.py | 511 ++++++++---------- 16 files changed, 357 insertions(+), 1474 deletions(-) delete mode 100644 .github/workflows/python-arrow.yml create mode 100644 crates/pyspark-rs/src/arrow.rs delete mode 100644 crates/spark-python-arrow/Cargo.toml delete mode 100644 crates/spark-python-arrow/README.md delete mode 100644 crates/spark-python-arrow/src/bridge.rs delete mode 100644 crates/spark-python-arrow/src/lib.rs delete mode 100644 crates/spark-python-arrow/tests/bridge_regressions.rs delete mode 100644 crates/spark-python-arrow/tests/conversion.rs diff --git a/.github/workflows/python-arrow.yml b/.github/workflows/python-arrow.yml deleted file mode 100644 index c4bde403..00000000 --- a/.github/workflows/python-arrow.yml +++ /dev/null @@ -1,59 +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. -# - -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 - rustup default stable - - name: Install PyArrow - run: python -m pip install "pyarrow==26.0.0" - - name: Test Arrow table export - # Test separately from pyspark-rs: extension-module feature unification - # would disable libpython linking for the embedded Python tests. - run: | - cargo test -p spark-python-arrow --locked --features pyo3/abi3-py39 - - name: Check converter Clippy - run: | - cargo clippy -p spark-python-arrow \ - --locked --all-targets --features pyo3/abi3-py39 -- -D warnings diff --git a/Cargo.lock b/Cargo.lock index 92e73626..29e7d3fe 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -168,6 +168,7 @@ dependencies = [ "arrow-ipc", "arrow-json", "arrow-ord", + "arrow-pyarrow", "arrow-row", "arrow-schema", "arrow-select", @@ -322,6 +323,18 @@ dependencies = [ "arrow-select", ] +[[package]] +name = "arrow-pyarrow" +version = "58.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d29abdf672a81c1aeb57fd2661457f9918964d49aed0e9f18932535f2a9e49ce" +dependencies = [ + "arrow-array", + "arrow-data", + "arrow-schema", + "pyo3", +] + [[package]] name = "arrow-row" version = "58.3.0" @@ -3396,7 +3409,6 @@ dependencies = [ "prost", "pyo3", "serde_json", - "spark-python-arrow", "tokio", "tonic", ] @@ -3998,14 +4010,6 @@ dependencies = [ "windows-sys 0.61.2", ] -[[package]] -name = "spark-python-arrow" -version = "0.1.0" -dependencies = [ - "arrow", - "pyo3", -] - [[package]] name = "sqlparser" version = "0.60.0" diff --git a/Cargo.toml b/Cargo.toml index 31fe8090..5fe753a8 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -24,7 +24,6 @@ members = [ "crates/spark-connect-macros", "crates/spark-connect-build", "crates/pyspark-rs", - "crates/spark-python-arrow", "examples", ] exclude = [ @@ -35,8 +34,6 @@ exclude = [ # Rust-native `examples` crate both link the pure-Rust core directly; both are # workspace members but excluded from `default-members` so `cargo test`/`cargo build` # stay clean. Build them explicitly (`cargo build -p pyspark-rs` / `-p examples`). -# `spark-python-arrow` is also explicit: its tests embed Python and require PyArrow. -# Test it separately to avoid extension-module feature unification disabling libpython linking. # # The `examples/wasm-udfs` and `examples/wasm-udf-inline` crates are excluded from # the workspace entirely because they require the `wasm32-unknown-unknown` target, diff --git a/crates/pyspark-rs/Cargo.toml b/crates/pyspark-rs/Cargo.toml index 05278bad..24171693 100644 --- a/crates/pyspark-rs/Cargo.toml +++ b/crates/pyspark-rs/Cargo.toml @@ -2,6 +2,8 @@ name = "pyspark-rs" version.workspace = true edition.workspace = true +# Arrow 58.3.0, including arrow-pyarrow, requires Rust 1.85. +rust-version = "1.85" license.workspace = true description = "PyO3 bindings exposing the Rust Spark Connect client as a drop-in pyspark package" # The Python extension (cdylib) is distributed as a PyPI wheel, not on crates.io. @@ -22,9 +24,8 @@ non_snake_case = "allow" spark-connect = { package = "apache-spark-connect", path = "../spark-connect" } spark-connect-core = { package = "apache-spark-connect-core", path = "../spark-connect-core" } spark-connect-proto = { package = "apache-spark-connect-proto", path = "../spark-connect-proto" } -spark-python-arrow = { path = "../spark-python-arrow" } pyo3.workspace = true -arrow.workspace = true +arrow = { workspace = true, features = ["pyarrow"] } tokio.workspace = true prost.workspace = true tonic.workspace = true diff --git a/crates/pyspark-rs/src/arrow.rs b/crates/pyspark-rs/src/arrow.rs new file mode 100644 index 00000000..7b4425ea --- /dev/null +++ b/crates/pyspark-rs/src/arrow.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, ToPyArrow}; +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, +) -> PyResult> { + 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)) { + // C schema strings cannot contain NUL. Reuse the collected batches for + // the IPC compatibility path without executing the request again. + let bytes = py + .detach(|| -> arrow::error::Result> { + let mut bytes = Vec::new(); + let mut writer = FileWriter::try_new(&mut bytes, &schema)?; + for batch in &batches { + writer.write(batch)?; + } + writer.finish()?; + drop(writer); + Ok(bytes) + }) + .map_err(|error| PyRuntimeError::new_err(error.to_string()))?; + return Ok(py + .import("pyarrow")? + .getattr("ipc")? + .call_method1("open_file", (PyBytes::new(py, &bytes),))? + .call_method0("read_all")? + .unbind()); + } + + // Older PyArrow can abort on unsupported stream schemas. Use its fallible + // standalone schema importer before handing over the reader. + schema.to_pyarrow(py)?; + let reader: Box = Box::new(RecordBatchIterator::new( + batches.into_iter().map(Ok), + schema, + )); + Ok(reader.into_pyarrow(py)?.call_method0("read_all")?.unbind()) +} + +fn field_needs_ipc(field: &Field) -> bool { + field.name().contains('\0') || data_type_needs_ipc(field.data_type()) +} + +fn data_type_needs_ipc(data_type: &DataType) -> bool { + match data_type { + DataType::Timestamp(_, Some(timezone)) => timezone.contains('\0'), + DataType::List(field) + | DataType::LargeList(field) + | DataType::ListView(field) + | DataType::LargeListView(field) + | DataType::FixedSizeList(field, _) + | DataType::Map(field, _) => field_needs_ipc(field), + DataType::Struct(fields) => fields.iter().any(|field| field_needs_ipc(field)), + DataType::Union(fields, _) => fields.iter().any(|(_, field)| field_needs_ipc(field)), + DataType::RunEndEncoded(run_ends, values) => { + field_needs_ipc(run_ends) || field_needs_ipc(values) + } + DataType::Dictionary(keys, values) => { + data_type_needs_ipc(keys) || data_type_needs_ipc(values) + } + _ => false, + } +} diff --git a/crates/pyspark-rs/src/dataframe.rs b/crates/pyspark-rs/src/dataframe.rs index 25529563..27eb5dae 100644 --- a/crates/pyspark-rs/src/dataframe.rs +++ b/crates/pyspark-rs/src/dataframe.rs @@ -1166,15 +1166,15 @@ impl PyDataFrame { /// Collect the DataFrame into a `pyarrow.Table`. /// /// Mirrors `pyspark.sql.DataFrame.toArrow()`: the result is decoded from the - /// server's Arrow IPC stream and handed to `pyarrow` through the shared - /// converter, sharing compatible decoded Arrow value buffers. Bitmap alignment + /// server's Arrow IPC stream and handed to `pyarrow` through arrow-rs's + /// reader conversion, sharing compatible decoded Arrow value buffers. Bitmap alignment /// and schemas requiring the IPC compatibility fallback can still copy. #[pyo3(name = "toArrow")] fn to_arrow(&self, py: Python<'_>) -> PyResult> { let batches = py .detach(|| self.dataframe.collect_record_batches()) .to_pyerr()?; - spark_python_arrow::bridge::record_batches_to_pyarrow(py, batches) + crate::arrow::record_batches_to_pyarrow(py, batches) } /// Collect the DataFrame into a pandas DataFrame. diff --git a/crates/pyspark-rs/src/lib.rs b/crates/pyspark-rs/src/lib.rs index 43a850f4..726b4fb5 100644 --- a/crates/pyspark-rs/src/lib.rs +++ b/crates/pyspark-rs/src/lib.rs @@ -24,6 +24,7 @@ pub(crate) fn coerce_option_value(v: &Bound<'_, PyAny>) -> PyResult - -# Arrow table export to PyArrow - -`spark-python-arrow` exports decoded Rust Arrow record batches to PyArrow for -`pyspark-rs`. It has one public entry point, -`bridge::record_batches_to_pyarrow`, used by `DataFrame.toArrow()`. It does not -depend on Spark Connect transport, generated protobufs, or an async runtime. - -Python-to-Rust Arrow imports, standalone array/type exports, and worker-side UDF -conversion are outside this change. Scalar input conversion and row -materialization remain in `pyspark-rs` until they have another concrete consumer. - -## Export contract - -The Arrow C Stream Interface path preserves batch boundaries, schema and field -metadata, and zero-column row counts. It shares compatible Arrow value and offset -buffers with PyArrow. The caller supplies valid Rust record batches and owns its -Python interpreter; the library does not initialize Python or retain session state. - -`toArrow()` first collects all decoded batches into a Rust `Vec`. -The bridge exports a native Arrow reader over those batches, imports it once as -a PyArrow `RecordBatchReader`, and calls `read_all()`. PyArrow pulls each batch -through native callbacks, without a Python import call per batch or a temporary -Python list of batches. The callbacks do not call Python. C Stream is an -in-process ownership interface, not an IPC byte stream; it does not make result -collection network-streaming or bounded-memory. - -PyArrow consumes the exported stream and retains each batch's release callback, -so the table's buffers outlive the reader and Rust's temporary wrappers. If an -importer raises before consuming the stream, Rust releases it and its remaining -batches. PyArrow and Python code in the same interpreter are trusted, including -the public aliases used by the export path. - -Field names or timezone strings containing NUL use an IPC compatibility fallback -over the already collected batches, without another RPC. Length-prefixed schema -metadata can contain NUL without this fallback. The stream's `get_schema` callback -uses a schema descriptor that retains sorted-map and dictionary-order flags when -present; the other callbacks forward to Arrow's native stream implementation. - -Table export raises `ValueError` if batches have inconsistent schemas, including -field metadata or dictionary ordering. This is stricter than the previous IPC -export path, which wrote batches under the first batch's schema. An empty batch -vector still produces an empty table with an empty schema; preserving schemas -from schema-only IPC streams requires a separate change to the decoding path. - -Zero-copy describes compatible Arrow **value and offset buffers**, not schema -descriptors, Python scalar/row objects, pandas materialization, IPC decoding, or -the IPC fallback. The C Stream callbacks export batches using C Data descriptors; -sliced arrays whose value and null offsets differ can require copying the validity -bitmap during that export. - -## Integration - -This is an unpublished workspace implementation crate (`publish = false`), not -a separately released package. `pyspark-rs` uses a path dependency. The current -package name and `0.1.0` version do not identify a crates.io release; independent -publication would require Apache package naming and release-flow integration. -Source-based consumers can pin a repository revision and validate it in their -own dependency graph. The repository's root `LICENSE` covers this crate. - -The library does not enable PyO3's `extension-module` feature or select a minimum Python ABI; -`pyspark-rs` selects these features for the Python extension. - -PyO3 must resolve to one compatible version in each consuming dependency graph. -The manifest uses the same PyO3 0.28 line as the Python extension, with 0.28.3 -pinned in the workspace lockfile. This change does not upgrade PyO3. -This crate declares Rust 1.85, matching Arrow 58.3.0's declared minimum. -The existing workspace default of 1.82 is unchanged; it is not a claim that -these Python/Arrow packages build on 1.82. The minimum compiler has not been -independently tested locally. - -The Python Arrow export CI job tests this crate with the workspace lockfile. -These tests run separately from the extension build and are not repeated in coverage. - -## Tests - -With PyArrow installed in the Python environment used by PyO3: - -```sh -CARGO_NET_OFFLINE=true CARGO_REGISTRIES_CRATES_IO_PROTOCOL=git \ - cargo test -p spark-python-arrow --locked --features pyo3/abi3-py39 -``` - -Run this package separately from the extension's test/build command: Cargo -feature unification with an extension-module consumer can disable libpython -linking required by the embedded test executable. For a virtual environment, -set `PYO3_PYTHON` to its interpreter and expose its site-packages to the embedded -interpreter with `PYTHONPATH` when necessary. - -The standalone tests exercise the table export entry point: buffer identity and -lifetime, schema metadata, slices, chunks, empty shapes, schema mismatch errors, -and NUL fallback. The Python client's `test_arrow_handoff.py` exercises the public -`toArrow()` and `toPandas()` paths, including importer failures. Scalar regressions -live with the client tests. These are not substitutes for Spark's complete -Connect parity suite. diff --git a/crates/spark-python-arrow/src/bridge.rs b/crates/spark-python-arrow/src/bridge.rs deleted file mode 100644 index a791ddb6..00000000 --- a/crates/spark-python-arrow/src/bridge.rs +++ /dev/null @@ -1,416 +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. - -//! Export Arrow record batches to PyArrow through the Arrow C Stream 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::ffi::{c_char, c_int, CString}; -use std::sync::Arc; - -use arrow::datatypes::{DataType, Field, Schema}; -use arrow::error::ArrowError; -use arrow::ffi::{FFI_ArrowArray, FFI_ArrowSchema}; -use arrow::ffi_stream::FFI_ArrowArrayStream; -use arrow::ipc::writer::FileWriter; -use arrow::record_batch::{RecordBatch, RecordBatchIterator}; -use pyo3::exceptions::{PyRuntimeError, PyValueError}; -use pyo3::prelude::*; -use pyo3::types::PyBytes; - -// Called after table export has routed NUL names and timezones to the IPC fallback. -fn export_c_schema(schema: &Schema) -> Result { - let data_type = DataType::Struct(schema.fields().clone()); - export_c_data_type(&data_type)?.with_metadata(schema.metadata()) -} - -fn validate_pyarrow_schema(py: Python<'_>, schema: &Schema) -> PyResult<()> { - let mut schema = export_c_schema(schema).map_err(arrow_error)?; - let address = &mut schema as *mut FFI_ArrowSchema as usize; - py.import("pyarrow")? - .getattr("Schema")? - .call_method1("_import_from_c", (address,))?; - Ok(()) -} - -/// 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) -> PyResult> { - 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)) { - // C schema strings cannot contain NUL. Use the already collected - // batches for this rare compatibility path without repeating any RPC. - let bytes = py - .detach(|| -> arrow::error::Result> { - let mut bytes = Vec::new(); - let mut writer = FileWriter::try_new(&mut bytes, &schema)?; - for batch in &batches { - writer.write(batch)?; - } - writer.finish()?; - drop(writer); - Ok(bytes) - }) - .map_err(arrow_error)?; - return Ok(pa - .getattr("ipc")? - .call_method1("open_file", (PyBytes::new(py, &bytes),))? - .call_method0("read_all")? - .unbind()); - } - - // Older PyArrow versions can abort on unsupported stream schemas. Validate - // the same schema through the fallible standalone importer before streaming. - validate_pyarrow_schema(py, &schema)?; - let batches = RecordBatchIterator::new(batches.into_iter().map(Ok), schema.clone()); - let mut stream = with_exported_schema(FFI_ArrowArrayStream::new(Box::new(batches)), schema); - let address = &mut stream as *mut FFI_ArrowArrayStream as usize; - // The importer moves the stream and clears its release callback. Rust drops - // any unconsumed stream on failure; imported arrays retain their own buffers. - // The pointer importer is available at our minimum supported PyArrow (10). - let reader = pa - .getattr("RecordBatchReader")? - .call_method1("_import_from_c", (address,))?; - Ok(reader.call_method0("read_all")?.unbind()) -} - -fn arrow_error(error: ArrowError) -> PyErr { - PyRuntimeError::new_err(error.to_string()) -} - -// Arrow 58's stream exporter needs the same sorted-map schema repair as the -// standalone exporter below. Own the native stream and override only its schema -// conversion; array export, iteration, and remaining-batch cleanup stay in Arrow. -struct ExportStream { - inner: FFI_ArrowArrayStream, - schema: Arc, - schema_error: Option, -} - -fn with_exported_schema(inner: FFI_ArrowArrayStream, schema: Arc) -> FFI_ArrowArrayStream { - let state = Box::new(ExportStream { - inner, - schema, - schema_error: None, - }); - FFI_ArrowArrayStream { - get_schema: Some(stream_schema), - get_next: Some(stream_next), - get_last_error: Some(stream_error), - release: Some(stream_release), - private_data: Box::into_raw(state).cast(), - } -} - -// These callbacks run without the GIL and must not call Python. The C Stream -// contract gives each callback exclusive access to a live stream and, where -// applicable, writable output storage. private_data is our Box until release. -unsafe extern "C" fn stream_schema( - stream: *mut FFI_ArrowArrayStream, - out: *mut FFI_ArrowSchema, -) -> c_int { - // SAFETY: the consumer supplies the live stream constructed above. - let state = unsafe { &mut *(*stream).private_data.cast::() }; - state.schema_error = None; - match export_c_schema(&state.schema) { - Ok(schema) => { - // SAFETY: out is writable. Each call transfers a fresh descriptor, - // which the consumer releases independently of the stream. - unsafe { std::ptr::write(out, schema) }; - 0 - } - Err(error) => { - // Error strings may contain NUL; never panic across the C boundary. - state.schema_error = - Some(CString::new(error.to_string().replace('\0', "\\0")).unwrap_or_default()); - 22 // EINVAL, specified by the C Stream interface. - } - } -} - -unsafe extern "C" fn stream_next( - stream: *mut FFI_ArrowArrayStream, - out: *mut FFI_ArrowArray, -) -> c_int { - // SAFETY: the consumer supplies the live stream constructed above. - let state = unsafe { &mut *(*stream).private_data.cast::() }; - state.schema_error = None; - match state.inner.get_next { - // SAFETY: delegate with the inner stream's own private_data, not ours. - Some(next) => unsafe { next(&mut state.inner, out) }, - None => { - state.schema_error = Some(c"Arrow stream has no get_next callback".into()); - 22 - } - } -} - -unsafe extern "C" fn stream_error(stream: *mut FFI_ArrowArrayStream) -> *const c_char { - // SAFETY: the consumer supplies the live stream constructed above. - let state = unsafe { &mut *(*stream).private_data.cast::() }; - if let Some(error) = &state.schema_error { - return error.as_ptr(); - } - match state.inner.get_last_error { - // SAFETY: the inner error string stays owned by the inner stream. - Some(error) => unsafe { error(&mut state.inner) }, - None => std::ptr::null(), - } -} - -unsafe extern "C" fn stream_release(stream: *mut FFI_ArrowArrayStream) { - if stream.is_null() { - return; - } - // SAFETY: the consumer supplies writable storage for this stream. - let stream = unsafe { &mut *stream }; - if stream.release.is_none() { - return; - } - let state = stream.private_data.cast::(); - stream.release = None; - stream.private_data = std::ptr::null_mut(); - stream.get_schema = None; - stream.get_next = None; - stream.get_last_error = None; - // SAFETY: this is the Box installed above, consumed exactly once. Dropping - // the inner stream releases unread batches, not independently exported arrays. - drop(unsafe { Box::from_raw(state) }); -} - -// Arrow 58's Field conversion overwrites a map's type flags with its field flags. -// Rebuild the descriptor tree using public constructors, retaining both. Only -// descriptors are rebuilt: array buffers are not involved in schema conversion. -fn export_c_data_type(data_type: &DataType) -> Result { - fn has_sorted_map(data_type: &DataType) -> bool { - matches!(data_type, DataType::Map(_, true)) - || child_fields(data_type) - .iter() - .any(|field| has_sorted_map(field.data_type())) - || matches!(data_type, DataType::Dictionary(_, values) if has_sorted_map(values)) - } - - fn rebuild( - source: &FFI_ArrowSchema, - data_type: &DataType, - ) -> Result { - let children = child_fields(data_type) - .into_iter() - .enumerate() - .map(|(index, field)| rebuild(source.child(index), field.data_type())) - .collect::, _>>()?; - let dictionary = match (source.dictionary(), data_type) { - (Some(schema), DataType::Dictionary(_, values)) => Some(rebuild(schema, values)?), - _ => None, - }; - let mut schema = FFI_ArrowSchema::try_new(source.format(), children, dictionary)? - .with_metadata(source.metadata()?)?; - if let Some(name) = source.name() { - schema = schema.with_name(name)?; - } - schema.flags = source.flags; - if matches!(data_type, DataType::Map(_, true)) { - // ARROW_FLAG_MAP_KEYS_SORTED in the public Arrow C Data Interface. - schema.flags |= 0b100; - } - Ok(schema) - } - - let schema = FFI_ArrowSchema::try_from(data_type)?; - if has_sorted_map(data_type) { - rebuild(&schema, data_type) - } else { - Ok(schema) - } -} - -fn child_fields(data_type: &DataType) -> Vec<&Field> { - match data_type { - DataType::List(field) - | DataType::LargeList(field) - | DataType::ListView(field) - | DataType::LargeListView(field) - | DataType::FixedSizeList(field, _) - | DataType::Map(field, _) => vec![field], - DataType::Struct(fields) => fields.iter().map(AsRef::as_ref).collect(), - DataType::Union(fields, _) => fields.iter().map(|(_, field)| field.as_ref()).collect(), - DataType::RunEndEncoded(run_ends, values) => vec![run_ends, values], - _ => vec![], - } -} - -fn schemas_equal(left: &Schema, right: &Schema) -> bool { - fn same_order(left: &DataType, right: &DataType) -> bool { - let fields_match = - child_fields(left) - .into_iter() - .zip(child_fields(right)) - .all(|(left, right)| { - left.dict_is_ordered() == right.dict_is_ordered() - && same_order(left.data_type(), right.data_type()) - }); - fields_match - && match (left, right) { - (DataType::Dictionary(_, left), DataType::Dictionary(_, right)) => { - same_order(left, right) - } - _ => true, - } - } - - // Arrow's structural equality intentionally excludes dictionary-order flags. - // PyArrow includes them, so a table must not relabel later batches' ordering. - left == right - && same_order( - &DataType::Struct(left.fields().clone()), - &DataType::Struct(right.fields().clone()), - ) -} - -fn field_needs_ipc(field: &Field) -> bool { - field.name().contains('\0') || data_type_needs_ipc(field.data_type()) -} - -fn data_type_needs_ipc(data_type: &DataType) -> bool { - match data_type { - DataType::Timestamp(_, Some(timezone)) => timezone.contains('\0'), - DataType::List(field) - | DataType::LargeList(field) - | DataType::ListView(field) - | DataType::LargeListView(field) - | DataType::FixedSizeList(field, _) - | DataType::Map(field, _) => field_needs_ipc(field), - DataType::Struct(fields) => fields.iter().any(|field| field_needs_ipc(field)), - DataType::Union(fields, _) => fields.iter().any(|(_, field)| field_needs_ipc(field)), - DataType::RunEndEncoded(run_ends, values) => { - field_needs_ipc(run_ends) || field_needs_ipc(values) - } - DataType::Dictionary(keys, values) => { - data_type_needs_ipc(keys) || data_type_needs_ipc(values) - } - _ => false, - } -} - -#[cfg(test)] -mod tests { - use std::collections::HashMap; - use std::ffi::CStr; - - use arrow::datatypes::TimeUnit; - - use super::*; - - #[test] - fn invalid_time_type_is_rejected_without_panicking() { - assert!(export_c_data_type(&DataType::Time32(TimeUnit::Nanosecond)).is_err()); - } - - #[test] - fn stream_schema_can_be_requested_and_released_repeatedly() { - let entries = Field::new( - "entries", - DataType::Struct( - vec![ - Field::new("key", DataType::Utf8, false), - Field::new("value", DataType::Int64, true), - ] - .into(), - ), - false, - ); - let schema = Arc::new(Schema::new_with_metadata( - vec![Field::new( - "maps", - DataType::Map(Arc::new(entries), true), - true, - )], - HashMap::from([("key".into(), "value\0metadata".into())]), - )); - let batches = RecordBatchIterator::new(std::iter::empty(), schema.clone()); - let mut stream = with_exported_schema(FFI_ArrowArrayStream::new(Box::new(batches)), schema); - let mut exported = Vec::new(); - for _ in 0..3 { - let mut out = FFI_ArrowSchema::empty(); - // SAFETY: both descriptors are live and out is writable. - assert_eq!(unsafe { stream_schema(&mut stream, &mut out) }, 0); - assert_ne!(out.child(0).flags & 0b100, 0); - assert_eq!(out.metadata().unwrap()["key"], "value\0metadata"); - exported.push(out); - } - drop(stream); - // Schema descriptors have their own owners after the stream is gone. - for schema in exported { - assert_eq!(schema.child(0).name(), Some("maps")); - } - } - - #[test] - fn stream_schema_errors_have_stable_storage_until_next_operation() { - let schema = Arc::new(Schema::new(vec![Field::new( - "invalid", - DataType::Time32(TimeUnit::Nanosecond), - true, - )])); - let batches = RecordBatchIterator::new(std::iter::empty(), schema.clone()); - let mut stream = with_exported_schema(FFI_ArrowArrayStream::new(Box::new(batches)), schema); - let mut out = FFI_ArrowSchema::empty(); - // SAFETY: the stream and output are live, writable descriptors. - assert_eq!(unsafe { stream_schema(&mut stream, &mut out) }, 22); - // SAFETY: the error pointer is owned by the stream until the next operation. - let error = unsafe { stream_error(&mut stream) }; - assert!(!error.is_null()); - assert_eq!(error, unsafe { stream_error(&mut stream) }); - assert!(!unsafe { CStr::from_ptr(error) }.to_bytes().is_empty()); - let mut array = FFI_ArrowArray::empty(); - assert_eq!(unsafe { stream_next(&mut stream, &mut array) }, 0); - assert!(unsafe { stream_error(&mut stream) }.is_null()); - } - - #[test] - fn stream_forwards_inner_iteration_errors() { - let schema = Arc::new(Schema::empty()); - let batches = RecordBatchIterator::new( - std::iter::once(Err(ArrowError::InvalidArgumentError( - "test read failure".into(), - ))), - schema.clone(), - ); - let mut stream = with_exported_schema(FFI_ArrowArrayStream::new(Box::new(batches)), schema); - let mut array = FFI_ArrowArray::empty(); - // SAFETY: the live stream owns the error string, and array is writable. - assert_eq!(unsafe { stream_next(&mut stream, &mut array) }, 22); - let message = unsafe { CStr::from_ptr(stream_error(&mut stream)) }; - assert!(message.to_str().unwrap().contains("test read failure")); - } -} diff --git a/crates/spark-python-arrow/src/lib.rs b/crates/spark-python-arrow/src/lib.rs deleted file mode 100644 index 60c5d084..00000000 --- a/crates/spark-python-arrow/src/lib.rs +++ /dev/null @@ -1,23 +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. - -//! Arrow record batch export to PyArrow without a dependency on Spark transport. -//! -//! The bridge shares compatible Arrow value buffers. Scalar conversion and row materialization -//! remain in the client binding. The caller owns its Python interpreter. - -pub mod bridge; diff --git a/crates/spark-python-arrow/tests/bridge_regressions.rs b/crates/spark-python-arrow/tests/bridge_regressions.rs deleted file mode 100644 index 4939b81d..00000000 --- a/crates/spark-python-arrow/tests/bridge_regressions.rs +++ /dev/null @@ -1,218 +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. - -use std::ffi::CStr; -use std::sync::Arc; - -use arrow::array::{ - Array, ArrayRef, DictionaryArray, Int64Array, Int8Array, ListArray, MapArray, StringArray, - StructArray, -}; -use arrow::buffer::{NullBuffer, OffsetBuffer}; -use arrow::datatypes::{DataType, Field, Int8Type, Schema}; -use arrow::record_batch::RecordBatch; -use pyo3::exceptions::PyValueError; -use pyo3::prelude::*; -use pyo3::types::PyDict; -use spark_python_arrow::bridge::record_batches_to_pyarrow; - -fn assert_export_equal(py: Python<'_>, batch: RecordBatch, expected_code: &CStr) -> PyResult<()> { - let locals = PyDict::new(py); - locals.set_item("pa", py.import("pyarrow")?)?; - py.run(expected_code, Some(&locals), Some(&locals))?; - let expected = locals.get_item("expected")?.expect("expected table"); - let actual = record_batches_to_pyarrow(py, vec![batch])?; - assert!(actual - .bind(py) - .call_method1("equals", (&expected,))? - .is_truthy()?); - assert!(actual - .bind(py) - .getattr("schema")? - .call_method1("equals", (expected.getattr("schema")?, true))? - .is_truthy()?); - Ok(()) -} - -#[test] -fn table_export_preserves_sorted_map_flags_in_nested_fields() { - Python::initialize(); - Python::attach(|py| -> PyResult<()> { - let entries = StructArray::new( - vec![ - Field::new("key", DataType::Utf8, false), - Field::new("value", DataType::Int64, true), - ] - .into(), - vec![ - Arc::new(StringArray::from(vec!["a", "b"])), - Arc::new(Int64Array::from(vec![1, 2])), - ], - None, - ); - let maps: ArrayRef = Arc::new(MapArray::new( - Arc::new(Field::new("entries", entries.data_type().clone(), false)), - OffsetBuffer::new(vec![0, 2, 2, 2].into()), - entries, - Some(NullBuffer::from(vec![true, false, true])), - true, - )); - let lists: ArrayRef = Arc::new(ListArray::new( - Arc::new(Field::new("item", maps.data_type().clone(), true)), - OffsetBuffer::new(vec![0, 1, 2, 3].into()), - maps.clone(), - None, - )); - let nested: ArrayRef = Arc::new(StructArray::new( - vec![Field::new("maps", lists.data_type().clone(), true)].into(), - vec![lists], - None, - )); - let schema = Arc::new(Schema::new(vec![ - Field::new("maps", maps.data_type().clone(), true), - Field::new("nested", nested.data_type().clone(), true), - ])); - let batch = RecordBatch::try_new(schema, vec![maps, nested]).unwrap(); - assert_export_equal( - py, - batch, - c"map_type = pa.map_(pa.string(), pa.int64(), keys_sorted=True) -maps = pa.array([[('a', 1), ('b', 2)], None, []], type=map_type) -nested_type = pa.struct([pa.field('maps', pa.list_(map_type))]) -nested = pa.array([ - {'maps': [[('a', 1), ('b', 2)]]}, {'maps': [None]}, {'maps': [[]]} -], type=nested_type) -expected = pa.Table.from_arrays([maps, nested], names=['maps', 'nested'])", - ) - }) - .unwrap(); -} - -#[test] -fn table_export_preserves_ordered_dictionaries_in_nested_fields() { - Python::initialize(); - Python::attach(|py| -> PyResult<()> { - let dictionary: ArrayRef = Arc::new(DictionaryArray::::new( - Int8Array::from(vec![Some(1), None, Some(0)]), - Arc::new(StringArray::from(vec!["a", "b"])), - )); - let field = Field::new("dictionary", dictionary.data_type().clone(), true) - .with_dict_is_ordered(true); - let nested: ArrayRef = Arc::new(StructArray::new( - vec![field.clone()].into(), - vec![dictionary.clone()], - None, - )); - let lists: ArrayRef = Arc::new(ListArray::new( - Arc::new(field.clone().with_name("item")), - OffsetBuffer::new(vec![0, 1, 2, 3].into()), - dictionary.clone(), - None, - )); - let schema = Arc::new(Schema::new(vec![ - field, - Field::new("nested", nested.data_type().clone(), true), - Field::new("lists", lists.data_type().clone(), true), - ])); - let batch = RecordBatch::try_new(schema, vec![dictionary, nested, lists]).unwrap(); - assert_export_equal( - py, - batch, - c"dictionary = pa.DictionaryArray.from_arrays( - pa.array([1, None, 0], type=pa.int8()), pa.array(['a', 'b']), ordered=True -) -nested = pa.StructArray.from_arrays([dictionary], names=['dictionary']) -lists = pa.ListArray.from_arrays(pa.array([0, 1, 2, 3], type=pa.int32()), dictionary) -expected = pa.Table.from_arrays([dictionary, nested, lists], - names=['dictionary', 'nested', 'lists'])", - ) - }) - .unwrap(); -} - -#[test] -fn table_export_rejects_inconsistent_nested_dictionary_order() { - Python::initialize(); - Python::attach(|py| { - let fields = |ordered| { - let dictionary = Field::new( - "dictionary", - DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)), - true, - ) - .with_dict_is_ordered(ordered); - [ - dictionary.clone(), - Field::new( - "nested", - DataType::Struct(vec![dictionary.clone()].into()), - true, - ), - Field::new("lists", DataType::List(Arc::new(dictionary)), true), - ] - }; - for (ordered, unordered) in fields(true).into_iter().zip(fields(false)) { - let batches = vec![ - RecordBatch::new_empty(Arc::new(Schema::new(vec![ordered]))), - RecordBatch::new_empty(Arc::new(Schema::new(vec![unordered]))), - ]; - let error = record_batches_to_pyarrow(py, batches).unwrap_err(); - assert!(error.is_instance_of::(py)); - assert!(error.to_string().contains("inconsistent schemas")); - } - }); -} - -#[test] -fn sliced_validity_shares_values_without_requiring_bitmap_identity() { - Python::initialize(); - Python::attach(|py| -> PyResult<()> { - let original = Int64Array::from(vec![Some(1), Some(2), None, Some(4), None]); - let slice = original.slice(1, 3); - assert_eq!(slice.nulls().unwrap().offset(), 1); - assert_eq!(slice.to_data().offset(), 0); - let address = slice.values().as_ptr() as usize; - let schema = Arc::new(Schema::new(vec![Field::new("x", DataType::Int64, true)])); - let batch = RecordBatch::try_new(schema, vec![Arc::new(slice)]).unwrap(); - let table = record_batches_to_pyarrow(py, vec![batch])?; - drop(original); - py.import("gc")?.call_method0("collect")?; - let array = table - .bind(py) - .call_method1("column", (0,))? - .call_method1("chunk", (0,))?; - assert_eq!( - array - .call_method0("buffers")? - .get_item(1)? - .getattr("address")? - .extract::()?, - address - ); - // Arrow may realign validity to the value buffer's offset. The bitmap's - // identity is not part of the value-buffer sharing promise. - assert_eq!(array.getattr("offset")?.extract::()?, 0); - assert_eq!( - array - .call_method0("to_pylist")? - .extract::>>()?, - vec![Some(2), None, Some(4)] - ); - Ok(()) - }) - .unwrap(); -} diff --git a/crates/spark-python-arrow/tests/conversion.rs b/crates/spark-python-arrow/tests/conversion.rs deleted file mode 100644 index b5773e41..00000000 --- a/crates/spark-python-arrow/tests/conversion.rs +++ /dev/null @@ -1,281 +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. - -//! Table export checks against the installed PyArrow. No Spark server is needed. - -use std::collections::HashMap; -use std::ffi::CStr; -use std::sync::Arc; - -use arrow::array::{Array, Int64Array, StringArray, StructArray}; -use arrow::datatypes::{DataType, Field, Schema, TimeUnit}; -use arrow::record_batch::{RecordBatch, RecordBatchOptions}; -use pyo3::exceptions::PyValueError; -use pyo3::prelude::*; -use pyo3::types::PyDict; -use spark_python_arrow::bridge::record_batches_to_pyarrow; - -fn evaluate<'py>(py: Python<'py>, expression: &CStr) -> PyResult> { - let locals = PyDict::new(py); - locals.set_item("pa", py.import("pyarrow")?)?; - py.eval(expression, Some(&locals), Some(&locals)) -} - -fn assert_schema_equal(actual: &Bound<'_, PyAny>, expected: &Bound<'_, PyAny>) -> PyResult<()> { - assert!(actual - .call_method1("equals", (expected, true))? - .is_truthy()?); - Ok(()) -} - -#[test] -fn table_export_shares_value_and_offset_buffers_after_rust_wrappers_drop() { - Python::initialize(); - Python::attach(|py| -> PyResult<()> { - let numbers = Int64Array::from(vec![Some(12), None, Some(34)]); - let strings = StringArray::from(vec![Some("first"), None, Some("last")]); - let addresses = [ - (0, 1, numbers.values().as_ptr() as usize), - (1, 1, strings.value_offsets().as_ptr() as usize), - (1, 2, strings.value_data().as_ptr() as usize), - ]; - let schema = Arc::new(Schema::new(vec![ - Field::new("number", DataType::Int64, true), - Field::new("text", DataType::Utf8, true), - ])); - let batch = - RecordBatch::try_new(schema, vec![Arc::new(numbers), Arc::new(strings)]).unwrap(); - // Export consumes the batch; only PyArrow owns the buffers afterward. - let table = record_batches_to_pyarrow(py, vec![batch])?; - py.import("gc")?.call_method0("collect")?; - for (column, buffer, address) in addresses { - let chunk = table - .bind(py) - .call_method1("column", (column,))? - .call_method1("chunk", (0,))?; - assert_eq!( - chunk - .call_method0("buffers")? - .get_item(buffer)? - .getattr("address")? - .extract::()?, - address - ); - } - let expected = evaluate( - py, - c"pa.table({'number': pa.array([12, None, 34], type=pa.int64()), - 'text': ['first', None, 'last']})", - )?; - assert!(table - .bind(py) - .call_method1("equals", (expected,))? - .is_truthy()?); - Ok(()) - }) - .unwrap(); -} - -#[test] -fn table_export_preserves_chunks_and_shares_each_buffer() { - Python::initialize(); - Python::attach(|py| -> PyResult<()> { - let schema = Arc::new(Schema::new(vec![Field::new("x", DataType::Int64, true)])); - let arrays = [ - Arc::new(Int64Array::from(vec![Some(1), None])), - Arc::new(Int64Array::from(vec![Some(3)])), - ]; - let addresses: Vec<_> = arrays - .iter() - .map(|array| array.values().as_ptr() as usize) - .collect(); - let batches = arrays - .into_iter() - .map(|array| RecordBatch::try_new(schema.clone(), vec![array]).unwrap()) - .collect(); - let table = record_batches_to_pyarrow(py, batches)?; - let column = table.bind(py).call_method1("column", (0,))?; - assert_eq!(column.getattr("num_chunks")?.extract::()?, 2); - for (index, address) in addresses.iter().enumerate() { - let chunk = column.call_method1("chunk", (index,))?; - assert_eq!( - chunk - .call_method0("buffers")? - .get_item(1)? - .getattr("address")? - .extract::()?, - *address - ); - } - assert_eq!( - column - .call_method0("to_pylist")? - .extract::>>()?, - vec![Some(1), None, Some(3)] - ); - Ok(()) - }) - .unwrap(); -} - -#[test] -fn table_export_preserves_nested_values_and_nul_metadata() { - Python::initialize(); - Python::attach(|py| -> PyResult<()> { - let field = Arc::new( - Field::new("value", DataType::Int64, true) - .with_metadata(HashMap::from([("field\0key".into(), "a\0b".into())])), - ); - let nested = StructArray::new( - vec![field].into(), - vec![Arc::new(Int64Array::from(vec![Some(1), None, Some(3)]))], - None, - ); - let schema = Arc::new(Schema::new_with_metadata( - vec![Field::new("nested", nested.data_type().clone(), true)], - HashMap::from([("schema_key".into(), "c\0d".into())]), - )); - let batch = RecordBatch::try_new(schema, vec![Arc::new(nested)]).unwrap(); - let table = record_batches_to_pyarrow(py, vec![batch])?; - let expected_schema = evaluate( - py, - c"pa.schema([pa.field('nested', pa.struct([ - pa.field('value', pa.int64(), metadata={b'field\\x00key': b'a\\x00b'}) - ]))], metadata={b'schema_key': b'c\\x00d'})", - )?; - assert_schema_equal(&table.bind(py).getattr("schema")?, &expected_schema)?; - let actual = table.bind(py).call_method0("to_pylist")?; - let expected = evaluate( - py, - c"[{'nested': {'value': 1}}, {'nested': {'value': None}}, - {'nested': {'value': 3}}]", - )?; - assert!(actual.eq(expected)?); - Ok(()) - }) - .unwrap(); -} - -#[test] -fn table_export_preserves_empty_schemas_and_zero_column_rows() { - Python::initialize(); - Python::attach(|py| -> PyResult<()> { - let empty = record_batches_to_pyarrow(py, vec![])?; - assert_eq!(empty.bind(py).getattr("num_rows")?.extract::()?, 0); - assert_eq!( - empty.bind(py).getattr("num_columns")?.extract::()?, - 0 - ); - assert_schema_equal( - &empty.bind(py).getattr("schema")?, - &evaluate(py, c"pa.schema([])")?, - )?; - - let schema = Arc::new(Schema::new(vec![Field::new("x", DataType::Int64, true)])); - let empty = record_batches_to_pyarrow(py, vec![RecordBatch::new_empty(schema)])?; - assert_eq!(empty.bind(py).getattr("num_rows")?.extract::()?, 0); - assert_schema_equal( - &empty.bind(py).getattr("schema")?, - &evaluate(py, c"pa.schema([('x', pa.int64())])")?, - )?; - - for rows in [0, 4] { - let schema = - Schema::empty().with_metadata(HashMap::from([("empty".into(), "metadata".into())])); - let options = RecordBatchOptions::new().with_row_count(Some(rows)); - let batch = - RecordBatch::try_new_with_options(Arc::new(schema), vec![], &options).unwrap(); - let table = record_batches_to_pyarrow(py, vec![batch])?; - assert_eq!( - table.bind(py).getattr("num_rows")?.extract::()?, - rows - ); - assert_eq!( - table.bind(py).getattr("num_columns")?.extract::()?, - 0 - ); - assert_schema_equal( - &table.bind(py).getattr("schema")?, - &evaluate(py, c"pa.schema([], metadata={b'empty': b'metadata'})")?, - )?; - } - Ok(()) - }) - .unwrap(); -} - -#[test] -fn table_export_rejects_inconsistent_fields_and_metadata() { - Python::initialize(); - Python::attach(|py| { - let field = Field::new("x", DataType::Int64, true); - let schema = Schema::new(vec![field.clone()]); - for other in [ - Schema::new(vec![Field::new("other", DataType::Int64, true)]), - schema - .clone() - .with_metadata(HashMap::from([("schema".into(), "other".into())])), - Schema::new(vec![field - .clone() - .with_metadata(HashMap::from([("field".into(), "other".into())]))]), - ] { - let batches = vec![ - RecordBatch::new_empty(Arc::new(schema.clone())), - RecordBatch::new_empty(Arc::new(other)), - ]; - let error = record_batches_to_pyarrow(py, batches).unwrap_err(); - assert!(error.is_instance_of::(py)); - assert!(error.to_string().contains("inconsistent schemas")); - } - }); -} - -#[test] -fn table_export_falls_back_to_ipc_for_nul_names_and_timezones() { - Python::initialize(); - Python::attach(|py| -> PyResult<()> { - let cases: [(Field, &CStr); 3] = [ - ( - Field::new("bad\0name", DataType::Int64, true), - c"pa.schema([('bad\\x00name', pa.int64())])", - ), - ( - Field::new( - "nested", - DataType::Struct(vec![Field::new("bad\0name", DataType::Int64, true)].into()), - true, - ), - c"pa.schema([('nested', pa.struct([('bad\\x00name', pa.int64())]))])", - ), - ( - Field::new( - "time", - DataType::Timestamp(TimeUnit::Microsecond, Some("UTC\0suffix".into())), - true, - ), - c"pa.schema([('time', pa.timestamp('us', tz='UTC\\x00suffix'))])", - ), - ]; - for (field, expected) in cases { - let batch = RecordBatch::new_empty(Arc::new(Schema::new(vec![field]))); - let table = record_batches_to_pyarrow(py, vec![batch])?; - assert_schema_equal(&table.bind(py).getattr("schema")?, &evaluate(py, expected)?)?; - } - Ok(()) - }) - .unwrap(); -} diff --git a/dev/design/architecture.md b/dev/design/architecture.md index 09cf031d..9df58a65 100644 --- a/dev/design/architecture.md +++ b/dev/design/architecture.md @@ -47,7 +47,6 @@ crates/ column, plan, dataframe, functions, session, readwriter, catalog, group, window, udf/udtf, streaming, observation pyspark-rs PyO3 bindings -> native module `_pyspark` - spark-python-arrow Arrow C Stream table export to PyArrow; no transport python/ pyspark/… thin drop-in package; public modules re-export from `_pyspark` so `from pyspark.sql import SparkSession` works @@ -69,8 +68,8 @@ Server-streaming RPCs (`ExecutePlan`) are consumed into Arrow record batches. ## Data path Results arrive as Arrow IPC streams. `arrow-rs` decodes them; `toArrow()` collects -all batches into a Rust vector and passes it to `spark-python-arrow`. The bridge -exports one native Arrow C Stream reader, imports it as a PyArrow +all batches into a Rust vector. `pyspark-rs` uses arrow-rs's `IntoPyArrow` for +one boxed `RecordBatchReader`, importing it as a PyArrow `RecordBatchReader`, and calls `read_all()`. PyArrow pulls batches through native callbacks that do not call Python, avoiding a per-batch Python import loop. This is an in-process buffer handoff, not network streaming or bounded-memory @@ -78,16 +77,15 @@ result collection. Compatible Arrow value and offset buffers are shared with PyArrow; sliced validity bitmaps may need realignment. The bridge validates batch schema -consistency, including metadata and dictionary ordering. Schemas with embedded +consistency, including metadata. Schemas with embedded NUL field names or timezones use an IPC compatibility fallback, without another -RPC. Schema metadata and ordering flags are preserved. `toPandas()` then calls +RPC. Spark results use unsorted maps and are not dictionary-encoded; this path +does not repair generic Arrow ordering flags. `toPandas()` then calls PyArrow's pandas conversion, which may copy. -Scalar input conversion and row materialization stay in `pyspark-rs`. The export -crate does not own schema inference, transport, or worker execution. `collect()` +Scalar input conversion and row materialization stay in `pyspark-rs`. `collect()` still uses the native client's owned rows; `createDataFrame(pandas)` still -materializes Python rows. Neither path is claimed to be zero-copy. See the -[table export contract](../../crates/spark-python-arrow/README.md). +materializes Python rows. Neither path is claimed to be zero-copy. ## UDF strategy (full compatibility) diff --git a/docs/architecture.md b/docs/architecture.md index 83e33c8c..ddd0cb90 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -63,7 +63,6 @@ A Cargo workspace of library crates, plus a Python skin: | `apache-spark-connect-core` | `crates/spark-connect-core` | Transport: channel, retries, reattach, artifacts, errors | | `apache-spark-connect` | `crates/spark-connect` | The DataFrame API: session, dataframe, column, functions, plan, group, catalog, window, readwriter, streaming, types | | `pyspark-rs` | `crates/pyspark-rs` | PyO3 bindings - builds the `_pyspark` extension module | -| `spark-python-arrow` (unpublished) | `crates/spark-python-arrow` | Arrow C Stream table export to PyArrow | | Python skin | `python/pyspark` | The drop-in `pyspark` package (+ vendored `pyspark.pandas`, `cloudpickle`, `pyspark.testing`) | ```mermaid @@ -71,7 +70,6 @@ flowchart LR PROTO["spark-connect-proto"] --> CORE["spark-connect-core"] CORE --> SC["spark-connect (DataFrame API)"] SC --> PYO3["pyspark-rs (PyO3)"] - PYARROW["spark-python-arrow (unpublished)"] --> PYO3 PYO3 --> SKIN["python/pyspark skin"] SC --> RUSTAPP["your Rust app"] ``` @@ -106,8 +104,9 @@ results survive transient disconnects. ### Python Arrow table handoff -`DataFrame.toArrow()` collects the decoded Arrow batches in Rust, then exports -one reader through the Arrow C Stream Interface. PyArrow imports that reader +`DataFrame.toArrow()` collects the decoded Arrow batches in Rust. Its `pyspark-rs` +wrapper uses arrow-rs's `IntoPyArrow` to export one boxed `RecordBatchReader` +through the Arrow C Stream Interface. PyArrow imports that reader once and calls `read_all()` to build the table. Native callbacks transfer batch descriptors and ownership without calling Python for each batch. C Stream is not an IPC byte stream, and this path still holds all result batches in memory @@ -119,12 +118,19 @@ need realignment, and `toPandas()` may copy during pandas conversion. The server IPC responses must still be decoded. Row materialization in `collect()` is unchanged and is not zero-copy. -The bridge preserves batch boundaries, schema metadata, and ordering flags. It -rejects batches with inconsistent schemas. Field names or timezones containing +The bridge preserves batch boundaries and schema metadata, and rejects batches +with inconsistent schemas, including metadata. It targets Spark results: maps +are unsorted and fields are not dictionary-encoded. It does not repair ordering +flags for generic Arrow inputs. Field names or timezones containing NUL use an IPC compatibility fallback over the collected batches, without another RPC; length-prefixed schema metadata can contain NUL without fallback. An empty batch vector still produces an empty table with an empty schema. +The schema is imported separately before the reader so older PyArrow versions +can reject unsupported schemas without entering their stream importer. The +reader's C callbacks and release handling are provided by arrow-rs, not custom +client callbacks. PyO3 remains at 0.28.3. + ## Two front-ends, one core === "Rust (native)" diff --git a/python/tests/test_arrow_handoff.py b/python/tests/test_arrow_handoff.py index ce92faca..695986a6 100644 --- a/python/tests/test_arrow_handoff.py +++ b/python/tests/test_arrow_handoff.py @@ -14,15 +14,17 @@ # 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. +"""Check public Arrow handoff behavior against a deterministic mock RPC service. -The localhost service returns deterministic Arrow IPC results. Buffer addresses -are captured after Rust decodes IPC, at its C Stream export boundary: addresses -in the server's Python arrays would not prove anything about this handoff. +Public cases cover Spark-output values, schemas, chunks, and lifetimes. Isolated +checks observe buffers at Rust's C Stream export boundary and verify schema +preflight. Server-side Python addresses cannot prove handoff buffer sharing. """ import ctypes import gc +import subprocess +import sys from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager from datetime import datetime, timezone @@ -166,159 +168,25 @@ class _ArrowArray(ctypes.Structure): class _ArrowArrayStream(ctypes.Structure): _fields_ = [ (name, ctypes.c_void_p) - for name in ( - "get_schema", - "get_next", - "get_last_error", - "release", - "private_data", - ) + for name in ("get_schema", "get_next", "get_last_error", "release", "private_data") ] -_ArrayRelease = ctypes.CFUNCTYPE(None, ctypes.POINTER(_ArrowArray)) _GetNext = ctypes.CFUNCTYPE( - ctypes.c_int, - ctypes.POINTER(_ArrowArrayStream), - ctypes.POINTER(_ArrowArray), + ctypes.c_int, ctypes.POINTER(_ArrowArrayStream), ctypes.POINTER(_ArrowArray) ) -_GetLastError = ctypes.CFUNCTYPE(ctypes.c_void_p, ctypes.POINTER(_ArrowArrayStream)) -_StreamRelease = ctypes.CFUNCTYPE(None, ctypes.POINTER(_ArrowArrayStream)) - -# ctypes owns the executable callback trampolines, but Arrow only retains raw -# function pointers. Keep these test callbacks alive even if a failed assertion -# leaves an imported stream or array in a pytest traceback. No Arrow buffers -# are retained by these callbacks. -_FFI_CALLBACK_KEEPALIVE = [] - - -@pytest.fixture(autouse=True) -def _forbid_per_batch_c_data_import(monkeypatch): - real_batch = pa.RecordBatch - - class ProxyMeta(type): - def __getattr__(cls, name): - return getattr(real_batch, name) - - def __instancecheck__(cls, instance): - return isinstance(instance, real_batch) - - class RecordBatchProxy(metaclass=ProxyMeta): - @staticmethod - def _import_from_c(*args): - raise AssertionError("The C Stream handoff must not import batches in Python") - - _import_from_c_capsule = _import_from_c - - monkeypatch.setattr(pa, "RecordBatch", RecordBatchProxy) def _snapshot(array): - """Copy only addresses/shape while the producer's exported array is live.""" + """Copy only addresses and shape while Rust's exported array is live.""" return { "length": array.length, "offset": array.offset, "buffers": [array.buffers[i] for i in range(array.n_buffers)], "children": [_snapshot(array.children[i].contents) for i in range(array.n_children)], - "dictionary": _snapshot(array.dictionary.contents) if array.dictionary else None, } -class _StreamSpy: - def __init__( - self, monkeypatch, *, reject=False, consume_before_reject=False, fail_next_at=None - ): - self.batches = [] - self.errors = [] - self.release_counts = [] - self.import_count = 0 - self.next_count = 0 - self.stream_releases = 0 - self.callbacks = [] - real_reader = pa.RecordBatchReader - spy = self - - class ProxyMeta(type): - def __getattr__(cls, name): - return getattr(real_reader, name) - - def __instancecheck__(cls, instance): - return isinstance(instance, real_reader) - - class ReaderProxy(metaclass=ProxyMeta): - @staticmethod - def _import_from_c(address): - stream = ctypes.cast(address, ctypes.POINTER(_ArrowArrayStream)) - original_next = _GetNext(stream.contents.get_next) - original_release = _StreamRelease(stream.contents.release) - - @_GetNext - def get_next(current, output): - spy.next_count += 1 - if spy.next_count == fail_next_at: - return 22 # EINVAL: PyArrow must propagate a failed stream read. - result = original_next(current, output) - if result == 0 and output.contents.release: - try: - spy._observe_array(output) - except Exception as error: - # Exceptions cannot cross a ctypes callback boundary. - spy.errors.append(error) - return result - - @_StreamRelease - def release(current): - spy.stream_releases += 1 - original_release(current) - - spy._keep_alive(original_next, original_release, get_next, release) - stream.contents.get_next = ctypes.cast(get_next, ctypes.c_void_p).value - stream.contents.release = ctypes.cast(release, ctypes.c_void_p).value - if fail_next_at is not None: - original_error = _GetLastError(stream.contents.get_last_error) - message = ctypes.create_string_buffer(b"test stream read failure") - - @_GetLastError - def get_last_error(current): - if spy.next_count == fail_next_at: - return ctypes.addressof(message) - return original_error(current) - - spy._keep_alive(original_error, get_last_error, message) - stream.contents.get_last_error = ctypes.cast( - get_last_error, ctypes.c_void_p - ).value - if reject and not consume_before_reject: - raise RuntimeError("test importer rejected the stream") - imported = real_reader._import_from_c(address) - spy.import_count += 1 - if reject: - # PyArrow owns the stream now. The exception traceback may - # retain this local until gc.collect(), just like user code. - raise RuntimeError("test importer rejected the consumed stream") - return imported - - monkeypatch.setattr(pa, "RecordBatchReader", ReaderProxy) - - def _keep_alive(self, *callbacks): - self.callbacks.extend(callbacks) - _FFI_CALLBACK_KEEPALIVE.extend(callbacks) - - def _observe_array(self, array): - batch_index = len(self.batches) - self.batches.append(_snapshot(array.contents)) - self.release_counts.append(0) - original_release = _ArrayRelease(array.contents.release) - - @_ArrayRelease - def release(current): - self.release_counts[batch_index] += 1 - original_release(current) - - self._keep_alive(original_release, release) - array.contents.release = ctypes.cast(release, ctypes.c_void_p).value - - def _assert_buffers(array, exported): assert len(array) == exported["length"] assert array.offset == exported["offset"] @@ -328,9 +196,7 @@ def _assert_buffers(array, exported): assert address is None elif buffer.size: assert buffer.address == address - if pa.types.is_dictionary(array.type): - _assert_buffers(array.dictionary, exported["dictionary"]) - elif pa.types.is_struct(array.type): + if pa.types.is_struct(array.type): assert array.type.num_fields == len(exported["children"]) for index, child in enumerate(exported["children"]): _assert_buffers(array.field(index), child) @@ -345,7 +211,7 @@ 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(), keys_sorted=True)), + pa.field("labels", pa.map_(pa.string(), pa.int64())), ] ) arrays = [ @@ -374,11 +240,6 @@ def _mixed_batch(): ], type=nested_type, ), - pa.DictionaryArray.from_arrays( - pa.array([0, None, 1, 0], type=pa.int8()), - pa.array(["red", "blue"]), - ordered=True, - ), pa.nulls(4), ] fields = [ @@ -389,62 +250,164 @@ def _mixed_batch(): return pa.RecordBatch.from_arrays(arrays, schema=schema) -def test_handoff_preserves_values_metadata_chunks_and_buffer_addresses(monkeypatch): +def _check_stream_buffers(): + # These proxies must be installed before arrow-pyarrow caches its classes. + # Each check has a fresh process, so the caches cannot affect another test. + batch = _mixed_batch() + schema = pa.schema( + [ + field.with_metadata({**field.metadata, b"nul\x00key": b"nul\x00value"}) + for field in batch.schema + ], + metadata={**batch.schema.metadata, b"nul\x00key": b"nul\x00value"}, + ) + batch = pa.RecordBatch.from_arrays(batch.columns, schema=schema) + batches = [batch.slice(1, 2), batch.slice(0, 1), batch.slice(3, 1)] + expected = pa.Table.from_batches(batches) + observed, errors, callbacks, imports, next_calls = [], [], [], [], [] + real_reader = pa.RecordBatchReader + + class ReaderProxy: + @staticmethod + def _import_from_c(address): + imports.append(address) + stream = ctypes.cast(address, ctypes.POINTER(_ArrowArrayStream)) + original_next = _GetNext(stream.contents.get_next) + + @_GetNext + def get_next(current, output): + next_calls.append(None) + result = original_next(current, output) + if result == 0 and output.contents.release: + try: + observed.append(_snapshot(output.contents)) + except Exception as error: + # Exceptions cannot cross a ctypes callback boundary. + errors.append(error) + return result + + # Arrow keeps raw function pointers, not the Python callback owners. + # Keep both alive until the native reader has consumed the stream. + callbacks.extend((original_next, get_next)) + stream.contents.get_next = ctypes.cast(get_next, ctypes.c_void_p).value + return real_reader._import_from_c(address) + + calls = {} + + def collect(): + with _dataframe(batches, calls=calls) as (frame, session): + table = frame.toArrow() + session.stop() + return table + + with pytest.MonkeyPatch.context() as patch: + patch.setattr(pa, "RecordBatchReader", ReaderProxy) + actual = collect() + gc.collect() + assert actual.equals(expected, check_metadata=True) + assert calls == {"execute": 1} + assert len(imports) == 1 + assert len(next_calls) == len(batches) + 1 + assert len(observed) == len(batches) + assert not errors + for batch_index, exported in enumerate(observed): + assert exported["length"] == batches[batch_index].num_rows + assert len(exported["children"]) == actual.num_columns + for column, child in zip(actual.columns, exported["children"]): + assert column.num_chunks == len(batches) + _assert_buffers(column.chunk(batch_index), child) + numeric = actual.column(0).chunk(0).slice(1, 1) + strings = actual.column(2).chunk(2).slice(0, 1) + addresses = [numeric.buffers()[1].address, strings.buffers()[2].address] + del column, actual + gc.collect() + assert numeric.to_pylist() == [-3] + assert strings.to_pylist() == ["omega"] + assert [numeric.buffers()[1].address, strings.buffers()[2].address] == addresses + + +def _check_schema_preflight(): + batch = _mixed_batch() + seen = [] + + class SchemaProxy: + @staticmethod + def _import_from_c(address): + seen.append("schema") + raise RuntimeError("test schema importer rejected the schema") + + class ReaderProxy: + @staticmethod + def _import_from_c(address): + seen.append("stream") + raise AssertionError("A rejected schema must not reach the stream importer") + + calls = {} + with pytest.MonkeyPatch.context() as patch: + patch.setattr(pa, "Schema", SchemaProxy) + patch.setattr(pa, "RecordBatchReader", ReaderProxy) + with _dataframe([batch], calls=calls) as (frame, _): + with pytest.raises(RuntimeError, match="test schema importer rejected"): + frame.toArrow() + assert seen == ["schema"] + assert calls == {"execute": 1} + + +@pytest.mark.parametrize("check", ["_check_stream_buffers", "_check_schema_preflight"]) +def test_isolated_handoff_boundaries(check): + # arrow-pyarrow uses process-global PyOnceLock caches. A subprocess keeps + # test proxies out of other tests, regardless of collection/execution order. + code = ( + "import runpy, sys\n" + f"sys.path[:] = {sys.path!r}\n" + f"namespace = runpy.run_path({__file__!r})\n" + f"namespace[{check!r}]()\n" + ) + result = subprocess.run( + [sys.executable, "-c", code], capture_output=True, text=True, timeout=60 + ) + assert result.returncode == 0, ( + f"{check} exited with {result.returncode}\n{result.stdout}\n{result.stderr}" + ) + + +def test_handoff_preserves_values_metadata_and_chunks(): batch = _mixed_batch() - # Sliced fixtures exercise nontrivial IPC input. The measured zero-copy - # boundary starts after Rust decodes it, not before IPC serialization. batches = [batch.slice(1, 2), batch.slice(0, 1), batch.slice(3, 1)] expected = pa.Table.from_batches(batches) - spy = _StreamSpy(monkeypatch) - with _dataframe(batches) as (frame, _): + calls = {} + with _dataframe(batches, calls=calls) as (frame, _): actual = frame.toArrow() assert actual.equals(expected, check_metadata=True) - assert len(spy.batches) == len(batches) - assert spy.import_count == 1 - assert spy.next_count == len(batches) + 1 - assert spy.stream_releases == 1 - assert not spy.errors - assert spy.release_counts == [0] * len(batches) - for batch_index, exported in enumerate(spy.batches): - assert exported["length"] == batches[batch_index].num_rows - assert len(exported["children"]) == actual.num_columns - for column_index, child in enumerate(exported["children"]): - column = actual.column(column_index) - assert column.num_chunks == len(batches) - _assert_buffers(column.chunk(batch_index), child) - del column, actual - gc.collect() - assert spy.release_counts == [1] * len(batches) + 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(monkeypatch): +def test_handoff_buffers_outlive_dataframe_session_and_table(): batch = _mixed_batch() expected = batch.to_pydict() - spy = _StreamSpy(monkeypatch) + calls = {} def collect(): - with _dataframe([batch]) as (frame, session): + with _dataframe([batch], calls=calls) as (frame, session): table = frame.toArrow() session.stop() return table table = collect() gc.collect() - assert spy.import_count == 1 - assert spy.stream_releases == 1 - assert not spy.errors - assert spy.release_counts == [0] assert table.to_pydict() == expected - chunk = table.column(0).chunk(0).slice(1, 3) - address = chunk.buffers()[1].address + 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 chunk.to_pylist() == [None, -3, 4] - assert chunk.buffers()[1].address == address - assert spy.release_counts == [0] - del chunk - gc.collect() - assert spy.release_counts == [1] + 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} @pytest.mark.parametrize("kind", ["no-batches", "empty-typed", "zero-columns"]) @@ -456,161 +419,111 @@ def test_handoff_empty_and_zero_column_results(kind): batches = [_mixed_batch().slice(0, 0)] expected = pa.Table.from_batches(batches) else: - structs = pa.array([{}, {}, {}], type=pa.struct([])) + structs = pa.array([{}] * 5, type=pa.struct([])) batch = pa.RecordBatch.from_struct_array(structs) - batches = [batch.replace_schema_metadata({b"zero-column": b"preserved"})] + 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) - with _dataframe(batches) as (frame, _): + assert expected.num_rows == 5 + calls = {} + with _dataframe(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(monkeypatch): +def test_to_pandas_uses_the_same_handoff(): batch = _mixed_batch() - indices = [0, 2, 4, 5, 7] + 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] - batches = [pa.RecordBatch.from_arrays(arrays, schema=schema)] + 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() - spy = _StreamSpy(monkeypatch) - with _dataframe(batches) as (frame, _): + calls = {} + with _dataframe(batches, calls=calls) as (frame, _): actual = frame.toPandas() pd.testing.assert_frame_equal(actual, expected) - assert len(spy.batches) == 1 - assert spy.import_count == 1 - assert spy.stream_releases == 1 - assert not spy.errors - del actual - gc.collect() - assert spy.release_counts == [1] - - -@pytest.mark.parametrize("consume_before_reject", [False, True]) -def test_import_rejection_releases_stream_once(monkeypatch, consume_before_reject): - batch = _mixed_batch() - spy = _StreamSpy(monkeypatch, reject=True, consume_before_reject=consume_before_reject) - with _dataframe([batch]) as (frame, _): - with pytest.raises(RuntimeError, match="test importer rejected"): - frame.toArrow() - gc.collect() - assert spy.import_count == int(consume_before_reject) - assert spy.next_count == 0 - assert spy.stream_releases == 1 - assert not spy.batches - assert not spy.release_counts - assert not spy.errors - - -@pytest.mark.parametrize("fail_next_at", [1, 2]) -def test_stream_read_failure_releases_stream_and_arrays_once(monkeypatch, fail_next_at): - batch = _mixed_batch() - spy = _StreamSpy(monkeypatch, fail_next_at=fail_next_at) - with _dataframe([batch, batch, batch]) as (frame, _): - with pytest.raises((pa.ArrowInvalid, OSError), match="test stream read failure"): - frame.toArrow() - gc.collect() - assert len(spy.batches) == fail_next_at - 1 - assert spy.import_count == 1 - assert spy.next_count == fail_next_at - assert spy.stream_releases == 1 - assert not spy.errors - # A batch already yielded to PyArrow has its own release callback. Releasing - # the failed stream also drops the Rust batches that were never exported. - assert spy.release_counts == [1] * (fail_next_at - 1) - - -def test_schema_import_rejection_prevents_stream_import(monkeypatch): - batch = _mixed_batch() - spy = _StreamSpy(monkeypatch) - - class SchemaProxy: - @staticmethod - def _import_from_c(address): - raise RuntimeError("test schema importer rejected the schema") - - monkeypatch.setattr(pa, "Schema", SchemaProxy) - with _dataframe([batch]) as (frame, _): - with pytest.raises(RuntimeError, match="test schema importer rejected"): - frame.toArrow() - assert not spy.callbacks + assert calls == {"execute": 1} -@pytest.mark.parametrize("kind", ["field", "nested-field", "timezone", "dictionary-timezone"]) -def test_nul_schema_strings_fall_back_without_executing_again(monkeypatch, kind): +@pytest.mark.parametrize("kind", ["field", "nested-field", "timezone"]) +def test_nul_schema_strings_fall_back_without_executing_again(kind): if kind == "field": - batch = pa.record_batch([pa.array([1, None])], names=["bad\x00name"]) + 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], type=data_type) - batch = pa.record_batch([values], names=["nested"]) + values = pa.array( + [{"bad\x00name": 1}, None, {"bad\x00name": None}, {"bad\x00name": 4}], + type=data_type, + ) else: - values = pa.array([0, 1], type=pa.timestamp("us", tz="UTC\x00suffix")) - if kind == "dictionary-timezone": - indices = pa.array([0, None, 1], type=pa.int8()) - values = pa.DictionaryArray.from_arrays(indices, values) - batch = pa.record_batch([values], names=["timestamp"]) - expected = pa.Table.from_batches([batch]) - spy = _StreamSpy(monkeypatch) + 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([batch], calls=calls) as (frame, _): + with _dataframe(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} - assert not spy.callbacks -def test_nul_metadata_still_uses_zero_copy_handoff(monkeypatch): +def test_nul_schema_and_field_metadata_are_preserved(): 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) - spy = _StreamSpy(monkeypatch) - with _dataframe([batch]) as (frame, _): + calls = {} + with _dataframe([batch], calls=calls) as (frame, _): actual = frame.toArrow() assert actual.equals(pa.Table.from_batches([batch]), check_metadata=True) - assert len(spy.batches) == 1 - assert spy.import_count == 1 - assert spy.stream_releases == 1 - assert not spy.errors - _assert_buffers(actual.column(0).chunk(0), spy.batches[0]["children"][0]) - del actual - gc.collect() - assert spy.release_counts == [1] + assert calls == {"execute": 1} -def test_mismatched_batch_schemas_are_rejected_before_import(monkeypatch): +@pytest.mark.parametrize("difference", ["type", "name", "schema-metadata", "field-metadata"]) +@pytest.mark.parametrize("nul_name", [False, True]) +def test_mismatched_batch_schemas_are_rejected_before_fallback(difference, nul_name): + name = "bad\x00name" if nul_name else "value" + first_field = pa.field(name, pa.int64(), metadata={b"field": b"original"}) + first_schema = pa.schema([first_field], metadata={b"schema": b"original"}) + second_values = pa.array(["text", "other"]) if difference == "type" else pa.array([3, 4]) + second_name = name + "-different" if difference == "name" else name + field_metadata = {b"field": b"different"} if difference == "field-metadata" else None + schema_metadata = {b"schema": b"different"} if difference == "schema-metadata" else None + second_field = pa.field( + second_name, second_values.type, metadata=field_metadata or first_field.metadata + ) + second_schema = pa.schema( + [second_field], metadata=schema_metadata or first_schema.metadata + ) batches = [ - pa.record_batch([pa.array([1])], names=["value"]), - pa.record_batch([pa.array(["text"])], names=["value"]), + pa.RecordBatch.from_arrays([pa.array([1, 2])], schema=first_schema), + pa.RecordBatch.from_arrays([second_values], schema=second_schema), ] - spy = _StreamSpy(monkeypatch) - with _dataframe(batches) as (frame, _): - with pytest.raises(ValueError, match="inconsistent schemas"): - frame.toArrow() - assert not spy.callbacks - - -def test_mismatched_dictionary_order_is_rejected_before_import(monkeypatch): - ordered = pa.DictionaryArray.from_arrays( - pa.array([0, 1], type=pa.int8()), pa.array(["a", "b"]), ordered=True - ) - unordered = pa.DictionaryArray.from_arrays(ordered.indices, ordered.dictionary) - batches = [pa.record_batch([array], names=["value"]) for array in (ordered, unordered)] - spy = _StreamSpy(monkeypatch) - with _dataframe(batches) as (frame, _): + calls = {} + with _dataframe(batches, calls=calls) as (frame, _): with pytest.raises(ValueError, match="inconsistent schemas"): frame.toArrow() - assert not spy.callbacks + assert calls == {"execute": 1} -def test_malformed_ipc_is_reported_before_import(monkeypatch): - spy = _StreamSpy(monkeypatch) - with _dataframe([], payloads=[(1, b"not arrow ipc")]) as (frame, _): +def test_malformed_ipc_is_reported(): + calls = {} + with _dataframe([], payloads=[(1, b"not arrow ipc")], calls=calls) as (frame, _): with pytest.raises(RuntimeError, match="Arrow"): frame.toArrow() - assert not spy.callbacks + assert calls == {"execute": 1} From d89c01b56a0a1a412c0d5f0718290ae88e09ceaf Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Wed, 28 Oct 2026 17:24:48 -0700 Subject: [PATCH 3/7] Simplify Arrow handoff tests and add live-server smoke cases Keep public behavior coverage and compile mock protocol descriptors from source. Add native toArrow smoke cases, remove the unproven schema preflight, and avoid the arrow module name collision. Leave CI workflows unchanged. Co-authored-by: Isaac --- crates/pyspark-rs/src/dataframe.rs | 2 +- crates/pyspark-rs/src/lib.rs | 2 +- .../src/{arrow.rs => pyarrow_export.rs} | 9 +- docs/architecture.md | 33 +- python/tests/test_arrow_e2e.py | 67 ++++ python/tests/test_arrow_handoff.py | 346 ++++-------------- 6 files changed, 159 insertions(+), 300 deletions(-) rename crates/pyspark-rs/src/{arrow.rs => pyarrow_export.rs} (91%) create mode 100644 python/tests/test_arrow_e2e.py diff --git a/crates/pyspark-rs/src/dataframe.rs b/crates/pyspark-rs/src/dataframe.rs index 27eb5dae..5ae82b0c 100644 --- a/crates/pyspark-rs/src/dataframe.rs +++ b/crates/pyspark-rs/src/dataframe.rs @@ -1174,7 +1174,7 @@ impl PyDataFrame { let batches = py .detach(|| self.dataframe.collect_record_batches()) .to_pyerr()?; - crate::arrow::record_batches_to_pyarrow(py, batches) + crate::pyarrow_export::record_batches_to_pyarrow(py, batches) } /// Collect the DataFrame into a pandas DataFrame. diff --git a/crates/pyspark-rs/src/lib.rs b/crates/pyspark-rs/src/lib.rs index 726b4fb5..fec1a7d0 100644 --- a/crates/pyspark-rs/src/lib.rs +++ b/crates/pyspark-rs/src/lib.rs @@ -24,7 +24,6 @@ pub(crate) fn coerce_option_value(v: &Bound<'_, PyAny>) -> PyResult arrow::error::Result> { let mut bytes = Vec::new(); @@ -70,9 +70,6 @@ pub(crate) fn record_batches_to_pyarrow( .unbind()); } - // Older PyArrow can abort on unsupported stream schemas. Use its fallible - // standalone schema importer before handing over the reader. - schema.to_pyarrow(py)?; let reader: Box = Box::new(RecordBatchIterator::new( batches.into_iter().map(Ok), schema, diff --git a/docs/architecture.md b/docs/architecture.md index ddd0cb90..0fd1506e 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -104,32 +104,11 @@ results survive transient disconnects. ### Python Arrow table handoff -`DataFrame.toArrow()` collects the decoded Arrow batches in Rust. Its `pyspark-rs` -wrapper uses arrow-rs's `IntoPyArrow` to export one boxed `RecordBatchReader` -through the Arrow C Stream Interface. PyArrow imports that reader -once and calls `read_all()` to build the table. Native callbacks transfer batch -descriptors and ownership without calling Python for each batch. C Stream is -not an IPC byte stream, and this path still holds all result batches in memory -before the handoff; it does not stream results from the network into Python. - -Compatible Arrow value and offset buffers are shared rather than serialized -again. Schema descriptors still need allocation, sliced validity bitmaps may -need realignment, and `toPandas()` may copy during pandas conversion. The server's -IPC responses must still be decoded. Row materialization in `collect()` is -unchanged and is not zero-copy. - -The bridge preserves batch boundaries and schema metadata, and rejects batches -with inconsistent schemas, including metadata. It targets Spark results: maps -are unsorted and fields are not dictionary-encoded. It does not repair ordering -flags for generic Arrow inputs. Field names or timezones containing -NUL use an IPC compatibility fallback over the collected batches, without -another RPC; length-prefixed schema metadata can contain NUL without fallback. -An empty batch vector still produces an empty table with an empty schema. - -The schema is imported separately before the reader so older PyArrow versions -can reject unsupported schemas without entering their stream importer. The -reader's C callbacks and release handling are provided by arrow-rs, not custom -client callbacks. PyO3 remains at 0.28.3. +`toArrow()` collects decoded batches and uses arrow-rs's `IntoPyArrow` for one +boxed `RecordBatchReader`. Compatible value and offset buffers are shared; +network decoding, sliced validity bitmaps, and pandas conversion may still copy. +Batch schema consistency is validated before export. Field names or timezones +containing NUL use IPC over the collected batches, without another RPC. Empty batch vectors retain the empty-schema behavior. ## Two front-ends, one core @@ -181,7 +160,7 @@ See [Rust UDFs via WebAssembly](udfs.md) for the full story. 3. **Correctness is gated in CI** - golden-proto tests guard plan building; the official Apache Spark Connect test suite guards transport and Arrow paths (see the CI notes under `dev/design/`). -4. **Version = Spark version** - the published crate/wheel version tracks the Spark release +4. **Version = Spark version** - the crate/wheel version tracks the Spark release it targets, so the number tells you which Spark it speaks. For the deeper design write-ups (design notes, CI, the official test suite, and diff --git a/python/tests/test_arrow_e2e.py b/python/tests/test_arrow_e2e.py new file mode 100644 index 00000000..dc76866d --- /dev/null +++ b/python/tests/test_arrow_e2e.py @@ -0,0 +1,67 @@ +# +# 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. +# +"""Native toArrow smoke tests against a live Spark Connect server.""" + +import os +from datetime import date, datetime, timezone + +import pyarrow as pa +import pytest +from pyspark.sql import SparkSession + + +@pytest.fixture +def session(): + remote = os.environ.get("SPARK_REMOTE") + if not remote: + pytest.skip("SPARK_REMOTE is required for live-server tests") + session = SparkSession.builder.remote(remote).getOrCreate() + try: + yield session + finally: + session.stop() + + +def test_to_arrow_nested_null_binary_and_temporal_values(session): + table = session.sql(""" + SELECT array(1, CAST(NULL AS INT), 3) AS items, + named_struct('name', 'spark', 'value', CAST(NULL AS BIGINT)) AS nested, + X'00FF' AS binary_value, DATE '2020-01-02' AS day, + TIMESTAMP '2020-01-02 03:04:05+00:00' AS instant + """).toArrow() + assert isinstance(table, pa.Table) + values = table.to_pylist() + assert len(values) == 1 + assert values[0] == { + "items": [1, None, 3], + "nested": {"name": "spark", "value": None}, + "binary_value": b"\x00\xff", + "day": date(2020, 1, 2), + "instant": datetime(2020, 1, 2, 3, 4, 5, tzinfo=timezone.utc), + } + + +def test_to_arrow_numeric_range(session): + table = session.range(10000).toArrow() + assert table.num_rows == 10000 + assert table.column("id").to_pylist() == list(range(10000)) + + +def test_to_arrow_empty_result(session): + table = session.range(0).toArrow() + assert isinstance(table, pa.Table) + assert table.num_rows == 0 diff --git a/python/tests/test_arrow_handoff.py b/python/tests/test_arrow_handoff.py index 695986a6..5989025a 100644 --- a/python/tests/test_arrow_handoff.py +++ b/python/tests/test_arrow_handoff.py @@ -16,80 +16,59 @@ # """Check public Arrow handoff behavior against a deterministic mock RPC service. -Public cases cover Spark-output values, schemas, chunks, and lifetimes. Isolated -checks observe buffers at Rust's C Stream export boundary and verify schema -preflight. Server-side Python addresses cannot prove handoff buffer sharing. +Cases cover Spark-output values, schemas, chunks, and lifetimes through public +APIs. Server-side Python addresses cannot prove Rust-to-Python buffer sharing. """ -import ctypes import gc +import os import subprocess -import sys 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 -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 +@pytest.fixture(scope="module") +def protocol(tmp_path_factory): + """Compile the repository protocol instead of hard-coding protobuf field numbers.""" + proto_root = Path(__file__).resolve().parents[2] / "crates/spark-connect-proto/proto" + descriptor = tmp_path_factory.mktemp("connect-proto") / "connect.pb" + command = [os.environ.get("PROTOC", "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, *, payloads=None, calls=None): - """Serve just ExecutePlan and release RPCs; unexpected RPCs fail closed. - Field numbers are from spark/connect/base.proto. Using the wire envelope - avoids depending on official pyspark's generated protobuf package, which - cannot be imported alongside the Rust-backed drop-in under the same name. - """ +@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: @@ -97,26 +76,44 @@ def _dataframe(batches, *, payloads=None, calls=None): 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 = b"00000000-0000-0000-0000-000000000001" + server_id = "00000000-0000-0000-0000-000000000001" def execute(request, context): if calls is not None: calls["execute"] = calls.get("execute", 0) + 1 - fields = _request_fields(request) - header = ( - _bytes_field(1, fields[1]) + _bytes_field(12, fields[6]) + _bytes_field(15, server_id) + 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): - arrow_batch = _varint(8) + _varint(rows) + _bytes_field(2, payload) - yield (header + _bytes_field(13, str(index).encode()) + _bytes_field(2, arrow_batch)) - yield header + _bytes_field(13, b"complete") + _bytes_field(14, b"") + 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): - fields = _request_fields(request) - return _bytes_field(1, fields[1]) + _bytes_field(2, fields[3]) + _bytes_field(3, server_id) + 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): - return _bytes_field(1, _request_fields(request)[1]) + _bytes_field(2, server_id) + 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) @@ -147,66 +144,6 @@ def release_session(request, context): server.stop(0).wait(timeout=5) -class _ArrowArray(ctypes.Structure): - pass - - -_ArrowArray._fields_ = [ - ("length", ctypes.c_int64), - ("null_count", ctypes.c_int64), - ("offset", ctypes.c_int64), - ("n_buffers", ctypes.c_int64), - ("n_children", ctypes.c_int64), - ("buffers", ctypes.POINTER(ctypes.c_void_p)), - ("children", ctypes.POINTER(ctypes.POINTER(_ArrowArray))), - ("dictionary", ctypes.POINTER(_ArrowArray)), - ("release", ctypes.c_void_p), - ("private_data", ctypes.c_void_p), -] - - -class _ArrowArrayStream(ctypes.Structure): - _fields_ = [ - (name, ctypes.c_void_p) - for name in ("get_schema", "get_next", "get_last_error", "release", "private_data") - ] - - -_GetNext = ctypes.CFUNCTYPE( - ctypes.c_int, ctypes.POINTER(_ArrowArrayStream), ctypes.POINTER(_ArrowArray) -) - - -def _snapshot(array): - """Copy only addresses and shape while Rust's exported array is live.""" - return { - "length": array.length, - "offset": array.offset, - "buffers": [array.buffers[i] for i in range(array.n_buffers)], - "children": [_snapshot(array.children[i].contents) for i in range(array.n_children)], - } - - -def _assert_buffers(array, exported): - assert len(array) == exported["length"] - assert array.offset == exported["offset"] - own_buffers = array.buffers()[: len(exported["buffers"])] - for buffer, address in zip(own_buffers, exported["buffers"]): - if buffer is None: - assert address is None - elif buffer.size: - assert buffer.address == address - if pa.types.is_struct(array.type): - assert array.type.num_fields == len(exported["children"]) - for index, child in enumerate(exported["children"]): - _assert_buffers(array.field(index), child) - elif pa.types.is_list(array.type) or pa.types.is_map(array.type): - assert len(exported["children"]) == 1 - _assert_buffers(array.values, exported["children"][0]) - else: - assert not exported["children"] - - def _mixed_batch(): nested_type = pa.struct( [ @@ -250,133 +187,12 @@ def _mixed_batch(): return pa.RecordBatch.from_arrays(arrays, schema=schema) -def _check_stream_buffers(): - # These proxies must be installed before arrow-pyarrow caches its classes. - # Each check has a fresh process, so the caches cannot affect another test. +def test_handoff_preserves_values_metadata_and_chunks(protocol): batch = _mixed_batch() - schema = pa.schema( - [ - field.with_metadata({**field.metadata, b"nul\x00key": b"nul\x00value"}) - for field in batch.schema - ], - metadata={**batch.schema.metadata, b"nul\x00key": b"nul\x00value"}, - ) - batch = pa.RecordBatch.from_arrays(batch.columns, schema=schema) batches = [batch.slice(1, 2), batch.slice(0, 1), batch.slice(3, 1)] expected = pa.Table.from_batches(batches) - observed, errors, callbacks, imports, next_calls = [], [], [], [], [] - real_reader = pa.RecordBatchReader - - class ReaderProxy: - @staticmethod - def _import_from_c(address): - imports.append(address) - stream = ctypes.cast(address, ctypes.POINTER(_ArrowArrayStream)) - original_next = _GetNext(stream.contents.get_next) - - @_GetNext - def get_next(current, output): - next_calls.append(None) - result = original_next(current, output) - if result == 0 and output.contents.release: - try: - observed.append(_snapshot(output.contents)) - except Exception as error: - # Exceptions cannot cross a ctypes callback boundary. - errors.append(error) - return result - - # Arrow keeps raw function pointers, not the Python callback owners. - # Keep both alive until the native reader has consumed the stream. - callbacks.extend((original_next, get_next)) - stream.contents.get_next = ctypes.cast(get_next, ctypes.c_void_p).value - return real_reader._import_from_c(address) - calls = {} - - def collect(): - with _dataframe(batches, calls=calls) as (frame, session): - table = frame.toArrow() - session.stop() - return table - - with pytest.MonkeyPatch.context() as patch: - patch.setattr(pa, "RecordBatchReader", ReaderProxy) - actual = collect() - gc.collect() - assert actual.equals(expected, check_metadata=True) - assert calls == {"execute": 1} - assert len(imports) == 1 - assert len(next_calls) == len(batches) + 1 - assert len(observed) == len(batches) - assert not errors - for batch_index, exported in enumerate(observed): - assert exported["length"] == batches[batch_index].num_rows - assert len(exported["children"]) == actual.num_columns - for column, child in zip(actual.columns, exported["children"]): - assert column.num_chunks == len(batches) - _assert_buffers(column.chunk(batch_index), child) - numeric = actual.column(0).chunk(0).slice(1, 1) - strings = actual.column(2).chunk(2).slice(0, 1) - addresses = [numeric.buffers()[1].address, strings.buffers()[2].address] - del column, actual - gc.collect() - assert numeric.to_pylist() == [-3] - assert strings.to_pylist() == ["omega"] - assert [numeric.buffers()[1].address, strings.buffers()[2].address] == addresses - - -def _check_schema_preflight(): - batch = _mixed_batch() - seen = [] - - class SchemaProxy: - @staticmethod - def _import_from_c(address): - seen.append("schema") - raise RuntimeError("test schema importer rejected the schema") - - class ReaderProxy: - @staticmethod - def _import_from_c(address): - seen.append("stream") - raise AssertionError("A rejected schema must not reach the stream importer") - - calls = {} - with pytest.MonkeyPatch.context() as patch: - patch.setattr(pa, "Schema", SchemaProxy) - patch.setattr(pa, "RecordBatchReader", ReaderProxy) - with _dataframe([batch], calls=calls) as (frame, _): - with pytest.raises(RuntimeError, match="test schema importer rejected"): - frame.toArrow() - assert seen == ["schema"] - assert calls == {"execute": 1} - - -@pytest.mark.parametrize("check", ["_check_stream_buffers", "_check_schema_preflight"]) -def test_isolated_handoff_boundaries(check): - # arrow-pyarrow uses process-global PyOnceLock caches. A subprocess keeps - # test proxies out of other tests, regardless of collection/execution order. - code = ( - "import runpy, sys\n" - f"sys.path[:] = {sys.path!r}\n" - f"namespace = runpy.run_path({__file__!r})\n" - f"namespace[{check!r}]()\n" - ) - result = subprocess.run( - [sys.executable, "-c", code], capture_output=True, text=True, timeout=60 - ) - assert result.returncode == 0, ( - f"{check} exited with {result.returncode}\n{result.stdout}\n{result.stderr}" - ) - - -def test_handoff_preserves_values_metadata_and_chunks(): - 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(batches, calls=calls) as (frame, _): + 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: @@ -384,13 +200,13 @@ def test_handoff_preserves_values_metadata_and_chunks(): assert calls == {"execute": 1} -def test_handoff_buffers_outlive_dataframe_session_and_table(): +def test_handoff_buffers_outlive_dataframe_session_and_table(protocol): batch = _mixed_batch() expected = batch.to_pydict() calls = {} def collect(): - with _dataframe([batch], calls=calls) as (frame, session): + with _dataframe(protocol=protocol, batches=[batch], calls=calls) as (frame, session): table = frame.toArrow() session.stop() return table @@ -411,7 +227,7 @@ def collect(): @pytest.mark.parametrize("kind", ["no-batches", "empty-typed", "zero-columns"]) -def test_handoff_empty_and_zero_column_results(kind): +def test_handoff_empty_and_zero_column_results(protocol, kind): if kind == "no-batches": batches = [] expected = pa.table({}) @@ -426,7 +242,7 @@ def test_handoff_empty_and_zero_column_results(kind): expected = pa.Table.from_batches(batches) assert expected.num_rows == 5 calls = {} - with _dataframe(batches, calls=calls) as (frame, _): + 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 @@ -434,7 +250,7 @@ def test_handoff_empty_and_zero_column_results(kind): assert calls == {"execute": 1} -def test_to_pandas_uses_the_same_handoff(): +def test_to_pandas_uses_the_same_handoff(protocol): batch = _mixed_batch() indices = [0, 2, 4, 5] schema = pa.schema( @@ -445,14 +261,14 @@ def test_to_pandas_uses_the_same_handoff(): batches = [selected.slice(0, 2), selected.slice(2, 2)] expected = pa.Table.from_batches(batches).to_pandas() calls = {} - with _dataframe(batches, calls=calls) as (frame, _): + with _dataframe(protocol=protocol, batches=batches, calls=calls) as (frame, _): actual = frame.toPandas() pd.testing.assert_frame_equal(actual, expected) assert calls == {"execute": 1} @pytest.mark.parametrize("kind", ["field", "nested-field", "timezone"]) -def test_nul_schema_strings_fall_back_without_executing_again(kind): +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]) @@ -474,7 +290,7 @@ def test_nul_schema_strings_fall_back_without_executing_again(kind): batches = [batch.slice(0, 2), batch.slice(2, 2)] expected = pa.Table.from_batches(batches) calls = {} - with _dataframe(batches, calls=calls) as (frame, _): + 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. @@ -483,12 +299,12 @@ def test_nul_schema_strings_fall_back_without_executing_again(kind): assert calls == {"execute": 1} -def test_nul_schema_and_field_metadata_are_preserved(): +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([batch], calls=calls) as (frame, _): + 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} @@ -496,7 +312,7 @@ def test_nul_schema_and_field_metadata_are_preserved(): @pytest.mark.parametrize("difference", ["type", "name", "schema-metadata", "field-metadata"]) @pytest.mark.parametrize("nul_name", [False, True]) -def test_mismatched_batch_schemas_are_rejected_before_fallback(difference, nul_name): +def test_mismatched_batch_schemas_are_rejected_before_fallback(protocol, difference, nul_name): name = "bad\x00name" if nul_name else "value" first_field = pa.field(name, pa.int64(), metadata={b"field": b"original"}) first_schema = pa.schema([first_field], metadata={b"schema": b"original"}) @@ -507,23 +323,23 @@ def test_mismatched_batch_schemas_are_rejected_before_fallback(difference, nul_n second_field = pa.field( second_name, second_values.type, metadata=field_metadata or first_field.metadata ) - second_schema = pa.schema( - [second_field], metadata=schema_metadata or first_schema.metadata - ) + second_schema = pa.schema([second_field], metadata=schema_metadata or first_schema.metadata) batches = [ pa.RecordBatch.from_arrays([pa.array([1, 2])], schema=first_schema), pa.RecordBatch.from_arrays([second_values], schema=second_schema), ] calls = {} - with _dataframe(batches, calls=calls) as (frame, _): + with _dataframe(protocol=protocol, batches=batches, calls=calls) as (frame, _): with pytest.raises(ValueError, match="inconsistent schemas"): frame.toArrow() assert calls == {"execute": 1} -def test_malformed_ipc_is_reported(): +def test_malformed_ipc_is_reported(protocol): calls = {} - with _dataframe([], payloads=[(1, b"not arrow ipc")], calls=calls) as (frame, _): + with _dataframe( + protocol=protocol, batches=[], payloads=[(1, b"not arrow ipc")], calls=calls + ) as (frame, _): with pytest.raises(RuntimeError, match="Arrow"): frame.toArrow() assert calls == {"execute": 1} From fe79d54db610cc181f0b8c5ddb1a0cf6258546a5 Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Wed, 28 Oct 2026 17:47:28 -0700 Subject: [PATCH 4/7] Skip Arrow handoff tests when protoc is unavailable Document the Arrow 58.3 NUL export panic at its upstream source and wrap the architecture paragraph. Keep CI workflows unchanged. Co-authored-by: Isaac --- crates/pyspark-rs/src/pyarrow_export.rs | 3 +++ docs/architecture.md | 3 ++- python/tests/test_arrow_handoff.py | 6 +++++- 3 files changed, 10 insertions(+), 2 deletions(-) diff --git a/crates/pyspark-rs/src/pyarrow_export.rs b/crates/pyspark-rs/src/pyarrow_export.rs index f439d07a..b94c31ff 100644 --- a/crates/pyspark-rs/src/pyarrow_export.rs +++ b/crates/pyspark-rs/src/pyarrow_export.rs @@ -50,6 +50,9 @@ pub(crate) fn record_batches_to_pyarrow( if schema.fields().iter().any(|field| field_needs_ipc(field)) { // Arrow 58.3 panics when exporting NUL-containing C strings. Reuse // the collected batches for IPC without executing the request again. + // CString::new(...).unwrap() in FFI_ArrowSchema::try_new / with_name: + // https://github.com/apache/arrow-rs/blob/58.3.0/arrow-schema/src/ffi.rs#L148 + // https://github.com/apache/arrow-rs/blob/58.3.0/arrow-schema/src/ffi.rs#L174 let bytes = py .detach(|| -> arrow::error::Result> { let mut bytes = Vec::new(); diff --git a/docs/architecture.md b/docs/architecture.md index 0fd1506e..17559ed9 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -108,7 +108,8 @@ results survive transient disconnects. boxed `RecordBatchReader`. Compatible value and offset buffers are shared; network decoding, sliced validity bitmaps, and pandas conversion may still copy. Batch schema consistency is validated before export. Field names or timezones -containing NUL use IPC over the collected batches, without another RPC. Empty batch vectors retain the empty-schema behavior. +containing NUL use IPC over the collected batches, without another RPC. Empty +batch vectors retain the empty-schema behavior. ## Two front-ends, one core diff --git a/python/tests/test_arrow_handoff.py b/python/tests/test_arrow_handoff.py index 5989025a..5deef93d 100644 --- a/python/tests/test_arrow_handoff.py +++ b/python/tests/test_arrow_handoff.py @@ -22,6 +22,7 @@ import gc import os +import shutil import subprocess from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager @@ -40,9 +41,12 @@ @pytest.fixture(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 = [os.environ.get("PROTOC", "protoc"), f"-I{proto_root}"] + command = [protoc, f"-I{proto_root}"] if include := os.environ.get("PROTOC_INCLUDE"): command.append(f"-I{include}") subprocess.run( From c4bbfd344d603d067804432338a216cb4e021e22 Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Wed, 28 Oct 2026 20:31:56 -0700 Subject: [PATCH 5/7] [SPARK-60675][CONNECT] Support PyArrow Table input through the C Stream interface Normalize Arrow schemas and timestamps with the existing PySpark helpers, then serialize imported Rust batches directly for local relations without Python row materialization. Co-authored-by: Isaac --- crates/pyspark-rs/src/lib.rs | 1 + crates/pyspark-rs/src/pyarrow_import.rs | 87 +++++++++++++++++++ crates/pyspark-rs/src/session.rs | 6 ++ crates/spark-connect/src/session.rs | 85 +++++++++++++++++++ docs/architecture.md | 7 ++ python/pyspark/sql/_arrow_input.py | 66 +++++++++++++++ python/tests/test_dropin_offline.py | 108 ++++++++++++++++++++++++ 7 files changed, 360 insertions(+) create mode 100644 crates/pyspark-rs/src/pyarrow_import.rs create mode 100644 python/pyspark/sql/_arrow_input.py diff --git a/crates/pyspark-rs/src/lib.rs b/crates/pyspark-rs/src/lib.rs index fec1a7d0..68507bd5 100644 --- a/crates/pyspark-rs/src/lib.rs +++ b/crates/pyspark-rs/src/lib.rs @@ -38,6 +38,7 @@ mod observation; mod pipelines; mod profiler; mod pyarrow_export; +mod pyarrow_import; mod readwriter; mod resource; mod row; diff --git a/crates/pyspark-rs/src/pyarrow_import.rs b/crates/pyspark-rs/src/pyarrow_import.rs new file mode 100644 index 00000000..f1b083be --- /dev/null +++ b/crates/pyspark-rs/src/pyarrow_import.rs @@ -0,0 +1,87 @@ +// 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. + +//! PyArrow Table input through arrow-rs's owned C Stream importer. + +use arrow::pyarrow::{FromPyArrow, Table}; +use pyo3::prelude::*; +use spark_connect::session::SparkSession; + +use crate::dataframe::PyDataFrame; +use crate::errors::ResultExt; + +pub(crate) fn create_dataframe( + py: Python<'_>, + session: &SparkSession, + table: &Bound<'_, PyAny>, + schema: Option<&Bound<'_, PyAny>>, +) -> PyResult { + let helper = py.import("pyspark.sql._arrow_input")?; + // Resolve DDL with the same converter as the existing local-data path. + let schema = match schema { + Some(s) if s.extract::().is_ok() => Some(crate::types::data_type_to_py( + py, + &crate::types::py_to_data_type(s)?, + )?), + Some(s) => Some(s.clone().unbind()), + None => None, + }; + let schema = schema.unwrap_or_else(|| py.None()); + // Empty input needs only the explicit schema, and no server configuration. + let empty: bool = table.call_method0("__len__")?.extract::()? == 0; + let (timestamp_type, timezone, large_types) = if empty { + (None, None, None) + } else { + py.detach(|| -> PyResult<_> { + let conf = session.conf(); + Ok(( + conf.get("spark.sql.timestampType").to_pyerr()?, + conf.get("spark.sql.session.timeZone").to_pyerr()?, + conf.get("spark.sql.execution.arrow.useLargeVarTypes") + .to_pyerr()?, + )) + })? + }; + let result = helper.getattr("prepare_table")?.call1(( + table, + schema, + timestamp_type, + timezone, + large_types, + ))?; + let normalized = result.get_item(0)?; + let spark_schema = crate::types::py_to_data_type(&result.get_item(1)?)?; + let names: Option> = result.get_item(2)?.extract()?; + let batches = if empty { + Vec::new() + } else { + // PyArrow 10 readers support _export_to_c, even without PyCapsules. + Table::from_pyarrow_bound(&normalized.call_method0("to_reader")?)? + .into_inner() + .0 + }; + let dataframe = py.detach(|| { + session + .create_dataframe_from_arrow(batches, spark_schema) + .to_pyerr() + })?; + let dataframe = match names { + Some(names) => dataframe.to_df(names.iter().map(String::as_str).collect()), + None => dataframe, + }; + Ok(PyDataFrame::new(dataframe)) +} diff --git a/crates/pyspark-rs/src/session.rs b/crates/pyspark-rs/src/session.rs index 5f6302e3..0e271dbd 100644 --- a/crates/pyspark-rs/src/session.rs +++ b/crates/pyspark-rs/src/session.rs @@ -348,6 +348,12 @@ impl PySparkSession { // signature parity. See PARITY_GAP_ANALYSIS.md. let _ = (samplingRatio, verifySchema); + if let Ok(pa) = py.import("pyarrow") { + if data.is_instance(&pa.getattr("Table")?)? { + return crate::pyarrow_import::create_dataframe(py, &self.session, data, schema); + } + } + // Accept a pandas DataFrame (as pyspark does): convert it to a list of row lists // with NaN/NaT mapped to None, and take its column names as the default schema. // pandas-on-Spark always passes an explicit `schema`, so this mainly needs to diff --git a/crates/spark-connect/src/session.rs b/crates/spark-connect/src/session.rs index 7eb29ace..31843d53 100644 --- a/crates/spark-connect/src/session.rs +++ b/crates/spark-connect/src/session.rs @@ -290,6 +290,44 @@ impl SparkSession { Ok(DataFrame::new(self.clone(), plan)) } + /// Create a local DataFrame from Arrow batches without materializing rows. + /// + /// `schema` is the Spark schema applied by the server. All batches must share + /// the same Arrow schema; IPC serialization is still required for transport. + pub fn create_dataframe_from_arrow( + &self, + batches: Vec, + schema: DataType, + ) -> Result { + let data = if let Some(first) = batches.first() { + let arrow_schema = first.schema(); + let mut buffer = Vec::new(); + { + let mut writer = + arrow::ipc::writer::StreamWriter::try_new(&mut buffer, &arrow_schema) + .map_err(|e| SparkError::connect_msg(e.to_string()))?; + for batch in &batches { + if batch.schema() != arrow_schema { + return Err(SparkError::connect_msg("Arrow batch schemas must match")); + } + writer + .write(batch) + .map_err(|e| SparkError::connect_msg(e.to_string()))?; + } + writer + .finish() + .map_err(|e| SparkError::connect_msg(e.to_string()))?; + } + Some(buffer) + } else { + None + }; + Ok(DataFrame::new( + self.clone(), + LogicalPlan::LocalRelation { schema, data }, + )) + } + /// Create a DataFrameReader for reading data from various sources. /// /// Mirrors `pyspark.sql.SparkSession.read`. @@ -1224,6 +1262,53 @@ mod tests { .expect("failed to build session") } + #[test] + fn arrow_local_relation_preserves_batches_and_rejects_schema_mismatch() { + use arrow::array::Array; + use arrow::datatypes::{DataType as ArrowType, Field, Schema}; + let arrow_schema = Arc::new(Schema::new(vec![Field::new("x", ArrowType::Int32, true)])); + let batch = RecordBatch::try_new( + arrow_schema, + vec![Arc::new(Int32Array::from(vec![Some(1), None]))], + ) + .unwrap(); + let schema = DataType::Struct { + fields: vec![struct_field("x", DataType::Integer)], + }; + let spark = session(); + let df = spark + .create_dataframe_from_arrow(vec![batch.clone(), batch.clone()], schema.clone()) + .unwrap(); + let LogicalPlan::LocalRelation { + data: Some(data), .. + } = df.plan() + else { + panic!("expected Arrow local relation") + }; + let batches = StreamReader::try_new(Cursor::new(data), None) + .unwrap() + .collect::, _>>() + .unwrap(); + assert_eq!(batches.len(), 2); + assert_eq!(batches[0], batch); + assert!(batches[1].column(0).is_null(1)); + let other = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new("y", ArrowType::Int32, true)])), + batch.columns().to_vec(), + ) + .unwrap(); + assert!(spark + .create_dataframe_from_arrow(vec![batch, other], schema.clone()) + .is_err()); + assert!(matches!( + spark + .create_dataframe_from_arrow(vec![], schema) + .unwrap() + .plan(), + LogicalPlan::LocalRelation { data: None, .. } + )); + } + #[test] fn session_is_not_stopped_initially() { let spark = session(); diff --git a/docs/architecture.md b/docs/architecture.md index 17559ed9..1fa9261c 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -111,6 +111,13 @@ Batch schema consistency is validated before export. Field names or timezones containing NUL use IPC over the collected batches, without another RPC. Empty batch vectors retain the empty-schema behavior. +`createDataFrame(pyarrow_table)` uses the reverse boundary: Spark's schema and +timestamp conversion helpers normalize the table, then arrow-rs imports its +`RecordBatchReader` through the C Stream interface. The Rust core serializes +those batches to IPC for the local relation sent to Spark. No Python rows are +materialized; casts, buffer alignment, and IPC serialization may still copy. +Empty tables require an explicit schema, matching PySpark Connect. + ## Two front-ends, one core === "Rust (native)" diff --git a/python/pyspark/sql/_arrow_input.py b/python/pyspark/sql/_arrow_input.py new file mode 100644 index 00000000..f77540c8 --- /dev/null +++ b/python/pyspark/sql/_arrow_input.py @@ -0,0 +1,66 @@ +# +# 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. +# + + +"""Schema normalization for the native Arrow Table input path.""" + +from pyspark.errors import PySparkValueError +from pyspark.sql.pandas.types import ( + _check_arrow_table_timestamps_localize, + from_arrow_schema, + to_arrow_schema, +) +from pyspark.sql.types import StructType + + +def prepare_table(table, schema, timestamp_type, timezone, large_types): + """Match Connect's Arrow casts without converting any column to Python rows.""" + names = list(schema) if isinstance(schema, (list, tuple)) else None + explicit = schema is not None and names is None + if len(table) == 0: + if not explicit: + raise PySparkValueError(errorClass="CANNOT_INFER_EMPTY_SCHEMA", messageParameters={}) + return table, schema, None + + if names is not None: + names.extend(f"_{i + 1}" for i in range(len(names), table.num_columns)) + inferred = ( + schema + if isinstance(schema, StructType) + else from_arrow_schema(table.schema, prefer_timestamp_ntz=timestamp_type == "TIMESTAMP_NTZ") + ) + normalized = ( + _check_arrow_table_timestamps_localize(table, inferred, True, timezone) + .cast( + to_arrow_schema( + inferred, + error_on_duplicated_field_names_in_struct=True, + timezone="UTC", + prefers_large_types=large_types == "true", + ) + ) + .rename_columns(inferred.names) + ) + if names is not None and len(names) != normalized.num_columns: + raise PySparkValueError( + errorClass="AXIS_LENGTH_MISMATCH", + messageParameters={ + "expected_length": str(len(names)), + "actual_length": str(normalized.num_columns), + }, + ) + return normalized, schema if explicit else inferred, names diff --git a/python/tests/test_dropin_offline.py b/python/tests/test_dropin_offline.py index 67ccaa4d..1aab0918 100644 --- a/python/tests/test_dropin_offline.py +++ b/python/tests/test_dropin_offline.py @@ -2755,3 +2755,111 @@ def test_create_dataframe_width_mismatch_is_a_normal_exception(): else: raise AssertionError("expected a width validation error") """], check=True) + + +@pytest.mark.parametrize("names,expected", [(None, None), (["a"], ["a", "_2"]), ([], ["_1", "_2"])]) +def test_arrow_input_normalization(names, expected): + pa = pytest.importorskip("pyarrow") + from pyspark.sql._arrow_input import prepare_table + + table = pa.table({"x": pa.chunked_array([[1], [2]]), "y": ["a", None]}) + normalized, schema, columns = prepare_table(table, names, "TIMESTAMP_LTZ", "UTC", "false") + assert normalized.equals(table) + assert schema.names == ["x", "y"] + assert columns == expected + + +def test_arrow_input_empty_and_width_errors(): + pa = pytest.importorskip("pyarrow") + from pyspark.errors import PySparkValueError + from pyspark.sql._arrow_input import prepare_table + + empty = pa.table({"x": pa.array([], type=pa.int64())}) + with pytest.raises(PySparkValueError, match="CANNOT_INFER_EMPTY_SCHEMA"): + prepare_table(empty, None, None, None, None) + schema = T.StructType([T.StructField("x", T.LongType())]) + assert prepare_table(empty, schema, None, None, None)[1] is schema + with pytest.raises(PySparkValueError, match="AXIS_LENGTH_MISMATCH"): + prepare_table(pa.table({"x": [1]}), ["a", "b"], "TIMESTAMP_LTZ", "UTC", "false") + + +@pytest.mark.parametrize("timestamp_type,expected", [("TIMESTAMP_NTZ", T.TimestampNTZType), ("TIMESTAMP_LTZ", T.TimestampType)]) +def test_arrow_input_timestamps(timestamp_type, expected): + pa = pytest.importorskip("pyarrow") + from pyspark.sql._arrow_input import prepare_table + + table = pa.table({"t": pa.array([1001], type=pa.timestamp("ns"))}) + normalized, schema, _ = prepare_table(table, None, timestamp_type, "UTC", "false") + assert isinstance(schema.fields[0].dataType, expected) + assert normalized.column(0).type.unit == "us" + assert normalized.column(0).cast(pa.int64()).to_pylist() == [1] + + +def test_arrow_input_explicit_cast_and_nested_values(): + pa = pytest.importorskip("pyarrow") + from pyspark.sql._arrow_input import prepare_table + + table = pa.table({"x": [1, 2], "nested": [[1, None], []]}) + schema = T.StructType([T.StructField("x", T.IntegerType()), + T.StructField("nested", T.ArrayType(T.LongType()))]) + normalized, result, _ = prepare_table(table, schema, "TIMESTAMP_LTZ", "UTC", "false") + assert result is schema + assert normalized.column(0).type == pa.int32() + assert normalized.to_pylist() == table.to_pylist() + with pytest.raises(pa.ArrowInvalid): + prepare_table(pa.table({"x": [2**40]}), T.StructType([T.StructField("x", T.IntegerType())]), + "TIMESTAMP_LTZ", "UTC", "false") + + +def test_arrow_input_roundtrip_live_server(): + """Run explicitly with SPARK_REMOTE; also checks imported buffers own their lifetime.""" + import gc + import os + pa = pytest.importorskip("pyarrow") + remote = os.environ.get("SPARK_REMOTE") + if not remote: + pytest.skip("requires SPARK_REMOTE") + from pyspark.sql import SparkSession + + spark = SparkSession.builder.remote(remote).getOrCreate() + try: + table = pa.table({"x": pa.chunked_array([[1, 2], [3]]), "nested": [[1], None, []]}) + df = spark.createDataFrame(table, ["id", "items"]) + del table + gc.collect() + assert df.columns == ["id", "items"] + assert [(r[0], r[1]) for r in df.collect()] == [(1, [1]), (2, None), (3, [])] + empty = spark.createDataFrame(pa.table({"x": pa.array([], type=pa.int64())}), + T.StructType([T.StructField("x", T.LongType())])) + assert empty.count() == 0 + assert empty.columns == ["x"] + finally: + spark.stop() + + +@pytest.mark.parametrize("kind", ["slice", "dictionary", "decimal", "null", "struct", "nul"]) +def test_arrow_input_types_live_server(kind): + import os + from decimal import Decimal + pa = pytest.importorskip("pyarrow") + remote = os.environ.get("SPARK_REMOTE") + if not remote: + pytest.skip("requires SPARK_REMOTE") + from pyspark.sql import SparkSession + + tables = { + "slice": pa.table({"x": pa.array([0, 1, 2, 3]).slice(1, 2)}), + "dictionary": pa.table({"x": pa.array(["a", None, "b", "a"]).dictionary_encode()}), + "decimal": pa.table({"x": pa.array([Decimal("1.23"), None], type=pa.decimal128(10, 2))}), + "null": pa.table({"x": pa.nulls(2)}), + "struct": pa.table({"x": [{"a": [1, None]}, None]}), + "nul": pa.table({"a\0b": [1, 2]}), + } + spark = SparkSession.builder.remote(remote).getOrCreate() + try: + table = tables[kind] + df = spark.createDataFrame(table) + assert df.columns == table.column_names + assert df.toArrow().to_pylist() == table.to_pylist() + finally: + spark.stop() From ce89b8dd496bbb1529f1b5531b1eaaa6d739bef0 Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Wed, 28 Oct 2026 21:02:38 -0700 Subject: [PATCH 6/7] [SPARK-60675][CONNECT] Upload large Arrow inputs as cached relations Honor server local-relation settings, split Arrow IPC into independently readable chunks, and upload session-scoped cache artifacts with SHA-256 deduplication and CRC validation. Keep large input payloads out of query plans and test uploads above the gRPC message limit. Co-authored-by: Isaac --- Cargo.lock | 1 + crates/spark-connect-core/Cargo.toml | 1 + crates/spark-connect-core/src/client.rs | 67 +++++++ crates/spark-connect/src/arrow_input.rs | 235 ++++++++++++++++++++++++ crates/spark-connect/src/dataframe.rs | 5 +- crates/spark-connect/src/lib.rs | 1 + crates/spark-connect/src/plan.rs | 16 ++ crates/spark-connect/src/session.rs | 39 +--- docs/architecture.md | 6 + python/tests/test_arrow_handoff.py | 210 +++++++++++++++++++++ python/tests/test_dropin_offline.py | 21 +++ 11 files changed, 570 insertions(+), 32 deletions(-) create mode 100644 crates/spark-connect/src/arrow_input.rs diff --git a/Cargo.lock b/Cargo.lock index 29e7d3fe..6238cb96 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -87,6 +87,7 @@ dependencies = [ "prost-types", "regex", "serde_json", + "sha2", "tempfile", "thiserror", "tokio", diff --git a/crates/spark-connect-core/Cargo.toml b/crates/spark-connect-core/Cargo.toml index c6a0c088..6cdf0e01 100644 --- a/crates/spark-connect-core/Cargo.toml +++ b/crates/spark-connect-core/Cargo.toml @@ -26,6 +26,7 @@ serde_json = "1" regex = "1" prost-types.workspace = true crc32fast = "1" +sha2 = "0.10" [dev-dependencies] tempfile = "3" diff --git a/crates/spark-connect-core/src/client.rs b/crates/spark-connect-core/src/client.rs index 31820ed5..8f658bdc 100644 --- a/crates/spark-connect-core/src/client.rs +++ b/crates/spark-connect-core/src/client.rs @@ -800,6 +800,73 @@ impl SparkConnectClient { Ok(()) } + /// Upload session-scoped cache blobs, returning SHA-256 hashes in input order. + /// Existing artifacts and duplicate blobs are uploaded only once. + pub async fn cache_artifacts(&self, blobs: Vec>) -> Result> { + use crate::artifact::InMemoryArtifact; + use sha2::{Digest, Sha256}; + use std::collections::HashSet; + + let hashes: Vec = blobs + .iter() + .map(|blob| format!("{:x}", Sha256::digest(blob))) + .collect(); + let names: Vec = hashes.iter().map(|hash| format!("cache/{hash}")).collect(); + let unique: HashSet<&str> = names.iter().map(String::as_str).collect(); + if unique.is_empty() { + return Ok(hashes); + } + let cached = self + .artifact_status(&unique.into_iter().collect::>()) + .await?; + let mut seen = HashSet::new(); + let mut artifacts = Vec::new(); + for (blob, name) in blobs.into_iter().zip(names) { + if !cached.get(&name).copied().unwrap_or(false) && seen.insert(name.clone()) { + artifacts.push(NamedArtifact::new( + name, + Box::new(InMemoryArtifact::new(blob)), + )); + } + } + if artifacts.is_empty() { + return Ok(hashes); + } + let expected: HashSet = artifacts.iter().map(|a| a.name.clone()).collect(); + let mut requests = build_artifact_request_stream(self.session_id.clone(), artifacts)?; + for request in &mut requests { + request.user_context = Some(UserContext { + user_id: self.user_id.clone().unwrap_or_default(), + ..Default::default() + }); + } + let mut req = Request::new(futures::stream::iter(requests)); + self._attach_metadata(&mut req); + let response = self + .stub + .clone() + .add_artifacts(req) + .await + .map_err(SparkError::from_grpc_status)? + .into_inner(); + let mut received = HashSet::new(); + for summary in response.artifacts { + if !summary.is_crc_successful { + return Err(SparkError::connect_msg(format!( + "CRC check failed for artifact: {}", + summary.name + ))); + } + received.insert(summary.name); + } + if !expected.is_subset(&received) { + return Err(SparkError::connect_msg( + "Missing cache artifact upload acknowledgement", + )); + } + Ok(hashes) + } + /// Check artifact status on the server. /// /// Mirrors functionality from `pyspark.sql.connect.client.artifact.ArtifactManager.is_cached_artifact`. diff --git a/crates/spark-connect/src/arrow_input.rs b/crates/spark-connect/src/arrow_input.rs new file mode 100644 index 00000000..8ef091e3 --- /dev/null +++ b/crates/spark-connect/src/arrow_input.rs @@ -0,0 +1,235 @@ +//! Arrow local relations and session-scoped chunked uploads. + +use arrow::datatypes::SchemaRef; +use arrow::ipc::writer::StreamWriter; +use arrow::record_batch::RecordBatch; +use spark_connect_core::client::SparkConnectClient; +use spark_connect_core::error::{Result, SparkError, SparkErrorKind}; +use spark_connect_core::runtime::block_on; + +use crate::plan::LogicalPlan; +use crate::types::DataType; + +struct Options { + cache_threshold: usize, + size_limit: usize, + chunk_rows: usize, + chunk_bytes: usize, + upload_batch_bytes: usize, +} + +impl Options { + fn from_client(client: &SparkConnectClient) -> Result { + let keys = [ + "spark.sql.session.localRelationCacheThreshold", + "spark.sql.session.localRelationSizeLimit", + "spark.sql.session.localRelationChunkSizeRows", + "spark.sql.session.localRelationChunkSizeBytes", + "spark.sql.session.localRelationBatchOfChunksSizeBytes", + ]; + let values = block_on(client.get_configs(&keys))?; + let mut parsed = Vec::new(); + for (i, key) in keys.iter().enumerate() { + let value = values + .get(i) + .and_then(Option::as_deref) + .and_then(|v| v.parse::().ok()) + .filter(|v| i == 0 || *v > 0) + .ok_or_else(|| { + SparkError::connect_msg(format!("Invalid or missing configuration: {key}")) + })?; + parsed.push(value); + } + Ok(Self { + cache_threshold: parsed[0], + size_limit: parsed[1], + chunk_rows: parsed[2], + chunk_bytes: parsed[3].min(parsed[4]), + upload_batch_bytes: parsed[4], + }) + } +} + +fn arrow_error(error: arrow::error::ArrowError) -> SparkError { + SparkError::connect_msg(error.to_string()) +} + +fn validate(batches: &[RecordBatch]) -> Result<()> { + if let Some(first) = batches.first() { + if batches.iter().any(|batch| batch.schema() != first.schema()) { + return Err(SparkError::connect_msg("Arrow batch schemas must match")); + } + } + Ok(()) +} + +pub(crate) fn inline_relation(batches: &[RecordBatch], schema: DataType) -> Result { + validate(batches)?; + let data = if let Some(first) = batches.first() { + let mut writer = StreamWriter::try_new(Vec::new(), &first.schema()).map_err(arrow_error)?; + for batch in batches { + writer.write(batch).map_err(arrow_error)?; + } + Some(writer.into_inner().map_err(arrow_error)?) + } else { + None + }; + Ok(LogicalPlan::LocalRelation { schema, data }) +} + +pub(crate) fn create_relation( + client: &SparkConnectClient, + batches: Vec, + schema: DataType, +) -> Result { + // Validate before any RPC, and preserve the offline empty-input path. + validate(&batches)?; + if batches.iter().all(|batch| batch.num_rows() == 0) { + return inline_relation(&batches, schema); + } + let options = Options::from_client(client)?; + // Count the logical slices rather than their potentially much larger backing buffers. + let bytes: usize = batches + .iter() + .flat_map(|b| b.columns()) + .map(|a| { + a.to_data() + .get_slice_memory_size() + .unwrap_or_else(|_| a.get_buffer_memory_size()) + }) + .sum(); + if bytes < options.cache_threshold { + return inline_relation(&batches, schema); + } + + let mut pending = Vec::new(); + let mut pending_size = 0; + let mut total_size = 0; + let mut hashes = Vec::new(); + let mut enqueue = |chunk: Vec| -> Result<()> { + total_size += chunk.len(); + if total_size > options.size_limit { + return Err(SparkError::classed( + SparkErrorKind::Analysis, + "LOCAL_RELATION_SIZE_LIMIT_EXCEEDED", + &[ + ("actualSize", &total_size.to_string()), + ("sizeLimit", &options.size_limit.to_string()), + ], + )); + } + if !pending.is_empty() && pending_size + chunk.len() > options.upload_batch_bytes { + hashes.extend(block_on( + client.cache_artifacts(std::mem::take(&mut pending)), + )?); + pending_size = 0; + } + pending_size += chunk.len(); + pending.push(chunk); + Ok(()) + }; + enqueue(schema.json().into_bytes())?; + serialize_chunks(&batches, &options, &mut enqueue)?; + hashes.extend(block_on(client.cache_artifacts(pending))?); + let schema_hash = hashes.remove(0); + Ok(LogicalPlan::ChunkedCachedLocalRelation { + data_hashes: hashes, + schema_hash: Some(schema_hash), + }) +} + +fn new_writer(schema: &SchemaRef) -> Result>> { + StreamWriter::try_new(Vec::new(), schema).map_err(arrow_error) +} + +/// Each chunk is an independently readable IPC stream. Check limits at most every +/// 1024 rows, as reference PySpark does; a single oversized row remains indivisible. +fn serialize_chunks( + batches: &[RecordBatch], + options: &Options, + emit: &mut impl FnMut(Vec) -> Result<()>, +) -> Result<()> { + let schema = batches[0].schema(); + let mut writer = new_writer(&schema)?; + let mut rows = 0; + for batch in batches { + let mut offset = 0; + while offset < batch.num_rows() { + let length = (batch.num_rows() - offset) + .min(1024) + .min(options.chunk_rows - rows); + writer + .write(&batch.slice(offset, length)) + .map_err(arrow_error)?; + offset += length; + rows += length; + if rows >= options.chunk_rows || writer.get_ref().len() >= options.chunk_bytes { + emit(writer.into_inner().map_err(arrow_error)?)?; + writer = new_writer(&schema)?; + rows = 0; + } + } + } + if rows > 0 { + emit(writer.into_inner().map_err(arrow_error)?)?; + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::Int64Array; + use arrow::datatypes::{DataType as ArrowType, Field, Schema}; + use arrow::ipc::reader::StreamReader; + use std::io::Cursor; + use std::sync::Arc; + + #[test] + fn chunks_preserve_rows_schema_and_order() { + let schema = Arc::new(Schema::new(vec![Field::new("id", ArrowType::Int64, true)])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int64Array::from_iter_values(0..2500))], + ) + .unwrap(); + for (chunk_rows, chunk_bytes) in [(700, usize::MAX), (10000, 1024)] { + let options = Options { + cache_threshold: 0, + size_limit: usize::MAX, + chunk_rows, + chunk_bytes, + upload_batch_bytes: usize::MAX, + }; + let mut chunks = Vec::new(); + serialize_chunks(std::slice::from_ref(&batch), &options, &mut |c| { + chunks.push(c); + Ok(()) + }) + .unwrap(); + assert!(chunks.len() > 1); + let mut actual = Vec::new(); + for chunk in chunks { + let reader = StreamReader::try_new(Cursor::new(chunk), None).unwrap(); + assert_eq!(reader.schema(), schema); + let mut rows = 0; + for batch in reader { + let batch = batch.unwrap(); + rows += batch.num_rows(); + actual.extend( + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .iter() + .copied(), + ); + } + assert!(rows <= chunk_rows); + } + assert_eq!(actual, (0..2500).collect::>()); + } + } +} diff --git a/crates/spark-connect/src/dataframe.rs b/crates/spark-connect/src/dataframe.rs index ea0e2c62..d66deb6a 100644 --- a/crates/spark-connect/src/dataframe.rs +++ b/crates/spark-connect/src/dataframe.rs @@ -1609,7 +1609,10 @@ impl DataFrame { /// Check if this DataFrame is local (collected). pub fn is_local(&self) -> bool { - matches!(self.plan.as_ref(), LogicalPlan::LocalRelation { .. }) + matches!( + self.plan.as_ref(), + LogicalPlan::LocalRelation { .. } | LogicalPlan::ChunkedCachedLocalRelation { .. } + ) } /// Check if this DataFrame is streaming. diff --git a/crates/spark-connect/src/lib.rs b/crates/spark-connect/src/lib.rs index 5f1a8faf..180bf089 100644 --- a/crates/spark-connect/src/lib.rs +++ b/crates/spark-connect/src/lib.rs @@ -5,6 +5,7 @@ #![allow(clippy::result_large_err)] // SparkError is large; changing the type would break the API +mod arrow_input; pub mod catalog; pub mod column; pub mod conf; diff --git a/crates/spark-connect/src/plan.rs b/crates/spark-connect/src/plan.rs index b24309bf..86dacf27 100644 --- a/crates/spark-connect/src/plan.rs +++ b/crates/spark-connect/src/plan.rs @@ -228,6 +228,11 @@ pub enum LogicalPlan { schema: DataType, data: Option>, }, + /// Local Arrow data uploaded as session-scoped cache artifacts. + ChunkedCachedLocalRelation { + data_hashes: Vec, + schema_hash: Option, + }, /// CachedRemoteRelation: a reference to a cached remote relation. CachedRemoteRelation { relation_id: String }, /// Read: `spark.read.format(...).load(...)` or `spark.read.table(...)`. @@ -944,6 +949,17 @@ impl LogicalPlan { Some(proto::relation::RelType::SubqueryAlias(Box::new(sq_alias))); } + LogicalPlan::ChunkedCachedLocalRelation { + data_hashes, + schema_hash, + } => { + relation.rel_type = Some(proto::relation::RelType::ChunkedCachedLocalRelation( + proto::ChunkedCachedLocalRelation { + data_hashes: data_hashes.clone(), + schema_hash: schema_hash.clone(), + }, + )); + } LogicalPlan::LocalRelation { schema, data } => { let mut local = proto::LocalRelation::default(); if let Some(d) = data { diff --git a/crates/spark-connect/src/session.rs b/crates/spark-connect/src/session.rs index 31843d53..ce92949e 100644 --- a/crates/spark-connect/src/session.rs +++ b/crates/spark-connect/src/session.rs @@ -294,38 +294,15 @@ impl SparkSession { /// /// `schema` is the Spark schema applied by the server. All batches must share /// the same Arrow schema; IPC serialization is still required for transport. + /// Non-empty inputs read the server's local-relation configuration. Large inputs + /// are uploaded as cache artifacts and referenced by hash in the resulting plan. pub fn create_dataframe_from_arrow( &self, batches: Vec, schema: DataType, ) -> Result { - let data = if let Some(first) = batches.first() { - let arrow_schema = first.schema(); - let mut buffer = Vec::new(); - { - let mut writer = - arrow::ipc::writer::StreamWriter::try_new(&mut buffer, &arrow_schema) - .map_err(|e| SparkError::connect_msg(e.to_string()))?; - for batch in &batches { - if batch.schema() != arrow_schema { - return Err(SparkError::connect_msg("Arrow batch schemas must match")); - } - writer - .write(batch) - .map_err(|e| SparkError::connect_msg(e.to_string()))?; - } - writer - .finish() - .map_err(|e| SparkError::connect_msg(e.to_string()))?; - } - Some(buffer) - } else { - None - }; - Ok(DataFrame::new( - self.clone(), - LogicalPlan::LocalRelation { schema, data }, - )) + let plan = crate::arrow_input::create_relation(&self.client, batches, schema)?; + Ok(DataFrame::new(self.clone(), plan)) } /// Create a DataFrameReader for reading data from various sources. @@ -1276,12 +1253,12 @@ mod tests { fields: vec![struct_field("x", DataType::Integer)], }; let spark = session(); - let df = spark - .create_dataframe_from_arrow(vec![batch.clone(), batch.clone()], schema.clone()) - .unwrap(); + let plan = + crate::arrow_input::inline_relation(&[batch.clone(), batch.clone()], schema.clone()) + .unwrap(); let LogicalPlan::LocalRelation { data: Some(data), .. - } = df.plan() + } = &plan else { panic!("expected Arrow local relation") }; diff --git a/docs/architecture.md b/docs/architecture.md index 1fa9261c..93107c3e 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -117,6 +117,12 @@ timestamp conversion helpers normalize the table, then arrow-rs imports its those batches to IPC for the local relation sent to Spark. No Python rows are materialized; casts, buffer alignment, and IPC serialization may still copy. Empty tables require an explicit schema, matching PySpark Connect. +Inputs at or above `spark.sql.session.localRelationCacheThreshold` are serialized +into independent IPC streams and uploaded as session-scoped cache artifacts. +The plan carries `ChunkedCachedLocalRelation` hashes instead of embedding the +entire payload. Chunk sizes, upload batch sizes, and the total relation size +limit follow the server's local-relation settings; artifact RPC messages use +32 KiB pieces, with SHA-256 deduplication and CRC checks. ## Two front-ends, one core diff --git a/python/tests/test_arrow_handoff.py b/python/tests/test_arrow_handoff.py index 5deef93d..a9a3e513 100644 --- a/python/tests/test_arrow_handoff.py +++ b/python/tests/test_arrow_handoff.py @@ -347,3 +347,213 @@ def test_malformed_ipc_is_reported(protocol): with pytest.raises(RuntimeError, match="Arrow"): frame.toArrow() assert calls == {"execute": 1} + + +@contextmanager +def _arrow_input_service(protocol, *, upload_error=None, **overrides): + """Resolve cache artifacts at a mock service with a 4 MiB receive limit.""" + import hashlib + import zlib + + config = { + "spark.sql.timestampType": "TIMESTAMP_LTZ", + "spark.sql.session.timeZone": "UTC", + "spark.sql.execution.arrow.useLargeVarTypes": "false", + "spark.sql.session.localRelationCacheThreshold": "1048576", + "spark.sql.session.localRelationSizeLimit": str(3 * 1024**3), + "spark.sql.session.localRelationChunkSizeRows": "10000", + "spark.sql.session.localRelationChunkSizeBytes": str(16 * 1024**2), + "spark.sql.session.localRelationBatchOfChunksSizeBytes": str(1024**3), + } + config.update(overrides) + state = {"cache": {}, "uploads": 0, "relations": [], "message_sizes": []} + server_id = "00000000-0000-0000-0000-000000000001" + + def configure(raw, context): + request = protocol("ConfigRequest").FromString(raw) + return protocol("ConfigResponse", session_id=request.session_id, + pairs=[dict(key=k, value=config[k]) for k in request.operation.get.keys] + ).SerializeToString() + + def statuses(raw, context): + request = protocol("ArtifactStatusesRequest").FromString(raw) + return protocol("ArtifactStatusesResponse", session_id=request.session_id, + statuses={name: dict(exists=name in state["cache"]) + for name in request.names}).SerializeToString() + + def upload(requests, context): + summaries = [] + name, expected, buffer = None, 0, bytearray() + session_id = None + + def save(name, data): + assert name == "cache/" + hashlib.sha256(data).hexdigest() + assert name not in state["cache"], "cached artifacts should not be uploaded again" + state["cache"][name] = bytes(data) + summaries.append(dict(name=name, is_crc_successful=True)) + + def checked(chunk): + assert zlib.crc32(chunk.data) == chunk.crc + return chunk.data + + for raw in requests: + state["message_sizes"].append(len(raw)) + request = protocol("AddArtifactsRequest").FromString(raw) + session_id = request.session_id + if request.HasField("batch"): + for artifact in request.batch.artifacts: + save(artifact.name, checked(artifact.data)) + elif request.HasField("begin_chunk"): + begin = request.begin_chunk + name, expected = begin.name, begin.total_bytes + buffer = bytearray(checked(begin.initial_chunk)) + else: + buffer.extend(checked(request.chunk)) + if name is not None and len(buffer) == expected: + save(name, buffer) + name, buffer = None, bytearray() + assert name is None + state["uploads"] += 1 + if upload_error == "crc": + summaries[0]["is_crc_successful"] = False + elif upload_error == "missing": + summaries = [] + return protocol("AddArtifactsResponse", session_id=session_id, + artifacts=summaries).SerializeToString() + + def execute(raw, context): + request = protocol("ExecutePlanRequest").FromString(raw) + relation = request.plan.root + state["relations"].append(relation) + if relation.HasField("chunked_cached_local_relation"): + cached = relation.chunked_cached_local_relation + payloads = [state["cache"]["cache/" + h] for h in cached.dataHashes] + state["spark_schema"] = state["cache"]["cache/" + cached.schemaHash].decode() + else: + payloads = [relation.local_relation.data] + rows, first, last = 0, None, None + for payload in payloads: + for batch in pa.ipc.open_stream(payload): + rows += batch.num_rows + if batch.num_rows: + if first is None: + first = batch.column(0)[0].as_py() + last = batch.column(0)[-1].as_py() + state["result"] = (rows, first, last) + result = pa.table({"rows": [rows]}) + sink = pa.BufferOutputStream() + with pa.ipc.new_stream(sink, result.schema) as writer: + writer.write_table(result) + header = dict(session_id=request.session_id, operation_id=request.operation_id, + server_side_session_id=server_id) + yield protocol("ExecutePlanResponse", **header, response_id="0", + arrow_batch=dict(row_count=1, data=sink.getvalue().to_pybytes()) + ).SerializeToString() + yield protocol("ExecutePlanResponse", **header, response_id="1", + result_complete={}).SerializeToString() + + def release(raw, context): + request = protocol("ReleaseExecuteRequest").FromString(raw) + return protocol("ReleaseExecuteResponse", session_id=request.session_id, + operation_id=request.operation_id, + server_side_session_id=server_id).SerializeToString() + + def stop(raw, context): + request = protocol("ReleaseSessionRequest").FromString(raw) + 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, options=[("grpc.max_receive_message_length", 4 * 1024**2)]) + server.add_generic_rpc_handlers([grpc.method_handlers_generic_handler( + "spark.connect.SparkConnectService", { + "Config": grpc.unary_unary_rpc_method_handler(configure), + "ArtifactStatus": grpc.unary_unary_rpc_method_handler(statuses), + "AddArtifacts": grpc.stream_unary_rpc_method_handler(upload), + "ExecutePlan": grpc.unary_stream_rpc_method_handler(execute), + "ReleaseExecute": grpc.unary_unary_rpc_method_handler(release), + "ReleaseSession": grpc.unary_unary_rpc_method_handler(stop), + })]) + port = server.add_insecure_port("127.0.0.1:0") + server.start() + session = SparkSession.builder.remote(f"sc://127.0.0.1:{port}").getOrCreate() + try: + yield session, state + finally: + session.stop() + server.stop(0).wait(timeout=5) + + +@pytest.mark.parametrize("rows", [2, 17000000]) +def test_arrow_input_upload_above_grpc_limit(protocol, rows): + import numpy as np + import json + + table = pa.table({"id": np.arange(rows, dtype=np.int64)}) + with _arrow_input_service(protocol) as (spark, state): + df = spark.createDataFrame(table) + assert df.isLocal() + assert df.toArrow().column(0).to_pylist() == [rows] + assert state["result"] == (rows, 0, rows - 1) + if rows == 2: + assert state["uploads"] == 0 + assert state["relations"][0].HasField("local_relation") + else: + assert table.nbytes > 128 * 1024**2 + assert state["relations"][0].HasField("chunked_cached_local_relation") + assert state["relations"][0].ByteSize() < 1024**2 + assert max(state["message_sizes"]) < 33 * 1024 + assert json.loads(state["spark_schema"])["fields"][0]["name"] == "id" + uploads = state["uploads"] + spark.createDataFrame(table).toArrow() + assert state["uploads"] == uploads + + +def test_arrow_input_upload_limits_and_chunk_order(protocol): + import numpy as np + from pyspark.errors import AnalysisException + table = pa.table({"id": np.arange(100, dtype=np.int64)}) + options = { + "spark.sql.session.localRelationCacheThreshold": "0", + "spark.sql.session.localRelationChunkSizeRows": "7", + "spark.sql.session.localRelationChunkSizeBytes": "1000", + "spark.sql.session.localRelationBatchOfChunksSizeBytes": "1000", + } + with _arrow_input_service(protocol, **options) as (spark, state): + df = spark.createDataFrame(table) + df.toArrow() + assert state["result"] == (100, 0, 99) + assert state["uploads"] > 1 + cached = state["relations"][0].chunked_cached_local_relation + assert len(cached.dataHashes) > 1 + for hash in cached.dataHashes: + assert pa.ipc.open_stream(state["cache"]["cache/" + hash]).read_all().num_rows <= 7 + + options["spark.sql.session.localRelationSizeLimit"] = "10" + with _arrow_input_service(protocol, **options) as (spark, state): + with pytest.raises(AnalysisException, match="LOCAL_RELATION_SIZE_LIMIT_EXCEEDED"): + spark.createDataFrame(table) + assert not state["relations"] + assert state["uploads"] == 0 + + +@pytest.mark.parametrize("failure,match", [("crc", "CRC check failed"), + ("missing", "Missing cache artifact")]) +def test_arrow_input_upload_failure(protocol, failure, match): + with _arrow_input_service(protocol, upload_error=failure, + **{"spark.sql.session.localRelationCacheThreshold": "0"}) as (spark, state): + with pytest.raises(Exception, match=match): + spark.createDataFrame(pa.table({"id": [1, 2]})) + assert not state["relations"] + + +def test_arrow_input_duplicate_chunks_preserve_multiplicity(protocol): + options = {"spark.sql.session.localRelationCacheThreshold": "0", + "spark.sql.session.localRelationChunkSizeRows": "5"} + with _arrow_input_service(protocol, **options) as (spark, state): + spark.createDataFrame(pa.table({"id": [1] * 20})).toArrow() + cached = state["relations"][0].chunked_cached_local_relation + assert len(cached.dataHashes) == 4 + assert len(set(cached.dataHashes)) == 1 + assert len(state["cache"]) == 2 # One schema and one unique data chunk. + assert state["result"] == (20, 1, 1) diff --git a/python/tests/test_dropin_offline.py b/python/tests/test_dropin_offline.py index 1aab0918..168ce7b7 100644 --- a/python/tests/test_dropin_offline.py +++ b/python/tests/test_dropin_offline.py @@ -2863,3 +2863,24 @@ def test_arrow_input_types_live_server(kind): assert df.toArrow().to_pylist() == table.to_pylist() finally: spark.stop() + + +def test_arrow_input_large_upload_live_server(): + """A table larger than one gRPC message must use uploaded cache artifacts.""" + import os + pa = pytest.importorskip("pyarrow") + remote = os.environ.get("SPARK_REMOTE") + if not remote: + pytest.skip("requires SPARK_REMOTE") + from pyspark.sql import SparkSession + + rows = 140000 + table = pa.table({"payload": [b"x" * 1024] * rows}) + assert table.nbytes > 128 * 1024**2 + spark = SparkSession.builder.remote(remote).getOrCreate() + try: + df = spark.createDataFrame(table) + assert df.count() == rows + assert df.selectExpr("sum(length(payload)) AS total").collect()[0][0] == rows * 1024 + finally: + spark.stop() From 911e23cbe21409fbb530894939b4f0134e69d6af Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Thu, 29 Oct 2026 01:33:49 -0700 Subject: [PATCH 7/7] [SPARK-60698][CONNECT] Batch Arrow input normalization config reads Read the three mutable Arrow normalization settings in one Get request per nonempty construction while retaining empty-input and transport behavior. Co-authored-by: Isaac --- crates/pyspark-rs/src/pyarrow_import.rs | 17 ++- crates/spark-connect/src/conf.rs | 10 +- python/tests/test_arrow_handoff.py | 140 +++++++++++++++++++++++- 3 files changed, 157 insertions(+), 10 deletions(-) diff --git a/crates/pyspark-rs/src/pyarrow_import.rs b/crates/pyspark-rs/src/pyarrow_import.rs index f1b083be..058fb99f 100644 --- a/crates/pyspark-rs/src/pyarrow_import.rs +++ b/crates/pyspark-rs/src/pyarrow_import.rs @@ -47,12 +47,19 @@ pub(crate) fn create_dataframe( (None, None, None) } else { py.detach(|| -> PyResult<_> { - let conf = session.conf(); + let values = session + .conf() + .get_many(&[ + "spark.sql.timestampType", + "spark.sql.session.timeZone", + "spark.sql.execution.arrow.useLargeVarTypes", + ]) + .to_pyerr()?; + let mut values = values.into_iter(); Ok(( - conf.get("spark.sql.timestampType").to_pyerr()?, - conf.get("spark.sql.session.timeZone").to_pyerr()?, - conf.get("spark.sql.execution.arrow.useLargeVarTypes") - .to_pyerr()?, + values.next().flatten(), + values.next().flatten(), + values.next().flatten(), )) })? }; diff --git a/crates/spark-connect/src/conf.rs b/crates/spark-connect/src/conf.rs index d43d611f..17a16ec6 100644 --- a/crates/spark-connect/src/conf.rs +++ b/crates/spark-connect/src/conf.rs @@ -58,10 +58,18 @@ impl RuntimeConf { /// println!("Shuffle partitions: {:?}", value); /// ``` pub fn get(&self, key: &str) -> Result> { - let results = block_on(self.client.get_configs(&[key]))?; + let results = self.get_many(&[key])?; Ok(results.into_iter().next().flatten()) } + /// Get configuration values in one request, in the same order as `keys`. + /// + /// Missing keys return `None`, as with [`Self::get`]. Values are read from + /// the server on every call; they are not cached locally. + pub fn get_many(&self, keys: &[&str]) -> Result>> { + block_on(self.client.get_configs(keys)) + } + /// Get a configuration value, returning `default` when the key is unset — instead of /// erroring as `get` does for an unknown key. Uses the server's `GetWithDefault` /// operation, mirroring `RuntimeConfig.get(key, default)`. diff --git a/python/tests/test_arrow_handoff.py b/python/tests/test_arrow_handoff.py index a9a3e513..43382b55 100644 --- a/python/tests/test_arrow_handoff.py +++ b/python/tests/test_arrow_handoff.py @@ -366,13 +366,20 @@ def _arrow_input_service(protocol, *, upload_error=None, **overrides): "spark.sql.session.localRelationBatchOfChunksSizeBytes": str(1024**3), } config.update(overrides) - state = {"cache": {}, "uploads": 0, "relations": [], "message_sizes": []} + state = {"cache": {}, "uploads": 0, "relations": [], "message_sizes": [], + "config": config, "config_requests": [], "input_tables": []} server_id = "00000000-0000-0000-0000-000000000001" def configure(raw, context): request = protocol("ConfigRequest").FromString(raw) + keys = list(request.operation.get.keys) + state["config_requests"].append(keys) + if error := state.pop("config_error", None): + context.abort(error, "configuration read failed") + # Deliberately reverse responses and omit missing keys to test key mapping. return protocol("ConfigResponse", session_id=request.session_id, - pairs=[dict(key=k, value=config[k]) for k in request.operation.get.keys] + pairs=[dict(key=k, value=config[k]) for k in reversed(keys) + if k in config] ).SerializeToString() def statuses(raw, context): @@ -432,8 +439,10 @@ def execute(raw, context): else: payloads = [relation.local_relation.data] rows, first, last = 0, None, None - for payload in payloads: - for batch in pa.ipc.open_stream(payload): + tables = [pa.ipc.open_stream(payload).read_all() for payload in payloads] + state["input_tables"].append(pa.concat_tables(tables)) + for table in tables: + for batch in table.to_batches(): rows += batch.num_rows if batch.num_rows: if first is None: @@ -557,3 +566,126 @@ def test_arrow_input_duplicate_chunks_preserve_multiplicity(protocol): assert len(set(cached.dataHashes)) == 1 assert len(state["cache"]) == 2 # One schema and one unique data chunk. assert state["result"] == (20, 1, 1) + + +_NORMALIZATION_KEYS = [ + "spark.sql.timestampType", + "spark.sql.session.timeZone", + "spark.sql.execution.arrow.useLargeVarTypes", +] +_UPLOAD_KEYS = [ + "spark.sql.session.localRelationCacheThreshold", + "spark.sql.session.localRelationSizeLimit", + "spark.sql.session.localRelationChunkSizeRows", + "spark.sql.session.localRelationChunkSizeBytes", + "spark.sql.session.localRelationBatchOfChunksSizeBytes", +] + + +@pytest.mark.parametrize("timestamp_type", ["TIMESTAMP_LTZ", "TIMESTAMP_NTZ"]) +@pytest.mark.parametrize("zone", ["UTC", "America/Los_Angeles"]) +@pytest.mark.parametrize("large", ["false", "true"]) +def test_arrow_input_batched_normalization_configs(protocol, timestamp_type, zone, large): + import json + from zoneinfo import ZoneInfo + + value = datetime(2024, 1, 15, 12, 34, 56, 123456) + table = pa.table({"t": [value, None], "s": ["hello", None], "b": [b"xyz", None]}) + with _arrow_input_service(protocol, **dict(zip( + _NORMALIZATION_KEYS, [timestamp_type, zone, large] + ))) as (spark, state): + frame = spark.createDataFrame(table) + assert state["config_requests"] == [_NORMALIZATION_KEYS, _UPLOAD_KEYS] + frame.collect() + expected_time = (value if timestamp_type == "TIMESTAMP_NTZ" else + value.replace(tzinfo=ZoneInfo(zone)).astimezone(timezone.utc)) + expected = pa.table({ + "t": pa.array([expected_time, None], type=pa.timestamp( + "us", tz=None if timestamp_type == "TIMESTAMP_NTZ" else "UTC")), + "s": pa.array(["hello", None], type=pa.large_string() if large == "true" + else pa.string()), + "b": pa.array([b"xyz", None], type=pa.large_binary() if large == "true" + else pa.binary()), + }) + actual = state["input_tables"][0] + assert actual.schema == expected.schema + assert actual.to_pydict() == expected.to_pydict() + fields = json.loads(state["relations"][0].local_relation.schema)["fields"] + assert [f["type"] for f in fields] == [ + "timestamp_ntz" if timestamp_type == "TIMESTAMP_NTZ" else "timestamp", + "string", "binary", + ] + + +@pytest.mark.parametrize("missing", _NORMALIZATION_KEYS) +@pytest.mark.parametrize("present_null", [False, True]) +def test_arrow_input_missing_normalization_config(protocol, missing, present_null): + with _arrow_input_service(protocol) as (spark, state): + if present_null: + state["config"][missing] = None + else: + del state["config"][missing] + spark.createDataFrame(pa.table({"x": [1, None], "s": ["hello", None]})).collect() + assert state["config_requests"] == [_NORMALIZATION_KEYS, _UPLOAD_KEYS] + expected = pa.table({"x": pa.array([1, None], type=pa.int64()), + "s": ["hello", None]}) + assert state["input_tables"][0].equals(expected) + + +def test_arrow_input_normalization_config_failure_then_fresh_read(protocol): + with _arrow_input_service(protocol) as (spark, state): + state["config_error"] = grpc.StatusCode.INVALID_ARGUMENT + table = pa.table({"x": [1]}) + with pytest.raises(Exception, match="configuration read failed"): + spark.createDataFrame(table) + assert state["config_requests"] == [_NORMALIZATION_KEYS] + assert state["uploads"] == 0 + spark.createDataFrame(table).collect() + assert state["config_requests"] == [ + _NORMALIZATION_KEYS, _NORMALIZATION_KEYS, _UPLOAD_KEYS] + assert state["input_tables"][0].to_pydict() == {"x": [1]} + + +def test_arrow_input_normalization_config_transient_retry(protocol): + with _arrow_input_service(protocol) as (spark, state): + state["config_error"] = grpc.StatusCode.UNAVAILABLE + spark.createDataFrame(pa.table({"x": [1]})).collect() + assert state["config_requests"] == [ + _NORMALIZATION_KEYS, _NORMALIZATION_KEYS, _UPLOAD_KEYS] + assert state["input_tables"][0].to_pydict() == {"x": [1]} + + +def test_arrow_input_normalization_configs_are_fresh(protocol): + import json + + value = datetime(2024, 1, 15, 12) + table = pa.table({"t": [value], "s": ["hello"]}) + with _arrow_input_service(protocol) as (spark, state): + spark.createDataFrame(table).collect() + state["config"].update(dict(zip(_NORMALIZATION_KEYS, [ + "TIMESTAMP_NTZ", "America/Los_Angeles", "true"]))) + spark.createDataFrame(table).collect() + state["config"].update(dict(zip(_NORMALIZATION_KEYS, [ + "TIMESTAMP_LTZ", "America/Los_Angeles", "false"]))) + spark.createDataFrame(table).collect() + assert state["config_requests"] == [_NORMALIZATION_KEYS, _UPLOAD_KEYS] * 3 + tables = state["input_tables"] + assert [t.column("t")[0].as_py() for t in tables] == [ + value.replace(tzinfo=timezone.utc), value, + datetime(2024, 1, 15, 20, tzinfo=timezone.utc)] + assert [t.schema.field("s").type for t in tables] == [ + pa.string(), pa.large_string(), pa.string()] + assert [json.loads(r.local_relation.schema)["fields"][0]["type"] + for r in state["relations"]] == ["timestamp", "timestamp_ntz", "timestamp"] + + +def test_arrow_input_empty_skips_all_configs(protocol): + from pyspark.sql.types import LongType, StructField, StructType + + with _arrow_input_service(protocol) as (spark, state): + table = pa.table({"x": pa.array([], type=pa.int64())}) + spark.createDataFrame(table, StructType([StructField("x", LongType())])) + with pytest.raises(Exception, match="CANNOT_INFER_EMPTY_SCHEMA"): + spark.createDataFrame(table) + assert state["config_requests"] == [] + assert state["uploads"] == 0