viirya commented on PR #108:
URL:
https://github.com/apache/spark-connect-rust/pull/108#issuecomment-6462230186
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:
```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 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).
--
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]
