sunchao commented on code in PR #7258: URL: https://github.com/apache/datafusion-comet/pull/7258#discussion_r4406375947
########## spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/NativeBatchDecoderIteratorSuite.scala: ########## @@ -0,0 +1,122 @@ +/* + * 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.spark.sql.comet.execution.shuffle + +import java.io.{ByteArrayInputStream, EOFException} + +import org.apache.spark.sql.CometTestBase +import org.apache.spark.sql.execution.metric.SQLMetrics + +import org.apache.comet.Native +import org.apache.comet.vector.NativeUtil + +class NativeBatchDecoderIteratorSuite extends CometTestBase { + + import NativeShuffleTestFrames.intFrame + + private final class TrackingInputStream(bytes: Array[Byte]) Review Comment: Moved `TrackingInputStream` into `NativeShuffleTestFrames` and reused it from both suites. Updated in `48e1024bfc`. ########## spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala: ########## @@ -73,6 +74,88 @@ class CometNativeShuffleSuite extends CometTestBase with AdaptiveSparkPlanHelper import testImplicits._ + test("broadcast producer batching combines local native shuffle blocks with and without AQE") { + withSQLConf( + CometConf.COMET_BATCH_SIZE.key -> "4096", + SQLConf.SHUFFLE_PARTITIONS.key -> "4", + SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "true", + SQLConf.COALESCE_PARTITIONS_MIN_PARTITION_NUM.key -> "1", + SQLConf.COALESCE_PARTITIONS_MIN_PARTITION_SIZE.key -> "1") { + for (adaptive <- Seq(false, true)) { + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString) { + val streams = scala.collection.mutable.ArrayBuffer.empty[Long] + Seq(false, true).foreach { enabled => + withSQLConf( + CometConf.COMET_BROADCAST_PRODUCER_COALESCE_ENABLED.key -> enabled.toString) { + val build = spark + .range(0, 512, 1, 16) + .selectExpr( + "CAST(id AS INT) AS k", + "CASE WHEN id % 3 = 0 THEN NULL ELSE concat('v', id) END AS v", + "CASE WHEN id % 5 = 0 THEN NULL ELSE id + 4294967296 END AS l", + "CASE WHEN id % 7 = 0 THEN NULL ELSE id % 2 = 0 END AS b", + "CASE WHEN id % 11 = 0 THEN NULL ELSE CAST(concat('b', id) AS BINARY) END AS bytes") + .repartition(col("k")) + val probe = spark.range(0, 512, 1, 2).selectExpr("CAST(id AS INT) AS k") + val df = probe.join(broadcast(build), "k") + checkAnswer( + df, + (0 until 512).map { i => + Row( + i, + if (i % 3 == 0) null else s"v$i", + if (i % 5 == 0) null else i.toLong + 4294967296L, + if (i % 7 == 0) null else i % 2 == 0, + if (i % 11 == 0) null else s"b$i".getBytes(UTF_8)) + }) + val plan = df.queryExecution.executedPlan + val exchanges = collect(plan) { case exchange: CometBroadcastExchangeExec => + exchange + } + assert(exchanges.size == 1, plan.treeString) + val exchange = exchanges.head + assert(exchange.metrics("numOutputRows").value == 512L) + assert( + collect(plan) { + case shuffle: CometShuffleExchangeExec + if shuffle.shuffleType == CometNativeShuffle => + shuffle + }.nonEmpty, + plan.treeString) + streams += exchange.metrics("numCoalescedBatches").value + () + } + } + assert(streams(1) > 0L && streams(1) < streams(0), s"AQE=$adaptive, batches=$streams") + } + } + } + } + + test("broadcast producer batching leaves unsupported schemas on the native shuffle reader") { Review Comment: Agreed: the DECIMAL test did not distinguish the routing gate from native passthrough. Removed it without extracting another production predicate solely for the test. Updated in `48e1024bfc`. ########## spark/src/test/scala/org/apache/spark/sql/comet/util/UtilsSuite.scala: ########## @@ -42,6 +42,33 @@ import org.apache.comet.vector.CometVector class UtilsSuite extends CometTestBase { + test("serializeBatches round trips binary values across output chunks") { Review Comment: Replaced this in the separate serializer PR #7310 with a one-row, 256-column case: vector data stays below 4 KiB, while varied field names make the compressed schema exceed that estimate. It asserts multiple 4 KiB output chunks and round-trips every value. The new test fails against the exact main serializer because it produces one chunk; all 16 UtilsSuite tests pass with the change. ########## spark/src/test/scala/org/apache/spark/sql/benchmark/CometBroadcastProducerBenchmark.scala: ########## @@ -0,0 +1,229 @@ +/* + * 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.spark.sql.benchmark + +import java.io.{ByteArrayInputStream, ByteArrayOutputStream} +import java.nio.{ByteBuffer, ByteOrder} +import java.nio.channels.Channels +import java.nio.charset.StandardCharsets + +import org.apache.arrow.memory.RootAllocator +import org.apache.arrow.vector.{BigIntVector, IntVector, VarCharVector, VectorSchemaRoot} +import org.apache.arrow.vector.ipc.ArrowStreamWriter +import org.apache.spark.benchmark.Benchmark +import org.apache.spark.sql.comet.execution.shuffle.NativeBatchDecoderIterator +import org.apache.spark.sql.comet.util.Utils +import org.apache.spark.sql.execution.metric.SQLMetric + +import org.apache.comet.Native +import org.apache.comet.vector.NativeUtil + +/** + * Measures native shuffle decoding, Arrow export and broadcast serialization together. Input + * frames are prepared before timing; output validation and driver decoding are also untimed. Both + * arms use the same artifact and differ only in producer batching. This is a producer + * microbenchmark, not a complete Spark query or a shuffle transport benchmark. + * + * Pass --on-first to reverse the case order, --decode-only to include the decoder/FFI cost + * without serialization, and an optional scenario-name substring to filter. + * + * To run: + * {{{ + * SPARK_GENERATE_BENCHMARK_FILES=1 make \ + * benchmark-org.apache.spark.sql.benchmark.CometBroadcastProducerBenchmark BENCH_HEAP=4g + * }}} + */ +object CometBroadcastProducerBenchmark extends CometBenchmarkBase { + + private val targetRows = 8192 + private val wideValue = "x" * (6 * 1024) + + private case class Scenario(name: String, rows: Int, rowsPerFrame: Int, wide: Boolean = false) + private case class Result(rows: Long, streams: Long, bytes: Long) + + private val scenarios = Seq( + Scenario("tiny: 1 row per frame", 8192, 1), + Scenario("tiny: 8 rows per frame", 8192, 8), + Scenario("full: 8192 rows per frame", 8192, 8192), + Scenario("byte cap: 64 rows of 6 KiB strings", 8192, 64, wide = true)) + + private def frames(scenario: Scenario): Array[Byte] = { + val allocator = new RootAllocator(Long.MaxValue) + val ids = new BigIntVector("id", allocator) + val values = new IntVector("value", allocator) + val strings = new VarCharVector("payload", allocator) + val root = + if (scenario.wide) VectorSchemaRoot.of(ids, values, strings) + else VectorSchemaRoot.of(ids, values) + val output = new ByteArrayOutputStream() + val payload = wideValue.getBytes(StandardCharsets.UTF_8) + try { + ids.allocateNew(scenario.rowsPerFrame) + values.allocateNew(scenario.rowsPerFrame) + if (scenario.wide) strings.allocateNew() + var start = 0 + while (start < scenario.rows) { + val count = math.min(scenario.rowsPerFrame, scenario.rows - start) + var row = 0 + while (row < count) { + val id = start + row + ids.setSafe(row, id.toLong) + if (id % 7 == 0) values.setNull(row) else values.setSafe(row, id * 3) + if (scenario.wide) strings.setSafe(row, payload) + row += 1 + } + root.setRowCount(count) + val ipc = new ByteArrayOutputStream() + val writer = new ArrowStreamWriter(root, null, Channels.newChannel(ipc)) + try { + writer.start() + writer.writeBatch() + writer.end() + } finally writer.close() + val bytes = ipc.toByteArray + val header = ByteBuffer Review Comment: The producer benchmark now reuses `NativeShuffleTestFrames.arrowFrame`, with the fixture scoped `private[sql]`. Shared decoder construction and the OFF/ON ordering are defined once, and both benchmark comments now describe what runs concisely. Measurement caveats remain in the PR description. Updated in `48e1024bfc`. -- 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]
