mbutrovich commented on code in PR #7016:
URL: https://github.com/apache/datafusion-comet/pull/7016#discussion_r4338766221
##########
spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala:
##########
@@ -341,10 +341,65 @@ object Utils extends CometTypeShim with Logging {
// use Spark's compression codec (LZ4 by default) and not Comet's
compression
val codec = CompressionCodec.createCodec(SparkEnv.get.conf)
- val cbbis = bytes.toInputStream()
- val ins = new DataInputStream(codec.compressedInputStream(cbbis))
// batches are in Arrow IPC format
- new ArrowReaderIterator(Channels.newChannel(ins), source)
+ new ArrowReaderIterator(ipcChannel(bytes, codec), source)
+ }
+
+ /**
+ * A channel over the Arrow IPC stream in `bytes`, which
[[serializeBatches]] wrote through
+ * `codec`.
+ */
+ def ipcChannel(bytes: ChunkedByteBuffer, codec: CompressionCodec):
ReadableByteChannel =
+ new InputStreamChannel(codec.compressedInputStream(bytes.toInputStream()))
+
+ /** Largest piece that [[InputStreamChannel]] reads from its stream at a
time. */
+ private[util] val MaxReadPiece = 64 * 1024
+
+ /**
+ * Reads `in` into each buffer it is asked to fill until the buffer is full
or `in` ends, up to
+ * [[MaxReadPiece]] bytes at a time through one reused array.
`Channels.newChannel` reads 8 KiB
+ * at a time, and enters and leaves the JDK's interruptible I/O machinery
for each piece. Like
+ * that channel, this one closes and throws `ClosedByInterruptException`
before it reads a piece
+ * on an interrupted thread, so a cancelled decode stops within a piece of
where it was.
+ */
+ private final class InputStreamChannel(in: InputStream) extends
ReadableByteChannel {
Review Comment:
What do you think about moving `InputStreamChannel` to a PR of its own? The
sharing doesn't depend on it, since `DecodedBroadcastCache.decode` could read
through `Channels.newChannel` as `coalesceBroadcastBatches` did.
The evidence for it so far is one timing in the description: a single decode
of the 1M-row broadcast on the driver, 16.4 ms through `Channels.newChannel`
against 16.0 ms through this channel. The profile in #6806 puts 8% of a probe
task in the `Channels.newChannel` copy loop, but this channel still copies each
piece from a heap array into the destination buffer ([line
388](https://github.com/apache/datafusion-comet/blob/04f329d3f49908dddb4e46f509c21ce6e313bd48/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala#L388)),
as the JDK's does. What it saves is the per-piece `begin()`/`end()` and the
smaller pieces.
On the other side, it reimplements the JDK's interrupt behavior, and that
took a second commit (733a74b3c) to match: the first version checked the
interrupt once per `read`, and Arrow reads a whole coalesced message body in
one `read`. It also sits on the path that
`spark.comet.exec.broadcastExchange.shareDecoded.enabled=false` takes, so
turning the setting off doesn't give back the decode path of `branch-1.1`.
In its own PR it could come with a benchmark of the per-task decode with
sharing off, which is where a faster channel would show up.
--
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]
