apache / apache/datafusion-comet
Allocate Comet's parquet reader buffers from ArrowUtils.rootAllocator to enable zero-copy PyArrow UDF runner
- Dominant language
- Scala
- Stars
- 1.3k
- Forks
- 373
- Avg merge
- 2d 6h
- Merged PRs (30d)
- 190
Description
## What is the problem the feature request solves?
PR #4234 replaces the row-based input path in `CometMapInBatchExec` with a `CometArrowPythonRunner` that consumes columnar batches directly. The end-to-end speedup on Spark 4.0 is **1.92x narrow / 1.79x mixed / 1.42x wide** (vs the prior row-path numbers of 1.32x / 1.32x / 1.09x).
The PR is _bulk-copy_, not true zero-copy. For every batch, Comet's buffers are `Unsafe.copyMemory`'d into Spark's allocator before the Arrow IPC stream is unloaded. The reason is that Comet's Parquet readers each construct their own independent `RootAllocator`, separate from `ArrowUtils.rootAllocator`. Arrow Java rejects cross-root transfers (`A buffer can only be associated between two allocators that share the same root`), so `TransferPair.transfer()` cannot move buffers from Comet's allocator into the Python runner's `VectorSchemaRoot`.
If Comet's reader allocators shared a root with `ArrowUtils.rootAllocator`, the bulk-copy step in `CometColumnarPythonInput.writeNextBatchToArrowStream` would collapse to a pointer-swap, and the wide-row speedup should climb further (the wide-row case is where the copy cost is most concentrated; narrow rows are bottlenecked on Python fork/IPC).
Related: #957, #4234.
## Describe the potential solution
### Change 1: allocator parent
The five files under `common/src/main/java/org/apache/comet/parquet/` that today declare `new RootAllocator()` would instead allocate as children of `ArrowUtils.rootAllocator`:
- `ColumnReader.java:55`
- `NativeColumnReader.java:41`
- `NativeBatchReader.java:154`
- `ArrowConstantColumnReader.java:46`
- `ArrowRowIndexColumnReader.java:39`
Plus anywhere else \`new RootAllocator()\` appears in the read path.
Result: every Comet buffer is in `ArrowUtils.rootAllocator`'s family. `makeTransferPair(target).transfer()` succeeds.
### Change 2: trait simplification (downstream)
`CometColumnarPythonInput.copyVector` collapses from a tree walk + per-buffer memcpy to a single \`TransferPair.transfer()\` call per column. `CometVectorIpcCopier` in \`common\` is no longer needed and can be deleted. About 100 lines of code go away.
### Coupled question: shading
`comet-common` currently shades `org.apache.arrow.*` into `org.apache.comet.shaded.arrow.*`. This was introduced so Comet's bundled Arrow (currently 18.x) doesn't conflict with Spark's bundled Arrow (~12 on Spark 3.4, ~17 on Spark 4.0). With the allocator change, Comet's reader code would need to call `ArrowUtils.rootAllocator()` which returns a Spark-classpath (unshaded) allocator. Options:
- **Un-shade `org.apache.arrow.memory.*`** in `common/pom.xml`. Memory classes are stable across Arrow versions, so the cross-version risk is contained. The vector / IPC packages can stay shaded if needed.
- **Bridge via reflection** at the alloc site. Works but pollutes the perf-critical alloc path.
Un-shading is the cleaner solution; the risk is bounded because we're only un-shading the memory subpackage, not the full Arrow tree.
### Coupled question: merge common and spark modules
@andygrove notes that the original justification for `common` being a separate module (Java Iceberg integration) no longer applies. Merging `common` into `spark` would naturally force a decision on shading:
- **Don't shade**: everything resolves against Spark's runtime Arrow. The allocator change above becomes trivial. Cost: Comet needs to be source-compatible with whatever Arrow API Spark ships on the oldest supported minor.
- **Shade everything**: not viable because the merged module references Spark's unshaded `BasePythonRunner` etc.
- **Partial shade**: same complexity as today, just in one module instead of two.
If we go the un-shade route on the merged module, the bulk-copy → zero-copy transition becomes a one-commit refactor. The trade is that Comet's code paths that use Arrow features unique to newer versions (currently 18.x) need backports or guards for older Spark versions.
This may be a bigger move than just the allocator change, but worth considering alongside since the two concerns intersect at exactly the shading boundary.
## Additional context
Catches that need attention either way:
1. **Reader lifetime.** Today each reader owns its allocator and closes it on `close()`. With shared-root allocators, ownership semantics change: closing a reader still frees its child allocator, but the root persists for the executor lifetime. Verify no reader code assumes `close()` deallocates the root.
2. **Native FFI handoff.** Comet's native code allocates Arrow buffers and exports them via FFI to JVM Arrow `c.*` classes (the relocation excludes the C bridge: `org.apache.arrow.c.jni.JniWrapper` etc.). The buffers come out of the native allocator; the JVM side wraps them in JVM Arrow vectors via FFI. Confirm the wrap uses the shared root so the chain remains consistent end-to-end.
3. **Cross-version Arrow API drift.** Spark 3.4 ships Arrow ~12, 4.0 ships ~17. If Comet uses Arrow 18-only APIs in the reader path, those would need backports for the old-Spark-on-shared-allocator world. Most of Comet's Arrow use is via the memory + vector primitives that have been stable since Arrow 10.x; targeted audit needed.
4. **Benchmark target.** Re-run `spark/src/test/resources/pyspark/benchmark_pyarrow_udf.py` after the change. Headline numbers to beat:
| workload | bulk-copy (current) | zero-copy (target) |
|---|---|---|
| narrow primitives | 1.92x | TBD; minor improvement expected (Python IPC dominates) |
| mixed with strings | 1.79x | TBD |
| wide rows (50 cols) | 1.42x | TBD; biggest improvement expected here |
Once this lands, the PR #4234 user guide can drop the "Not zero-copy; bulk per-buffer memcpy" limitation paragraph.
Contributor guide
Assessment
This issue has not been assessed yet.