[dbtsai]

What changes were proposed in this pull request?

Under 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>:

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 and names, 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, 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.

RSS control methodology and reproducer

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.

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}))

Remaining validation and follow-ups

• 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.

[viirya]

The core of this version is the same as arrow-rs's own reader export. arrow::pyarrow implements IntoPyArrow for Box<dyn RecordBatchReader + Send>: it wraps the reader in FFI_ArrowArrayStream::new and calls pa.RecordBatchReader._import_from_c. That is exactly what record_batches_to_pyarrow does here before read_all(), so the main path could be:

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 rest of bridge.rs is there for inputs Spark doesn't send. ExportStream and the four unsafe extern "C" callbacks (stream_schema/stream_next/stream_error/stream_release) exist only to replace the stream's schema export so that sorted-map flags survive. Spark always writes maps with keysSorted = false, and Connect results have no dictionary-encoded fields, so neither the sorted-map repair nor the dictionary-order comparison in schemas_equal can be triggered. As it stands, this version adds a hand-written C Stream implementation in unsafe Rust (private-data ownership, error-string lifetimes, errno returns) plus unit tests for it, which #103 didn't need. That is a lot to maintain for a case that can't happen, and the upstream bug is already fixed by apache/arrow-rs#11630.

The upstream export does leave two things to handle here, and both are small:

• RecordBatchIterator doesn't check batches against the declared schema, so the inconsistent-schema ValueError needs a plain batch.schema() != schema check before export.
• The NUL name/timezone IPC fallback still applies.

If the schema preflight for older PyArrow is needed, schema.to_pyarrow(py)? does the same thing in one line.

On performance, the benchmarks here show the stream path ahead only when there are many small batches (about 5% at 1049 batches) and slightly behind at 128 MiB. Whatever gain there is comes from importing all batches in one call, which the upstream reader export also does. I'd suggest whichever variant moves forward use the upstream reader export, the schema check and the NUL fallback, kept in pyspark-rs. That would drop the separate crate, its workflow, and most of the tests (the sorted-map and dictionary-order tests, and the stream callback tests).

[dbtsai]

Thanks @viirya, implemented your suggestion in d3eabc8.

The handoff now uses Box<dyn RecordBatchReader + Send>::into_pyarrow() followed by read_all(), directly in pyspark-rs. I removed the custom callbacks, separate converter crate/workflow, and generic ordering/callback tests. The schema consistency check, NUL name/timezone IPC fallback, and schema.to_pyarrow(py)? preflight remain. PyO3 stays at 0.28.3.

The updated wheel passed 294 tests on each of Python 3.12/PyArrow 26 and Python 3.9/PyArrow 10, including buffer-sharing/lifetime checks. All 360 samples in the fresh C Data / previous C Stream / upstream-reader comparison passed. Performance was broadly comparable, rather than an established general speedup; cold-wide medians were essentially unchanged.

One memory caveat remains: a 128 MiB workload sample reached 285.7 MiB RSS with the upstream reader versus a 239.6 MiB maximum for the previous stream in the same round, despite close medians. That sample passed validation and is retained; the cause is unresolved. These are shared-desktop, canned-localhost receive/conversion measurements, not Spark execution.

I've updated the description with the results and limitations. Keeping this draft while the RSS tail, full official parity, and minimum-Rust validation remain open. This is substantially simpler to maintain—thanks for pointing out the upstream reader path.

[viirya]

Thanks for the rewrite. This is the shape I was hoping for: about 100 lines in pyspark-rs, with stream export, ownership and release left to arrow-rs's IntoPyArrow for Box<dyn RecordBatchReader + Send>. There is no custom unsafe code, and PyO3 stays at a single version. The NUL pre-check is the right call. Arrow 58.3 panics in CString::new(...).unwrap() and does not return an error, so the "try the export and fall back on error" approach I suggested on #103 would not have worked. Checking up front is the correct way to handle it.

A few remaining items, mostly about making sure the claims stay tested:

• Run test_arrow_e2e.py in CI. As it stands nothing runs it, and the official parity gate drives official pyspark over our transport, so it never reaches the Rust toArrow. The parity job in build_python_connect.yml already starts a Connect server and builds the extension into ./python, so one more step there would cover it:

  cd python && SPARK_REMOTE=sc://localhost:15002 PYTHONPATH=$(pwd) python -m pytest tests/test_arrow_e2e.py
  

• Nothing pins the new path. See the inline comment on test_handoff_preserves_values_metadata_and_chunks.
• RSS growth across repeated calls. Could you run the same 28-call probe on master? If master shows the same 416 → 481 MiB growth, it comes from the allocator or the mock server and not from this change, and the open question is settled.

[viirya]

Could this say where the panic happens? With the stream export, the schema is converted inside the extern "C" get_schema callback, which has no catch_unwind, so the panic aborts the Python process and does not surface as an exception. That is why the check has to stay even if someone later considers catching the panic instead. Something like "panics inside the C stream get_schema callback, which aborts the process" would help.

[viirya]

Apart from the mismatched-schema cases, these handoff tests pass unchanged against master's IPC implementation, so a later change that routed toArrow back through IPC would go unnoticed. The docstring on the buffer-lifetime test acknowledges it can't prove Rust-to-Python sharing either. One cheap guard is to monkeypatch pyarrow.ipc.open_file to raise in the non-NUL tests, or in a dedicated test. That asserts the property this PR exists for, and the NUL tests then show the fallback is the only path that uses IPC.

[viirya]

The coverage workflow installs protobuf-compiler today, but if that ever changes, all of these tests would skip and the job would stay green. Could this fail instead of skipping when running in CI (e.g. when CI is set)?

[tdcmeehan]

These cases reach the fallback through a top-level name, a struct child and a top-level timezone, but none needs the List/Map arm of data_type_needs_ipc (pyarrow_export.rs L90-95) to find the NUL. If that arm regressed, a list of structs, or a map whose value struct has a NUL field name, would reach CString::new(name).unwrap() (arrow-schema ffi.rs#L174) inside the stream's extern "C" get_schema callback (ffi_stream.rs#L127, #L220), and the process would abort instead of raising. Could kind add a list-of-struct case and a map-value-struct case?

[tdcmeehan]

The description says nullability differences are rejected too, and Field::eq does compare nullable (field.rs#L108-L115), but none of these cases changes it, so a comparison that ignored nullability would still pass all eight. Could difference add a case where the second batch's field is nullable=False?

[tdcmeehan]

Carrying over from #103: pyproject.toml still declares pyarrow>=10.0.0 (L30), but this step is the only CI run of the handoff tests, and it gets the latest PyArrow from L93 (26.0.0 in this head's run). The preflight for older PyArrow is gone, and toArrow() now goes through PyArrow's private RecordBatchReader._import_from_c (arrow-pyarrow lib.rs#L460-L466) instead of the public pa.ipc.open_file, so the local PyArrow 10 run is the only evidence for the floor. Could CI add a leg pinned to pyarrow==10.0.0? If that isn't practical, could the docs state which PyArrow versions are tested?

[tdcmeehan]

Nit: numeric and strings keep their buffers alive for the whole test, so these addresses can't change and the two asserts can't fail. The to_pylist() checks after del table; gc.collect() (L226-227) are what exercise the lifetime. Could these lines and addresses (L223) be dropped?

[dbtsai]

Thanks @viirya and @tdcmeehan. Addressed the regression-test and documentation items in e60aa53:

• Added a normal-handoff guard that makes pyarrow.ipc.open_file fail. It passes on this revision and fails against master d884e34, so returning to the old IPC handoff is detected. The NUL tests separately assert that IPC is used exactly once.
• Added list-of-struct and map-value-struct NUL-name cases, plus nullability mismatch cases with and without NUL names.
• Missing protoc now fails when CI is set and skips locally; two tests cover those outcomes.
• Clarified that the NUL panic occurs inside the C Stream get_schema callback and aborts Python, rather than producing a Python exception.
• Removed the redundant buffer-address assertions; retained the value checks after dropping the table/session and forcing GC.
• Documented the tested combinations in docs/architecture.md: PyArrow 10.0.0/Python 3.9 and PyArrow 26.0.0/Python 3.12. CI currently tests the latest PyArrow, not a pinned minimum.

I also ran the identical 28-call RSS probe on rebuilt master and this revision in fresh processes, with the same 128 MiB input, environment and four-call warmup. Post-GC RSS on master rose 1129.23 → 3976.06 MiB; C Stream rose 418.81 → 569.47 MiB. All calls passed validation. Growth is not unique to this PR, but the combined mock-server/client process still prevents attributing it to the client, server, or allocator. I am not claiming flat RSS or leak freedom. The methodology and reproducer are committed in dev/design/arrow-handoff-rss.md.

Validation: 299 passed, 1 skipped on each of PyArrow 10 and 26, including 26 handoff/fixture cases; 392 Rust unit tests passed; all integrations compiled with and without wasm-udf; 3 native live-server smoke tests passed; official collection parity 15 passed, 1 skipped. Known baseline Clippy failures and the existing SparkSession.Hook API gap remain.

Still open: wiring test_arrow_e2e.py into CI. Per the author's request to leave .github/workflows unchanged, this follow-up documents those tests as manual-only rather than claiming that the official parity gate covers native toArrow(). The live-server tests were rerun locally; this does not close the CI-wiring request. The PyArrow-version documentation follows the alternative tdcmeehan suggested to a pinned CI job.

[dbtsai]

Follow-up to the remaining CI-wiring item: pushed e0eb002, which adds test_arrow_e2e.py to the Connect parity job in .github/workflows/build_python_connect.yml.

The step uses the already-built native extension and existing live server, with SPARK_REMOTE=sc://localhost:15002 and PYTHONPATH pointing at this repository's python directory. It runs after the official suite and before the local-cluster restart, preserving the pristine-server start for official parity. These three tests now directly exercise the native toArrow() path in CI instead of relying on the official-PySpark transport gate.

Validated workflow YAML parsing and ran the exact pytest command locally against Spark Connect: 3 passed. The architecture docs and PR description now reflect the CI wiring. This addresses the previously open wiring request; the new CI run still needs to finish. PyArrow 10 remains covered by the documented local runs, rather than a pinned CI job.

[HyukjinKwon]

Review summary

I checked the production handoff in crates/pyspark-rs/src/pyarrow_export.rs and found it consistent with the stated design. The exact batch-schema check (metadata and nullability included) runs before both export paths. The recursive NUL guard covers field names at every nesting depth and timestamp timezones, and falls back to IPC over the already-collected batches with no second ExecutePlan. The default path relies on arrow-rs's own IntoPyArrow reader export, with no custom unsafe code. The tests pin both routes: pyarrow.ipc.open_file is made to fail on the normal path, and the NUL cases assert exactly one IPC call. They also cover the list-of-struct, map-value-struct and nullability-only cases raised earlier. The live-server smoke tests now run in the Connect parity job.

The one remaining item is in the new RSS design note: its reproduction steps point PYTHONPATH at the wrong directory, so rerunning the master-versus-PR control as written may not load the build under test. Details are inline.

Findings

1 total: 0 P0, 0 P1, 0 P2, 1 P3.

Nit (P3)

• RSS reproducer names the wrong PYTHONPATH directory — dev/design/arrow-handoff-rss.md:31 — see inline.

[HyukjinKwon]

Nit (P3): Followed literally, these two instructions don't load the extension under test. The harness does from pyspark.sql import SparkSession, so PYTHONPATH has to be the directory that contains the pyspark package, i.e. <checkout>/python. That's what the new parity step uses (PYTHONPATH: ${{ github.workspace }}/python). "The Python package containing the extension" reads as python/pyspark, and "point PYTHONPATH at that checkout" (next line) is the checkout root, which has no top-level pyspark. In both cases the import fails or silently falls back to another installed pyspark, and the master-versus-PR RSS control can measure the wrong build. Could both lines say <PR checkout>/python and <master checkout>/python explicitly? The working-directory note below can stay as is.

Verification:

• Inspection: The reproduction instructions name the checkout's python directory, the one containing the pyspark package, as PYTHONPATH for both the PR build and the master control, and no longer refer to the package directory itself or the checkout root.

[dbtsai]

Thanks @HyukjinKwon . cc @viirya @tdcmeehan to see if your questions were all addressed.

[viirya]

Thanks for working through all of these. The IPC guard and the exactly-once IPC assertion on the NUL cases now pin both routes, and the live-server smoke tests are passing in the parity job (3 passed in the latest run). The master control settles the RSS question: master grew much more in the same probe, so the growth isn't coming from this change.

Two small doc-placement nits inline. Neither blocks merging.

[viirya]

This reads as a one-off measurement record for this PR (dated, single runs on a shared desktop) rather than design documentation. The reproducer also depends on test internals (runpy.run_path on the test file, protocol.__wrapped__, _dataframe), so it will quietly stop working when the tests change, and nothing checks it. Could the numbers and the script live in the PR description instead of the repo?

[viirya]

docs/ is the published user site, and this paragraph is maintainer information (which PyArrow versions were tested locally, which CI job runs the smoke tests, what the parity gate covers). Could it move to dev/design/architecture.md? There's also an extra blank line after it (L121-122).