Skip to content

[WIP][CONNECT] Resume reattachable execution after clean EOF - #139

Open
WindsorWZZ wants to merge 5 commits into
apache:masterfrom
WindsorWZZ:windsorwzz/reattach-clean-eof
Open

[WIP][CONNECT] Resume reattachable execution after clean EOF#139
WindsorWZZ wants to merge 5 commits into
apache:masterfrom
WindsorWZZ:windsorwzz/reattach-clean-eof

Conversation

@WindsorWZZ

@WindsorWZZ WindsorWZZ commented Oct 29, 2026 •
edited
Loading

Copy link
Copy Markdown

What changes were proposed in this pull request?

Resume a reattachable ExecutePlan stream with ReattachExecute when the server closes it cleanly before ResultComplete. Reuse the last observed response ID and the existing retry handling.

Draft proposal: the public Apache SPARK issue key is pending.

Why are the changes needed?

Spark can close an individual response stream at its duration or size limit while the execution is still running. Treating that EOF as completion silently truncates results. The protocol requires the client to reattach until it receives ResultComplete.

Does this PR introduce any user-facing change?

Yes. Reattachable executions continue receiving results after a clean stream closure instead of returning incomplete results. Non-retryable errors and retry exhaustion still surface to the caller.

How was this patch tested?

  • Eight regression tests use real gRPC request/response framing over in-memory I/O; no Spark server or TCP port is required. They cover repeated EOFs, an empty initial stream, response-ID continuity, normal completion, transient stream errors, reattach RPC errors, and retry exhaustion.
  • Two additional unit tests cover graceful EOF with retries disabled, completion and stream-error classification.
  • Before the fix, five of the eight transport regressions fail; after the fix, all 105 core unit tests pass.
  • Checked-in-lockfile validation: 105 core unit tests and 392 client unit tests pass. All client integration targets compile with and without wasm-udf. The three existing core live-server tests remain ignored.
  • cargo fmt --all --check passes. Core Clippy (--all-targets --locked) completes successfully with existing warnings.

The existing Rust CI workflow runs these tests through cargo test -p apache-spark-connect -p apache-spark-connect-core --release.

Was this patch authored or co-authored using generative AI tooling?

Generated-by: OpenAI Codex (GPT-6)

This pull request and its description were written by Isaac.

WindsorWZZ and others added 5 commits October 29, 2026 08:11
@WindsorWZZ Isaac
Resume reattachable execution after clean stream EOF
c259779 
Treat a clean EOF before ResultComplete as a stream boundary and resume from the last response ID. Cover EOF, completion, transient failures, and retry exhaustion with in-memory gRPC regressions.

Co-authored-by: Isaac <no-reply@databricks.com>
@dbtsai
[SPARK-60586][CONNECT] Export Arrow tables through the Arrow C Stream…
d9defae 
… 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. apache#103 and apache#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 apache#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 apache#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 apache#108 from dbtsai/codex/arrow-c-stream.

Authored-by: DB Tsai <dbtsai@dbtsai.com>
Signed-off-by: DB Tsai <dbtsai@dbtsai.com>
@WindsorWZZ Isaac
Resume reattachable execution after clean stream EOF
d5d9a42 
Treat a clean EOF before ResultComplete as a stream boundary and resume from the last response ID. Cover EOF, completion, transient failures, and retry exhaustion with in-memory gRPC regressions.

Co-authored-by: Isaac <no-reply@databricks.com>
@WindsorWZZ Isaac
Cover graceful stream completion and error classification
b890a1f 
Use the explicit stream-end decision helper and cover graceful EOF with retries disabled, completion, and retryable/non-retryable status handling.

Co-authored-by: Isaac <no-reply@databricks.com>
@WindsorWZZ Isaac
Preserve published clean-EOF branch history
465013f 
Retain the updated stream-end helper and regression coverage while integrating the published branch ancestry.

Co-authored-by: Isaac <no-reply@databricks.com>
@WindsorWZZ
WindsorWZZ marked this pull request as ready for review October 30, 2026 07:12

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@WindsorWZZ @dbtsai