tdcmeehan commented on code in PR #103: URL: https://github.com/apache/spark-connect-rust/pull/103#discussion_r4396392306
########## crates/spark-python-arrow/README.md: ########## @@ -0,0 +1,103 @@ +<!-- +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 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 Review Comment: This command differs from what CI runs (`cargo test -p spark-python-arrow --locked --features pyo3/abi3-py39` in `python-arrow.yml`). `CARGO_NET_OFFLINE=true` makes it fail unless every dependency is already in the local Cargo cache, and `CARGO_REGISTRIES_CRATES_IO_PROTOCOL=git` switches away from Cargo's default sparse index. Could the README show the CI command, with the offline variables as an optional note for offline environments? ########## .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" Review Comment: `pyproject.toml` declares `pyarrow>=10.0.0`, and `toArrow()` now goes through PyArrow's private `Schema._import_from_c` / `RecordBatch._import_from_c` instead of the public IPC reader. CI pins only 26.0.0 here, and the other workflows install the latest PyArrow, so the PyArrow 10 run in the PR description isn't repeated in CI. Could this job add a leg at the declared floor? If that isn't practical, the README could state which PyArrow versions are tested. ########## crates/spark-python-arrow/src/bridge.rs: ########## @@ -0,0 +1,237 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Export Arrow record batches to PyArrow through the Arrow C Data Interface. +//! +//! The caller supplies valid Rust record batches and owns the Python interpreter. +//! PyArrow's importers consume the exported descriptors and retain their buffers. + +use std::sync::Arc; + +use arrow::array::{Array, StructArray}; +use arrow::datatypes::{DataType, Field, Schema}; +use arrow::error::ArrowError; +use arrow::ffi::{FFI_ArrowArray, FFI_ArrowSchema}; +use arrow::ipc::writer::FileWriter; +use arrow::record_batch::RecordBatch; +use pyo3::exceptions::{PyRuntimeError, PyValueError}; +use pyo3::prelude::*; +use pyo3::types::{PyBytes, PyList}; + +// Called after table export has routed NUL names and timezones to the IPC fallback. +fn export_schema(py: Python<'_>, schema: &Schema) -> PyResult<Py<PyAny>> { + let data_type = DataType::Struct(schema.fields().clone()); + let mut schema = export_c_data_type(&data_type) + .and_then(|exported| exported.with_metadata(schema.metadata())) + .map_err(arrow_error)?; + let address = &mut schema as *mut FFI_ArrowSchema as usize; + Ok(py + .import("pyarrow")? + .getattr("Schema")? + .call_method1("_import_from_c", (address,))? + .unbind()) +} + +/// Build a PyArrow table, sharing data buffers for C-compatible schemas. +/// +/// Batch boundaries, schema metadata, and zero-column row counts are preserved. +/// An empty batch vector produces an empty table with an empty schema. +/// Value and offset buffers are shared; sliced validity bitmaps may be realigned. +/// Inconsistent batch schemas, including metadata and ordering flags, raise `ValueError`. +pub fn record_batches_to_pyarrow(py: Python<'_>, batches: Vec<RecordBatch>) -> PyResult<Py<PyAny>> { + let schema = batches + .first() + .map(|batch| batch.schema()) + .unwrap_or_else(|| Arc::new(Schema::empty())); + if batches.iter().any(|batch| { + !Arc::ptr_eq(batch.schema_ref(), &schema) && !schemas_equal(batch.schema_ref(), &schema) + }) { + return Err(PyValueError::new_err( + "Arrow record batches have inconsistent schemas", + )); + } + let pa = py.import("pyarrow")?; + if schema.fields().iter().any(|field| field_needs_ipc(field)) { + // 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<Vec<u8>> { + 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. Review Comment: Could this link the upstream issue (apache/arrow-rs#11603, fixed on `main` by apache/arrow-rs#11630)? That makes it clear when `export_c_data_type`/`rebuild` and the `0b100` flag below can be removed after an Arrow upgrade. ########## 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<Bound<'py, PyAny>> { + 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::<usize>()?, + 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::<usize>()?, 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::<usize>()?, + *address + ); + } + assert_eq!( + column + .call_method0("to_pylist")? + .extract::<Vec<Option<i64>>>()?, + 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::<usize>()?, 0); + assert_eq!( + empty.bind(py).getattr("num_columns")?.extract::<usize>()?, + 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::<usize>()?, 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::<usize>()?, + rows + ); + assert_eq!( + table.bind(py).getattr("num_columns")?.extract::<usize>()?, + 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::<PyValueError>(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)?)?; + } Review Comment: Both NUL-fallback tests use a single batch: this one (empty) and `test_nul_schema_strings_fall_back_without_executing_again`. Two more cases would cover the rest of the fallback's contract: (1) two non-empty batches, asserting `num_chunks == 2` and the values; (2) a NUL field name with mismatched batch schemas, expecting `ValueError`. Case (2) pins the schema check in `bridge.rs` (L60-66) ahead of the fallback branch (L68). -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
