This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-6191-67803a7a422c44de07af1e5d25c1dbeae8df68d4 in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
commit 02e84047a423a08debe3c6b6daab19fd9233427c Author: Andy Grove <[email protected]> AuthorDate: Thu Sep 24 15:41:51 2026 +0000 fix: read shuffle write buffer, spill limit and off-heap sizes in bytes (#6191) * fix: read shuffle write buffer, spill limit and off-heap sizes in bytes spark.comet.shuffle.native.writeBufferSize was declared in MiB but its value was sent to native code as a byte count, so the native shuffle writer ran with a 1-byte write buffer by default. Declare it in bytes with a 1 MiB default; a bare number keeps its meaning. Native code parsed spark.comet.maxTempDirectorySize and the boolean flags it reads from the raw session strings, so a size with a unit or an upper-case boolean silently fell back to the native default. Resolve every config that native code reads on the JVM before it crosses JNI, and document that the spill limit applies per native plan rather than per task. getMemoryConfig read spark.memory.offHeap.size as MiB, while Spark reads a bare number as bytes. Closes #6183, #6184, #6185. * docs: give a real example of a task running several native plans A Spark operator that Comet does not support does not split a stage into two native plans: Comet does not resume native execution above it. The native operators on either side of a union or a coalesce do run as separate plans in the same task. * Update spark/src/main/scala/org/apache/comet/CometConf.scala Co-authored-by: Matt Butrovich <[email protected]> --------- Co-authored-by: Matt Butrovich <[email protected]> --- .github/workflows/pr_build_linux.yml | 1 + .github/workflows/pr_build_macos.yml | 1 + docs/source/user-guide/latest/tuning.md | 14 ++-- native/core/src/execution/spark_config.rs | 4 ++ .../main/scala/org/apache/comet/CometConf.scala | 16 +++-- .../scala/org/apache/comet/CometExecIterator.scala | 22 ++++-- .../org/apache/comet/exec/CometExecSuite.scala | 27 ++++++-- .../spark/CometExecIteratorLifecycleSuite.scala | 14 ++++ .../shuffle/CometNativeShuffleWriterSuite.scala | 79 ++++++++++++++++++++++ 9 files changed, 155 insertions(+), 23 deletions(-) diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index 8abd6c9249..66a87ae913 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -507,6 +507,7 @@ jobs: org.apache.spark.sql.comet.execution.shuffle.CometCelebornShuffleReaderSuite org.apache.spark.sql.comet.execution.shuffle.CometCelebornShufflePlanningSuite org.apache.spark.sql.comet.execution.shuffle.CometNativeShuffleInputRDDSuite + org.apache.spark.sql.comet.execution.shuffle.CometNativeShuffleWriterSuite org.apache.spark.sql.comet.execution.shuffle.CometNativePositionalRoundRobinSuite org.apache.spark.sql.comet.execution.shuffle.CometDiskBlockWriterSuite org.apache.comet.exec.CometShuffleEncryptionSuite diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml index af6bbc48c8..8740931a5c 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.github/workflows/pr_build_macos.yml @@ -155,6 +155,7 @@ jobs: org.apache.spark.sql.comet.execution.shuffle.CometCelebornShuffleReaderSuite org.apache.spark.sql.comet.execution.shuffle.CometCelebornShufflePlanningSuite org.apache.spark.sql.comet.execution.shuffle.CometNativeShuffleInputRDDSuite + org.apache.spark.sql.comet.execution.shuffle.CometNativeShuffleWriterSuite org.apache.spark.sql.comet.execution.shuffle.CometNativePositionalRoundRobinSuite org.apache.spark.sql.comet.execution.shuffle.CometDiskBlockWriterSuite org.apache.comet.exec.CometShuffleEncryptionSuite diff --git a/docs/source/user-guide/latest/tuning.md b/docs/source/user-guide/latest/tuning.md index 2ebe8987c5..e1257012ec 100644 --- a/docs/source/user-guide/latest/tuning.md +++ b/docs/source/user-guide/latest/tuning.md @@ -283,12 +283,14 @@ flushes sorted spill files. It must not exceed `spark.comet.batchSize`. ### Limiting Spill Disk Usage -Native operators that spill to disk (aggregate, sort, shuffle) are collectively bounded by -`spark.comet.maxTempDirectorySize` (default 100 GB). The limit is applied per Spark task, so an -executor running `N` concurrent tasks may use up to `N` times this value on shared local disks. -If the limit is reached, further spills fail and the query errors out. Raise this on workloads -with large sort/aggregate/shuffle spills, or lower it to protect executors on shared disks -(remembering to divide by task concurrency to reason about the aggregate). +Native operators that spill to disk (aggregate, sort, shuffle) are bounded by +`spark.comet.maxTempDirectorySize` (default 100 GB). The operators of one Comet native plan share +the limit. A Spark task can run more than one native plan at a time, for example the native +operators on either side of a union or a coalesce, so an executor running `N` concurrent tasks may +use more than `N` times this value on shared local disks. If the limit is reached, further spills +fail and the query errors out. Raise this on workloads with large sort/aggregate/shuffle spills, or +lower it to protect executors on shared disks, remembering that the total across an executor is a +multiple of this value. ## Parquet Reader Tuning diff --git a/native/core/src/execution/spark_config.rs b/native/core/src/execution/spark_config.rs index 4c2811cb5d..7e5fc3c6ba 100644 --- a/native/core/src/execution/spark_config.rs +++ b/native/core/src/execution/spark_config.rs @@ -26,6 +26,10 @@ pub(crate) const COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED: &str = "spark.comet.parquet.rowFilterPushdown.enabled"; pub(crate) const SPARK_EXECUTOR_CORES: &str = "spark.executor.cores"; +/// Comet configs read through this trait must be resolved by the JVM first: +/// `CometExecIterator.serializeCometSQLConfs` sends booleans as `true` or `false` and sizes as a +/// bare byte count. A config missing from its list arrives exactly as the user wrote it, and a +/// value such as `10g` or `TRUE` then silently parses as the default. pub(crate) trait SparkConfig { fn get_bool(&self, name: &str) -> bool; fn get_u64(&self, name: &str, default_value: u64) -> u64; diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 9ee128e46f..a83bef6b45 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -727,9 +727,11 @@ object CometConf extends ShimCometConf { "shuffle data to disk. Larger values may improve write performance by reducing " + "the number of system calls, but will use more memory. " + "The default is 1MB which provides a good balance between performance and memory usage.") - .bytesConf(ByteUnit.MiB) - .checkValue(v => v > 0, "Write buffer size must be positive") - .createWithDefault(1) + .bytesConf(ByteUnit.BYTE) + .checkValue( + v => v > 0 && v <= Int.MaxValue, + s"Write buffer size must be between 1 and ${Int.MaxValue} bytes") + .createWithDefault(1024 * 1024) val COMET_SHUFFLE_JVM_PREFER_DICTIONARY_RATIO: ConfigEntry[Double] = conf( "spark.comet.shuffle.jvm.preferDictionary.ratio") @@ -1102,9 +1104,11 @@ object CometConf extends ShimCometConf { .category(CATEGORY_TUNING) .doc( "The maximum amount of data (in bytes) stored inside the temporary directories " + - "used by native operators when spilling. Applied per Spark task, so an executor " + - "running N concurrent tasks may use up to N times this value on shared local disks. " + - "Once the limit is reached, further spills will fail and the query will error out.") + "used by native operators when spilling. Applied to each Comet native plan " + + "separately, and a Spark task can run more than one native plan at a time, so an " + + "executor running N concurrent tasks may use more than N times this value on shared " + + "local disks. Once the limit is reached, further spills will fail and the query will " + + "error out.") .bytesConf(ByteUnit.BYTE) .createWithDefault(100L * 1024 * 1024 * 1024) // 100 GB diff --git a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala index 6a51cd05ea..7de6998c4e 100644 --- a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala +++ b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala @@ -597,12 +597,20 @@ object CometExecIterator extends Logging { val executorCores = numDriverOrExecutorCores(SparkEnv.get.conf) builder.putEntries("spark.executor.cores", executorCores.toString) - // Any Comet config that the native side reads must be added here manually. - // `cometSqlConfs` only carries values that were explicitly set, so defaults - // from `createWithDefault(...)` would otherwise not cross JNI. - builder.putEntries( - CometConf.COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED.key, - CometConf.COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED.get(SQLConf.get).toString) + // Any Comet config that the native side reads must be added here manually, resolved. + // `cometSqlConfs` only carries values that were explicitly set, exactly as they were + // written, so defaults from `createWithDefault(...)` would otherwise not cross JNI, and + // native code, which parses only a bare number or a lowercase boolean, would silently fall + // back to its own default for a value such as `10g` or `TRUE`. + Seq[ConfigEntry[_]]( + CometConf.COMET_DEBUG_ENABLED, + CometConf.COMET_DEBUG_MEMORY_ENABLED, + CometConf.COMET_EXPLAIN_NATIVE_ENABLED, + CometConf.COMET_MAX_TEMP_DIRECTORY_SIZE, + CometConf.COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED, + CometConf.COMET_TRACING_ENABLED).foreach { entry => + builder.putEntries(entry.key, entry.get(SQLConf.get).toString) + } builder.build().toByteArray } @@ -614,7 +622,7 @@ object CometExecIterator extends Logging { val offHeapMode = CometSparkSessionExtensions.isOffHeapEnabled(conf) if (offHeapMode) { // in off-heap mode, Comet uses unified memory management to share off-heap memory with Spark - val offHeapSize = ByteUnit.MiB.toBytes(conf.getSizeAsMb("spark.memory.offHeap.size")) + val offHeapSize = conf.getSizeAsBytes("spark.memory.offHeap.size") val memoryFraction = CometConf.COMET_OFFHEAP_MEMORY_POOL_FRACTION.get() val memoryLimit = (offHeapSize * memoryFraction).toLong val memoryLimitPerTask = (memoryLimit.toDouble * coresPerTask / numCores).toLong diff --git a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala index 387785dfc5..973260b69f 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala @@ -76,10 +76,6 @@ class CometExecSuite extends CometTestBase { ConfigMap.parseFrom(protobuf) } - // test not setting the config - val deserialized: ConfigMap = roundtrip - assert(null == deserialized.getEntriesMap.get(CometConf.COMET_EXPLAIN_NATIVE_ENABLED.key)) - // test explicitly setting the config for (value <- Seq("true", "false")) { withSQLConf(CometConf.COMET_EXPLAIN_NATIVE_ENABLED.key -> value) { @@ -90,6 +86,29 @@ class CometExecSuite extends CometTestBase { } } + test("SQLConf serde resolves the configs that native code parses") { + def entries = ConfigMap.parseFrom(CometExecIterator.serializeCometSQLConfs()).getEntriesMap + val flags = Seq( + CometConf.COMET_DEBUG_ENABLED, + CometConf.COMET_DEBUG_MEMORY_ENABLED, + CometConf.COMET_EXPLAIN_NATIVE_ENABLED, + CometConf.COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED, + CometConf.COMET_TRACING_ENABLED) + + // Native code parses only a bare byte count or a lowercase boolean and silently falls back + // to its own default otherwise, so these cross JNI resolved, defaults included. + val defaults = entries + assert(defaults.get(CometConf.COMET_MAX_TEMP_DIRECTORY_SIZE.key) == "107374182400") + flags.foreach(flag => assert(defaults.get(flag.key) == "false", flag.key)) + + withSQLConf( + (CometConf.COMET_MAX_TEMP_DIRECTORY_SIZE.key -> "10g") +: flags.map(_.key -> "TRUE"): _*) { + val resolved = entries + assert(resolved.get(CometConf.COMET_MAX_TEMP_DIRECTORY_SIZE.key) == "10737418240") + flags.foreach(flag => assert(resolved.get(flag.key) == "true", flag.key)) + } + } + test("sample without replacement") { withParquetTable((0 until 1000).map(i => (i, i + 1)), "tbl") { val df = sql("SELECT * FROM tbl").sample(withReplacement = false, fraction = 0.3, seed = 42) diff --git a/spark/src/test/scala/org/apache/spark/CometExecIteratorLifecycleSuite.scala b/spark/src/test/scala/org/apache/spark/CometExecIteratorLifecycleSuite.scala index c1f29c7447..6dad42b06e 100644 --- a/spark/src/test/scala/org/apache/spark/CometExecIteratorLifecycleSuite.scala +++ b/spark/src/test/scala/org/apache/spark/CometExecIteratorLifecycleSuite.scala @@ -289,6 +289,20 @@ class CometExecIteratorLifecycleSuite extends CometTestBase { assert(nativeMemoryLimitWarning(Array(100 * mib, reserved, 1L, 1L), reserved, limit).isEmpty) } + test("the memory pool limit reads a bare off-heap size as bytes, as Spark does") { + import CometExecIterator.getMemoryConfig + val fourGiB = 4L * 1024 * 1024 * 1024 + val offHeap = new SparkConf(false) + .set("spark.master", "local[4]") + .set("spark.memory.offHeap.enabled", "true") + assert( + getMemoryConfig(offHeap.clone.set("spark.memory.offHeap.size", "4294967296")).memoryLimit + == fourGiB) + assert( + getMemoryConfig(offHeap.clone.set("spark.memory.offHeap.size", "4g")).memoryLimit + == fourGiB) + } + test("the native memory limit is the off-heap size plus the memory overhead") { import CometExecIterator.nativeMemoryLimit val mib = 1024L * 1024 diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriterSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriterSuite.scala new file mode 100644 index 0000000000..8e72e37384 --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriterSuite.scala @@ -0,0 +1,79 @@ +/* + * 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 org.apache.spark.TaskContext +import org.apache.spark.sql.CometTestBase +import org.apache.spark.sql.vectorized.ColumnarBatch + +import org.apache.comet.CometConf + +/** + * Checks the plan [[CometNativeShuffleWriter]] hands to native code. Lives in the + * `execution.shuffle` package so it can call the `private[shuffle]` `buildUnifiedPlan`. + */ +class CometNativeShuffleWriterSuite extends CometTestBase { + + import testImplicits._ + + test("the write buffer size reaches native code in bytes") { + withSQLConf( + CometConf.COMET_SHUFFLE_MODE.key -> "native", + "spark.sql.adaptive.enabled" -> "false") { + withParquetTable((0 until 10).map(i => (i, s"row-$i")), "tbl") { + val exchange = sql("SELECT * FROM tbl") + .repartition(3, $"_1") + .queryExecution + .executedPlan + .collectFirst { case value: CometShuffleExchangeExec => value } + .getOrElse(fail("Expected a native Comet shuffle exchange")) + val dependency = exchange.shuffleDependency + .asInstanceOf[CometShuffleDependency[Int, ColumnarBatch, ColumnarBatch]] + val context: TaskContext = TaskContext.empty() + + // Native code sizes the writers for the data and spill files from this field, in bytes. + def writeBufferSize: Int = + new CometNativeShuffleWriter[Int, ColumnarBatch]( + dependency.nativeShuffleSpec.get, + dependency.outputPartitioning.get, + dependency.outputAttributes, + dependency.shuffleWriteMetrics, + dependency.numParts, + dependency.shuffleId, + context.taskAttemptId(), + context, + context.taskMetrics().shuffleWriteMetrics, + dependency.rangePartitionBounds) + .buildUnifiedPlan("unused.data") + .getShuffleWriter + .getWriteBufferSize + + assert(writeBufferSize == 1024 * 1024) + withSQLConf(CometConf.COMET_SHUFFLE_NATIVE_WRITE_BUFFER_SIZE.key -> "8m") { + assert(writeBufferSize == 8 * 1024 * 1024) + } + // A bare number is a byte count. + withSQLConf(CometConf.COMET_SHUFFLE_NATIVE_WRITE_BUFFER_SIZE.key -> "65536") { + assert(writeBufferSize == 65536) + } + } + } + } +} --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
