From 234353b5e0350587d736ecdd21e1c90e97c0fd7f Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Mon, 26 Oct 2026 22:50:32 -0700 Subject: [PATCH 1/8] Extract shared Python and Arrow converter Share Arrow C Data handoff, scalar conversion, and explicit UDF conversion policy between Python-facing Spark consumers. Adopt the bridge in toArrow and shared scalar helpers in client input and row conversion. Generated-by: OpenAI Codex Co-authored-by: Isaac --- .github/workflows/coverage.yml | 5 +- Cargo.lock | 10 + Cargo.toml | 3 + crates/pyspark-rs/Cargo.toml | 1 + crates/pyspark-rs/src/dataframe.rs | 17 +- crates/pyspark-rs/src/row.rs | 32 +- crates/pyspark-rs/src/session.rs | 17 +- crates/spark-python-arrow/Cargo.toml | 16 + crates/spark-python-arrow/LICENSE | 201 ++++ crates/spark-python-arrow/README.md | 90 ++ crates/spark-python-arrow/src/bridge.rs | 546 ++++++++++ crates/spark-python-arrow/src/lib.rs | 11 + crates/spark-python-arrow/src/scalar.rs | 189 ++++ crates/spark-python-arrow/src/udf/input.rs | 437 ++++++++ crates/spark-python-arrow/src/udf/mod.rs | 91 ++ crates/spark-python-arrow/src/udf/result.rs | 881 ++++++++++++++++ crates/spark-python-arrow/tests/conversion.rs | 976 ++++++++++++++++++ dev/design/architecture.md | 18 +- python/tests/test_arrow_handoff.py | 506 +++++++++ 19 files changed, 4004 insertions(+), 43 deletions(-) create mode 100644 crates/spark-python-arrow/Cargo.toml create mode 100644 crates/spark-python-arrow/LICENSE 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/src/scalar.rs create mode 100644 crates/spark-python-arrow/src/udf/input.rs create mode 100644 crates/spark-python-arrow/src/udf/mod.rs create mode 100644 crates/spark-python-arrow/src/udf/result.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..88560f9c 100644 --- a/.github/workflows/coverage.yml +++ b/.github/workflows/coverage.yml @@ -92,6 +92,9 @@ jobs: python -m pip install maturin pytest coverage python -m pip install pandas pyarrow numpy grpcio grpcio-status googleapis-common-protos cloudpickle zstandard py4j + - name: Test standalone Python and Arrow converter + run: cargo test -p spark-python-arrow --locked --features pyo3/abi3-py39 + - name: Cache Spark ${{ env.SPARK_VERSION }} distribution uses: actions/cache@v4 with: @@ -226,7 +229,7 @@ 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/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..135d2b1a 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 = { version = "0.1.0", 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..c409c780 100644 --- a/crates/pyspark-rs/src/dataframe.rs +++ b/crates/pyspark-rs/src/dataframe.rs @@ -1166,19 +1166,14 @@ 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, without copying the decoded Arrow data buffers. #[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/pyspark-rs/src/row.rs b/crates/pyspark-rs/src/row.rs index 524bb73a..28332cd2 100644 --- a/crates/pyspark-rs/src/row.rs +++ b/crates/pyspark-rs/src/row.rs @@ -2,8 +2,9 @@ use pyo3::exceptions::{PyAttributeError, PyIndexError, PyKeyError}; use pyo3::prelude::*; -use pyo3::types::{PyBytes, PyDict, PyList}; +use pyo3::types::{PyDict, PyList}; use spark_connect::row::{Row, Value}; +use spark_python_arrow::scalar::Scalar; /// Python wrapper for a Spark Row. #[pyclass(name = "Row", module = "pyspark.sql.types")] @@ -23,24 +24,17 @@ pub fn value_to_py<'py>(py: Python<'py>, v: &Value) -> PyResult py.None().into_bound(py), - Value::Bool(b) => b.into_pyobject(py)?.to_owned().into_any(), - Value::Byte(n) => (*n as i64).into_pyobject(py)?.into_any(), - Value::Short(n) => (*n as i64).into_pyobject(py)?.into_any(), - Value::Integer(n) => (*n as i64).into_pyobject(py)?.into_any(), - Value::Long(n) => (*n).into_pyobject(py)?.into_any(), - Value::Float(f) => (*f as f64).into_pyobject(py)?.into_any(), - Value::Double(f) => (*f).into_pyobject(py)?.into_any(), - Value::String(s) => s.as_str().into_pyobject(py)?.into_any(), - Value::Binary(b) => PyBytes::new(py, b).into_any(), - Value::Date(d) => { - // Convert days-since-epoch (i32) to datetime.date - let date_mod = py.import("datetime")?; - let date_cls = date_mod.getattr("date")?; - // 719163 is the ordinal of 1970-01-01 - let ordinal = (*d as i64) + 719163i64; - date_cls.call_method1("fromordinal", (ordinal,))?.into_any() - } + Value::Null => Scalar::Null.into_python(py)?, + Value::Bool(b) => Scalar::Bool(*b).into_python(py)?, + Value::Byte(n) => Scalar::Int(i64::from(*n)).into_python(py)?, + Value::Short(n) => Scalar::Int(i64::from(*n)).into_python(py)?, + Value::Integer(n) => Scalar::Int(i64::from(*n)).into_python(py)?, + Value::Long(n) => Scalar::Int(*n).into_python(py)?, + Value::Float(f) => Scalar::Float(f64::from(*f)).into_python(py)?, + Value::Double(f) => Scalar::Float(*f).into_python(py)?, + Value::String(s) => Scalar::String(s).into_python(py)?, + Value::Binary(b) => Scalar::Binary(b).into_python(py)?, + Value::Date(d) => Scalar::Date(*d).into_python(py)?, Value::Timestamp(t) => { // Convert microseconds-since-epoch (i64) to datetime.datetime let datetime_mod = py.import("datetime")?; diff --git a/crates/pyspark-rs/src/session.rs b/crates/pyspark-rs/src/session.rs index 5a967bc3..1a3d049d 100644 --- a/crates/pyspark-rs/src/session.rs +++ b/crates/pyspark-rs/src/session.rs @@ -945,13 +945,15 @@ fn value_to_datatype(v: &Value) -> DataType { /// Convert a Python value to a Rust Value. pub(crate) fn py_to_value(obj: &Bound<'_, PyAny>) -> PyResult { + use spark_python_arrow::scalar; + // Check for None if obj.is_none() { return Ok(Value::Null); } // Check for bool (before int because bool is a subclass of int in Python) - if let Ok(b) = obj.extract::() { + if let Ok(b) = scalar::extract_bool(obj) { return Ok(Value::Bool(b)); } @@ -962,27 +964,24 @@ pub(crate) fn py_to_value(obj: &Bound<'_, PyAny>) -> PyResult { } // Check for int - if let Ok(i) = obj.extract::() { + if let Ok(i) = scalar::extract_i64(obj) { return Ok(Value::Long(i)); } // Check for float - if let Ok(f) = obj.extract::() { + if let Ok(f) = scalar::extract_f64(obj) { return Ok(Value::Double(f)); } // Check for str - if let Ok(s) = obj.extract::() { + if let Ok(s) = scalar::extract_string(obj) { return Ok(Value::String(s)); } // Check for bytes/bytearray SPECIFICALLY (not `extract::>()`, which also // matches a Python list of small ints like [1,2,3] and would mis-tag arrays as binary). - if let Ok(b) = obj.cast::() { - return Ok(Value::Binary(b.as_bytes().to_vec())); - } - if let Ok(ba) = obj.cast::() { - return Ok(Value::Binary(ba.to_vec())); + if let Some(bytes) = scalar::extract_binary(obj) { + return Ok(Value::Binary(bytes)); } // list / tuple -> array (recursively converted) diff --git a/crates/spark-python-arrow/Cargo.toml b/crates/spark-python-arrow/Cargo.toml new file mode 100644 index 00000000..e7383b6b --- /dev/null +++ b/crates/spark-python-arrow/Cargo.toml @@ -0,0 +1,16 @@ +# SPDX-License-Identifier: Apache-2.0 +[package] +name = "spark-python-arrow" +version = "0.1.0" +edition = "2021" +license = "Apache-2.0" +description = "Arrow and Python conversion shared by Spark clients and Python workers" +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, <0.30" } + +[lints.rust] +unsafe_op_in_unsafe_fn = "deny" diff --git a/crates/spark-python-arrow/LICENSE b/crates/spark-python-arrow/LICENSE new file mode 100644 index 00000000..261eeb9e --- /dev/null +++ b/crates/spark-python-arrow/LICENSE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed 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. diff --git a/crates/spark-python-arrow/README.md b/crates/spark-python-arrow/README.md new file mode 100644 index 00000000..5b206c21 --- /dev/null +++ b/crates/spark-python-arrow/README.md @@ -0,0 +1,90 @@ + + +# Shared Python and Arrow conversion + +`spark-python-arrow` is a Rust library for Python-facing Spark clients and workers. +It depends on Arrow and PyO3, not Spark Connect transport, generated protobufs, +an async runtime, or a worker executable. Consumers use the same versioned source; +they should not maintain copied converter modules. + +## Interfaces + +- `bridge`: import and export Arrow schemas, types, arrays, and record batches + through PyArrow's C Data Interface. Table export preserves batch boundaries and + shares decoded Arrow data buffers. Field names or timezone strings containing + NUL use a documented IPC fallback at the table boundary; low-level C exports + reject them. Length-prefixed schema metadata can contain NUL. +- `scalar`: borrowed scalar materialization and Python extraction. Exact-built-in + extraction is separate from coercive extraction so callers cannot accidentally + apply UDF coercion rules to client inputs. +- `udf`: the `SQL_ARROW_BATCHED_UDF` conversion policy. `input_spec` and + `result_spec` select supported native paths. `InputColumns` and `ResultColumn` + preserve the Python worker's conversion order, mutable-result handling, and + ArrowInvalid retry behavior. Python classes and converter callbacks come from + the caller; the library does not import a particular PySpark distribution. + `result_spec` resolves the caller's Variant class lazily, only after accepting + a native Variant conversion. + +Native UDF conversion covers signed integers, floats, booleans, strings, binary, +and tagged Variant structs, plus dates on input. Nested, decimal, timestamp, and +date-result columns retain the caller's Python fallback. Unsupported conversion +is `Ok(Err(reason))`; an error inspecting a supported conversion is `Err(PyErr)`. +Callers must not turn execution errors into unsupported-type fallback. + +## Ownership and caller responsibilities + +The library does not initialize Python, acquire a process-global interpreter, +create threads, or retain session state. The consumer chooses embedding versus +extension-module linking and owns all Python callbacks and their interpreter. +Array imports accept actual PyArrow objects and use their native exporters. +Bare extension types/arrays are rejected because Arrow's bare type cannot retain +their extension metadata. Schemas and record batches retain it on their fields; +UDF extension results select the caller's fallback. +As with the Arrow C Data Interface itself, producers are responsible for valid +buffer contents. Imports check structural invariants without scanning payloads. + +Zero-copy describes compatible Arrow **data buffers**, not schema descriptors, +Python scalar/row objects, pandas materialization, IPC decoding, or fallback +conversion. Import/export release callbacks retain the buffers after either +language's temporary wrapper has been dropped. + +The client's owned `Value` and `Row` types remain in the Python-free Rust client. +The Python binding adapts their scalar fields to `scalar::Scalar`; its row class, +schema inference, and temporal/nested compatibility rules remain caller-owned. +Worker execution, cancellation, transport, and session settings likewise stay +outside this crate. + +## Integration + +Cargo consumers use the crate as a normal Rust dependency. The library does not +enable PyO3's `extension-module` feature or select a minimum Python ABI: a Python +extension chooses those features in the host package, while an embedded worker +links its chosen Python runtime. +PyO3 must resolve to one compatible version in each consuming dependency graph. + +Bazel consumers should pin the OSS revision or release using their existing +third-party Rust dependency mechanism and depend on this library target. There +must be no internal copy of `src/bridge.rs`, `src/scalar.rs`, or `src/udf`. +An unpublished local checkout may be used for development, but is not a release +pin and must not become a checked-in machine-specific path. + +## 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 cover buffer identity and lifetime, schema/type round trips, +empty shapes, native scalar behavior, and UDF fallback/error ordering. The +Python client's `test_arrow_handoff.py` exercises the public `toArrow()` and +`toPandas()` paths. These are not substitutes for downstream worker differential +tests or 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..be8bc9a3 --- /dev/null +++ b/crates/spark-python-arrow/src/bridge.rs @@ -0,0 +1,546 @@ +// 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. + +//! Share Arrow buffers with PyArrow through the Arrow C Data Interface. +//! +//! Imports accept real PyArrow objects, not arbitrary objects implementing an +//! export method. PyArrow remains responsible for the validity of its C data. +//! These functions check structural invariants but do not scan values or change +//! when a caller observes conversion errors such as invalid string contents. + +use std::sync::Arc; + +use arrow::array::{make_array, Array, ArrayData, ArrayRef, StructArray}; +use arrow::datatypes::{DataType, Field, Schema}; +use arrow::error::ArrowError; +use arrow::ffi::{from_ffi, FFI_ArrowArray, FFI_ArrowSchema}; +use arrow::ipc::writer::FileWriter; +use arrow::record_batch::{RecordBatch, RecordBatchOptions}; +use pyo3::exceptions::{PyRuntimeError, PyTypeError, PyValueError}; +use pyo3::prelude::*; +use pyo3::types::{PyBytes, PyList}; + +/// Import a PyArrow schema, preserving schema and field metadata. +pub fn import_schema(value: &Bound<'_, PyAny>) -> PyResult { + let class = require_pyarrow_class(value, "Schema")?; + check_python_schema(value)?; + let mut schema = FFI_ArrowSchema::empty(); + let address = &mut schema as *mut FFI_ArrowSchema as usize; + class.call_method1("_export_to_c", (value, address))?; + Schema::try_from(&schema).map_err(arrow_error) +} + +/// Import a PyArrow data type. +/// +/// Root extension types require a schema or record batch: a bare Rust `DataType` +/// has nowhere to retain their extension metadata, so this function rejects them. +pub fn import_data_type(value: &Bound<'_, PyAny>) -> PyResult { + let class = require_pyarrow_class(value, "DataType")?; + reject_root_extension_type(value)?; + check_python_data_type(value)?; + let mut schema = FFI_ArrowSchema::empty(); + let address = &mut schema as *mut FFI_ArrowSchema as usize; + class.call_method1("_export_to_c", (value, address))?; + DataType::try_from(&schema).map_err(arrow_error) +} + +/// Import a PyArrow array without copying its data buffers. +/// +/// Root extension arrays are rejected because an `ArrayRef` cannot retain their +/// extension metadata. Import a record batch to preserve that field metadata. +pub fn import_array(value: &Bound<'_, PyAny>) -> PyResult { + let class = require_pyarrow_class(value, "Array")?; + let data_type = value.getattr("type")?; + reject_root_extension_type(&data_type)?; + check_python_data_type(&data_type)?; + let (array, schema) = export_python_array(value, &class)?; + // SAFETY: the base PyArrow exporter, not an arbitrary Python method, has + // initialized both descriptors and transferred ownership of valid C data. + let data = unsafe { from_ffi(array, &schema) }.map_err(arrow_error)?; + validate_structure(&data)?; + Ok(make_array(data)) +} + +/// Import a PyArrow record batch, preserving metadata and zero-column rows. +pub fn import_batch(value: &Bound<'_, PyAny>) -> PyResult { + let class = require_pyarrow_class(value, "RecordBatch")?; + check_python_schema(&value.getattr("schema")?)?; + let (array, ffi_schema) = export_python_array(value, &class)?; + let schema = Schema::try_from(&ffi_schema).map_err(arrow_error)?; + // SAFETY: only the real PyArrow RecordBatch exporter writes the descriptors. + let data = unsafe { from_ffi(array, &ffi_schema) }.map_err(arrow_error)?; + validate_structure(&data)?; + if !matches!(data.data_type(), DataType::Struct(_)) || data.null_count() != 0 { + return Err(PyValueError::new_err( + "A record batch must be a struct array without top-level nulls", + )); + } + let array = make_array(data); + let array = array + .as_any() + .downcast_ref::() + .ok_or_else(|| PyValueError::new_err("A record batch must be a struct array"))?; + let options = RecordBatchOptions::new().with_row_count(Some(array.len())); + // Unlike RecordBatch::from(StructArray), this validates non-nullable columns + // and retains the schema-level metadata from the exported C schema. + RecordBatch::try_new_with_options(Arc::new(schema), array.columns().to_vec(), &options) + .map_err(arrow_error) +} + +/// Export an Arrow array without copying its data buffers. +pub fn export_array(py: Python<'_>, array: ArrayRef) -> PyResult> { + let array_type = py.import("pyarrow")?.getattr("Array")?; + Ok(export_array_with_type(&array_type, array)?.unbind()) +} + +pub(crate) fn export_array_with_type<'py>( + array_type: &Bound<'py, PyAny>, + array: ArrayRef, +) -> PyResult> { + check_c_data_type(array.data_type())?; + let mut schema = FFI_ArrowSchema::try_from(array.data_type()).map_err(arrow_error)?; + let schema_address = &mut schema as *mut FFI_ArrowSchema as usize; + let mut array = FFI_ArrowArray::new(&array.to_data()); + let array_address = &mut array as *mut FFI_ArrowArray as usize; + // Both descriptors are owned here until the importer consumes them. Their + // release callbacks also cover importer errors before either transfer. + array_type.call_method1("_import_from_c", (array_address, schema_address)) +} + +/// Export an Arrow schema, preserving schema and field metadata. +/// +/// C schema names and timezone strings cannot contain NUL. This low-level API +/// returns an error for such schemas; table export provides an IPC fallback. +pub fn export_schema(py: Python<'_>, schema: &Schema) -> PyResult> { + check_c_schema(schema)?; + let mut schema = FFI_ArrowSchema::try_from(schema).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()) +} + +/// Export an Arrow data type, rejecting strings not representable in C schemas. +pub fn export_data_type(py: Python<'_>, data_type: &DataType) -> PyResult> { + check_c_data_type(data_type)?; + let mut schema = FFI_ArrowSchema::try_from(data_type).map_err(arrow_error)?; + let address = &mut schema as *mut FFI_ArrowSchema as usize; + Ok(py + .import("pyarrow")? + .getattr("DataType")? + .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. +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| 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()) +} + +pub(crate) fn arrow_error(error: ArrowError) -> PyErr { + PyRuntimeError::new_err(error.to_string()) +} + +pub(crate) fn is_extension_type(data_type: &Bound<'_, PyAny>) -> PyResult { + let class = data_type + .py() + .import("pyarrow.lib")? + .getattr("BaseExtensionType")?; + data_type.is_instance(&class) +} + +fn reject_root_extension_type(data_type: &Bound<'_, PyAny>) -> PyResult<()> { + if is_extension_type(data_type)? { + return Err(PyTypeError::new_err( + "Arrow extension types require a schema or record batch to preserve metadata", + )); + } + Ok(()) +} + +fn require_pyarrow_class<'py>( + value: &Bound<'py, PyAny>, + name: &str, +) -> PyResult> { + // Use the extension's classes, not replaceable pyarrow package aliases. + // Calling their descriptors also bypasses a subclass's Python override. + let class = value.py().import("pyarrow.lib")?.getattr(name)?; + if !value.is_instance(&class)? { + return Err(PyTypeError::new_err(format!("Expected a pyarrow.{name}"))); + } + Ok(class) +} + +fn export_python_array( + value: &Bound<'_, PyAny>, + class: &Bound<'_, PyAny>, +) -> PyResult<(FFI_ArrowArray, FFI_ArrowSchema)> { + let mut array = FFI_ArrowArray::empty(); + let mut schema = FFI_ArrowSchema::empty(); + let array_address = &mut array as *mut FFI_ArrowArray as usize; + let schema_address = &mut schema as *mut FFI_ArrowSchema as usize; + class.call_method1("_export_to_c", (value, array_address, schema_address))?; + Ok((array, schema)) +} + +fn validate_structure(data: &ArrayData) -> PyResult<()> { + // Arrow's inexpensive validation recurses through child descriptors. It + // checks sizes and offset endpoints, not all offsets, UTF-8 contents, or + // dictionary indices whose conversion errors callers may defer. + data.validate().map_err(arrow_error) +} + +fn nul_schema_error() -> PyErr { + PyValueError::new_err("Arrow C schema names and timezones cannot contain NUL") +} + +fn check_c_schema(schema: &Schema) -> PyResult<()> { + if schema.fields().iter().any(|field| field_needs_ipc(field)) { + return Err(nul_schema_error()); + } + Ok(()) +} + +fn check_c_data_type(data_type: &DataType) -> PyResult<()> { + if data_type_needs_ipc(data_type) { + return Err(nul_schema_error()); + } + Ok(()) +} + +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, + } +} + +fn check_python_schema(schema: &Bound<'_, PyAny>) -> PyResult<()> { + for index in 0..schema.len()? { + check_python_field(&schema.call_method1("field", (index,))?)?; + } + Ok(()) +} + +fn check_python_field(field: &Bound<'_, PyAny>) -> PyResult<()> { + if field.getattr("name")?.extract::()?.contains('\0') { + return Err(nul_schema_error()); + } + check_python_data_type(&field.getattr("type")?) +} + +fn check_python_data_type(data_type: &Bound<'_, PyAny>) -> PyResult<()> { + // Check the original Python strings before C export can truncate at NUL. + // Metadata uses length-prefixed strings and does not need this restriction. + if let Ok(timezone) = data_type.getattr("tz") { + if !timezone.is_none() && timezone.extract::()?.contains('\0') { + return Err(nul_schema_error()); + } + } + let num_fields = data_type.getattr("num_fields")?.extract::()?; + for index in 0..num_fields { + check_python_field(&data_type.call_method1("field", (index,))?)?; + } + // Dictionary and extension value/storage types are not child fields. + // Lists expose both a child field and value_type; do not visit those twice. + if num_fields == 0 { + for name in ["value_type", "storage_type"] { + if let Ok(child) = data_type.getattr(name) { + check_python_data_type(&child)?; + } + } + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + use std::ffi::CStr; + + use arrow::array::Int64Array; + use arrow::datatypes::TimeUnit; + use pyo3::types::PyDict; + + use super::*; + + fn eval<'py>(py: Python<'py>, code: &CStr) -> Bound<'py, PyAny> { + let globals = PyDict::new(py); + globals + .set_item("pa", py.import("pyarrow").unwrap()) + .unwrap(); + py.eval(code, Some(&globals), None).unwrap() + } + + #[test] + fn schema_roundtrip_preserves_metadata() { + Python::initialize(); + Python::attach(|py| { + let field = Field::new("value", DataType::Int64, true) + .with_metadata(HashMap::from([("field\0key".into(), "value\0".into())])); + let schema = Schema::new(vec![field]) + .with_metadata(HashMap::from([("schema".into(), "metadata".into())])); + let exported = export_schema(py, &schema).unwrap(); + assert_eq!(import_schema(exported.bind(py)).unwrap(), schema); + }); + } + + #[test] + fn array_roundtrip_shares_buffers_and_preserves_lifetime() { + Python::initialize(); + Python::attach(|py| { + let values: ArrayRef = Arc::new(Int64Array::from(vec![Some(1), None, Some(3)])); + let address = values.to_data().buffers()[0].as_ptr(); + let exported = export_array(py, values).unwrap(); + let imported = import_array(exported.bind(py)).unwrap(); + drop(exported); + assert_eq!(imported.to_data().buffers()[0].as_ptr(), address); + assert_eq!(imported.len(), 3); + assert!(imported.is_null(1)); + assert_eq!( + imported + .as_any() + .downcast_ref::() + .unwrap() + .value(2), + 3 + ); + }); + } + + #[test] + fn imports_preserve_slices_and_batch_metadata() { + Python::initialize(); + Python::attach(|py| { + let batch = eval( + py, + c"pa.RecordBatch.from_arrays( + [pa.array([1, None, 3])], + schema=pa.schema( + [pa.field('v', pa.int64(), metadata={b'f': b'v'})], + metadata={b's': b'v'} + ) + ).slice(1, 2)", + ); + let imported = import_batch(&batch).unwrap(); + assert_eq!(imported.num_rows(), 2); + assert!(imported.column(0).is_null(0)); + assert_eq!(imported.schema().metadata().get("s").unwrap(), "v"); + assert_eq!(imported.schema().field(0).metadata().get("f").unwrap(), "v"); + let table = record_batches_to_pyarrow(py, vec![imported]).unwrap(); + assert!(table + .bind(py) + .getattr("schema") + .unwrap() + .call_method1("equals", (batch.getattr("schema").unwrap(), true)) + .unwrap() + .is_truthy() + .unwrap()); + }); + } + + #[test] + fn imports_preserve_zero_column_rows_and_empty_schema_metadata() { + Python::initialize(); + Python::attach(|py| { + for length in [0, 3] { + let array = StructArray::new_empty_fields(length, None); + let batch = RecordBatch::from(array); + let schema = batch + .schema() + .as_ref() + .clone() + .with_metadata(HashMap::from([("schema".into(), "empty".into())])); + let batch = batch.with_schema(Arc::new(schema)).unwrap(); + let table = record_batches_to_pyarrow(py, vec![batch.clone()]).unwrap(); + let batches = table.bind(py).call_method0("to_batches").unwrap(); + if length == 0 { + assert_eq!( + table + .bind(py) + .getattr("num_rows") + .unwrap() + .extract::() + .unwrap(), + 0 + ); + assert_eq!( + import_schema(&table.bind(py).getattr("schema").unwrap()).unwrap(), + *batch.schema() + ); + } else { + let imported = import_batch(&batches.get_item(0).unwrap()).unwrap(); + assert_eq!(imported.num_rows(), length); + assert_eq!(imported.num_columns(), 0); + assert_eq!(imported.schema(), batch.schema()); + } + } + }); + } + + #[test] + fn imports_reject_wrong_pyarrow_classes() { + Python::initialize(); + Python::attach(|py| { + let array = eval(py, c"pa.array([1, 2])"); + assert!(import_schema(&array) + .unwrap_err() + .is_instance_of::(py)); + assert!(import_data_type(&array) + .unwrap_err() + .is_instance_of::(py)); + assert!(import_batch(&array) + .unwrap_err() + .is_instance_of::(py)); + let schema = eval(py, c"pa.schema([('v', pa.int64())])"); + assert!(import_array(&schema) + .unwrap_err() + .is_instance_of::(py)); + let fake = eval( + py, + c"type('Fake', (), {'_export_to_c': lambda *args: 1 / 0})()", + ); + assert!(import_schema(&fake) + .unwrap_err() + .is_instance_of::(py)); + assert!(import_data_type(&fake) + .unwrap_err() + .is_instance_of::(py)); + assert!(import_array(&fake) + .unwrap_err() + .is_instance_of::(py)); + assert!(import_batch(&fake) + .unwrap_err() + .is_instance_of::(py)); + }); + } + + #[test] + fn batch_import_rejects_nulls_in_non_nullable_columns() { + Python::initialize(); + Python::attach(|py| { + let batch = eval( + py, + c"pa.RecordBatch.from_arrays( + [pa.array([None], type=pa.int64())], + schema=pa.schema([pa.field('v', pa.int64(), nullable=False)]) + )", + ); + assert!(import_batch(&batch) + .unwrap_err() + .to_string() + .contains("non-nullable")); + }); + } + + #[test] + fn nul_schema_strings_are_rejected_without_panicking() { + let field = Field::new("bad\0name", DataType::Int64, true); + assert!(check_c_schema(&Schema::new(vec![field.clone()])).is_err()); + assert!(check_c_data_type(&DataType::List(Arc::new(field))).is_err()); + let timestamp = DataType::Timestamp(TimeUnit::Microsecond, Some("UTC\0suffix".into())); + assert!(check_c_data_type(×tamp).is_err()); + assert!(check_c_data_type(&DataType::Dictionary( + Box::new(DataType::Int8), + Box::new(timestamp) + )) + .is_err()); + Python::initialize(); + Python::attach(|py| { + assert!(export_data_type(py, &DataType::Time32(TimeUnit::Nanosecond)).is_err()); + let schema = eval(py, c"pa.schema([('bad\\x00name', pa.int64())])"); + assert!(import_schema(&schema) + .unwrap_err() + .is_instance_of::(py)); + let data_type = eval(py, c"pa.timestamp('us', tz='UTC\\x00suffix')"); + assert!(import_data_type(&data_type) + .unwrap_err() + .is_instance_of::(py)); + }); + } +} diff --git a/crates/spark-python-arrow/src/lib.rs b/crates/spark-python-arrow/src/lib.rs new file mode 100644 index 00000000..8c2045f9 --- /dev/null +++ b/crates/spark-python-arrow/src/lib.rs @@ -0,0 +1,11 @@ +// SPDX-License-Identifier: Apache-2.0 + +//! Python/Arrow conversion without a dependency on Spark transport or worker execution. +//! +//! The bridge shares Arrow data buffers; scalar and row materialization still allocates Python +//! objects. The caller owns its interpreter, Python classes, session settings, and error policy. +//! UDF-specific coercion and end-of-batch result conversion are explicit in the [`udf`] API. + +pub mod bridge; +pub mod scalar; +pub mod udf; diff --git a/crates/spark-python-arrow/src/scalar.rs b/crates/spark-python-arrow/src/scalar.rs new file mode 100644 index 00000000..8b287428 --- /dev/null +++ b/crates/spark-python-arrow/src/scalar.rs @@ -0,0 +1,189 @@ +/* + * 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. + */ + +//! Scalar conversion shared by client row materialization and UDF columns. +//! +//! Exact extraction never invokes a subclass's coercion methods. The separate +//! `extract_*` helpers retain PyO3's conversion semantics; callers choose their order, notably +//! checking booleans before integers and Decimal before numeric coercion. This module does not +//! infer a schema or define nested, decimal, timestamp, or application-specific row types. + +use std::borrow::Cow; + +use pyo3::prelude::*; +use pyo3::types::{PyBool, PyByteArray, PyBytes, PyFloat, PyInt, PyString}; +use pyo3::IntoPyObjectExt; + +/// The ordinal of 1970-01-01 in Python's `datetime.date` calendar. +const EPOCH_ORDINAL: i64 = 719_163; + +/// A borrowed scalar, independent of a client's owned value representation. +#[derive(Clone, Copy, Debug, PartialEq)] +pub enum Scalar<'a> { + Null, + Bool(bool), + Int(i64), + Float(f64), + String(&'a str), + /// Materialized as immutable Python `bytes`. + Binary(&'a [u8]), + /// Days since 1970-01-01, materialized as `datetime.date`. + Date(i32), +} + +impl Scalar<'_> { + /// Materializes this scalar in the caller's attached interpreter. + pub fn into_python<'py>(self, py: Python<'py>) -> PyResult> { + match self { + Self::Null => Ok(py.None().into_bound(py)), + Self::Bool(value) => value.into_bound_py_any(py), + Self::Int(value) => value.into_bound_py_any(py), + Self::Float(value) => value.into_bound_py_any(py), + Self::String(value) => utf8_into_python(py, value.as_bytes()), + Self::Binary(value) => Ok(PyBytes::new(py, value).into_any()), + Self::Date(days) => { + let fromordinal = py + .import("datetime")? + .getattr("date")? + .getattr("fromordinal")?; + date_into_python(days, &fromordinal) + } + } + } +} + +/// Decodes UTF-8 with Python's strict decoder, preserving its exception on invalid bytes. +/// +/// This accepts raw bytes because an Arrow C Data import can contain unvalidated string data; +/// forming a Rust `&str` first would be unsound for such an input. +pub fn utf8_into_python<'py>(py: Python<'py>, bytes: &[u8]) -> PyResult> { + Ok(PyString::from_bytes(py, bytes)?.into_any()) +} + +/// Whether days since 1970-01-01 fall within Python's date range (years 1 through 9999). +pub fn is_date_in_range(days: i32) -> bool { + (-719_162..=2_932_896).contains(&days) +} + +/// Materializes a date using an explicitly supplied, optionally cached `date.fromordinal`. +/// +/// The constructor defines out-of-range errors. A caller reproducing another conversion path's +/// errors can check [`is_date_in_range`] before invoking it. +pub fn date_into_python<'py>( + days: i32, + fromordinal: &Bound<'py, PyAny>, +) -> PyResult> { + fromordinal.call1((i64::from(days) + EPOCH_ORDINAL,)) +} + +/// An exact built-in boolean, excluding truthiness coercion. +pub fn exact_bool(value: &Bound<'_, PyAny>) -> Option { + value + .is_exact_instance_of::() + .then(|| value.extract().ok()) + .flatten() +} + +/// An exact built-in integer fitting in 64 bits, excluding booleans and integer subclasses. +pub fn exact_i64(value: &Bound<'_, PyAny>) -> Option { + value + .is_exact_instance_of::() + .then(|| value.extract().ok()) + .flatten() +} + +/// An exact built-in float, without integer or user-defined numeric coercion. +pub fn exact_f64(value: &Bound<'_, PyAny>) -> Option { + value + .is_exact_instance_of::() + .then(|| value.extract().ok()) + .flatten() +} + +/// An exact built-in string with a valid UTF-8 representation. +/// +/// PyO3 uses a supported ABI conversion and may allocate when the selected limited API cannot +/// borrow CPython's UTF-8 cache. Lone surrogates return `None`, leaving error/fallback policy to +/// the caller, rather than replacing characters or retaining a Python error indicator. +pub fn exact_str<'a>(value: &'a Bound<'_, PyAny>) -> Option> { + if !value.is_exact_instance_of::() { + return None; + } + value.cast::().ok()?.to_cow().ok() +} + +/// The borrowed contents of exact built-in `bytes`, excluding subclasses and mutable buffers. +pub fn exact_bytes<'a>(value: &'a Bound<'_, PyAny>) -> Option<&'a [u8]> { + if !value.is_exact_instance_of::() { + return None; + } + value.cast::().ok().map(|bytes| bytes.as_bytes()) +} + +/// Retains `value` only when it is exact built-in `bytes`. +pub fn into_exact_bytes(value: Bound<'_, PyAny>) -> Option> { + if !value.is_exact_instance_of::() { + return None; + } + value.cast_into::().ok() +} + +/// Extracts with PyO3's boolean rules, not Python truthiness. +pub fn extract_bool(value: &Bound<'_, PyAny>) -> PyResult { + value.extract() +} + +/// Extracts with PyO3's integer rules, including accepted integer-like objects. +pub fn extract_i64(value: &Bound<'_, PyAny>) -> PyResult { + value.extract() +} + +/// Extracts with PyO3's float rules, including numeric coercion. +pub fn extract_f64(value: &Bound<'_, PyAny>) -> PyResult { + value.extract() +} + +/// Extracts with PyO3's string rules, including string subclasses but not arbitrary `str()`. +pub fn extract_string(value: &Bound<'_, PyAny>) -> PyResult { + value.extract() +} + +/// Copies bytes or bytearray contents, including subclasses, but never a list of integers. +pub fn extract_binary(value: &Bound<'_, PyAny>) -> Option> { + if let Ok(bytes) = value.cast::() { + Some(bytes.as_bytes().to_vec()) + } else if let Ok(bytes) = value.cast::() { + Some(bytes.to_vec()) + } else { + None + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn date_range_matches_python_calendar_boundaries() { + assert!(is_date_in_range(-719_162)); + assert!(is_date_in_range(2_932_896)); + assert!(!is_date_in_range(-719_163)); + assert!(!is_date_in_range(2_932_897)); + assert!(!is_date_in_range(i32::MIN)); + assert!(!is_date_in_range(i32::MAX)); + } +} diff --git a/crates/spark-python-arrow/src/udf/input.rs b/crates/spark-python-arrow/src/udf/input.rs new file mode 100644 index 00000000..ce54bb2d --- /dev/null +++ b/crates/spark-python-arrow/src/udf/input.rs @@ -0,0 +1,437 @@ +/* + * 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 arrow::array::{ + Array, ArrayRef, AsArray, BinaryArray, BooleanArray, Date32Array, Float32Array, Float64Array, + GenericBinaryArray, GenericStringArray, Int16Array, Int32Array, Int64Array, Int8Array, + LargeBinaryArray, LargeStringArray, OffsetSizeTrait, StringArray, StructArray, +}; +use arrow::datatypes::DataType; +use pyo3::exceptions::{PyIndexError, PyRuntimeError, PyUnicodeDecodeError}; +use pyo3::prelude::*; +use pyo3::types::{PyByteArray, PyBytes, PyDate, PyDict}; +use pyo3::{intern, PyTypeInfo}; + +use crate::scalar::{date_into_python, is_date_in_range, utf8_into_python, Scalar}; +use crate::udf::{import_batch, is_spark_type, variant_fields}; + +/// Accepts an input column of `data_type`, the Arrow type of `spark_type`, given PySpark's input +/// converter for it or `None` and the backend package's `sql.types` module; or returns the reason +/// the native UDF input converter does not support it. +pub fn input_spec( + data_type: &DataType, + spark_type: &Bound<'_, PyAny>, + converter: &Bound<'_, PyAny>, + sql_types: &Bound<'_, PyAny>, +) -> PyResult> { + let py = spark_type.py(); + let read = if converter.is_none() { + if !Column::reads(data_type) { + return Ok(Err(format!( + "{data_type} is not supported by the native UDF input converter" + ))); + } + Read::Plain + } else if matches!(data_type, DataType::Binary | DataType::LargeBinary) + && is_spark_type(spark_type, sql_types, "BinaryType")? + { + // PySpark's BinaryType converter returns the bytes PyArrow gives, or a bytearray copy + // when spark.sql.execution.pyspark.binaryAsBytes is off. + let sample = converter.call1((PyBytes::new(py, b""),))?; + if sample.is_exact_instance_of::() { + Read::Plain + } else if sample.is_exact_instance_of::() { + Read::ByteArray + } else { + return Ok(Err(format!( + "{spark_type}'s converter gives {}", + sample.get_type() + ))); + } + } else if is_variant_struct(data_type) && is_spark_type(spark_type, sql_types, "VariantType")? { + // The converter can return a VariantVal whose class differs from the backend module's. + // Take the class from its output to preserve that identity. + let sample = PyDict::new(py); + sample.set_item("value", PyBytes::new(py, b""))?; + sample.set_item("metadata", PyBytes::new(py, b""))?; + Read::Variant { + class: converter.call1((sample,))?.get_type().into_any().unbind(), + } + } else { + return Ok(Err(format!("{spark_type} needs PySpark's converter"))); + }; + Ok(Ok(InputSpec { + data_type: data_type.clone(), + read, + converter: converter.clone().unbind(), + })) +} + +/// Whether `data_type` is the struct Spark gives `VariantType`, with parts of either width. +fn is_variant_struct(data_type: &DataType) -> bool { + [DataType::Binary, DataType::LargeBinary] + .into_iter() + .any(|part| *data_type == DataType::Struct(variant_fields(part))) +} + +/// How [`InputColumns`] reads a column [`input_spec`] accepted. +pub struct InputSpec { + data_type: DataType, + read: Read, + /// PySpark's input converter for the column, or `None`. + converter: Py, +} + +impl InputSpec { + fn data_type(&self) -> &DataType { + &self.data_type + } + + pub fn clone_ref(&self, py: Python<'_>) -> InputSpec { + InputSpec { + data_type: self.data_type.clone(), + read: match &self.read { + Read::Plain => Read::Plain, + Read::ByteArray => Read::ByteArray, + Read::Variant { class } => Read::Variant { + class: class.clone_ref(py), + }, + }, + converter: self.converter.clone_ref(py), + } + } +} + +/// What a column's values become, as PySpark makes them. +pub(crate) enum Read { + /// What PyArrow's `to_pylist` gives, which PySpark keeps. + Plain, + /// A `bytearray` copy of each binary value. + ByteArray, + /// A VariantVal of `class` for each struct of the variant's two binaries. + Variant { class: Py }, +} + +/// A batch's input columns, read one cell at a time. +pub struct InputColumns { + batch: Py, + columns: Vec, + converters: Vec>, + convert_input: Py, + rows: usize, +} + +impl InputColumns { + /// Imports `batch`, a `pa.RecordBatch`, without copying it. Its columns must have the types + /// of `specs`. `convert_input` is PySpark's input step, `_convert_arrow_udf_input`, which + /// [`InputColumns::value`] runs to raise PySpark's error. + pub fn import( + batch: &Bound<'_, PyAny>, + specs: &[InputSpec], + convert_input: &Bound<'_, PyAny>, + ) -> PyResult { + let py = batch.py(); + let input = import_batch(batch)?; + let actual: Vec<&DataType> = input.columns().iter().map(|c| c.data_type()).collect(); + if actual.len() != specs.len() + || actual + .iter() + .zip(specs) + .any(|(actual, spec)| **actual != spec.data_type) + { + let expected: Vec<&DataType> = specs.iter().map(InputSpec::data_type).collect(); + return Err(PyRuntimeError::new_err(format!( + "the batch has types {actual:?}, but {expected:?} were expected" + ))); + } + let columns = input + .columns() + .iter() + .zip(specs) + .map(|(array, spec)| { + Column::new(py, array, &spec.read)?.ok_or_else(|| { + PyRuntimeError::new_err(format!("no reader for {}", array.data_type())) + }) + }) + .collect::>>()?; + Ok(InputColumns { + batch: batch.clone().unbind(), + columns, + converters: specs + .iter() + .map(|spec| spec.converter.clone_ref(py)) + .collect(), + convert_input: convert_input.clone().unbind(), + rows: input.len(), + }) + } + + pub fn num_rows(&self) -> usize { + self.rows + } + + /// What PySpark gives a UDF for `row` of `column`, as a new object. `column` indexes the + /// imported types, and `row` is below [`InputColumns::num_rows`]. + /// + /// A cell PySpark fails to convert, such as a string that is not UTF-8, raises what PySpark's + /// input step raises for its column. PySpark converts every input column, in order, before it + /// calls a UDF, so to fail as it does, read every cell, column by column in order, before + /// calling any UDF. + pub fn value<'py>( + &self, + py: Python<'py>, + column: usize, + row: usize, + ) -> PyResult> { + let Some(input) = self.columns.get(column) else { + return Err(PyIndexError::new_err(format!( + "column index {column} is out of bounds for {} columns", + self.columns.len() + ))); + }; + if row >= self.rows { + return Err(PyIndexError::new_err(format!( + "row index {row} is out of bounds for {} rows", + self.rows + ))); + } + input.value(py, row, || self.pyspark_error(py, column)) + } + + /// What PySpark's input step raises for `index`, a column with a cell it fails to convert. + fn pyspark_error(&self, py: Python<'_>, index: usize) -> PyErr { + let converted = (|| -> PyResult<_> { + let column = self + .batch + .bind(py) + .call_method1(intern!(py, "column"), (index,))?; + self.convert_input + .bind(py) + .call1((column, self.converters[index].bind(py))) + })(); + match converted { + Err(error) => error, + Ok(_) => PyRuntimeError::new_err(format!( + "PySpark converted input {index}, which the native converter cannot read" + )), + } + } +} + +/// A column, read as PySpark reads it. +pub(crate) enum Column { + Int8(Int8Array), + Int16(Int16Array), + Int32(Int32Array), + Int64(Int64Array), + Float32(Float32Array), + Float64(Float64Array), + Boolean(BooleanArray), + Utf8(StringArray), + LargeUtf8(LargeStringArray), + /// With `datetime.date.fromordinal`. + Date32(Date32Array, Py), + /// With whether each value is read as a `bytearray`. + Binary(BinaryArray, bool), + LargeBinary(LargeBinaryArray, bool), + Variant(Box>), + LargeVariant(Box>), +} + +/// A variant column, with its two parts and the class of VariantVal its values are made of. +pub(crate) struct VariantColumn { + array: StructArray, + value: GenericBinaryArray, + metadata: GenericBinaryArray, + class: Py, +} + +impl VariantColumn { + fn new(py: Python<'_>, array: &ArrayRef, class: &Py) -> Self { + let array = array.as_struct().clone(); + VariantColumn { + value: array.column(0).as_binary().clone(), + metadata: array.column(1).as_binary().clone(), + array, + class: class.clone_ref(py), + } + } + + fn value<'py>( + &self, + py: Python<'py>, + row: usize, + unconvertible: impl FnOnce() -> PyErr, + ) -> PyResult> { + if self.array.is_null(row) { + return Ok(py.None().into_bound(py)); + } + if self.value.is_null(row) || self.metadata.is_null(row) { + return Err(unconvertible()); + } + self.class.bind(py).call1(( + PyBytes::new(py, self.value.value(row)), + PyBytes::new(py, self.metadata.value(row)), + )) + } +} + +impl Column { + /// The types this native reader supports without an additional PySpark converter. + fn reads(data_type: &DataType) -> bool { + matches!( + data_type, + DataType::Int8 + | DataType::Int16 + | DataType::Int32 + | DataType::Int64 + | DataType::Float32 + | DataType::Float64 + | DataType::Boolean + | DataType::Utf8 + | DataType::LargeUtf8 + | DataType::Date32 + ) + } + + pub(crate) fn new(py: Python<'_>, array: &ArrayRef, read: &Read) -> PyResult> { + Ok(Some(match (array.data_type(), read) { + (DataType::Int8, Read::Plain) => Column::Int8(array.as_primitive().clone()), + (DataType::Int16, Read::Plain) => Column::Int16(array.as_primitive().clone()), + (DataType::Int32, Read::Plain) => Column::Int32(array.as_primitive().clone()), + (DataType::Int64, Read::Plain) => Column::Int64(array.as_primitive().clone()), + (DataType::Float32, Read::Plain) => Column::Float32(array.as_primitive().clone()), + (DataType::Float64, Read::Plain) => Column::Float64(array.as_primitive().clone()), + (DataType::Boolean, Read::Plain) => Column::Boolean(array.as_boolean().clone()), + (DataType::Utf8, Read::Plain) => Column::Utf8(array.as_string().clone()), + (DataType::LargeUtf8, Read::Plain) => Column::LargeUtf8(array.as_string().clone()), + (DataType::Date32, Read::Plain) => Column::Date32( + array.as_primitive().clone(), + PyDate::type_object(py).getattr("fromordinal")?.unbind(), + ), + (DataType::Binary, Read::Plain | Read::ByteArray) => { + Column::Binary(array.as_binary().clone(), matches!(read, Read::ByteArray)) + } + (DataType::LargeBinary, Read::Plain | Read::ByteArray) => { + Column::LargeBinary(array.as_binary().clone(), matches!(read, Read::ByteArray)) + } + (DataType::Struct(fields), Read::Variant { class }) + if *fields == variant_fields(DataType::Binary) => + { + Column::Variant(Box::new(VariantColumn::new(py, array, class))) + } + (DataType::Struct(fields), Read::Variant { class }) + if *fields == variant_fields(DataType::LargeBinary) => + { + Column::LargeVariant(Box::new(VariantColumn::new(py, array, class))) + } + _ => return Ok(None), + })) + } + + /// A cell PySpark fails to convert raises `unconvertible()`: a string that is not UTF-8, a date + /// outside `datetime.date`'s range, or a variant with a null part. + pub(crate) fn value<'py>( + &self, + py: Python<'py>, + row: usize, + unconvertible: impl FnOnce() -> PyErr, + ) -> PyResult> { + fn read<'a, 'py, A: Array>( + py: Python<'py>, + array: &A, + row: usize, + value: impl FnOnce() -> Scalar<'a>, + ) -> PyResult> { + if array.is_null(row) { + Scalar::Null.into_python(py) + } else { + value().into_python(py) + } + } + /// Decodes the slot's bytes rather than reading a `&str`: arrow-rs imports strings + /// through the C Data Interface without checking them, so CPython's decode is the check. + fn string<'py, O: OffsetSizeTrait>( + py: Python<'py>, + array: &GenericStringArray, + row: usize, + unconvertible: impl FnOnce() -> PyErr, + ) -> PyResult> { + if array.is_null(row) { + return Ok(py.None().into_bound(py)); + } + let offsets = array.value_offsets(); + let bytes = &array.value_data()[offsets[row].as_usize()..offsets[row + 1].as_usize()]; + match utf8_into_python(py, bytes) { + Ok(string) => Ok(string), + Err(error) if error.is_instance_of::(py) => { + Err(unconvertible()) + } + Err(error) => Err(error), + } + } + fn binary<'py, O: OffsetSizeTrait>( + py: Python<'py>, + array: &GenericBinaryArray, + row: usize, + bytearray: bool, + ) -> PyResult> { + if array.is_null(row) { + return Scalar::Null.into_python(py); + } + let bytes = array.value(row); + if bytearray { + Ok(PyByteArray::new(py, bytes).into_any()) + } else { + Scalar::Binary(bytes).into_python(py) + } + } + match self { + Column::Int8(array) => { + read(py, array, row, || Scalar::Int(i64::from(array.value(row)))) + } + Column::Int16(array) => { + read(py, array, row, || Scalar::Int(i64::from(array.value(row)))) + } + Column::Int32(array) => { + read(py, array, row, || Scalar::Int(i64::from(array.value(row)))) + } + Column::Int64(array) => read(py, array, row, || Scalar::Int(array.value(row))), + Column::Float32(array) => read(py, array, row, || { + Scalar::Float(f64::from(array.value(row))) + }), + Column::Float64(array) => read(py, array, row, || Scalar::Float(array.value(row))), + Column::Boolean(array) => read(py, array, row, || Scalar::Bool(array.value(row))), + Column::Utf8(array) => string(py, array, row, unconvertible), + Column::LargeUtf8(array) => string(py, array, row, unconvertible), + Column::Date32(array, fromordinal) => { + if array.is_null(row) { + return Ok(py.None().into_bound(py)); + } + let days = array.value(row); + if !is_date_in_range(days) { + return Err(unconvertible()); + } + date_into_python(days, fromordinal.bind(py)) + } + Column::Binary(array, bytearray) => binary(py, array, row, *bytearray), + Column::LargeBinary(array, bytearray) => binary(py, array, row, *bytearray), + Column::Variant(variant) => variant.value(py, row, unconvertible), + Column::LargeVariant(variant) => variant.value(py, row, unconvertible), + } + } +} diff --git a/crates/spark-python-arrow/src/udf/mod.rs b/crates/spark-python-arrow/src/udf/mod.rs new file mode 100644 index 00000000..dc481466 --- /dev/null +++ b/crates/spark-python-arrow/src/udf/mod.rs @@ -0,0 +1,91 @@ +/* + * 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 <-> Python conversion one value at a time, for a per-row UDF loop that has to match the +//! `SQL_ARROW_BATCHED_UDF` loop in `pyspark.worker.read_udfs`. +//! +//! [`InputColumns`] reads a batch's input columns cell by cell, as PySpark's input step, +//! `_convert_arrow_udf_input`, reads them: `ArrowTableToRowsConversion._to_pylist`, then the +//! column's converter. [`ResultColumn`] builds one UDF's results into an Arrow array, as PySpark's +//! output step does: `pa.array(values, type=T)`, retried as +//! `pa.array(values).cast(T, safe=safecheck)` on `ArrowInvalid`. It appends built-in results, and +//! variant results whose parts are `bytes`, to an arrow-rs builder and hands any other result to +//! PyArrow. A caller keeps no list of types of its own: it asks [`input_spec`] about each input +//! column and [`result_spec`] about each result. +//! +//! A spec returns `Ok(Err(reason))` for an unsupported conversion; the caller may select its +//! existing Python implementation in that case. `Err(PyErr)` instead means inspecting the +//! supplied type/converter failed and must not silently become an unsupported-type fallback. +//! Runtime conversion failures retain Python's errors and result ordering. Nested, decimal, +//! and timestamp conversion is not implemented natively by this UDF policy. + +use std::collections::HashMap; + +use arrow::array::{ArrayRef, StructArray}; +use arrow::datatypes::{DataType, Field, Fields, Schema}; +use pyo3::prelude::*; + +mod input; +mod result; + +pub use input::{input_spec, InputColumns, InputSpec}; +pub use result::{result_spec, Finished, ResultColumn, ResultSpec}; + +/// The struct Spark gives `VariantType`, with parts of `part`: `Binary` from PySpark's +/// `to_arrow_type`, or `LargeBinary` from the JVM under `useLargeVarTypes`. Its tagged `metadata` +/// field tells it apart from a struct of two binaries. +fn variant_fields(part: DataType) -> Fields { + Fields::from(vec![ + Field::new("value", part.clone(), false), + Field::new("metadata", part, false) + .with_metadata(HashMap::from([("variant".to_owned(), "true".to_owned())])), + ]) +} + +/// Whether `spark_type` is exactly `sql_types.`, a type with a supported native converter. +/// `sql_types` is the caller-supplied backend module that parsed `spark_type`. +fn is_spark_type( + spark_type: &Bound<'_, PyAny>, + sql_types: &Bound<'_, PyAny>, + name: &str, +) -> PyResult { + Ok(spark_type.get_type().is(&sql_types.getattr(name)?)) +} + +/// Imports `schema`, a `pa.Schema`, such as the input schema a UDF's batches follow. +pub fn import_schema(schema: &Bound<'_, PyAny>) -> PyResult { + crate::bridge::import_schema(schema) +} + +/// Imports `data_type`, a `pa.DataType`, rejecting root extension types whose metadata +/// cannot be represented by a bare Rust data type. +pub fn import_data_type(data_type: &Bound<'_, PyAny>) -> PyResult { + crate::bridge::import_data_type(data_type) +} + +fn import_batch(batch: &Bound<'_, PyAny>) -> PyResult { + Ok(StructArray::from(crate::bridge::import_batch(batch)?)) +} + +/// Hands `array` to pyarrow, whose `pa.Array` is `array_type`, without copying it. pyarrow moves +/// both structs' contents out, so dropping them afterwards releases nothing. +fn export_array<'py>( + array_type: &Bound<'py, PyAny>, + array: ArrayRef, +) -> PyResult> { + crate::bridge::export_array_with_type(array_type, array) +} diff --git a/crates/spark-python-arrow/src/udf/result.rs b/crates/spark-python-arrow/src/udf/result.rs new file mode 100644 index 00000000..30a87516 --- /dev/null +++ b/crates/spark-python-arrow/src/udf/result.rs @@ -0,0 +1,881 @@ +/* + * 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::sync::Arc; + +use arrow::array::{ + Array, ArrayRef, AsArray, BinaryBuilder, BooleanBufferBuilder, BooleanBuilder, Float32Array, + Float64Builder, Int16Builder, Int32Builder, Int64Builder, Int8Builder, LargeBinaryBuilder, + LargeStringBuilder, NullBufferBuilder, PrimitiveBuilder, StringBuilder, StructArray, +}; +use arrow::buffer::BooleanBuffer; +use arrow::datatypes::{ArrowPrimitiveType, DataType, Float64Type}; +use pyo3::exceptions::PyRuntimeError; +use pyo3::intern; +use pyo3::prelude::*; +use pyo3::types::{PyBool, PyBytes, PyDict, PyFloat, PyInt, PyList, PyString, PyType}; + +use crate::bridge::arrow_error; +use crate::scalar::{exact_bool, exact_bytes, exact_f64, exact_i64, exact_str, into_exact_bytes}; +use crate::udf::input::{Column, Read}; +use crate::udf::{export_array, import_data_type, is_spark_type, variant_fields}; + +/// The largest integer magnitudes PyArrow converts to float64 and float32, which it refuses +/// above 2^53 and 2^24. +const MAX_EXACT_F64_INT: u64 = 1 << 53; +const MAX_EXACT_F32_INT: u64 = 1 << 24; + +/// Accepts a UDF's results for [`ResultColumn`], or returns why its native builder is unsupported. +/// `arrow_type` is the `pa.DataType` PySpark builds for `spark_return_type`, `converter` is +/// PySpark's result converter or `None`, `safecheck` is +/// `spark.sql.execution.pandas.convertToArrowArraySafely`, which PySpark's retry casts with, and +/// `sql_types` is the backend package's `sql.types` module, whose classes `spark_return_type` +/// comes from. For supported Variant results, `variant_class` resolves the class UDF code +/// constructs, which need not be the backend module's class. The resolver runs only after the +/// Arrow type and converter are accepted; its errors propagate. No package or class is inferred. +pub fn result_spec<'py>( + py: Python<'py>, + arrow_type: &Bound<'_, PyAny>, + spark_return_type: &Bound<'_, PyAny>, + converter: &Bound<'_, PyAny>, + safecheck: bool, + sql_types: &Bound<'_, PyAny>, + variant_class: impl FnOnce() -> PyResult>>, +) -> PyResult> { + // Lowering an extension type to its storage type would select the primitive + // builder and export an array without the requested extension identity. + if crate::bridge::is_extension_type(arrow_type)? { + return Ok(Err( + "Arrow extension types require the caller's Python converter".to_owned(), + )); + } + let data_type = import_data_type(arrow_type)?; + let Some(kind) = OutputKind::of(&data_type) else { + return Ok(Err(format!( + "{data_type} is not supported by the native UDF result converter" + ))); + }; + // Build a variant only as PySpark's VariantType converter makes it. + let converter = match (converter.is_none(), kind.converter_type()) { + (true, _) if kind != OutputKind::Variant => None, + (false, Some(name)) if is_spark_type(spark_return_type, sql_types, name)? => { + Some(converter.clone().unbind()) + } + _ => { + return Ok(Err(format!( + "{spark_return_type} needs PySpark's converter" + ))) + } + }; + let variant = match kind { + // The caller knows the class UDF code constructs; it need not belong to `sql_types`. + OutputKind::Variant => match variant_class()? { + Some(class) => Some(class.into_any().unbind()), + None => { + return Ok(Err( + "VariantType requires the caller's VariantVal class".to_owned() + )) + } + }, + _ => None, + }; + let pyarrow = py.import("pyarrow")?; + Ok(Ok(ResultSpec { + kind, + arrow_type: arrow_type.clone().unbind(), + converter, + variant, + safecheck, + value_capacity: i32::MAX as usize - 1, + pyarrow_array: pyarrow.getattr("array")?.unbind(), + concat_arrays: pyarrow.getattr("concat_arrays")?.unbind(), + array_type: pyarrow.getattr("Array")?.unbind(), + arrow_invalid: pyarrow.getattr("lib")?.getattr("ArrowInvalid")?.unbind(), + })) +} + +/// A UDF's accepted result type, and what PySpark's output step needs for it. +pub struct ResultSpec { + kind: OutputKind, + arrow_type: Py, + /// PySpark's `StringType`, `BinaryType` or `VariantType` converter; every other result type + /// handled by this native builder has none. + converter: Option>, + /// PySpark's `VariantVal`, for a variant result. + variant: Option>, + safecheck: bool, + /// The most bytes the values of a string, binary or variant-part array can take, as its + /// offsets are 32-bit: PyArrow's limit, one under `i32::MAX`, past which `pa.array` chunks. + value_capacity: usize, + pyarrow_array: Py, + concat_arrays: Py, + array_type: Py, + arrow_invalid: Py, +} + +impl ResultSpec { + /// Lowers the short-offset buffer threshold before selecting the PyArrow fallback. + /// + /// This also permits testing chunk/fallback behavior without allocating 2 GiB. The threshold + /// cannot be raised above its current value or the maximum supported by 32-bit offsets. + pub fn with_value_capacity(self, value_capacity: usize) -> ResultSpec { + ResultSpec { + value_capacity: self.value_capacity.min(value_capacity), + ..self + } + } + + /// `value` through PySpark's result converter. Each returns `None` unchanged, and the + /// `StringType` and `BinaryType` converters return exact strs and bytes unchanged, so it is not + /// called for them. + fn convert<'py>( + &self, + py: Python<'py>, + value: Bound<'py, PyAny>, + ) -> PyResult> { + match &self.converter { + Some(converter) if !(value.is_none() || self.kind.converts_unchanged(&value)) => { + converter.bind(py).call1((value,)) + } + _ => Ok(value), + } + } + + /// The `except pa.lib.ArrowInvalid` branch of PySpark's output step: + /// `pa.array(values).cast(target_type=T, safe=safecheck)`. + fn retry<'py>( + &self, + py: Python<'py>, + values: Bound<'py, PyList>, + ) -> PyResult> { + let kwargs = PyDict::new(py); + kwargs.set_item(intern!(py, "target_type"), self.arrow_type.bind(py))?; + kwargs.set_item(intern!(py, "safe"), self.safecheck)?; + let inferred = self.pyarrow_array.bind(py).call1((values,))?; + inferred.call_method(intern!(py, "cast"), (), Some(&kwargs)) + } +} + +/// The result types this UDF policy builds natively. Date results use the caller's Python path; +/// they are not part of this native result policy. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum OutputKind { + Int8, + Int16, + Int32, + Int64, + Float32, + Float64, + Boolean, + Utf8, + LargeUtf8, + Binary, + LargeBinary, + Variant, +} + +impl OutputKind { + fn of(data_type: &DataType) -> Option { + Some(match data_type { + DataType::Int8 => OutputKind::Int8, + DataType::Int16 => OutputKind::Int16, + DataType::Int32 => OutputKind::Int32, + DataType::Int64 => OutputKind::Int64, + DataType::Float32 => OutputKind::Float32, + DataType::Float64 => OutputKind::Float64, + DataType::Boolean => OutputKind::Boolean, + DataType::Utf8 => OutputKind::Utf8, + DataType::LargeUtf8 => OutputKind::LargeUtf8, + DataType::Binary => OutputKind::Binary, + DataType::LargeBinary => OutputKind::LargeBinary, + DataType::Struct(fields) if *fields == variant_fields(DataType::Binary) => { + OutputKind::Variant + } + _ => return None, + }) + } + + fn is_string(self) -> bool { + matches!(self, OutputKind::Utf8 | OutputKind::LargeUtf8) + } + + /// The Spark type whose PySpark converter results of this kind may go through. + fn converter_type(self) -> Option<&'static str> { + match self { + OutputKind::Utf8 | OutputKind::LargeUtf8 => Some("StringType"), + OutputKind::Binary | OutputKind::LargeBinary => Some("BinaryType"), + OutputKind::Variant => Some("VariantType"), + _ => None, + } + } + + /// Whether this kind's converter returns `value` unchanged. + fn converts_unchanged(self, value: &Bound<'_, PyAny>) -> bool { + match self { + OutputKind::Utf8 | OutputKind::LargeUtf8 => value.is_exact_instance_of::(), + OutputKind::Binary | OutputKind::LargeBinary => value.is_exact_instance_of::(), + _ => false, + } + } +} + +enum Builder { + Int8(Int8Builder), + Int16(Int16Builder), + Int32(Int32Builder), + Int64(Int64Builder), + /// Each result as the float64 PySpark's retry has to be given, which `finish` narrows as + /// PyArrow does, with the rows that held an exact `int`. + Float32(Float64Builder, BooleanBufferBuilder), + /// With the rows that held an exact `int`, which PySpark's retry has to be given back as one. + Float64(Float64Builder, BooleanBufferBuilder), + Boolean(BooleanBuilder), + Utf8(StringBuilder), + LargeUtf8(LargeStringBuilder), + Binary(BinaryBuilder), + LargeBinary(LargeBinaryBuilder), + /// The variants' two parts and which rows hold one. It takes no result when the UDF returns + /// it, as a VariantVal can change afterwards, only every result once the UDF has run on + /// every row: see [`Builder::append_variants`]. + Variant(BinaryBuilder, BinaryBuilder, NullBufferBuilder), +} + +/// What a [`Builder`] appended. +struct Appended { + /// As an array of the spec's type. + array: ArrayRef, + /// The same rows as the objects PySpark's retry is given: for float32, the float64 each was, + /// and for a string column, what the converter made of a built-in. + returned: ArrayRef, + /// For floats, the rows that held an exact `int`. + int_rows: Option, +} + +impl Builder { + fn new(kind: OutputKind, rows: usize) -> Builder { + match kind { + OutputKind::Int8 => Builder::Int8(Int8Builder::with_capacity(rows)), + OutputKind::Int16 => Builder::Int16(Int16Builder::with_capacity(rows)), + OutputKind::Int32 => Builder::Int32(Int32Builder::with_capacity(rows)), + OutputKind::Int64 => Builder::Int64(Int64Builder::with_capacity(rows)), + OutputKind::Float32 => Builder::Float32( + Float64Builder::with_capacity(rows), + BooleanBufferBuilder::new(rows), + ), + OutputKind::Float64 => Builder::Float64( + Float64Builder::with_capacity(rows), + BooleanBufferBuilder::new(rows), + ), + OutputKind::Boolean => Builder::Boolean(BooleanBuilder::with_capacity(rows)), + OutputKind::Utf8 => Builder::Utf8(StringBuilder::with_capacity(rows, rows * 8)), + OutputKind::LargeUtf8 => { + Builder::LargeUtf8(LargeStringBuilder::with_capacity(rows, rows * 8)) + } + OutputKind::Binary => Builder::Binary(BinaryBuilder::with_capacity(rows, rows * 8)), + OutputKind::LargeBinary => { + Builder::LargeBinary(LargeBinaryBuilder::with_capacity(rows, rows * 8)) + } + OutputKind::Variant => Builder::Variant( + BinaryBuilder::with_capacity(rows, rows * 8), + BinaryBuilder::with_capacity(rows, rows * 8), + NullBufferBuilder::new(rows), + ), + } + } + + /// Appends `value` if it is `None` or an exact `int`, `float`, `bool`, `str` or `bytes` that + /// PyArrow is known to convert to this type unchanged; appends nothing and returns false + /// otherwise, or if a `Utf8` or `Binary` builder would then hold more than `capacity` bytes. A + /// variant builder takes no value, not even `None`: see [`Builder::append_variants`]. + fn append_exact(&mut self, value: &Bound<'_, PyAny>, capacity: usize) -> bool { + if value.is_none() { + match self { + Builder::Int8(builder) => builder.append_null(), + Builder::Int16(builder) => builder.append_null(), + Builder::Int32(builder) => builder.append_null(), + Builder::Int64(builder) => builder.append_null(), + Builder::Float32(builder, ints) | Builder::Float64(builder, ints) => { + builder.append_null(); + ints.append(false); + } + Builder::Boolean(builder) => builder.append_null(), + Builder::Utf8(builder) => builder.append_null(), + Builder::LargeUtf8(builder) => builder.append_null(), + Builder::Binary(builder) => builder.append_null(), + Builder::LargeBinary(builder) => builder.append_null(), + Builder::Variant(..) => return false, + } + return true; + } + match self { + Builder::Int8(builder) => append_int(builder, value), + Builder::Int16(builder) => append_int(builder, value), + Builder::Int32(builder) => append_int(builder, value), + Builder::Int64(builder) => append_int(builder, value), + Builder::Float32(builder, ints) => { + append_float(builder, ints, value, MAX_EXACT_F32_INT) + } + Builder::Float64(builder, ints) => { + append_float(builder, ints, value, MAX_EXACT_F64_INT) + } + Builder::Boolean(builder) => { + if let Some(value) = exact_bool(value) { + builder.append_value(value); + true + } else { + false + } + } + Builder::Utf8(builder) => match exact_str(value) { + Some(text) if builder.values_slice().len() + text.len() <= capacity => { + builder.append_value(text); + true + } + _ => false, + }, + Builder::LargeUtf8(builder) => match exact_str(value) { + Some(text) => { + builder.append_value(text); + true + } + None => false, + }, + Builder::Binary(builder) => match exact_bytes(value) { + Some(bytes) if builder.values_slice().len() + bytes.len() <= capacity => { + builder.append_value(bytes); + true + } + _ => false, + }, + Builder::LargeBinary(builder) => match exact_bytes(value) { + Some(bytes) => { + builder.append_value(bytes); + true + } + None => false, + }, + Builder::Variant(..) => false, + } + } + + /// Appends `values`, all of a variant column's results, given once the UDF has run on every + /// row, if each is `None` or a VariantVal of exactly `class` whose two parts are exact + /// `bytes`: PySpark's converter then gives PyArrow those two, which it stores unchanged. + /// Appends nothing and returns false otherwise, if either part's values would take more than + /// `capacity` bytes, and for any other column. + fn append_variants( + &mut self, + values: &Bound<'_, PyList>, + class: &Bound<'_, PyAny>, + capacity: usize, + ) -> bool { + let Builder::Variant(value_parts, metadata_parts, nulls) = self else { + return false; + }; + let py = values.py(); + let mut parts = Vec::with_capacity(values.len()); + for value in values.iter() { + if value.is_none() { + parts.push(None); + continue; + } + if !value.get_type().is(class) { + return false; + } + // In the order PySpark's converter reads them. + let part = value + .getattr(intern!(py, "value")) + .ok() + .and_then(into_exact_bytes); + let metadata = value + .getattr(intern!(py, "metadata")) + .ok() + .and_then(into_exact_bytes); + let (Some(part), Some(metadata)) = (part, metadata) else { + return false; + }; + parts.push(Some((part, metadata))); + } + let (value_bytes, metadata_bytes) = parts.iter().flatten().fold( + (0, 0), + |(value_bytes, metadata_bytes), (part, metadata)| { + ( + value_bytes + part.as_bytes().len(), + metadata_bytes + metadata.as_bytes().len(), + ) + }, + ); + if value_bytes > capacity || metadata_bytes > capacity { + return false; + } + for part in &parts { + match part { + // PyArrow gives a null variant's two parts empty values, not nulls. + None => { + value_parts.append_value(b""); + metadata_parts.append_value(b""); + nulls.append_null(); + } + Some((part, metadata)) => { + value_parts.append_value(part.as_bytes()); + metadata_parts.append_value(metadata.as_bytes()); + nulls.append_non_null(); + } + } + } + true + } + + /// The bytes of values a `Utf8` or `Binary` builder holds; `None` for every other builder. + fn value_bytes(&self) -> Option { + match self { + Builder::Utf8(builder) => Some(builder.values_slice().len()), + Builder::Binary(builder) => Some(builder.values_slice().len()), + _ => None, + } + } + + /// The values so far. + fn finish(&mut self) -> PyResult { + fn as_returned(array: ArrayRef) -> Appended { + Appended { + returned: array.clone(), + array, + int_rows: None, + } + } + Ok(match self { + Builder::Int8(builder) => as_returned(Arc::new(builder.finish())), + Builder::Int16(builder) => as_returned(Arc::new(builder.finish())), + Builder::Int32(builder) => as_returned(Arc::new(builder.finish())), + Builder::Int64(builder) => as_returned(Arc::new(builder.finish())), + Builder::Float32(builder, ints) => { + let returned = builder.finish(); + // As PyArrow's static_cast does. + let narrowed: Float32Array = returned.unary(|value| value as f32); + Appended { + array: Arc::new(narrowed), + returned: Arc::new(returned), + int_rows: Some(ints.finish()), + } + } + Builder::Float64(builder, ints) => Appended { + int_rows: Some(ints.finish()), + ..as_returned(Arc::new(builder.finish())) + }, + Builder::Boolean(builder) => as_returned(Arc::new(builder.finish())), + Builder::Utf8(builder) => as_returned(Arc::new(builder.finish())), + Builder::LargeUtf8(builder) => as_returned(Arc::new(builder.finish())), + Builder::Binary(builder) => as_returned(Arc::new(builder.finish())), + Builder::LargeBinary(builder) => as_returned(Arc::new(builder.finish())), + Builder::Variant(value_parts, metadata_parts, nulls) => as_returned(Arc::new( + StructArray::try_new( + variant_fields(DataType::Binary), + vec![ + Arc::new(value_parts.finish()), + Arc::new(metadata_parts.finish()), + ], + nulls.finish(), + ) + .map_err(arrow_error)?, + )), + }) + } +} + +/// Appends an exact `float`, or an exact `int` of at most `max_int` in magnitude, which PyArrow +/// converts without rounding. +fn append_float( + builder: &mut Float64Builder, + ints: &mut BooleanBufferBuilder, + value: &Bound<'_, PyAny>, + max_int: u64, +) -> bool { + if let Some(value) = exact_f64(value) { + builder.append_value(value); + ints.append(false); + true + } else if let Some(int) = exact_i64(value).filter(|int| int.unsigned_abs() <= max_int) { + builder.append_value(int as f64); + ints.append(true); + true + } else { + false + } +} + +fn append_int(builder: &mut PrimitiveBuilder, value: &Bound<'_, PyAny>) -> bool +where + T: ArrowPrimitiveType, + T::Native: TryFrom, +{ + match exact_i64(value).and_then(|int| T::Native::try_from(int).ok()) { + Some(int) => { + builder.append_value(int); + true + } + None => false, + } +} + +/// Whether `value` is `None` or an exact `int`, `float`, `bool` or `str`, none of which can change +/// after the UDF returns it. +fn is_builtin(value: &Bound<'_, PyAny>) -> bool { + value.is_none() + || value.is_exact_instance_of::() + || value.is_exact_instance_of::() + || value.is_exact_instance_of::() + || value.is_exact_instance_of::() +} + +/// One UDF's results for one batch, taken through PySpark's steps in order: every result through +/// the converter, then all of them through `pa.array(values, type=T)`, retried as +/// `pa.array(values).cast(T, safe=safecheck)` on `ArrowInvalid`. +/// +/// PySpark runs the UDF on every row of the batch before it converts any result, so a result can +/// change after the UDF returns it, and a UDF's error on a later row wins over a result that +/// cannot be converted. [`ResultColumn::push`] therefore converts only built-ins, and the caller +/// must run the UDF on every row, giving each result to `push` or [`ResultColumn::extend`] in row +/// order, before it calls [`ResultColumn::finish`]. +pub struct ResultColumn<'s> { + spec: &'s ResultSpec, + builder: Builder, + state: State, +} + +enum State { + Building, + /// The results from the first one the builder could not take, unconverted, in row order: + /// those in `listed`, then those in `pushed`. `listed` has the held part of each list given + /// to `extend` and the results pushed before the last such call; `pushed` has the results + /// pushed after it. + Holding { + listed: Option>, + pushed: Vec>, + }, + /// The converter raised, which PySpark does before it gives PyArrow any value, or a result + /// could not be held. + Failed(PyErr), +} + +/// A finished [`ResultColumn`]. +pub struct Finished<'py> { + /// The results as a `pa.Array` of the spec's type, or past 2 GiB the `pa.ChunkedArray` + /// PyArrow makes of them. + pub array: Bound<'py, PyAny>, + /// How many results the first `pa.array` call was given. + pub pyarrow_values: usize, + /// Whether PySpark's retry built the array. + pub retried: bool, +} + +impl<'s> ResultColumn<'s> { + /// A column for a batch of `rows` rows. + pub fn new(spec: &'s ResultSpec, rows: usize) -> ResultColumn<'s> { + ResultColumn { + spec, + builder: Builder::new(spec.kind, rows), + state: State::Building, + } + } + + /// Adds the next row's result. A built-in the builder takes is converted and appended now, as + /// it cannot change; from the first result that is not, every result is held for + /// [`ResultColumn::finish`]. Never raises: a converter error is raised by `finish`. + pub fn push<'py>(&mut self, py: Python<'py>, value: Bound<'py, PyAny>) { + match &mut self.state { + State::Building => {} + State::Holding { pushed, .. } => return pushed.push(value.unbind()), + State::Failed(_) => return, + } + if self.builder.append_exact(&value, self.spec.value_capacity) + || self.append_converted(py, &value) + { + return; + } + self.state = State::Holding { + listed: None, + pushed: vec![value.unbind()], + }; + } + + /// Adds the next rows' results, as pushing each of `values` in turn would, for a caller that + /// collected them into a list. From the first result the builder cannot take, the column keeps + /// a copy of the rest of `values` as one slice, so the caller may change `values` afterwards. + pub fn extend<'py>(&mut self, py: Python<'py>, values: &Bound<'py, PyList>) { + if let State::Building = self.state { + if let Err(error) = self.extend_building(py, values) { + self.state = State::Failed(error); + } + } else if let State::Holding { listed, pushed } = &mut self.state { + if let Err(error) = extend_held(py, listed, pushed, values) { + self.state = State::Failed(error); + } + } + } + + fn extend_building<'py>( + &mut self, + py: Python<'py>, + values: &Bound<'py, PyList>, + ) -> PyResult<()> { + let mut appended = 0; + // The iterator protocol reads a list with one call per item, and the abi3 list API with + // three. + for value in values.try_iter()? { + let value = value?; + if !(self.builder.append_exact(&value, self.spec.value_capacity) + || self.append_converted(py, &value)) + { + break; + } + if let State::Failed(_) = self.state { + return Ok(()); + } + appended += 1; + } + if appended < values.len() { + self.state = State::Holding { + listed: Some(values.get_slice(appended, values.len()).unbind()), + pushed: Vec::new(), + }; + } + Ok(()) + } + + /// Appends `value`, which the builder did not take, if it is a float, a bool or an int that + /// fits in 64 bits, which the `StringType` converter turns into a str it takes. Returns false + /// if `value` has to be held instead. A converter error fails the column. + fn append_converted<'py>(&mut self, py: Python<'py>, value: &Bound<'py, PyAny>) -> bool { + // A built-in cannot change after the UDF returns it, so converting it now matches PySpark, + // which converts once the UDF has run on every row. For a string column the builder takes + // only None and exact strs, which the StringType converter returns unchanged. + if !self.spec.kind.is_string() || self.spec.converter.is_none() || !is_builtin(value) { + return false; + } + // A later row may call sys.set_int_max_str_digits, which decides whether str() takes an + // int of more than 640 digits. + // Keep large integers until every UDF call has completed so that limit changes are seen. + if value.is_exact_instance_of::() && exact_i64(value).is_none() { + return false; + } + match self.spec.convert(py, value.clone()) { + Ok(converted) => self + .builder + .append_exact(&converted, self.spec.value_capacity), + Err(error) => { + self.state = State::Failed(error); + true + } + } + } + + /// The results as a `pa.Array`. Held results go through the converter and one + /// `pa.array(held, type=T)` call, and are joined to the appended ones; if PyArrow rejects one + /// with `ArrowInvalid`, PySpark's retry builds the whole batch instead. If the two would not + /// fit in one array, every result goes into that call, as in PySpark, where past 2 GiB it + /// returns a `pa.ChunkedArray` that fails the batch once every UDF has run. A converter error + /// is raised first, as PySpark converts every result before PyArrow sees one. A variant + /// column's results are all held, and are appended now if the builder takes every one. + pub fn finish<'py>(mut self, py: Python<'py>) -> PyResult> { + let mut held = match self.state { + State::Building => None, + State::Holding { listed, pushed } => Some(match listed { + None => PyList::new(py, pushed)?, + Some(listed) => { + let listed = listed.into_bound(py); + for value in pushed { + listed.append(value)?; + } + listed + } + }), + State::Failed(error) => return Err(error), + }; + if let (Some(values), Some(class)) = (&held, &self.spec.variant) { + if self + .builder + .append_variants(values, class.bind(py), self.spec.value_capacity) + { + held = None; + } + } + let head_bytes = self.builder.value_bytes(); + let appended = self.builder.finish()?; + let array_type = self.spec.array_type.bind(py); + let Some(held) = held else { + let array = export_array(array_type, appended.array)?; + return Ok(Finished { + array, + pyarrow_values: 0, + retried: false, + }); + }; + let converted = if self.spec.converter.is_some() { + let converted = held + .iter() + .map(|value| self.spec.convert(py, value)) + .collect::>>()?; + PyList::new(py, converted)? + } else { + held + }; + let pyarrow_array = self.spec.pyarrow_array.bind(py); + let arrow_invalid = self.spec.arrow_invalid.bind(py); + if appended.array.is_empty() + || head_bytes.is_some_and(|bytes| !fits(bytes, &converted, self.spec.value_capacity)) + { + let values = appended_values(py, &appended)?; + values.call_method1(intern!(py, "extend"), (&converted,))?; + let pyarrow_values = values.len(); + return match pyarrow_array.call1((&values, self.spec.arrow_type.bind(py))) { + Ok(array) => Ok(Finished { + array, + pyarrow_values, + retried: false, + }), + Err(error) if error.is_instance(py, arrow_invalid) => Ok(Finished { + array: self.spec.retry(py, values)?, + pyarrow_values, + retried: true, + }), + Err(error) => Err(error), + }; + } + let pyarrow_values = converted.len(); + match pyarrow_array.call1((&converted, self.spec.arrow_type.bind(py))) { + Ok(tail) => { + let head = export_array(array_type, appended.array)?; + let array = self + .spec + .concat_arrays + .bind(py) + .call1((PyList::new(py, [head, tail])?,))?; + Ok(Finished { + array, + pyarrow_values, + retried: false, + }) + } + Err(error) if error.is_instance(py, arrow_invalid) => { + let values = appended_values(py, &appended)?; + for value in converted.iter() { + values.append(value)?; + } + Ok(Finished { + array: self.spec.retry(py, values)?, + pyarrow_values, + retried: true, + }) + } + Err(error) => Err(error), + } + } +} + +/// Whether `values`, results through the converter, fit after `used` of `capacity` bytes in a +/// string or binary array. A value that is not `None`, an exact `str` or exact `bytes` is taken +/// not to. +fn fits(mut used: usize, values: &Bound<'_, PyList>, capacity: usize) -> bool { + for value in values.iter() { + used += if value.is_none() { + 0 + } else if let Some(text) = exact_str(&value) { + text.len() + } else if let Some(bytes) = exact_bytes(&value) { + bytes.len() + } else { + return false; + }; + if used > capacity { + return false; + } + } + true +} + +/// Appends `values` to the held results, after the pushed ones. +fn extend_held<'py>( + py: Python<'py>, + listed: &mut Option>, + pushed: &mut Vec>, + values: &Bound<'py, PyList>, +) -> PyResult<()> { + let list = listed + .get_or_insert_with(|| PyList::empty(py).unbind()) + .bind(py); + for value in pushed.drain(..) { + list.append(value)?; + } + list.call_method1(intern!(py, "extend"), (values,))?; + Ok(()) +} + +/// The appended results as the objects PySpark's retry is given, with the float rows that held +/// an exact `int` given back as one. +fn appended_values<'py>(py: Python<'py>, appended: &Appended) -> PyResult> { + let values = PyList::empty(py); + let returned = &appended.returned; + // A variant column appends nothing before it holds every result. + if returned.is_empty() { + return Ok(values); + } + let column = Column::new(py, returned, &Read::Plain)?.ok_or_else(|| { + PyRuntimeError::new_err(format!("no reader for {}", returned.data_type())) + })?; + for row in 0..returned.len() { + if appended + .int_rows + .as_ref() + .is_some_and(|ints| ints.value(row)) + { + values.append(returned.as_primitive::().value(row) as i64)?; + } else { + values.append(column.value(py, row, || { + PyRuntimeError::new_err(format!("appended result {row} cannot be read back")) + })?)?; + } + } + Ok(values) +} + +#[cfg(test)] +mod tests { + use arrow::datatypes::{Field, Fields, TimeUnit}; + + use super::*; + + #[test] + fn builds_a_variant_only_as_pyspark_types_it() { + assert_eq!( + OutputKind::of(&DataType::Struct(variant_fields(DataType::Binary))), + Some(OutputKind::Variant) + ); + let untagged = Fields::from(vec![ + Field::new("value", DataType::Binary, false), + Field::new("metadata", DataType::Binary, false), + ]); + for data_type in [ + DataType::Struct(untagged), + DataType::Decimal128(10, 2), + DataType::Timestamp(TimeUnit::Microsecond, None), + DataType::Date64, + ] { + assert_eq!(OutputKind::of(&data_type), None, "{data_type}"); + } + } +} diff --git a/crates/spark-python-arrow/tests/conversion.rs b/crates/spark-python-arrow/tests/conversion.rs new file mode 100644 index 00000000..28f76c3a --- /dev/null +++ b/crates/spark-python-arrow/tests/conversion.rs @@ -0,0 +1,976 @@ +// 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. + +//! Standalone conversion checks against the installed PyArrow. No Spark server is needed. +//! These exercise the shared interfaces, not the full worker or client conformance suites. + +use std::collections::HashMap; +use std::ffi::CString; +use std::sync::Arc; + +use arrow::array::{Array, Int64Array}; +use arrow::datatypes::{DataType, Field, Schema, TimeUnit}; +use arrow::record_batch::{RecordBatch, RecordBatchOptions}; +use pyo3::exceptions::{ + PyImportError, PyIndexError, PyTypeError, PyUnicodeDecodeError, PyValueError, +}; +use pyo3::prelude::*; +use pyo3::types::{PyBytes, PyDict, PyList, PyType}; +use spark_python_arrow::bridge::{ + export_array, export_data_type, export_schema, import_array, import_batch, import_data_type, + import_schema, record_batches_to_pyarrow, +}; +use spark_python_arrow::scalar::{ + exact_bool, exact_bytes, exact_f64, exact_i64, exact_str, extract_binary, utf8_into_python, + Scalar, +}; +use spark_python_arrow::udf::{ + input_spec, result_spec, Finished, InputColumns, ResultColumn, ResultSpec, +}; + +fn context(py: Python<'_>) -> PyResult> { + let locals = PyDict::new(py); + locals.set_item("pa", py.import("pyarrow")?)?; + locals.set_item("datetime", py.import("datetime")?)?; + Ok(locals) +} + +fn evaluate<'py>( + py: Python<'py>, + locals: &Bound<'py, PyDict>, + expression: &str, +) -> PyResult> { + let code = CString::new(expression).expect("test expressions contain no NUL"); + py.eval(&code, Some(locals), Some(locals)) +} + +fn extension_context<'py>(py: Python<'py>, name: &str) -> PyResult> { + let locals = context(py)?; + locals.set_item("extension_name", name)?; + py.run( + c"class TestExtensionType(pa.ExtensionType): + def __init__(self, storage_type=pa.string()): + super().__init__(storage_type, extension_name) + + def __arrow_ext_serialize__(self): + return b'extension metadata' + + @classmethod + def __arrow_ext_deserialize__(cls, storage_type, serialized): + assert serialized == b'extension metadata' + return cls(storage_type) + +extension_type = TestExtensionType() +extension_array = pa.ExtensionArray.from_storage( + extension_type, pa.array(['first', None, 'last']) +)", + Some(&locals), + Some(&locals), + )?; + Ok(locals) +} + +fn variant_context(py: Python<'_>) -> PyResult> { + let locals = context(py)?; + py.run( + c"class VariantType: pass + +class VariantVal: + def __init__(self, value, metadata): + self.value = value + self.metadata = metadata + +class SqlTypes: pass +SqlTypes.VariantType = VariantType +spark_type = VariantType() +converter = lambda value: {'value': value.value, 'metadata': value.metadata} +arrow_type = pa.struct([ + pa.field('value', pa.binary(), nullable=False), + pa.field('metadata', pa.binary(), nullable=False, metadata={b'variant': b'true'}), +])", + Some(&locals), + Some(&locals), + )?; + Ok(locals) +} + +fn assert_array_equal(actual: &Bound<'_, PyAny>, expected: &Bound<'_, PyAny>) -> PyResult<()> { + assert!( + actual + .call_method1("equals", (expected,))? + .extract::()?, + "actual {actual}; expected {expected}" + ); + Ok(()) +} + +fn primitive_spec<'py>( + py: Python<'py>, + arrow_type: &Bound<'py, PyAny>, + safe: bool, +) -> PyResult { + let none = py.None().into_bound(py); + Ok( + result_spec(py, arrow_type, &none, &none, safe, &none, || Ok(None))? + .expect("the primitive type must be supported"), + ) +} + +fn build_results<'py>( + py: Python<'py>, + arrow_type: &Bound<'py, PyAny>, + values: &Bound<'py, PyList>, + safe: bool, +) -> PyResult> { + let spec = primitive_spec(py, arrow_type, safe)?; + let mut column = ResultColumn::new(&spec, values.len()); + column.extend(py, values); + column.finish(py) +} + +// The documented worker output sequence, including its ArrowInvalid-only retry. +fn reference_results<'py>( + py: Python<'py>, + arrow_type: &Bound<'py, PyAny>, + values: &Bound<'py, PyList>, + safe: bool, +) -> PyResult> { + let pa = py.import("pyarrow")?; + match pa.call_method1("array", (values, arrow_type)) { + Ok(array) => Ok(array), + Err(error) if error.is_instance(py, &pa.getattr("ArrowInvalid")?) => { + let kwargs = PyDict::new(py); + kwargs.set_item("safe", safe)?; + pa.call_method1("array", (values,))? + .call_method("cast", (arrow_type,), Some(&kwargs)) + } + Err(error) => Err(error), + } +} + +#[test] +fn exported_array_keeps_rust_buffers_alive_without_copying() { + Python::initialize(); + Python::attach(|py| -> PyResult<()> { + let rust_array = Arc::new(Int64Array::from(vec![Some(12), None, Some(34)])); + let address = rust_array.values().as_ptr() as usize; + let array = export_array(py, rust_array.clone())?; + drop(rust_array); + py.import("gc")?.call_method0("collect")?; + + let array = array.bind(py); + let actual_address: usize = array + .call_method0("buffers")? + .get_item(1)? + .getattr("address")? + .extract()?; + assert_eq!(actual_address, address); + assert_eq!( + array + .call_method0("to_pylist")? + .extract::>>()?, + vec![Some(12), None, Some(34)] + ); + Ok(()) + }) + .unwrap(); +} + +#[test] +fn imported_slice_keeps_python_buffers_alive_without_copying() { + Python::initialize(); + Python::attach(|py| -> PyResult<()> { + let locals = context(py)?; + let original = evaluate(py, &locals, "pa.array([10, 20, None, 40], type=pa.int64())")?; + let address: usize = original + .call_method0("buffers")? + .get_item(1)? + .getattr("address")? + .extract()?; + let slice = original.call_method1("slice", (1, 2))?; + let imported = import_array(&slice)?; + drop(slice); + drop(original); + py.import("gc")?.call_method0("collect")?; + + let values = imported.as_any().downcast_ref::().unwrap(); + assert_eq!( + values.values().as_ptr() as usize, + address + std::mem::size_of::() + ); + assert_eq!(values.len(), 2); + assert_eq!(values.value(0), 20); + assert!(values.is_null(1)); + + let exported = export_array(py, imported)?; + let expected = evaluate(py, &locals, "pa.array([20, None], type=pa.int64())")?; + assert_array_equal(exported.bind(py), &expected) + }) + .unwrap(); +} + +#[test] +fn schema_and_nested_type_round_trip_preserves_metadata() { + Python::initialize(); + Python::attach(|py| -> PyResult<()> { + let field = Field::new("value", DataType::Int64, true) + .with_metadata(HashMap::from([("field_key".to_owned(), "a\0b".to_owned())])); + let schema = Schema::new_with_metadata( + vec![field], + HashMap::from([("schema_key".to_owned(), "c\0d".to_owned())]), + ); + let exported = export_schema(py, &schema)?; + assert_eq!(import_schema(exported.bind(py))?, schema); + + for data_type in [ + DataType::Int64, + DataType::Utf8, + DataType::List(Arc::new(Field::new("item", DataType::Int32, true))), + DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())), + DataType::Struct(schema.fields().clone()), + ] { + let exported = export_data_type(py, &data_type)?; + assert_eq!(import_data_type(exported.bind(py))?, data_type); + } + Ok(()) + }) + .unwrap(); +} + +#[test] +fn record_batch_import_preserves_schema_and_values() { + Python::initialize(); + Python::attach(|py| -> PyResult<()> { + let locals = context(py)?; + let batch = evaluate( + py, + &locals, + "pa.RecordBatch.from_arrays([pa.array([1, None, 3], type=pa.int64())], \ + schema=pa.schema([pa.field('x', pa.int64(), metadata={b'f': b'v'})], \ + metadata={b's': b'v'}))", + )?; + let schema = import_schema(&batch.getattr("schema")?)?; + let imported = import_batch(&batch)?; + assert_eq!(imported.schema().as_ref(), &schema); + assert_eq!(imported.num_rows(), 3); + let exported = record_batches_to_pyarrow(py, vec![imported])?; + let expected = py + .import("pyarrow")? + .getattr("Table")? + .call_method1("from_batches", (PyList::new(py, [batch])?,))?; + assert_array_equal(exported.bind(py), &expected) + }) + .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(|a| a.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 empty_batches_and_zero_column_rows_round_trip() { + 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 + ); + + 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.clone())])?; + assert_eq!(empty.bind(py).getattr("num_rows")?.extract::()?, 0); + assert_eq!(import_schema(&empty.bind(py).getattr("schema")?)?, *schema); + + let options = RecordBatchOptions::new().with_row_count(Some(4)); + let batch = + RecordBatch::try_new_with_options(Arc::new(Schema::empty()), vec![], &options).unwrap(); + let table = record_batches_to_pyarrow(py, vec![batch])?; + assert_eq!(table.bind(py).getattr("num_rows")?.extract::()?, 4); + assert_eq!( + table.bind(py).getattr("num_columns")?.extract::()?, + 0 + ); + + let locals = context(py)?; + let batch = evaluate( + py, + &locals, + "pa.record_batch([pa.array([1, 2, 3])], names=['x']).select([])", + )?; + let imported = import_batch(&batch)?; + assert_eq!(imported.num_rows(), 3); + assert_eq!(imported.num_columns(), 0); + Ok(()) + }) + .unwrap(); +} + +#[test] +fn bridge_rejects_wrong_objects_without_calling_their_exporter() { + Python::initialize(); + Python::attach(|py| -> PyResult<()> { + let locals = context(py)?; + py.run( + c"class Fake: + def _export_to_c(self, *args): + raise AssertionError('called')", + Some(&locals), + Some(&locals), + )?; + let fake = evaluate(py, &locals, "Fake()")?; + for error in [ + import_schema(&fake).unwrap_err(), + import_data_type(&fake).unwrap_err(), + import_array(&fake).unwrap_err(), + import_batch(&fake).unwrap_err(), + ] { + assert!(error.is_instance_of::(py), "{error}"); + } + Ok(()) + }) + .unwrap(); +} + +#[test] +fn bare_extension_imports_fail_instead_of_losing_type_identity() { + Python::initialize(); + Python::attach(|py| -> PyResult<()> { + let locals = extension_context(py, "spark_python_arrow.tests.bare_extension")?; + let data_type = evaluate(py, &locals, "extension_type")?; + let array = evaluate(py, &locals, "extension_array")?; + for error in [ + import_data_type(&data_type).unwrap_err(), + spark_python_arrow::udf::import_data_type(&data_type).unwrap_err(), + import_array(&array).unwrap_err(), + ] { + assert!(error.is_instance_of::(py), "{error}"); + assert!(error.to_string().contains("extension types"), "{error}"); + } + Ok(()) + }) + .unwrap(); +} + +#[test] +fn extension_results_select_python_fallback_without_lowering_the_type() { + Python::initialize(); + Python::attach(|py| -> PyResult<()> { + let locals = extension_context(py, "spark_python_arrow.tests.result_extension")?; + let data_type = evaluate(py, &locals, "extension_type")?; + let none = py.None().into_bound(py); + let unsupported = result_spec(py, &data_type, &none, &none, true, &none, || { + panic!("an unsupported extension must not resolve the Variant class") + })?; + let reason = match unsupported { + Ok(_) => panic!("an extension result must not use its primitive storage builder"), + Err(reason) => reason, + }; + assert!(reason.contains("extension types"), "{reason}"); + let wrong_type = evaluate(py, &locals, "object()")?; + let failure = result_spec(py, &wrong_type, &none, &none, true, &none, || { + panic!("an invalid Arrow type must not resolve the Variant class") + }); + assert!(matches!(failure, Err(error) if error.is_instance_of::(py))); + Ok(()) + }) + .unwrap(); +} + +#[test] +fn field_metadata_preserves_extensions_in_schemas_and_record_batches() { + Python::initialize(); + Python::attach(|py| -> PyResult<()> { + let name = "spark_python_arrow.tests.field_extension"; + let locals = extension_context(py, name)?; + let data_type = evaluate(py, &locals, "extension_type")?; + let pa = py.import("pyarrow")?; + pa.call_method1("register_extension_type", (&data_type,))?; + let result = (|| -> PyResult<()> { + let schema = evaluate( + py, + &locals, + "pa.schema([pa.field('value', extension_type, metadata={b'f': b'v'})], \ + metadata={b's': b'v'})", + )?; + let imported_schema = import_schema(&schema)?; + assert_eq!( + imported_schema + .field(0) + .metadata() + .get("ARROW:extension:name") + .map(String::as_str), + Some(name) + ); + let exported_schema = export_schema(py, &imported_schema)?; + assert!(exported_schema + .bind(py) + .call_method1("equals", (&schema, true))? + .is_truthy()?); + + // Nested extension fields have a metadata owner even when the outer + // value is a bare DataType, so only root extension types are rejected. + let nested = evaluate( + py, + &locals, + "pa.struct([pa.field('value', extension_type)])", + )?; + let exported_type = export_data_type(py, &import_data_type(&nested)?)?; + assert!(exported_type + .bind(py) + .call_method1("equals", (&nested,))? + .is_truthy()?); + + let kwargs = PyDict::new(py); + kwargs.set_item("schema", &schema)?; + let batch = pa.getattr("RecordBatch")?.call_method( + "from_arrays", + (PyList::new( + py, + [evaluate(py, &locals, "extension_array")?], + )?,), + Some(&kwargs), + )?; + let imported = import_batch(&batch)?; + assert_eq!(imported.schema().as_ref(), &imported_schema); + let exported = record_batches_to_pyarrow(py, vec![imported])?; + let expected = pa + .getattr("Table")? + .call_method1("from_batches", (PyList::new(py, [batch])?,))?; + assert_array_equal(exported.bind(py), &expected)?; + assert!(exported + .bind(py) + .getattr("schema")? + .call_method1("equals", (&schema, true))? + .is_truthy()?); + Ok(()) + })(); + pa.call_method1("unregister_extension_type", (name,))?; + result + }) + .unwrap(); +} + +#[test] +fn non_variant_and_unsupported_policies_do_not_resolve_the_variant_class() { + Python::initialize(); + Python::attach(|py| -> PyResult<()> { + let locals = variant_context(py)?; + let none = py.None().into_bound(py); + for (expression, supported) in [("pa.int64()", true), ("pa.list_(pa.int64())", false)] { + let arrow_type = evaluate(py, &locals, expression)?; + let selected = result_spec(py, &arrow_type, &none, &none, true, &none, || { + panic!("non-Variant and unsupported Arrow types must not resolve VariantVal") + })?; + assert_eq!(selected.is_ok(), supported); + } + let arrow_type = evaluate(py, &locals, "arrow_type")?; + let spark_type = evaluate(py, &locals, "spark_type")?; + let converter = evaluate(py, &locals, "converter")?; + let sql_types = evaluate(py, &locals, "SqlTypes")?; + for (spark_type, converter) in [(&spark_type, &none), (&none, &converter)] { + let selected = result_spec( + py, + &arrow_type, + spark_type, + converter, + true, + &sql_types, + || panic!("an unsupported converter must not resolve VariantVal"), + )?; + assert!(selected.is_err()); + } + Ok(()) + }) + .unwrap(); +} + +#[test] +fn supported_variant_policy_resolves_once_and_propagates_resolver_errors() { + Python::initialize(); + Python::attach(|py| -> PyResult<()> { + let locals = variant_context(py)?; + let arrow_type = evaluate(py, &locals, "arrow_type")?; + let spark_type = evaluate(py, &locals, "spark_type")?; + let converter = evaluate(py, &locals, "converter")?; + let sql_types = evaluate(py, &locals, "SqlTypes")?; + let mut calls = 0; + let failure = result_spec( + py, + &arrow_type, + &spark_type, + &converter, + true, + &sql_types, + || { + calls += 1; + Err(PyImportError::new_err("VariantVal resolver failed")) + }, + ); + assert_eq!(calls, 1); + let error = match failure { + Err(error) => error, + Ok(_) => panic!("a resolver failure must remain a Python error"), + }; + assert!(error.is_instance_of::(py)); + assert!(error.to_string().contains("VariantVal resolver failed")); + + let missing = result_spec( + py, + &arrow_type, + &spark_type, + &converter, + true, + &sql_types, + || Ok(None), + )?; + assert!(matches!(missing, Err(reason) if reason.contains("VariantVal class"))); + + let selected = result_spec( + py, + &arrow_type, + &spark_type, + &converter, + true, + &sql_types, + || { + calls += 1; + Ok(Some( + evaluate(py, &locals, "VariantVal")?.cast_into::()?, + )) + }, + )?; + assert!(selected.is_ok()); + assert_eq!(calls, 2); + Ok(()) + }) + .unwrap(); +} + +#[test] +fn table_export_rejects_inconsistent_schemas() { + Python::initialize(); + Python::attach(|py| { + let batches = ["a", "b"] + .into_iter() + .map(|name| { + RecordBatch::new_empty(Arc::new(Schema::new(vec![Field::new( + name, + DataType::Int64, + true, + )]))) + }) + .collect(); + let error = record_batches_to_pyarrow(py, batches).unwrap_err(); + assert!(error.is_instance_of::(py)); + }); +} + +#[test] +fn scalar_materialization_preserves_values_and_python_types() { + Python::initialize(); + Python::attach(|py| -> PyResult<()> { + let locals = context(py)?; + let cases = [ + (Scalar::Null, "None"), + (Scalar::Bool(true), "True"), + (Scalar::Int(i64::MIN), "-(2**63)"), + (Scalar::Int(i64::MAX), "2**63 - 1"), + (Scalar::Float(1.25), "1.25"), + (Scalar::String("a\0b"), "'a\\x00b'"), + (Scalar::String("\u{03bb}"), "'\\u03bb'"), + (Scalar::Binary(b"a\0b"), "b'a\\x00b'"), + (Scalar::Date(0), "datetime.date(1970, 1, 1)"), + (Scalar::Date(-719_162), "datetime.date.min"), + (Scalar::Date(2_932_896), "datetime.date.max"), + ]; + for (value, expression) in cases { + let actual = value.into_python(py)?; + let expected = evaluate(py, &locals, expression)?; + assert!(actual.eq(&expected)?, "{value:?}"); + assert!(actual.get_type().is(expected.get_type()), "{value:?}"); + } + let bytes = Scalar::Binary(b"value").into_python(py)?; + assert!(bytes.is_exact_instance_of::()); + for days in [-719_163, 2_932_897, i32::MIN, i32::MAX] { + let error = Scalar::Date(days).into_python(py).unwrap_err(); + let expected = evaluate( + py, + &locals, + &format!("datetime.date.fromordinal({})", i64::from(days) + 719_163), + ) + .unwrap_err(); + assert!(error.get_type(py).is(expected.get_type(py))); + } + Ok(()) + }) + .unwrap(); +} + +#[test] +fn exact_scalar_extraction_rejects_coercion_and_clears_rejected_errors() { + Python::initialize(); + Python::attach(|py| -> PyResult<()> { + let locals = context(py)?; + py.run( + c"class IntSubclass(int): pass\nclass StrSubclass(str): pass", + Some(&locals), + Some(&locals), + )?; + let truth = evaluate(py, &locals, "True")?; + assert_eq!(exact_bool(&truth), Some(true)); + assert_eq!(exact_i64(&truth), None); + let integer = evaluate(py, &locals, "7")?; + assert_eq!(exact_i64(&integer), Some(7)); + assert_eq!(exact_bool(&integer), None); + assert_eq!(exact_f64(&integer), None); + assert_eq!(exact_f64(&evaluate(py, &locals, "1.25")?), Some(1.25)); + for expression in ["IntSubclass(7)", "2**100", "-(2**100)"] { + assert_eq!(exact_i64(&evaluate(py, &locals, expression)?), None); + assert!(!PyErr::occurred(py)); + } + for expression in ["StrSubclass('text')", "'\\ud800'"] { + assert_eq!(exact_str(&evaluate(py, &locals, expression)?), None); + assert!(!PyErr::occurred(py)); + } + let text = evaluate(py, &locals, "'a\\x00b'")?; + assert_eq!(exact_str(&text).as_deref(), Some("a\0b")); + let bytes = evaluate(py, &locals, "b'ab'")?; + assert_eq!(exact_bytes(&bytes), Some(b"ab".as_slice())); + assert_eq!(extract_binary(&bytes), Some(b"ab".to_vec())); + let mutable = evaluate(py, &locals, "bytearray(b'ab')")?; + assert_eq!(exact_bytes(&mutable), None); + assert_eq!(extract_binary(&mutable), Some(b"ab".to_vec())); + for expression in ["[97, 98]", "(97, 98)", "'ab'"] { + assert_eq!(extract_binary(&evaluate(py, &locals, expression)?), None); + } + let error = utf8_into_python(py, &[0xff]).unwrap_err(); + assert!(error.is_instance_of::(py)); + Ok(()) + }) + .unwrap(); +} + +#[test] +fn exact_primitive_results_match_pyarrow_without_python_value_fallback() { + Python::initialize(); + Python::attach(|py| -> PyResult<()> { + let locals = context(py)?; + for (data_type, values) in [ + ("pa.int8()", "[-128, None, 127]"), + ("pa.int16()", "[-32768, None, 32767]"), + ("pa.int32()", "[-(2**31), None, 2**31 - 1]"), + ("pa.int64()", "[-(2**63), None, 2**63 - 1]"), + ("pa.float32()", "[-1.25, None, 2**24, 0, -0.0]"), + ("pa.float64()", "[-1.25, None, 2**53, 0, -0.0]"), + ("pa.bool_()", "[True, None, False]"), + ("pa.string()", "['hello', None, 'a\\x00b', '']"), + ("pa.large_string()", "['hello', None, 'a\\x00b', '']"), + ("pa.binary()", "[b'hello', None, b'a\\x00b', b'']"), + ("pa.large_binary()", "[b'hello', None, b'a\\x00b', b'']"), + ] { + let arrow_type = evaluate(py, &locals, data_type)?; + let values = evaluate(py, &locals, values)?.cast_into::()?; + let expected = reference_results(py, &arrow_type, &values, true)?; + let actual = build_results(py, &arrow_type, &values, true)?; + assert_eq!(actual.pyarrow_values, 0, "{data_type}"); + assert!(!actual.retried, "{data_type}"); + assert_array_equal(&actual.array, &expected)?; + } + Ok(()) + }) + .unwrap(); +} + +#[test] +fn subclass_results_use_fallback_and_match_pyarrow() { + Python::initialize(); + Python::attach(|py| -> PyResult<()> { + let locals = context(py)?; + py.run( + c"class IntSubclass(int): pass\nclass StrSubclass(str): pass", + Some(&locals), + Some(&locals), + )?; + // For variable-width columns, the capacity check cannot prove that a subclass fits. + // It sends the prefix and suffix to PyArrow together, rather than concatenating them. + for (data_type, values, pyarrow_values) in [ + ("pa.int64()", "[1, IntSubclass(2), None, 3]", 3), + ("pa.string()", "['a', StrSubclass('b'), None, 'c']", 4), + ] { + let arrow_type = evaluate(py, &locals, data_type)?; + let values = evaluate(py, &locals, values)?.cast_into::()?; + let expected = reference_results(py, &arrow_type, &values, true)?; + let actual = build_results(py, &arrow_type, &values, true)?; + assert_eq!(actual.pyarrow_values, pyarrow_values, "{data_type}"); + assert!(!actual.retried); + assert_array_equal(&actual.array, &expected)?; + } + Ok(()) + }) + .unwrap(); +} + +#[test] +fn overflowing_results_follow_pyarrow_safe_and_unsafe_retry() { + Python::initialize(); + Python::attach(|py| -> PyResult<()> { + let locals = context(py)?; + for (data_type, values) in [ + ("pa.int8()", "[1, 128]"), + ("pa.int64()", "[1, 2**63]"), + ("pa.float32()", "[1, 2**24 + 1]"), + ("pa.float64()", "[1, 2**53 + 1]"), + ] { + let arrow_type = evaluate(py, &locals, data_type)?; + let values = evaluate(py, &locals, values)?.cast_into::()?; + for safe in [true, false] { + let expected = reference_results(py, &arrow_type, &values, safe); + let actual = build_results(py, &arrow_type, &values, safe); + match (actual, expected) { + (Ok(actual), Ok(expected)) => { + assert!(actual.pyarrow_values > 0, "{data_type}"); + assert_array_equal(&actual.array, &expected)?; + } + (Err(actual), Err(expected)) => { + assert!( + actual.get_type(py).is(expected.get_type(py)), + "{data_type}, safe={safe}: {actual}; expected {expected}" + ); + } + (actual, expected) => panic!( + "{data_type}, safe={safe}: actual success {}, expected success {}", + actual.is_ok(), + expected.is_ok() + ), + } + } + } + Ok(()) + }) + .unwrap(); +} + +#[test] +fn held_mutable_results_are_read_at_finish_and_extend_copies_the_list() { + Python::initialize(); + Python::attach(|py| -> PyResult<()> { + let locals = context(py)?; + let arrow_type = evaluate(py, &locals, "pa.binary()")?; + let values = evaluate(py, &locals, "[b'prefix', bytearray(b'b'), b'tail']")? + .cast_into::()?; + let spec = primitive_spec(py, &arrow_type, true)?; + let mut column = ResultColumn::new(&spec, 4); + column.extend(py, &values); + values.get_item(1)?.set_item(0, 99)?; + values.set_item(2, PyBytes::new(py, b"replacement"))?; + column.push(py, py.None().into_bound(py)); + let actual = column.finish(py)?; + let expected = evaluate(py, &locals, "pa.array([b'prefix', b'c', b'tail', None])")?; + assert_array_equal(&actual.array, &expected) + }) + .unwrap(); +} + +#[test] +fn lowered_variable_width_capacity_falls_back_for_the_whole_column() { + Python::initialize(); + Python::attach(|py| -> PyResult<()> { + let locals = context(py)?; + let arrow_type = evaluate(py, &locals, "pa.string()")?; + let values = evaluate(py, &locals, "['ab', 'cd', None]")?.cast_into::()?; + let spec = primitive_spec(py, &arrow_type, true)?.with_value_capacity(3); + let mut column = ResultColumn::new(&spec, values.len()); + column.extend(py, &values); + let actual = column.finish(py)?; + assert_eq!(actual.pyarrow_values, 3); + assert!(!actual.retried); + let expected = reference_results(py, &arrow_type, &values, true)?; + assert_array_equal(&actual.array, &expected) + }) + .unwrap(); +} + +#[test] +fn input_columns_match_pyarrow_cells_for_supported_primitive_types() { + Python::initialize(); + Python::attach(|py| -> PyResult<()> { + let locals = context(py)?; + let none = py.None().into_bound(py); + let convert_input = evaluate(py, &locals, "lambda column, converter: column.to_pylist()")?; + for (data_type, values) in [ + ("pa.int8()", "[-128, None, 127]"), + ("pa.int16()", "[-32768, None, 32767]"), + ("pa.int32()", "[-(2**31), None, 2**31 - 1]"), + ("pa.int64()", "[-(2**63), None, 2**63 - 1]"), + ("pa.float32()", "[-1.25, None, 2.5]"), + ("pa.float64()", "[-1.25, None, 2.5]"), + ("pa.bool_()", "[True, None, False]"), + ("pa.string()", "['hello', None, 'a\\x00b']"), + ("pa.large_string()", "['hello', None, 'a\\x00b']"), + ( + "pa.date32()", + "[datetime.date.min, None, datetime.date.max]", + ), + ] { + let array = evaluate( + py, + &locals, + &format!("pa.array({values}, type={data_type})"), + )?; + let expected = array.call_method0("to_pylist")?; + let rust_type = import_data_type(&array.getattr("type")?)?; + let spec = input_spec(&rust_type, &none, &none, &none)? + .expect("the primitive type must be supported"); + let batch = py + .import("pyarrow")? + .call_method1("record_batch", (PyList::new(py, [array])?, vec!["x"]))?; + let input = InputColumns::import(&batch, &[spec], &convert_input)?; + assert_eq!(input.num_rows(), 3); + drop(batch); + for row in 0..input.num_rows() { + let actual = input.value(py, 0, row)?; + let expected = expected.get_item(row)?; + assert!(actual.eq(&expected)?, "{data_type}, row={row}"); + assert!( + actual.get_type().is(expected.get_type()), + "{data_type}, row={row}" + ); + } + } + Ok(()) + }) + .unwrap(); +} + +#[test] +fn binary_input_policy_selects_bytes_or_bytearray() { + Python::initialize(); + Python::attach(|py| -> PyResult<()> { + let locals = context(py)?; + let none = py.None().into_bound(py); + // Binary policy belongs to the caller: provide its type identity and converter. + let binary_class = evaluate(py, &locals, "type('BinaryType', (), {})")?; + let binary_type = binary_class.call0()?; + let sql_types = py.import("types")?.getattr("SimpleNamespace")?.call0()?; + sql_types.setattr("BinaryType", binary_class)?; + let convert_input = evaluate( + py, + &locals, + "lambda column, converter: [converter(v) for v in column.to_pylist()]", + )?; + for data_type in ["pa.binary()", "pa.large_binary()"] { + let array = evaluate( + py, + &locals, + &format!("pa.array([b'hello', None, b'a\\x00b'], type={data_type})"), + )?; + let expected = array.call_method0("to_pylist")?; + let rust_type = import_data_type(&array.getattr("type")?)?; + assert!(input_spec(&rust_type, &none, &none, &none)?.is_err()); + let batch = py + .import("pyarrow")? + .call_method1("record_batch", (PyList::new(py, [array])?, vec!["x"]))?; + for policy in [ + "lambda v: v", + "lambda v: None if v is None else bytearray(v)", + ] { + let converter = evaluate(py, &locals, policy)?; + let spec = input_spec(&rust_type, &binary_type, &converter, &sql_types)? + .expect("both configured binary policies are supported"); + let input = InputColumns::import(&batch, &[spec], &convert_input)?; + for row in 0..input.num_rows() { + let actual = input.value(py, 0, row)?; + let expected = converter.call1((expected.get_item(row)?,))?; + assert!(actual.eq(&expected)?, "{data_type}, {policy}, row={row}"); + assert!( + actual.get_type().is(expected.get_type()), + "{data_type}, {policy}, row={row}" + ); + } + } + } + Ok(()) + }) + .unwrap(); +} + +#[test] +fn input_columns_reject_mismatched_specs_and_out_of_range_cells() { + Python::initialize(); + Python::attach(|py| -> PyResult<()> { + let locals = context(py)?; + let none = py.None().into_bound(py); + let convert_input = evaluate(py, &locals, "lambda column, converter: column.to_pylist()")?; + let batch = evaluate( + py, + &locals, + "pa.record_batch([pa.array([1, 2], type=pa.int64())], names=['x'])", + )?; + let wrong = input_spec(&DataType::Int32, &none, &none, &none)?.expect("int32 is supported"); + assert!(InputColumns::import(&batch, &[wrong], &convert_input).is_err()); + assert!(InputColumns::import(&batch, &[], &convert_input).is_err()); + let right = input_spec(&DataType::Int64, &none, &none, &none)?.expect("int64 is supported"); + let columns = InputColumns::import(&batch, &[right], &convert_input)?; + for (column, row) in [(1, 0), (0, 2), (usize::MAX, 0), (0, usize::MAX)] { + let error = columns.value(py, column, row).unwrap_err(); + assert!(error.is_instance_of::(py)); + } + assert_eq!(columns.value(py, 0, 1)?.extract::()?, 2); + Ok(()) + }) + .unwrap(); +} diff --git a/dev/design/architecture.md b/dev/design/architecture.md index 87303768..cabd8ce4 100644 --- a/dev/design/architecture.md +++ b/dev/design/architecture.md @@ -47,6 +47,8 @@ crates/ column, plan, dataframe, functions, session, readwriter, catalog, group, window, udf/udtf, streaming, observation pyspark-rs PyO3 bindings -> native module `_pyspark` + spark-python-arrow shared Python/Arrow bridge, scalar conversion, and + explicit Python UDF conversion policy; no transport python/ pyspark/… thin drop-in package; public modules re-export from `_pyspark` so `from pyspark.sql import SparkSession` works @@ -67,9 +69,19 @@ 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. + +The shared crate also owns the scalar routines used by the Python binding's +input conversion and row materialization, and an explicit Python UDF conversion +policy for worker consumers. It does not own the client `Row` representation, +schema inference, transport, or a worker's execution loop. `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 +[shared converter contract](../../crates/spark-python-arrow/README.md). ## UDF strategy (full compatibility) diff --git a/python/tests/test_arrow_handoff.py b/python/tests/test_arrow_handoff.py new file mode 100644 index 00000000..dac57b7e --- /dev/null +++ b/python/tests/test_arrow_handoff.py @@ -0,0 +1,506 @@ +# +# 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())), + ] + ) + 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"]) + ), + 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_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 From d1b9f80e95863b21239935c142fbbe8bf12ffe79 Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Tue, 27 Oct 2026 12:42:48 -0700 Subject: [PATCH 2/8] Cover shared converter policies and preserve Arrow schema flags Add UDF converter, fallback, mutation and error-ordering regressions, retain converted capacity fallbacks, and preserve sorted-map and dictionary-order metadata. Document nullable callback and buffer ownership contracts and add isolated PyO3 compatibility CI. Generated-by: OpenAI Codex Co-authored-by: Isaac --- .github/workflows/python-arrow.yml | 73 ++ crates/pyspark-rs/src/dataframe.rs | 3 +- crates/spark-python-arrow/Cargo.toml | 18 +- crates/spark-python-arrow/README.md | 80 +- crates/spark-python-arrow/src/bridge.rs | 241 +++++- crates/spark-python-arrow/src/lib.rs | 21 +- crates/spark-python-arrow/src/scalar.rs | 32 +- crates/spark-python-arrow/src/udf/input.rs | 37 +- crates/spark-python-arrow/src/udf/mod.rs | 37 +- crates/spark-python-arrow/src/udf/result.rs | 185 ++++- .../tests/bridge_regressions.rs | 307 ++++++++ crates/spark-python-arrow/tests/udf_inputs.rs | 294 ++++++++ .../spark-python-arrow/tests/udf_results.rs | 707 ++++++++++++++++++ python/tests/test_arrow_handoff.py | 19 +- 14 files changed, 1928 insertions(+), 126 deletions(-) create mode 100644 .github/workflows/python-arrow.yml create mode 100644 crates/spark-python-arrow/tests/bridge_regressions.rs create mode 100644 crates/spark-python-arrow/tests/udf_inputs.rs create mode 100644 crates/spark-python-arrow/tests/udf_results.rs diff --git a/.github/workflows/python-arrow.yml b/.github/workflows/python-arrow.yml new file mode 100644 index 00000000..73c87a7e --- /dev/null +++ b/.github/workflows/python-arrow.yml @@ -0,0 +1,73 @@ +# +# 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 compatibility + +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 (PyO3 ${{ matrix.pyo3 }}) + runs-on: ubuntu-latest + timeout-minutes: 20 + strategy: + fail-fast: false + matrix: + pyo3: ["0.28.3", "0.29.0"] + env: + SPARK_PYTHON_ARROW_TEST_PYSPARK: "1" + 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 reference Python converters + run: python -m pip install "pyspark-client==4.2.0" "pyarrow==26.0.0" + - name: Isolate the standalone crate and select PyO3 + # The workspace lockfile pins 0.28.3. Build outside it so each matrix leg + # actually compiles the converter against the requested minor version. + run: | + cp -R crates/spark-python-arrow "$RUNNER_TEMP/spark-python-arrow" + cargo generate-lockfile --manifest-path "$RUNNER_TEMP/spark-python-arrow/Cargo.toml" + cargo update --manifest-path "$RUNNER_TEMP/spark-python-arrow/Cargo.toml" \ + -p pyo3 --precise "${{ matrix.pyo3 }}" + cargo tree --manifest-path "$RUNNER_TEMP/spark-python-arrow/Cargo.toml" --locked -i pyo3 + - name: Test native and official PySpark conversion semantics + run: | + cargo test --manifest-path "$RUNNER_TEMP/spark-python-arrow/Cargo.toml" \ + --locked --features pyo3/abi3-py39 + - name: Check converter Clippy + run: | + cargo clippy --manifest-path "$RUNNER_TEMP/spark-python-arrow/Cargo.toml" \ + --locked --all-targets --features pyo3/abi3-py39 -- -D warnings diff --git a/crates/pyspark-rs/src/dataframe.rs b/crates/pyspark-rs/src/dataframe.rs index c409c780..25529563 100644 --- a/crates/pyspark-rs/src/dataframe.rs +++ b/crates/pyspark-rs/src/dataframe.rs @@ -1167,7 +1167,8 @@ impl PyDataFrame { /// /// Mirrors `pyspark.sql.DataFrame.toArrow()`: the result is decoded from the /// server's Arrow IPC stream and handed to `pyarrow` through the shared - /// converter, without copying the decoded Arrow data buffers. + /// 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> { let batches = py diff --git a/crates/spark-python-arrow/Cargo.toml b/crates/spark-python-arrow/Cargo.toml index e7383b6b..9ce9d119 100644 --- a/crates/spark-python-arrow/Cargo.toml +++ b/crates/spark-python-arrow/Cargo.toml @@ -1,4 +1,19 @@ -# SPDX-License-Identifier: Apache-2.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. [package] name = "spark-python-arrow" version = "0.1.0" @@ -10,6 +25,7 @@ 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. +# Compatibility CI resolves each supported minor independently of the workspace lockfile. pyo3 = { version = ">=0.28, <0.30" } [lints.rust] diff --git a/crates/spark-python-arrow/README.md b/crates/spark-python-arrow/README.md index 5b206c21..0e70460c 100644 --- a/crates/spark-python-arrow/README.md +++ b/crates/spark-python-arrow/README.md @@ -1,4 +1,21 @@ - + # Shared Python and Arrow conversion @@ -31,22 +48,56 @@ date-result columns retain the caller's Python fallback. Unsupported conversion is `Ok(Err(reason))`; an error inspecting a supported conversion is `Err(PyErr)`. Callers must not turn execution errors into unsupported-type fallback. +### UDF converter contract + +The result policy accepts the stock, pure PySpark converters for **nullable +top-level UDF results**, not arbitrary callbacks or converters for non-nullable +fields. It can bypass those converters for `None` and exact strings/bytes whose +conversion is an identity. The Spark type check does not prove a callback has +these semantics. Callers must supply the matching PySpark converter; no probe is +performed because invoking a callback during policy selection can change error +timing and cannot prove its behavior for every value. + +Run the UDF for every row before calling `ResultColumn::finish`. Mutable results +are retained until then, and conversion errors are deferred so a later UDF error +still takes precedence. Once a result has been converted, a capacity fallback +retains that converted value instead of invoking its converter again. + ## Ownership and caller responsibilities The library does not initialize Python, acquire a process-global interpreter, create threads, or retain session state. The consumer chooses embedding versus extension-module linking and owns all Python callbacks and their interpreter. Array imports accept actual PyArrow objects and use their native exporters. -Bare extension types/arrays are rejected because Arrow's bare type cannot retain -their extension metadata. Schemas and record batches retain it on their fields; -UDF extension results select the caller's fallback. +Bare extension types/arrays and bare ordered dictionaries are rejected because +Arrow's bare type cannot retain their field metadata or ordering flag. Schemas +and record batches retain these properties on their fields. Extension types in +dictionary-value positions are rejected, even within schemas, because that +position has no metadata-owning Rust field. UDF extension results select the +caller's fallback. As with the Arrow C Data Interface itself, producers are responsible for valid -buffer contents. Imports check structural invariants without scanning payloads. - -Zero-copy describes compatible Arrow **data buffers**, not schema descriptors, -Python scalar/row objects, pandas materialization, IPC decoding, or fallback -conversion. Import/export release callbacks retain the buffers after either -language's temporary wrapper has been dropped. +buffer contents and all offsets. Imports perform limited structural checks, not +full Arrow validation or a scan of every interior offset. A real PyArrow class +does not by itself prove complete buffer validity. CPython's strict UTF-8 decoder +checks the selected string bytes, not the Arrow offsets used to select them. + +Imports call native `pyarrow.lib` class descriptors to avoid dispatching an +arbitrary object's exporter. Exports and Python fallback operations use the +public `pyarrow` aliases, allowing normal Python instrumentation, including the +ownership/failure-injection tests. This is not isolation from arbitrary Python +code running in the same interpreter. + +Zero-copy describes compatible Arrow **value and offset buffers**, not schema +descriptors, Python scalar/row objects, pandas materialization, IPC decoding, or +fallback conversion. Sliced arrays whose value and null offsets differ can +require copying the validity bitmap during C Data export. Import/export release +callbacks retain shared buffers after either language's temporary wrapper has +been dropped. + +UDF result builders allocate new Arrow buffers. A mixed native/fallback result +uses `pyarrow.concat_arrays` to join the prefix and suffix into the single array +the worker expects; that concatenation copies the combined column. Only the +handoff of a completed, compatible native array shares its buffers. The client's owned `Value` and `Row` types remain in the Python-free Rust client. The Python binding adapts their scalar fields to `scalar::Scalar`; its row class, @@ -61,6 +112,10 @@ enable PyO3's `extension-module` feature or select a minimum Python ABI: a Pytho extension chooses those features in the host package, while an embedded worker links its chosen Python runtime. PyO3 must resolve to one compatible version in each consuming dependency graph. +The manifest permits PyO3 0.28 and 0.29 so consumers on either line can use the +same source. The workspace lockfile selects 0.28.3 and does not validate 0.29. +The isolated converter compatibility CI job tests both versions independently; +0.29 support must not be considered verified until that job passes. Bazel consumers should pin the OSS revision or release using their existing third-party Rust dependency mechanism and depend on this library target. There @@ -83,6 +138,11 @@ 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. +With official `pyspark-client==4.2.0` also installed, set +`SPARK_PYTHON_ARROW_TEST_PYSPARK=1` to enable comparisons against its nullable +String, Binary, and Variant result converters. The compatibility CI matrix sets +this flag; the default local command skips that optional comparison. + The standalone tests cover buffer identity and lifetime, schema/type round trips, empty shapes, native scalar behavior, and UDF fallback/error ordering. The Python client's `test_arrow_handoff.py` exercises the public `toArrow()` and diff --git a/crates/spark-python-arrow/src/bridge.rs b/crates/spark-python-arrow/src/bridge.rs index be8bc9a3..a1877217 100644 --- a/crates/spark-python-arrow/src/bridge.rs +++ b/crates/spark-python-arrow/src/bridge.rs @@ -19,15 +19,16 @@ //! //! Imports accept real PyArrow objects, not arbitrary objects implementing an //! export method. PyArrow remains responsible for the validity of its C data. -//! These functions check structural invariants but do not scan values or change -//! when a caller observes conversion errors such as invalid string contents. +//! These functions perform limited structural checks, not full validation of all +//! offsets, dictionary indices, or payloads. Valid offsets and buffer bounds are +//! producer preconditions. String decoding errors remain observable at conversion. use std::sync::Arc; use arrow::array::{make_array, Array, ArrayData, ArrayRef, StructArray}; -use arrow::datatypes::{DataType, Field, Schema}; +use arrow::datatypes::{DataType, Field, Schema, UnionFields}; use arrow::error::ArrowError; -use arrow::ffi::{from_ffi, FFI_ArrowArray, FFI_ArrowSchema}; +use arrow::ffi::{from_ffi_and_data_type, FFI_ArrowArray, FFI_ArrowSchema}; use arrow::ipc::writer::FileWriter; use arrow::record_batch::{RecordBatch, RecordBatchOptions}; use pyo3::exceptions::{PyRuntimeError, PyTypeError, PyValueError}; @@ -41,36 +42,39 @@ pub fn import_schema(value: &Bound<'_, PyAny>) -> PyResult { let mut schema = FFI_ArrowSchema::empty(); let address = &mut schema as *mut FFI_ArrowSchema as usize; class.call_method1("_export_to_c", (value, address))?; - Schema::try_from(&schema).map_err(arrow_error) + import_c_schema(&schema).map_err(arrow_error) } /// Import a PyArrow data type. /// -/// Root extension types require a schema or record batch: a bare Rust `DataType` -/// has nowhere to retain their extension metadata, so this function rejects them. +/// Root extension types and ordered dictionaries require a schema or record batch: +/// a bare Rust `DataType` has nowhere to retain these field properties. pub fn import_data_type(value: &Bound<'_, PyAny>) -> PyResult { let class = require_pyarrow_class(value, "DataType")?; reject_root_extension_type(value)?; + reject_root_dictionary_order(value)?; check_python_data_type(value)?; let mut schema = FFI_ArrowSchema::empty(); let address = &mut schema as *mut FFI_ArrowSchema as usize; class.call_method1("_export_to_c", (value, address))?; - DataType::try_from(&schema).map_err(arrow_error) + import_c_data_type(&schema).map_err(arrow_error) } /// Import a PyArrow array without copying its data buffers. /// -/// Root extension arrays are rejected because an `ArrayRef` cannot retain their -/// extension metadata. Import a record batch to preserve that field metadata. +/// Root extensions and ordered dictionaries are rejected because an `ArrayRef` +/// cannot retain their field properties. Import a record batch to preserve them. pub fn import_array(value: &Bound<'_, PyAny>) -> PyResult { let class = require_pyarrow_class(value, "Array")?; let data_type = value.getattr("type")?; reject_root_extension_type(&data_type)?; + reject_root_dictionary_order(&data_type)?; check_python_data_type(&data_type)?; let (array, schema) = export_python_array(value, &class)?; + let data_type = import_c_data_type(&schema).map_err(arrow_error)?; // SAFETY: the base PyArrow exporter, not an arbitrary Python method, has // initialized both descriptors and transferred ownership of valid C data. - let data = unsafe { from_ffi(array, &schema) }.map_err(arrow_error)?; + let data = unsafe { from_ffi_and_data_type(array, data_type) }.map_err(arrow_error)?; validate_structure(&data)?; Ok(make_array(data)) } @@ -80,9 +84,10 @@ pub fn import_batch(value: &Bound<'_, PyAny>) -> PyResult { let class = require_pyarrow_class(value, "RecordBatch")?; check_python_schema(&value.getattr("schema")?)?; let (array, ffi_schema) = export_python_array(value, &class)?; - let schema = Schema::try_from(&ffi_schema).map_err(arrow_error)?; + let schema = import_c_schema(&ffi_schema).map_err(arrow_error)?; + let data_type = DataType::Struct(schema.fields().clone()); // SAFETY: only the real PyArrow RecordBatch exporter writes the descriptors. - let data = unsafe { from_ffi(array, &ffi_schema) }.map_err(arrow_error)?; + let data = unsafe { from_ffi_and_data_type(array, data_type) }.map_err(arrow_error)?; validate_structure(&data)?; if !matches!(data.data_type(), DataType::Struct(_)) || data.null_count() != 0 { return Err(PyValueError::new_err( @@ -101,7 +106,9 @@ pub fn import_batch(value: &Bound<'_, PyAny>) -> PyResult { .map_err(arrow_error) } -/// Export an Arrow array without copying its data buffers. +/// Export an Arrow array sharing its value and offset buffers. +/// +/// Arrow may copy a sliced validity bitmap to align its offset with the data. pub fn export_array(py: Python<'_>, array: ArrayRef) -> PyResult> { let array_type = py.import("pyarrow")?.getattr("Array")?; Ok(export_array_with_type(&array_type, array)?.unbind()) @@ -112,7 +119,7 @@ pub(crate) fn export_array_with_type<'py>( array: ArrayRef, ) -> PyResult> { check_c_data_type(array.data_type())?; - let mut schema = FFI_ArrowSchema::try_from(array.data_type()).map_err(arrow_error)?; + let mut schema = export_c_data_type(array.data_type()).map_err(arrow_error)?; let schema_address = &mut schema as *mut FFI_ArrowSchema as usize; let mut array = FFI_ArrowArray::new(&array.to_data()); let array_address = &mut array as *mut FFI_ArrowArray as usize; @@ -127,7 +134,10 @@ pub(crate) fn export_array_with_type<'py>( /// returns an error for such schemas; table export provides an IPC fallback. pub fn export_schema(py: Python<'_>, schema: &Schema) -> PyResult> { check_c_schema(schema)?; - let mut schema = FFI_ArrowSchema::try_from(schema).map_err(arrow_error)?; + 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")? @@ -139,7 +149,7 @@ pub fn export_schema(py: Python<'_>, schema: &Schema) -> PyResult> { /// Export an Arrow data type, rejecting strings not representable in C schemas. pub fn export_data_type(py: Python<'_>, data_type: &DataType) -> PyResult> { check_c_data_type(data_type)?; - let mut schema = FFI_ArrowSchema::try_from(data_type).map_err(arrow_error)?; + let mut schema = export_c_data_type(data_type).map_err(arrow_error)?; let address = &mut schema as *mut FFI_ArrowSchema as usize; Ok(py .import("pyarrow")? @@ -152,12 +162,15 @@ pub fn export_data_type(py: Python<'_>, data_type: &DataType) -> PyResult, batches: Vec) -> PyResult> { let schema = batches .first() .map(|batch| batch.schema()) .unwrap_or_else(|| Arc::new(Schema::empty())); - if batches.iter().any(|batch| batch.schema_ref() != &schema) { + 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", )); @@ -225,6 +238,186 @@ fn reject_root_extension_type(data_type: &Bound<'_, PyAny>) -> PyResult<()> { Ok(()) } +fn reject_root_dictionary_order(data_type: &Bound<'_, PyAny>) -> PyResult<()> { + if let Ok(ordered) = data_type.getattr("ordered") { + if ordered.is_truthy()? { + return Err(PyTypeError::new_err( + "Ordered Arrow dictionaries require a schema or record batch to preserve ordering", + )); + } + } + Ok(()) +} + +// 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 import_c_schema(schema: &FFI_ArrowSchema) -> Result { + let DataType::Struct(fields) = import_c_data_type(schema)? else { + return Err(ArrowError::CDataInterface( + "A schema must be a struct type".into(), + )); + }; + Ok(Schema::new(fields).with_metadata(schema.metadata()?)) +} + +fn import_c_data_type(schema: &FFI_ArrowSchema) -> Result { + fn has_ordered_dictionary(schema: &FFI_ArrowSchema) -> bool { + schema.dictionary_ordered() + || schema.children().any(has_ordered_dictionary) + || schema.dictionary().is_some_and(has_ordered_dictionary) + } + + // Arrow 58 drops DICTIONARY_ORDERED when importing Fields. Restore it at + // every field, including those nested in array types and dictionary values. + fn restore(data_type: DataType, schema: &FFI_ArrowSchema) -> Result { + let field = |field: &Field, index: usize| -> Result, ArrowError> { + let child = schema.child(index); + Ok(Arc::new( + field + .clone() + .with_data_type(restore(field.data_type().clone(), child)?) + .with_dict_is_ordered(child.dictionary_ordered()), + )) + }; + Ok(match data_type { + DataType::List(child) => DataType::List(field(&child, 0)?), + DataType::LargeList(child) => DataType::LargeList(field(&child, 0)?), + DataType::ListView(child) => DataType::ListView(field(&child, 0)?), + DataType::LargeListView(child) => DataType::LargeListView(field(&child, 0)?), + DataType::FixedSizeList(child, size) => { + DataType::FixedSizeList(field(&child, 0)?, size) + } + DataType::Map(child, sorted) => DataType::Map(field(&child, 0)?, sorted), + DataType::Struct(fields) => DataType::Struct( + fields + .iter() + .enumerate() + .map(|(index, child)| field(child, index)) + .collect::, _>>()? + .into(), + ), + DataType::Union(fields, mode) => DataType::Union( + UnionFields::try_new( + fields.iter().map(|(id, _)| id), + fields + .iter() + .enumerate() + .map(|(index, (_, child))| field(child, index)) + .collect::, _>>()?, + )?, + mode, + ), + DataType::RunEndEncoded(run_ends, values) => { + DataType::RunEndEncoded(field(&run_ends, 0)?, field(&values, 1)?) + } + DataType::Dictionary(keys, values) => DataType::Dictionary( + keys, + Box::new(restore( + *values, + schema.dictionary().ok_or_else(|| { + ArrowError::CDataInterface("A dictionary must have a value type".into()) + })?, + )?), + ), + other => other, + }) + } + + let data_type = DataType::try_from(schema)?; + if has_ordered_dictionary(schema) { + restore(data_type, schema) + } else { + Ok(data_type) + } +} + fn require_pyarrow_class<'py>( value: &Bound<'py, PyAny>, name: &str, @@ -252,8 +445,8 @@ fn export_python_array( fn validate_structure(data: &ArrayData) -> PyResult<()> { // Arrow's inexpensive validation recurses through child descriptors. It - // checks sizes and offset endpoints, not all offsets, UTF-8 contents, or - // dictionary indices whose conversion errors callers may defer. + // checks sizes and offset endpoints, not every offset or dictionary index; + // those remain producer preconditions. UTF-8 is decoded when a value is read. data.validate().map_err(arrow_error) } @@ -331,6 +524,14 @@ fn check_python_data_type(data_type: &Bound<'_, PyAny>) -> PyResult<()> { if num_fields == 0 { for name in ["value_type", "storage_type"] { if let Ok(child) = data_type.getattr(name) { + if name == "value_type" && is_extension_type(&child)? { + return Err(PyTypeError::new_err( + "Arrow dictionary value extension types cannot preserve their metadata", + )); + } + if name == "value_type" { + reject_root_dictionary_order(&child)?; + } check_python_data_type(&child)?; } } diff --git a/crates/spark-python-arrow/src/lib.rs b/crates/spark-python-arrow/src/lib.rs index 8c2045f9..946b1f58 100644 --- a/crates/spark-python-arrow/src/lib.rs +++ b/crates/spark-python-arrow/src/lib.rs @@ -1,9 +1,24 @@ -// SPDX-License-Identifier: Apache-2.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. //! Python/Arrow conversion without a dependency on Spark transport or worker execution. //! -//! The bridge shares Arrow data buffers; scalar and row materialization still allocates Python -//! objects. The caller owns its interpreter, Python classes, session settings, and error policy. +//! The bridge shares compatible Arrow value buffers; scalar and row materialization still +//! allocates Python objects. The caller owns its interpreter, Python classes, and error policy. //! UDF-specific coercion and end-of-batch result conversion are explicit in the [`udf`] API. pub mod bridge; diff --git a/crates/spark-python-arrow/src/scalar.rs b/crates/spark-python-arrow/src/scalar.rs index 8b287428..36c2fb92 100644 --- a/crates/spark-python-arrow/src/scalar.rs +++ b/crates/spark-python-arrow/src/scalar.rs @@ -1,19 +1,19 @@ -/* - * 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. - */ +// 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. //! Scalar conversion shared by client row materialization and UDF columns. //! diff --git a/crates/spark-python-arrow/src/udf/input.rs b/crates/spark-python-arrow/src/udf/input.rs index ce54bb2d..0349e822 100644 --- a/crates/spark-python-arrow/src/udf/input.rs +++ b/crates/spark-python-arrow/src/udf/input.rs @@ -1,19 +1,19 @@ -/* - * 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. - */ +// 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 arrow::array::{ Array, ArrayRef, AsArray, BinaryArray, BooleanArray, Date32Array, Float32Array, Float64Array, @@ -363,8 +363,9 @@ impl Column { value().into_python(py) } } - /// Decodes the slot's bytes rather than reading a `&str`: arrow-rs imports strings - /// through the C Data Interface without checking them, so CPython's decode is the check. + /// Decodes the slot's bytes rather than forming a `&str` from unvalidated UTF-8. + /// CPython validates only the UTF-8 payload. Offsets and buffer ranges remain the + /// Arrow producer's responsibility; this is not validation of arbitrary C data. fn string<'py, O: OffsetSizeTrait>( py: Python<'py>, array: &GenericStringArray, diff --git a/crates/spark-python-arrow/src/udf/mod.rs b/crates/spark-python-arrow/src/udf/mod.rs index dc481466..950eef53 100644 --- a/crates/spark-python-arrow/src/udf/mod.rs +++ b/crates/spark-python-arrow/src/udf/mod.rs @@ -1,19 +1,19 @@ -/* - * 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. - */ +// 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 <-> Python conversion one value at a time, for a per-row UDF loop that has to match the //! `SQL_ARROW_BATCHED_UDF` loop in `pyspark.worker.read_udfs`. @@ -81,8 +81,9 @@ fn import_batch(batch: &Bound<'_, PyAny>) -> PyResult { Ok(StructArray::from(crate::bridge::import_batch(batch)?)) } -/// Hands `array` to pyarrow, whose `pa.Array` is `array_type`, without copying it. pyarrow moves -/// both structs' contents out, so dropping them afterwards releases nothing. +/// Hands `array` to PyArrow, whose `pa.Array` is `array_type`, sharing compatible value buffers. +/// Validity-bitmap alignment may copy. PyArrow moves both descriptors' contents out, so dropping +/// them afterwards releases nothing. fn export_array<'py>( array_type: &Bound<'py, PyAny>, array: ArrayRef, diff --git a/crates/spark-python-arrow/src/udf/result.rs b/crates/spark-python-arrow/src/udf/result.rs index 30a87516..8100f36c 100644 --- a/crates/spark-python-arrow/src/udf/result.rs +++ b/crates/spark-python-arrow/src/udf/result.rs @@ -1,19 +1,19 @@ -/* - * 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. - */ +// 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::sync::Arc; @@ -47,6 +47,15 @@ const MAX_EXACT_F32_INT: u64 = 1 << 24; /// comes from. For supported Variant results, `variant_class` resolves the class UDF code /// constructs, which need not be the backend module's class. The resolver runs only after the /// Arrow type and converter are accepted; its errors propagate. No package or class is inferred. +/// +/// # Converter contract +/// +/// `converter` must be PySpark's pure, top-level nullable result converter (as returned by +/// `LocalDataToArrowConversion._create_converter(..., nullable=True)`), or `None` for identity. +/// Non-nullable field converters and arbitrary stateful callbacks are not supported: this policy +/// keeps `None` unchanged and may bypass the converter for exact strings, bytes, and Variant +/// parts, or convert immutable built-ins before the remaining UDF calls. No callback is probed to +/// validate this contract, since probing would change its invocation and error timing. pub fn result_spec<'py>( py: Python<'py>, arrow_type: &Bound<'_, PyAny>, @@ -63,6 +72,42 @@ pub fn result_spec<'py>( "Arrow extension types require the caller's Python converter".to_owned(), )); } + let arrow_lib = py.import("pyarrow.lib")?; + if arrow_type.is_instance(&arrow_lib.getattr("DictionaryType")?)? { + return Ok(Err( + "Arrow dictionary types require the caller's Python converter".to_owned(), + )); + } + if arrow_type.is_instance(&arrow_lib.getattr("DataType")?)? { + let fields = arrow_type.getattr("num_fields")?.extract::()?; + if fields > 0 { + // Only the tagged Variant struct can use a nested native builder. Select the + // caller's fallback before importing other nested types, whose child metadata + // may not be representable by a bare Rust DataType (dictionary extensions, etc.). + let mut variant_candidate = + fields == 2 && arrow_type.is_instance(&arrow_lib.getattr("StructType")?)?; + if variant_candidate { + let binary = arrow_lib.call_method0("binary")?; + for index in 0..fields { + let child = arrow_type + .call_method1("field", (index,))? + .getattr("type")?; + if !child + .call_method1("equals", (&binary,))? + .extract::()? + { + variant_candidate = false; + break; + } + } + } + if !variant_candidate { + return Ok(Err(format!( + "{arrow_type} is not supported by the native UDF result converter" + ))); + } + } + } let data_type = import_data_type(arrow_type)?; let Some(kind) = OutputKind::of(&data_type) else { return Ok(Err(format!( @@ -559,7 +604,8 @@ pub struct ResultColumn<'s> { enum State { Building, - /// The results from the first one the builder could not take, unconverted, in row order: + /// The results from the first one the builder could not take, in row order. The first may be + /// an already-converted string that exceeded capacity; the remainder is unconverted: /// those in `listed`, then those in `pushed`. `listed` has the held part of each list given /// to `extend` and the results pushed before the last such call; `pushed` has the results /// pushed after it. @@ -596,14 +642,14 @@ impl<'s> ResultColumn<'s> { /// Adds the next row's result. A built-in the builder takes is converted and appended now, as /// it cannot change; from the first result that is not, every result is held for /// [`ResultColumn::finish`]. Never raises: a converter error is raised by `finish`. - pub fn push<'py>(&mut self, py: Python<'py>, value: Bound<'py, PyAny>) { + pub fn push<'py>(&mut self, py: Python<'py>, mut value: Bound<'py, PyAny>) { match &mut self.state { State::Building => {} State::Holding { pushed, .. } => return pushed.push(value.unbind()), State::Failed(_) => return, } if self.builder.append_exact(&value, self.spec.value_capacity) - || self.append_converted(py, &value) + || self.append_converted(py, &mut value) { return; } @@ -633,34 +679,35 @@ impl<'s> ResultColumn<'s> { py: Python<'py>, values: &Bound<'py, PyList>, ) -> PyResult<()> { - let mut appended = 0; // The iterator protocol reads a list with one call per item, and the abi3 list API with // three. - for value in values.try_iter()? { - let value = value?; + for (index, value) in values.try_iter()?.enumerate() { + let mut value = value?; if !(self.builder.append_exact(&value, self.spec.value_capacity) - || self.append_converted(py, &value)) + || self.append_converted(py, &mut value)) { - break; + let held = values.get_slice(index, values.len()); + // The capacity fallback can follow a successful conversion. Retain that + // result rather than invoking its converter a second time at finish. + held.set_item(0, value)?; + self.state = State::Holding { + listed: Some(held.unbind()), + pushed: Vec::new(), + }; + return Ok(()); } if let State::Failed(_) = self.state { return Ok(()); } - appended += 1; - } - if appended < values.len() { - self.state = State::Holding { - listed: Some(values.get_slice(appended, values.len()).unbind()), - pushed: Vec::new(), - }; } Ok(()) } /// Appends `value`, which the builder did not take, if it is a float, a bool or an int that - /// fits in 64 bits, which the `StringType` converter turns into a str it takes. Returns false - /// if `value` has to be held instead. A converter error fails the column. - fn append_converted<'py>(&mut self, py: Python<'py>, value: &Bound<'py, PyAny>) -> bool { + /// fits in 64 bits, which the `StringType` converter turns into a str it takes. Replaces + /// `value` with a successful conversion even when capacity requires holding it for `finish`. + /// Returns false if `value` has to be held instead. A converter error fails the column. + fn append_converted<'py>(&mut self, py: Python<'py>, value: &mut Bound<'py, PyAny>) -> bool { // A built-in cannot change after the UDF returns it, so converting it now matches PySpark, // which converts once the UDF has run on every row. For a string column the builder takes // only None and exact strs, which the StringType converter returns unchanged. @@ -674,9 +721,10 @@ impl<'s> ResultColumn<'s> { return false; } match self.spec.convert(py, value.clone()) { - Ok(converted) => self - .builder - .append_exact(&converted, self.spec.value_capacity), + Ok(converted) => { + *value = converted; + self.builder.append_exact(value, self.spec.value_capacity) + } Err(error) => { self.state = State::Failed(error); true @@ -856,6 +904,7 @@ fn appended_values<'py>(py: Python<'py>, appended: &Appended) -> PyResult PyResult<()> { + for extend in [false, true] { + let locals = PyDict::new(py); + py.run( + c"from types import SimpleNamespace +class StringType: pass +sql_types = SimpleNamespace(StringType=StringType) +spark_type = StringType() +events = [] +def converter(value): + events.append(('converter', value)) + raise ValueError('first conversion failed') +def array_spy(*args, **kwargs): + events.append(('array',)) + raise AssertionError('PyArrow must not run after a converter error') +", + Some(&locals), + Some(&locals), + )?; + let arrow_type = py.import("pyarrow")?.call_method0("string")?; + let mut spec = result_spec( + py, + &arrow_type, + &locals.get_item("spark_type")?.unwrap(), + &locals.get_item("converter")?.unwrap(), + true, + &locals.get_item("sql_types")?.unwrap(), + || Ok(None), + )? + .expect("nullable string policy is supported"); + // Replace only this spec's cached callback, not a process-wide PyArrow alias. + spec.pyarrow_array = locals.get_item("array_spy")?.unwrap().unbind(); + let mut column = ResultColumn::new(&spec, 4); + let values = PyList::new(py, [1, 2])?; + if extend { + column.extend(py, &values); + } else { + for value in values.iter() { + column.push(py, value); + } + } + column.extend(py, &values); + column.push(py, py.None().into_bound(py)); + let error = column.finish(py).err().expect("converter must fail"); + assert!(error.is_instance_of::(py)); + assert_eq!(error.to_string(), "ValueError: first conversion failed"); + assert!(py + .eval( + c"events == [('converter', 1)]", + Some(&locals), + Some(&locals) + )? + .extract::()?); + } + Ok(()) + }) + .unwrap(); + } } 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..2ae8a27d --- /dev/null +++ b/crates/spark-python-arrow/tests/bridge_regressions.rs @@ -0,0 +1,307 @@ +// 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, Int64Array}; +use arrow::datatypes::DataType; +use pyo3::exceptions::{PyTypeError, PyValueError}; +use pyo3::prelude::*; +use pyo3::types::PyDict; +use spark_python_arrow::bridge::{ + export_array, export_data_type, export_schema, import_array, import_batch, import_data_type, + import_schema, record_batches_to_pyarrow, +}; + +fn setup<'py>(py: Python<'py>, code: &CStr) -> PyResult> { + let locals = PyDict::new(py); + locals.set_item("pa", py.import("pyarrow")?)?; + py.run(code, Some(&locals), Some(&locals))?; + Ok(locals) +} + +fn value<'py>(locals: &Bound<'py, PyDict>, name: &str) -> PyResult> { + Ok(locals.get_item(name)?.expect("test variable must exist")) +} + +fn assert_equal(actual: &Bound<'_, PyAny>, expected: &Bound<'_, PyAny>) -> PyResult<()> { + assert!(actual.call_method1("equals", (expected,))?.is_truthy()?); + Ok(()) +} + +fn assert_schema_equal(actual: &Bound<'_, PyAny>, expected: &Bound<'_, PyAny>) -> PyResult<()> { + assert!(actual + .call_method1("equals", (expected, true))? + .is_truthy()?); + Ok(()) +} + +#[test] +fn sorted_maps_keep_type_flags_through_schema_array_and_batch_paths() { + Python::initialize(); + Python::attach(|py| -> PyResult<()> { + let locals = setup( + py, + 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), metadata={b'nested': b'yes'}) +]) +nested = pa.array([ + {'maps': [[('a', 1)]]}, None, {'maps': []} +], type=nested_type) +schema = pa.schema([ + pa.field('maps', map_type, metadata={b'field': b'yes'}), + pa.field('nested', nested_type) +], metadata={b'schema': b'yes'}) +batch = pa.RecordBatch.from_arrays([maps, nested], schema=schema) +expected = pa.Table.from_batches([batch])", + )?; + + let schema = value(&locals, "schema")?; + let imported_schema = import_schema(&schema)?; + let exported_schema = export_schema(py, &imported_schema)?; + assert_schema_equal(exported_schema.bind(py), &schema)?; + for name in ["map_type", "nested_type"] { + let data_type = value(&locals, name)?; + let exported = export_data_type(py, &import_data_type(&data_type)?)?; + assert_equal(exported.bind(py), &data_type)?; + } + for name in ["maps", "nested"] { + let array = value(&locals, name)?; + let exported = export_array(py, import_array(&array)?)?; + assert_equal(exported.bind(py), &array)?; + assert_equal(&exported.bind(py).getattr("type")?, &array.getattr("type")?)?; + } + let batch = value(&locals, "batch")?; + let imported = import_batch(&batch)?; + let exported = record_batches_to_pyarrow(py, vec![imported])?; + assert_equal(exported.bind(py), &value(&locals, "expected")?)?; + assert_schema_equal(&exported.bind(py).getattr("schema")?, &schema) + }) + .unwrap(); +} + +#[test] +fn ordered_dictionaries_keep_field_flags_in_nested_types_and_batches() { + Python::initialize(); + Python::attach(|py| -> PyResult<()> { + let locals = setup( + py, + 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) +schema = pa.schema([ + pa.field('dictionary', dictionary.type), + pa.field('nested', nested.type), + pa.field('lists', lists.type) +], metadata={b'schema': b'ordered'}) +batch = pa.RecordBatch.from_arrays([dictionary, nested, lists], schema=schema) +expected = pa.Table.from_batches([batch])", + )?; + let schema = value(&locals, "schema")?; + let imported_schema = import_schema(&schema)?; + // Field equality deliberately ignores dictionary-order flags in Arrow Rust. + // Check them explicitly, then also compare the exported Python schema. + assert_eq!(imported_schema.field(0).dict_is_ordered(), Some(true)); + let DataType::Struct(fields) = imported_schema.field(1).data_type() else { + panic!("expected a struct") + }; + assert_eq!(fields[0].dict_is_ordered(), Some(true)); + let DataType::List(field) = imported_schema.field(2).data_type() else { + panic!("expected a list") + }; + assert_eq!(field.dict_is_ordered(), Some(true)); + let exported = export_schema(py, &imported_schema)?; + assert_schema_equal(exported.bind(py), &schema)?; + + for name in ["nested", "lists"] { + let array = value(&locals, name)?; + let data_type = array.getattr("type")?; + let exported_type = export_data_type(py, &import_data_type(&data_type)?)?; + assert_equal(exported_type.bind(py), &data_type)?; + let exported_array = export_array(py, import_array(&array)?)?; + assert_equal(exported_array.bind(py), &array)?; + assert_equal(&exported_array.bind(py).getattr("type")?, &data_type)?; + } + + let imported = import_batch(&value(&locals, "batch")?)?; + assert_eq!(imported.schema().field(0).dict_is_ordered(), Some(true)); + let DataType::Struct(fields) = imported.column(1).data_type() else { + panic!("expected a struct column") + }; + assert_eq!(fields[0].dict_is_ordered(), Some(true)); + let exported = record_batches_to_pyarrow(py, vec![imported])?; + assert_equal(exported.bind(py), &value(&locals, "expected")?)?; + assert_schema_equal(&exported.bind(py).getattr("schema")?, &schema) + }) + .unwrap(); +} + +#[test] +fn bare_ordered_dictionaries_require_a_field_to_preserve_ordering() { + Python::initialize(); + Python::attach(|py| -> PyResult<()> { + let locals = setup( + py, + c"ordered = pa.DictionaryArray.from_arrays( + pa.array([0, 1], type=pa.int8()), pa.array(['a', 'b']), ordered=True +) +unordered = pa.DictionaryArray.from_arrays( + pa.array([0, 1], type=pa.int8()), pa.array(['a', 'b']) +)", + )?; + let ordered = value(&locals, "ordered")?; + for error in [ + import_array(&ordered).unwrap_err(), + import_data_type(&ordered.getattr("type")?).unwrap_err(), + ] { + assert!(error.is_instance_of::(py)); + assert!(error.to_string().contains("preserve ordering")); + } + let unordered = value(&locals, "unordered")?; + let exported = export_array(py, import_array(&unordered)?)?; + assert_equal(exported.bind(py), &unordered) + }) + .unwrap(); +} + +#[test] +fn batches_with_inconsistent_nested_dictionary_order_are_rejected() { + Python::initialize(); + Python::attach(|py| -> PyResult<()> { + let locals = setup( + py, + c"ordered = pa.DictionaryArray.from_arrays( + pa.array([1, None, 0], type=pa.int8()), pa.array(['a', 'b']), ordered=True +) +unordered = pa.DictionaryArray.from_arrays(ordered.indices, ordered.dictionary) +ordered_struct = pa.StructArray.from_arrays([ordered], names=['dictionary']) +unordered_struct = pa.StructArray.from_arrays([unordered], names=['dictionary']) +ordered_list = pa.ListArray.from_arrays(pa.array([0, 1, 2, 3], type=pa.int32()), ordered) +unordered_list = pa.ListArray.from_arrays(pa.array([0, 1, 2, 3], type=pa.int32()), unordered) +pairs = [ + (ordered, unordered), + (ordered_struct, unordered_struct), + (ordered_list, unordered_list) +] +batches = [ + (pa.record_batch([left], names=['x']), pa.record_batch([right], names=['x'])) + for left, right in pairs +]", + )?; + let batches = value(&locals, "batches")?; + for index in 0..batches.len()? { + let pair = batches.get_item(index)?; + let left = import_batch(&pair.get_item(0)?)?; + let right = import_batch(&pair.get_item(1)?)?; + let error = record_batches_to_pyarrow(py, vec![left, right]).unwrap_err(); + assert!(error.is_instance_of::(py)); + assert!(error.to_string().contains("inconsistent schemas")); + } + Ok(()) + }) + .unwrap(); +} + +#[test] +fn dictionary_value_extensions_are_rejected_instead_of_losing_identity() { + Python::initialize(); + Python::attach(|py| -> PyResult<()> { + let locals = setup( + py, + c"class TestExtension(pa.ExtensionType): + def __init__(self): + super().__init__(pa.int64(), 'spark_python_arrow.dictionary_extension') + + def __arrow_ext_serialize__(self): + return b'extension' + + @classmethod + def __arrow_ext_deserialize__(cls, storage_type, serialized): + return cls() + +extension = pa.ExtensionArray.from_storage(TestExtension(), pa.array([10, 20])) +dictionary = pa.DictionaryArray.from_arrays(pa.array([0, 1], type=pa.int8()), extension) +nested = pa.StructArray.from_arrays([dictionary], names=['dictionary']) +schema = pa.schema([pa.field('dictionary', dictionary.type)]) +batch = pa.RecordBatch.from_arrays([dictionary], schema=schema) +nested_schema = pa.schema([pa.field('nested', nested.type)]) +nested_batch = pa.RecordBatch.from_arrays([nested], schema=nested_schema)", + )?; + for name in ["dictionary", "nested"] { + let array = value(&locals, name)?; + for error in [ + import_array(&array).unwrap_err(), + import_data_type(&array.getattr("type")?).unwrap_err(), + ] { + assert!(error.is_instance_of::(py)); + assert!(error.to_string().contains("dictionary value extension")); + } + } + for name in ["schema", "nested_schema"] { + let error = import_schema(&value(&locals, name)?).unwrap_err(); + assert!(error.is_instance_of::(py)); + assert!(error.to_string().contains("dictionary value extension")); + } + for name in ["batch", "nested_batch"] { + let error = import_batch(&value(&locals, name)?).unwrap_err(); + assert!(error.is_instance_of::(py)); + assert!(error.to_string().contains("dictionary value extension")); + } + Ok(()) + }) + .unwrap(); +} + +#[test] +fn nullable_non_byte_aligned_slice_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 exported = export_array(py, Arc::new(slice))?; + drop(original); + py.import("gc")?.call_method0("collect")?; + let array = exported.bind(py); + assert_eq!( + array + .call_method0("buffers")? + .get_item(1)? + .getattr("address")? + .extract::()?, + address + ); + // Arrow may realign the validity bitmap to the value buffer's offset. + // Its identity is not part of the bridge's 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/udf_inputs.rs b/crates/spark-python-arrow/tests/udf_inputs.rs new file mode 100644 index 00000000..40a69cc8 --- /dev/null +++ b/crates/spark-python-arrow/tests/udf_inputs.rs @@ -0,0 +1,294 @@ +// 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. + +//! UDF input conversion and deferred Python errors, using valid Arrow buffer layouts. + +use arrow::datatypes::DataType; +use pyo3::exceptions::{PyRuntimeError, PyUnicodeDecodeError}; +use pyo3::prelude::*; +use pyo3::types::{PyDict, PyList}; +use spark_python_arrow::bridge::import_data_type; +use spark_python_arrow::udf::{input_spec, InputColumns}; + +fn context(py: Python<'_>) -> PyResult> { + let locals = PyDict::new(py); + locals.set_item("pa", py.import("pyarrow")?)?; + // This is the worker's input sequence for the primitive types tested here: + // materialize the whole column, then apply its optional caller converter. + py.run( + c"calls = [] +def convert_input(column, converter): + calls.append((column, converter)) + values = column.to_pylist() + return values if converter is None else [converter(value) for value in values] +", + Some(&locals), + Some(&locals), + )?; + Ok(locals) +} + +fn get<'py>(locals: &Bound<'py, PyDict>, name: &str) -> PyResult> { + Ok(locals.get_item(name)?.expect("fixture name is defined")) +} + +fn primitive_input<'py>( + py: Python<'py>, + array: &Bound<'py, PyAny>, + convert_input: &Bound<'py, PyAny>, +) -> PyResult { + let none = py.None().into_bound(py); + let data_type = import_data_type(&array.getattr("type")?)?; + let first = input_spec(&DataType::Int64, &none, &none, &none)? + .expect("the leading integer column is supported"); + let second = input_spec(&data_type, &none, &none, &none)? + .expect("the tested primitive column is supported"); + let pa = py.import("pyarrow")?; + let leading = pa.call_method1("array", (vec![1_i64; array.len()?],))?; + let batch = pa.call_method1( + "record_batch", + ( + PyList::new(py, [leading, array.clone()])?, + vec!["first", "second"], + ), + )?; + InputColumns::import(&batch, &[first, second], convert_input) +} + +fn assert_same_error(py: Python<'_>, actual: &PyErr, expected: &PyErr) -> PyResult<()> { + assert!(actual.get_type(py).is(expected.get_type(py))); + assert_eq!( + actual.value(py).str()?.extract::()?, + expected.value(py).str()?.extract::()? + ); + Ok(()) +} + +fn assert_one_callback(locals: &Bound<'_, PyDict>, array: &Bound<'_, PyAny>) -> PyResult<()> { + let calls = get(locals, "calls")?; + assert_eq!(calls.len()?, 1); + let call = calls.get_item(0)?; + assert!(call + .get_item(0)? + .call_method1("equals", (array,))? + .extract::()?); + assert!(call.get_item(1)?.is_none()); + Ok(()) +} + +#[test] +fn out_of_python_range_date_uses_the_original_column_callback() { + Python::initialize(); + Python::attach(|py| -> PyResult<()> { + let locals = context(py)?; + let convert_input = get(&locals, "convert_input")?; + let none = py.None().into_bound(py); + for days in [-719_163_i32, 2_932_897, i32::MIN, i32::MAX] { + locals.set_item("days", days)?; + // Every day value fits Arrow date32, even when Python cannot represent it. + let array = py.eval( + c"pa.array([0, days, None], type=pa.date32())", + Some(&locals), + None, + )?; + let kwargs = PyDict::new(py); + kwargs.set_item("full", true)?; + array.call_method("validate", (), Some(&kwargs))?; + let expected = convert_input.call1((&array, &none)).unwrap_err(); + get(&locals, "calls")?.call_method0("clear")?; + let input = primitive_input(py, &array, &convert_input)?; + assert_eq!(input.value(py, 0, 0)?.extract::()?, 1); + assert_eq!( + input.value(py, 1, 0)?.str()?.extract::()?, + "1970-01-01" + ); + assert!(input.value(py, 1, 2)?.is_none()); + assert_eq!(get(&locals, "calls")?.len()?, 0); + let actual = input.value(py, 1, 1).unwrap_err(); + assert_same_error(py, &actual, &expected)?; + assert_one_callback(&locals, &array)?; + } + Ok(()) + }) + .unwrap(); +} + +#[test] +fn invalid_utf8_payload_defers_to_the_python_column_error() { + Python::initialize(); + Python::attach(|py| -> PyResult<()> { + let locals = context(py)?; + locals.set_item("struct", py.import("struct")?)?; + let convert_input = get(&locals, "convert_input")?; + let none = py.None().into_bound(py); + for (type_name, offset_format) in [("string", "<4i"), ("large_string", "<4q")] { + locals.set_item( + "data_type", + py.import("pyarrow")?.getattr(type_name)?.call0()?, + )?; + locals.set_item("offset_format", offset_format)?; + // All offsets are monotonic and within this seven-byte payload. Only the text + // in the second cell is invalid UTF-8; no malformed offset or pointer is used. + let array = py.eval( + cr"pa.Array.from_buffers(data_type, 3, [ + None, + pa.py_buffer(struct.pack(offset_format, 0, 2, 3, 7)), + pa.py_buffer(b'ok\xfflast'), +])", + Some(&locals), + Some(&locals), + )?; + array.call_method0("validate")?; + let expected = convert_input.call1((&array, &none)).unwrap_err(); + assert!(expected.is_instance_of::(py)); + get(&locals, "calls")?.call_method0("clear")?; + let input = primitive_input(py, &array, &convert_input)?; + assert_eq!(input.value(py, 1, 0)?.extract::()?, "ok"); + assert_eq!(input.value(py, 1, 2)?.extract::()?, "last"); + assert_eq!(get(&locals, "calls")?.len()?, 0); + let actual = input.value(py, 1, 1).unwrap_err(); + assert_same_error(py, &actual, &expected)?; + assert_one_callback(&locals, &array)?; + } + Ok(()) + }) + .unwrap(); +} + +#[test] +fn callback_exception_is_propagated_without_replacement() { + Python::initialize(); + Python::attach(|py| -> PyResult<()> { + let locals = context(py)?; + py.run( + c"sentinel = LookupError('caller conversion failed') +def convert_input(column, converter): + calls.append((column, converter)) + raise sentinel +array = pa.array([2932897], type=pa.date32()) +", + Some(&locals), + Some(&locals), + )?; + let array = get(&locals, "array")?; + let input = primitive_input(py, &array, &get(&locals, "convert_input")?)?; + let error = input.value(py, 1, 0).unwrap_err(); + assert!(error.value(py).is(get(&locals, "sentinel")?)); + assert_one_callback(&locals, &array) + }) + .unwrap(); +} + +#[test] +fn unexpected_callback_success_is_not_silently_accepted() { + Python::initialize(); + Python::attach(|py| -> PyResult<()> { + let locals = context(py)?; + py.run( + c"def convert_input(column, converter): + calls.append((column, converter)) + return ['unexpected success'] +array = pa.array([2932897], type=pa.date32()) +", + Some(&locals), + Some(&locals), + )?; + let array = get(&locals, "array")?; + let input = primitive_input(py, &array, &get(&locals, "convert_input")?)?; + let error = input.value(py, 1, 0).unwrap_err(); + assert!(error.is_instance_of::(py)); + assert_eq!( + error.value(py).str()?.extract::()?, + "PySpark converted input 1, which the native converter cannot read" + ); + assert_one_callback(&locals, &array) + }) + .unwrap(); +} + +#[test] +fn variant_input_preserves_class_values_and_masked_null_parts() { + Python::initialize(); + Python::attach(|py| -> PyResult<()> { + let locals = context(py)?; + py.run( + c"class VariantType: pass +class VariantVal: + def __init__(self, value, metadata): + self.value = value + self.metadata = metadata + def __eq__(self, other): + return type(self) is type(other) and (self.value, self.metadata) == ( + other.value, other.metadata + ) +class SqlTypes: pass +SqlTypes.VariantType = VariantType +spark_type = VariantType() +converter = lambda value: None if value is None else VariantVal( + value['value'], value['metadata'] +) +", + Some(&locals), + Some(&locals), + )?; + for width in ["binary", "large_binary"] { + locals.set_item("part_type", py.import("pyarrow")?.getattr(width)?.call0()?)?; + py.run( + cr"fields = [ + pa.field('value', part_type, nullable=False), + pa.field('metadata', part_type, nullable=False, metadata={b'variant': b'true'}), +] +array = pa.StructArray.from_arrays([ + pa.array([b'a', None, b'\x00'], type=part_type), + pa.array([b'm', None, b'\xff'], type=part_type), +], fields=fields, mask=pa.array([False, True, False])) +array.validate(full=True) +expected = convert_input(array, converter) +calls.clear() +batch = pa.record_batch([array], names=['variant']) +", + Some(&locals), + Some(&locals), + )?; + let data_type = import_data_type(&get(&locals, "array")?.getattr("type")?)?; + let spec = input_spec( + &data_type, + &get(&locals, "spark_type")?, + &get(&locals, "converter")?, + &get(&locals, "SqlTypes")?, + )? + .expect("tagged Variant inputs are supported"); + let input = InputColumns::import( + &get(&locals, "batch")?, + &[spec], + &get(&locals, "convert_input")?, + )?; + for row in 0..input.num_rows() { + let actual = input.value(py, 0, row)?; + let expected = get(&locals, "expected")?.get_item(row)?; + assert!(actual.eq(&expected)?); + assert!(actual.get_type().is(expected.get_type())); + } + // The parent-null row is valid even though its non-nullable parts are null: + // those child nulls are masked, so no fallback or child access is required. + assert!(input.value(py, 0, 1)?.is_none()); + assert_eq!(get(&locals, "calls")?.len()?, 0); + } + Ok(()) + }) + .unwrap(); +} diff --git a/crates/spark-python-arrow/tests/udf_results.rs b/crates/spark-python-arrow/tests/udf_results.rs new file mode 100644 index 00000000..e1b2434c --- /dev/null +++ b/crates/spark-python-arrow/tests/udf_results.rs @@ -0,0 +1,707 @@ +// 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. + +//! UDF result-policy checks, without a Spark server. The fixtures model PySpark's nullable +//! top-level String/Binary/Variant converters. Logging and injected failures make scheduling +//! observable; they are test instrumentation, not support for arbitrary stateful converters. +//! Set SPARK_PYTHON_ARROW_TEST_PYSPARK=1 to additionally compare against installed PySpark. + +use std::ffi::CString; + +use pyo3::exceptions::{PyRuntimeError, PyValueError}; +use pyo3::prelude::*; +use pyo3::types::{PyDict, PyList, PyType}; +use spark_python_arrow::udf::{result_spec, Finished, ResultColumn, ResultSpec}; + +#[derive(Clone, Copy, Debug)] +enum Mode { + Push, + Extend, +} + +const MODES: [Mode; 2] = [Mode::Push, Mode::Extend]; + +fn context(py: Python<'_>) -> PyResult> { + let locals = PyDict::new(py); + py.run( + c"import pyarrow as pa +from types import SimpleNamespace + +class StringType: pass +class BinaryType: pass +class VariantType: pass +sql_types = SimpleNamespace( + StringType=StringType, BinaryType=BinaryType, VariantType=VariantType +) + +class VariantVal: + def __init__(self, value, metadata): + self.value = value + self.metadata = metadata + +class VariantSubclass(VariantVal): pass +class BytesSubclass(bytes): pass + +class Text: + def __init__(self, text): + self.text = text + def __str__(self): + return self.text + +class BadText: + def __str__(self): + raise ValueError('text conversion failed') + +variant_type = pa.struct([ + pa.field('value', pa.binary(), nullable=False), + pa.field('metadata', pa.binary(), nullable=False, metadata={b'variant': b'true'}), +]) + +def convert_string(value): + if value is None: + return None + return str(value).lower() if isinstance(value, bool) else str(value) + +def convert_binary(value): + if value is None: + return None + assert isinstance(value, (bytes, bytearray)) + return bytes(value) + +def convert_variant(value): + if value is None: + return None + if isinstance(value, VariantVal): + return {'value': value.value, 'metadata': value.metadata} + raise ValueError('malformed variant') + +events = [] + +def tracked_string(value): + events.append(value) + return convert_string(value) + +def tracked_binary(value): + events.append(value) + return convert_binary(value) + +def tracked_variant(value): + events.append(value) + return convert_variant(value) + +def faulty_string(value): + events.append(value) + if value == 7: + raise ValueError('injected converter failure') + return convert_string(value) + +def later_udf_error(value): + if value == 2: + raise RuntimeError('later UDF failure') + return 7 + +def reference_results(values, converter, arrow_type): + converted = [converter(value) for value in values] + try: + return pa.array(converted, type=arrow_type) + except pa.ArrowInvalid: + return pa.array(converted).cast(arrow_type, safe=True) +", + Some(&locals), + Some(&locals), + )?; + Ok(locals) +} + +fn evaluate<'py>( + py: Python<'py>, + locals: &Bound<'py, PyDict>, + expression: &str, +) -> PyResult> { + let code = CString::new(expression).expect("test expression contains no NUL"); + py.eval(&code, Some(locals), Some(locals)) +} + +fn policy<'py>( + py: Python<'py>, + locals: &Bound<'py, PyDict>, + arrow_type: &Bound<'py, PyAny>, + spark_type: &Bound<'py, PyAny>, + converter: &Bound<'py, PyAny>, +) -> PyResult { + let sql_types = evaluate(py, locals, "sql_types")?; + let spec = result_spec( + py, + arrow_type, + spark_type, + converter, + true, + &sql_types, + || { + Ok(Some( + evaluate(py, locals, "VariantVal")?.cast_into::()?, + )) + }, + )?; + Ok(spec.expect("fixture type is supported")) +} + +fn fixture_policy<'py>( + py: Python<'py>, + locals: &Bound<'py, PyDict>, + arrow_type: &str, + spark_type: &str, + converter: &str, +) -> PyResult { + policy( + py, + locals, + &evaluate(py, locals, arrow_type)?, + &evaluate(py, locals, spark_type)?, + &evaluate(py, locals, converter)?, + ) +} + +fn add_values<'py>( + py: Python<'py>, + column: &mut ResultColumn<'_>, + values: &Bound<'py, PyList>, + mode: Mode, +) { + match mode { + Mode::Push => { + for value in values.iter() { + column.push(py, value); + } + } + Mode::Extend => column.extend(py, values), + } +} + +fn reference<'py>( + py: Python<'py>, + locals: &Bound<'py, PyDict>, + values: &Bound<'py, PyList>, + converter: &Bound<'py, PyAny>, + arrow_type: &Bound<'py, PyAny>, +) -> PyResult> { + evaluate(py, locals, "reference_results")?.call1((values, converter, arrow_type)) +} + +fn assert_same_result( + py: Python<'_>, + actual: PyResult>, + expected: PyResult>, +) -> PyResult<()> { + match (actual, expected) { + (Ok(actual), Ok(expected)) => assert!( + actual + .array + .call_method1("equals", (&expected,))? + .extract::()?, + "actual {}; expected {expected}", + actual.array + ), + (Err(actual), Err(expected)) => { + assert!(actual.get_type(py).is(expected.get_type(py))); + assert_eq!(actual.to_string(), expected.to_string()); + } + (actual, expected) => panic!( + "actual success {}, expected success {}", + actual.is_ok(), + expected.is_ok() + ), + } + Ok(()) +} + +#[test] +fn string_conversion_matches_nullable_pyspark_policy_in_both_ingestion_modes() { + Python::initialize(); + Python::attach(|py| -> PyResult<()> { + for mode in MODES { + let locals = context(py)?; + let values = evaluate(py, &locals, "[None, 'plain', True, False, 1, -7, 1.25]")? + .cast_into::()?; + for arrow_expression in ["pa.string()", "pa.large_string()"] { + evaluate(py, &locals, "events.clear()")?; + let spec = fixture_policy( + py, + &locals, + arrow_expression, + "StringType()", + "tracked_string", + )?; + assert!(evaluate(py, &locals, "not events")?.extract::()?); + let mut column = ResultColumn::new(&spec, values.len()); + add_values(py, &mut column, &values, mode); + let actual = column.finish(py)?; + assert_eq!(actual.pyarrow_values, 0); + assert!( + evaluate(py, &locals, "events == [True, False, 1, -7, 1.25]")? + .extract::()? + ); + let expected = reference( + py, + &locals, + &values, + &evaluate(py, &locals, "convert_string")?, + &evaluate(py, &locals, arrow_expression)?, + ); + assert_same_result(py, Ok(actual), expected)?; + } + } + Ok(()) + }) + .unwrap(); +} + +#[test] +fn capacity_fallback_retains_converted_value_and_calls_converter_once() { + Python::initialize(); + Python::attach(|py| -> PyResult<()> { + for mode in MODES { + for capacity in [0, 1] { + let locals = context(py)?; + let values = evaluate(py, &locals, "[1, 2, 3]")?.cast_into::()?; + let spec = + fixture_policy(py, &locals, "pa.string()", "StringType()", "tracked_string")? + .with_value_capacity(capacity); + let mut column = ResultColumn::new(&spec, values.len()); + add_values(py, &mut column, &values, mode); + let actual = column.finish(py)?; + assert_eq!(actual.pyarrow_values, 3); + assert!(evaluate(py, &locals, "events == [1, 2, 3]")?.extract::()?); + let expected = reference( + py, + &locals, + &values, + &evaluate(py, &locals, "convert_string")?, + &evaluate(py, &locals, "pa.string()")?, + ); + assert_same_result(py, Ok(actual), expected)?; + } + } + Ok(()) + }) + .unwrap(); +} + +#[test] +fn held_string_and_binary_callbacks_observe_final_mutable_values() { + Python::initialize(); + Python::attach(|py| -> PyResult<()> { + for mode in MODES { + for (arrow_type, spark_type, converter, values, mutation) in [ + ( + "pa.string()", + "StringType()", + "tracked_string", + "['prefix', Text('before'), None, True]", + "values[1].text = 'after'", + ), + ( + "pa.binary()", + "BinaryType()", + "tracked_binary", + "[b'prefix', bytearray(b'before'), None, BytesSubclass(b'tail')]", + "values[1][:] = b'after'", + ), + ] { + let locals = context(py)?; + let values = evaluate(py, &locals, values)?.cast_into::()?; + locals.set_item("values", &values)?; + let spec = fixture_policy(py, &locals, arrow_type, spark_type, converter)?; + let mut column = ResultColumn::new(&spec, values.len()); + add_values(py, &mut column, &values, mode); + assert!(evaluate(py, &locals, "not events")?.extract::()?); + py.run( + &CString::new(mutation).unwrap(), + Some(&locals), + Some(&locals), + )?; + let actual = column.finish(py)?; + assert!(actual.pyarrow_values > 0); + assert!( + evaluate(py, &locals, "len(events) == 2 and events[0] is values[1]")? + .extract::()? + ); + let pure_converter = converter.replacen("tracked_", "convert_", 1); + let expected = reference( + py, + &locals, + &values, + &evaluate(py, &locals, &pure_converter)?, + &evaluate(py, &locals, arrow_type)?, + ); + assert_same_result(py, Ok(actual), expected)?; + } + } + Ok(()) + }) + .unwrap(); +} + +#[test] +fn converter_failures_are_delayed_and_match_reference_error_order() { + Python::initialize(); + Python::attach(|py| -> PyResult<()> { + for mode in MODES { + for (arrow_type, spark_type, converter, values) in [ + ( + "pa.string()", + "StringType()", + "faulty_string", + "[1, 7, BadText()]", + ), + ( + "pa.string()", + "StringType()", + "convert_string", + "['prefix', BadText(), 7]", + ), + ( + "pa.binary()", + "BinaryType()", + "convert_binary", + "[b'prefix', object(), None]", + ), + ] { + let locals = context(py)?; + let values = evaluate(py, &locals, values)?.cast_into::()?; + let spec = fixture_policy(py, &locals, arrow_type, spark_type, converter)?; + let mut column = ResultColumn::new(&spec, values.len()); + add_values(py, &mut column, &values, mode); + // Neither ingestion API raises. Failed-state ingestion also preserves the + // first error rather than running a later callback. + column.push(py, py.None().into_bound(py)); + column.extend(py, &PyList::empty(py)); + let actual = column.finish(py); + assert!(actual.is_err()); + let expected = reference( + py, + &locals, + &values, + &evaluate(py, &locals, converter)?, + &evaluate(py, &locals, arrow_type)?, + ); + assert_same_result(py, actual, expected)?; + } + } + Ok(()) + }) + .unwrap(); +} + +#[test] +fn later_udf_error_wins_when_caller_finishes_only_after_all_udf_calls() { + Python::initialize(); + Python::attach(|py| -> PyResult<()> { + for mode in MODES { + let locals = context(py)?; + let spec = fixture_policy(py, &locals, "pa.string()", "StringType()", "faulty_string")?; + let udf = evaluate(py, &locals, "later_udf_error")?; + let outcome = (|| -> PyResult> { + let mut column = ResultColumn::new(&spec, 3); + let collected = PyList::empty(py); + for row in 0..3 { + let value = udf.call1((row,))?; + match mode { + Mode::Push => column.push(py, value), + Mode::Extend => collected.append(value)?, + } + } + if matches!(mode, Mode::Extend) { + column.extend(py, &collected); + } + column.finish(py) + })(); + let error = outcome.err().expect("the later UDF call must fail"); + assert!(error.is_instance_of::(py)); + assert!(error.to_string().contains("later UDF failure")); + } + Ok(()) + }) + .unwrap(); +} + +#[test] +fn large_integer_string_conversion_is_deferred_until_finish() { + Python::initialize(); + Python::attach(|py| -> PyResult<()> { + for mode in MODES { + for expression in ["[2**80]", "[2**16000]"] { + let locals = context(py)?; + let values = evaluate(py, &locals, expression)?.cast_into::()?; + let spec = + fixture_policy(py, &locals, "pa.string()", "StringType()", "tracked_string")?; + let mut column = ResultColumn::new(&spec, values.len()); + add_values(py, &mut column, &values, mode); + assert!(evaluate(py, &locals, "not events")?.extract::()?); + let actual = column.finish(py); + assert!(evaluate(py, &locals, "len(events) == 1")?.extract::()?); + // Some Python versions reject the larger value under their digit limit. + // Either outcome must match the same nullable converter at finish time. + let expected = reference( + py, + &locals, + &values, + &evaluate(py, &locals, "convert_string")?, + &evaluate(py, &locals, "pa.string()")?, + ); + assert_same_result(py, actual, expected)?; + } + } + Ok(()) + }) + .unwrap(); +} + +#[test] +fn native_variant_results_include_nulls_and_observe_mutation_before_finish() { + Python::initialize(); + Python::attach(|py| -> PyResult<()> { + for mode in MODES { + let locals = context(py)?; + let value = evaluate(py, &locals, "VariantVal(b'old', b'metadata')")?; + let values = PyList::new(py, [value.clone(), py.None().into_bound(py), value.clone()])?; + let spec = fixture_policy( + py, + &locals, + "variant_type", + "VariantType()", + "tracked_variant", + )?; + let mut column = ResultColumn::new(&spec, values.len()); + add_values(py, &mut column, &values, mode); + value.setattr("value", evaluate(py, &locals, "b'new'")?)?; + let actual = column.finish(py)?; + assert_eq!(actual.pyarrow_values, 0); + assert!(!actual.retried); + assert!(evaluate(py, &locals, "not events")?.extract::()?); + let expected = reference( + py, + &locals, + &values, + &evaluate(py, &locals, "convert_variant")?, + &evaluate(py, &locals, "variant_type")?, + ); + assert_same_result(py, Ok(actual), expected)?; + } + Ok(()) + }) + .unwrap(); +} + +#[test] +fn variant_subclass_non_bytes_and_capacity_use_fallback() { + Python::initialize(); + Python::attach(|py| -> PyResult<()> { + for mode in MODES { + for (expression, capacity) in [ + ("[VariantSubclass(b'v', b'm'), None]", usize::MAX), + ( + "[VariantVal(bytearray(b'v'), bytearray(b'm')), None]", + usize::MAX, + ), + ("[VariantVal(b'vv', b'mm'), None]", 1), + ] { + let locals = context(py)?; + let values = evaluate(py, &locals, expression)?.cast_into::()?; + let spec = fixture_policy( + py, + &locals, + "variant_type", + "VariantType()", + "tracked_variant", + )? + .with_value_capacity(capacity); + let mut column = ResultColumn::new(&spec, values.len()); + add_values(py, &mut column, &values, mode); + let actual = column.finish(py)?; + assert_eq!(actual.pyarrow_values, values.len()); + assert!(evaluate(py, &locals, "len(events) == 1")?.extract::()?); + let expected = reference( + py, + &locals, + &values, + &evaluate(py, &locals, "convert_variant")?, + &evaluate(py, &locals, "variant_type")?, + ); + assert_same_result(py, Ok(actual), expected)?; + } + } + Ok(()) + }) + .unwrap(); +} + +#[test] +fn variant_converter_failure_precedes_arrow_rejection() { + Python::initialize(); + Python::attach(|py| -> PyResult<()> { + for mode in MODES { + let locals = context(py)?; + // The first result reaches PyArrow only if every converter succeeds. Its invalid + // binary part must not mask the second result's conversion failure. + let values = evaluate(py, &locals, "[VariantVal(object(), b'm'), object()]")? + .cast_into::()?; + let spec = fixture_policy( + py, + &locals, + "variant_type", + "VariantType()", + "convert_variant", + )?; + let mut column = ResultColumn::new(&spec, values.len()); + add_values(py, &mut column, &values, mode); + let actual = column.finish(py); + assert!(matches!(&actual, Err(error) if error.is_instance_of::(py))); + let expected = reference( + py, + &locals, + &values, + &evaluate(py, &locals, "convert_variant")?, + &evaluate(py, &locals, "variant_type")?, + ); + assert_same_result(py, actual, expected)?; + } + Ok(()) + }) + .unwrap(); +} + +#[test] +fn dictionary_results_select_python_fallback_including_ordered_dictionaries() { + Python::initialize(); + Python::attach(|py| -> PyResult<()> { + let locals = context(py)?; + let none = py.None().into_bound(py); + for ordered in ["False", "True"] { + let arrow_type = evaluate( + py, + &locals, + &format!("pa.dictionary(pa.int8(), pa.string(), ordered={ordered})"), + )?; + let selected = result_spec(py, &arrow_type, &none, &none, true, &none, || { + panic!("dictionary fallback must not resolve VariantVal") + })?; + assert!(matches!(selected, Err(reason) if reason.contains("dictionary"))); + } + Ok(()) + }) + .unwrap(); +} + +#[test] +fn nested_dictionary_extension_results_select_python_fallback_before_import() { + Python::initialize(); + Python::attach(|py| -> PyResult<()> { + let locals = context(py)?; + py.run( + c"class BinaryExtension(pa.ExtensionType): + def __init__(self): + super().__init__(pa.binary(), 'spark.test.udf_result_binary') + def __arrow_ext_serialize__(self): + return b'' + @classmethod + def __arrow_ext_deserialize__(cls, storage_type, serialized): + return cls() + +dictionary_extension = pa.dictionary(pa.int8(), BinaryExtension()) +types = [ + pa.list_(dictionary_extension), + pa.struct([pa.field('nested', dictionary_extension)]), + pa.map_(pa.string(), dictionary_extension), + pa.struct([ + pa.field('value', dictionary_extension, nullable=False), + pa.field('metadata', pa.binary(), nullable=False, metadata={b'variant': b'true'}), + ]), +] +", + Some(&locals), + Some(&locals), + )?; + let none = py.None().into_bound(py); + for arrow_type in evaluate(py, &locals, "types")?.try_iter()? { + let arrow_type = arrow_type?; + let selected = result_spec(py, &arrow_type, &none, &none, true, &none, || { + panic!("unsupported nested types must not resolve VariantVal") + })?; + assert!(matches!(selected, Err(reason) if reason.contains("not supported"))); + } + Ok(()) + }) + .unwrap(); +} + +#[test] +fn official_pyspark_result_converters_match_native_and_fallback_paths() { + if std::env::var("SPARK_PYTHON_ARROW_TEST_PYSPARK").as_deref() != Ok("1") { + eprintln!("set SPARK_PYTHON_ARROW_TEST_PYSPARK=1 with official pyspark installed"); + return; + } + Python::initialize(); + Python::attach(|py| -> PyResult<()> { + let locals = context(py)?; + py.run( + c"from pyspark.sql import types as sql_types +from pyspark.sql.conversion import LocalDataToArrowConversion +from pyspark.sql.pandas.types import to_arrow_type +VariantVal = sql_types.VariantVal +class OfficialVariantSubclass(VariantVal): pass +cases = [ + (sql_types.StringType(), [None, 'plain', True, False, 1, -7, 1.25, Text('text')]), + (sql_types.StringType(), ['prefix', BadText()]), + (sql_types.StringType(), [2**80]), + (sql_types.StringType(), [2**16000]), + (sql_types.BinaryType(), [b'prefix', bytearray(b'value'), None, BytesSubclass(b'tail')]), + (sql_types.BinaryType(), [b'prefix', object()]), + (sql_types.VariantType(), [VariantVal(b'v', b'm'), None]), + (sql_types.VariantType(), [OfficialVariantSubclass(b'v', b'm'), None]), + (sql_types.VariantType(), [VariantVal(bytearray(b'v'), bytearray(b'm')), None]), + (sql_types.VariantType(), [VariantVal(object(), b'm'), object()]), +] +", + Some(&locals), + Some(&locals), + )?; + for case in evaluate(py, &locals, "cases")?.try_iter()? { + let case = case?; + let spark_type = case.get_item(0)?; + let values = case.get_item(1)?.cast_into::()?; + let converter = evaluate(py, &locals, "LocalDataToArrowConversion._create_converter")? + .call1((&spark_type,))?; + let arrow_type = evaluate(py, &locals, "to_arrow_type")?.call1((&spark_type,))?; + for mode in MODES { + for capacity in [1, usize::MAX] { + let spec = policy(py, &locals, &arrow_type, &spark_type, &converter)? + .with_value_capacity(capacity); + let mut column = ResultColumn::new(&spec, values.len()); + add_values(py, &mut column, &values, mode); + assert_same_result( + py, + column.finish(py), + reference(py, &locals, &values, &converter, &arrow_type), + )?; + } + } + } + Ok(()) + }) + .unwrap(); +} diff --git a/python/tests/test_arrow_handoff.py b/python/tests/test_arrow_handoff.py index dac57b7e..05ff69f7 100644 --- a/python/tests/test_arrow_handoff.py +++ b/python/tests/test_arrow_handoff.py @@ -269,7 +269,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())), + pa.field("labels", pa.map_(pa.string(), pa.int64(), keys_sorted=True)), ] ) arrays = [ @@ -299,7 +299,9 @@ def _mixed_batch(): type=nested_type, ), pa.DictionaryArray.from_arrays( - pa.array([0, None, 1, 0], type=pa.int8()), pa.array(["red", "blue"]) + pa.array([0, None, 1, 0], type=pa.int8()), + pa.array(["red", "blue"]), + ordered=True, ), pa.nulls(4), ] @@ -498,6 +500,19 @@ def test_mismatched_batch_schemas_are_rejected_before_import(monkeypatch): 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, _): From 7bcf8e1b5def3cc7b49dc324d32901774b6d6df5 Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Tue, 27 Oct 2026 12:58:55 -0700 Subject: [PATCH 3/8] Upgrade PyO3 to 0.30.1 Build the OSS Python extension with PyO3 0.30.1 and declare its Rust 1.85 requirement. Retain the shared converter's PyO3 0.28 and 0.29 compatibility and add 0.30.1 to its isolated CI matrix. Co-authored-by: Isaac --- .github/workflows/python-arrow.yml | 4 +- Cargo.lock | 98 +++++++++++++++------------- Cargo.toml | 2 +- crates/pyspark-rs/Cargo.toml | 1 + crates/pyspark-rs/src/row.rs | 2 +- crates/spark-python-arrow/Cargo.toml | 2 +- crates/spark-python-arrow/README.md | 9 +-- 7 files changed, 65 insertions(+), 53 deletions(-) diff --git a/.github/workflows/python-arrow.yml b/.github/workflows/python-arrow.yml index 73c87a7e..b80020bd 100644 --- a/.github/workflows/python-arrow.yml +++ b/.github/workflows/python-arrow.yml @@ -40,7 +40,7 @@ jobs: strategy: fail-fast: false matrix: - pyo3: ["0.28.3", "0.29.0"] + pyo3: ["0.28.3", "0.29.0", "0.30.1"] env: SPARK_PYTHON_ARROW_TEST_PYSPARK: "1" steps: @@ -55,7 +55,7 @@ jobs: - name: Install reference Python converters run: python -m pip install "pyspark-client==4.2.0" "pyarrow==26.0.0" - name: Isolate the standalone crate and select PyO3 - # The workspace lockfile pins 0.28.3. Build outside it so each matrix leg + # The workspace lockfile pins 0.30.1. Build outside it so each matrix leg # actually compiles the converter against the requested minor version. run: | cp -R crates/spark-python-arrow "$RUNNER_TEMP/spark-python-arrow" diff --git a/Cargo.lock b/Cargo.lock index 92e73626..8f91c977 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -102,7 +102,7 @@ version = "4.3.0" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -408,7 +408,7 @@ checksum = "c7c24de15d275a1ecfd47a380fb4d5ec9bfe0933f309ed5e705b775596a3574d" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -419,7 +419,7 @@ checksum = "9035ad2d096bed7955a320ee7e2230574d28fd3c3a0f186cbea1ff3c7eed5dbb" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -588,7 +588,7 @@ checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -1268,7 +1268,7 @@ checksum = "2e367e6a71051d0ebdd29b2f85d12059b38b1d1f172c6906e80016da662226bd" dependencies = [ "datafusion-doc", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -1450,7 +1450,7 @@ checksum = "97369cbbc041bc366949bc74d34658d6cda5621039731c6310521892a3a20ae0" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -1659,7 +1659,7 @@ checksum = "e835b70203e41293343137df5c0664546da5745f82ec9b84d40be8336958447b" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -2386,7 +2386,7 @@ checksum = "ed3955f1a9c7c0c15e092f9c887db08b1fc683305fdf6eb6684f22555355e202" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -2573,7 +2573,7 @@ checksum = "c96395f0a926bc13b1c17622aaddda1ecb55d49c8f1bf9777e4d877800a43f8b" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -3232,7 +3232,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b" dependencies = [ "proc-macro2", - "syn", + "syn 2.0.117", ] [[package]] @@ -3271,7 +3271,7 @@ dependencies = [ "pulldown-cmark", "pulldown-cmark-to-cmark", "regex", - "syn", + "syn 2.0.117", "tempfile", ] @@ -3285,7 +3285,7 @@ dependencies = [ "itertools", "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -3329,9 +3329,9 @@ dependencies = [ [[package]] name = "pyo3" -version = "0.28.3" +version = "0.30.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "91fd8e38a3b50ed1167fb981cd6fd60147e091784c427b8f7183a7ee32c31c12" +checksum = "229ca13b8a15bb0d1487b81ab975bfcd97513416499a8312011b6036e42e0de4" dependencies = [ "libc", "once_cell", @@ -3343,18 +3343,18 @@ dependencies = [ [[package]] name = "pyo3-build-config" -version = "0.28.3" +version = "0.30.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e368e7ddfdeb98c9bca7f8383be1648fd84ab466bf2bc015e94008db6d35611e" +checksum = "70130e2b6dd1ee589a8788026a429561ca5c40fa7a2fa7914c2282eddae24700" dependencies = [ "target-lexicon", ] [[package]] name = "pyo3-ffi" -version = "0.28.3" +version = "0.30.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7f29e10af80b1f7ccaf7f69eace800a03ecd13e883acfacc1e5d0988605f651e" +checksum = "5e7ff73c555c286d212d4ce7e5fbda845a973c4f61bd5f160f5f2c3c2a04d1c4" dependencies = [ "libc", "pyo3-build-config", @@ -3362,27 +3362,26 @@ dependencies = [ [[package]] name = "pyo3-macros" -version = "0.28.3" +version = "0.30.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "df6e520eff47c45997d2fc7dd8214b25dd1310918bbb2642156ef66a67f29813" +checksum = "c46d2e5862c7bb8207e3b639f713dda3c499933444784154c11d013c5a691b60" dependencies = [ "proc-macro2", "pyo3-macros-backend", "quote", - "syn", + "syn 3.0.9", ] [[package]] name = "pyo3-macros-backend" -version = "0.28.3" +version = "0.30.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c4cdc218d835738f81c2338f822078af45b4afdf8b2e33cbb5916f108b813acb" +checksum = "38b89f066f2ae72a2aef52d1537ca2fcd104440ebfadee06e97aa90738a25b9b" dependencies = [ "heck", "proc-macro2", - "pyo3-build-config", "quote", - "syn", + "syn 3.0.9", ] [[package]] @@ -3589,7 +3588,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "76009fbe0614077fc1a2ce255e3a1881a2e3a3527097d5dc6d8212c585e7e38b" dependencies = [ "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -3873,7 +3872,7 @@ checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -4025,7 +4024,7 @@ checksum = "028e551d5e270b31b9f3ea271778d9d827148d4287a5d96167b6bb9787f5cc38" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -4083,7 +4082,7 @@ dependencies = [ "heck", "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -4103,6 +4102,17 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "syn" +version = "3.0.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d78c8dee4c7bf0e14673097256fed6142ce9d3b85a408189d07482442145823b" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + [[package]] name = "sync_wrapper" version = "1.0.2" @@ -4120,7 +4130,7 @@ checksum = "728a70f3dbaf5bab7f0c4b1ac8d7ae5ea60a4b5549c8a5914361c99147a709d2" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -4173,7 +4183,7 @@ checksum = "ebc4ee7f67670e9b64d05fa4253e753e016c6c95ff35b89b7941d6b856dec1d5" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -4242,7 +4252,7 @@ checksum = "385a6cb71ab9ab790c5fe8d67f1645e6c450a7ce006a33de03daa956cf70a496" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -4320,7 +4330,7 @@ dependencies = [ "prettyplease", "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -4345,7 +4355,7 @@ dependencies = [ "prost-build", "prost-types", "quote", - "syn", + "syn 2.0.117", "tempfile", "tonic-build", ] @@ -4418,7 +4428,7 @@ checksum = "7490cfa5ec963746568740651ac6781f701c9c5ea257c58e057f3ba8cf69e8da" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -4635,7 +4645,7 @@ dependencies = [ "bumpalo", "proc-macro2", "quote", - "syn", + "syn 2.0.117", "wasm-bindgen-shared", ] @@ -4822,7 +4832,7 @@ checksum = "053e2e040ab57b9dc951b72c264860db7eb3b0200ba345b4e4c3b14f67855ddf" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -4833,7 +4843,7 @@ checksum = "3f316c4a2570ba26bbec722032c4099d8c8bc095efccdc15688708623367e358" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -5030,7 +5040,7 @@ dependencies = [ "heck", "indexmap", "prettyplease", - "syn", + "syn 2.0.117", "wasm-metadata", "wit-bindgen-core", "wit-component", @@ -5046,7 +5056,7 @@ dependencies = [ "prettyplease", "proc-macro2", "quote", - "syn", + "syn 2.0.117", "wit-bindgen-core", "wit-bindgen-rust", ] @@ -5119,7 +5129,7 @@ checksum = "de844c262c8848816172cef550288e7dc6c7b7814b4ee56b3e1553f275f1858e" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", "synstructure", ] @@ -5140,7 +5150,7 @@ checksum = "1ae7f38b72ec2a254e2b87ef277cf2cd4fb97cbebf944faa6f33354da0867930" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -5160,7 +5170,7 @@ checksum = "11532158c46691caf0f2593ea8358fed6bbf68a0315e80aae9bd41fbade684a1" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", "synstructure", ] @@ -5200,7 +5210,7 @@ checksum = "625dc425cab0dca6dc3c3319506e6593dcb08a9f387ea3b284dbd52a92c40555" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 31fe8090..9e90594d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -87,7 +87,7 @@ datafusion = { version = "53", default-features = false } polars = { version = "0.54", default-features = false, features = ["ipc"] } # PyO3 for the pyspark-rs CPython extension. -pyo3 = { version = "0.28", features = ["extension-module", "abi3-py39"] } +pyo3 = { version = "0.30.1", features = ["extension-module", "abi3-py39"] } # misc bytes = "1" diff --git a/crates/pyspark-rs/Cargo.toml b/crates/pyspark-rs/Cargo.toml index 135d2b1a..c9307c57 100644 --- a/crates/pyspark-rs/Cargo.toml +++ b/crates/pyspark-rs/Cargo.toml @@ -2,6 +2,7 @@ name = "pyspark-rs" version.workspace = true edition.workspace = true +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. diff --git a/crates/pyspark-rs/src/row.rs b/crates/pyspark-rs/src/row.rs index 28332cd2..406216af 100644 --- a/crates/pyspark-rs/src/row.rs +++ b/crates/pyspark-rs/src/row.rs @@ -18,7 +18,7 @@ impl PyRow { } } -/// Convert a Spark `Value` into a Python object (PyO3 0.28 `IntoPyObject`). +/// Convert a Spark `Value` into a Python object using PyO3's `IntoPyObject`. pub fn value_to_py<'py>(py: Python<'py>, v: &Value) -> PyResult> { // Propagate conversion failures as Python exceptions with `?` rather than // panicking across the FFI boundary (the conversions are near-infallible, diff --git a/crates/spark-python-arrow/Cargo.toml b/crates/spark-python-arrow/Cargo.toml index 9ce9d119..9fc769ab 100644 --- a/crates/spark-python-arrow/Cargo.toml +++ b/crates/spark-python-arrow/Cargo.toml @@ -26,7 +26,7 @@ repository = "https://github.com/apache/spark-connect-rust" arrow = { version = "58", default-features = false, features = ["ffi", "ipc"] } # The consumer chooses the Python ABI floor and embedding/extension-module linking. # Compatibility CI resolves each supported minor independently of the workspace lockfile. -pyo3 = { version = ">=0.28, <0.30" } +pyo3 = { version = ">=0.28, <0.31" } [lints.rust] unsafe_op_in_unsafe_fn = "deny" diff --git a/crates/spark-python-arrow/README.md b/crates/spark-python-arrow/README.md index 0e70460c..d1dc9444 100644 --- a/crates/spark-python-arrow/README.md +++ b/crates/spark-python-arrow/README.md @@ -112,10 +112,11 @@ enable PyO3's `extension-module` feature or select a minimum Python ABI: a Pytho extension chooses those features in the host package, while an embedded worker links its chosen Python runtime. PyO3 must resolve to one compatible version in each consuming dependency graph. -The manifest permits PyO3 0.28 and 0.29 so consumers on either line can use the -same source. The workspace lockfile selects 0.28.3 and does not validate 0.29. -The isolated converter compatibility CI job tests both versions independently; -0.29 support must not be considered verified until that job passes. +The manifest permits PyO3 0.28, 0.29, and 0.30 so consumers on these lines can use +the same source. The OSS Python extension selects 0.30.1, which requires Rust +1.85 or newer. The isolated converter compatibility CI job tests each minor +version independently of the workspace lockfile. A version is not considered +verified until its corresponding compatibility job passes. Bazel consumers should pin the OSS revision or release using their existing third-party Rust dependency mechanism and depend on this library target. There From 7c44352643e03a7393ee97b10fbe473694b4c7ea Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Tue, 27 Oct 2026 13:14:15 -0700 Subject: [PATCH 4/8] Report optional PySpark parity tests as ignored Replace the environment-gated early return with an explicit ignored test. Run the official converter comparison with --include-ignored in each PyO3 compatibility CI leg and document the local commands. Co-authored-by: Isaac --- .github/workflows/python-arrow.yml | 4 +--- crates/spark-python-arrow/README.md | 17 +++++++++++++---- crates/spark-python-arrow/tests/udf_results.rs | 7 ++----- 3 files changed, 16 insertions(+), 12 deletions(-) diff --git a/.github/workflows/python-arrow.yml b/.github/workflows/python-arrow.yml index b80020bd..ffafbd4c 100644 --- a/.github/workflows/python-arrow.yml +++ b/.github/workflows/python-arrow.yml @@ -41,8 +41,6 @@ jobs: fail-fast: false matrix: pyo3: ["0.28.3", "0.29.0", "0.30.1"] - env: - SPARK_PYTHON_ARROW_TEST_PYSPARK: "1" steps: - uses: actions/checkout@v4 - uses: actions/setup-python@v5 @@ -66,7 +64,7 @@ jobs: - name: Test native and official PySpark conversion semantics run: | cargo test --manifest-path "$RUNNER_TEMP/spark-python-arrow/Cargo.toml" \ - --locked --features pyo3/abi3-py39 + --locked --features pyo3/abi3-py39 -- --include-ignored - name: Check converter Clippy run: | cargo clippy --manifest-path "$RUNNER_TEMP/spark-python-arrow/Cargo.toml" \ diff --git a/crates/spark-python-arrow/README.md b/crates/spark-python-arrow/README.md index d1dc9444..1c017b1f 100644 --- a/crates/spark-python-arrow/README.md +++ b/crates/spark-python-arrow/README.md @@ -139,10 +139,19 @@ 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. -With official `pyspark-client==4.2.0` also installed, set -`SPARK_PYTHON_ARROW_TEST_PYSPARK=1` to enable comparisons against its nullable -String, Binary, and Variant result converters. The compatibility CI matrix sets -this flag; the default local command skips that optional comparison. +The comparison against official PySpark's nullable String, Binary, and Variant +result converters is marked `#[ignore]`: the default command reports it as +ignored, not passed. With official `pyspark-client==4.2.0` also installed, run all +tests including that comparison: + +```sh +CARGO_NET_OFFLINE=true CARGO_REGISTRIES_CRATES_IO_PROTOCOL=git \ + cargo test -p spark-python-arrow --locked --features pyo3/abi3-py39 -- --include-ignored +``` + +The compatibility CI matrix uses `--include-ignored` for every PyO3 version. +Explicitly running the comparison without its Python dependencies fails instead +of silently skipping it. The standalone tests cover buffer identity and lifetime, schema/type round trips, empty shapes, native scalar behavior, and UDF fallback/error ordering. The diff --git a/crates/spark-python-arrow/tests/udf_results.rs b/crates/spark-python-arrow/tests/udf_results.rs index e1b2434c..2aae74ee 100644 --- a/crates/spark-python-arrow/tests/udf_results.rs +++ b/crates/spark-python-arrow/tests/udf_results.rs @@ -18,7 +18,7 @@ //! UDF result-policy checks, without a Spark server. The fixtures model PySpark's nullable //! top-level String/Binary/Variant converters. Logging and injected failures make scheduling //! observable; they are test instrumentation, not support for arbitrary stateful converters. -//! Set SPARK_PYTHON_ARROW_TEST_PYSPARK=1 to additionally compare against installed PySpark. +//! Use --include-ignored with official PySpark installed to run the optional comparison. use std::ffi::CString; @@ -650,11 +650,8 @@ types = [ } #[test] +#[ignore = "requires official pyspark-client==4.2.0 and PyArrow"] fn official_pyspark_result_converters_match_native_and_fallback_paths() { - if std::env::var("SPARK_PYTHON_ARROW_TEST_PYSPARK").as_deref() != Ok("1") { - eprintln!("set SPARK_PYTHON_ARROW_TEST_PYSPARK=1 with official pyspark installed"); - return; - } Python::initialize(); Python::attach(|py| -> PyResult<()> { let locals = context(py)?; From a9a7169eb957ca2a1622d09c3fb4af616a71747a Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Tue, 27 Oct 2026 13:30:34 -0700 Subject: [PATCH 5/8] Narrow shared Python Arrow conversion to client use Defer the worker-only UDF policy and helpers until a concrete consumer can accompany them. Keep the adopted bridge and scalar paths, mark the implementation crate unpublished, use the repository license, and document schema-mismatch errors and the current compatibility contract. Co-authored-by: Isaac --- .github/workflows/python-arrow.yml | 8 +- crates/pyspark-rs/Cargo.toml | 2 +- crates/spark-python-arrow/Cargo.toml | 3 +- crates/spark-python-arrow/LICENSE | 201 ---- crates/spark-python-arrow/README.md | 102 +- crates/spark-python-arrow/src/bridge.rs | 1 + crates/spark-python-arrow/src/lib.rs | 6 +- crates/spark-python-arrow/src/scalar.rs | 84 +- crates/spark-python-arrow/src/udf/input.rs | 438 -------- crates/spark-python-arrow/src/udf/mod.rs | 92 -- crates/spark-python-arrow/src/udf/result.rs | 992 ------------------ crates/spark-python-arrow/tests/conversion.rs | 581 ++-------- crates/spark-python-arrow/tests/udf_inputs.rs | 294 ------ .../spark-python-arrow/tests/udf_results.rs | 704 ------------- dev/design/architecture.md | 10 +- docs/architecture.md | 4 +- 16 files changed, 156 insertions(+), 3366 deletions(-) delete mode 100644 crates/spark-python-arrow/LICENSE delete mode 100644 crates/spark-python-arrow/src/udf/input.rs delete mode 100644 crates/spark-python-arrow/src/udf/mod.rs delete mode 100644 crates/spark-python-arrow/src/udf/result.rs delete mode 100644 crates/spark-python-arrow/tests/udf_inputs.rs delete mode 100644 crates/spark-python-arrow/tests/udf_results.rs diff --git a/.github/workflows/python-arrow.yml b/.github/workflows/python-arrow.yml index ffafbd4c..a8119f6d 100644 --- a/.github/workflows/python-arrow.yml +++ b/.github/workflows/python-arrow.yml @@ -50,8 +50,8 @@ jobs: run: | rustup toolchain install stable --profile minimal --component clippy rustup default stable - - name: Install reference Python converters - run: python -m pip install "pyspark-client==4.2.0" "pyarrow==26.0.0" + - name: Install PyArrow + run: python -m pip install "pyarrow==26.0.0" - name: Isolate the standalone crate and select PyO3 # The workspace lockfile pins 0.30.1. Build outside it so each matrix leg # actually compiles the converter against the requested minor version. @@ -61,10 +61,10 @@ jobs: cargo update --manifest-path "$RUNNER_TEMP/spark-python-arrow/Cargo.toml" \ -p pyo3 --precise "${{ matrix.pyo3 }}" cargo tree --manifest-path "$RUNNER_TEMP/spark-python-arrow/Cargo.toml" --locked -i pyo3 - - name: Test native and official PySpark conversion semantics + - name: Test Arrow handoff and scalar conversion run: | cargo test --manifest-path "$RUNNER_TEMP/spark-python-arrow/Cargo.toml" \ - --locked --features pyo3/abi3-py39 -- --include-ignored + --locked --features pyo3/abi3-py39 - name: Check converter Clippy run: | cargo clippy --manifest-path "$RUNNER_TEMP/spark-python-arrow/Cargo.toml" \ diff --git a/crates/pyspark-rs/Cargo.toml b/crates/pyspark-rs/Cargo.toml index c9307c57..cb6bb62c 100644 --- a/crates/pyspark-rs/Cargo.toml +++ b/crates/pyspark-rs/Cargo.toml @@ -23,7 +23,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 = { version = "0.1.0", path = "../spark-python-arrow" } +spark-python-arrow = { path = "../spark-python-arrow" } pyo3.workspace = true arrow.workspace = true tokio.workspace = true diff --git a/crates/spark-python-arrow/Cargo.toml b/crates/spark-python-arrow/Cargo.toml index 9fc769ab..fdcbd9a3 100644 --- a/crates/spark-python-arrow/Cargo.toml +++ b/crates/spark-python-arrow/Cargo.toml @@ -19,7 +19,8 @@ name = "spark-python-arrow" version = "0.1.0" edition = "2021" license = "Apache-2.0" -description = "Arrow and Python conversion shared by Spark clients and Python workers" +publish = false +description = "Arrow and Python conversion for Spark clients" repository = "https://github.com/apache/spark-connect-rust" [dependencies] diff --git a/crates/spark-python-arrow/LICENSE b/crates/spark-python-arrow/LICENSE deleted file mode 100644 index 261eeb9e..00000000 --- a/crates/spark-python-arrow/LICENSE +++ /dev/null @@ -1,201 +0,0 @@ - Apache License - Version 2.0, January 2004 - http://www.apache.org/licenses/ - - TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION - - 1. Definitions. - - "License" shall mean the terms and conditions for use, reproduction, - and distribution as defined by Sections 1 through 9 of this document. - - "Licensor" shall mean the copyright owner or entity authorized by - the copyright owner that is granting the License. - - "Legal Entity" shall mean the union of the acting entity and all - other entities that control, are controlled by, or are under common - control with that entity. For the purposes of this definition, - "control" means (i) the power, direct or indirect, to cause the - direction or management of such entity, whether by contract or - otherwise, or (ii) ownership of fifty percent (50%) or more of the - outstanding shares, or (iii) beneficial ownership of such entity. - - "You" (or "Your") shall mean an individual or Legal Entity - exercising permissions granted by this License. - - "Source" form shall mean the preferred form for making modifications, - including but not limited to software source code, documentation - source, and configuration files. - - "Object" form shall mean any form resulting from mechanical - transformation or translation of a Source form, including but - not limited to compiled object code, generated documentation, - and conversions to other media types. - - "Work" shall mean the work of authorship, whether in Source or - Object form, made available under the License, as indicated by a - copyright notice that is included in or attached to the work - (an example is provided in the Appendix below). - - "Derivative Works" shall mean any work, whether in Source or Object - form, that is based on (or derived from) the Work and for which the - editorial revisions, annotations, elaborations, or other modifications - represent, as a whole, an original work of authorship. For the purposes - of this License, Derivative Works shall not include works that remain - separable from, or merely link (or bind by name) to the interfaces of, - the Work and Derivative Works thereof. - - "Contribution" shall mean any work of authorship, including - the original version of the Work and any modifications or additions - to that Work or Derivative Works thereof, that is intentionally - submitted to Licensor for inclusion in the Work by the copyright owner - or by an individual or Legal Entity authorized to submit on behalf of - the copyright owner. For the purposes of this definition, "submitted" - means any form of electronic, verbal, or written communication sent - to the Licensor or its representatives, including but not limited to - communication on electronic mailing lists, source code control systems, - and issue tracking systems that are managed by, or on behalf of, the - Licensor for the purpose of discussing and improving the Work, but - excluding communication that is conspicuously marked or otherwise - designated in writing by the copyright owner as "Not a Contribution." - - "Contributor" shall mean Licensor and any individual or Legal Entity - on behalf of whom a Contribution has been received by Licensor and - subsequently incorporated within the Work. - - 2. Grant of Copyright License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - copyright license to reproduce, prepare Derivative Works of, - publicly display, publicly perform, sublicense, and distribute the - Work and such Derivative Works in Source or Object form. - - 3. Grant of Patent License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - (except as stated in this section) patent license to make, have made, - use, offer to sell, sell, import, and otherwise transfer the Work, - where such license applies only to those patent claims licensable - by such Contributor that are necessarily infringed by their - Contribution(s) alone or by combination of their Contribution(s) - with the Work to which such Contribution(s) was submitted. If You - institute patent litigation against any entity (including a - cross-claim or counterclaim in a lawsuit) alleging that the Work - or a Contribution incorporated within the Work constitutes direct - or contributory patent infringement, then any patent licenses - granted to You under this License for that Work shall terminate - as of the date such litigation is filed. - - 4. Redistribution. You may reproduce and distribute copies of the - Work or Derivative Works thereof in any medium, with or without - modifications, and in Source or Object form, provided that You - meet the following conditions: - - (a) You must give any other recipients of the Work or - Derivative Works a copy of this License; and - - (b) You must cause any modified files to carry prominent notices - stating that You changed the files; and - - (c) You must retain, in the Source form of any Derivative Works - that You distribute, all copyright, patent, trademark, and - attribution notices from the Source form of the Work, - excluding those notices that do not pertain to any part of - the Derivative Works; and - - (d) If the Work includes a "NOTICE" text file as part of its - distribution, then any Derivative Works that You distribute must - include a readable copy of the attribution notices contained - within such NOTICE file, excluding those notices that do not - pertain to any part of the Derivative Works, in at least one - of the following places: within a NOTICE text file distributed - as part of the Derivative Works; within the Source form or - documentation, if provided along with the Derivative Works; or, - within a display generated by the Derivative Works, if and - wherever such third-party notices normally appear. The contents - of the NOTICE file are for informational purposes only and - do not modify the License. You may add Your own attribution - notices within Derivative Works that You distribute, alongside - or as an addendum to the NOTICE text from the Work, provided - that such additional attribution notices cannot be construed - as modifying the License. - - You may add Your own copyright statement to Your modifications and - may provide additional or different license terms and conditions - for use, reproduction, or distribution of Your modifications, or - for any such Derivative Works as a whole, provided Your use, - reproduction, and distribution of the Work otherwise complies with - the conditions stated in this License. - - 5. Submission of Contributions. Unless You explicitly state otherwise, - any Contribution intentionally submitted for inclusion in the Work - by You to the Licensor shall be under the terms and conditions of - this License, without any additional terms or conditions. - Notwithstanding the above, nothing herein shall supersede or modify - the terms of any separate license agreement you may have executed - with Licensor regarding such Contributions. - - 6. Trademarks. This License does not grant permission to use the trade - names, trademarks, service marks, or product names of the Licensor, - except as required for reasonable and customary use in describing the - origin of the Work and reproducing the content of the NOTICE file. - - 7. Disclaimer of Warranty. Unless required by applicable law or - agreed to in writing, Licensor provides the Work (and each - Contributor provides its Contributions) on an "AS IS" BASIS, - WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or - implied, including, without limitation, any warranties or conditions - of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A - PARTICULAR PURPOSE. You are solely responsible for determining the - appropriateness of using or redistributing the Work and assume any - risks associated with Your exercise of permissions under this License. - - 8. Limitation of Liability. In no event and under no legal theory, - whether in tort (including negligence), contract, or otherwise, - unless required by applicable law (such as deliberate and grossly - negligent acts) or agreed to in writing, shall any Contributor be - liable to You for damages, including any direct, indirect, special, - incidental, or consequential damages of any character arising as a - result of this License or out of the use or inability to use the - Work (including but not limited to damages for loss of goodwill, - work stoppage, computer failure or malfunction, or any and all - other commercial damages or losses), even if such Contributor - has been advised of the possibility of such damages. - - 9. Accepting Warranty or Additional Liability. While redistributing - the Work or Derivative Works thereof, You may choose to offer, - and charge a fee for, acceptance of support, warranty, indemnity, - or other liability obligations and/or rights consistent with this - License. However, in accepting such obligations, You may act only - on Your own behalf and on Your sole responsibility, not on behalf - of any other Contributor, and only if You agree to indemnify, - defend, and hold each Contributor harmless for any liability - incurred by, or claims asserted against, such Contributor by reason - of your accepting any such warranty or additional liability. - - END OF TERMS AND CONDITIONS - - APPENDIX: How to apply the Apache License to your work. - - To apply the Apache License to your work, attach the following - boilerplate notice, with the fields enclosed by brackets "[]" - replaced with your own identifying information. (Don't include - the brackets!) The text should be enclosed in the appropriate - comment syntax for the file format. We also recommend that a - file or class name and description of purpose be included on the - same "printed page" as the copyright notice for easier - identification within third-party archives. - - Copyright [yyyy] [name of copyright owner] - - Licensed 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. diff --git a/crates/spark-python-arrow/README.md b/crates/spark-python-arrow/README.md index 1c017b1f..f4ade527 100644 --- a/crates/spark-python-arrow/README.md +++ b/crates/spark-python-arrow/README.md @@ -19,10 +19,10 @@ under the License. # Shared Python and Arrow conversion -`spark-python-arrow` is a Rust library for Python-facing Spark clients and workers. -It depends on Arrow and PyO3, not Spark Connect transport, generated protobufs, -an async runtime, or a worker executable. Consumers use the same versioned source; -they should not maintain copied converter modules. +`spark-python-arrow` contains the Arrow bridge and scalar conversion routines +used by `pyspark-rs`. It depends on Arrow and PyO3, not Spark Connect transport, +generated protobufs, or an async runtime. Worker-side UDF conversion is outside +this crate's scope. ## Interfaces @@ -31,37 +31,19 @@ they should not maintain copied converter modules. shares decoded Arrow data buffers. Field names or timezone strings containing NUL use a documented IPC fallback at the table boundary; low-level C exports reject them. Length-prefixed schema metadata can contain NUL. -- `scalar`: borrowed scalar materialization and Python extraction. Exact-built-in - extraction is separate from coercive extraction so callers cannot accidentally - apply UDF coercion rules to client inputs. -- `udf`: the `SQL_ARROW_BATCHED_UDF` conversion policy. `input_spec` and - `result_spec` select supported native paths. `InputColumns` and `ResultColumn` - preserve the Python worker's conversion order, mutable-result handling, and - ArrowInvalid retry behavior. Python classes and converter callbacks come from - the caller; the library does not import a particular PySpark distribution. - `result_spec` resolves the caller's Variant class lazily, only after accepting - a native Variant conversion. - -Native UDF conversion covers signed integers, floats, booleans, strings, binary, -and tagged Variant structs, plus dates on input. Nested, decimal, timestamp, and -date-result columns retain the caller's Python fallback. Unsupported conversion -is `Ok(Err(reason))`; an error inspecting a supported conversion is `Err(PyErr)`. -Callers must not turn execution errors into unsupported-type fallback. - -### UDF converter contract - -The result policy accepts the stock, pure PySpark converters for **nullable -top-level UDF results**, not arbitrary callbacks or converters for non-nullable -fields. It can bypass those converters for `None` and exact strings/bytes whose -conversion is an identity. The Spark type check does not prove a callback has -these semantics. Callers must supply the matching PySpark converter; no probe is -performed because invoking a callback during policy selection can change error -timing and cannot prove its behavior for every value. - -Run the UDF for every row before calling `ResultColumn::finish`. Mutable results -are retained until then, and conversion errors are deferred so a later UDF error -still takes precedence. Once a result has been converted, a capacity fallback -retains that converted value instead of invoking its converter again. +- `scalar`: borrowed scalar materialization for collected rows and Python + extraction for client inputs. Extraction retains PyO3's coercion behavior; + the client chooses the order of type checks and handles nested values. + +`pyspark-rs` uses `bridge::record_batches_to_pyarrow` in `DataFrame.toArrow()`, +`scalar::Scalar` in row materialization, and `scalar::extract_*` in session input +conversion. These are the crate's current production consumers. + +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. ## Ownership and caller responsibilities @@ -73,8 +55,7 @@ Bare extension types/arrays and bare ordered dictionaries are rejected because Arrow's bare type cannot retain their field metadata or ordering flag. Schemas and record batches retain these properties on their fields. Extension types in dictionary-value positions are rejected, even within schemas, because that -position has no metadata-owning Rust field. UDF extension results select the -caller's fallback. +position has no metadata-owning Rust field. As with the Arrow C Data Interface itself, producers are responsible for valid buffer contents and all offsets. Imports perform limited structural checks, not full Arrow validation or a scan of every interior offset. A real PyArrow class @@ -94,36 +75,33 @@ require copying the validity bitmap during C Data export. Import/export release callbacks retain shared buffers after either language's temporary wrapper has been dropped. -UDF result builders allocate new Arrow buffers. A mixed native/fallback result -uses `pyarrow.concat_arrays` to join the prefix and suffix into the single array -the worker expects; that concatenation copies the combined column. Only the -handoff of a completed, compatible native array shares its buffers. - The client's owned `Value` and `Row` types remain in the Python-free Rust client. The Python binding adapts their scalar fields to `scalar::Scalar`; its row class, schema inference, and temporal/nested compatibility rules remain caller-owned. -Worker execution, cancellation, transport, and session settings likewise stay -outside this crate. +Cancellation, transport, and session settings likewise stay outside this crate. ## Integration -Cargo consumers use the crate as a normal Rust dependency. The library does not +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 manifest keeps its package metadata self-contained so the compatibility +tests can build it outside the workspace lockfile. The library does not enable PyO3's `extension-module` feature or select a minimum Python ABI: a Python extension chooses those features in the host package, while an embedded worker links its chosen Python runtime. PyO3 must resolve to one compatible version in each consuming dependency graph. The manifest permits PyO3 0.28, 0.29, and 0.30 so consumers on these lines can use the same source. The OSS Python extension selects 0.30.1, which requires Rust -1.85 or newer. The isolated converter compatibility CI job tests each minor +1.85 or newer; `pyspark-rs` declares that minimum independently of the workspace's +existing 1.82 default. The isolated converter compatibility CI job tests each minor version independently of the workspace lockfile. A version is not considered verified until its corresponding compatibility job passes. -Bazel consumers should pin the OSS revision or release using their existing -third-party Rust dependency mechanism and depend on this library target. There -must be no internal copy of `src/bridge.rs`, `src/scalar.rs`, or `src/udf`. -An unpublished local checkout may be used for development, but is not a release -pin and must not become a checked-in machine-specific path. - ## Tests With PyArrow installed in the Python environment used by PyO3: @@ -139,22 +117,8 @@ 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 comparison against official PySpark's nullable String, Binary, and Variant -result converters is marked `#[ignore]`: the default command reports it as -ignored, not passed. With official `pyspark-client==4.2.0` also installed, run all -tests including that comparison: - -```sh -CARGO_NET_OFFLINE=true CARGO_REGISTRIES_CRATES_IO_PROTOCOL=git \ - cargo test -p spark-python-arrow --locked --features pyo3/abi3-py39 -- --include-ignored -``` - -The compatibility CI matrix uses `--include-ignored` for every PyO3 version. -Explicitly running the comparison without its Python dependencies fails instead -of silently skipping it. - The standalone tests cover buffer identity and lifetime, schema/type round trips, -empty shapes, native scalar behavior, and UDF fallback/error ordering. The +empty shapes, schema mismatch errors, and scalar extraction/materialization. The Python client's `test_arrow_handoff.py` exercises the public `toArrow()` and -`toPandas()` paths. These are not substitutes for downstream worker differential -tests or Spark's complete Connect parity suite. +`toPandas()` paths. 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 index a1877217..3f85d14c 100644 --- a/crates/spark-python-arrow/src/bridge.rs +++ b/crates/spark-python-arrow/src/bridge.rs @@ -163,6 +163,7 @@ pub fn export_data_type(py: Python<'_>, data_type: &DataType) -> PyResult, batches: Vec) -> PyResult> { let schema = batches .first() diff --git a/crates/spark-python-arrow/src/lib.rs b/crates/spark-python-arrow/src/lib.rs index 946b1f58..825bd91d 100644 --- a/crates/spark-python-arrow/src/lib.rs +++ b/crates/spark-python-arrow/src/lib.rs @@ -15,12 +15,10 @@ // specific language governing permissions and limitations // under the License. -//! Python/Arrow conversion without a dependency on Spark transport or worker execution. +//! Python/Arrow conversion without a dependency on Spark transport. //! //! The bridge shares compatible Arrow value buffers; scalar and row materialization still -//! allocates Python objects. The caller owns its interpreter, Python classes, and error policy. -//! UDF-specific coercion and end-of-batch result conversion are explicit in the [`udf`] API. +//! allocates Python objects. The caller owns its interpreter and row representation. pub mod bridge; pub mod scalar; -pub mod udf; diff --git a/crates/spark-python-arrow/src/scalar.rs b/crates/spark-python-arrow/src/scalar.rs index 36c2fb92..6ad09326 100644 --- a/crates/spark-python-arrow/src/scalar.rs +++ b/crates/spark-python-arrow/src/scalar.rs @@ -15,17 +15,14 @@ // specific language governing permissions and limitations // under the License. -//! Scalar conversion shared by client row materialization and UDF columns. +//! Scalar conversion for client input extraction and row materialization. //! -//! Exact extraction never invokes a subclass's coercion methods. The separate -//! `extract_*` helpers retain PyO3's conversion semantics; callers choose their order, notably +//! The `extract_*` helpers retain PyO3's conversion semantics; callers choose their order, notably //! checking booleans before integers and Decimal before numeric coercion. This module does not //! infer a schema or define nested, decimal, timestamp, or application-specific row types. -use std::borrow::Cow; - use pyo3::prelude::*; -use pyo3::types::{PyBool, PyByteArray, PyBytes, PyFloat, PyInt, PyString}; +use pyo3::types::{PyByteArray, PyBytes, PyString}; use pyo3::IntoPyObjectExt; /// The ordinal of 1970-01-01 in Python's `datetime.date` calendar. @@ -74,15 +71,9 @@ pub fn utf8_into_python<'py>(py: Python<'py>, bytes: &[u8]) -> PyResult bool { - (-719_162..=2_932_896).contains(&days) -} - /// Materializes a date using an explicitly supplied, optionally cached `date.fromordinal`. /// -/// The constructor defines out-of-range errors. A caller reproducing another conversion path's -/// errors can check [`is_date_in_range`] before invoking it. +/// The constructor defines out-of-range errors. pub fn date_into_python<'py>( days: i32, fromordinal: &Bound<'py, PyAny>, @@ -90,58 +81,6 @@ pub fn date_into_python<'py>( fromordinal.call1((i64::from(days) + EPOCH_ORDINAL,)) } -/// An exact built-in boolean, excluding truthiness coercion. -pub fn exact_bool(value: &Bound<'_, PyAny>) -> Option { - value - .is_exact_instance_of::() - .then(|| value.extract().ok()) - .flatten() -} - -/// An exact built-in integer fitting in 64 bits, excluding booleans and integer subclasses. -pub fn exact_i64(value: &Bound<'_, PyAny>) -> Option { - value - .is_exact_instance_of::() - .then(|| value.extract().ok()) - .flatten() -} - -/// An exact built-in float, without integer or user-defined numeric coercion. -pub fn exact_f64(value: &Bound<'_, PyAny>) -> Option { - value - .is_exact_instance_of::() - .then(|| value.extract().ok()) - .flatten() -} - -/// An exact built-in string with a valid UTF-8 representation. -/// -/// PyO3 uses a supported ABI conversion and may allocate when the selected limited API cannot -/// borrow CPython's UTF-8 cache. Lone surrogates return `None`, leaving error/fallback policy to -/// the caller, rather than replacing characters or retaining a Python error indicator. -pub fn exact_str<'a>(value: &'a Bound<'_, PyAny>) -> Option> { - if !value.is_exact_instance_of::() { - return None; - } - value.cast::().ok()?.to_cow().ok() -} - -/// The borrowed contents of exact built-in `bytes`, excluding subclasses and mutable buffers. -pub fn exact_bytes<'a>(value: &'a Bound<'_, PyAny>) -> Option<&'a [u8]> { - if !value.is_exact_instance_of::() { - return None; - } - value.cast::().ok().map(|bytes| bytes.as_bytes()) -} - -/// Retains `value` only when it is exact built-in `bytes`. -pub fn into_exact_bytes(value: Bound<'_, PyAny>) -> Option> { - if !value.is_exact_instance_of::() { - return None; - } - value.cast_into::().ok() -} - /// Extracts with PyO3's boolean rules, not Python truthiness. pub fn extract_bool(value: &Bound<'_, PyAny>) -> PyResult { value.extract() @@ -172,18 +111,3 @@ pub fn extract_binary(value: &Bound<'_, PyAny>) -> Option> { None } } - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn date_range_matches_python_calendar_boundaries() { - assert!(is_date_in_range(-719_162)); - assert!(is_date_in_range(2_932_896)); - assert!(!is_date_in_range(-719_163)); - assert!(!is_date_in_range(2_932_897)); - assert!(!is_date_in_range(i32::MIN)); - assert!(!is_date_in_range(i32::MAX)); - } -} diff --git a/crates/spark-python-arrow/src/udf/input.rs b/crates/spark-python-arrow/src/udf/input.rs deleted file mode 100644 index 0349e822..00000000 --- a/crates/spark-python-arrow/src/udf/input.rs +++ /dev/null @@ -1,438 +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 arrow::array::{ - Array, ArrayRef, AsArray, BinaryArray, BooleanArray, Date32Array, Float32Array, Float64Array, - GenericBinaryArray, GenericStringArray, Int16Array, Int32Array, Int64Array, Int8Array, - LargeBinaryArray, LargeStringArray, OffsetSizeTrait, StringArray, StructArray, -}; -use arrow::datatypes::DataType; -use pyo3::exceptions::{PyIndexError, PyRuntimeError, PyUnicodeDecodeError}; -use pyo3::prelude::*; -use pyo3::types::{PyByteArray, PyBytes, PyDate, PyDict}; -use pyo3::{intern, PyTypeInfo}; - -use crate::scalar::{date_into_python, is_date_in_range, utf8_into_python, Scalar}; -use crate::udf::{import_batch, is_spark_type, variant_fields}; - -/// Accepts an input column of `data_type`, the Arrow type of `spark_type`, given PySpark's input -/// converter for it or `None` and the backend package's `sql.types` module; or returns the reason -/// the native UDF input converter does not support it. -pub fn input_spec( - data_type: &DataType, - spark_type: &Bound<'_, PyAny>, - converter: &Bound<'_, PyAny>, - sql_types: &Bound<'_, PyAny>, -) -> PyResult> { - let py = spark_type.py(); - let read = if converter.is_none() { - if !Column::reads(data_type) { - return Ok(Err(format!( - "{data_type} is not supported by the native UDF input converter" - ))); - } - Read::Plain - } else if matches!(data_type, DataType::Binary | DataType::LargeBinary) - && is_spark_type(spark_type, sql_types, "BinaryType")? - { - // PySpark's BinaryType converter returns the bytes PyArrow gives, or a bytearray copy - // when spark.sql.execution.pyspark.binaryAsBytes is off. - let sample = converter.call1((PyBytes::new(py, b""),))?; - if sample.is_exact_instance_of::() { - Read::Plain - } else if sample.is_exact_instance_of::() { - Read::ByteArray - } else { - return Ok(Err(format!( - "{spark_type}'s converter gives {}", - sample.get_type() - ))); - } - } else if is_variant_struct(data_type) && is_spark_type(spark_type, sql_types, "VariantType")? { - // The converter can return a VariantVal whose class differs from the backend module's. - // Take the class from its output to preserve that identity. - let sample = PyDict::new(py); - sample.set_item("value", PyBytes::new(py, b""))?; - sample.set_item("metadata", PyBytes::new(py, b""))?; - Read::Variant { - class: converter.call1((sample,))?.get_type().into_any().unbind(), - } - } else { - return Ok(Err(format!("{spark_type} needs PySpark's converter"))); - }; - Ok(Ok(InputSpec { - data_type: data_type.clone(), - read, - converter: converter.clone().unbind(), - })) -} - -/// Whether `data_type` is the struct Spark gives `VariantType`, with parts of either width. -fn is_variant_struct(data_type: &DataType) -> bool { - [DataType::Binary, DataType::LargeBinary] - .into_iter() - .any(|part| *data_type == DataType::Struct(variant_fields(part))) -} - -/// How [`InputColumns`] reads a column [`input_spec`] accepted. -pub struct InputSpec { - data_type: DataType, - read: Read, - /// PySpark's input converter for the column, or `None`. - converter: Py, -} - -impl InputSpec { - fn data_type(&self) -> &DataType { - &self.data_type - } - - pub fn clone_ref(&self, py: Python<'_>) -> InputSpec { - InputSpec { - data_type: self.data_type.clone(), - read: match &self.read { - Read::Plain => Read::Plain, - Read::ByteArray => Read::ByteArray, - Read::Variant { class } => Read::Variant { - class: class.clone_ref(py), - }, - }, - converter: self.converter.clone_ref(py), - } - } -} - -/// What a column's values become, as PySpark makes them. -pub(crate) enum Read { - /// What PyArrow's `to_pylist` gives, which PySpark keeps. - Plain, - /// A `bytearray` copy of each binary value. - ByteArray, - /// A VariantVal of `class` for each struct of the variant's two binaries. - Variant { class: Py }, -} - -/// A batch's input columns, read one cell at a time. -pub struct InputColumns { - batch: Py, - columns: Vec, - converters: Vec>, - convert_input: Py, - rows: usize, -} - -impl InputColumns { - /// Imports `batch`, a `pa.RecordBatch`, without copying it. Its columns must have the types - /// of `specs`. `convert_input` is PySpark's input step, `_convert_arrow_udf_input`, which - /// [`InputColumns::value`] runs to raise PySpark's error. - pub fn import( - batch: &Bound<'_, PyAny>, - specs: &[InputSpec], - convert_input: &Bound<'_, PyAny>, - ) -> PyResult { - let py = batch.py(); - let input = import_batch(batch)?; - let actual: Vec<&DataType> = input.columns().iter().map(|c| c.data_type()).collect(); - if actual.len() != specs.len() - || actual - .iter() - .zip(specs) - .any(|(actual, spec)| **actual != spec.data_type) - { - let expected: Vec<&DataType> = specs.iter().map(InputSpec::data_type).collect(); - return Err(PyRuntimeError::new_err(format!( - "the batch has types {actual:?}, but {expected:?} were expected" - ))); - } - let columns = input - .columns() - .iter() - .zip(specs) - .map(|(array, spec)| { - Column::new(py, array, &spec.read)?.ok_or_else(|| { - PyRuntimeError::new_err(format!("no reader for {}", array.data_type())) - }) - }) - .collect::>>()?; - Ok(InputColumns { - batch: batch.clone().unbind(), - columns, - converters: specs - .iter() - .map(|spec| spec.converter.clone_ref(py)) - .collect(), - convert_input: convert_input.clone().unbind(), - rows: input.len(), - }) - } - - pub fn num_rows(&self) -> usize { - self.rows - } - - /// What PySpark gives a UDF for `row` of `column`, as a new object. `column` indexes the - /// imported types, and `row` is below [`InputColumns::num_rows`]. - /// - /// A cell PySpark fails to convert, such as a string that is not UTF-8, raises what PySpark's - /// input step raises for its column. PySpark converts every input column, in order, before it - /// calls a UDF, so to fail as it does, read every cell, column by column in order, before - /// calling any UDF. - pub fn value<'py>( - &self, - py: Python<'py>, - column: usize, - row: usize, - ) -> PyResult> { - let Some(input) = self.columns.get(column) else { - return Err(PyIndexError::new_err(format!( - "column index {column} is out of bounds for {} columns", - self.columns.len() - ))); - }; - if row >= self.rows { - return Err(PyIndexError::new_err(format!( - "row index {row} is out of bounds for {} rows", - self.rows - ))); - } - input.value(py, row, || self.pyspark_error(py, column)) - } - - /// What PySpark's input step raises for `index`, a column with a cell it fails to convert. - fn pyspark_error(&self, py: Python<'_>, index: usize) -> PyErr { - let converted = (|| -> PyResult<_> { - let column = self - .batch - .bind(py) - .call_method1(intern!(py, "column"), (index,))?; - self.convert_input - .bind(py) - .call1((column, self.converters[index].bind(py))) - })(); - match converted { - Err(error) => error, - Ok(_) => PyRuntimeError::new_err(format!( - "PySpark converted input {index}, which the native converter cannot read" - )), - } - } -} - -/// A column, read as PySpark reads it. -pub(crate) enum Column { - Int8(Int8Array), - Int16(Int16Array), - Int32(Int32Array), - Int64(Int64Array), - Float32(Float32Array), - Float64(Float64Array), - Boolean(BooleanArray), - Utf8(StringArray), - LargeUtf8(LargeStringArray), - /// With `datetime.date.fromordinal`. - Date32(Date32Array, Py), - /// With whether each value is read as a `bytearray`. - Binary(BinaryArray, bool), - LargeBinary(LargeBinaryArray, bool), - Variant(Box>), - LargeVariant(Box>), -} - -/// A variant column, with its two parts and the class of VariantVal its values are made of. -pub(crate) struct VariantColumn { - array: StructArray, - value: GenericBinaryArray, - metadata: GenericBinaryArray, - class: Py, -} - -impl VariantColumn { - fn new(py: Python<'_>, array: &ArrayRef, class: &Py) -> Self { - let array = array.as_struct().clone(); - VariantColumn { - value: array.column(0).as_binary().clone(), - metadata: array.column(1).as_binary().clone(), - array, - class: class.clone_ref(py), - } - } - - fn value<'py>( - &self, - py: Python<'py>, - row: usize, - unconvertible: impl FnOnce() -> PyErr, - ) -> PyResult> { - if self.array.is_null(row) { - return Ok(py.None().into_bound(py)); - } - if self.value.is_null(row) || self.metadata.is_null(row) { - return Err(unconvertible()); - } - self.class.bind(py).call1(( - PyBytes::new(py, self.value.value(row)), - PyBytes::new(py, self.metadata.value(row)), - )) - } -} - -impl Column { - /// The types this native reader supports without an additional PySpark converter. - fn reads(data_type: &DataType) -> bool { - matches!( - data_type, - DataType::Int8 - | DataType::Int16 - | DataType::Int32 - | DataType::Int64 - | DataType::Float32 - | DataType::Float64 - | DataType::Boolean - | DataType::Utf8 - | DataType::LargeUtf8 - | DataType::Date32 - ) - } - - pub(crate) fn new(py: Python<'_>, array: &ArrayRef, read: &Read) -> PyResult> { - Ok(Some(match (array.data_type(), read) { - (DataType::Int8, Read::Plain) => Column::Int8(array.as_primitive().clone()), - (DataType::Int16, Read::Plain) => Column::Int16(array.as_primitive().clone()), - (DataType::Int32, Read::Plain) => Column::Int32(array.as_primitive().clone()), - (DataType::Int64, Read::Plain) => Column::Int64(array.as_primitive().clone()), - (DataType::Float32, Read::Plain) => Column::Float32(array.as_primitive().clone()), - (DataType::Float64, Read::Plain) => Column::Float64(array.as_primitive().clone()), - (DataType::Boolean, Read::Plain) => Column::Boolean(array.as_boolean().clone()), - (DataType::Utf8, Read::Plain) => Column::Utf8(array.as_string().clone()), - (DataType::LargeUtf8, Read::Plain) => Column::LargeUtf8(array.as_string().clone()), - (DataType::Date32, Read::Plain) => Column::Date32( - array.as_primitive().clone(), - PyDate::type_object(py).getattr("fromordinal")?.unbind(), - ), - (DataType::Binary, Read::Plain | Read::ByteArray) => { - Column::Binary(array.as_binary().clone(), matches!(read, Read::ByteArray)) - } - (DataType::LargeBinary, Read::Plain | Read::ByteArray) => { - Column::LargeBinary(array.as_binary().clone(), matches!(read, Read::ByteArray)) - } - (DataType::Struct(fields), Read::Variant { class }) - if *fields == variant_fields(DataType::Binary) => - { - Column::Variant(Box::new(VariantColumn::new(py, array, class))) - } - (DataType::Struct(fields), Read::Variant { class }) - if *fields == variant_fields(DataType::LargeBinary) => - { - Column::LargeVariant(Box::new(VariantColumn::new(py, array, class))) - } - _ => return Ok(None), - })) - } - - /// A cell PySpark fails to convert raises `unconvertible()`: a string that is not UTF-8, a date - /// outside `datetime.date`'s range, or a variant with a null part. - pub(crate) fn value<'py>( - &self, - py: Python<'py>, - row: usize, - unconvertible: impl FnOnce() -> PyErr, - ) -> PyResult> { - fn read<'a, 'py, A: Array>( - py: Python<'py>, - array: &A, - row: usize, - value: impl FnOnce() -> Scalar<'a>, - ) -> PyResult> { - if array.is_null(row) { - Scalar::Null.into_python(py) - } else { - value().into_python(py) - } - } - /// Decodes the slot's bytes rather than forming a `&str` from unvalidated UTF-8. - /// CPython validates only the UTF-8 payload. Offsets and buffer ranges remain the - /// Arrow producer's responsibility; this is not validation of arbitrary C data. - fn string<'py, O: OffsetSizeTrait>( - py: Python<'py>, - array: &GenericStringArray, - row: usize, - unconvertible: impl FnOnce() -> PyErr, - ) -> PyResult> { - if array.is_null(row) { - return Ok(py.None().into_bound(py)); - } - let offsets = array.value_offsets(); - let bytes = &array.value_data()[offsets[row].as_usize()..offsets[row + 1].as_usize()]; - match utf8_into_python(py, bytes) { - Ok(string) => Ok(string), - Err(error) if error.is_instance_of::(py) => { - Err(unconvertible()) - } - Err(error) => Err(error), - } - } - fn binary<'py, O: OffsetSizeTrait>( - py: Python<'py>, - array: &GenericBinaryArray, - row: usize, - bytearray: bool, - ) -> PyResult> { - if array.is_null(row) { - return Scalar::Null.into_python(py); - } - let bytes = array.value(row); - if bytearray { - Ok(PyByteArray::new(py, bytes).into_any()) - } else { - Scalar::Binary(bytes).into_python(py) - } - } - match self { - Column::Int8(array) => { - read(py, array, row, || Scalar::Int(i64::from(array.value(row)))) - } - Column::Int16(array) => { - read(py, array, row, || Scalar::Int(i64::from(array.value(row)))) - } - Column::Int32(array) => { - read(py, array, row, || Scalar::Int(i64::from(array.value(row)))) - } - Column::Int64(array) => read(py, array, row, || Scalar::Int(array.value(row))), - Column::Float32(array) => read(py, array, row, || { - Scalar::Float(f64::from(array.value(row))) - }), - Column::Float64(array) => read(py, array, row, || Scalar::Float(array.value(row))), - Column::Boolean(array) => read(py, array, row, || Scalar::Bool(array.value(row))), - Column::Utf8(array) => string(py, array, row, unconvertible), - Column::LargeUtf8(array) => string(py, array, row, unconvertible), - Column::Date32(array, fromordinal) => { - if array.is_null(row) { - return Ok(py.None().into_bound(py)); - } - let days = array.value(row); - if !is_date_in_range(days) { - return Err(unconvertible()); - } - date_into_python(days, fromordinal.bind(py)) - } - Column::Binary(array, bytearray) => binary(py, array, row, *bytearray), - Column::LargeBinary(array, bytearray) => binary(py, array, row, *bytearray), - Column::Variant(variant) => variant.value(py, row, unconvertible), - Column::LargeVariant(variant) => variant.value(py, row, unconvertible), - } - } -} diff --git a/crates/spark-python-arrow/src/udf/mod.rs b/crates/spark-python-arrow/src/udf/mod.rs deleted file mode 100644 index 950eef53..00000000 --- a/crates/spark-python-arrow/src/udf/mod.rs +++ /dev/null @@ -1,92 +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 <-> Python conversion one value at a time, for a per-row UDF loop that has to match the -//! `SQL_ARROW_BATCHED_UDF` loop in `pyspark.worker.read_udfs`. -//! -//! [`InputColumns`] reads a batch's input columns cell by cell, as PySpark's input step, -//! `_convert_arrow_udf_input`, reads them: `ArrowTableToRowsConversion._to_pylist`, then the -//! column's converter. [`ResultColumn`] builds one UDF's results into an Arrow array, as PySpark's -//! output step does: `pa.array(values, type=T)`, retried as -//! `pa.array(values).cast(T, safe=safecheck)` on `ArrowInvalid`. It appends built-in results, and -//! variant results whose parts are `bytes`, to an arrow-rs builder and hands any other result to -//! PyArrow. A caller keeps no list of types of its own: it asks [`input_spec`] about each input -//! column and [`result_spec`] about each result. -//! -//! A spec returns `Ok(Err(reason))` for an unsupported conversion; the caller may select its -//! existing Python implementation in that case. `Err(PyErr)` instead means inspecting the -//! supplied type/converter failed and must not silently become an unsupported-type fallback. -//! Runtime conversion failures retain Python's errors and result ordering. Nested, decimal, -//! and timestamp conversion is not implemented natively by this UDF policy. - -use std::collections::HashMap; - -use arrow::array::{ArrayRef, StructArray}; -use arrow::datatypes::{DataType, Field, Fields, Schema}; -use pyo3::prelude::*; - -mod input; -mod result; - -pub use input::{input_spec, InputColumns, InputSpec}; -pub use result::{result_spec, Finished, ResultColumn, ResultSpec}; - -/// The struct Spark gives `VariantType`, with parts of `part`: `Binary` from PySpark's -/// `to_arrow_type`, or `LargeBinary` from the JVM under `useLargeVarTypes`. Its tagged `metadata` -/// field tells it apart from a struct of two binaries. -fn variant_fields(part: DataType) -> Fields { - Fields::from(vec![ - Field::new("value", part.clone(), false), - Field::new("metadata", part, false) - .with_metadata(HashMap::from([("variant".to_owned(), "true".to_owned())])), - ]) -} - -/// Whether `spark_type` is exactly `sql_types.`, a type with a supported native converter. -/// `sql_types` is the caller-supplied backend module that parsed `spark_type`. -fn is_spark_type( - spark_type: &Bound<'_, PyAny>, - sql_types: &Bound<'_, PyAny>, - name: &str, -) -> PyResult { - Ok(spark_type.get_type().is(&sql_types.getattr(name)?)) -} - -/// Imports `schema`, a `pa.Schema`, such as the input schema a UDF's batches follow. -pub fn import_schema(schema: &Bound<'_, PyAny>) -> PyResult { - crate::bridge::import_schema(schema) -} - -/// Imports `data_type`, a `pa.DataType`, rejecting root extension types whose metadata -/// cannot be represented by a bare Rust data type. -pub fn import_data_type(data_type: &Bound<'_, PyAny>) -> PyResult { - crate::bridge::import_data_type(data_type) -} - -fn import_batch(batch: &Bound<'_, PyAny>) -> PyResult { - Ok(StructArray::from(crate::bridge::import_batch(batch)?)) -} - -/// Hands `array` to PyArrow, whose `pa.Array` is `array_type`, sharing compatible value buffers. -/// Validity-bitmap alignment may copy. PyArrow moves both descriptors' contents out, so dropping -/// them afterwards releases nothing. -fn export_array<'py>( - array_type: &Bound<'py, PyAny>, - array: ArrayRef, -) -> PyResult> { - crate::bridge::export_array_with_type(array_type, array) -} diff --git a/crates/spark-python-arrow/src/udf/result.rs b/crates/spark-python-arrow/src/udf/result.rs deleted file mode 100644 index 8100f36c..00000000 --- a/crates/spark-python-arrow/src/udf/result.rs +++ /dev/null @@ -1,992 +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::sync::Arc; - -use arrow::array::{ - Array, ArrayRef, AsArray, BinaryBuilder, BooleanBufferBuilder, BooleanBuilder, Float32Array, - Float64Builder, Int16Builder, Int32Builder, Int64Builder, Int8Builder, LargeBinaryBuilder, - LargeStringBuilder, NullBufferBuilder, PrimitiveBuilder, StringBuilder, StructArray, -}; -use arrow::buffer::BooleanBuffer; -use arrow::datatypes::{ArrowPrimitiveType, DataType, Float64Type}; -use pyo3::exceptions::PyRuntimeError; -use pyo3::intern; -use pyo3::prelude::*; -use pyo3::types::{PyBool, PyBytes, PyDict, PyFloat, PyInt, PyList, PyString, PyType}; - -use crate::bridge::arrow_error; -use crate::scalar::{exact_bool, exact_bytes, exact_f64, exact_i64, exact_str, into_exact_bytes}; -use crate::udf::input::{Column, Read}; -use crate::udf::{export_array, import_data_type, is_spark_type, variant_fields}; - -/// The largest integer magnitudes PyArrow converts to float64 and float32, which it refuses -/// above 2^53 and 2^24. -const MAX_EXACT_F64_INT: u64 = 1 << 53; -const MAX_EXACT_F32_INT: u64 = 1 << 24; - -/// Accepts a UDF's results for [`ResultColumn`], or returns why its native builder is unsupported. -/// `arrow_type` is the `pa.DataType` PySpark builds for `spark_return_type`, `converter` is -/// PySpark's result converter or `None`, `safecheck` is -/// `spark.sql.execution.pandas.convertToArrowArraySafely`, which PySpark's retry casts with, and -/// `sql_types` is the backend package's `sql.types` module, whose classes `spark_return_type` -/// comes from. For supported Variant results, `variant_class` resolves the class UDF code -/// constructs, which need not be the backend module's class. The resolver runs only after the -/// Arrow type and converter are accepted; its errors propagate. No package or class is inferred. -/// -/// # Converter contract -/// -/// `converter` must be PySpark's pure, top-level nullable result converter (as returned by -/// `LocalDataToArrowConversion._create_converter(..., nullable=True)`), or `None` for identity. -/// Non-nullable field converters and arbitrary stateful callbacks are not supported: this policy -/// keeps `None` unchanged and may bypass the converter for exact strings, bytes, and Variant -/// parts, or convert immutable built-ins before the remaining UDF calls. No callback is probed to -/// validate this contract, since probing would change its invocation and error timing. -pub fn result_spec<'py>( - py: Python<'py>, - arrow_type: &Bound<'_, PyAny>, - spark_return_type: &Bound<'_, PyAny>, - converter: &Bound<'_, PyAny>, - safecheck: bool, - sql_types: &Bound<'_, PyAny>, - variant_class: impl FnOnce() -> PyResult>>, -) -> PyResult> { - // Lowering an extension type to its storage type would select the primitive - // builder and export an array without the requested extension identity. - if crate::bridge::is_extension_type(arrow_type)? { - return Ok(Err( - "Arrow extension types require the caller's Python converter".to_owned(), - )); - } - let arrow_lib = py.import("pyarrow.lib")?; - if arrow_type.is_instance(&arrow_lib.getattr("DictionaryType")?)? { - return Ok(Err( - "Arrow dictionary types require the caller's Python converter".to_owned(), - )); - } - if arrow_type.is_instance(&arrow_lib.getattr("DataType")?)? { - let fields = arrow_type.getattr("num_fields")?.extract::()?; - if fields > 0 { - // Only the tagged Variant struct can use a nested native builder. Select the - // caller's fallback before importing other nested types, whose child metadata - // may not be representable by a bare Rust DataType (dictionary extensions, etc.). - let mut variant_candidate = - fields == 2 && arrow_type.is_instance(&arrow_lib.getattr("StructType")?)?; - if variant_candidate { - let binary = arrow_lib.call_method0("binary")?; - for index in 0..fields { - let child = arrow_type - .call_method1("field", (index,))? - .getattr("type")?; - if !child - .call_method1("equals", (&binary,))? - .extract::()? - { - variant_candidate = false; - break; - } - } - } - if !variant_candidate { - return Ok(Err(format!( - "{arrow_type} is not supported by the native UDF result converter" - ))); - } - } - } - let data_type = import_data_type(arrow_type)?; - let Some(kind) = OutputKind::of(&data_type) else { - return Ok(Err(format!( - "{data_type} is not supported by the native UDF result converter" - ))); - }; - // Build a variant only as PySpark's VariantType converter makes it. - let converter = match (converter.is_none(), kind.converter_type()) { - (true, _) if kind != OutputKind::Variant => None, - (false, Some(name)) if is_spark_type(spark_return_type, sql_types, name)? => { - Some(converter.clone().unbind()) - } - _ => { - return Ok(Err(format!( - "{spark_return_type} needs PySpark's converter" - ))) - } - }; - let variant = match kind { - // The caller knows the class UDF code constructs; it need not belong to `sql_types`. - OutputKind::Variant => match variant_class()? { - Some(class) => Some(class.into_any().unbind()), - None => { - return Ok(Err( - "VariantType requires the caller's VariantVal class".to_owned() - )) - } - }, - _ => None, - }; - let pyarrow = py.import("pyarrow")?; - Ok(Ok(ResultSpec { - kind, - arrow_type: arrow_type.clone().unbind(), - converter, - variant, - safecheck, - value_capacity: i32::MAX as usize - 1, - pyarrow_array: pyarrow.getattr("array")?.unbind(), - concat_arrays: pyarrow.getattr("concat_arrays")?.unbind(), - array_type: pyarrow.getattr("Array")?.unbind(), - arrow_invalid: pyarrow.getattr("lib")?.getattr("ArrowInvalid")?.unbind(), - })) -} - -/// A UDF's accepted result type, and what PySpark's output step needs for it. -pub struct ResultSpec { - kind: OutputKind, - arrow_type: Py, - /// PySpark's `StringType`, `BinaryType` or `VariantType` converter; every other result type - /// handled by this native builder has none. - converter: Option>, - /// PySpark's `VariantVal`, for a variant result. - variant: Option>, - safecheck: bool, - /// The most bytes the values of a string, binary or variant-part array can take, as its - /// offsets are 32-bit: PyArrow's limit, one under `i32::MAX`, past which `pa.array` chunks. - value_capacity: usize, - pyarrow_array: Py, - concat_arrays: Py, - array_type: Py, - arrow_invalid: Py, -} - -impl ResultSpec { - /// Lowers the short-offset buffer threshold before selecting the PyArrow fallback. - /// - /// This also permits testing chunk/fallback behavior without allocating 2 GiB. The threshold - /// cannot be raised above its current value or the maximum supported by 32-bit offsets. - pub fn with_value_capacity(self, value_capacity: usize) -> ResultSpec { - ResultSpec { - value_capacity: self.value_capacity.min(value_capacity), - ..self - } - } - - /// `value` through PySpark's result converter. Each returns `None` unchanged, and the - /// `StringType` and `BinaryType` converters return exact strs and bytes unchanged, so it is not - /// called for them. - fn convert<'py>( - &self, - py: Python<'py>, - value: Bound<'py, PyAny>, - ) -> PyResult> { - match &self.converter { - Some(converter) if !(value.is_none() || self.kind.converts_unchanged(&value)) => { - converter.bind(py).call1((value,)) - } - _ => Ok(value), - } - } - - /// The `except pa.lib.ArrowInvalid` branch of PySpark's output step: - /// `pa.array(values).cast(target_type=T, safe=safecheck)`. - fn retry<'py>( - &self, - py: Python<'py>, - values: Bound<'py, PyList>, - ) -> PyResult> { - let kwargs = PyDict::new(py); - kwargs.set_item(intern!(py, "target_type"), self.arrow_type.bind(py))?; - kwargs.set_item(intern!(py, "safe"), self.safecheck)?; - let inferred = self.pyarrow_array.bind(py).call1((values,))?; - inferred.call_method(intern!(py, "cast"), (), Some(&kwargs)) - } -} - -/// The result types this UDF policy builds natively. Date results use the caller's Python path; -/// they are not part of this native result policy. -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -enum OutputKind { - Int8, - Int16, - Int32, - Int64, - Float32, - Float64, - Boolean, - Utf8, - LargeUtf8, - Binary, - LargeBinary, - Variant, -} - -impl OutputKind { - fn of(data_type: &DataType) -> Option { - Some(match data_type { - DataType::Int8 => OutputKind::Int8, - DataType::Int16 => OutputKind::Int16, - DataType::Int32 => OutputKind::Int32, - DataType::Int64 => OutputKind::Int64, - DataType::Float32 => OutputKind::Float32, - DataType::Float64 => OutputKind::Float64, - DataType::Boolean => OutputKind::Boolean, - DataType::Utf8 => OutputKind::Utf8, - DataType::LargeUtf8 => OutputKind::LargeUtf8, - DataType::Binary => OutputKind::Binary, - DataType::LargeBinary => OutputKind::LargeBinary, - DataType::Struct(fields) if *fields == variant_fields(DataType::Binary) => { - OutputKind::Variant - } - _ => return None, - }) - } - - fn is_string(self) -> bool { - matches!(self, OutputKind::Utf8 | OutputKind::LargeUtf8) - } - - /// The Spark type whose PySpark converter results of this kind may go through. - fn converter_type(self) -> Option<&'static str> { - match self { - OutputKind::Utf8 | OutputKind::LargeUtf8 => Some("StringType"), - OutputKind::Binary | OutputKind::LargeBinary => Some("BinaryType"), - OutputKind::Variant => Some("VariantType"), - _ => None, - } - } - - /// Whether this kind's converter returns `value` unchanged. - fn converts_unchanged(self, value: &Bound<'_, PyAny>) -> bool { - match self { - OutputKind::Utf8 | OutputKind::LargeUtf8 => value.is_exact_instance_of::(), - OutputKind::Binary | OutputKind::LargeBinary => value.is_exact_instance_of::(), - _ => false, - } - } -} - -enum Builder { - Int8(Int8Builder), - Int16(Int16Builder), - Int32(Int32Builder), - Int64(Int64Builder), - /// Each result as the float64 PySpark's retry has to be given, which `finish` narrows as - /// PyArrow does, with the rows that held an exact `int`. - Float32(Float64Builder, BooleanBufferBuilder), - /// With the rows that held an exact `int`, which PySpark's retry has to be given back as one. - Float64(Float64Builder, BooleanBufferBuilder), - Boolean(BooleanBuilder), - Utf8(StringBuilder), - LargeUtf8(LargeStringBuilder), - Binary(BinaryBuilder), - LargeBinary(LargeBinaryBuilder), - /// The variants' two parts and which rows hold one. It takes no result when the UDF returns - /// it, as a VariantVal can change afterwards, only every result once the UDF has run on - /// every row: see [`Builder::append_variants`]. - Variant(BinaryBuilder, BinaryBuilder, NullBufferBuilder), -} - -/// What a [`Builder`] appended. -struct Appended { - /// As an array of the spec's type. - array: ArrayRef, - /// The same rows as the objects PySpark's retry is given: for float32, the float64 each was, - /// and for a string column, what the converter made of a built-in. - returned: ArrayRef, - /// For floats, the rows that held an exact `int`. - int_rows: Option, -} - -impl Builder { - fn new(kind: OutputKind, rows: usize) -> Builder { - match kind { - OutputKind::Int8 => Builder::Int8(Int8Builder::with_capacity(rows)), - OutputKind::Int16 => Builder::Int16(Int16Builder::with_capacity(rows)), - OutputKind::Int32 => Builder::Int32(Int32Builder::with_capacity(rows)), - OutputKind::Int64 => Builder::Int64(Int64Builder::with_capacity(rows)), - OutputKind::Float32 => Builder::Float32( - Float64Builder::with_capacity(rows), - BooleanBufferBuilder::new(rows), - ), - OutputKind::Float64 => Builder::Float64( - Float64Builder::with_capacity(rows), - BooleanBufferBuilder::new(rows), - ), - OutputKind::Boolean => Builder::Boolean(BooleanBuilder::with_capacity(rows)), - OutputKind::Utf8 => Builder::Utf8(StringBuilder::with_capacity(rows, rows * 8)), - OutputKind::LargeUtf8 => { - Builder::LargeUtf8(LargeStringBuilder::with_capacity(rows, rows * 8)) - } - OutputKind::Binary => Builder::Binary(BinaryBuilder::with_capacity(rows, rows * 8)), - OutputKind::LargeBinary => { - Builder::LargeBinary(LargeBinaryBuilder::with_capacity(rows, rows * 8)) - } - OutputKind::Variant => Builder::Variant( - BinaryBuilder::with_capacity(rows, rows * 8), - BinaryBuilder::with_capacity(rows, rows * 8), - NullBufferBuilder::new(rows), - ), - } - } - - /// Appends `value` if it is `None` or an exact `int`, `float`, `bool`, `str` or `bytes` that - /// PyArrow is known to convert to this type unchanged; appends nothing and returns false - /// otherwise, or if a `Utf8` or `Binary` builder would then hold more than `capacity` bytes. A - /// variant builder takes no value, not even `None`: see [`Builder::append_variants`]. - fn append_exact(&mut self, value: &Bound<'_, PyAny>, capacity: usize) -> bool { - if value.is_none() { - match self { - Builder::Int8(builder) => builder.append_null(), - Builder::Int16(builder) => builder.append_null(), - Builder::Int32(builder) => builder.append_null(), - Builder::Int64(builder) => builder.append_null(), - Builder::Float32(builder, ints) | Builder::Float64(builder, ints) => { - builder.append_null(); - ints.append(false); - } - Builder::Boolean(builder) => builder.append_null(), - Builder::Utf8(builder) => builder.append_null(), - Builder::LargeUtf8(builder) => builder.append_null(), - Builder::Binary(builder) => builder.append_null(), - Builder::LargeBinary(builder) => builder.append_null(), - Builder::Variant(..) => return false, - } - return true; - } - match self { - Builder::Int8(builder) => append_int(builder, value), - Builder::Int16(builder) => append_int(builder, value), - Builder::Int32(builder) => append_int(builder, value), - Builder::Int64(builder) => append_int(builder, value), - Builder::Float32(builder, ints) => { - append_float(builder, ints, value, MAX_EXACT_F32_INT) - } - Builder::Float64(builder, ints) => { - append_float(builder, ints, value, MAX_EXACT_F64_INT) - } - Builder::Boolean(builder) => { - if let Some(value) = exact_bool(value) { - builder.append_value(value); - true - } else { - false - } - } - Builder::Utf8(builder) => match exact_str(value) { - Some(text) if builder.values_slice().len() + text.len() <= capacity => { - builder.append_value(text); - true - } - _ => false, - }, - Builder::LargeUtf8(builder) => match exact_str(value) { - Some(text) => { - builder.append_value(text); - true - } - None => false, - }, - Builder::Binary(builder) => match exact_bytes(value) { - Some(bytes) if builder.values_slice().len() + bytes.len() <= capacity => { - builder.append_value(bytes); - true - } - _ => false, - }, - Builder::LargeBinary(builder) => match exact_bytes(value) { - Some(bytes) => { - builder.append_value(bytes); - true - } - None => false, - }, - Builder::Variant(..) => false, - } - } - - /// Appends `values`, all of a variant column's results, given once the UDF has run on every - /// row, if each is `None` or a VariantVal of exactly `class` whose two parts are exact - /// `bytes`: PySpark's converter then gives PyArrow those two, which it stores unchanged. - /// Appends nothing and returns false otherwise, if either part's values would take more than - /// `capacity` bytes, and for any other column. - fn append_variants( - &mut self, - values: &Bound<'_, PyList>, - class: &Bound<'_, PyAny>, - capacity: usize, - ) -> bool { - let Builder::Variant(value_parts, metadata_parts, nulls) = self else { - return false; - }; - let py = values.py(); - let mut parts = Vec::with_capacity(values.len()); - for value in values.iter() { - if value.is_none() { - parts.push(None); - continue; - } - if !value.get_type().is(class) { - return false; - } - // In the order PySpark's converter reads them. - let part = value - .getattr(intern!(py, "value")) - .ok() - .and_then(into_exact_bytes); - let metadata = value - .getattr(intern!(py, "metadata")) - .ok() - .and_then(into_exact_bytes); - let (Some(part), Some(metadata)) = (part, metadata) else { - return false; - }; - parts.push(Some((part, metadata))); - } - let (value_bytes, metadata_bytes) = parts.iter().flatten().fold( - (0, 0), - |(value_bytes, metadata_bytes), (part, metadata)| { - ( - value_bytes + part.as_bytes().len(), - metadata_bytes + metadata.as_bytes().len(), - ) - }, - ); - if value_bytes > capacity || metadata_bytes > capacity { - return false; - } - for part in &parts { - match part { - // PyArrow gives a null variant's two parts empty values, not nulls. - None => { - value_parts.append_value(b""); - metadata_parts.append_value(b""); - nulls.append_null(); - } - Some((part, metadata)) => { - value_parts.append_value(part.as_bytes()); - metadata_parts.append_value(metadata.as_bytes()); - nulls.append_non_null(); - } - } - } - true - } - - /// The bytes of values a `Utf8` or `Binary` builder holds; `None` for every other builder. - fn value_bytes(&self) -> Option { - match self { - Builder::Utf8(builder) => Some(builder.values_slice().len()), - Builder::Binary(builder) => Some(builder.values_slice().len()), - _ => None, - } - } - - /// The values so far. - fn finish(&mut self) -> PyResult { - fn as_returned(array: ArrayRef) -> Appended { - Appended { - returned: array.clone(), - array, - int_rows: None, - } - } - Ok(match self { - Builder::Int8(builder) => as_returned(Arc::new(builder.finish())), - Builder::Int16(builder) => as_returned(Arc::new(builder.finish())), - Builder::Int32(builder) => as_returned(Arc::new(builder.finish())), - Builder::Int64(builder) => as_returned(Arc::new(builder.finish())), - Builder::Float32(builder, ints) => { - let returned = builder.finish(); - // As PyArrow's static_cast does. - let narrowed: Float32Array = returned.unary(|value| value as f32); - Appended { - array: Arc::new(narrowed), - returned: Arc::new(returned), - int_rows: Some(ints.finish()), - } - } - Builder::Float64(builder, ints) => Appended { - int_rows: Some(ints.finish()), - ..as_returned(Arc::new(builder.finish())) - }, - Builder::Boolean(builder) => as_returned(Arc::new(builder.finish())), - Builder::Utf8(builder) => as_returned(Arc::new(builder.finish())), - Builder::LargeUtf8(builder) => as_returned(Arc::new(builder.finish())), - Builder::Binary(builder) => as_returned(Arc::new(builder.finish())), - Builder::LargeBinary(builder) => as_returned(Arc::new(builder.finish())), - Builder::Variant(value_parts, metadata_parts, nulls) => as_returned(Arc::new( - StructArray::try_new( - variant_fields(DataType::Binary), - vec![ - Arc::new(value_parts.finish()), - Arc::new(metadata_parts.finish()), - ], - nulls.finish(), - ) - .map_err(arrow_error)?, - )), - }) - } -} - -/// Appends an exact `float`, or an exact `int` of at most `max_int` in magnitude, which PyArrow -/// converts without rounding. -fn append_float( - builder: &mut Float64Builder, - ints: &mut BooleanBufferBuilder, - value: &Bound<'_, PyAny>, - max_int: u64, -) -> bool { - if let Some(value) = exact_f64(value) { - builder.append_value(value); - ints.append(false); - true - } else if let Some(int) = exact_i64(value).filter(|int| int.unsigned_abs() <= max_int) { - builder.append_value(int as f64); - ints.append(true); - true - } else { - false - } -} - -fn append_int(builder: &mut PrimitiveBuilder, value: &Bound<'_, PyAny>) -> bool -where - T: ArrowPrimitiveType, - T::Native: TryFrom, -{ - match exact_i64(value).and_then(|int| T::Native::try_from(int).ok()) { - Some(int) => { - builder.append_value(int); - true - } - None => false, - } -} - -/// Whether `value` is `None` or an exact `int`, `float`, `bool` or `str`, none of which can change -/// after the UDF returns it. -fn is_builtin(value: &Bound<'_, PyAny>) -> bool { - value.is_none() - || value.is_exact_instance_of::() - || value.is_exact_instance_of::() - || value.is_exact_instance_of::() - || value.is_exact_instance_of::() -} - -/// One UDF's results for one batch, taken through PySpark's steps in order: every result through -/// the converter, then all of them through `pa.array(values, type=T)`, retried as -/// `pa.array(values).cast(T, safe=safecheck)` on `ArrowInvalid`. -/// -/// PySpark runs the UDF on every row of the batch before it converts any result, so a result can -/// change after the UDF returns it, and a UDF's error on a later row wins over a result that -/// cannot be converted. [`ResultColumn::push`] therefore converts only built-ins, and the caller -/// must run the UDF on every row, giving each result to `push` or [`ResultColumn::extend`] in row -/// order, before it calls [`ResultColumn::finish`]. -pub struct ResultColumn<'s> { - spec: &'s ResultSpec, - builder: Builder, - state: State, -} - -enum State { - Building, - /// The results from the first one the builder could not take, in row order. The first may be - /// an already-converted string that exceeded capacity; the remainder is unconverted: - /// those in `listed`, then those in `pushed`. `listed` has the held part of each list given - /// to `extend` and the results pushed before the last such call; `pushed` has the results - /// pushed after it. - Holding { - listed: Option>, - pushed: Vec>, - }, - /// The converter raised, which PySpark does before it gives PyArrow any value, or a result - /// could not be held. - Failed(PyErr), -} - -/// A finished [`ResultColumn`]. -pub struct Finished<'py> { - /// The results as a `pa.Array` of the spec's type, or past 2 GiB the `pa.ChunkedArray` - /// PyArrow makes of them. - pub array: Bound<'py, PyAny>, - /// How many results the first `pa.array` call was given. - pub pyarrow_values: usize, - /// Whether PySpark's retry built the array. - pub retried: bool, -} - -impl<'s> ResultColumn<'s> { - /// A column for a batch of `rows` rows. - pub fn new(spec: &'s ResultSpec, rows: usize) -> ResultColumn<'s> { - ResultColumn { - spec, - builder: Builder::new(spec.kind, rows), - state: State::Building, - } - } - - /// Adds the next row's result. A built-in the builder takes is converted and appended now, as - /// it cannot change; from the first result that is not, every result is held for - /// [`ResultColumn::finish`]. Never raises: a converter error is raised by `finish`. - pub fn push<'py>(&mut self, py: Python<'py>, mut value: Bound<'py, PyAny>) { - match &mut self.state { - State::Building => {} - State::Holding { pushed, .. } => return pushed.push(value.unbind()), - State::Failed(_) => return, - } - if self.builder.append_exact(&value, self.spec.value_capacity) - || self.append_converted(py, &mut value) - { - return; - } - self.state = State::Holding { - listed: None, - pushed: vec![value.unbind()], - }; - } - - /// Adds the next rows' results, as pushing each of `values` in turn would, for a caller that - /// collected them into a list. From the first result the builder cannot take, the column keeps - /// a copy of the rest of `values` as one slice, so the caller may change `values` afterwards. - pub fn extend<'py>(&mut self, py: Python<'py>, values: &Bound<'py, PyList>) { - if let State::Building = self.state { - if let Err(error) = self.extend_building(py, values) { - self.state = State::Failed(error); - } - } else if let State::Holding { listed, pushed } = &mut self.state { - if let Err(error) = extend_held(py, listed, pushed, values) { - self.state = State::Failed(error); - } - } - } - - fn extend_building<'py>( - &mut self, - py: Python<'py>, - values: &Bound<'py, PyList>, - ) -> PyResult<()> { - // The iterator protocol reads a list with one call per item, and the abi3 list API with - // three. - for (index, value) in values.try_iter()?.enumerate() { - let mut value = value?; - if !(self.builder.append_exact(&value, self.spec.value_capacity) - || self.append_converted(py, &mut value)) - { - let held = values.get_slice(index, values.len()); - // The capacity fallback can follow a successful conversion. Retain that - // result rather than invoking its converter a second time at finish. - held.set_item(0, value)?; - self.state = State::Holding { - listed: Some(held.unbind()), - pushed: Vec::new(), - }; - return Ok(()); - } - if let State::Failed(_) = self.state { - return Ok(()); - } - } - Ok(()) - } - - /// Appends `value`, which the builder did not take, if it is a float, a bool or an int that - /// fits in 64 bits, which the `StringType` converter turns into a str it takes. Replaces - /// `value` with a successful conversion even when capacity requires holding it for `finish`. - /// Returns false if `value` has to be held instead. A converter error fails the column. - fn append_converted<'py>(&mut self, py: Python<'py>, value: &mut Bound<'py, PyAny>) -> bool { - // A built-in cannot change after the UDF returns it, so converting it now matches PySpark, - // which converts once the UDF has run on every row. For a string column the builder takes - // only None and exact strs, which the StringType converter returns unchanged. - if !self.spec.kind.is_string() || self.spec.converter.is_none() || !is_builtin(value) { - return false; - } - // A later row may call sys.set_int_max_str_digits, which decides whether str() takes an - // int of more than 640 digits. - // Keep large integers until every UDF call has completed so that limit changes are seen. - if value.is_exact_instance_of::() && exact_i64(value).is_none() { - return false; - } - match self.spec.convert(py, value.clone()) { - Ok(converted) => { - *value = converted; - self.builder.append_exact(value, self.spec.value_capacity) - } - Err(error) => { - self.state = State::Failed(error); - true - } - } - } - - /// The results as a `pa.Array`. Held results go through the converter and one - /// `pa.array(held, type=T)` call, and are joined to the appended ones; if PyArrow rejects one - /// with `ArrowInvalid`, PySpark's retry builds the whole batch instead. If the two would not - /// fit in one array, every result goes into that call, as in PySpark, where past 2 GiB it - /// returns a `pa.ChunkedArray` that fails the batch once every UDF has run. A converter error - /// is raised first, as PySpark converts every result before PyArrow sees one. A variant - /// column's results are all held, and are appended now if the builder takes every one. - pub fn finish<'py>(mut self, py: Python<'py>) -> PyResult> { - let mut held = match self.state { - State::Building => None, - State::Holding { listed, pushed } => Some(match listed { - None => PyList::new(py, pushed)?, - Some(listed) => { - let listed = listed.into_bound(py); - for value in pushed { - listed.append(value)?; - } - listed - } - }), - State::Failed(error) => return Err(error), - }; - if let (Some(values), Some(class)) = (&held, &self.spec.variant) { - if self - .builder - .append_variants(values, class.bind(py), self.spec.value_capacity) - { - held = None; - } - } - let head_bytes = self.builder.value_bytes(); - let appended = self.builder.finish()?; - let array_type = self.spec.array_type.bind(py); - let Some(held) = held else { - let array = export_array(array_type, appended.array)?; - return Ok(Finished { - array, - pyarrow_values: 0, - retried: false, - }); - }; - let converted = if self.spec.converter.is_some() { - let converted = held - .iter() - .map(|value| self.spec.convert(py, value)) - .collect::>>()?; - PyList::new(py, converted)? - } else { - held - }; - let pyarrow_array = self.spec.pyarrow_array.bind(py); - let arrow_invalid = self.spec.arrow_invalid.bind(py); - if appended.array.is_empty() - || head_bytes.is_some_and(|bytes| !fits(bytes, &converted, self.spec.value_capacity)) - { - let values = appended_values(py, &appended)?; - values.call_method1(intern!(py, "extend"), (&converted,))?; - let pyarrow_values = values.len(); - return match pyarrow_array.call1((&values, self.spec.arrow_type.bind(py))) { - Ok(array) => Ok(Finished { - array, - pyarrow_values, - retried: false, - }), - Err(error) if error.is_instance(py, arrow_invalid) => Ok(Finished { - array: self.spec.retry(py, values)?, - pyarrow_values, - retried: true, - }), - Err(error) => Err(error), - }; - } - let pyarrow_values = converted.len(); - match pyarrow_array.call1((&converted, self.spec.arrow_type.bind(py))) { - Ok(tail) => { - let head = export_array(array_type, appended.array)?; - let array = self - .spec - .concat_arrays - .bind(py) - .call1((PyList::new(py, [head, tail])?,))?; - Ok(Finished { - array, - pyarrow_values, - retried: false, - }) - } - Err(error) if error.is_instance(py, arrow_invalid) => { - let values = appended_values(py, &appended)?; - for value in converted.iter() { - values.append(value)?; - } - Ok(Finished { - array: self.spec.retry(py, values)?, - pyarrow_values, - retried: true, - }) - } - Err(error) => Err(error), - } - } -} - -/// Whether `values`, results through the converter, fit after `used` of `capacity` bytes in a -/// string or binary array. A value that is not `None`, an exact `str` or exact `bytes` is taken -/// not to. -fn fits(mut used: usize, values: &Bound<'_, PyList>, capacity: usize) -> bool { - for value in values.iter() { - used += if value.is_none() { - 0 - } else if let Some(text) = exact_str(&value) { - text.len() - } else if let Some(bytes) = exact_bytes(&value) { - bytes.len() - } else { - return false; - }; - if used > capacity { - return false; - } - } - true -} - -/// Appends `values` to the held results, after the pushed ones. -fn extend_held<'py>( - py: Python<'py>, - listed: &mut Option>, - pushed: &mut Vec>, - values: &Bound<'py, PyList>, -) -> PyResult<()> { - let list = listed - .get_or_insert_with(|| PyList::empty(py).unbind()) - .bind(py); - for value in pushed.drain(..) { - list.append(value)?; - } - list.call_method1(intern!(py, "extend"), (values,))?; - Ok(()) -} - -/// The appended results as the objects PySpark's retry is given, with the float rows that held -/// an exact `int` given back as one. -fn appended_values<'py>(py: Python<'py>, appended: &Appended) -> PyResult> { - let values = PyList::empty(py); - let returned = &appended.returned; - // A variant column appends nothing before it holds every result. - if returned.is_empty() { - return Ok(values); - } - let column = Column::new(py, returned, &Read::Plain)?.ok_or_else(|| { - PyRuntimeError::new_err(format!("no reader for {}", returned.data_type())) - })?; - for row in 0..returned.len() { - if appended - .int_rows - .as_ref() - .is_some_and(|ints| ints.value(row)) - { - values.append(returned.as_primitive::().value(row) as i64)?; - } else { - values.append(column.value(py, row, || { - PyRuntimeError::new_err(format!("appended result {row} cannot be read back")) - })?)?; - } - } - Ok(values) -} - -#[cfg(test)] -mod tests { - use arrow::datatypes::{Field, Fields, TimeUnit}; - use pyo3::exceptions::PyValueError; - - use super::*; - - #[test] - fn builds_a_variant_only_as_pyspark_types_it() { - assert_eq!( - OutputKind::of(&DataType::Struct(variant_fields(DataType::Binary))), - Some(OutputKind::Variant) - ); - let untagged = Fields::from(vec![ - Field::new("value", DataType::Binary, false), - Field::new("metadata", DataType::Binary, false), - ]); - for data_type in [ - DataType::Struct(untagged), - DataType::Decimal128(10, 2), - DataType::Timestamp(TimeUnit::Microsecond, None), - DataType::Date64, - ] { - assert_eq!(OutputKind::of(&data_type), None, "{data_type}"); - } - } - - #[test] - fn failed_state_raises_before_calling_pyarrow_and_retains_the_first_error() { - Python::initialize(); - Python::attach(|py| -> PyResult<()> { - for extend in [false, true] { - let locals = PyDict::new(py); - py.run( - c"from types import SimpleNamespace -class StringType: pass -sql_types = SimpleNamespace(StringType=StringType) -spark_type = StringType() -events = [] -def converter(value): - events.append(('converter', value)) - raise ValueError('first conversion failed') -def array_spy(*args, **kwargs): - events.append(('array',)) - raise AssertionError('PyArrow must not run after a converter error') -", - Some(&locals), - Some(&locals), - )?; - let arrow_type = py.import("pyarrow")?.call_method0("string")?; - let mut spec = result_spec( - py, - &arrow_type, - &locals.get_item("spark_type")?.unwrap(), - &locals.get_item("converter")?.unwrap(), - true, - &locals.get_item("sql_types")?.unwrap(), - || Ok(None), - )? - .expect("nullable string policy is supported"); - // Replace only this spec's cached callback, not a process-wide PyArrow alias. - spec.pyarrow_array = locals.get_item("array_spy")?.unwrap().unbind(); - let mut column = ResultColumn::new(&spec, 4); - let values = PyList::new(py, [1, 2])?; - if extend { - column.extend(py, &values); - } else { - for value in values.iter() { - column.push(py, value); - } - } - column.extend(py, &values); - column.push(py, py.None().into_bound(py)); - let error = column.finish(py).err().expect("converter must fail"); - assert!(error.is_instance_of::(py)); - assert_eq!(error.to_string(), "ValueError: first conversion failed"); - assert!(py - .eval( - c"events == [('converter', 1)]", - Some(&locals), - Some(&locals) - )? - .extract::()?); - } - Ok(()) - }) - .unwrap(); - } -} diff --git a/crates/spark-python-arrow/tests/conversion.rs b/crates/spark-python-arrow/tests/conversion.rs index 28f76c3a..87b7532d 100644 --- a/crates/spark-python-arrow/tests/conversion.rs +++ b/crates/spark-python-arrow/tests/conversion.rs @@ -16,7 +16,7 @@ // under the License. //! Standalone conversion checks against the installed PyArrow. No Spark server is needed. -//! These exercise the shared interfaces, not the full worker or client conformance suites. +//! These exercise the shared interfaces, not the full client conformance suite. use std::collections::HashMap; use std::ffi::CString; @@ -26,21 +26,18 @@ use arrow::array::{Array, Int64Array}; use arrow::datatypes::{DataType, Field, Schema, TimeUnit}; use arrow::record_batch::{RecordBatch, RecordBatchOptions}; use pyo3::exceptions::{ - PyImportError, PyIndexError, PyTypeError, PyUnicodeDecodeError, PyValueError, + PyOverflowError, PyTypeError, PyUnicodeDecodeError, PyUnicodeEncodeError, PyValueError, }; use pyo3::prelude::*; -use pyo3::types::{PyBytes, PyDict, PyList, PyType}; +use pyo3::types::{PyBytes, PyDict, PyList}; use spark_python_arrow::bridge::{ export_array, export_data_type, export_schema, import_array, import_batch, import_data_type, import_schema, record_batches_to_pyarrow, }; use spark_python_arrow::scalar::{ - exact_bool, exact_bytes, exact_f64, exact_i64, exact_str, extract_binary, utf8_into_python, + extract_binary, extract_bool, extract_f64, extract_i64, extract_string, utf8_into_python, Scalar, }; -use spark_python_arrow::udf::{ - input_spec, result_spec, Finished, InputColumns, ResultColumn, ResultSpec, -}; fn context(py: Python<'_>) -> PyResult> { let locals = PyDict::new(py); @@ -84,30 +81,6 @@ extension_array = pa.ExtensionArray.from_storage( Ok(locals) } -fn variant_context(py: Python<'_>) -> PyResult> { - let locals = context(py)?; - py.run( - c"class VariantType: pass - -class VariantVal: - def __init__(self, value, metadata): - self.value = value - self.metadata = metadata - -class SqlTypes: pass -SqlTypes.VariantType = VariantType -spark_type = VariantType() -converter = lambda value: {'value': value.value, 'metadata': value.metadata} -arrow_type = pa.struct([ - pa.field('value', pa.binary(), nullable=False), - pa.field('metadata', pa.binary(), nullable=False, metadata={b'variant': b'true'}), -])", - Some(&locals), - Some(&locals), - )?; - Ok(locals) -} - fn assert_array_equal(actual: &Bound<'_, PyAny>, expected: &Bound<'_, PyAny>) -> PyResult<()> { assert!( actual @@ -118,50 +91,6 @@ fn assert_array_equal(actual: &Bound<'_, PyAny>, expected: &Bound<'_, PyAny>) -> Ok(()) } -fn primitive_spec<'py>( - py: Python<'py>, - arrow_type: &Bound<'py, PyAny>, - safe: bool, -) -> PyResult { - let none = py.None().into_bound(py); - Ok( - result_spec(py, arrow_type, &none, &none, safe, &none, || Ok(None))? - .expect("the primitive type must be supported"), - ) -} - -fn build_results<'py>( - py: Python<'py>, - arrow_type: &Bound<'py, PyAny>, - values: &Bound<'py, PyList>, - safe: bool, -) -> PyResult> { - let spec = primitive_spec(py, arrow_type, safe)?; - let mut column = ResultColumn::new(&spec, values.len()); - column.extend(py, values); - column.finish(py) -} - -// The documented worker output sequence, including its ArrowInvalid-only retry. -fn reference_results<'py>( - py: Python<'py>, - arrow_type: &Bound<'py, PyAny>, - values: &Bound<'py, PyList>, - safe: bool, -) -> PyResult> { - let pa = py.import("pyarrow")?; - match pa.call_method1("array", (values, arrow_type)) { - Ok(array) => Ok(array), - Err(error) if error.is_instance(py, &pa.getattr("ArrowInvalid")?) => { - let kwargs = PyDict::new(py); - kwargs.set_item("safe", safe)?; - pa.call_method1("array", (values,))? - .call_method("cast", (arrow_type,), Some(&kwargs)) - } - Err(error) => Err(error), - } -} - #[test] fn exported_array_keeps_rust_buffers_alive_without_copying() { Python::initialize(); @@ -394,7 +323,6 @@ fn bare_extension_imports_fail_instead_of_losing_type_identity() { let array = evaluate(py, &locals, "extension_array")?; for error in [ import_data_type(&data_type).unwrap_err(), - spark_python_arrow::udf::import_data_type(&data_type).unwrap_err(), import_array(&array).unwrap_err(), ] { assert!(error.is_instance_of::(py), "{error}"); @@ -405,31 +333,6 @@ fn bare_extension_imports_fail_instead_of_losing_type_identity() { .unwrap(); } -#[test] -fn extension_results_select_python_fallback_without_lowering_the_type() { - Python::initialize(); - Python::attach(|py| -> PyResult<()> { - let locals = extension_context(py, "spark_python_arrow.tests.result_extension")?; - let data_type = evaluate(py, &locals, "extension_type")?; - let none = py.None().into_bound(py); - let unsupported = result_spec(py, &data_type, &none, &none, true, &none, || { - panic!("an unsupported extension must not resolve the Variant class") - })?; - let reason = match unsupported { - Ok(_) => panic!("an extension result must not use its primitive storage builder"), - Err(reason) => reason, - }; - assert!(reason.contains("extension types"), "{reason}"); - let wrong_type = evaluate(py, &locals, "object()")?; - let failure = result_spec(py, &wrong_type, &none, &none, true, &none, || { - panic!("an invalid Arrow type must not resolve the Variant class") - }); - assert!(matches!(failure, Err(error) if error.is_instance_of::(py))); - Ok(()) - }) - .unwrap(); -} - #[test] fn field_metadata_preserves_extensions_in_schemas_and_record_batches() { Python::initialize(); @@ -504,102 +407,6 @@ fn field_metadata_preserves_extensions_in_schemas_and_record_batches() { .unwrap(); } -#[test] -fn non_variant_and_unsupported_policies_do_not_resolve_the_variant_class() { - Python::initialize(); - Python::attach(|py| -> PyResult<()> { - let locals = variant_context(py)?; - let none = py.None().into_bound(py); - for (expression, supported) in [("pa.int64()", true), ("pa.list_(pa.int64())", false)] { - let arrow_type = evaluate(py, &locals, expression)?; - let selected = result_spec(py, &arrow_type, &none, &none, true, &none, || { - panic!("non-Variant and unsupported Arrow types must not resolve VariantVal") - })?; - assert_eq!(selected.is_ok(), supported); - } - let arrow_type = evaluate(py, &locals, "arrow_type")?; - let spark_type = evaluate(py, &locals, "spark_type")?; - let converter = evaluate(py, &locals, "converter")?; - let sql_types = evaluate(py, &locals, "SqlTypes")?; - for (spark_type, converter) in [(&spark_type, &none), (&none, &converter)] { - let selected = result_spec( - py, - &arrow_type, - spark_type, - converter, - true, - &sql_types, - || panic!("an unsupported converter must not resolve VariantVal"), - )?; - assert!(selected.is_err()); - } - Ok(()) - }) - .unwrap(); -} - -#[test] -fn supported_variant_policy_resolves_once_and_propagates_resolver_errors() { - Python::initialize(); - Python::attach(|py| -> PyResult<()> { - let locals = variant_context(py)?; - let arrow_type = evaluate(py, &locals, "arrow_type")?; - let spark_type = evaluate(py, &locals, "spark_type")?; - let converter = evaluate(py, &locals, "converter")?; - let sql_types = evaluate(py, &locals, "SqlTypes")?; - let mut calls = 0; - let failure = result_spec( - py, - &arrow_type, - &spark_type, - &converter, - true, - &sql_types, - || { - calls += 1; - Err(PyImportError::new_err("VariantVal resolver failed")) - }, - ); - assert_eq!(calls, 1); - let error = match failure { - Err(error) => error, - Ok(_) => panic!("a resolver failure must remain a Python error"), - }; - assert!(error.is_instance_of::(py)); - assert!(error.to_string().contains("VariantVal resolver failed")); - - let missing = result_spec( - py, - &arrow_type, - &spark_type, - &converter, - true, - &sql_types, - || Ok(None), - )?; - assert!(matches!(missing, Err(reason) if reason.contains("VariantVal class"))); - - let selected = result_spec( - py, - &arrow_type, - &spark_type, - &converter, - true, - &sql_types, - || { - calls += 1; - Ok(Some( - evaluate(py, &locals, "VariantVal")?.cast_into::()?, - )) - }, - )?; - assert!(selected.is_ok()); - assert_eq!(calls, 2); - Ok(()) - }) - .unwrap(); -} - #[test] fn table_export_rejects_inconsistent_schemas() { Python::initialize(); @@ -661,315 +468,131 @@ fn scalar_materialization_preserves_values_and_python_types() { } #[test] -fn exact_scalar_extraction_rejects_coercion_and_clears_rejected_errors() { +fn client_scalar_extraction_preserves_coercion_and_binary_copy_semantics() { Python::initialize(); Python::attach(|py| -> PyResult<()> { let locals = context(py)?; py.run( - c"class IntSubclass(int): pass\nclass StrSubclass(str): pass", + c"class IntSubclass(int): pass +class FloatSubclass(float): pass +class StrSubclass(str): pass +class BytesSubclass(bytes): pass +class ByteArraySubclass(bytearray): pass +class IndexValue: + def __index__(self): + return 42 +class FloatValue: + def __float__(self): + return 1.5 +class TruthValue: + def __bool__(self): + raise AssertionError('truthiness must not be invoked') +class StringValue: + def __str__(self): + raise AssertionError('stringification must not be invoked') +", Some(&locals), Some(&locals), )?; let truth = evaluate(py, &locals, "True")?; - assert_eq!(exact_bool(&truth), Some(true)); - assert_eq!(exact_i64(&truth), None); - let integer = evaluate(py, &locals, "7")?; - assert_eq!(exact_i64(&integer), Some(7)); - assert_eq!(exact_bool(&integer), None); - assert_eq!(exact_f64(&integer), None); - assert_eq!(exact_f64(&evaluate(py, &locals, "1.25")?), Some(1.25)); - for expression in ["IntSubclass(7)", "2**100", "-(2**100)"] { - assert_eq!(exact_i64(&evaluate(py, &locals, expression)?), None); - assert!(!PyErr::occurred(py)); + assert!(extract_bool(&truth)?); + assert!(!extract_bool(&evaluate(py, &locals, "False")?)?); + // The client tries bool before int because integer extraction also accepts bool. + assert_eq!(extract_i64(&truth)?, 1); + for expression in ["1", "None", "TruthValue()"] { + let error = extract_bool(&evaluate(py, &locals, expression)?).unwrap_err(); + assert!( + error.is_instance_of::(py), + "{expression}: {error}" + ); } - for expression in ["StrSubclass('text')", "'\\ud800'"] { - assert_eq!(exact_str(&evaluate(py, &locals, expression)?), None); - assert!(!PyErr::occurred(py)); + for (expression, expected) in [ + ("7", 7), + ("IntSubclass(7)", 7), + ("IndexValue()", 42), + ("-(2**63)", i64::MIN), + ("2**63 - 1", i64::MAX), + ] { + assert_eq!(extract_i64(&evaluate(py, &locals, expression)?)?, expected); } - let text = evaluate(py, &locals, "'a\\x00b'")?; - assert_eq!(exact_str(&text).as_deref(), Some("a\0b")); - let bytes = evaluate(py, &locals, "b'ab'")?; - assert_eq!(exact_bytes(&bytes), Some(b"ab".as_slice())); - assert_eq!(extract_binary(&bytes), Some(b"ab".to_vec())); - let mutable = evaluate(py, &locals, "bytearray(b'ab')")?; - assert_eq!(exact_bytes(&mutable), None); - assert_eq!(extract_binary(&mutable), Some(b"ab".to_vec())); - for expression in ["[97, 98]", "(97, 98)", "'ab'"] { - assert_eq!(extract_binary(&evaluate(py, &locals, expression)?), None); + for expression in ["1.5", "'7'"] { + let error = extract_i64(&evaluate(py, &locals, expression)?).unwrap_err(); + assert!( + error.is_instance_of::(py), + "{expression}: {error}" + ); } - let error = utf8_into_python(py, &[0xff]).unwrap_err(); - assert!(error.is_instance_of::(py)); - Ok(()) - }) - .unwrap(); -} - -#[test] -fn exact_primitive_results_match_pyarrow_without_python_value_fallback() { - Python::initialize(); - Python::attach(|py| -> PyResult<()> { - let locals = context(py)?; - for (data_type, values) in [ - ("pa.int8()", "[-128, None, 127]"), - ("pa.int16()", "[-32768, None, 32767]"), - ("pa.int32()", "[-(2**31), None, 2**31 - 1]"), - ("pa.int64()", "[-(2**63), None, 2**63 - 1]"), - ("pa.float32()", "[-1.25, None, 2**24, 0, -0.0]"), - ("pa.float64()", "[-1.25, None, 2**53, 0, -0.0]"), - ("pa.bool_()", "[True, None, False]"), - ("pa.string()", "['hello', None, 'a\\x00b', '']"), - ("pa.large_string()", "['hello', None, 'a\\x00b', '']"), - ("pa.binary()", "[b'hello', None, b'a\\x00b', b'']"), - ("pa.large_binary()", "[b'hello', None, b'a\\x00b', b'']"), - ] { - let arrow_type = evaluate(py, &locals, data_type)?; - let values = evaluate(py, &locals, values)?.cast_into::()?; - let expected = reference_results(py, &arrow_type, &values, true)?; - let actual = build_results(py, &arrow_type, &values, true)?; - assert_eq!(actual.pyarrow_values, 0, "{data_type}"); - assert!(!actual.retried, "{data_type}"); - assert_array_equal(&actual.array, &expected)?; + for expression in ["2**100", "-(2**100)"] { + let error = extract_i64(&evaluate(py, &locals, expression)?).unwrap_err(); + assert!( + error.is_instance_of::(py), + "{expression}: {error}" + ); } - Ok(()) - }) - .unwrap(); -} - -#[test] -fn subclass_results_use_fallback_and_match_pyarrow() { - Python::initialize(); - Python::attach(|py| -> PyResult<()> { - let locals = context(py)?; - py.run( - c"class IntSubclass(int): pass\nclass StrSubclass(str): pass", - Some(&locals), - Some(&locals), - )?; - // For variable-width columns, the capacity check cannot prove that a subclass fits. - // It sends the prefix and suffix to PyArrow together, rather than concatenating them. - for (data_type, values, pyarrow_values) in [ - ("pa.int64()", "[1, IntSubclass(2), None, 3]", 3), - ("pa.string()", "['a', StrSubclass('b'), None, 'c']", 4), + for (expression, expected) in [ + ("1.25", 1.25), + ("FloatSubclass(1.25)", 1.25), + ("7", 7.0), + ("FloatValue()", 1.5), + ("IndexValue()", 42.0), ] { - let arrow_type = evaluate(py, &locals, data_type)?; - let values = evaluate(py, &locals, values)?.cast_into::()?; - let expected = reference_results(py, &arrow_type, &values, true)?; - let actual = build_results(py, &arrow_type, &values, true)?; - assert_eq!(actual.pyarrow_values, pyarrow_values, "{data_type}"); - assert!(!actual.retried); - assert_array_equal(&actual.array, &expected)?; + assert_eq!(extract_f64(&evaluate(py, &locals, expression)?)?, expected); } - Ok(()) - }) - .unwrap(); -} - -#[test] -fn overflowing_results_follow_pyarrow_safe_and_unsafe_retry() { - Python::initialize(); - Python::attach(|py| -> PyResult<()> { - let locals = context(py)?; - for (data_type, values) in [ - ("pa.int8()", "[1, 128]"), - ("pa.int64()", "[1, 2**63]"), - ("pa.float32()", "[1, 2**24 + 1]"), - ("pa.float64()", "[1, 2**53 + 1]"), - ] { - let arrow_type = evaluate(py, &locals, data_type)?; - let values = evaluate(py, &locals, values)?.cast_into::()?; - for safe in [true, false] { - let expected = reference_results(py, &arrow_type, &values, safe); - let actual = build_results(py, &arrow_type, &values, safe); - match (actual, expected) { - (Ok(actual), Ok(expected)) => { - assert!(actual.pyarrow_values > 0, "{data_type}"); - assert_array_equal(&actual.array, &expected)?; - } - (Err(actual), Err(expected)) => { - assert!( - actual.get_type(py).is(expected.get_type(py)), - "{data_type}, safe={safe}: {actual}; expected {expected}" - ); - } - (actual, expected) => panic!( - "{data_type}, safe={safe}: actual success {}, expected success {}", - actual.is_ok(), - expected.is_ok() - ), - } - } + let error = extract_f64(&evaluate(py, &locals, "'1.25'")?).unwrap_err(); + assert!(error.is_instance_of::(py)); + for expression in ["'text'", "StrSubclass('text')"] { + assert_eq!(extract_string(&evaluate(py, &locals, expression)?)?, "text"); } - Ok(()) - }) - .unwrap(); -} - -#[test] -fn held_mutable_results_are_read_at_finish_and_extend_copies_the_list() { - Python::initialize(); - Python::attach(|py| -> PyResult<()> { - let locals = context(py)?; - let arrow_type = evaluate(py, &locals, "pa.binary()")?; - let values = evaluate(py, &locals, "[b'prefix', bytearray(b'b'), b'tail']")? - .cast_into::()?; - let spec = primitive_spec(py, &arrow_type, true)?; - let mut column = ResultColumn::new(&spec, 4); - column.extend(py, &values); - values.get_item(1)?.set_item(0, 99)?; - values.set_item(2, PyBytes::new(py, b"replacement"))?; - column.push(py, py.None().into_bound(py)); - let actual = column.finish(py)?; - let expected = evaluate(py, &locals, "pa.array([b'prefix', b'c', b'tail', None])")?; - assert_array_equal(&actual.array, &expected) - }) - .unwrap(); -} - -#[test] -fn lowered_variable_width_capacity_falls_back_for_the_whole_column() { - Python::initialize(); - Python::attach(|py| -> PyResult<()> { - let locals = context(py)?; - let arrow_type = evaluate(py, &locals, "pa.string()")?; - let values = evaluate(py, &locals, "['ab', 'cd', None]")?.cast_into::()?; - let spec = primitive_spec(py, &arrow_type, true)?.with_value_capacity(3); - let mut column = ResultColumn::new(&spec, values.len()); - column.extend(py, &values); - let actual = column.finish(py)?; - assert_eq!(actual.pyarrow_values, 3); - assert!(!actual.retried); - let expected = reference_results(py, &arrow_type, &values, true)?; - assert_array_equal(&actual.array, &expected) - }) - .unwrap(); -} - -#[test] -fn input_columns_match_pyarrow_cells_for_supported_primitive_types() { - Python::initialize(); - Python::attach(|py| -> PyResult<()> { - let locals = context(py)?; - let none = py.None().into_bound(py); - let convert_input = evaluate(py, &locals, "lambda column, converter: column.to_pylist()")?; - for (data_type, values) in [ - ("pa.int8()", "[-128, None, 127]"), - ("pa.int16()", "[-32768, None, 32767]"), - ("pa.int32()", "[-(2**31), None, 2**31 - 1]"), - ("pa.int64()", "[-(2**63), None, 2**63 - 1]"), - ("pa.float32()", "[-1.25, None, 2.5]"), - ("pa.float64()", "[-1.25, None, 2.5]"), - ("pa.bool_()", "[True, None, False]"), - ("pa.string()", "['hello', None, 'a\\x00b']"), - ("pa.large_string()", "['hello', None, 'a\\x00b']"), - ( - "pa.date32()", - "[datetime.date.min, None, datetime.date.max]", - ), + let text = evaluate(py, &locals, "'a\\x00b'")?; + assert_eq!(extract_string(&text)?, "a\0b"); + for expression in ["7", "StringValue()", "b'text'"] { + let error = extract_string(&evaluate(py, &locals, expression)?).unwrap_err(); + assert!( + error.is_instance_of::(py), + "{expression}: {error}" + ); + } + let error = extract_string(&evaluate(py, &locals, "'\\ud800'")?).unwrap_err(); + assert!(error.is_instance_of::(py)); + for expression in [ + "b'ab'", + "BytesSubclass(b'ab')", + "bytearray(b'ab')", + "ByteArraySubclass(b'ab')", ] { - let array = evaluate( - py, - &locals, - &format!("pa.array({values}, type={data_type})"), - )?; - let expected = array.call_method0("to_pylist")?; - let rust_type = import_data_type(&array.getattr("type")?)?; - let spec = input_spec(&rust_type, &none, &none, &none)? - .expect("the primitive type must be supported"); - let batch = py - .import("pyarrow")? - .call_method1("record_batch", (PyList::new(py, [array])?, vec!["x"]))?; - let input = InputColumns::import(&batch, &[spec], &convert_input)?; - assert_eq!(input.num_rows(), 3); - drop(batch); - for row in 0..input.num_rows() { - let actual = input.value(py, 0, row)?; - let expected = expected.get_item(row)?; - assert!(actual.eq(&expected)?, "{data_type}, row={row}"); - assert!( - actual.get_type().is(expected.get_type()), - "{data_type}, row={row}" - ); - } + assert_eq!( + extract_binary(&evaluate(py, &locals, expression)?), + Some(b"ab".to_vec()) + ); } - Ok(()) - }) - .unwrap(); -} - -#[test] -fn binary_input_policy_selects_bytes_or_bytearray() { - Python::initialize(); - Python::attach(|py| -> PyResult<()> { - let locals = context(py)?; - let none = py.None().into_bound(py); - // Binary policy belongs to the caller: provide its type identity and converter. - let binary_class = evaluate(py, &locals, "type('BinaryType', (), {})")?; - let binary_type = binary_class.call0()?; - let sql_types = py.import("types")?.getattr("SimpleNamespace")?.call0()?; - sql_types.setattr("BinaryType", binary_class)?; - let convert_input = evaluate( - py, - &locals, - "lambda column, converter: [converter(v) for v in column.to_pylist()]", - )?; - for data_type in ["pa.binary()", "pa.large_binary()"] { - let array = evaluate( - py, - &locals, - &format!("pa.array([b'hello', None, b'a\\x00b'], type={data_type})"), - )?; - let expected = array.call_method0("to_pylist")?; - let rust_type = import_data_type(&array.getattr("type")?)?; - assert!(input_spec(&rust_type, &none, &none, &none)?.is_err()); - let batch = py - .import("pyarrow")? - .call_method1("record_batch", (PyList::new(py, [array])?, vec!["x"]))?; - for policy in [ - "lambda v: v", - "lambda v: None if v is None else bytearray(v)", - ] { - let converter = evaluate(py, &locals, policy)?; - let spec = input_spec(&rust_type, &binary_type, &converter, &sql_types)? - .expect("both configured binary policies are supported"); - let input = InputColumns::import(&batch, &[spec], &convert_input)?; - for row in 0..input.num_rows() { - let actual = input.value(py, 0, row)?; - let expected = converter.call1((expected.get_item(row)?,))?; - assert!(actual.eq(&expected)?, "{data_type}, {policy}, row={row}"); - assert!( - actual.get_type().is(expected.get_type()), - "{data_type}, {policy}, row={row}" - ); - } - } + let mutable = evaluate(py, &locals, "bytearray(b'ab')")?; + let snapshot = extract_binary(&mutable).expect("bytearray is supported"); + mutable.set_item(0, 99)?; + assert_eq!(snapshot.as_slice(), b"ab"); + assert_eq!(extract_binary(&mutable), Some(b"cb".to_vec())); + for expression in ["[97, 98]", "(97, 98)", "'ab'", "memoryview(b'ab')"] { + assert_eq!(extract_binary(&evaluate(py, &locals, expression)?), None); } + assert!(!PyErr::occurred(py)); Ok(()) }) .unwrap(); } #[test] -fn input_columns_reject_mismatched_specs_and_out_of_range_cells() { +fn utf8_scalar_decoding_preserves_embedded_nul_and_rejects_invalid_bytes() { Python::initialize(); Python::attach(|py| -> PyResult<()> { - let locals = context(py)?; - let none = py.None().into_bound(py); - let convert_input = evaluate(py, &locals, "lambda column, converter: column.to_pylist()")?; - let batch = evaluate( - py, - &locals, - "pa.record_batch([pa.array([1, 2], type=pa.int64())], names=['x'])", - )?; - let wrong = input_spec(&DataType::Int32, &none, &none, &none)?.expect("int32 is supported"); - assert!(InputColumns::import(&batch, &[wrong], &convert_input).is_err()); - assert!(InputColumns::import(&batch, &[], &convert_input).is_err()); - let right = input_spec(&DataType::Int64, &none, &none, &none)?.expect("int64 is supported"); - let columns = InputColumns::import(&batch, &[right], &convert_input)?; - for (column, row) in [(1, 0), (0, 2), (usize::MAX, 0), (0, usize::MAX)] { - let error = columns.value(py, column, row).unwrap_err(); - assert!(error.is_instance_of::(py)); + for value in ["", "a\0b", "\u{03bb}"] { + assert_eq!( + utf8_into_python(py, value.as_bytes())?.extract::()?, + value + ); + } + for bytes in [&[0xff][..], &[0xc3][..]] { + let error = utf8_into_python(py, bytes).unwrap_err(); + assert!(error.is_instance_of::(py)); } - assert_eq!(columns.value(py, 0, 1)?.extract::()?, 2); Ok(()) }) .unwrap(); diff --git a/crates/spark-python-arrow/tests/udf_inputs.rs b/crates/spark-python-arrow/tests/udf_inputs.rs deleted file mode 100644 index 40a69cc8..00000000 --- a/crates/spark-python-arrow/tests/udf_inputs.rs +++ /dev/null @@ -1,294 +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. - -//! UDF input conversion and deferred Python errors, using valid Arrow buffer layouts. - -use arrow::datatypes::DataType; -use pyo3::exceptions::{PyRuntimeError, PyUnicodeDecodeError}; -use pyo3::prelude::*; -use pyo3::types::{PyDict, PyList}; -use spark_python_arrow::bridge::import_data_type; -use spark_python_arrow::udf::{input_spec, InputColumns}; - -fn context(py: Python<'_>) -> PyResult> { - let locals = PyDict::new(py); - locals.set_item("pa", py.import("pyarrow")?)?; - // This is the worker's input sequence for the primitive types tested here: - // materialize the whole column, then apply its optional caller converter. - py.run( - c"calls = [] -def convert_input(column, converter): - calls.append((column, converter)) - values = column.to_pylist() - return values if converter is None else [converter(value) for value in values] -", - Some(&locals), - Some(&locals), - )?; - Ok(locals) -} - -fn get<'py>(locals: &Bound<'py, PyDict>, name: &str) -> PyResult> { - Ok(locals.get_item(name)?.expect("fixture name is defined")) -} - -fn primitive_input<'py>( - py: Python<'py>, - array: &Bound<'py, PyAny>, - convert_input: &Bound<'py, PyAny>, -) -> PyResult { - let none = py.None().into_bound(py); - let data_type = import_data_type(&array.getattr("type")?)?; - let first = input_spec(&DataType::Int64, &none, &none, &none)? - .expect("the leading integer column is supported"); - let second = input_spec(&data_type, &none, &none, &none)? - .expect("the tested primitive column is supported"); - let pa = py.import("pyarrow")?; - let leading = pa.call_method1("array", (vec![1_i64; array.len()?],))?; - let batch = pa.call_method1( - "record_batch", - ( - PyList::new(py, [leading, array.clone()])?, - vec!["first", "second"], - ), - )?; - InputColumns::import(&batch, &[first, second], convert_input) -} - -fn assert_same_error(py: Python<'_>, actual: &PyErr, expected: &PyErr) -> PyResult<()> { - assert!(actual.get_type(py).is(expected.get_type(py))); - assert_eq!( - actual.value(py).str()?.extract::()?, - expected.value(py).str()?.extract::()? - ); - Ok(()) -} - -fn assert_one_callback(locals: &Bound<'_, PyDict>, array: &Bound<'_, PyAny>) -> PyResult<()> { - let calls = get(locals, "calls")?; - assert_eq!(calls.len()?, 1); - let call = calls.get_item(0)?; - assert!(call - .get_item(0)? - .call_method1("equals", (array,))? - .extract::()?); - assert!(call.get_item(1)?.is_none()); - Ok(()) -} - -#[test] -fn out_of_python_range_date_uses_the_original_column_callback() { - Python::initialize(); - Python::attach(|py| -> PyResult<()> { - let locals = context(py)?; - let convert_input = get(&locals, "convert_input")?; - let none = py.None().into_bound(py); - for days in [-719_163_i32, 2_932_897, i32::MIN, i32::MAX] { - locals.set_item("days", days)?; - // Every day value fits Arrow date32, even when Python cannot represent it. - let array = py.eval( - c"pa.array([0, days, None], type=pa.date32())", - Some(&locals), - None, - )?; - let kwargs = PyDict::new(py); - kwargs.set_item("full", true)?; - array.call_method("validate", (), Some(&kwargs))?; - let expected = convert_input.call1((&array, &none)).unwrap_err(); - get(&locals, "calls")?.call_method0("clear")?; - let input = primitive_input(py, &array, &convert_input)?; - assert_eq!(input.value(py, 0, 0)?.extract::()?, 1); - assert_eq!( - input.value(py, 1, 0)?.str()?.extract::()?, - "1970-01-01" - ); - assert!(input.value(py, 1, 2)?.is_none()); - assert_eq!(get(&locals, "calls")?.len()?, 0); - let actual = input.value(py, 1, 1).unwrap_err(); - assert_same_error(py, &actual, &expected)?; - assert_one_callback(&locals, &array)?; - } - Ok(()) - }) - .unwrap(); -} - -#[test] -fn invalid_utf8_payload_defers_to_the_python_column_error() { - Python::initialize(); - Python::attach(|py| -> PyResult<()> { - let locals = context(py)?; - locals.set_item("struct", py.import("struct")?)?; - let convert_input = get(&locals, "convert_input")?; - let none = py.None().into_bound(py); - for (type_name, offset_format) in [("string", "<4i"), ("large_string", "<4q")] { - locals.set_item( - "data_type", - py.import("pyarrow")?.getattr(type_name)?.call0()?, - )?; - locals.set_item("offset_format", offset_format)?; - // All offsets are monotonic and within this seven-byte payload. Only the text - // in the second cell is invalid UTF-8; no malformed offset or pointer is used. - let array = py.eval( - cr"pa.Array.from_buffers(data_type, 3, [ - None, - pa.py_buffer(struct.pack(offset_format, 0, 2, 3, 7)), - pa.py_buffer(b'ok\xfflast'), -])", - Some(&locals), - Some(&locals), - )?; - array.call_method0("validate")?; - let expected = convert_input.call1((&array, &none)).unwrap_err(); - assert!(expected.is_instance_of::(py)); - get(&locals, "calls")?.call_method0("clear")?; - let input = primitive_input(py, &array, &convert_input)?; - assert_eq!(input.value(py, 1, 0)?.extract::()?, "ok"); - assert_eq!(input.value(py, 1, 2)?.extract::()?, "last"); - assert_eq!(get(&locals, "calls")?.len()?, 0); - let actual = input.value(py, 1, 1).unwrap_err(); - assert_same_error(py, &actual, &expected)?; - assert_one_callback(&locals, &array)?; - } - Ok(()) - }) - .unwrap(); -} - -#[test] -fn callback_exception_is_propagated_without_replacement() { - Python::initialize(); - Python::attach(|py| -> PyResult<()> { - let locals = context(py)?; - py.run( - c"sentinel = LookupError('caller conversion failed') -def convert_input(column, converter): - calls.append((column, converter)) - raise sentinel -array = pa.array([2932897], type=pa.date32()) -", - Some(&locals), - Some(&locals), - )?; - let array = get(&locals, "array")?; - let input = primitive_input(py, &array, &get(&locals, "convert_input")?)?; - let error = input.value(py, 1, 0).unwrap_err(); - assert!(error.value(py).is(get(&locals, "sentinel")?)); - assert_one_callback(&locals, &array) - }) - .unwrap(); -} - -#[test] -fn unexpected_callback_success_is_not_silently_accepted() { - Python::initialize(); - Python::attach(|py| -> PyResult<()> { - let locals = context(py)?; - py.run( - c"def convert_input(column, converter): - calls.append((column, converter)) - return ['unexpected success'] -array = pa.array([2932897], type=pa.date32()) -", - Some(&locals), - Some(&locals), - )?; - let array = get(&locals, "array")?; - let input = primitive_input(py, &array, &get(&locals, "convert_input")?)?; - let error = input.value(py, 1, 0).unwrap_err(); - assert!(error.is_instance_of::(py)); - assert_eq!( - error.value(py).str()?.extract::()?, - "PySpark converted input 1, which the native converter cannot read" - ); - assert_one_callback(&locals, &array) - }) - .unwrap(); -} - -#[test] -fn variant_input_preserves_class_values_and_masked_null_parts() { - Python::initialize(); - Python::attach(|py| -> PyResult<()> { - let locals = context(py)?; - py.run( - c"class VariantType: pass -class VariantVal: - def __init__(self, value, metadata): - self.value = value - self.metadata = metadata - def __eq__(self, other): - return type(self) is type(other) and (self.value, self.metadata) == ( - other.value, other.metadata - ) -class SqlTypes: pass -SqlTypes.VariantType = VariantType -spark_type = VariantType() -converter = lambda value: None if value is None else VariantVal( - value['value'], value['metadata'] -) -", - Some(&locals), - Some(&locals), - )?; - for width in ["binary", "large_binary"] { - locals.set_item("part_type", py.import("pyarrow")?.getattr(width)?.call0()?)?; - py.run( - cr"fields = [ - pa.field('value', part_type, nullable=False), - pa.field('metadata', part_type, nullable=False, metadata={b'variant': b'true'}), -] -array = pa.StructArray.from_arrays([ - pa.array([b'a', None, b'\x00'], type=part_type), - pa.array([b'm', None, b'\xff'], type=part_type), -], fields=fields, mask=pa.array([False, True, False])) -array.validate(full=True) -expected = convert_input(array, converter) -calls.clear() -batch = pa.record_batch([array], names=['variant']) -", - Some(&locals), - Some(&locals), - )?; - let data_type = import_data_type(&get(&locals, "array")?.getattr("type")?)?; - let spec = input_spec( - &data_type, - &get(&locals, "spark_type")?, - &get(&locals, "converter")?, - &get(&locals, "SqlTypes")?, - )? - .expect("tagged Variant inputs are supported"); - let input = InputColumns::import( - &get(&locals, "batch")?, - &[spec], - &get(&locals, "convert_input")?, - )?; - for row in 0..input.num_rows() { - let actual = input.value(py, 0, row)?; - let expected = get(&locals, "expected")?.get_item(row)?; - assert!(actual.eq(&expected)?); - assert!(actual.get_type().is(expected.get_type())); - } - // The parent-null row is valid even though its non-nullable parts are null: - // those child nulls are masked, so no fallback or child access is required. - assert!(input.value(py, 0, 1)?.is_none()); - assert_eq!(get(&locals, "calls")?.len()?, 0); - } - Ok(()) - }) - .unwrap(); -} diff --git a/crates/spark-python-arrow/tests/udf_results.rs b/crates/spark-python-arrow/tests/udf_results.rs deleted file mode 100644 index 2aae74ee..00000000 --- a/crates/spark-python-arrow/tests/udf_results.rs +++ /dev/null @@ -1,704 +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. - -//! UDF result-policy checks, without a Spark server. The fixtures model PySpark's nullable -//! top-level String/Binary/Variant converters. Logging and injected failures make scheduling -//! observable; they are test instrumentation, not support for arbitrary stateful converters. -//! Use --include-ignored with official PySpark installed to run the optional comparison. - -use std::ffi::CString; - -use pyo3::exceptions::{PyRuntimeError, PyValueError}; -use pyo3::prelude::*; -use pyo3::types::{PyDict, PyList, PyType}; -use spark_python_arrow::udf::{result_spec, Finished, ResultColumn, ResultSpec}; - -#[derive(Clone, Copy, Debug)] -enum Mode { - Push, - Extend, -} - -const MODES: [Mode; 2] = [Mode::Push, Mode::Extend]; - -fn context(py: Python<'_>) -> PyResult> { - let locals = PyDict::new(py); - py.run( - c"import pyarrow as pa -from types import SimpleNamespace - -class StringType: pass -class BinaryType: pass -class VariantType: pass -sql_types = SimpleNamespace( - StringType=StringType, BinaryType=BinaryType, VariantType=VariantType -) - -class VariantVal: - def __init__(self, value, metadata): - self.value = value - self.metadata = metadata - -class VariantSubclass(VariantVal): pass -class BytesSubclass(bytes): pass - -class Text: - def __init__(self, text): - self.text = text - def __str__(self): - return self.text - -class BadText: - def __str__(self): - raise ValueError('text conversion failed') - -variant_type = pa.struct([ - pa.field('value', pa.binary(), nullable=False), - pa.field('metadata', pa.binary(), nullable=False, metadata={b'variant': b'true'}), -]) - -def convert_string(value): - if value is None: - return None - return str(value).lower() if isinstance(value, bool) else str(value) - -def convert_binary(value): - if value is None: - return None - assert isinstance(value, (bytes, bytearray)) - return bytes(value) - -def convert_variant(value): - if value is None: - return None - if isinstance(value, VariantVal): - return {'value': value.value, 'metadata': value.metadata} - raise ValueError('malformed variant') - -events = [] - -def tracked_string(value): - events.append(value) - return convert_string(value) - -def tracked_binary(value): - events.append(value) - return convert_binary(value) - -def tracked_variant(value): - events.append(value) - return convert_variant(value) - -def faulty_string(value): - events.append(value) - if value == 7: - raise ValueError('injected converter failure') - return convert_string(value) - -def later_udf_error(value): - if value == 2: - raise RuntimeError('later UDF failure') - return 7 - -def reference_results(values, converter, arrow_type): - converted = [converter(value) for value in values] - try: - return pa.array(converted, type=arrow_type) - except pa.ArrowInvalid: - return pa.array(converted).cast(arrow_type, safe=True) -", - Some(&locals), - Some(&locals), - )?; - Ok(locals) -} - -fn evaluate<'py>( - py: Python<'py>, - locals: &Bound<'py, PyDict>, - expression: &str, -) -> PyResult> { - let code = CString::new(expression).expect("test expression contains no NUL"); - py.eval(&code, Some(locals), Some(locals)) -} - -fn policy<'py>( - py: Python<'py>, - locals: &Bound<'py, PyDict>, - arrow_type: &Bound<'py, PyAny>, - spark_type: &Bound<'py, PyAny>, - converter: &Bound<'py, PyAny>, -) -> PyResult { - let sql_types = evaluate(py, locals, "sql_types")?; - let spec = result_spec( - py, - arrow_type, - spark_type, - converter, - true, - &sql_types, - || { - Ok(Some( - evaluate(py, locals, "VariantVal")?.cast_into::()?, - )) - }, - )?; - Ok(spec.expect("fixture type is supported")) -} - -fn fixture_policy<'py>( - py: Python<'py>, - locals: &Bound<'py, PyDict>, - arrow_type: &str, - spark_type: &str, - converter: &str, -) -> PyResult { - policy( - py, - locals, - &evaluate(py, locals, arrow_type)?, - &evaluate(py, locals, spark_type)?, - &evaluate(py, locals, converter)?, - ) -} - -fn add_values<'py>( - py: Python<'py>, - column: &mut ResultColumn<'_>, - values: &Bound<'py, PyList>, - mode: Mode, -) { - match mode { - Mode::Push => { - for value in values.iter() { - column.push(py, value); - } - } - Mode::Extend => column.extend(py, values), - } -} - -fn reference<'py>( - py: Python<'py>, - locals: &Bound<'py, PyDict>, - values: &Bound<'py, PyList>, - converter: &Bound<'py, PyAny>, - arrow_type: &Bound<'py, PyAny>, -) -> PyResult> { - evaluate(py, locals, "reference_results")?.call1((values, converter, arrow_type)) -} - -fn assert_same_result( - py: Python<'_>, - actual: PyResult>, - expected: PyResult>, -) -> PyResult<()> { - match (actual, expected) { - (Ok(actual), Ok(expected)) => assert!( - actual - .array - .call_method1("equals", (&expected,))? - .extract::()?, - "actual {}; expected {expected}", - actual.array - ), - (Err(actual), Err(expected)) => { - assert!(actual.get_type(py).is(expected.get_type(py))); - assert_eq!(actual.to_string(), expected.to_string()); - } - (actual, expected) => panic!( - "actual success {}, expected success {}", - actual.is_ok(), - expected.is_ok() - ), - } - Ok(()) -} - -#[test] -fn string_conversion_matches_nullable_pyspark_policy_in_both_ingestion_modes() { - Python::initialize(); - Python::attach(|py| -> PyResult<()> { - for mode in MODES { - let locals = context(py)?; - let values = evaluate(py, &locals, "[None, 'plain', True, False, 1, -7, 1.25]")? - .cast_into::()?; - for arrow_expression in ["pa.string()", "pa.large_string()"] { - evaluate(py, &locals, "events.clear()")?; - let spec = fixture_policy( - py, - &locals, - arrow_expression, - "StringType()", - "tracked_string", - )?; - assert!(evaluate(py, &locals, "not events")?.extract::()?); - let mut column = ResultColumn::new(&spec, values.len()); - add_values(py, &mut column, &values, mode); - let actual = column.finish(py)?; - assert_eq!(actual.pyarrow_values, 0); - assert!( - evaluate(py, &locals, "events == [True, False, 1, -7, 1.25]")? - .extract::()? - ); - let expected = reference( - py, - &locals, - &values, - &evaluate(py, &locals, "convert_string")?, - &evaluate(py, &locals, arrow_expression)?, - ); - assert_same_result(py, Ok(actual), expected)?; - } - } - Ok(()) - }) - .unwrap(); -} - -#[test] -fn capacity_fallback_retains_converted_value_and_calls_converter_once() { - Python::initialize(); - Python::attach(|py| -> PyResult<()> { - for mode in MODES { - for capacity in [0, 1] { - let locals = context(py)?; - let values = evaluate(py, &locals, "[1, 2, 3]")?.cast_into::()?; - let spec = - fixture_policy(py, &locals, "pa.string()", "StringType()", "tracked_string")? - .with_value_capacity(capacity); - let mut column = ResultColumn::new(&spec, values.len()); - add_values(py, &mut column, &values, mode); - let actual = column.finish(py)?; - assert_eq!(actual.pyarrow_values, 3); - assert!(evaluate(py, &locals, "events == [1, 2, 3]")?.extract::()?); - let expected = reference( - py, - &locals, - &values, - &evaluate(py, &locals, "convert_string")?, - &evaluate(py, &locals, "pa.string()")?, - ); - assert_same_result(py, Ok(actual), expected)?; - } - } - Ok(()) - }) - .unwrap(); -} - -#[test] -fn held_string_and_binary_callbacks_observe_final_mutable_values() { - Python::initialize(); - Python::attach(|py| -> PyResult<()> { - for mode in MODES { - for (arrow_type, spark_type, converter, values, mutation) in [ - ( - "pa.string()", - "StringType()", - "tracked_string", - "['prefix', Text('before'), None, True]", - "values[1].text = 'after'", - ), - ( - "pa.binary()", - "BinaryType()", - "tracked_binary", - "[b'prefix', bytearray(b'before'), None, BytesSubclass(b'tail')]", - "values[1][:] = b'after'", - ), - ] { - let locals = context(py)?; - let values = evaluate(py, &locals, values)?.cast_into::()?; - locals.set_item("values", &values)?; - let spec = fixture_policy(py, &locals, arrow_type, spark_type, converter)?; - let mut column = ResultColumn::new(&spec, values.len()); - add_values(py, &mut column, &values, mode); - assert!(evaluate(py, &locals, "not events")?.extract::()?); - py.run( - &CString::new(mutation).unwrap(), - Some(&locals), - Some(&locals), - )?; - let actual = column.finish(py)?; - assert!(actual.pyarrow_values > 0); - assert!( - evaluate(py, &locals, "len(events) == 2 and events[0] is values[1]")? - .extract::()? - ); - let pure_converter = converter.replacen("tracked_", "convert_", 1); - let expected = reference( - py, - &locals, - &values, - &evaluate(py, &locals, &pure_converter)?, - &evaluate(py, &locals, arrow_type)?, - ); - assert_same_result(py, Ok(actual), expected)?; - } - } - Ok(()) - }) - .unwrap(); -} - -#[test] -fn converter_failures_are_delayed_and_match_reference_error_order() { - Python::initialize(); - Python::attach(|py| -> PyResult<()> { - for mode in MODES { - for (arrow_type, spark_type, converter, values) in [ - ( - "pa.string()", - "StringType()", - "faulty_string", - "[1, 7, BadText()]", - ), - ( - "pa.string()", - "StringType()", - "convert_string", - "['prefix', BadText(), 7]", - ), - ( - "pa.binary()", - "BinaryType()", - "convert_binary", - "[b'prefix', object(), None]", - ), - ] { - let locals = context(py)?; - let values = evaluate(py, &locals, values)?.cast_into::()?; - let spec = fixture_policy(py, &locals, arrow_type, spark_type, converter)?; - let mut column = ResultColumn::new(&spec, values.len()); - add_values(py, &mut column, &values, mode); - // Neither ingestion API raises. Failed-state ingestion also preserves the - // first error rather than running a later callback. - column.push(py, py.None().into_bound(py)); - column.extend(py, &PyList::empty(py)); - let actual = column.finish(py); - assert!(actual.is_err()); - let expected = reference( - py, - &locals, - &values, - &evaluate(py, &locals, converter)?, - &evaluate(py, &locals, arrow_type)?, - ); - assert_same_result(py, actual, expected)?; - } - } - Ok(()) - }) - .unwrap(); -} - -#[test] -fn later_udf_error_wins_when_caller_finishes_only_after_all_udf_calls() { - Python::initialize(); - Python::attach(|py| -> PyResult<()> { - for mode in MODES { - let locals = context(py)?; - let spec = fixture_policy(py, &locals, "pa.string()", "StringType()", "faulty_string")?; - let udf = evaluate(py, &locals, "later_udf_error")?; - let outcome = (|| -> PyResult> { - let mut column = ResultColumn::new(&spec, 3); - let collected = PyList::empty(py); - for row in 0..3 { - let value = udf.call1((row,))?; - match mode { - Mode::Push => column.push(py, value), - Mode::Extend => collected.append(value)?, - } - } - if matches!(mode, Mode::Extend) { - column.extend(py, &collected); - } - column.finish(py) - })(); - let error = outcome.err().expect("the later UDF call must fail"); - assert!(error.is_instance_of::(py)); - assert!(error.to_string().contains("later UDF failure")); - } - Ok(()) - }) - .unwrap(); -} - -#[test] -fn large_integer_string_conversion_is_deferred_until_finish() { - Python::initialize(); - Python::attach(|py| -> PyResult<()> { - for mode in MODES { - for expression in ["[2**80]", "[2**16000]"] { - let locals = context(py)?; - let values = evaluate(py, &locals, expression)?.cast_into::()?; - let spec = - fixture_policy(py, &locals, "pa.string()", "StringType()", "tracked_string")?; - let mut column = ResultColumn::new(&spec, values.len()); - add_values(py, &mut column, &values, mode); - assert!(evaluate(py, &locals, "not events")?.extract::()?); - let actual = column.finish(py); - assert!(evaluate(py, &locals, "len(events) == 1")?.extract::()?); - // Some Python versions reject the larger value under their digit limit. - // Either outcome must match the same nullable converter at finish time. - let expected = reference( - py, - &locals, - &values, - &evaluate(py, &locals, "convert_string")?, - &evaluate(py, &locals, "pa.string()")?, - ); - assert_same_result(py, actual, expected)?; - } - } - Ok(()) - }) - .unwrap(); -} - -#[test] -fn native_variant_results_include_nulls_and_observe_mutation_before_finish() { - Python::initialize(); - Python::attach(|py| -> PyResult<()> { - for mode in MODES { - let locals = context(py)?; - let value = evaluate(py, &locals, "VariantVal(b'old', b'metadata')")?; - let values = PyList::new(py, [value.clone(), py.None().into_bound(py), value.clone()])?; - let spec = fixture_policy( - py, - &locals, - "variant_type", - "VariantType()", - "tracked_variant", - )?; - let mut column = ResultColumn::new(&spec, values.len()); - add_values(py, &mut column, &values, mode); - value.setattr("value", evaluate(py, &locals, "b'new'")?)?; - let actual = column.finish(py)?; - assert_eq!(actual.pyarrow_values, 0); - assert!(!actual.retried); - assert!(evaluate(py, &locals, "not events")?.extract::()?); - let expected = reference( - py, - &locals, - &values, - &evaluate(py, &locals, "convert_variant")?, - &evaluate(py, &locals, "variant_type")?, - ); - assert_same_result(py, Ok(actual), expected)?; - } - Ok(()) - }) - .unwrap(); -} - -#[test] -fn variant_subclass_non_bytes_and_capacity_use_fallback() { - Python::initialize(); - Python::attach(|py| -> PyResult<()> { - for mode in MODES { - for (expression, capacity) in [ - ("[VariantSubclass(b'v', b'm'), None]", usize::MAX), - ( - "[VariantVal(bytearray(b'v'), bytearray(b'm')), None]", - usize::MAX, - ), - ("[VariantVal(b'vv', b'mm'), None]", 1), - ] { - let locals = context(py)?; - let values = evaluate(py, &locals, expression)?.cast_into::()?; - let spec = fixture_policy( - py, - &locals, - "variant_type", - "VariantType()", - "tracked_variant", - )? - .with_value_capacity(capacity); - let mut column = ResultColumn::new(&spec, values.len()); - add_values(py, &mut column, &values, mode); - let actual = column.finish(py)?; - assert_eq!(actual.pyarrow_values, values.len()); - assert!(evaluate(py, &locals, "len(events) == 1")?.extract::()?); - let expected = reference( - py, - &locals, - &values, - &evaluate(py, &locals, "convert_variant")?, - &evaluate(py, &locals, "variant_type")?, - ); - assert_same_result(py, Ok(actual), expected)?; - } - } - Ok(()) - }) - .unwrap(); -} - -#[test] -fn variant_converter_failure_precedes_arrow_rejection() { - Python::initialize(); - Python::attach(|py| -> PyResult<()> { - for mode in MODES { - let locals = context(py)?; - // The first result reaches PyArrow only if every converter succeeds. Its invalid - // binary part must not mask the second result's conversion failure. - let values = evaluate(py, &locals, "[VariantVal(object(), b'm'), object()]")? - .cast_into::()?; - let spec = fixture_policy( - py, - &locals, - "variant_type", - "VariantType()", - "convert_variant", - )?; - let mut column = ResultColumn::new(&spec, values.len()); - add_values(py, &mut column, &values, mode); - let actual = column.finish(py); - assert!(matches!(&actual, Err(error) if error.is_instance_of::(py))); - let expected = reference( - py, - &locals, - &values, - &evaluate(py, &locals, "convert_variant")?, - &evaluate(py, &locals, "variant_type")?, - ); - assert_same_result(py, actual, expected)?; - } - Ok(()) - }) - .unwrap(); -} - -#[test] -fn dictionary_results_select_python_fallback_including_ordered_dictionaries() { - Python::initialize(); - Python::attach(|py| -> PyResult<()> { - let locals = context(py)?; - let none = py.None().into_bound(py); - for ordered in ["False", "True"] { - let arrow_type = evaluate( - py, - &locals, - &format!("pa.dictionary(pa.int8(), pa.string(), ordered={ordered})"), - )?; - let selected = result_spec(py, &arrow_type, &none, &none, true, &none, || { - panic!("dictionary fallback must not resolve VariantVal") - })?; - assert!(matches!(selected, Err(reason) if reason.contains("dictionary"))); - } - Ok(()) - }) - .unwrap(); -} - -#[test] -fn nested_dictionary_extension_results_select_python_fallback_before_import() { - Python::initialize(); - Python::attach(|py| -> PyResult<()> { - let locals = context(py)?; - py.run( - c"class BinaryExtension(pa.ExtensionType): - def __init__(self): - super().__init__(pa.binary(), 'spark.test.udf_result_binary') - def __arrow_ext_serialize__(self): - return b'' - @classmethod - def __arrow_ext_deserialize__(cls, storage_type, serialized): - return cls() - -dictionary_extension = pa.dictionary(pa.int8(), BinaryExtension()) -types = [ - pa.list_(dictionary_extension), - pa.struct([pa.field('nested', dictionary_extension)]), - pa.map_(pa.string(), dictionary_extension), - pa.struct([ - pa.field('value', dictionary_extension, nullable=False), - pa.field('metadata', pa.binary(), nullable=False, metadata={b'variant': b'true'}), - ]), -] -", - Some(&locals), - Some(&locals), - )?; - let none = py.None().into_bound(py); - for arrow_type in evaluate(py, &locals, "types")?.try_iter()? { - let arrow_type = arrow_type?; - let selected = result_spec(py, &arrow_type, &none, &none, true, &none, || { - panic!("unsupported nested types must not resolve VariantVal") - })?; - assert!(matches!(selected, Err(reason) if reason.contains("not supported"))); - } - Ok(()) - }) - .unwrap(); -} - -#[test] -#[ignore = "requires official pyspark-client==4.2.0 and PyArrow"] -fn official_pyspark_result_converters_match_native_and_fallback_paths() { - Python::initialize(); - Python::attach(|py| -> PyResult<()> { - let locals = context(py)?; - py.run( - c"from pyspark.sql import types as sql_types -from pyspark.sql.conversion import LocalDataToArrowConversion -from pyspark.sql.pandas.types import to_arrow_type -VariantVal = sql_types.VariantVal -class OfficialVariantSubclass(VariantVal): pass -cases = [ - (sql_types.StringType(), [None, 'plain', True, False, 1, -7, 1.25, Text('text')]), - (sql_types.StringType(), ['prefix', BadText()]), - (sql_types.StringType(), [2**80]), - (sql_types.StringType(), [2**16000]), - (sql_types.BinaryType(), [b'prefix', bytearray(b'value'), None, BytesSubclass(b'tail')]), - (sql_types.BinaryType(), [b'prefix', object()]), - (sql_types.VariantType(), [VariantVal(b'v', b'm'), None]), - (sql_types.VariantType(), [OfficialVariantSubclass(b'v', b'm'), None]), - (sql_types.VariantType(), [VariantVal(bytearray(b'v'), bytearray(b'm')), None]), - (sql_types.VariantType(), [VariantVal(object(), b'm'), object()]), -] -", - Some(&locals), - Some(&locals), - )?; - for case in evaluate(py, &locals, "cases")?.try_iter()? { - let case = case?; - let spark_type = case.get_item(0)?; - let values = case.get_item(1)?.cast_into::()?; - let converter = evaluate(py, &locals, "LocalDataToArrowConversion._create_converter")? - .call1((&spark_type,))?; - let arrow_type = evaluate(py, &locals, "to_arrow_type")?.call1((&spark_type,))?; - for mode in MODES { - for capacity in [1, usize::MAX] { - let spec = policy(py, &locals, &arrow_type, &spark_type, &converter)? - .with_value_capacity(capacity); - let mut column = ResultColumn::new(&spec, values.len()); - add_values(py, &mut column, &values, mode); - assert_same_result( - py, - column.finish(py), - reference(py, &locals, &values, &converter, &arrow_type), - )?; - } - } - } - Ok(()) - }) - .unwrap(); -} diff --git a/dev/design/architecture.md b/dev/design/architecture.md index cabd8ce4..e0a9f406 100644 --- a/dev/design/architecture.md +++ b/dev/design/architecture.md @@ -47,8 +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 shared Python/Arrow bridge, scalar conversion, and - explicit Python UDF conversion policy; no transport + spark-python-arrow Python/Arrow bridge and scalar conversion; no transport python/ pyspark/… thin drop-in package; public modules re-export from `_pyspark` so `from pyspark.sql import SparkSession` works @@ -76,10 +75,9 @@ NUL field names or timezones use an IPC compatibility fallback, without another RPC. `toPandas()` then calls PyArrow's pandas conversion, which may copy. The shared crate also owns the scalar routines used by the Python binding's -input conversion and row materialization, and an explicit Python UDF conversion -policy for worker consumers. It does not own the client `Row` representation, -schema inference, transport, or a worker's execution loop. `collect()` still uses -the native client's owned rows; `createDataFrame(pandas)` still materializes +input conversion and row materialization. It does not own the client `Row` +representation, 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 [shared converter contract](../../crates/spark-python-arrow/README.md). diff --git a/docs/architecture.md b/docs/architecture.md index f7fdb920..6bd77262 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/PyArrow C Data bridge and client scalar conversion | | 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 From 8e368d41d6322671773e474178259eaaad08d97d Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Tue, 27 Oct 2026 13:49:09 -0700 Subject: [PATCH 6/8] Remove unused scalar conversion helpers Inline client string and date materialization, remove stale worker and decoder wording, and exercise the remaining Scalar string API. Co-authored-by: Isaac --- crates/spark-python-arrow/README.md | 8 +++---- crates/spark-python-arrow/src/bridge.rs | 8 +++---- crates/spark-python-arrow/src/scalar.rs | 22 ++--------------- crates/spark-python-arrow/tests/conversion.rs | 24 +++++++------------ 4 files changed, 17 insertions(+), 45 deletions(-) diff --git a/crates/spark-python-arrow/README.md b/crates/spark-python-arrow/README.md index f4ade527..5ac65774 100644 --- a/crates/spark-python-arrow/README.md +++ b/crates/spark-python-arrow/README.md @@ -59,8 +59,7 @@ position has no metadata-owning Rust field. As with the Arrow C Data Interface itself, producers are responsible for valid buffer contents and all offsets. Imports perform limited structural checks, not full Arrow validation or a scan of every interior offset. A real PyArrow class -does not by itself prove complete buffer validity. CPython's strict UTF-8 decoder -checks the selected string bytes, not the Arrow offsets used to select them. +does not by itself prove complete buffer validity. Imports call native `pyarrow.lib` class descriptors to avoid dispatching an arbitrary object's exporter. Exports and Python fallback operations use the @@ -91,9 +90,8 @@ own dependency graph. The repository's root `LICENSE` covers this crate. The manifest keeps its package metadata self-contained so the compatibility tests can build it outside the workspace lockfile. The library does not -enable PyO3's `extension-module` feature or select a minimum Python ABI: a Python -extension chooses those features in the host package, while an embedded worker -links its chosen Python runtime. +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 permits PyO3 0.28, 0.29, and 0.30 so consumers on these lines can use the same source. The OSS Python extension selects 0.30.1, which requires Rust diff --git a/crates/spark-python-arrow/src/bridge.rs b/crates/spark-python-arrow/src/bridge.rs index 3f85d14c..486f1344 100644 --- a/crates/spark-python-arrow/src/bridge.rs +++ b/crates/spark-python-arrow/src/bridge.rs @@ -20,8 +20,8 @@ //! Imports accept real PyArrow objects, not arbitrary objects implementing an //! export method. PyArrow remains responsible for the validity of its C data. //! These functions perform limited structural checks, not full validation of all -//! offsets, dictionary indices, or payloads. Valid offsets and buffer bounds are -//! producer preconditions. String decoding errors remain observable at conversion. +//! offsets, dictionary indices, or payloads. Valid offsets, buffer bounds, and UTF-8 +//! are producer preconditions. use std::sync::Arc; @@ -446,8 +446,8 @@ fn export_python_array( fn validate_structure(data: &ArrayData) -> PyResult<()> { // Arrow's inexpensive validation recurses through child descriptors. It - // checks sizes and offset endpoints, not every offset or dictionary index; - // those remain producer preconditions. UTF-8 is decoded when a value is read. + // checks sizes and offset endpoints, not every offset or dictionary index. + // These and valid UTF-8 remain producer preconditions. data.validate().map_err(arrow_error) } diff --git a/crates/spark-python-arrow/src/scalar.rs b/crates/spark-python-arrow/src/scalar.rs index 6ad09326..94ee1106 100644 --- a/crates/spark-python-arrow/src/scalar.rs +++ b/crates/spark-python-arrow/src/scalar.rs @@ -50,37 +50,19 @@ impl Scalar<'_> { Self::Bool(value) => value.into_bound_py_any(py), Self::Int(value) => value.into_bound_py_any(py), Self::Float(value) => value.into_bound_py_any(py), - Self::String(value) => utf8_into_python(py, value.as_bytes()), + Self::String(value) => Ok(PyString::new(py, value).into_any()), Self::Binary(value) => Ok(PyBytes::new(py, value).into_any()), Self::Date(days) => { let fromordinal = py .import("datetime")? .getattr("date")? .getattr("fromordinal")?; - date_into_python(days, &fromordinal) + fromordinal.call1((i64::from(days) + EPOCH_ORDINAL,)) } } } } -/// Decodes UTF-8 with Python's strict decoder, preserving its exception on invalid bytes. -/// -/// This accepts raw bytes because an Arrow C Data import can contain unvalidated string data; -/// forming a Rust `&str` first would be unsound for such an input. -pub fn utf8_into_python<'py>(py: Python<'py>, bytes: &[u8]) -> PyResult> { - Ok(PyString::from_bytes(py, bytes)?.into_any()) -} - -/// Materializes a date using an explicitly supplied, optionally cached `date.fromordinal`. -/// -/// The constructor defines out-of-range errors. -pub fn date_into_python<'py>( - days: i32, - fromordinal: &Bound<'py, PyAny>, -) -> PyResult> { - fromordinal.call1((i64::from(days) + EPOCH_ORDINAL,)) -} - /// Extracts with PyO3's boolean rules, not Python truthiness. pub fn extract_bool(value: &Bound<'_, PyAny>) -> PyResult { value.extract() diff --git a/crates/spark-python-arrow/tests/conversion.rs b/crates/spark-python-arrow/tests/conversion.rs index 87b7532d..d3fe5912 100644 --- a/crates/spark-python-arrow/tests/conversion.rs +++ b/crates/spark-python-arrow/tests/conversion.rs @@ -25,18 +25,15 @@ use std::sync::Arc; use arrow::array::{Array, Int64Array}; use arrow::datatypes::{DataType, Field, Schema, TimeUnit}; use arrow::record_batch::{RecordBatch, RecordBatchOptions}; -use pyo3::exceptions::{ - PyOverflowError, PyTypeError, PyUnicodeDecodeError, PyUnicodeEncodeError, PyValueError, -}; +use pyo3::exceptions::{PyOverflowError, PyTypeError, PyUnicodeEncodeError, PyValueError}; use pyo3::prelude::*; -use pyo3::types::{PyBytes, PyDict, PyList}; +use pyo3::types::{PyBytes, PyDict, PyList, PyString}; use spark_python_arrow::bridge::{ export_array, export_data_type, export_schema, import_array, import_batch, import_data_type, import_schema, record_batches_to_pyarrow, }; use spark_python_arrow::scalar::{ - extract_binary, extract_bool, extract_f64, extract_i64, extract_string, utf8_into_python, - Scalar, + extract_binary, extract_bool, extract_f64, extract_i64, extract_string, Scalar, }; fn context(py: Python<'_>) -> PyResult> { @@ -580,18 +577,13 @@ class StringValue: } #[test] -fn utf8_scalar_decoding_preserves_embedded_nul_and_rejects_invalid_bytes() { +fn string_scalar_materialization_preserves_unicode_and_embedded_nul() { Python::initialize(); Python::attach(|py| -> PyResult<()> { - for value in ["", "a\0b", "\u{03bb}"] { - assert_eq!( - utf8_into_python(py, value.as_bytes())?.extract::()?, - value - ); - } - for bytes in [&[0xff][..], &[0xc3][..]] { - let error = utf8_into_python(py, bytes).unwrap_err(); - assert!(error.is_instance_of::(py)); + for value in ["", "ascii", "a\0b", "\u{00e9}", "\u{03bb}", "\u{1f600}"] { + let actual = Scalar::String(value).into_python(py)?; + assert!(actual.is_exact_instance_of::()); + assert_eq!(actual.extract::()?, value); } Ok(()) }) From ad7679c24eb438e02302a73b1a40981812310d3a Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Tue, 27 Oct 2026 13:59:49 -0700 Subject: [PATCH 7/8] Limit Python Arrow conversion to table export Remove unused inbound and standalone conversion APIs, keep scalar conversion in the client binding, declare the Arrow/PyO3 Rust minimum, and focus tests and documentation on the production export path. Co-authored-by: Isaac --- .github/workflows/coverage.yml | 7 +- .github/workflows/python-arrow.yml | 2 +- crates/pyspark-rs/src/row.rs | 32 +- crates/pyspark-rs/src/session.rs | 19 +- crates/spark-python-arrow/Cargo.toml | 3 +- crates/spark-python-arrow/README.md | 101 ++- crates/spark-python-arrow/src/bridge.rs | 537 +-------------- crates/spark-python-arrow/src/lib.rs | 7 +- crates/spark-python-arrow/src/scalar.rs | 95 --- .../tests/bridge_regressions.rs | 343 ++++------ crates/spark-python-arrow/tests/conversion.rs | 622 +++++------------- dev/design/architecture.md | 13 +- docs/architecture.md | 2 +- python/tests/test_scalar_conversion.py | 131 ++++ 14 files changed, 514 insertions(+), 1400 deletions(-) delete mode 100644 crates/spark-python-arrow/src/scalar.rs create mode 100644 python/tests/test_scalar_conversion.py diff --git a/.github/workflows/coverage.yml b/.github/workflows/coverage.yml index 88560f9c..41b9bb22 100644 --- a/.github/workflows/coverage.yml +++ b/.github/workflows/coverage.yml @@ -92,9 +92,6 @@ jobs: python -m pip install maturin pytest coverage python -m pip install pandas pyarrow numpy grpcio grpcio-status googleapis-common-protos cloudpickle zstandard py4j - - name: Test standalone Python and Arrow converter - run: cargo test -p spark-python-arrow --locked --features pyo3/abi3-py39 - - name: Cache Spark ${{ env.SPARK_VERSION }} distribution uses: actions/cache@v4 with: @@ -229,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 tests/test_arrow_handoff.py + cd python + PYTHONPATH=$(pwd) python -m coverage run --rcfile=.coveragerc -m pytest \ + tests/test_dropin_offline.py tests/test_arrow_handoff.py tests/test_scalar_conversion.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 index a8119f6d..e1efda1d 100644 --- a/.github/workflows/python-arrow.yml +++ b/.github/workflows/python-arrow.yml @@ -61,7 +61,7 @@ jobs: cargo update --manifest-path "$RUNNER_TEMP/spark-python-arrow/Cargo.toml" \ -p pyo3 --precise "${{ matrix.pyo3 }}" cargo tree --manifest-path "$RUNNER_TEMP/spark-python-arrow/Cargo.toml" --locked -i pyo3 - - name: Test Arrow handoff and scalar conversion + - name: Test Arrow table export run: | cargo test --manifest-path "$RUNNER_TEMP/spark-python-arrow/Cargo.toml" \ --locked --features pyo3/abi3-py39 diff --git a/crates/pyspark-rs/src/row.rs b/crates/pyspark-rs/src/row.rs index 406216af..fe3b1fdf 100644 --- a/crates/pyspark-rs/src/row.rs +++ b/crates/pyspark-rs/src/row.rs @@ -2,9 +2,9 @@ use pyo3::exceptions::{PyAttributeError, PyIndexError, PyKeyError}; use pyo3::prelude::*; -use pyo3::types::{PyDict, PyList}; +use pyo3::types::{PyBytes, PyDict, PyList, PyString}; +use pyo3::IntoPyObjectExt; use spark_connect::row::{Row, Value}; -use spark_python_arrow::scalar::Scalar; /// Python wrapper for a Spark Row. #[pyclass(name = "Row", module = "pyspark.sql.types")] @@ -24,17 +24,23 @@ pub fn value_to_py<'py>(py: Python<'py>, v: &Value) -> PyResult Scalar::Null.into_python(py)?, - Value::Bool(b) => Scalar::Bool(*b).into_python(py)?, - Value::Byte(n) => Scalar::Int(i64::from(*n)).into_python(py)?, - Value::Short(n) => Scalar::Int(i64::from(*n)).into_python(py)?, - Value::Integer(n) => Scalar::Int(i64::from(*n)).into_python(py)?, - Value::Long(n) => Scalar::Int(*n).into_python(py)?, - Value::Float(f) => Scalar::Float(f64::from(*f)).into_python(py)?, - Value::Double(f) => Scalar::Float(*f).into_python(py)?, - Value::String(s) => Scalar::String(s).into_python(py)?, - Value::Binary(b) => Scalar::Binary(b).into_python(py)?, - Value::Date(d) => Scalar::Date(*d).into_python(py)?, + Value::Null => py.None().into_bound(py), + Value::Bool(b) => b.into_bound_py_any(py)?, + Value::Byte(n) => n.into_bound_py_any(py)?, + Value::Short(n) => n.into_bound_py_any(py)?, + Value::Integer(n) => n.into_bound_py_any(py)?, + Value::Long(n) => n.into_bound_py_any(py)?, + Value::Float(f) => f.into_bound_py_any(py)?, + Value::Double(f) => f.into_bound_py_any(py)?, + Value::String(s) => PyString::new(py, s).into_any(), + Value::Binary(b) => PyBytes::new(py, b).into_any(), + Value::Date(d) => { + // 719163 is the ordinal of 1970-01-01. Use i64 to avoid overflow + // before Python reports dates outside its supported range. + py.import("datetime")? + .getattr("date")? + .call_method1("fromordinal", (i64::from(*d) + 719_163,))? + } Value::Timestamp(t) => { // Convert microseconds-since-epoch (i64) to datetime.datetime let datetime_mod = py.import("datetime")?; diff --git a/crates/pyspark-rs/src/session.rs b/crates/pyspark-rs/src/session.rs index 1a3d049d..7140393b 100644 --- a/crates/pyspark-rs/src/session.rs +++ b/crates/pyspark-rs/src/session.rs @@ -1,7 +1,7 @@ //! PyO3 wrapper for Spark Connect sessions. use pyo3::prelude::*; -use pyo3::types::{PyDict, PyList}; +use pyo3::types::{PyByteArray, PyBytes, PyDict, PyList}; use spark_connect::row::{Row, Value}; use spark_connect::session::SparkSession; use spark_connect::types::{DataType, StructField}; @@ -945,15 +945,13 @@ fn value_to_datatype(v: &Value) -> DataType { /// Convert a Python value to a Rust Value. pub(crate) fn py_to_value(obj: &Bound<'_, PyAny>) -> PyResult { - use spark_python_arrow::scalar; - // Check for None if obj.is_none() { return Ok(Value::Null); } // Check for bool (before int because bool is a subclass of int in Python) - if let Ok(b) = scalar::extract_bool(obj) { + if let Ok(b) = obj.extract::() { return Ok(Value::Bool(b)); } @@ -964,24 +962,27 @@ pub(crate) fn py_to_value(obj: &Bound<'_, PyAny>) -> PyResult { } // Check for int - if let Ok(i) = scalar::extract_i64(obj) { + if let Ok(i) = obj.extract::() { return Ok(Value::Long(i)); } // Check for float - if let Ok(f) = scalar::extract_f64(obj) { + if let Ok(f) = obj.extract::() { return Ok(Value::Double(f)); } // Check for str - if let Ok(s) = scalar::extract_string(obj) { + if let Ok(s) = obj.extract::() { return Ok(Value::String(s)); } // Check for bytes/bytearray SPECIFICALLY (not `extract::>()`, which also // matches a Python list of small ints like [1,2,3] and would mis-tag arrays as binary). - if let Some(bytes) = scalar::extract_binary(obj) { - return Ok(Value::Binary(bytes)); + if let Ok(bytes) = obj.cast::() { + return Ok(Value::Binary(bytes.as_bytes().to_vec())); + } + if let Ok(bytes) = obj.cast::() { + return Ok(Value::Binary(bytes.to_vec())); } // list / tuple -> array (recursively converted) diff --git a/crates/spark-python-arrow/Cargo.toml b/crates/spark-python-arrow/Cargo.toml index fdcbd9a3..18b401e9 100644 --- a/crates/spark-python-arrow/Cargo.toml +++ b/crates/spark-python-arrow/Cargo.toml @@ -18,9 +18,10 @@ name = "spark-python-arrow" version = "0.1.0" edition = "2021" +rust-version = "1.85" license = "Apache-2.0" publish = false -description = "Arrow and Python conversion for Spark clients" +description = "Arrow record batch export to PyArrow for Spark clients" repository = "https://github.com/apache/spark-connect-rust" [dependencies] diff --git a/crates/spark-python-arrow/README.md b/crates/spark-python-arrow/README.md index 5ac65774..66b24d9f 100644 --- a/crates/spark-python-arrow/README.md +++ b/crates/spark-python-arrow/README.md @@ -17,27 +17,34 @@ specific language governing permissions and limitations under the License. --> -# Shared Python and Arrow conversion +# Arrow table export to PyArrow -`spark-python-arrow` contains the Arrow bridge and scalar conversion routines -used by `pyspark-rs`. It depends on Arrow and PyO3, not Spark Connect transport, -generated protobufs, or an async runtime. Worker-side UDF conversion is outside -this crate's scope. +`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. -## Interfaces +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. -- `bridge`: import and export Arrow schemas, types, arrays, and record batches - through PyArrow's C Data Interface. Table export preserves batch boundaries and - shares decoded Arrow data buffers. Field names or timezone strings containing - NUL use a documented IPC fallback at the table boundary; low-level C exports - reject them. Length-prefixed schema metadata can contain NUL. -- `scalar`: borrowed scalar materialization for collected rows and Python - extraction for client inputs. Extraction retains PyO3's coercion behavior; - the client chooses the order of type checks and handles nested values. +## Export contract -`pyspark-rs` uses `bridge::record_batches_to_pyarrow` in `DataFrame.toArrow()`, -`scalar::Scalar` in row materialization, and `scalar::extract_*` in session input -conversion. These are the crate's current production consumers. +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 @@ -45,39 +52,10 @@ 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. -## Ownership and caller responsibilities - -The library does not initialize Python, acquire a process-global interpreter, -create threads, or retain session state. The consumer chooses embedding versus -extension-module linking and owns all Python callbacks and their interpreter. -Array imports accept actual PyArrow objects and use their native exporters. -Bare extension types/arrays and bare ordered dictionaries are rejected because -Arrow's bare type cannot retain their field metadata or ordering flag. Schemas -and record batches retain these properties on their fields. Extension types in -dictionary-value positions are rejected, even within schemas, because that -position has no metadata-owning Rust field. -As with the Arrow C Data Interface itself, producers are responsible for valid -buffer contents and all offsets. Imports perform limited structural checks, not -full Arrow validation or a scan of every interior offset. A real PyArrow class -does not by itself prove complete buffer validity. - -Imports call native `pyarrow.lib` class descriptors to avoid dispatching an -arbitrary object's exporter. Exports and Python fallback operations use the -public `pyarrow` aliases, allowing normal Python instrumentation, including the -ownership/failure-injection tests. This is not isolation from arbitrary Python -code running in the same interpreter. - Zero-copy describes compatible Arrow **value and offset buffers**, not schema descriptors, Python scalar/row objects, pandas materialization, IPC decoding, or -fallback conversion. Sliced arrays whose value and null offsets differ can -require copying the validity bitmap during C Data export. Import/export release -callbacks retain shared buffers after either language's temporary wrapper has -been dropped. - -The client's owned `Value` and `Row` types remain in the Python-free Rust client. -The Python binding adapts their scalar fields to `scalar::Scalar`; its row class, -schema inference, and temporal/nested compatibility rules remain caller-owned. -Cancellation, transport, and session settings likewise stay outside this crate. +the IPC fallback. Sliced arrays whose value and null offsets differ can require +copying the validity bitmap during C Data export. ## Integration @@ -92,13 +70,17 @@ The manifest keeps its package metadata self-contained so the compatibility tests can build it outside the workspace lockfile. 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 permits PyO3 0.28, 0.29, and 0.30 so consumers on these lines can use -the same source. The OSS Python extension selects 0.30.1, which requires Rust -1.85 or newer; `pyspark-rs` declares that minimum independently of the workspace's -existing 1.82 default. The isolated converter compatibility CI job tests each minor -version independently of the workspace lockfile. A version is not considered -verified until its corresponding compatibility job passes. +The manifest permits PyO3 0.28, 0.29, and 0.30; the Python extension selects +0.30.1. Both this crate and `pyspark-rs` declare Rust 1.85, matching the declared +minimum of Arrow 58.3.0 and PyO3 0.30.1. 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 isolated compatibility CI job tests each PyO3 minor independently of the +workspace lockfile. A version is not considered verified until its corresponding +job passes. These tests run in the compatibility workflow, not again in coverage. ## Tests @@ -115,8 +97,9 @@ 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 cover buffer identity and lifetime, schema/type round trips, -empty shapes, schema mismatch errors, and scalar extraction/materialization. The -Python client's `test_arrow_handoff.py` exercises the public `toArrow()` and -`toPandas()` paths. These are not substitutes for Spark's complete Connect -parity suite. +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 index 486f1344..e03fd79f 100644 --- a/crates/spark-python-arrow/src/bridge.rs +++ b/crates/spark-python-arrow/src/bridge.rs @@ -15,125 +15,25 @@ // specific language governing permissions and limitations // under the License. -//! Share Arrow buffers with PyArrow through the Arrow C Data Interface. +//! Export Arrow record batches to PyArrow through the Arrow C Data Interface. //! -//! Imports accept real PyArrow objects, not arbitrary objects implementing an -//! export method. PyArrow remains responsible for the validity of its C data. -//! These functions perform limited structural checks, not full validation of all -//! offsets, dictionary indices, or payloads. Valid offsets, buffer bounds, and UTF-8 -//! are producer preconditions. +//! 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::{make_array, Array, ArrayData, ArrayRef, StructArray}; -use arrow::datatypes::{DataType, Field, Schema, UnionFields}; +use arrow::array::{Array, StructArray}; +use arrow::datatypes::{DataType, Field, Schema}; use arrow::error::ArrowError; -use arrow::ffi::{from_ffi_and_data_type, FFI_ArrowArray, FFI_ArrowSchema}; +use arrow::ffi::{FFI_ArrowArray, FFI_ArrowSchema}; use arrow::ipc::writer::FileWriter; -use arrow::record_batch::{RecordBatch, RecordBatchOptions}; -use pyo3::exceptions::{PyRuntimeError, PyTypeError, PyValueError}; +use arrow::record_batch::RecordBatch; +use pyo3::exceptions::{PyRuntimeError, PyValueError}; use pyo3::prelude::*; use pyo3::types::{PyBytes, PyList}; -/// Import a PyArrow schema, preserving schema and field metadata. -pub fn import_schema(value: &Bound<'_, PyAny>) -> PyResult { - let class = require_pyarrow_class(value, "Schema")?; - check_python_schema(value)?; - let mut schema = FFI_ArrowSchema::empty(); - let address = &mut schema as *mut FFI_ArrowSchema as usize; - class.call_method1("_export_to_c", (value, address))?; - import_c_schema(&schema).map_err(arrow_error) -} - -/// Import a PyArrow data type. -/// -/// Root extension types and ordered dictionaries require a schema or record batch: -/// a bare Rust `DataType` has nowhere to retain these field properties. -pub fn import_data_type(value: &Bound<'_, PyAny>) -> PyResult { - let class = require_pyarrow_class(value, "DataType")?; - reject_root_extension_type(value)?; - reject_root_dictionary_order(value)?; - check_python_data_type(value)?; - let mut schema = FFI_ArrowSchema::empty(); - let address = &mut schema as *mut FFI_ArrowSchema as usize; - class.call_method1("_export_to_c", (value, address))?; - import_c_data_type(&schema).map_err(arrow_error) -} - -/// Import a PyArrow array without copying its data buffers. -/// -/// Root extensions and ordered dictionaries are rejected because an `ArrayRef` -/// cannot retain their field properties. Import a record batch to preserve them. -pub fn import_array(value: &Bound<'_, PyAny>) -> PyResult { - let class = require_pyarrow_class(value, "Array")?; - let data_type = value.getattr("type")?; - reject_root_extension_type(&data_type)?; - reject_root_dictionary_order(&data_type)?; - check_python_data_type(&data_type)?; - let (array, schema) = export_python_array(value, &class)?; - let data_type = import_c_data_type(&schema).map_err(arrow_error)?; - // SAFETY: the base PyArrow exporter, not an arbitrary Python method, has - // initialized both descriptors and transferred ownership of valid C data. - let data = unsafe { from_ffi_and_data_type(array, data_type) }.map_err(arrow_error)?; - validate_structure(&data)?; - Ok(make_array(data)) -} - -/// Import a PyArrow record batch, preserving metadata and zero-column rows. -pub fn import_batch(value: &Bound<'_, PyAny>) -> PyResult { - let class = require_pyarrow_class(value, "RecordBatch")?; - check_python_schema(&value.getattr("schema")?)?; - let (array, ffi_schema) = export_python_array(value, &class)?; - let schema = import_c_schema(&ffi_schema).map_err(arrow_error)?; - let data_type = DataType::Struct(schema.fields().clone()); - // SAFETY: only the real PyArrow RecordBatch exporter writes the descriptors. - let data = unsafe { from_ffi_and_data_type(array, data_type) }.map_err(arrow_error)?; - validate_structure(&data)?; - if !matches!(data.data_type(), DataType::Struct(_)) || data.null_count() != 0 { - return Err(PyValueError::new_err( - "A record batch must be a struct array without top-level nulls", - )); - } - let array = make_array(data); - let array = array - .as_any() - .downcast_ref::() - .ok_or_else(|| PyValueError::new_err("A record batch must be a struct array"))?; - let options = RecordBatchOptions::new().with_row_count(Some(array.len())); - // Unlike RecordBatch::from(StructArray), this validates non-nullable columns - // and retains the schema-level metadata from the exported C schema. - RecordBatch::try_new_with_options(Arc::new(schema), array.columns().to_vec(), &options) - .map_err(arrow_error) -} - -/// Export an Arrow array sharing its value and offset buffers. -/// -/// Arrow may copy a sliced validity bitmap to align its offset with the data. -pub fn export_array(py: Python<'_>, array: ArrayRef) -> PyResult> { - let array_type = py.import("pyarrow")?.getattr("Array")?; - Ok(export_array_with_type(&array_type, array)?.unbind()) -} - -pub(crate) fn export_array_with_type<'py>( - array_type: &Bound<'py, PyAny>, - array: ArrayRef, -) -> PyResult> { - check_c_data_type(array.data_type())?; - let mut schema = export_c_data_type(array.data_type()).map_err(arrow_error)?; - let schema_address = &mut schema as *mut FFI_ArrowSchema as usize; - let mut array = FFI_ArrowArray::new(&array.to_data()); - let array_address = &mut array as *mut FFI_ArrowArray as usize; - // Both descriptors are owned here until the importer consumes them. Their - // release callbacks also cover importer errors before either transfer. - array_type.call_method1("_import_from_c", (array_address, schema_address)) -} - -/// Export an Arrow schema, preserving schema and field metadata. -/// -/// C schema names and timezone strings cannot contain NUL. This low-level API -/// returns an error for such schemas; table export provides an IPC fallback. -pub fn export_schema(py: Python<'_>, schema: &Schema) -> PyResult> { - check_c_schema(schema)?; +// 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())) @@ -146,18 +46,6 @@ pub fn export_schema(py: Python<'_>, schema: &Schema) -> PyResult> { .unbind()) } -/// Export an Arrow data type, rejecting strings not representable in C schemas. -pub fn export_data_type(py: Python<'_>, data_type: &DataType) -> PyResult> { - check_c_data_type(data_type)?; - let mut schema = export_c_data_type(data_type).map_err(arrow_error)?; - let address = &mut schema as *mut FFI_ArrowSchema as usize; - Ok(py - .import("pyarrow")? - .getattr("DataType")? - .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. @@ -218,38 +106,10 @@ pub fn record_batches_to_pyarrow(py: Python<'_>, batches: Vec) -> P .unbind()) } -pub(crate) fn arrow_error(error: ArrowError) -> PyErr { +fn arrow_error(error: ArrowError) -> PyErr { PyRuntimeError::new_err(error.to_string()) } -pub(crate) fn is_extension_type(data_type: &Bound<'_, PyAny>) -> PyResult { - let class = data_type - .py() - .import("pyarrow.lib")? - .getattr("BaseExtensionType")?; - data_type.is_instance(&class) -} - -fn reject_root_extension_type(data_type: &Bound<'_, PyAny>) -> PyResult<()> { - if is_extension_type(data_type)? { - return Err(PyTypeError::new_err( - "Arrow extension types require a schema or record batch to preserve metadata", - )); - } - Ok(()) -} - -fn reject_root_dictionary_order(data_type: &Bound<'_, PyAny>) -> PyResult<()> { - if let Ok(ordered) = data_type.getattr("ordered") { - if ordered.is_truthy()? { - return Err(PyTypeError::new_err( - "Ordered Arrow dictionaries require a schema or record batch to preserve ordering", - )); - } - } - Ok(()) -} - // 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. @@ -339,136 +199,6 @@ fn schemas_equal(left: &Schema, right: &Schema) -> bool { ) } -fn import_c_schema(schema: &FFI_ArrowSchema) -> Result { - let DataType::Struct(fields) = import_c_data_type(schema)? else { - return Err(ArrowError::CDataInterface( - "A schema must be a struct type".into(), - )); - }; - Ok(Schema::new(fields).with_metadata(schema.metadata()?)) -} - -fn import_c_data_type(schema: &FFI_ArrowSchema) -> Result { - fn has_ordered_dictionary(schema: &FFI_ArrowSchema) -> bool { - schema.dictionary_ordered() - || schema.children().any(has_ordered_dictionary) - || schema.dictionary().is_some_and(has_ordered_dictionary) - } - - // Arrow 58 drops DICTIONARY_ORDERED when importing Fields. Restore it at - // every field, including those nested in array types and dictionary values. - fn restore(data_type: DataType, schema: &FFI_ArrowSchema) -> Result { - let field = |field: &Field, index: usize| -> Result, ArrowError> { - let child = schema.child(index); - Ok(Arc::new( - field - .clone() - .with_data_type(restore(field.data_type().clone(), child)?) - .with_dict_is_ordered(child.dictionary_ordered()), - )) - }; - Ok(match data_type { - DataType::List(child) => DataType::List(field(&child, 0)?), - DataType::LargeList(child) => DataType::LargeList(field(&child, 0)?), - DataType::ListView(child) => DataType::ListView(field(&child, 0)?), - DataType::LargeListView(child) => DataType::LargeListView(field(&child, 0)?), - DataType::FixedSizeList(child, size) => { - DataType::FixedSizeList(field(&child, 0)?, size) - } - DataType::Map(child, sorted) => DataType::Map(field(&child, 0)?, sorted), - DataType::Struct(fields) => DataType::Struct( - fields - .iter() - .enumerate() - .map(|(index, child)| field(child, index)) - .collect::, _>>()? - .into(), - ), - DataType::Union(fields, mode) => DataType::Union( - UnionFields::try_new( - fields.iter().map(|(id, _)| id), - fields - .iter() - .enumerate() - .map(|(index, (_, child))| field(child, index)) - .collect::, _>>()?, - )?, - mode, - ), - DataType::RunEndEncoded(run_ends, values) => { - DataType::RunEndEncoded(field(&run_ends, 0)?, field(&values, 1)?) - } - DataType::Dictionary(keys, values) => DataType::Dictionary( - keys, - Box::new(restore( - *values, - schema.dictionary().ok_or_else(|| { - ArrowError::CDataInterface("A dictionary must have a value type".into()) - })?, - )?), - ), - other => other, - }) - } - - let data_type = DataType::try_from(schema)?; - if has_ordered_dictionary(schema) { - restore(data_type, schema) - } else { - Ok(data_type) - } -} - -fn require_pyarrow_class<'py>( - value: &Bound<'py, PyAny>, - name: &str, -) -> PyResult> { - // Use the extension's classes, not replaceable pyarrow package aliases. - // Calling their descriptors also bypasses a subclass's Python override. - let class = value.py().import("pyarrow.lib")?.getattr(name)?; - if !value.is_instance(&class)? { - return Err(PyTypeError::new_err(format!("Expected a pyarrow.{name}"))); - } - Ok(class) -} - -fn export_python_array( - value: &Bound<'_, PyAny>, - class: &Bound<'_, PyAny>, -) -> PyResult<(FFI_ArrowArray, FFI_ArrowSchema)> { - let mut array = FFI_ArrowArray::empty(); - let mut schema = FFI_ArrowSchema::empty(); - let array_address = &mut array as *mut FFI_ArrowArray as usize; - let schema_address = &mut schema as *mut FFI_ArrowSchema as usize; - class.call_method1("_export_to_c", (value, array_address, schema_address))?; - Ok((array, schema)) -} - -fn validate_structure(data: &ArrayData) -> PyResult<()> { - // Arrow's inexpensive validation recurses through child descriptors. It - // checks sizes and offset endpoints, not every offset or dictionary index. - // These and valid UTF-8 remain producer preconditions. - data.validate().map_err(arrow_error) -} - -fn nul_schema_error() -> PyErr { - PyValueError::new_err("Arrow C schema names and timezones cannot contain NUL") -} - -fn check_c_schema(schema: &Schema) -> PyResult<()> { - if schema.fields().iter().any(|field| field_needs_ipc(field)) { - return Err(nul_schema_error()); - } - Ok(()) -} - -fn check_c_data_type(data_type: &DataType) -> PyResult<()> { - if data_type_needs_ipc(data_type) { - return Err(nul_schema_error()); - } - Ok(()) -} - fn field_needs_ipc(field: &Field) -> bool { field.name().contains('\0') || data_type_needs_ipc(field.data_type()) } @@ -494,255 +224,14 @@ fn data_type_needs_ipc(data_type: &DataType) -> bool { } } -fn check_python_schema(schema: &Bound<'_, PyAny>) -> PyResult<()> { - for index in 0..schema.len()? { - check_python_field(&schema.call_method1("field", (index,))?)?; - } - Ok(()) -} - -fn check_python_field(field: &Bound<'_, PyAny>) -> PyResult<()> { - if field.getattr("name")?.extract::()?.contains('\0') { - return Err(nul_schema_error()); - } - check_python_data_type(&field.getattr("type")?) -} - -fn check_python_data_type(data_type: &Bound<'_, PyAny>) -> PyResult<()> { - // Check the original Python strings before C export can truncate at NUL. - // Metadata uses length-prefixed strings and does not need this restriction. - if let Ok(timezone) = data_type.getattr("tz") { - if !timezone.is_none() && timezone.extract::()?.contains('\0') { - return Err(nul_schema_error()); - } - } - let num_fields = data_type.getattr("num_fields")?.extract::()?; - for index in 0..num_fields { - check_python_field(&data_type.call_method1("field", (index,))?)?; - } - // Dictionary and extension value/storage types are not child fields. - // Lists expose both a child field and value_type; do not visit those twice. - if num_fields == 0 { - for name in ["value_type", "storage_type"] { - if let Ok(child) = data_type.getattr(name) { - if name == "value_type" && is_extension_type(&child)? { - return Err(PyTypeError::new_err( - "Arrow dictionary value extension types cannot preserve their metadata", - )); - } - if name == "value_type" { - reject_root_dictionary_order(&child)?; - } - check_python_data_type(&child)?; - } - } - } - Ok(()) -} - #[cfg(test)] mod tests { - use std::collections::HashMap; - use std::ffi::CStr; - - use arrow::array::Int64Array; use arrow::datatypes::TimeUnit; - use pyo3::types::PyDict; use super::*; - fn eval<'py>(py: Python<'py>, code: &CStr) -> Bound<'py, PyAny> { - let globals = PyDict::new(py); - globals - .set_item("pa", py.import("pyarrow").unwrap()) - .unwrap(); - py.eval(code, Some(&globals), None).unwrap() - } - - #[test] - fn schema_roundtrip_preserves_metadata() { - Python::initialize(); - Python::attach(|py| { - let field = Field::new("value", DataType::Int64, true) - .with_metadata(HashMap::from([("field\0key".into(), "value\0".into())])); - let schema = Schema::new(vec![field]) - .with_metadata(HashMap::from([("schema".into(), "metadata".into())])); - let exported = export_schema(py, &schema).unwrap(); - assert_eq!(import_schema(exported.bind(py)).unwrap(), schema); - }); - } - - #[test] - fn array_roundtrip_shares_buffers_and_preserves_lifetime() { - Python::initialize(); - Python::attach(|py| { - let values: ArrayRef = Arc::new(Int64Array::from(vec![Some(1), None, Some(3)])); - let address = values.to_data().buffers()[0].as_ptr(); - let exported = export_array(py, values).unwrap(); - let imported = import_array(exported.bind(py)).unwrap(); - drop(exported); - assert_eq!(imported.to_data().buffers()[0].as_ptr(), address); - assert_eq!(imported.len(), 3); - assert!(imported.is_null(1)); - assert_eq!( - imported - .as_any() - .downcast_ref::() - .unwrap() - .value(2), - 3 - ); - }); - } - - #[test] - fn imports_preserve_slices_and_batch_metadata() { - Python::initialize(); - Python::attach(|py| { - let batch = eval( - py, - c"pa.RecordBatch.from_arrays( - [pa.array([1, None, 3])], - schema=pa.schema( - [pa.field('v', pa.int64(), metadata={b'f': b'v'})], - metadata={b's': b'v'} - ) - ).slice(1, 2)", - ); - let imported = import_batch(&batch).unwrap(); - assert_eq!(imported.num_rows(), 2); - assert!(imported.column(0).is_null(0)); - assert_eq!(imported.schema().metadata().get("s").unwrap(), "v"); - assert_eq!(imported.schema().field(0).metadata().get("f").unwrap(), "v"); - let table = record_batches_to_pyarrow(py, vec![imported]).unwrap(); - assert!(table - .bind(py) - .getattr("schema") - .unwrap() - .call_method1("equals", (batch.getattr("schema").unwrap(), true)) - .unwrap() - .is_truthy() - .unwrap()); - }); - } - - #[test] - fn imports_preserve_zero_column_rows_and_empty_schema_metadata() { - Python::initialize(); - Python::attach(|py| { - for length in [0, 3] { - let array = StructArray::new_empty_fields(length, None); - let batch = RecordBatch::from(array); - let schema = batch - .schema() - .as_ref() - .clone() - .with_metadata(HashMap::from([("schema".into(), "empty".into())])); - let batch = batch.with_schema(Arc::new(schema)).unwrap(); - let table = record_batches_to_pyarrow(py, vec![batch.clone()]).unwrap(); - let batches = table.bind(py).call_method0("to_batches").unwrap(); - if length == 0 { - assert_eq!( - table - .bind(py) - .getattr("num_rows") - .unwrap() - .extract::() - .unwrap(), - 0 - ); - assert_eq!( - import_schema(&table.bind(py).getattr("schema").unwrap()).unwrap(), - *batch.schema() - ); - } else { - let imported = import_batch(&batches.get_item(0).unwrap()).unwrap(); - assert_eq!(imported.num_rows(), length); - assert_eq!(imported.num_columns(), 0); - assert_eq!(imported.schema(), batch.schema()); - } - } - }); - } - - #[test] - fn imports_reject_wrong_pyarrow_classes() { - Python::initialize(); - Python::attach(|py| { - let array = eval(py, c"pa.array([1, 2])"); - assert!(import_schema(&array) - .unwrap_err() - .is_instance_of::(py)); - assert!(import_data_type(&array) - .unwrap_err() - .is_instance_of::(py)); - assert!(import_batch(&array) - .unwrap_err() - .is_instance_of::(py)); - let schema = eval(py, c"pa.schema([('v', pa.int64())])"); - assert!(import_array(&schema) - .unwrap_err() - .is_instance_of::(py)); - let fake = eval( - py, - c"type('Fake', (), {'_export_to_c': lambda *args: 1 / 0})()", - ); - assert!(import_schema(&fake) - .unwrap_err() - .is_instance_of::(py)); - assert!(import_data_type(&fake) - .unwrap_err() - .is_instance_of::(py)); - assert!(import_array(&fake) - .unwrap_err() - .is_instance_of::(py)); - assert!(import_batch(&fake) - .unwrap_err() - .is_instance_of::(py)); - }); - } - - #[test] - fn batch_import_rejects_nulls_in_non_nullable_columns() { - Python::initialize(); - Python::attach(|py| { - let batch = eval( - py, - c"pa.RecordBatch.from_arrays( - [pa.array([None], type=pa.int64())], - schema=pa.schema([pa.field('v', pa.int64(), nullable=False)]) - )", - ); - assert!(import_batch(&batch) - .unwrap_err() - .to_string() - .contains("non-nullable")); - }); - } - #[test] - fn nul_schema_strings_are_rejected_without_panicking() { - let field = Field::new("bad\0name", DataType::Int64, true); - assert!(check_c_schema(&Schema::new(vec![field.clone()])).is_err()); - assert!(check_c_data_type(&DataType::List(Arc::new(field))).is_err()); - let timestamp = DataType::Timestamp(TimeUnit::Microsecond, Some("UTC\0suffix".into())); - assert!(check_c_data_type(×tamp).is_err()); - assert!(check_c_data_type(&DataType::Dictionary( - Box::new(DataType::Int8), - Box::new(timestamp) - )) - .is_err()); - Python::initialize(); - Python::attach(|py| { - assert!(export_data_type(py, &DataType::Time32(TimeUnit::Nanosecond)).is_err()); - let schema = eval(py, c"pa.schema([('bad\\x00name', pa.int64())])"); - assert!(import_schema(&schema) - .unwrap_err() - .is_instance_of::(py)); - let data_type = eval(py, c"pa.timestamp('us', tz='UTC\\x00suffix')"); - assert!(import_data_type(&data_type) - .unwrap_err() - .is_instance_of::(py)); - }); + 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 index 825bd91d..60c5d084 100644 --- a/crates/spark-python-arrow/src/lib.rs +++ b/crates/spark-python-arrow/src/lib.rs @@ -15,10 +15,9 @@ // specific language governing permissions and limitations // under the License. -//! Python/Arrow conversion without a dependency on Spark transport. +//! Arrow record batch export to PyArrow without a dependency on Spark transport. //! -//! The bridge shares compatible Arrow value buffers; scalar and row materialization still -//! allocates Python objects. The caller owns its interpreter and row representation. +//! 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; -pub mod scalar; diff --git a/crates/spark-python-arrow/src/scalar.rs b/crates/spark-python-arrow/src/scalar.rs deleted file mode 100644 index 94ee1106..00000000 --- a/crates/spark-python-arrow/src/scalar.rs +++ /dev/null @@ -1,95 +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. - -//! Scalar conversion for client input extraction and row materialization. -//! -//! The `extract_*` helpers retain PyO3's conversion semantics; callers choose their order, notably -//! checking booleans before integers and Decimal before numeric coercion. This module does not -//! infer a schema or define nested, decimal, timestamp, or application-specific row types. - -use pyo3::prelude::*; -use pyo3::types::{PyByteArray, PyBytes, PyString}; -use pyo3::IntoPyObjectExt; - -/// The ordinal of 1970-01-01 in Python's `datetime.date` calendar. -const EPOCH_ORDINAL: i64 = 719_163; - -/// A borrowed scalar, independent of a client's owned value representation. -#[derive(Clone, Copy, Debug, PartialEq)] -pub enum Scalar<'a> { - Null, - Bool(bool), - Int(i64), - Float(f64), - String(&'a str), - /// Materialized as immutable Python `bytes`. - Binary(&'a [u8]), - /// Days since 1970-01-01, materialized as `datetime.date`. - Date(i32), -} - -impl Scalar<'_> { - /// Materializes this scalar in the caller's attached interpreter. - pub fn into_python<'py>(self, py: Python<'py>) -> PyResult> { - match self { - Self::Null => Ok(py.None().into_bound(py)), - Self::Bool(value) => value.into_bound_py_any(py), - Self::Int(value) => value.into_bound_py_any(py), - Self::Float(value) => value.into_bound_py_any(py), - Self::String(value) => Ok(PyString::new(py, value).into_any()), - Self::Binary(value) => Ok(PyBytes::new(py, value).into_any()), - Self::Date(days) => { - let fromordinal = py - .import("datetime")? - .getattr("date")? - .getattr("fromordinal")?; - fromordinal.call1((i64::from(days) + EPOCH_ORDINAL,)) - } - } - } -} - -/// Extracts with PyO3's boolean rules, not Python truthiness. -pub fn extract_bool(value: &Bound<'_, PyAny>) -> PyResult { - value.extract() -} - -/// Extracts with PyO3's integer rules, including accepted integer-like objects. -pub fn extract_i64(value: &Bound<'_, PyAny>) -> PyResult { - value.extract() -} - -/// Extracts with PyO3's float rules, including numeric coercion. -pub fn extract_f64(value: &Bound<'_, PyAny>) -> PyResult { - value.extract() -} - -/// Extracts with PyO3's string rules, including string subclasses but not arbitrary `str()`. -pub fn extract_string(value: &Bound<'_, PyAny>) -> PyResult { - value.extract() -} - -/// Copies bytes or bytearray contents, including subclasses, but never a list of integers. -pub fn extract_binary(value: &Bound<'_, PyAny>) -> Option> { - if let Ok(bytes) = value.cast::() { - Some(bytes.as_bytes().to_vec()) - } else if let Ok(bytes) = value.cast::() { - Some(bytes.to_vec()) - } else { - None - } -} diff --git a/crates/spark-python-arrow/tests/bridge_regressions.rs b/crates/spark-python-arrow/tests/bridge_regressions.rs index 2ae8a27d..4939b81d 100644 --- a/crates/spark-python-arrow/tests/bridge_regressions.rs +++ b/crates/spark-python-arrow/tests/bridge_regressions.rs @@ -18,261 +18,167 @@ use std::ffi::CStr; use std::sync::Arc; -use arrow::array::{Array, Int64Array}; -use arrow::datatypes::DataType; -use pyo3::exceptions::{PyTypeError, PyValueError}; +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::{ - export_array, export_data_type, export_schema, import_array, import_batch, import_data_type, - import_schema, record_batches_to_pyarrow, -}; +use spark_python_arrow::bridge::record_batches_to_pyarrow; -fn setup<'py>(py: Python<'py>, code: &CStr) -> PyResult> { +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(code, Some(&locals), Some(&locals))?; - Ok(locals) -} - -fn value<'py>(locals: &Bound<'py, PyDict>, name: &str) -> PyResult> { - Ok(locals.get_item(name)?.expect("test variable must exist")) -} - -fn assert_equal(actual: &Bound<'_, PyAny>, expected: &Bound<'_, PyAny>) -> PyResult<()> { - assert!(actual.call_method1("equals", (expected,))?.is_truthy()?); - Ok(()) -} - -fn assert_schema_equal(actual: &Bound<'_, PyAny>, expected: &Bound<'_, PyAny>) -> PyResult<()> { + 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 - .call_method1("equals", (expected, true))? + .bind(py) + .getattr("schema")? + .call_method1("equals", (expected.getattr("schema")?, true))? .is_truthy()?); Ok(()) } #[test] -fn sorted_maps_keep_type_flags_through_schema_array_and_batch_paths() { +fn table_export_preserves_sorted_map_flags_in_nested_fields() { Python::initialize(); Python::attach(|py| -> PyResult<()> { - let locals = setup( + 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), metadata={b'nested': b'yes'}) -]) +nested_type = pa.struct([pa.field('maps', pa.list_(map_type))]) nested = pa.array([ - {'maps': [[('a', 1)]]}, None, {'maps': []} + {'maps': [[('a', 1), ('b', 2)]]}, {'maps': [None]}, {'maps': [[]]} ], type=nested_type) -schema = pa.schema([ - pa.field('maps', map_type, metadata={b'field': b'yes'}), - pa.field('nested', nested_type) -], metadata={b'schema': b'yes'}) -batch = pa.RecordBatch.from_arrays([maps, nested], schema=schema) -expected = pa.Table.from_batches([batch])", - )?; - - let schema = value(&locals, "schema")?; - let imported_schema = import_schema(&schema)?; - let exported_schema = export_schema(py, &imported_schema)?; - assert_schema_equal(exported_schema.bind(py), &schema)?; - for name in ["map_type", "nested_type"] { - let data_type = value(&locals, name)?; - let exported = export_data_type(py, &import_data_type(&data_type)?)?; - assert_equal(exported.bind(py), &data_type)?; - } - for name in ["maps", "nested"] { - let array = value(&locals, name)?; - let exported = export_array(py, import_array(&array)?)?; - assert_equal(exported.bind(py), &array)?; - assert_equal(&exported.bind(py).getattr("type")?, &array.getattr("type")?)?; - } - let batch = value(&locals, "batch")?; - let imported = import_batch(&batch)?; - let exported = record_batches_to_pyarrow(py, vec![imported])?; - assert_equal(exported.bind(py), &value(&locals, "expected")?)?; - assert_schema_equal(&exported.bind(py).getattr("schema")?, &schema) +expected = pa.Table.from_arrays([maps, nested], names=['maps', 'nested'])", + ) }) .unwrap(); } #[test] -fn ordered_dictionaries_keep_field_flags_in_nested_types_and_batches() { +fn table_export_preserves_ordered_dictionaries_in_nested_fields() { Python::initialize(); Python::attach(|py| -> PyResult<()> { - let locals = setup( + 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) -schema = pa.schema([ - pa.field('dictionary', dictionary.type), - pa.field('nested', nested.type), - pa.field('lists', lists.type) -], metadata={b'schema': b'ordered'}) -batch = pa.RecordBatch.from_arrays([dictionary, nested, lists], schema=schema) -expected = pa.Table.from_batches([batch])", - )?; - let schema = value(&locals, "schema")?; - let imported_schema = import_schema(&schema)?; - // Field equality deliberately ignores dictionary-order flags in Arrow Rust. - // Check them explicitly, then also compare the exported Python schema. - assert_eq!(imported_schema.field(0).dict_is_ordered(), Some(true)); - let DataType::Struct(fields) = imported_schema.field(1).data_type() else { - panic!("expected a struct") - }; - assert_eq!(fields[0].dict_is_ordered(), Some(true)); - let DataType::List(field) = imported_schema.field(2).data_type() else { - panic!("expected a list") - }; - assert_eq!(field.dict_is_ordered(), Some(true)); - let exported = export_schema(py, &imported_schema)?; - assert_schema_equal(exported.bind(py), &schema)?; - - for name in ["nested", "lists"] { - let array = value(&locals, name)?; - let data_type = array.getattr("type")?; - let exported_type = export_data_type(py, &import_data_type(&data_type)?)?; - assert_equal(exported_type.bind(py), &data_type)?; - let exported_array = export_array(py, import_array(&array)?)?; - assert_equal(exported_array.bind(py), &array)?; - assert_equal(&exported_array.bind(py).getattr("type")?, &data_type)?; - } - - let imported = import_batch(&value(&locals, "batch")?)?; - assert_eq!(imported.schema().field(0).dict_is_ordered(), Some(true)); - let DataType::Struct(fields) = imported.column(1).data_type() else { - panic!("expected a struct column") - }; - assert_eq!(fields[0].dict_is_ordered(), Some(true)); - let exported = record_batches_to_pyarrow(py, vec![imported])?; - assert_equal(exported.bind(py), &value(&locals, "expected")?)?; - assert_schema_equal(&exported.bind(py).getattr("schema")?, &schema) +expected = pa.Table.from_arrays([dictionary, nested, lists], + names=['dictionary', 'nested', 'lists'])", + ) }) .unwrap(); } #[test] -fn bare_ordered_dictionaries_require_a_field_to_preserve_ordering() { +fn table_export_rejects_inconsistent_nested_dictionary_order() { Python::initialize(); - Python::attach(|py| -> PyResult<()> { - let locals = setup( - py, - c"ordered = pa.DictionaryArray.from_arrays( - pa.array([0, 1], type=pa.int8()), pa.array(['a', 'b']), ordered=True -) -unordered = pa.DictionaryArray.from_arrays( - pa.array([0, 1], type=pa.int8()), pa.array(['a', 'b']) -)", - )?; - let ordered = value(&locals, "ordered")?; - for error in [ - import_array(&ordered).unwrap_err(), - import_data_type(&ordered.getattr("type")?).unwrap_err(), - ] { - assert!(error.is_instance_of::(py)); - assert!(error.to_string().contains("preserve ordering")); - } - let unordered = value(&locals, "unordered")?; - let exported = export_array(py, import_array(&unordered)?)?; - assert_equal(exported.bind(py), &unordered) - }) - .unwrap(); -} - -#[test] -fn batches_with_inconsistent_nested_dictionary_order_are_rejected() { - Python::initialize(); - Python::attach(|py| -> PyResult<()> { - let locals = setup( - py, - c"ordered = pa.DictionaryArray.from_arrays( - pa.array([1, None, 0], type=pa.int8()), pa.array(['a', 'b']), ordered=True -) -unordered = pa.DictionaryArray.from_arrays(ordered.indices, ordered.dictionary) -ordered_struct = pa.StructArray.from_arrays([ordered], names=['dictionary']) -unordered_struct = pa.StructArray.from_arrays([unordered], names=['dictionary']) -ordered_list = pa.ListArray.from_arrays(pa.array([0, 1, 2, 3], type=pa.int32()), ordered) -unordered_list = pa.ListArray.from_arrays(pa.array([0, 1, 2, 3], type=pa.int32()), unordered) -pairs = [ - (ordered, unordered), - (ordered_struct, unordered_struct), - (ordered_list, unordered_list) -] -batches = [ - (pa.record_batch([left], names=['x']), pa.record_batch([right], names=['x'])) - for left, right in pairs -]", - )?; - let batches = value(&locals, "batches")?; - for index in 0..batches.len()? { - let pair = batches.get_item(index)?; - let left = import_batch(&pair.get_item(0)?)?; - let right = import_batch(&pair.get_item(1)?)?; - let error = record_batches_to_pyarrow(py, vec![left, right]).unwrap_err(); + 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")); } - Ok(()) - }) - .unwrap(); -} - -#[test] -fn dictionary_value_extensions_are_rejected_instead_of_losing_identity() { - Python::initialize(); - Python::attach(|py| -> PyResult<()> { - let locals = setup( - py, - c"class TestExtension(pa.ExtensionType): - def __init__(self): - super().__init__(pa.int64(), 'spark_python_arrow.dictionary_extension') - - def __arrow_ext_serialize__(self): - return b'extension' - - @classmethod - def __arrow_ext_deserialize__(cls, storage_type, serialized): - return cls() - -extension = pa.ExtensionArray.from_storage(TestExtension(), pa.array([10, 20])) -dictionary = pa.DictionaryArray.from_arrays(pa.array([0, 1], type=pa.int8()), extension) -nested = pa.StructArray.from_arrays([dictionary], names=['dictionary']) -schema = pa.schema([pa.field('dictionary', dictionary.type)]) -batch = pa.RecordBatch.from_arrays([dictionary], schema=schema) -nested_schema = pa.schema([pa.field('nested', nested.type)]) -nested_batch = pa.RecordBatch.from_arrays([nested], schema=nested_schema)", - )?; - for name in ["dictionary", "nested"] { - let array = value(&locals, name)?; - for error in [ - import_array(&array).unwrap_err(), - import_data_type(&array.getattr("type")?).unwrap_err(), - ] { - assert!(error.is_instance_of::(py)); - assert!(error.to_string().contains("dictionary value extension")); - } - } - for name in ["schema", "nested_schema"] { - let error = import_schema(&value(&locals, name)?).unwrap_err(); - assert!(error.is_instance_of::(py)); - assert!(error.to_string().contains("dictionary value extension")); - } - for name in ["batch", "nested_batch"] { - let error = import_batch(&value(&locals, name)?).unwrap_err(); - assert!(error.is_instance_of::(py)); - assert!(error.to_string().contains("dictionary value extension")); - } - Ok(()) - }) - .unwrap(); + }); } #[test] -fn nullable_non_byte_aligned_slice_shares_values_without_requiring_bitmap_identity() { +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]); @@ -280,10 +186,15 @@ fn nullable_non_byte_aligned_slice_shares_values_without_requiring_bitmap_identi assert_eq!(slice.nulls().unwrap().offset(), 1); assert_eq!(slice.to_data().offset(), 0); let address = slice.values().as_ptr() as usize; - let exported = export_array(py, Arc::new(slice))?; + 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 = exported.bind(py); + let array = table + .bind(py) + .call_method1("column", (0,))? + .call_method1("chunk", (0,))?; assert_eq!( array .call_method0("buffers")? @@ -292,8 +203,8 @@ fn nullable_non_byte_aligned_slice_shares_values_without_requiring_bitmap_identi .extract::()?, address ); - // Arrow may realign the validity bitmap to the value buffer's offset. - // Its identity is not part of the bridge's value-buffer sharing promise. + // 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 diff --git a/crates/spark-python-arrow/tests/conversion.rs b/crates/spark-python-arrow/tests/conversion.rs index d3fe5912..b5773e41 100644 --- a/crates/spark-python-arrow/tests/conversion.rs +++ b/crates/spark-python-arrow/tests/conversion.rs @@ -15,190 +15,77 @@ // specific language governing permissions and limitations // under the License. -//! Standalone conversion checks against the installed PyArrow. No Spark server is needed. -//! These exercise the shared interfaces, not the full client conformance suite. +//! Table export checks against the installed PyArrow. No Spark server is needed. use std::collections::HashMap; -use std::ffi::CString; +use std::ffi::CStr; use std::sync::Arc; -use arrow::array::{Array, Int64Array}; +use arrow::array::{Array, Int64Array, StringArray, StructArray}; use arrow::datatypes::{DataType, Field, Schema, TimeUnit}; use arrow::record_batch::{RecordBatch, RecordBatchOptions}; -use pyo3::exceptions::{PyOverflowError, PyTypeError, PyUnicodeEncodeError, PyValueError}; +use pyo3::exceptions::PyValueError; use pyo3::prelude::*; -use pyo3::types::{PyBytes, PyDict, PyList, PyString}; -use spark_python_arrow::bridge::{ - export_array, export_data_type, export_schema, import_array, import_batch, import_data_type, - import_schema, record_batches_to_pyarrow, -}; -use spark_python_arrow::scalar::{ - extract_binary, extract_bool, extract_f64, extract_i64, extract_string, Scalar, -}; +use pyo3::types::PyDict; +use spark_python_arrow::bridge::record_batches_to_pyarrow; -fn context(py: Python<'_>) -> PyResult> { +fn evaluate<'py>(py: Python<'py>, expression: &CStr) -> PyResult> { let locals = PyDict::new(py); locals.set_item("pa", py.import("pyarrow")?)?; - locals.set_item("datetime", py.import("datetime")?)?; - Ok(locals) + py.eval(expression, Some(&locals), Some(&locals)) } -fn evaluate<'py>( - py: Python<'py>, - locals: &Bound<'py, PyDict>, - expression: &str, -) -> PyResult> { - let code = CString::new(expression).expect("test expressions contain no NUL"); - py.eval(&code, Some(locals), Some(locals)) -} - -fn extension_context<'py>(py: Python<'py>, name: &str) -> PyResult> { - let locals = context(py)?; - locals.set_item("extension_name", name)?; - py.run( - c"class TestExtensionType(pa.ExtensionType): - def __init__(self, storage_type=pa.string()): - super().__init__(storage_type, extension_name) - - def __arrow_ext_serialize__(self): - return b'extension metadata' - - @classmethod - def __arrow_ext_deserialize__(cls, storage_type, serialized): - assert serialized == b'extension metadata' - return cls(storage_type) - -extension_type = TestExtensionType() -extension_array = pa.ExtensionArray.from_storage( - extension_type, pa.array(['first', None, 'last']) -)", - Some(&locals), - Some(&locals), - )?; - Ok(locals) -} - -fn assert_array_equal(actual: &Bound<'_, PyAny>, expected: &Bound<'_, PyAny>) -> PyResult<()> { - assert!( - actual - .call_method1("equals", (expected,))? - .extract::()?, - "actual {actual}; expected {expected}" - ); +fn assert_schema_equal(actual: &Bound<'_, PyAny>, expected: &Bound<'_, PyAny>) -> PyResult<()> { + assert!(actual + .call_method1("equals", (expected, true))? + .is_truthy()?); Ok(()) } #[test] -fn exported_array_keeps_rust_buffers_alive_without_copying() { +fn table_export_shares_value_and_offset_buffers_after_rust_wrappers_drop() { Python::initialize(); Python::attach(|py| -> PyResult<()> { - let rust_array = Arc::new(Int64Array::from(vec![Some(12), None, Some(34)])); - let address = rust_array.values().as_ptr() as usize; - let array = export_array(py, rust_array.clone())?; - drop(rust_array); - py.import("gc")?.call_method0("collect")?; - - let array = array.bind(py); - let actual_address: usize = array - .call_method0("buffers")? - .get_item(1)? - .getattr("address")? - .extract()?; - assert_eq!(actual_address, address); - assert_eq!( - array - .call_method0("to_pylist")? - .extract::>>()?, - vec![Some(12), None, Some(34)] - ); - Ok(()) - }) - .unwrap(); -} - -#[test] -fn imported_slice_keeps_python_buffers_alive_without_copying() { - Python::initialize(); - Python::attach(|py| -> PyResult<()> { - let locals = context(py)?; - let original = evaluate(py, &locals, "pa.array([10, 20, None, 40], type=pa.int64())")?; - let address: usize = original - .call_method0("buffers")? - .get_item(1)? - .getattr("address")? - .extract()?; - let slice = original.call_method1("slice", (1, 2))?; - let imported = import_array(&slice)?; - drop(slice); - drop(original); + 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")?; - - let values = imported.as_any().downcast_ref::().unwrap(); - assert_eq!( - values.values().as_ptr() as usize, - address + std::mem::size_of::() - ); - assert_eq!(values.len(), 2); - assert_eq!(values.value(0), 20); - assert!(values.is_null(1)); - - let exported = export_array(py, imported)?; - let expected = evaluate(py, &locals, "pa.array([20, None], type=pa.int64())")?; - assert_array_equal(exported.bind(py), &expected) - }) - .unwrap(); -} - -#[test] -fn schema_and_nested_type_round_trip_preserves_metadata() { - Python::initialize(); - Python::attach(|py| -> PyResult<()> { - let field = Field::new("value", DataType::Int64, true) - .with_metadata(HashMap::from([("field_key".to_owned(), "a\0b".to_owned())])); - let schema = Schema::new_with_metadata( - vec![field], - HashMap::from([("schema_key".to_owned(), "c\0d".to_owned())]), - ); - let exported = export_schema(py, &schema)?; - assert_eq!(import_schema(exported.bind(py))?, schema); - - for data_type in [ - DataType::Int64, - DataType::Utf8, - DataType::List(Arc::new(Field::new("item", DataType::Int32, true))), - DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())), - DataType::Struct(schema.fields().clone()), - ] { - let exported = export_data_type(py, &data_type)?; - assert_eq!(import_data_type(exported.bind(py))?, data_type); + 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 + ); } - Ok(()) - }) - .unwrap(); -} - -#[test] -fn record_batch_import_preserves_schema_and_values() { - Python::initialize(); - Python::attach(|py| -> PyResult<()> { - let locals = context(py)?; - let batch = evaluate( + let expected = evaluate( py, - &locals, - "pa.RecordBatch.from_arrays([pa.array([1, None, 3], type=pa.int64())], \ - schema=pa.schema([pa.field('x', pa.int64(), metadata={b'f': b'v'})], \ - metadata={b's': b'v'}))", + c"pa.table({'number': pa.array([12, None, 34], type=pa.int64()), + 'text': ['first', None, 'last']})", )?; - let schema = import_schema(&batch.getattr("schema")?)?; - let imported = import_batch(&batch)?; - assert_eq!(imported.schema().as_ref(), &schema); - assert_eq!(imported.num_rows(), 3); - let exported = record_batches_to_pyarrow(py, vec![imported])?; - let expected = py - .import("pyarrow")? - .getattr("Table")? - .call_method1("from_batches", (PyList::new(py, [batch])?,))?; - assert_array_equal(exported.bind(py), &expected) + assert!(table + .bind(py) + .call_method1("equals", (expected,))? + .is_truthy()?); + Ok(()) }) .unwrap(); } @@ -214,7 +101,7 @@ fn table_export_preserves_chunks_and_shares_each_buffer() { ]; let addresses: Vec<_> = arrays .iter() - .map(|a| a.values().as_ptr() as usize) + .map(|array| array.values().as_ptr() as usize) .collect(); let batches = arrays .into_iter() @@ -246,218 +133,85 @@ fn table_export_preserves_chunks_and_shares_each_buffer() { } #[test] -fn empty_batches_and_zero_column_rows_round_trip() { +fn table_export_preserves_nested_values_and_nul_metadata() { 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 + let field = Arc::new( + Field::new("value", DataType::Int64, true) + .with_metadata(HashMap::from([("field\0key".into(), "a\0b".into())])), ); - - 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.clone())])?; - assert_eq!(empty.bind(py).getattr("num_rows")?.extract::()?, 0); - assert_eq!(import_schema(&empty.bind(py).getattr("schema")?)?, *schema); - - let options = RecordBatchOptions::new().with_row_count(Some(4)); - let batch = - RecordBatch::try_new_with_options(Arc::new(Schema::empty()), vec![], &options).unwrap(); - let table = record_batches_to_pyarrow(py, vec![batch])?; - assert_eq!(table.bind(py).getattr("num_rows")?.extract::()?, 4); - assert_eq!( - table.bind(py).getattr("num_columns")?.extract::()?, - 0 + let nested = StructArray::new( + vec![field].into(), + vec![Arc::new(Int64Array::from(vec![Some(1), None, Some(3)]))], + None, ); - - let locals = context(py)?; - let batch = evaluate( + 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, - &locals, - "pa.record_batch([pa.array([1, 2, 3])], names=['x']).select([])", + c"[{'nested': {'value': 1}}, {'nested': {'value': None}}, + {'nested': {'value': 3}}]", )?; - let imported = import_batch(&batch)?; - assert_eq!(imported.num_rows(), 3); - assert_eq!(imported.num_columns(), 0); + assert!(actual.eq(expected)?); Ok(()) }) .unwrap(); } #[test] -fn bridge_rejects_wrong_objects_without_calling_their_exporter() { +fn table_export_preserves_empty_schemas_and_zero_column_rows() { Python::initialize(); Python::attach(|py| -> PyResult<()> { - let locals = context(py)?; - py.run( - c"class Fake: - def _export_to_c(self, *args): - raise AssertionError('called')", - Some(&locals), - Some(&locals), + 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 fake = evaluate(py, &locals, "Fake()")?; - for error in [ - import_schema(&fake).unwrap_err(), - import_data_type(&fake).unwrap_err(), - import_array(&fake).unwrap_err(), - import_batch(&fake).unwrap_err(), - ] { - assert!(error.is_instance_of::(py), "{error}"); - } - Ok(()) - }) - .unwrap(); -} -#[test] -fn bare_extension_imports_fail_instead_of_losing_type_identity() { - Python::initialize(); - Python::attach(|py| -> PyResult<()> { - let locals = extension_context(py, "spark_python_arrow.tests.bare_extension")?; - let data_type = evaluate(py, &locals, "extension_type")?; - let array = evaluate(py, &locals, "extension_array")?; - for error in [ - import_data_type(&data_type).unwrap_err(), - import_array(&array).unwrap_err(), - ] { - assert!(error.is_instance_of::(py), "{error}"); - assert!(error.to_string().contains("extension types"), "{error}"); - } - Ok(()) - }) - .unwrap(); -} + 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())])")?, + )?; -#[test] -fn field_metadata_preserves_extensions_in_schemas_and_record_batches() { - Python::initialize(); - Python::attach(|py| -> PyResult<()> { - let name = "spark_python_arrow.tests.field_extension"; - let locals = extension_context(py, name)?; - let data_type = evaluate(py, &locals, "extension_type")?; - let pa = py.import("pyarrow")?; - pa.call_method1("register_extension_type", (&data_type,))?; - let result = (|| -> PyResult<()> { - let schema = evaluate( - py, - &locals, - "pa.schema([pa.field('value', extension_type, metadata={b'f': b'v'})], \ - metadata={b's': b'v'})", - )?; - let imported_schema = import_schema(&schema)?; + 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!( - imported_schema - .field(0) - .metadata() - .get("ARROW:extension:name") - .map(String::as_str), - Some(name) + table.bind(py).getattr("num_rows")?.extract::()?, + rows ); - let exported_schema = export_schema(py, &imported_schema)?; - assert!(exported_schema - .bind(py) - .call_method1("equals", (&schema, true))? - .is_truthy()?); - - // Nested extension fields have a metadata owner even when the outer - // value is a bare DataType, so only root extension types are rejected. - let nested = evaluate( - py, - &locals, - "pa.struct([pa.field('value', extension_type)])", - )?; - let exported_type = export_data_type(py, &import_data_type(&nested)?)?; - assert!(exported_type - .bind(py) - .call_method1("equals", (&nested,))? - .is_truthy()?); - - let kwargs = PyDict::new(py); - kwargs.set_item("schema", &schema)?; - let batch = pa.getattr("RecordBatch")?.call_method( - "from_arrays", - (PyList::new( - py, - [evaluate(py, &locals, "extension_array")?], - )?,), - Some(&kwargs), + 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'})")?, )?; - let imported = import_batch(&batch)?; - assert_eq!(imported.schema().as_ref(), &imported_schema); - let exported = record_batches_to_pyarrow(py, vec![imported])?; - let expected = pa - .getattr("Table")? - .call_method1("from_batches", (PyList::new(py, [batch])?,))?; - assert_array_equal(exported.bind(py), &expected)?; - assert!(exported - .bind(py) - .getattr("schema")? - .call_method1("equals", (&schema, true))? - .is_truthy()?); - Ok(()) - })(); - pa.call_method1("unregister_extension_type", (name,))?; - result - }) - .unwrap(); -} - -#[test] -fn table_export_rejects_inconsistent_schemas() { - Python::initialize(); - Python::attach(|py| { - let batches = ["a", "b"] - .into_iter() - .map(|name| { - RecordBatch::new_empty(Arc::new(Schema::new(vec![Field::new( - name, - DataType::Int64, - true, - )]))) - }) - .collect(); - let error = record_batches_to_pyarrow(py, batches).unwrap_err(); - assert!(error.is_instance_of::(py)); - }); -} - -#[test] -fn scalar_materialization_preserves_values_and_python_types() { - Python::initialize(); - Python::attach(|py| -> PyResult<()> { - let locals = context(py)?; - let cases = [ - (Scalar::Null, "None"), - (Scalar::Bool(true), "True"), - (Scalar::Int(i64::MIN), "-(2**63)"), - (Scalar::Int(i64::MAX), "2**63 - 1"), - (Scalar::Float(1.25), "1.25"), - (Scalar::String("a\0b"), "'a\\x00b'"), - (Scalar::String("\u{03bb}"), "'\\u03bb'"), - (Scalar::Binary(b"a\0b"), "b'a\\x00b'"), - (Scalar::Date(0), "datetime.date(1970, 1, 1)"), - (Scalar::Date(-719_162), "datetime.date.min"), - (Scalar::Date(2_932_896), "datetime.date.max"), - ]; - for (value, expression) in cases { - let actual = value.into_python(py)?; - let expected = evaluate(py, &locals, expression)?; - assert!(actual.eq(&expected)?, "{value:?}"); - assert!(actual.get_type().is(expected.get_type()), "{value:?}"); - } - let bytes = Scalar::Binary(b"value").into_python(py)?; - assert!(bytes.is_exact_instance_of::()); - for days in [-719_163, 2_932_897, i32::MIN, i32::MAX] { - let error = Scalar::Date(days).into_python(py).unwrap_err(); - let expected = evaluate( - py, - &locals, - &format!("datetime.date.fromordinal({})", i64::from(days) + 719_163), - ) - .unwrap_err(); - assert!(error.get_type(py).is(expected.get_type(py))); } Ok(()) }) @@ -465,125 +219,61 @@ fn scalar_materialization_preserves_values_and_python_types() { } #[test] -fn client_scalar_extraction_preserves_coercion_and_binary_copy_semantics() { +fn table_export_rejects_inconsistent_fields_and_metadata() { Python::initialize(); - Python::attach(|py| -> PyResult<()> { - let locals = context(py)?; - py.run( - c"class IntSubclass(int): pass -class FloatSubclass(float): pass -class StrSubclass(str): pass -class BytesSubclass(bytes): pass -class ByteArraySubclass(bytearray): pass -class IndexValue: - def __index__(self): - return 42 -class FloatValue: - def __float__(self): - return 1.5 -class TruthValue: - def __bool__(self): - raise AssertionError('truthiness must not be invoked') -class StringValue: - def __str__(self): - raise AssertionError('stringification must not be invoked') -", - Some(&locals), - Some(&locals), - )?; - let truth = evaluate(py, &locals, "True")?; - assert!(extract_bool(&truth)?); - assert!(!extract_bool(&evaluate(py, &locals, "False")?)?); - // The client tries bool before int because integer extraction also accepts bool. - assert_eq!(extract_i64(&truth)?, 1); - for expression in ["1", "None", "TruthValue()"] { - let error = extract_bool(&evaluate(py, &locals, expression)?).unwrap_err(); - assert!( - error.is_instance_of::(py), - "{expression}: {error}" - ); - } - for (expression, expected) in [ - ("7", 7), - ("IntSubclass(7)", 7), - ("IndexValue()", 42), - ("-(2**63)", i64::MIN), - ("2**63 - 1", i64::MAX), - ] { - assert_eq!(extract_i64(&evaluate(py, &locals, expression)?)?, expected); - } - for expression in ["1.5", "'7'"] { - let error = extract_i64(&evaluate(py, &locals, expression)?).unwrap_err(); - assert!( - error.is_instance_of::(py), - "{expression}: {error}" - ); - } - for expression in ["2**100", "-(2**100)"] { - let error = extract_i64(&evaluate(py, &locals, expression)?).unwrap_err(); - assert!( - error.is_instance_of::(py), - "{expression}: {error}" - ); - } - for (expression, expected) in [ - ("1.25", 1.25), - ("FloatSubclass(1.25)", 1.25), - ("7", 7.0), - ("FloatValue()", 1.5), - ("IndexValue()", 42.0), + 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())]))]), ] { - assert_eq!(extract_f64(&evaluate(py, &locals, expression)?)?, expected); - } - let error = extract_f64(&evaluate(py, &locals, "'1.25'")?).unwrap_err(); - assert!(error.is_instance_of::(py)); - for expression in ["'text'", "StrSubclass('text')"] { - assert_eq!(extract_string(&evaluate(py, &locals, expression)?)?, "text"); + 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")); } - let text = evaluate(py, &locals, "'a\\x00b'")?; - assert_eq!(extract_string(&text)?, "a\0b"); - for expression in ["7", "StringValue()", "b'text'"] { - let error = extract_string(&evaluate(py, &locals, expression)?).unwrap_err(); - assert!( - error.is_instance_of::(py), - "{expression}: {error}" - ); - } - let error = extract_string(&evaluate(py, &locals, "'\\ud800'")?).unwrap_err(); - assert!(error.is_instance_of::(py)); - for expression in [ - "b'ab'", - "BytesSubclass(b'ab')", - "bytearray(b'ab')", - "ByteArraySubclass(b'ab')", - ] { - assert_eq!( - extract_binary(&evaluate(py, &locals, expression)?), - Some(b"ab".to_vec()) - ); - } - let mutable = evaluate(py, &locals, "bytearray(b'ab')")?; - let snapshot = extract_binary(&mutable).expect("bytearray is supported"); - mutable.set_item(0, 99)?; - assert_eq!(snapshot.as_slice(), b"ab"); - assert_eq!(extract_binary(&mutable), Some(b"cb".to_vec())); - for expression in ["[97, 98]", "(97, 98)", "'ab'", "memoryview(b'ab')"] { - assert_eq!(extract_binary(&evaluate(py, &locals, expression)?), None); - } - assert!(!PyErr::occurred(py)); - Ok(()) - }) - .unwrap(); + }); } #[test] -fn string_scalar_materialization_preserves_unicode_and_embedded_nul() { +fn table_export_falls_back_to_ipc_for_nul_names_and_timezones() { Python::initialize(); Python::attach(|py| -> PyResult<()> { - for value in ["", "ascii", "a\0b", "\u{00e9}", "\u{03bb}", "\u{1f600}"] { - let actual = Scalar::String(value).into_python(py)?; - assert!(actual.is_exact_instance_of::()); - assert_eq!(actual.extract::()?, value); + 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(()) }) diff --git a/dev/design/architecture.md b/dev/design/architecture.md index e0a9f406..5ca6104c 100644 --- a/dev/design/architecture.md +++ b/dev/design/architecture.md @@ -47,7 +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 Python/Arrow bridge and scalar conversion; no transport + 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 @@ -74,12 +74,11 @@ 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. -The shared crate also owns the scalar routines used by the Python binding's -input conversion and row materialization. It does not own the client `Row` -representation, 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 -[shared converter contract](../../crates/spark-python-arrow/README.md). +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 6bd77262..09358778 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -63,7 +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/PyArrow C Data bridge and client scalar conversion | +| `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 diff --git a/python/tests/test_scalar_conversion.py b/python/tests/test_scalar_conversion.py new file mode 100644 index 00000000..52e3e782 --- /dev/null +++ b/python/tests/test_scalar_conversion.py @@ -0,0 +1,131 @@ +# +# 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. +# +"""Test client scalar conversion through Row without a Spark server.""" + +from datetime import date +from decimal import Decimal + +import pyarrow as pa +import pytest +from pyspark.sql import Row + +from test_arrow_handoff import _dataframe + + +@pytest.mark.parametrize( + "value", + [ + None, + True, + False, + -(2**63), + 2**63 - 1, + 1.25, + "", + "ascii", + "a\0b", + "\u00e9", + "\u03bb", + "\U0001f600", + b"a\0b", + Decimal("1.25"), + date(1970, 1, 1), + date.min, + date.max, + ], +) +def test_row_scalar_round_trip_preserves_value_and_type(value): + actual = Row(value=value).value + assert actual == value + assert type(actual) is type(value) + + +def test_row_extraction_preserves_numeric_coercion_and_builtin_subclasses(): + class IndexValue: + def __index__(self): + return 42 + + class FloatValue: + def __float__(self): + return 1.5 + + class IntSubclass(int): + pass + + class FloatSubclass(float): + pass + + class StrSubclass(str): + pass + + for value, expected in [ + (IndexValue(), 42), + (FloatValue(), 1.5), + (IntSubclass(7), 7), + (FloatSubclass(1.25), 1.25), + (StrSubclass("a\0b"), "a\0b"), + ]: + actual = Row(value=value).value + assert actual == expected + assert type(actual) is type(expected) + + +def test_row_binary_values_are_snapshots_but_integer_sequences_stay_lists(): + class BytesSubclass(bytes): + pass + + class ByteArraySubclass(bytearray): + pass + + for value in [b"ab", BytesSubclass(b"ab"), bytearray(b"ab"), ByteArraySubclass(b"ab")]: + row = Row(value=value) + if isinstance(value, bytearray): + value[0] = ord("c") + assert row.value == b"ab" + assert type(row.value) is bytes + + for value in [[97, 98], (97, 98)]: + actual = Row(value=value).value + assert actual == [97, 98] + assert type(actual) is list + + +def test_row_rejects_arbitrary_truthiness_stringification_and_buffers(): + class TruthValue: + def __bool__(self): + raise AssertionError("truthiness must not be invoked") + + class StringValue: + def __str__(self): + raise AssertionError("stringification must not be invoked") + + for value in [TruthValue(), StringValue(), memoryview(b"ab"), "\ud800"]: + with pytest.raises(TypeError, match="unsupported value type"): + Row(value=value) + + +@pytest.mark.parametrize("days", [-719_163, 2_932_897, -(2**31), 2**31 - 1]) +def test_collect_rejects_date32_values_outside_python_range(days): + batch = pa.record_batch([pa.array([days], type=pa.date32())], names=["value"]) + with _dataframe([batch]) as (dataframe, _): + row = dataframe.collect()[0] + + # Row materialization is lazy, so the exception occurs on field access. + with pytest.raises((ValueError, OverflowError)) as expected: + date.fromordinal(days + 719_163) + with pytest.raises(type(expected.value), match=str(expected.value)): + _ = row.value From a2d251b798d1c2bc854bb40e3233bb3efec8df2e Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Tue, 27 Oct 2026 16:02:33 -0700 Subject: [PATCH 8/8] Keep Arrow export on upstream PyO3 0.28 Defer the PyO3 upgrade and restore upstream row/scalar conversion. Keep the Arrow export optimization and its regression coverage, testing the converter against the workspace's locked PyO3 0.28.3. Generated-by: OpenAI Codex --- .github/workflows/coverage.yml | 2 +- .github/workflows/python-arrow.yml | 24 ++--- Cargo.lock | 98 +++++++++--------- Cargo.toml | 2 +- crates/pyspark-rs/Cargo.toml | 1 - crates/pyspark-rs/src/row.rs | 32 +++--- crates/pyspark-rs/src/session.rs | 10 +- crates/spark-python-arrow/Cargo.toml | 3 +- crates/spark-python-arrow/README.md | 22 ++--- python/tests/test_scalar_conversion.py | 131 ------------------------- 10 files changed, 84 insertions(+), 241 deletions(-) delete mode 100644 python/tests/test_scalar_conversion.py diff --git a/.github/workflows/coverage.yml b/.github/workflows/coverage.yml index 41b9bb22..5afc7c0d 100644 --- a/.github/workflows/coverage.yml +++ b/.github/workflows/coverage.yml @@ -228,7 +228,7 @@ jobs: run: | cd python PYTHONPATH=$(pwd) python -m coverage run --rcfile=.coveragerc -m pytest \ - tests/test_dropin_offline.py tests/test_arrow_handoff.py tests/test_scalar_conversion.py + 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 index e1efda1d..c4bde403 100644 --- a/.github/workflows/python-arrow.yml +++ b/.github/workflows/python-arrow.yml @@ -17,7 +17,7 @@ # under the License. # -name: Python Arrow compatibility +name: Python Arrow export on: push: @@ -34,13 +34,9 @@ concurrency: jobs: converter: - name: Python Arrow (PyO3 ${{ matrix.pyo3 }}) + name: Python Arrow export runs-on: ubuntu-latest timeout-minutes: 20 - strategy: - fail-fast: false - matrix: - pyo3: ["0.28.3", "0.29.0", "0.30.1"] steps: - uses: actions/checkout@v4 - uses: actions/setup-python@v5 @@ -52,20 +48,12 @@ jobs: rustup default stable - name: Install PyArrow run: python -m pip install "pyarrow==26.0.0" - - name: Isolate the standalone crate and select PyO3 - # The workspace lockfile pins 0.30.1. Build outside it so each matrix leg - # actually compiles the converter against the requested minor version. - run: | - cp -R crates/spark-python-arrow "$RUNNER_TEMP/spark-python-arrow" - cargo generate-lockfile --manifest-path "$RUNNER_TEMP/spark-python-arrow/Cargo.toml" - cargo update --manifest-path "$RUNNER_TEMP/spark-python-arrow/Cargo.toml" \ - -p pyo3 --precise "${{ matrix.pyo3 }}" - cargo tree --manifest-path "$RUNNER_TEMP/spark-python-arrow/Cargo.toml" --locked -i pyo3 - 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 --manifest-path "$RUNNER_TEMP/spark-python-arrow/Cargo.toml" \ - --locked --features pyo3/abi3-py39 + cargo test -p spark-python-arrow --locked --features pyo3/abi3-py39 - name: Check converter Clippy run: | - cargo clippy --manifest-path "$RUNNER_TEMP/spark-python-arrow/Cargo.toml" \ + cargo clippy -p spark-python-arrow \ --locked --all-targets --features pyo3/abi3-py39 -- -D warnings diff --git a/Cargo.lock b/Cargo.lock index 8f91c977..92e73626 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -102,7 +102,7 @@ version = "4.3.0" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn", ] [[package]] @@ -408,7 +408,7 @@ checksum = "c7c24de15d275a1ecfd47a380fb4d5ec9bfe0933f309ed5e705b775596a3574d" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn", ] [[package]] @@ -419,7 +419,7 @@ checksum = "9035ad2d096bed7955a320ee7e2230574d28fd3c3a0f186cbea1ff3c7eed5dbb" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn", ] [[package]] @@ -588,7 +588,7 @@ checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn", ] [[package]] @@ -1268,7 +1268,7 @@ checksum = "2e367e6a71051d0ebdd29b2f85d12059b38b1d1f172c6906e80016da662226bd" dependencies = [ "datafusion-doc", "quote", - "syn 2.0.117", + "syn", ] [[package]] @@ -1450,7 +1450,7 @@ checksum = "97369cbbc041bc366949bc74d34658d6cda5621039731c6310521892a3a20ae0" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn", ] [[package]] @@ -1659,7 +1659,7 @@ checksum = "e835b70203e41293343137df5c0664546da5745f82ec9b84d40be8336958447b" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn", ] [[package]] @@ -2386,7 +2386,7 @@ checksum = "ed3955f1a9c7c0c15e092f9c887db08b1fc683305fdf6eb6684f22555355e202" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn", ] [[package]] @@ -2573,7 +2573,7 @@ checksum = "c96395f0a926bc13b1c17622aaddda1ecb55d49c8f1bf9777e4d877800a43f8b" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn", ] [[package]] @@ -3232,7 +3232,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b" dependencies = [ "proc-macro2", - "syn 2.0.117", + "syn", ] [[package]] @@ -3271,7 +3271,7 @@ dependencies = [ "pulldown-cmark", "pulldown-cmark-to-cmark", "regex", - "syn 2.0.117", + "syn", "tempfile", ] @@ -3285,7 +3285,7 @@ dependencies = [ "itertools", "proc-macro2", "quote", - "syn 2.0.117", + "syn", ] [[package]] @@ -3329,9 +3329,9 @@ dependencies = [ [[package]] name = "pyo3" -version = "0.30.1" +version = "0.28.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "229ca13b8a15bb0d1487b81ab975bfcd97513416499a8312011b6036e42e0de4" +checksum = "91fd8e38a3b50ed1167fb981cd6fd60147e091784c427b8f7183a7ee32c31c12" dependencies = [ "libc", "once_cell", @@ -3343,18 +3343,18 @@ dependencies = [ [[package]] name = "pyo3-build-config" -version = "0.30.1" +version = "0.28.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "70130e2b6dd1ee589a8788026a429561ca5c40fa7a2fa7914c2282eddae24700" +checksum = "e368e7ddfdeb98c9bca7f8383be1648fd84ab466bf2bc015e94008db6d35611e" dependencies = [ "target-lexicon", ] [[package]] name = "pyo3-ffi" -version = "0.30.1" +version = "0.28.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5e7ff73c555c286d212d4ce7e5fbda845a973c4f61bd5f160f5f2c3c2a04d1c4" +checksum = "7f29e10af80b1f7ccaf7f69eace800a03ecd13e883acfacc1e5d0988605f651e" dependencies = [ "libc", "pyo3-build-config", @@ -3362,26 +3362,27 @@ dependencies = [ [[package]] name = "pyo3-macros" -version = "0.30.1" +version = "0.28.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c46d2e5862c7bb8207e3b639f713dda3c499933444784154c11d013c5a691b60" +checksum = "df6e520eff47c45997d2fc7dd8214b25dd1310918bbb2642156ef66a67f29813" dependencies = [ "proc-macro2", "pyo3-macros-backend", "quote", - "syn 3.0.9", + "syn", ] [[package]] name = "pyo3-macros-backend" -version = "0.30.1" +version = "0.28.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "38b89f066f2ae72a2aef52d1537ca2fcd104440ebfadee06e97aa90738a25b9b" +checksum = "c4cdc218d835738f81c2338f822078af45b4afdf8b2e33cbb5916f108b813acb" dependencies = [ "heck", "proc-macro2", + "pyo3-build-config", "quote", - "syn 3.0.9", + "syn", ] [[package]] @@ -3588,7 +3589,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "76009fbe0614077fc1a2ce255e3a1881a2e3a3527097d5dc6d8212c585e7e38b" dependencies = [ "quote", - "syn 2.0.117", + "syn", ] [[package]] @@ -3872,7 +3873,7 @@ checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn", ] [[package]] @@ -4024,7 +4025,7 @@ checksum = "028e551d5e270b31b9f3ea271778d9d827148d4287a5d96167b6bb9787f5cc38" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn", ] [[package]] @@ -4082,7 +4083,7 @@ dependencies = [ "heck", "proc-macro2", "quote", - "syn 2.0.117", + "syn", ] [[package]] @@ -4102,17 +4103,6 @@ dependencies = [ "unicode-ident", ] -[[package]] -name = "syn" -version = "3.0.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d78c8dee4c7bf0e14673097256fed6142ce9d3b85a408189d07482442145823b" -dependencies = [ - "proc-macro2", - "quote", - "unicode-ident", -] - [[package]] name = "sync_wrapper" version = "1.0.2" @@ -4130,7 +4120,7 @@ checksum = "728a70f3dbaf5bab7f0c4b1ac8d7ae5ea60a4b5549c8a5914361c99147a709d2" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn", ] [[package]] @@ -4183,7 +4173,7 @@ checksum = "ebc4ee7f67670e9b64d05fa4253e753e016c6c95ff35b89b7941d6b856dec1d5" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn", ] [[package]] @@ -4252,7 +4242,7 @@ checksum = "385a6cb71ab9ab790c5fe8d67f1645e6c450a7ce006a33de03daa956cf70a496" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn", ] [[package]] @@ -4330,7 +4320,7 @@ dependencies = [ "prettyplease", "proc-macro2", "quote", - "syn 2.0.117", + "syn", ] [[package]] @@ -4355,7 +4345,7 @@ dependencies = [ "prost-build", "prost-types", "quote", - "syn 2.0.117", + "syn", "tempfile", "tonic-build", ] @@ -4428,7 +4418,7 @@ checksum = "7490cfa5ec963746568740651ac6781f701c9c5ea257c58e057f3ba8cf69e8da" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn", ] [[package]] @@ -4645,7 +4635,7 @@ dependencies = [ "bumpalo", "proc-macro2", "quote", - "syn 2.0.117", + "syn", "wasm-bindgen-shared", ] @@ -4832,7 +4822,7 @@ checksum = "053e2e040ab57b9dc951b72c264860db7eb3b0200ba345b4e4c3b14f67855ddf" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn", ] [[package]] @@ -4843,7 +4833,7 @@ checksum = "3f316c4a2570ba26bbec722032c4099d8c8bc095efccdc15688708623367e358" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn", ] [[package]] @@ -5040,7 +5030,7 @@ dependencies = [ "heck", "indexmap", "prettyplease", - "syn 2.0.117", + "syn", "wasm-metadata", "wit-bindgen-core", "wit-component", @@ -5056,7 +5046,7 @@ dependencies = [ "prettyplease", "proc-macro2", "quote", - "syn 2.0.117", + "syn", "wit-bindgen-core", "wit-bindgen-rust", ] @@ -5129,7 +5119,7 @@ checksum = "de844c262c8848816172cef550288e7dc6c7b7814b4ee56b3e1553f275f1858e" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn", "synstructure", ] @@ -5150,7 +5140,7 @@ checksum = "1ae7f38b72ec2a254e2b87ef277cf2cd4fb97cbebf944faa6f33354da0867930" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn", ] [[package]] @@ -5170,7 +5160,7 @@ checksum = "11532158c46691caf0f2593ea8358fed6bbf68a0315e80aae9bd41fbade684a1" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn", "synstructure", ] @@ -5210,7 +5200,7 @@ checksum = "625dc425cab0dca6dc3c3319506e6593dcb08a9f387ea3b284dbd52a92c40555" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 9e90594d..31fe8090 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -87,7 +87,7 @@ datafusion = { version = "53", default-features = false } polars = { version = "0.54", default-features = false, features = ["ipc"] } # PyO3 for the pyspark-rs CPython extension. -pyo3 = { version = "0.30.1", features = ["extension-module", "abi3-py39"] } +pyo3 = { version = "0.28", features = ["extension-module", "abi3-py39"] } # misc bytes = "1" diff --git a/crates/pyspark-rs/Cargo.toml b/crates/pyspark-rs/Cargo.toml index cb6bb62c..05278bad 100644 --- a/crates/pyspark-rs/Cargo.toml +++ b/crates/pyspark-rs/Cargo.toml @@ -2,7 +2,6 @@ name = "pyspark-rs" version.workspace = true edition.workspace = true -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. diff --git a/crates/pyspark-rs/src/row.rs b/crates/pyspark-rs/src/row.rs index 91549e5a..dfdf506d 100644 --- a/crates/pyspark-rs/src/row.rs +++ b/crates/pyspark-rs/src/row.rs @@ -2,8 +2,7 @@ use pyo3::exceptions::{PyAttributeError, PyIndexError, PyKeyError}; use pyo3::prelude::*; -use pyo3::types::{PyBytes, PyDict, PyList, PyString}; -use pyo3::IntoPyObjectExt; +use pyo3::types::{PyBytes, PyDict, PyList}; use spark_connect::row::{Row, Value}; /// Python wrapper for a Spark Row. @@ -18,28 +17,29 @@ impl PyRow { } } -/// Convert a Spark `Value` into a Python object using PyO3's `IntoPyObject`. +/// Convert a Spark `Value` into a Python object (PyO3 0.28 `IntoPyObject`). pub fn value_to_py<'py>(py: Python<'py>, v: &Value) -> PyResult> { // Propagate conversion failures as Python exceptions with `?` rather than // panicking across the FFI boundary (the conversions are near-infallible, // but a panic here would unwind into CPython). Ok(match v { Value::Null => py.None().into_bound(py), - Value::Bool(b) => b.into_bound_py_any(py)?, - Value::Byte(n) => n.into_bound_py_any(py)?, - Value::Short(n) => n.into_bound_py_any(py)?, - Value::Integer(n) => n.into_bound_py_any(py)?, - Value::Long(n) => n.into_bound_py_any(py)?, - Value::Float(f) => f.into_bound_py_any(py)?, - Value::Double(f) => f.into_bound_py_any(py)?, - Value::String(s) => PyString::new(py, s).into_any(), + Value::Bool(b) => b.into_pyobject(py)?.to_owned().into_any(), + Value::Byte(n) => (*n as i64).into_pyobject(py)?.into_any(), + Value::Short(n) => (*n as i64).into_pyobject(py)?.into_any(), + Value::Integer(n) => (*n as i64).into_pyobject(py)?.into_any(), + Value::Long(n) => (*n).into_pyobject(py)?.into_any(), + Value::Float(f) => (*f as f64).into_pyobject(py)?.into_any(), + Value::Double(f) => (*f).into_pyobject(py)?.into_any(), + Value::String(s) => s.as_str().into_pyobject(py)?.into_any(), Value::Binary(b) => PyBytes::new(py, b).into_any(), Value::Date(d) => { - // 719163 is the ordinal of 1970-01-01. Use i64 to avoid overflow - // before Python reports dates outside its supported range. - py.import("datetime")? - .getattr("date")? - .call_method1("fromordinal", (i64::from(*d) + 719_163,))? + // Convert days-since-epoch (i32) to datetime.date + let date_mod = py.import("datetime")?; + let date_cls = date_mod.getattr("date")?; + // 719163 is the ordinal of 1970-01-01 + let ordinal = (*d as i64) + 719163i64; + date_cls.call_method1("fromordinal", (ordinal,))?.into_any() } Value::Timestamp(t) => { // Convert microseconds-since-epoch (i64) to datetime.datetime diff --git a/crates/pyspark-rs/src/session.rs b/crates/pyspark-rs/src/session.rs index 3f6a42bb..5650e606 100644 --- a/crates/pyspark-rs/src/session.rs +++ b/crates/pyspark-rs/src/session.rs @@ -1,7 +1,7 @@ //! PyO3 wrapper for Spark Connect sessions. use pyo3::prelude::*; -use pyo3::types::{PyByteArray, PyBytes, PyDict, PyList}; +use pyo3::types::{PyDict, PyList}; use spark_connect::row::{Row, Value}; use spark_connect::session::SparkSession; use spark_connect::types::{DataType, StructField}; @@ -980,11 +980,11 @@ pub(crate) fn py_to_value(obj: &Bound<'_, PyAny>) -> PyResult { // Check for bytes/bytearray SPECIFICALLY (not `extract::>()`, which also // matches a Python list of small ints like [1,2,3] and would mis-tag arrays as binary). - if let Ok(bytes) = obj.cast::() { - return Ok(Value::Binary(bytes.as_bytes().to_vec())); + if let Ok(b) = obj.cast::() { + return Ok(Value::Binary(b.as_bytes().to_vec())); } - if let Ok(bytes) = obj.cast::() { - return Ok(Value::Binary(bytes.to_vec())); + if let Ok(ba) = obj.cast::() { + return Ok(Value::Binary(ba.to_vec())); } // list / tuple -> array (recursively converted) diff --git a/crates/spark-python-arrow/Cargo.toml b/crates/spark-python-arrow/Cargo.toml index 18b401e9..ada92c39 100644 --- a/crates/spark-python-arrow/Cargo.toml +++ b/crates/spark-python-arrow/Cargo.toml @@ -27,8 +27,7 @@ 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. -# Compatibility CI resolves each supported minor independently of the workspace lockfile. -pyo3 = { version = ">=0.28, <0.31" } +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 index 66b24d9f..72876801 100644 --- a/crates/spark-python-arrow/README.md +++ b/crates/spark-python-arrow/README.md @@ -66,21 +66,19 @@ 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 manifest keeps its package metadata self-contained so the compatibility -tests can build it outside the workspace lockfile. The library does not -enable PyO3's `extension-module` feature or select a minimum Python ABI; +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 permits PyO3 0.28, 0.29, and 0.30; the Python extension selects -0.30.1. Both this crate and `pyspark-rs` declare Rust 1.85, matching the declared -minimum of Arrow 58.3.0 and PyO3 0.30.1. 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 isolated compatibility CI job tests each PyO3 minor independently of the -workspace lockfile. A version is not considered verified until its corresponding -job passes. These tests run in the compatibility workflow, not again in coverage. +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 diff --git a/python/tests/test_scalar_conversion.py b/python/tests/test_scalar_conversion.py deleted file mode 100644 index 52e3e782..00000000 --- a/python/tests/test_scalar_conversion.py +++ /dev/null @@ -1,131 +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. -# -"""Test client scalar conversion through Row without a Spark server.""" - -from datetime import date -from decimal import Decimal - -import pyarrow as pa -import pytest -from pyspark.sql import Row - -from test_arrow_handoff import _dataframe - - -@pytest.mark.parametrize( - "value", - [ - None, - True, - False, - -(2**63), - 2**63 - 1, - 1.25, - "", - "ascii", - "a\0b", - "\u00e9", - "\u03bb", - "\U0001f600", - b"a\0b", - Decimal("1.25"), - date(1970, 1, 1), - date.min, - date.max, - ], -) -def test_row_scalar_round_trip_preserves_value_and_type(value): - actual = Row(value=value).value - assert actual == value - assert type(actual) is type(value) - - -def test_row_extraction_preserves_numeric_coercion_and_builtin_subclasses(): - class IndexValue: - def __index__(self): - return 42 - - class FloatValue: - def __float__(self): - return 1.5 - - class IntSubclass(int): - pass - - class FloatSubclass(float): - pass - - class StrSubclass(str): - pass - - for value, expected in [ - (IndexValue(), 42), - (FloatValue(), 1.5), - (IntSubclass(7), 7), - (FloatSubclass(1.25), 1.25), - (StrSubclass("a\0b"), "a\0b"), - ]: - actual = Row(value=value).value - assert actual == expected - assert type(actual) is type(expected) - - -def test_row_binary_values_are_snapshots_but_integer_sequences_stay_lists(): - class BytesSubclass(bytes): - pass - - class ByteArraySubclass(bytearray): - pass - - for value in [b"ab", BytesSubclass(b"ab"), bytearray(b"ab"), ByteArraySubclass(b"ab")]: - row = Row(value=value) - if isinstance(value, bytearray): - value[0] = ord("c") - assert row.value == b"ab" - assert type(row.value) is bytes - - for value in [[97, 98], (97, 98)]: - actual = Row(value=value).value - assert actual == [97, 98] - assert type(actual) is list - - -def test_row_rejects_arbitrary_truthiness_stringification_and_buffers(): - class TruthValue: - def __bool__(self): - raise AssertionError("truthiness must not be invoked") - - class StringValue: - def __str__(self): - raise AssertionError("stringification must not be invoked") - - for value in [TruthValue(), StringValue(), memoryview(b"ab"), "\ud800"]: - with pytest.raises(TypeError, match="unsupported value type"): - Row(value=value) - - -@pytest.mark.parametrize("days", [-719_163, 2_932_897, -(2**31), 2**31 - 1]) -def test_collect_rejects_date32_values_outside_python_range(days): - batch = pa.record_batch([pa.array([days], type=pa.date32())], names=["value"]) - with _dataframe([batch]) as (dataframe, _): - row = dataframe.collect()[0] - - # Row materialization is lazy, so the exception occurs on field access. - with pytest.raises((ValueError, OverflowError)) as expected: - date.fromordinal(days + 719_163) - with pytest.raises(type(expected.value), match=str(expected.value)): - _ = row.value