sunchao commented on code in PR #5538:
URL: https://github.com/apache/datafusion-comet/pull/5538#discussion_r4043916900
##########
spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala:
##########
@@ -358,26 +400,36 @@ object Utils extends CometTypeShim with Logging {
* a single Arrow IPC stream.
*/
def coalesceBroadcastBatches(
- input: Iterator[ChunkedByteBuffer]): (Array[ChunkedByteBuffer], Long,
Long) = {
+ input: Iterator[ChunkedByteBuffer],
+ nativeIpc: Boolean = false): (Array[ChunkedByteBuffer], Long, Long) = {
val buffers = input.filterNot(_.size == 0).toArray
if (buffers.isEmpty) {
return (Array.empty, 0L, 0L)
}
+ val totalInputBytes = buffers.foldLeft(0L) { (total, buffer) =>
+ if (Long.MaxValue - total < buffer.size) Long.MaxValue else total +
buffer.size
+ }
+ if (shouldSkipDirectBroadcastCoalesce(nativeIpc, totalInputBytes)) {
Review Comment:
[P2] Check the coalesced output against the JNI block limit
The sum of the compressed inputs is not an upper bound on the compressed
output: concatenating Arrow columns can remove the short-distance matches LZ4
used within each input batch. With direct read enabled, a large hinted
broadcast can therefore pass this guard and still produce a single block that
`CometBroadcastBlockIterator.hasNext()` rejects.
I reproduced this with the same Arrow 18.3 / LZ4 appender and writer
operations on correlated BIGINT columns: the separate IPC streams total
**551,523,846 bytes**, but the coalesced stream is **2,147,620,180 bytes**,
above `Integer.MAX_VALUE`. The row count and resulting broadcast are below the
existing row and default 8 GiB broadcast limits. The unchanged iterator rejects
that measured size. This is a component reproduction, not a full Spark query
run.
Please also validate the actual compressed output size and return the
original blocks when it exceeds the JNI limit, or emit bounded independent IPC
batches. The current test only exercises the input-size predicate, so it misses
this case.
##########
spark/src/main/java/org/apache/comet/CometBroadcastBlockIterator.java:
##########
@@ -0,0 +1,109 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.comet;
+
+import java.io.ByteArrayInputStream;
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.Iterator;
+
+/** Provides codec-prefixed Arrow IPC broadcast blocks to native code via JNI.
*/
+public final class CometBroadcastBlockIterator extends
CometShuffleBlockIterator {
+
+ private static final int INITIAL_BUFFER_SIZE = 128 * 1024;
+
+ private Iterator<ByteBuffer[]> blocks;
+ private ByteBuffer dataBuf = ByteBuffer.allocateDirect(INITIAL_BUFFER_SIZE);
+ private boolean closed = false;
+ private int currentBlockLength = 0;
+
+ public CometBroadcastBlockIterator(Iterator<ByteBuffer[]> blocks) {
+ // Native uses the same block-iterator JNI protocol for shuffle and
broadcast inputs. The
+ // superclass stream is never read because every protocol method is
overridden here.
+ super(new ByteArrayInputStream(new byte[0]));
+ this.blocks = blocks;
+ }
+
+ @Override
+ public int hasNext() throws IOException {
+ if (closed) {
+ return -1;
+ }
+
+ ByteBuffer[] block = null;
+ long blockSize = 0;
+ while (blocks.hasNext() && block == null) {
+ ByteBuffer[] candidate = blocks.next();
+ long candidateSize = 0;
+ for (ByteBuffer chunk : candidate) {
+ candidateSize += chunk.remaining();
+ }
+ if (candidateSize > 0) {
+ block = candidate;
+ blockSize = candidateSize;
+ }
+ }
+ if (block == null) {
+ close();
+ return -1;
+ }
+
+ if (blockSize > Integer.MAX_VALUE) {
+ throw new IllegalStateException(
+ "Native broadcast block size of "
+ + blockSize
+ + " exceeds the direct-read maximum of "
+ + Integer.MAX_VALUE
+ + " bytes");
+ }
+
+ currentBlockLength = (int) blockSize;
+ if (dataBuf.capacity() < currentBlockLength) {
+ long doubled = Math.max((long) dataBuf.capacity() * 2L, blockSize);
+ dataBuf = ByteBuffer.allocateDirect((int) Math.min(doubled,
Integer.MAX_VALUE));
+ }
+
+ dataBuf.clear();
+ dataBuf.limit(currentBlockLength);
+ for (ByteBuffer chunk : block) {
+ dataBuf.put(chunk.duplicate());
+ }
+ return currentBlockLength;
+ }
+
+ @Override
+ public ByteBuffer getBuffer() {
+ return dataBuf;
+ }
+
+ @Override
+ public int getCurrentBlockLength() {
+ return currentBlockLength;
+ }
+
+ @Override
+ public void close() throws IOException {
+ if (!closed) {
+ closed = true;
+ blocks = null;
Review Comment:
[P2] Drop the direct broadcast buffer when the iterator closes
`close()` clears the source blocks but retains this class's `dataBuf`. After
a coalesced broadcast is consumed, that buffer holds a complete compressed copy
of the build relation per consumer task. `CometExecIterator` retains the block
iterator while the join continues probing, so this direct allocation stays live
alongside the decoded build data even after EOF; `super.close()` cannot clear
the subclass buffer.
Compiling the unchanged iterator classes and feeding a 32 MiB block leaves
**33,685,504 bytes** of direct memory after EOF and forced GC. Explicit close
has the same result. Dropping only the subclass's `dataBuf` reference, while
keeping the closed iterator alive, reclaims exactly **33,554,432 bytes**. This
verifies retention, not an end-to-end OOM.
Please drop the staging-buffer reference on close/EOF and cover buffer
growth followed by close. Native decoding finishes consuming the compressed
bytes before the next `hasNext()` call, so the buffer is no longer needed at
EOF.
--
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]
