Repository navigation
Conversation
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 <no-reply@databricks.com>
dbtsai
changed the title
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 <no-reply@databricks.com>
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 <no-reply@databricks.com>
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 <no-reply@databricks.com>
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 <no-reply@databricks.com>
Inline client string and date materialization, remove stale worker and decoder wording, and exercise the remaining Scalar string API. Co-authored-by: Isaac <no-reply@databricks.com>
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 <no-reply@databricks.com>
dbtsai
changed the title
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
dbtsai
changed the title
|
Thanks for this. Avoiding the IPC round-trip is a clear win, and the benchmarks make the case well. Before this lands, I'd like to question how much of the bridge needs to be custom. arrow-rs 58.3 already ships this handoff: But I don't think Spark can produce that case:
If that's right, A couple of questions:
|
| CARGO_NET_OFFLINE=true CARGO_REGISTRIES_CRATES_IO_PROTOCOL=git \ | ||
| cargo test -p spark-python-arrow --locked --features pyo3/abi3-py39 |
There was a problem hiding this comment.
This command differs from what CI runs (cargo test -p spark-python-arrow --locked --features pyo3/abi3-py39 in python-arrow.yml). CARGO_NET_OFFLINE=true makes it fail unless every dependency is already in the local Cargo cache, and CARGO_REGISTRIES_CRATES_IO_PROTOCOL=git switches away from Cargo's default sparse index. Could the README show the CI command, with the offline variables as an optional note for offline environments?
| - name: Install PyArrow | ||
| run: python -m pip install "pyarrow==26.0.0" |
There was a problem hiding this comment.
pyproject.toml declares pyarrow>=10.0.0, and toArrow() now goes through PyArrow's private Schema._import_from_c / RecordBatch._import_from_c instead of the public IPC reader. CI pins only 26.0.0 here, and the other workflows install the latest PyArrow, so the PyArrow 10 run in the PR description isn't repeated in CI. Could this job add a leg at the declared floor? If that isn't practical, the README could state which PyArrow versions are tested.
| // 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. |
There was a problem hiding this comment.
Could this link the upstream issue (apache/arrow-rs#11603, fixed on main by apache/arrow-rs#11630)? That makes it clear when export_c_data_type/rebuild and the 0b100 flag below can be removed after an Arrow upgrade.
| 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)?)?; | ||
| } |
There was a problem hiding this comment.
Both NUL-fallback tests use a single batch: this one (empty) and test_nul_schema_strings_fall_back_without_executing_again. Two more cases would cover the rest of the fallback's contract: (1) two non-empty batches, asserting num_chunks == 2 and the values; (2) a NUL field name with mismatched batch schemas, expecting ValueError. Case (2) pins the schema check in bridge.rs (L60-66) ahead of the fallback branch (L68).
viirya
left a comment
viirya
left a comment
There was a problem hiding this comment.
Thanks for splitting out the PyO3 upgrade and adding benchmarks. The memory numbers for large toArrow()/toPandas() results make a good case. I checked the FFI ownership and it looks right: the schema is imported before any array, rejected imports are released through Drop, and zero-column batches keep their row count. +1 to @tdcmeehan's question about arrow::pyarrow::Table. The answer decides how much of this crate needs to stay, so I'll avoid repeating those points.
One gap in the testing evidence: the official collection parity run goes through rust_transport_plugin.py, so official pyspark decodes the Arrow data and PyDataFrame.toArrow never runs. The pandas-parity gate works the same way. The three native toArrow() smoke cases you ran against a local server are the only real-server coverage of this path, and they aren't in the PR. Could you commit them as drop-in e2e tests so CI runs them on every change?
| )); | ||
| } | ||
| let pa = py.import("pyarrow")?; | ||
| if schema.fields().iter().any(|field| field_needs_ipc(field)) { |
There was a problem hiding this comment.
Instead of walking the type tree up front, could this call export_schema first and fall back to IPC if it returns an error? The schema is exported before any array, so nothing has been consumed when it fails. That would remove field_needs_ipc/data_type_needs_ipc (L202–225) and also cover any other schema the C exporter rejects. If you move to arrow::pyarrow::Table, the same approach works, but keep a clone of the batch Vec for the fallback (cheap, since the batches are Arc-backed), because into_pyarrow consumes the table.
| name = "spark-python-arrow" | ||
| version = "0.1.0" | ||
| edition = "2021" | ||
| rust-version = "1.85" |
There was a problem hiding this comment.
Now that the crate is no longer built outside the workspace, edition, license and arrow could use .workspace = true like every other crate. On rust-version: apache-spark-connect already depends on arrow 58.3, which declares Rust 1.85, so the workspace's 1.82 is already wrong for the core crate too, not just this one (README L76). A separate change that raises the workspace value to 1.85 would let this crate inherit it and make that README paragraph unnecessary.
| python-version: "3.11" | ||
| - name: Install Rust | ||
| run: | | ||
| rustup toolchain install stable --profile minimal --component clippy |
There was a problem hiding this comment.
This installs latest stable and then runs clippy with -D warnings (L59), so the job will start failing whenever a new stable release adds lints, even if this crate hasn't changed. Could it pin a toolchain, or match the workspace lint job's setup?
|
|
||
|
|
||
| @contextmanager | ||
| def _dataframe(batches, *, payloads=None, calls=None): |
There was a problem hiding this comment.
This file is 521 lines. It hand-encodes the protobuf wire format with hard-coded field numbers, mirrors ArrowArray in ctypes, hooks the release callbacks, and has the autouse _forbid_c_stream_import fixture @tdcmeehan asked about. That's a lot to maintain, and much of it depends on how the export is implemented. Buffer identity and lifetime are already covered by conversion.rs. Could the Python side be cut down to the public toArrow()/toPandas() behavior plus one importer-failure case, with the real-server e2e covering the rest?
|
Following up on the question about The core of let batches = py.detach(|| self.dataframe.collect_record_batches()).to_pyerr()?;
let schema = batches.first().map(|b| b.schema()).unwrap_or_else(|| Arc::new(Schema::empty()));
match Table::try_new(batches.clone(), schema.clone()) { // Arc clones, kept for the fallback
Ok(table) => match table.into_pyarrow(py) { ... } // NUL name/timezone -> IPC fallback
Err(e) => Err(PyValueError::new_err(e.to_string())),
}That is roughly 20 lines in
Most of the tests would go with the crate. In
Spark never produces either input: maps are always written with What I think is worth keeping is the NUL fallback tests and one buffer-address assertion, so a future change can't quietly reintroduce the IPC copy. Both fit in |
|
Closing this alternative in favor of #108. We favor its use of arrow-rs's upstream whole-result C Stream reader, which avoids both custom stream callbacks and per-batch reader setup. Further review and validation will continue on #108, where the remaining memory and validation caveats are documented. Thanks for the feedback on these alternatives. |
Provide an alternative to the per-batch C Data handoff in PR apache#103 while retaining PyO3 0.28.3. Add ownership, buffer-sharing, compatibility, and public client handoff tests. Generated-by: OpenAI Codex
… Interface ### What changes were proposed in this pull request? Under [SPARK-60586](https://issues.apache.org/jira/browse/SPARK-60586), replace the extra Rust Arrow → IPC bytes → PyArrow round trip in native `DataFrame.toArrow()` with arrow-rs's built-in export of one `Box<dyn RecordBatchReader + Send>`: ```rust let reader: Box<dyn RecordBatchReader + Send> = Box::new(RecordBatchIterator::new(batches.into_iter().map(Ok), schema)); Ok(reader.into_pyarrow(py)?.call_method0("read_all")?.unbind()) ``` The helper is `crates/pyspark-rs/src/pyarrow_export.rs`. Arrow owns the FFI callbacks, ownership transfer, and release handling; the client adds no custom unsafe code. The separate converter crate and custom callback adapter are gone. Before exporting, validate batch schema consistency and retain an IPC fallback for field names or timezones containing NUL. The standalone PyArrow schema preflight was removed: its claimed older-version abort rationale was unproven. Arrow 58.3.0 itself uses `CString::new(...).unwrap()` when exporting [formats](https://github.com/apache/arrow-rs/blob/58.3.0/arrow-schema/src/ffi.rs#L148) and [names](https://github.com/apache/arrow-rs/blob/58.3.0/arrow-schema/src/ffi.rs#L174), so the explicit NUL check remains necessary before export. PyO3 remains **0.28.3**, Arrow **58.3.0**; enabling Arrow's `pyarrow` feature adds `arrow-pyarrow 58.3.0`. The contract covers ordinary Spark results: unsorted maps and non-dictionary-encoded fields. Row/scalar conversion, UDF workers, and Python-to-Rust imports are unchanged. #103 and #117 are alternative implementations, not dependencies. ### Why are the changes needed? Already-decoded Arrow batches previously had to be serialized into another IPC file, copied into Python bytes, and decoded again. One upstream reader now consumes the entire result without a Python import call per batch. All decoded batches are still collected into a Rust vector first; this does not provide network streaming or bounded-memory collection. ### Does this PR introduce _any_ user-facing change? `toArrow()` shares compatible decoded value/offset buffers with PyArrow. `toPandas()` uses the same handoff, but pandas conversion may copy. Network IPC decoding, descriptors, Python rows, sliced validity-bitmap realignment, and the compatibility fallback are outside the buffer-sharing claim. **Intentional behavior change:** inconsistent batch schemas now raise `ValueError`, including differences in nullability or field/schema metadata. Previously batches were written under the first batch's schema. Validation happens before the IPC fallback. Spark normally produces batches from one result schema. Values, metadata, batch boundaries, and zero-column row counts are covered by public-behavior tests. An empty batch vector retains master's empty-schema behavior; retaining the typed schema when no batches arrive is a separate follow-up. NUL-containing field names/timezones use the already-collected batches for IPC without another RPC. ### How was this patch tested? Current tested source: **`e60aa53`**, based on upstream master `d884e34`. Rebuilt locally on macOS arm64 with Rust 1.99.1 and optimized release settings: - **299 passed, 1 skipped** on Python 3.12 / PyArrow 26.0.0. - **299 passed, 1 skipped** on Python 3.9 / PyArrow 10.0.0, after removing the schema preflight. Each run includes **26 handoff/fixture cases** plus the offline suite. - **24 fixture-dependent handoff cases skip locally with an actionable message when `protoc` is unavailable; CI fails instead of skipping. Two additional tests cover both missing-protoc outcomes.** With `protoc`, the mock uses descriptors compiled from the repository's protocol instead of hand-encoded field numbers. - Tests cover values, nested types, metadata, chunks, empty/zero-column results, `toPandas()`, NUL fallback, schema mismatch, and buffers remaining usable after their session/table is dropped. Post-deletion value checks establish retained-buffer lifetime; redundant address-equality assertions were removed. **These checks do not independently prove Rust-to-Python buffer identity.** The private FFI mirrors, callback hooks, schema monkeypatch, and subprocess-isolated boundary spies were removed. - The normal-handoff regression guard makes `pyarrow.ipc.open_file` fail; it passes on this head and fails against master `d884e34`. NUL cases explicitly verify the IPC fallback, including list-of-struct and map-value-struct names. Schema-mismatch cases now include nullability. - **392 core Rust unit tests passed**, all integrations compiled with and without `wasm-udf`, **3 native live-server smoke tests passed**, and official collection parity passed (**15 passed, 1 skipped**) on the review-fix source. - Formatting, Python lint, and whitespace checks pass. The preceding source `d89c01b` also passed **392 core Rust unit tests**, compiled all core integration tests with and without `wasm-udf`, passed **3 native `toArrow()` live-server smoke tests**, and passed official collection parity through the Rust transport (**15 passed, 1 skipped**). Collection parity required `spark.sql.timeType.enabled=true` on the local Spark 4.2 server. The current review follow-up strengthens regression coverage and documents a matched master RSS control; it makes no production behavior changes. **CI scope:** the coverage workflow invokes `test_arrow_handoff.py` and installs `protobuf-compiler`. Follow-up `e0eb002` adds a native smoke step in the Connect parity job of `build_python_connect.yml`, using the already-built extension and running `python -m pytest tests/test_arrow_e2e.py -q` with `SPARK_REMOTE` and the local package on `PYTHONPATH`. It runs after the official suite and before the server restarts in local-cluster mode, so the official suite still begins with a pristine server. Workflow YAML parsing and the exact pytest command passed locally (**3 passed**). CI results for this follow-up are separate from those local checks. PyArrow 10 validation remains local; there is no pinned PyArrow 10 CI job. #### Historical benchmark evidence — not measured on the current head Using [pyspark-connect-benchmark at `1bbdec9`](https://github.com/dbtsai/pyspark-connect-benchmark/tree/1bbdec9f3665fcd2ec2ee5129bbcdb95ee4b00c4), the earlier upstream-reader candidate `a9ee289` was compared with master `d884e34`. That candidate still had the standalone schema preflight and the older test suite. The PR has since been rebased and updated; these timings are **historical results, not current-head measurements**. The matched optimized builds used Rust 1.99.1, `opt-level=3`, `abi3-py39`, PyO3 0.28.3, Arrow 58.3.0, Python 3.12.14 and PyArrow 26.0.0 on an Apple M3 Max. **240/240 samples passed** across 10 workloads, 6 fresh processes per client, and 2 counterbalanced rounds. Validation checked shape and first/last values per column, not every cell. Representative historical medians, **master → earlier candidate**: | Workload | Round | Wall time (ms) | Sampled peak RSS (MiB) | |---|---|---:|---:| | 38 MiB numeric `toArrow()` | Forward | 24.843 → 14.255 | 268.94 → 146.92 | | 38 MiB numeric `toArrow()` | Reversed | 27.990 → 14.005 | 268.90 → 146.92 | | 128 MiB numeric `toArrow()` | Forward | 63.207 → 30.491 | 628.53 → 235.24 | | 128 MiB numeric `toArrow()` | Reversed | 63.300 → 31.891 | 629.55 → 236.22 | | 38 MiB numeric `toPandas()` | Forward | 30.023 → 15.494 | 308.38 → 186.51 | | 38 MiB numeric `toPandas()` | Reversed | 29.995 → 15.163 | 308.52 → 186.50 | These are shared-desktop measurements of client receive/conversion against canned localhost RPC responses, not Spark execution. RSS sampling can miss short peaks; timings and RSS do not establish buffer identity. **RSS control added during review.** The identical 28-call, 128 MiB probe was rerun in fresh processes against rebuilt master `d884e34` and this C Stream revision, using the same environment. After four warmup calls, master rose from **1129.23 to 3976.06 MiB**, while C Stream rose from **418.81 to 569.47 MiB**. All results passed validation. Growth is therefore not unique to this PR; the control grew more in this single run. Both mock server and client still share one process, so this does not identify the retention source or establish leak freedom or a general memory advantage. The earlier C Stream probe rose from 416.27 to 481.66 MiB, and run-to-run variation remains unexplained. The methodology and reproducer are included below. <details> <summary>RSS control methodology and reproducer</summary> This is a local diagnostic for PR #108, not a performance benchmark or proof of leak freedom. Both the mock gRPC server and the native client run in one Python process, so the measurement includes their allocations and allocator retention. On 2026-10-29, the same 28-call probe was run in fresh processes on master `d884e34` (IPC handoff) and the review revision of #108 based on `fe79d54` (C Stream handoff; production changes since that head are comments only). Both extensions were rebuilt with the same release Cargo configuration and Rust 1.99.1, then loaded using the same Python 3.12 / PyArrow 26 environment on macOS arm64. Master ran first. These are single runs on a shared desktop. Each call receives 128 one-MiB numeric batches from the same mock service, validates the 16,777,216-row table's endpoints, deletes the table, forces Python GC, and samples process RSS. The first four calls are warmup. All calls passed validation. | Handoff | First measured RSS (MiB) | Last measured RSS (MiB) | Measured minimum–maximum (MiB) | | --- | ---: | ---: | ---: | | Master IPC | 1129.23 | 3976.06 | 1129.23–3976.06 | | PR C Stream | 418.81 | 569.47 | 418.81–569.47 | Growth is not unique to the C Stream change; the master control grew much more in this run. This does not identify the source of retention, isolate client allocations, or establish a general memory advantage. The earlier C Stream run reported 416.27 → 481.66 MiB; variation between runs remains unexplained. A claim of flat post-GC RSS is not supported. To reproduce, run the following from the PR checkout with its tests available, setting `PYTHONPATH` to `<PR checkout>/python`, the directory containing `pyspark`. For the master control, build master and set `PYTHONPATH` to `<master checkout>/python`; keep this test harness and all other environment settings the same. `PROTOC` and `PROTOC_INCLUDE` must identify an installed protocol compiler and includes. The baseline IPC implementation does not include the new test file, so the working directory must remain the PR checkout for both runs. ```python import gc import json import runpy import tempfile from pathlib import Path import numpy as np import psutil import pyarrow as pa namespace = runpy.run_path("python/tests/test_arrow_handoff.py") class TempFactory: def mktemp(self, name): return Path(tempfile.mkdtemp(prefix=name)) protocol = namespace["protocol"].__wrapped__(TempFactory()) batch = pa.record_batch( [pa.array(np.arange(131072, dtype=np.int64))], names=["value"] ) process = psutil.Process() samples = [] with namespace["_dataframe"]([batch] * 128, protocol=protocol) as (frame, session): for iteration in range(28): table = frame.toArrow() assert table.num_rows == 16777216 assert table.column(0)[0].as_py() == 0 assert table.column(0)[-1].as_py() == 131071 del table gc.collect() samples.append(process.memory_info().rss / 1048576) print(json.dumps({"warmup": 4, "post_gc_rss_mib": samples})) ``` </details> #### Remaining validation and follow-ups - [x] Run a matched master control for the repeated-result RSS probe. - [ ] Isolate the source of retention and investigate the historical high-RSS tail. - [ ] Rerun comparative benchmarks on the current head. - [ ] Full official parity suite and minimum-Rust build. Existing workflows use stable Rust; `pyspark-rs` declares 1.85 while the workspace still declares 1.82. Workspace MSRV cleanup is separate. - [ ] Review CI and outstanding baseline lint/API gaps. Workspace Clippy fails in unchanged core/binding code; the existing `SparkSession.Hook` gap remains. Local checks are not a green-CI claim. ### Was this patch authored or co-authored using generative AI tooling? Generated-by: OpenAI Codex (application version unavailable). This pull request and its description were written by Isaac. Closes #108 from dbtsai/codex/arrow-c-stream. Authored-by: DB Tsai <dbtsai@dbtsai.com> Signed-off-by: DB Tsai <dbtsai@dbtsai.com>
What changes were proposed in this pull request?
Replace the extra Rust Arrow IPC file -> Python bytes -> PyArrow decoding step in native
DataFrame.toArrow()with the Arrow C Data Interface. An unpublishedspark-python-arrowworkspace crate exposes one table-export entry point, used bypyspark-rs.This PR now contains only the Arrow export optimization. PyO3 stays on 0.28.3, matching upstream; the dependency upgrade is deferred in #107. Row/scalar conversion code is unchanged from upstream. Python-to-Rust Arrow imports, standalone array/type exports, and worker-side UDF conversion are not included.
Why are the changes needed?
Avoid reserializing already decoded Arrow batches and allocating another full Python bytes copy before building a PyArrow table. This reduces handoff time and temporary memory without changing the Spark Connect wire protocol. Tracked by SPARK-60586.
Does this PR introduce any user-facing change?
toArrow()shares compatible Arrow value and offset buffers with PyArrow while preserving batch boundaries, metadata, and buffer lifetimes.toPandas()uses this path but may still copy during pandas materialization.Zero-copy applies to the decoded Rust Arrow -> PyArrow handoff, not network IPC decoding, schema descriptors, Python rows, or pandas conversion. Sliced validity bitmaps may need realignment. NUL-containing field names/timezones retain the IPC fallback without repeating the RPC.
Inconsistent batch schemas, including metadata and dictionary ordering, now raise
ValueErrorinstead of being written under the first schema. Empty batch vectors retain the old empty-schema behavior. Seecrates/spark-python-arrow/README.mdfor the export contract.How was this patch tested?
Local validation on macOS arm64, Rust 1.99.1, with an optimized release/
abi3-py39wheel:wasm-udf.toArrow()smoke cases passed against a local Spark Connect server, covering nested/null/binary/date/timestamp values, numeric ranges, and empty results.Before/after benchmark
Used dbtsai/pyspark-connect-benchmark, comparing upstream
fbd6af5763bae5cd572ba4438026e675165ba70awith this Arrow-only change. Both wheels include upstream's recent plan/row sharing, use PyO3 0.28.3 and Arrow 58.3.0, and were built offline with identical release optimization (opt-level=3) and ABI settings.Python 3.12.14, PyArrow 26.0.0, pandas 2.3.3, NumPy 2.5.4 and psutil 7.2.2 match. These public client operations receive canned localhost gRPC responses, not real Spark execution. Peak RSS is sampled client-process resident memory, not allocation counts; it includes allocator/warmup retention and can miss brief spikes.
Initial run: 12 cases, six independent fresh-process samples per client/case. Median times and sampled peak RSS:
core-return-long-w50-b10000-toArrowreturn-buffer-134217728-toArrowcore-return-string-w4-b1000-toArrowcore-return-string-w4-b10000-toArrowreturn-nullable50-toArrowreturn-nested3-toArrowallocation-first-4096x512-toArrowcore-return-long-w50-b10000-toPandasreturn-nested3-toPandascore-return-long-w4-b1000-collectcore-return-string-w4-b1000-collectinput-pandas-explicit-10000x3A separate confirmation used 12 samples per client/case with baseline/candidate harness slots swapped:
toArrow()toArrow()toPandas()collect()controlcollect()controlThe two collection controls have overlapping timing ranges in both runs; the earlier PyO3-upgrade slowdown is not reproduced here. The large Arrow cases retain substantial speed and memory gains. Small nullable results and nested pandas conversion show little timing change, so this is not a claim that every operation gets faster.
All 396 samples across the three clients passed validation (216 initial + 180 confirmation). Rust request counters matched, and source/wheel identities were unchanged during both runs. The runs are reported separately, not pooled. Buffer-address/lifetime tests establish zero-copy sharing; benchmark timings alone do not.
Remaining draft checks
SparkSession.HookAPI gap remains outside this change.Was this patch authored or co-authored using generative AI tooling?
Generated-by: OpenAI Codex (application version unavailable).