andygrove commented on code in PR #5538:
URL: https://github.com/apache/datafusion-comet/pull/5538#discussion_r3886893616
##########
spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala:
##########
@@ -49,6 +50,8 @@ import org.apache.comet.shims.CometTypeShim
import org.apache.comet.vector.CometVector
object Utils extends CometTypeShim with Logging {
+ private val NATIVE_IPC_LZ4_PREFIX = Array[Byte](0x4c, 0x5a, 0x34, 0x5f)
Review Comment:
Should the direct-read format follow the configured codec rather than always
writing LZ4? Native `read_ipc_compressed` handles `SNAP`, `ZSTD` and `NONE` as
well as `LZ4_`, and `spark.comet.exec.shuffle.compression.codec` already exists
with a `zstd`/`lz4`/`snappy` check. As it stands someone who configured zstd
silently gets LZ4 for broadcast.
##########
spark/src/main/scala/org/apache/comet/CometConf.scala:
##########
@@ -229,6 +229,16 @@ object CometConf extends ShimCometConf {
createExecEnabledConfig("broadcastHashJoin", defaultValue = true)
val COMET_EXEC_BROADCAST_EXCHANGE_ENABLED: ConfigEntry[Boolean] =
createExecEnabledConfig("broadcastExchange", defaultValue = true)
+ val COMET_EXEC_BROADCAST_DIRECT_READ_ENABLED: ConfigEntry[Boolean] =
+ conf(s"$COMET_EXEC_CONFIG_PREFIX.broadcast.directRead.enabled")
Review Comment:
The shuffle equivalent is `spark.comet.shuffle.directRead.enabled`. #4986
deliberately moved these out of `spark.comet.exec.*` and kept the old name only
as an alternative, so this lands back under the prefix we are moving away from.
Should it be `spark.comet.broadcast.directRead.enabled`? Easier to settle now
than to add a deprecation alias once it has shipped.
##########
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 {
Review Comment:
Extending `CometShuffleBlockIterator` and handing it a throwaway
`ByteArrayInputStream` means the superclass constructor still allocates a 128KB
direct `ByteBuffer` and a 16-byte header buffer per instance, neither of which
this class ever touches, and there is one iterator per consumer partition.
Since every protocol method is overridden anyway, would it be cleaner to
extract a small interface for the JNI block protocol and have both classes
implement it? The native `comet_shuffle_block_iterator` bridge would keep
working, and the inheritance would stop implying shared state that is not there.
##########
spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala:
##########
@@ -332,12 +348,38 @@ object Utils extends CometTypeShim with Logging {
return Iterator.empty
}
- // 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(Channels.newChannel(compressedIpcInputStream(bytes)),
source)
+ }
+
+ private def compressedIpcOutputStream(
+ output: ChunkedByteBufferOutputStream,
+ nativeIpc: Boolean): DataOutputStream = {
+ if (nativeIpc) {
+ output.write(NATIVE_IPC_LZ4_PREFIX)
+ new DataOutputStream(new net.jpountz.lz4.LZ4FrameOutputStream(output))
+ } else {
+ val codec = CompressionCodec.createCodec(SparkEnv.get.conf)
+ new DataOutputStream(codec.compressedOutputStream(output))
+ }
+ }
+
+ /**
+ * Opens either the existing Spark-codec stream or the codec-prefixed
direct-read stream. This
+ * keeps non-native broadcast consumers compatible when direct read is
enabled.
+ */
+ private def compressedIpcInputStream(bytes: ChunkedByteBuffer): InputStream
= {
+ val input = new PushbackInputStream(bytes.toInputStream(),
NATIVE_IPC_LZ4_PREFIX.length)
+ val prefix = new Array[Byte](NATIVE_IPC_LZ4_PREFIX.length)
+ val read = input.read(prefix)
Review Comment:
`input.read(prefix)` can come back with fewer than 4 bytes, and a short read
here falls through to the Spark codec branch on what is actually a native-IPC
stream. I do not think it can happen in practice because the chunks are 1MB,
but `input.readNBytes(4)` would take the question off the table.
##########
native/core/src/execution/operators/shuffle_scan.rs:
##########
@@ -219,6 +275,14 @@ impl ShuffleScanExec {
.map(|col| unpack_dictionary(col))
.collect();
+ debug_assert_eq!(
Review Comment:
This looks redundant. `check_column_count` already returns a real error on
the non-validating path, `validate_remote_schema` covers the validating one,
and `cast_and_stamp_schema` checks again in `poll_next`. In release builds the
assert compiles away entirely, and it is the only reason `kind` has to be
threaded down into `get_next`. Can it come out?
##########
spark/src/main/scala/org/apache/spark/sql/comet/CometBroadcastExchangeExec.scala:
##########
@@ -64,7 +65,8 @@ case class CometBroadcastExchangeExec(
originalPlan: SparkPlan,
override val output: Seq[Attribute],
mode: BroadcastMode,
- override val child: SparkPlan)
+ override val child: SparkPlan,
+ directRead: Boolean =
CometConf.COMET_EXEC_BROADCAST_DIRECT_READ_ENABLED.get())
Review Comment:
Two questions about the conf-reading default here.
There are two existing construction sites that do not pass this parameter
and will silently inherit whatever the conf says at that moment:
`CometExecRule.scala:527`, which builds the exchange for
`CometSubqueryBroadcastExec`, and
`CometPlanAdaptiveDynamicPruningFilters.scala:279`, whose comment says the
fresh exchange has to have the same canonical form as the join's exchange so
AQE's stage cache produces a `ReusedExchangeExec`. Now that `doCanonicalize`
includes `directRead`, that reuse depends on a session conf resolving
identically at two different points in planning, one of which is on an AQE
thread. Dropping the default so the compiler forces each site to state its
value would make the coupling explicit. My DPP test passes today so this is
latent rather than live, but it is cheap to close off.
Separately, what does including `directRead` in `doCanonicalize` protect
against? The serialized format is self-describing and `decodeBatches` sniffs
the prefix, so two exchanges with different modes look reusable to me anyway.
##########
spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala:
##########
@@ -332,12 +348,38 @@ object Utils extends CometTypeShim with Logging {
return Iterator.empty
}
- // 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(Channels.newChannel(compressedIpcInputStream(bytes)),
source)
+ }
+
+ private def compressedIpcOutputStream(
+ output: ChunkedByteBufferOutputStream,
+ nativeIpc: Boolean): DataOutputStream = {
+ if (nativeIpc) {
+ output.write(NATIVE_IPC_LZ4_PREFIX)
+ new DataOutputStream(new net.jpountz.lz4.LZ4FrameOutputStream(output))
Review Comment:
`LZ4FrameOutputStream(output)` defaults to `BLOCKSIZE.SIZE_4MB`. I measured
allocation against lz4-java 1.8.0 and it comes to about 12.6MB of heap per
instance, versus about 135KB with `SIZE_64KB`. `serializeBatches` builds one of
these per `ColumnarBatch`, and Spark's `LZ4CompressionCodec` on the path this
replaces uses a 32KB block by default, so this is roughly a 90x jump in
per-batch allocation.
The scaladoc on `coalesceBroadcastBatches` just below describes 200K tiny
batches arriving from a 400-task, 500-partition shuffle, which is exactly the
shape that would hurt. Could we pass `LZ4FrameOutputStream.BLOCKSIZE.SIZE_64KB`
here? That fixes the read side at the same time, since `LZ4FrameInputStream`
sizes its buffers from the frame's `BD` byte.
##########
spark/src/main/scala/org/apache/spark/sql/comet/CometBroadcastExchangeExec.scala:
##########
@@ -111,7 +113,11 @@ case class CometBroadcastExchangeExec(
private def getByteArrayRdd(plan: SparkPlan): RDD[(Long, ChunkedByteBuffer)]
= {
plan.executeColumnar().mapPartitionsInternal { iter =>
- Utils.serializeBatches(iter)
+ if (directRead) {
+ Utils.serializeBroadcastBatches(iter)
Review Comment:
Reading `directRead` inside the lambda makes it capture the plan node.
`mapPartitionsInternal` does not run `sc.clean`, and I confirmed against the
compiled class that the closure is now
`$anonfun$getByteArrayRdd$1(CometBroadcastExchangeExec, Iterator)` where before
it captured nothing. That ships the whole exchange, including `originalPlan`
and the child subtree, with every build-side task.
Could you hoist `val useNativeIpc = directRead` above the closure and branch
on that instead?
--
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]
