HyukjinKwon opened a new pull request, #141:
URL: https://github.com/apache/spark-connect-rust/pull/141
### What changes were proposed in this pull request?
**Reattachable execution** (`spark-connect-core`). Every DataFrame action of
the Rust API and of the drop-in uses it. It now matches
`ExecutePlanResponseReattachableIterator` of PySpark 4.3.0:
- The server sometimes ends the response stream before `ResultComplete`, on
purpose. It does so after
`spark.connect.execute.reattachable.senderMaxStreamDuration` (2m) or
`senderMaxStreamSize` (1g). The stream now reattaches from the last response
and keeps reading. Before, it took the end of the stream for the end of the
result.
- The stream sends `ReleaseExecute`:
- `ReleaseUntil` after each response it reads;
- `ReleaseAll` once the result is complete, the stream fails, or the
stream is dropped (for example an abandoned `toLocalIterator`).
Releases are sent in the background under the client's retry policy, and a
failure is only logged, as in the reference.
**Drop-in parity with PySpark 4.3.0**:
- `StructType.add()` updates the base `DataType` too. `mapInArrow`,
`mapInPandas`, the grouped `applyIn*` functions and `str()` read their schema
from it, so they missed added fields.
- `DataFrame.repartitionById(n, "col")` treats a string as a column name
(`F._to_col`). Before, it was a string literal.
- `RuntimeConf.get(key, None)` returns `None` for an unset key. An omitted
default still raises.
- `SparkSession.createDataFrame(pyarrow.Table)` is supported and converted
as the reference does it:
- **Names:** from a list `schema`, with extra columns named `_N`, or from
the table.
- **Types:** from a `StructType` or DDL `schema`, or from the Arrow schema
(`from_arrow_schema`, following `spark.sql.timestampType`).
- **Conversion:** timestamps are localized to the session time zone, the
table is cast with `to_arrow_schema`, and it's sent as an Arrow IPC stream.
- **Empty table:** a schema-only relation when a schema is given,
`CANNOT_INFER_EMPTY_SCHEMA` otherwise.
**Core additions:**
- `SparkSession::create_dataframe_from_arrow`: a `LocalRelation` from Arrow
IPC bytes.
- `RuntimeConf::get_many`: one config RPC for several keys, like the
reference's `get_config_dict`.
**CI:**
- The Rust e2e and coverage servers start with
`senderMaxStreamSize=1048576`. Every result over 1 MiB then spans several
streams, so the reattach path runs.
- The new live-server tests of the drop-in run in the parity job's existing
live-server step.
### Why are the changes needed?
- **Silently truncated results.**
- With a 1 MiB `senderMaxStreamSize`, `collect()` of `range(300000)`
returned 135,000 rows.
- With a 1s `senderMaxStreamDuration`, a 3-second query returned no rows,
and its `count()` returned 0.
- With the server defaults, this happens to any result that takes over 2
minutes or 1 GiB.
- **Executions are never released.**
- The server keeps every finished execution, with its request and up to 10
MiB of its responses, until `detachedTimeout` (5m) or session close.
- I found this running the Connect RDD of apache/spark#55888 on the
drop-in. After five `parallelize` calls of 10M elements, `spark.stop()` failed
with `Java heap space` on a 1 GB server.
- At session close, the server logs each remaining execution's whole
request. That logging is a separate Spark issue.
- **The four API gaps** turned up in the same work. The official parity
tests run the reference pyspark over our transport, so they never exercise the
drop-in's own classes. This PR adds live-server tests that do.
### Does this PR introduce _any_ user-facing change?
Yes:
- Results that were silently truncated are now complete.
- Finished executions no longer accumulate on the server.
- The four APIs above behave as in PySpark 4.3.0. Before:
- a `StructType` built with `add()` lost the added fields in
`mapInArrow`/`mapInPandas`;
- `repartitionById(n, "col")` failed analysis;
- `conf.get(key, None)` raised for an unset key;
- `createDataFrame(pyarrow.Table)` raised `TypeError`.
### How was this patch tested?
- **New Rust e2e tests** in `crates/spark-connect/tests/e2e_integration.rs`:
- a 300,000-row `collect` across stream ends;
- completed and dropped executions are released: `ReattachExecute` then
fails with `OPERATION_NOT_FOUND`.
All three fail on master (135,000 rows; the server still has the
operation) and pass with this change.
- **New `python/tests/test_dropin_e2e.py`** (11 live-server tests). They
pass both with this change and with the reference pyspark 4.3.0-rc3 client. 10
of them fail on master.
- **A new offline test** for `StructType.add()`.
- **Against a Spark 4.3.0-rc3 Connect server** (the connect-bundled RC3
distribution), with `senderMaxStreamSize=1048576`:
- all nine Rust e2e suites: 63 tests;
- `test_arrow_e2e.py` and `test_dropin_e2e.py`: 14 tests;
- a script that runs the fixed APIs through this client and through the
reference RC3 client. The results are identical in all 27 checks, except for
one pre-existing difference: the drop-in repeats the error class in the message
of server errors (`[SQL_CONF_NOT_FOUND] [SQL_CONF_NOT_FOUND] ...`).
- **Offline checks:**
- `cargo test -p apache-spark-connect -p apache-spark-connect-core
--release` (706 passed);
- `cargo test -p apache-spark-connect --release --features wasm-udf
--no-run`;
- `cargo fmt --all --check`;
- `test_dropin_offline.py` and `test_arrow_handoff.py` (301 passed);
- `scripts/audit_no_stubs.sh`.
- **`clippy`:** the only diagnostics on changed lines are `result_large_err`
on the new `py.detach(...).to_pyerr()` calls. The crate already has 119 of
those.
CI here still runs against Spark 4.2.0 until #99 lands. The reference code
that this mirrors is the same in 4.2.0 and 4.3.0-rc3: the release and reattach
logic of `reattach.py`, `conf.get`, `repartitionById`, and the `pyarrow.Table`
branch of `createDataFrame`.
Not in this PR:
- The reference re-sends `ExecutePlan` when `ReattachExecute` reports
`OPERATION_NOT_FOUND` before any response arrived. This PR doesn't add that.
- 4.3.0's per-RPC timeouts aren't added either.
- Like the list and pandas inputs, a `pyarrow.Table` is sent inline. The
reference uses a cached local relation above
`spark.sql.session.localRelationCacheThreshold`; that isn't implemented.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Opus 5.5
This pull request and its description were written by Isaac.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
