sunchao commented on code in PR #7258:
URL: https://github.com/apache/datafusion-comet/pull/7258#discussion_r4394917335
##########
spark/src/main/scala/org/apache/comet/Native.scala:
##########
@@ -289,6 +289,29 @@ class Native extends NativeBase {
schemaAddrs: Array[Long],
tracingEnabled: Boolean): Long
+ /** Own a bounded group of decoded local shuffle batches until export or
reader cleanup. */
+ @native def createShuffleBatchCoalescer(maxRows: Int): Long
+
+ /** Decode one frame and return the next output's column count, or -1 for
more input. */
+ @native def appendShuffleBatch(
+ handle: Long,
+ shuffleBlock: ByteBuffer,
+ length: Int,
+ numColumns: Int,
+ tracingEnabled: Boolean): Int
+
+ /** Return ready columns, flushing any pending group at EOF. */
+ @native def drainShuffleBatch(handle: Long, endOfInput: Boolean,
tracingEnabled: Boolean): Int
Review Comment:
Renamed the JNI declaration, symbol and iterator calls to `peekShuffleBatch`
in `b97daacb45`. `takeShuffleBatch` remains the operation that removes the next
output.
##########
spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala:
##########
@@ -73,6 +73,78 @@ class CometNativeShuffleSuite extends CometTestBase with
AdaptiveSparkPlanHelper
import testImplicits._
+ test("broadcast producer batching combines local native shuffle blocks with
and without AQE") {
Review Comment:
Extended the existing integration case in `b97daacb45` with nullable BIGINT,
BOOLEAN and BINARY columns, retaining the AQE ON/OFF, native build-side shuffle
and reduced-stream assertions. DECIMAL stays in a separate unsupported-schema
query so it does not disable batching for the supported-type case. I kept this
in the existing Scala test to avoid duplicating its route-sensitive setup in a
SQL fixture. The local decoder checks also moved to their own registered suite,
with an exact-order flush-plus-passthrough case.
--
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]
