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..72876801 --- /dev/null +++ b/crates/spark-python-arrow/README.md @@ -0,0 +1,103 @@ + + +# 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 C Data 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. + +PyArrow consumes the exported descriptors and retains their release callbacks, +so the data survives after Rust's temporary wrappers have been dropped. If an +importer raises before consuming a descriptor, Rust releases it. 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 outbound schema helpers also +retain sorted-map and dictionary-order flags when present. + +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. Sliced arrays whose value and null offsets differ can require +copying the validity bitmap during C Data 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..e03fd79f --- /dev/null +++ b/crates/spark-python-arrow/src/bridge.rs @@ -0,0 +1,237 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Export Arrow record batches to PyArrow through the Arrow C Data Interface. +//! +//! The caller supplies valid Rust record batches and owns the Python interpreter. +//! PyArrow's importers consume the exported descriptors and retain their buffers. + +use std::sync::Arc; + +use arrow::array::{Array, StructArray}; +use arrow::datatypes::{DataType, Field, Schema}; +use arrow::error::ArrowError; +use arrow::ffi::{FFI_ArrowArray, FFI_ArrowSchema}; +use arrow::ipc::writer::FileWriter; +use arrow::record_batch::RecordBatch; +use pyo3::exceptions::{PyRuntimeError, PyValueError}; +use pyo3::prelude::*; +use pyo3::types::{PyBytes, PyList}; + +// Called after table export has routed NUL names and timezones to the IPC fallback. +fn export_schema(py: Python<'_>, schema: &Schema) -> PyResult> { + let data_type = DataType::Struct(schema.fields().clone()); + let mut schema = export_c_data_type(&data_type) + .and_then(|exported| exported.with_metadata(schema.metadata())) + .map_err(arrow_error)?; + let address = &mut schema as *mut FFI_ArrowSchema as usize; + Ok(py + .import("pyarrow")? + .getattr("Schema")? + .call_method1("_import_from_c", (address,))? + .unbind()) +} + +/// Build a PyArrow table, sharing data buffers for C-compatible schemas. +/// +/// Batch boundaries, schema metadata, and zero-column row counts are preserved. +/// An empty batch vector produces an empty table with an empty schema. +/// Value and offset buffers are shared; sliced validity bitmaps may be realigned. +/// Inconsistent batch schemas, including metadata and ordering flags, raise `ValueError`. +pub fn record_batches_to_pyarrow(py: Python<'_>, batches: Vec) -> 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()); + } + + // Import once before any arrays. This also lets older PyArrow versions + // reject an unsupported schema without consuming any data descriptors. + let py_schema = export_schema(py, &schema)?; + let import_batch = pa.getattr("RecordBatch")?.getattr("_import_from_c")?; + let py_batches = PyList::empty(py); + for batch in batches { + let data = StructArray::from(batch).to_data(); + let mut array = FFI_ArrowArray::new(&data); + let address = &mut array as *mut FFI_ArrowArray as usize; + // Import consumes the array and clears its release callback. If import + // fails before consumption, Rust releases the remaining ownership. + py_batches.append(import_batch.call1((address, py_schema.bind(py)))?)?; + } + Ok(pa + .getattr("Table")? + .call_method1("from_batches", (py_batches, py_schema.bind(py)))? + .unbind()) +} + +fn arrow_error(error: ArrowError) -> PyErr { + PyRuntimeError::new_err(error.to_string()) +} + +// 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 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()); + } +} 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..5ca6104c 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 record batch 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,17 @@ 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()` hands +the record batches to `spark-python-arrow` using the Arrow C Data Interface. +Compatible Arrow data buffers are shared with PyArrow. Schemas with embedded +NUL field names or timezones use an IPC compatibility fallback, without another +RPC. `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..09358778 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 record batch 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"] ``` @@ -152,7 +154,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..05ff69f7 --- /dev/null +++ b/python/tests/test_arrow_handoff.py @@ -0,0 +1,521 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""Exercise the actual Rust drop-in's Arrow handoff, without a Spark server. + +The localhost service returns deterministic Arrow IPC results. Buffer addresses +are captured after Rust decodes IPC, at its C Data export boundary: addresses +in the server's Python arrays would not prove anything about this handoff. +""" + +import ctypes +import gc +from concurrent.futures import ThreadPoolExecutor +from contextlib import contextmanager +from datetime import datetime, timezone +from decimal import Decimal + +import grpc +import pandas as pd +import pyarrow as pa +import pytest +from pyspark.sql import SparkSession + + +def _varint(value): + encoded = bytearray() + while value > 127: + encoded.append((value & 127) | 128) + value >>= 7 + encoded.append(value) + return bytes(encoded) + + +def _bytes_field(number, value): + return _varint((number << 3) | 2) + _varint(len(value)) + value + + +def _request_fields(data): + """Read only the protobuf envelope, not the plan or other nested messages.""" + position = 0 + + def varint(): + nonlocal position + value = shift = 0 + while True: + byte = data[position] + position += 1 + value |= (byte & 127) << shift + if byte < 128: + return value + shift += 7 + + fields = {} + while position < len(data): + tag = varint() + wire_type = tag & 7 + if wire_type == 2: + length = varint() + fields[tag >> 3] = data[position : position + length] + position += length + elif wire_type == 0: + varint() + elif wire_type in (1, 5): + position += 8 if wire_type == 1 else 4 + else: + raise AssertionError(f"Unexpected protobuf wire type: {wire_type}") + return fields + + +@contextmanager +def _dataframe(batches, *, payloads=None, calls=None): + """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), +] + + +_Release = ctypes.CFUNCTYPE(None, ctypes.POINTER(_ArrowArray)) + +# 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 array in a pytest traceback. No Arrow buffers are retained. +_RELEASE_CALLBACK_KEEPALIVE = [] + + +@pytest.fixture(autouse=True) +def _forbid_c_stream_import(monkeypatch): + class NoStreamReader: + @staticmethod + def _import_from_c(*args): + raise AssertionError("The C Data handoff must not import a C Stream") + + _import_from_c_capsule = _import_from_c + + monkeypatch.setattr(pa, "RecordBatchReader", NoStreamReader) + + +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 _ArraySpy: + def __init__(self, monkeypatch, *, reject_at=None, consume_before_reject=False): + self.batches = [] + self.schemas = [] + self.release_counts = [] + self.import_count = 0 + self.callbacks = [] + real_batch = pa.RecordBatch + spy = self + + 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(address, schema): + array = ctypes.cast(address, ctypes.POINTER(_ArrowArray)) + batch_index = len(spy.batches) + spy.batches.append(_snapshot(array.contents)) + spy.schemas.append(schema) + spy.release_counts.append(0) + original_release = _Release(array.contents.release) + + @_Release + def release(current): + spy.release_counts[batch_index] += 1 + original_release(current) + + spy.callbacks.extend((original_release, release)) + _RELEASE_CALLBACK_KEEPALIVE.extend((original_release, release)) + array.contents.release = ctypes.cast(release, ctypes.c_void_p).value + reject = batch_index + 1 == reject_at + if reject and not consume_before_reject: + raise RuntimeError("test importer rejected the array") + imported = real_batch._import_from_c(address, schema) + spy.import_count += 1 + if reject: + # PyArrow owns the array now. The exception traceback may + # retain this local until gc.collect(), just like user code. + raise RuntimeError("test importer rejected the consumed array") + return imported + + monkeypatch.setattr(pa, "RecordBatch", RecordBatchProxy) + + +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 = _ArraySpy(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 == len(batches) + assert all(schema is spy.schemas[0] for schema in spy.schemas) + 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 = _ArraySpy(monkeypatch) + + def collect(): + with _dataframe([batch]) as (frame, session): + table = frame.toArrow() + session.stop() + return table + + table = collect() + gc.collect() + 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 = _ArraySpy(monkeypatch) + with _dataframe(batches) as (frame, _): + actual = frame.toPandas() + pd.testing.assert_frame_equal(actual, expected) + assert len(spy.batches) == 1 + del actual + gc.collect() + assert spy.release_counts == [1] + + +def test_import_rejection_releases_unconsumed_array(monkeypatch): + batch = _mixed_batch() + spy = _ArraySpy(monkeypatch, reject_at=1) + with _dataframe([batch]) as (frame, _): + with pytest.raises(RuntimeError, match="test importer rejected"): + frame.toArrow() + gc.collect() + assert spy.import_count == 0 + assert spy.release_counts == [1] + + +@pytest.mark.parametrize("consume_before_reject", [False, True]) +def test_second_import_failure_releases_arrays_once(monkeypatch, consume_before_reject): + batch = _mixed_batch() + spy = _ArraySpy(monkeypatch, reject_at=2, consume_before_reject=consume_before_reject) + with _dataframe([batch, batch, batch]) as (frame, _): + with pytest.raises(RuntimeError, match="test importer rejected"): + frame.toArrow() + gc.collect() + assert len(spy.batches) == 2 + assert spy.import_count == (2 if consume_before_reject else 1) + # The first batch has been imported and the second may or may not have + # transferred ownership. Neither can leak or be released twice on failure. + assert spy.release_counts == [1, 1] + + +def test_schema_import_rejection_prevents_array_import(monkeypatch): + batch = _mixed_batch() + spy = _ArraySpy(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 = _ArraySpy(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 = _ArraySpy(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_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 = _ArraySpy(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 = _ArraySpy(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 = _ArraySpy(monkeypatch) + with _dataframe([], payloads=[(1, b"not arrow ipc")]) as (frame, _): + with pytest.raises(RuntimeError, match="Arrow"): + frame.toArrow() + assert not spy.callbacks