perf: replace CometBatchIterator FFI input path with the Arrow C Stream Interface#4572
Draft
mbutrovich wants to merge 65 commits into
Draft
perf: replace CometBatchIterator FFI input path with the Arrow C Stream Interface#4572mbutrovich wants to merge 65 commits into
mbutrovich wants to merge 65 commits into
Conversation
…nonical Arrow C Stream Interface (JVM Data.exportArrayStream <-> native ArrowArrayStreamReader), eliminating the per-batch FFI deep copy and the arrow_ffi_safe flag.
…le_localtablescan
# Conflicts: # spark/src/main/scala/org/apache/spark/sql/comet/operators.scala
# Conflicts: # spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala # spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala
… on the Scala side.
Contributor
Author
|
Ran TPC-H SF1000 with Comet on Spark 3.5.8: 
Looks like it further improves on #4507. |
# Conflicts: # spark/src/main/scala/org/apache/spark/sql/comet/CometExecRDD.scala # spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDD.scala # spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala # spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala # spark/src/main/scala/org/apache/spark/sql/comet/operators.scala
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Which issue does this PR close?
Closes #3770.
Rationale for this change
CometBatchIteratorplus a per-batch FFI deep copy, guarded by anarrow_ffi_safeflag, because the JVM could reuse or mutate a batch's buffers after handing it off. Every batch crossing the boundary was copied.What changes are included in this PR?
Iterator[ColumnarBatch]as anorg.apache.arrow.c.ArrowArrayStream(Data.exportArrayStream); native takes ownership viafrom_raw.CometBatchIterator.javaand thearrow_ffi_safeproto field/plumbing are removed.CometExecIterator/CometExecRDDnow pass anArray[Object]of already-exportedArrowArrayStream(orCometShuffleBlockIterator) slots instead ofCometBatchIterator.ArrowReaderimplementations bridging Spark data to Arrow:RowArrowReader(InternalRow),SparkColumnarArrowReader(non-Arrow SparkColumnarBatch),ColumnarBatchArrowReader(Arrow-backedColumnarBatch, with VSR ownership transfer).CometNativeArrowSourcetrait: an operator supplies one per-partition reader and gets both the JVM columnar path (doExecuteColumnar) and the native C Stream path (doExecuteAsArrowStream). Implemented byCometLocalTableScanExecandCometSparkToColumnarExec.AlignedArrowStreamReaderwraps arrow-rs's stream reader to align buffers per imported batch (the JVM exports 8-byte-aligned buffers, which trip arrow-rs's alignment assertion). This is a temporary workaround: upstream Callalign_buffers()infrom_ffi, remove redundant call fromarrow-pyarrowarrow-rs#10030 fixes it and ships in arrow 59.0.0, after which this reader can be dropped.scan.rsdrops the per-batch deep copy.reconcileStreamSchemaadvertises the truthful first-batch Arrow schema (not the consumer's declared types) so nativeScanExec's boundary cast fires; logs one deduped warning per type drift (e.g.width_bucketreturn-type drift).CometLocalTableScanExecnow mixes inDataTypeSupportand runs a schema-level fallback inconvert, so aLocalTableScanExecwhose schema carries a type with noArrowWritercoverage (Spark 4.1TimeType, intervals, etc.) falls back to Spark instead of failing at the boundary.NullTypeis allow-listed sinceArrowWriterhandles it.How are these changes tested?
CometExecSuite,CometShuffleSuite,ParquetReadSuite, the fuzz suites).CometArrowStreamSuitecovering stream export and schema reconciliation, added to the Linux and macOS PR build workflows.CometExecSuitecases forCometLocalTableScanExec: aTimeTypeschema-level fallback check (Spark 4.1+), plus two Arrow-buffer leak checks (project consumer andcollect_list) that fail via the allocator leak detector if the per-batch buffers leak.