You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by ka...@apache.org on 2025/03/21 19:51:04 UTC
(datafusion-comet) branch main updated: chore: Fix some inconsistencies in memory pool configuration (#1561)
This is an automated email from the ASF dual-hosted git repository.
kazuyukitanimura pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
The following commit(s) were added to refs/heads/main by this push:
new badbd3768 chore: Fix some inconsistencies in memory pool configuration (#1561)
badbd3768 is described below
commit badbd376898cb4f80b5136a90697019ad11a2e90
Author: Andy Grove <ag...@apache.org>
AuthorDate: Fri Mar 21 13:51:00 2025 -0600
chore: Fix some inconsistencies in memory pool configuration (#1561)
## Which issue does this PR close?
Closes https://github.com/apache/datafusion-comet/issues/1560
## Rationale for this change
- Fix some mistakes I made in https://github.com/apache/datafusion-comet/pull/1525
- Make some changes to `fair_unified` pool memory calculations
---
.../main/scala/org/apache/comet/CometConf.scala | 25 ++++----
docs/source/user-guide/configs.md | 8 +--
docs/source/user-guide/tuning.md | 17 ++---
native/core/src/execution/jni_api.rs | 14 +++--
.../scala/org/apache/comet/CometExecIterator.scala | 19 ++++--
.../apache/comet/CometSparkSessionExtensions.scala | 70 ++++++++++++++-------
.../src/main/scala/org/apache/spark/Plugins.scala | 27 ++------
.../comet/CometSparkSessionExtensionsSuite.scala | 72 +++++++++++++++-------
8 files changed, 147 insertions(+), 105 deletions(-)
diff --git a/common/src/main/scala/org/apache/comet/CometConf.scala b/common/src/main/scala/org/apache/comet/CometConf.scala
index f354f3d15..b1b25ee09 100644
--- a/common/src/main/scala/org/apache/comet/CometConf.scala
+++ b/common/src/main/scala/org/apache/comet/CometConf.scala
@@ -236,7 +236,7 @@ object CometConf extends ShimCometConf {
val COMET_MEMORY_OVERHEAD: OptionalConfigEntry[Long] = conf("spark.comet.memoryOverhead")
.doc(
"The amount of additional memory to be allocated per executor process for Comet, in MiB, " +
- "when running in on-heap mode or when using the `fair_unified` pool in off-heap mode. " +
+ "when running Spark in on-heap mode. " +
"This config is optional. If this is not specified, it will be set to " +
s"`spark.comet.memory.overhead.factor` * `spark.executor.memory`. $TUNING_GUIDE.")
.bytesConf(ByteUnit.MiB)
@@ -244,9 +244,10 @@ object CometConf extends ShimCometConf {
val COMET_MEMORY_OVERHEAD_FACTOR: ConfigEntry[Double] =
conf("spark.comet.memory.overhead.factor")
- .doc("Fraction of executor memory to be allocated as additional memory for Comet " +
- "when running in on-heap mode or when using the `fair_unified` pool in off-heap mode. " +
- s"$TUNING_GUIDE.")
+ .doc(
+ "Fraction of executor memory to be allocated as additional memory for Comet " +
+ "when running Spark in on-heap mode. " +
+ s"$TUNING_GUIDE.")
.doubleConf
.checkValue(
factor => factor > 0,
@@ -255,8 +256,7 @@ object CometConf extends ShimCometConf {
val COMET_MEMORY_OVERHEAD_MIN_MIB: ConfigEntry[Long] = conf("spark.comet.memory.overhead.min")
.doc("Minimum amount of additional memory to be allocated per executor process for Comet, " +
- "in MiB, when running in on-heap mode or when using the `fair_unified` pool in off-heap " +
- s"mode. $TUNING_GUIDE.")
+ s"in MiB, when running Spark in on-heap mode. $TUNING_GUIDE.")
.bytesConf(ByteUnit.MiB)
.checkValue(
_ >= 0,
@@ -485,13 +485,14 @@ object CometConf extends ShimCometConf {
.createWithDefault(false)
val COMET_EXEC_MEMORY_POOL_TYPE: ConfigEntry[String] = conf("spark.comet.exec.memoryPool")
- .doc(
- "The type of memory pool to be used for Comet native execution. " +
- "Available memory pool types are 'greedy', 'fair_spill', 'greedy_task_shared', " +
- "'fair_spill_task_shared', 'greedy_global', 'fair_spill_global', and `unbounded`. " +
- "For off-heap types are 'unified' and `fair_unified`.")
+ .doc("The type of memory pool to be used for Comet native execution. " +
+ "When running Spark in on-heap mode, available pool types are 'greedy', 'fair_spill', " +
+ "'greedy_task_shared', 'fair_spill_task_shared', 'greedy_global', 'fair_spill_global', " +
+ "and `unbounded`. When running Spark in off-heap mode, available pool types are " +
+ "'unified' and `fair_unified`. The default pool type is `greedy_task_shared` for on-heap " +
+ s"mode and `unified` for off-heap mode. $TUNING_GUIDE.")
.stringConf
- .createWithDefault("greedy_task_shared")
+ .createWithDefault("default")
val COMET_SCAN_PREFETCH_ENABLED: ConfigEntry[Boolean] =
conf("spark.comet.scan.preFetch.enabled")
diff --git a/docs/source/user-guide/configs.md b/docs/source/user-guide/configs.md
index 32c6a3db8..d0f95480d 100644
--- a/docs/source/user-guide/configs.md
+++ b/docs/source/user-guide/configs.md
@@ -54,7 +54,7 @@ Comet provides the following configuration settings.
| spark.comet.exec.hashJoin.enabled | Whether to enable hashJoin by default. | true |
| spark.comet.exec.initCap.enabled | Whether to enable initCap by default. | false |
| spark.comet.exec.localLimit.enabled | Whether to enable localLimit by default. | true |
-| spark.comet.exec.memoryPool | The type of memory pool to be used for Comet native execution. Available memory pool types are 'greedy', 'fair_spill', 'greedy_task_shared', 'fair_spill_task_shared', 'greedy_global', 'fair_spill_global', and `unbounded`. For off-heap types are 'unified' and `fair_unified`. | greedy_task_shared |
+| spark.comet.exec.memoryPool | The type of memory pool to be used for Comet native execution. When running Spark in on-heap mode, available pool types are 'greedy', 'fair_spill', 'greedy_task_shared', 'fair_spill_task_shared', 'greedy_global', 'fair_spill_global', and `unbounded`. When running Spark in off-heap mode, available pool types are 'unified' and `fair_unified`. The default pool type is `greedy_task_shared` for on-heap mode and `unified` for off-heap mode. For more information, [...]
| spark.comet.exec.project.enabled | Whether to enable project by default. | true |
| spark.comet.exec.replaceSortMergeJoin | Experimental feature to force Spark to replace SortMergeJoin with ShuffledHashJoin for improved performance. This feature is not stable yet. For more information, refer to the Comet Tuning Guide (https://datafusion.apache.org/comet/user-guide/tuning.html). | false |
| spark.comet.exec.shuffle.compression.codec | The codec of Comet native shuffle used to compress shuffle data. lz4, zstd, and snappy are supported. Compression can be disabled by setting spark.shuffle.compress=false. | lz4 |
@@ -71,9 +71,9 @@ Comet provides the following configuration settings.
| spark.comet.explain.verbose.enabled | When this setting is enabled, Comet will provide a verbose tree representation of the extended information. | false |
| spark.comet.explainFallback.enabled | When this setting is enabled, Comet will provide logging explaining the reason(s) why a query stage cannot be executed natively. Set this to false to reduce the amount of logging. | false |
| spark.comet.expression.allowIncompatible | Comet is not currently fully compatible with Spark for all expressions. Set this config to true to allow them anyway. For more information, refer to the Comet Compatibility Guide (https://datafusion.apache.org/comet/user-guide/compatibility.html). | false |
-| spark.comet.memory.overhead.factor | Fraction of executor memory to be allocated as additional memory for Comet when running in on-heap mode or when using the `fair_unified` pool in off-heap mode. For more information, refer to the Comet Tuning Guide (https://datafusion.apache.org/comet/user-guide/tuning.html). | 0.2 |
-| spark.comet.memory.overhead.min | Minimum amount of additional memory to be allocated per executor process for Comet, in MiB, when running in on-heap mode or when using the `fair_unified` pool in off-heap mode. For more information, refer to the Comet Tuning Guide (https://datafusion.apache.org/comet/user-guide/tuning.html). | 402653184b |
-| spark.comet.memoryOverhead | The amount of additional memory to be allocated per executor process for Comet, in MiB, when running in on-heap mode or when using the `fair_unified` pool in off-heap mode. This config is optional. If this is not specified, it will be set to `spark.comet.memory.overhead.factor` * `spark.executor.memory`. For more information, refer to the Comet Tuning Guide (https://datafusion.apache.org/comet/user-guide/tuning.html). | |
+| spark.comet.memory.overhead.factor | Fraction of executor memory to be allocated as additional memory for Comet when running in on-heap mode. For more information, refer to the Comet Tuning Guide (https://datafusion.apache.org/comet/user-guide/tuning.html). | 0.2 |
+| spark.comet.memory.overhead.min | Minimum amount of additional memory to be allocated per executor process for Comet, in MiB, when running in on-heap mode. For more information, refer to the Comet Tuning Guide (https://datafusion.apache.org/comet/user-guide/tuning.html). | 402653184b |
+| spark.comet.memoryOverhead | The amount of additional memory to be allocated per executor process for Comet, in MiB, when running in on-heap mode. This config is optional. If this is not specified, it will be set to `spark.comet.memory.overhead.factor` * `spark.executor.memory`. For more information, refer to the Comet Tuning Guide (https://datafusion.apache.org/comet/user-guide/tuning.html). | |
| spark.comet.metrics.updateInterval | The interval in milliseconds to update metrics. If interval is negative, metrics will be updated upon task completion. | 3000 |
| spark.comet.nativeLoadRequired | Whether to require Comet native library to load successfully when Comet is enabled. If not, Comet will silently fallback to Spark when it fails to load the native lib. Otherwise, an error will be thrown and the Spark job will be aborted. | false |
| spark.comet.parquet.enable.directBuffer | Whether to use Java direct byte buffer when reading Parquet. | false |
diff --git a/docs/source/user-guide/tuning.md b/docs/source/user-guide/tuning.md
index 1a17f4ccc..7a745d5fc 100644
--- a/docs/source/user-guide/tuning.md
+++ b/docs/source/user-guide/tuning.md
@@ -108,34 +108,29 @@ resource managers respect Apache Spark memory configuration before starting the
Comet implements multiple memory pool implementations. The type of pool can be specified with `spark.comet.exec.memoryPool`.
-The valid pool types are:
+The valid pool types for off-heap mode are:
- `unified` (default when `spark.memory.offHeap.enabled=true` is set)
- `fair_unified`
+Both of these pools share off-heap memory between Spark and Comet. This approach is referred to as
+unified memory management. The size of the pool is specified by `spark.memory.offHeap.size`.
+
The `unified` pool type implements a greedy first-come first-serve limit. This pool works well for queries that do not
-need to spill or have a single spillable operator. The size of the pool is specified by `spark.memory.offHeap.size`
-and the pool interacts with Spark's memory pool, effectively sharing the off-heap memory between Spark and Comet. This
-approach is sometimes referred to as unified memory management.
+need to spill or have a single spillable operator.
The `fair_unified` pool type prevents operators from using more than an even fraction of the available memory
(i.e. `pool_size / num_reservations`). This pool works best when you know beforehand
the query has multiple operators that will likely all need to spill. Sometimes it will cause spills even
when there is sufficient memory in order to leave enough memory for other operators.
-The pool size configuration for the `fair_unified` pool, is a little more complex. The total pool size is computed by
-multiplying `spark.memory.offHeap.size` by `spark.comet.memory.overhead.factor` with the minimum amount being
-`spark.comet.memory.overhead.min`. It is also possible to manually specify `spark.comet.memoryOverhead` instead to
-override this default behavior. Note that the `fair_unified` pool does not use unified memory management to interact
-with Spark's memory pools, which is why the allocation defaults to a fraction of off-heap memory.
-
### Configuring On-Heap Memory Pools
When running in on-heap mode, Comet will use its own dedicated memory pools that are not shared with Spark.
The type of pool can be specified with `spark.comet.exec.memoryPool`. The default setting is `greedy_task_shared`.
-The valid pool types are:
+The valid pool types for on-heap mode are:
- `greedy`
- `greedy_global`
diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs
index 0aa1b113d..dd4d031b9 100644
--- a/native/core/src/execution/jni_api.rs
+++ b/native/core/src/execution/jni_api.rs
@@ -302,12 +302,18 @@ fn parse_memory_pool_config(
let memory_pool_config = if off_heap_mode {
match memory_pool_type.as_str() {
"fair_unified" => MemoryPoolConfig::new(MemoryPoolType::FairUnified, pool_size),
- _ => {
- // the Unified memory pool interacts with Spark's memory pool to allocate
+ "default" | "unified" => {
+ // the `unified` memory pool interacts with Spark's memory pool to allocate
// memory therefore does not need a size to be explicitly set. The pool size
// shared with Spark is set by `spark.memory.offHeap.size`.
MemoryPoolConfig::new(MemoryPoolType::Unified, 0)
}
+ _ => {
+ return Err(CometError::Config(format!(
+ "Unsupported memory pool type for off-heap mode: {}",
+ memory_pool_type
+ )))
+ }
}
} else {
// Use the memory pool from DF
@@ -316,7 +322,7 @@ fn parse_memory_pool_config(
"fair_spill_task_shared" => {
MemoryPoolConfig::new(MemoryPoolType::FairSpillTaskShared, pool_size_per_task)
}
- "greedy_task_shared" => {
+ "default" | "greedy_task_shared" => {
MemoryPoolConfig::new(MemoryPoolType::GreedyTaskShared, pool_size_per_task)
}
"fair_spill_global" => {
@@ -328,7 +334,7 @@ fn parse_memory_pool_config(
"unbounded" => MemoryPoolConfig::new(MemoryPoolType::Unbounded, 0),
_ => {
return Err(CometError::Config(format!(
- "Unsupported memory pool type: {}",
+ "Unsupported memory pool type for on-heap mode: {}",
memory_pool_type
)))
}
diff --git a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala
index e4befefa6..b9bb3b2b9 100644
--- a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala
+++ b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala
@@ -21,6 +21,7 @@ package org.apache.comet
import org.apache.spark._
import org.apache.spark.internal.Logging
+import org.apache.spark.network.util.ByteUnit
import org.apache.spark.sql.comet.CometMetricNode
import org.apache.spark.sql.vectorized._
@@ -63,9 +64,17 @@ class CometExecIterator(
}.toArray
private val plan = {
val conf = SparkEnv.get.conf
- // Only enable unified memory manager when off-heap mode is enabled. Otherwise,
- // we'll use the built-in memory pool from DF, and initializes with `memory_limit`
- // and `memory_fraction` below.
+
+ val offHeapMode = CometSparkSessionExtensions.isOffHeapEnabled(conf)
+ val memoryLimit = if (offHeapMode) {
+ // in unified mode we share off-heap memory with Spark
+ ByteUnit.MiB.toBytes(conf.getSizeAsMb("spark.memory.offHeap.size"))
+ } else {
+ // we'll use the built-in memory pool from DF, and initializes with `memory_limit`
+ // and `memory_fraction` below.
+ CometSparkSessionExtensions.getCometMemoryOverhead(conf)
+ }
+
nativeLib.createPlan(
id,
cometBatchIterators,
@@ -75,9 +84,9 @@ class CometExecIterator(
metricsUpdateInterval = COMET_METRICS_UPDATE_INTERVAL.get(),
new CometTaskMemoryManager(id),
batchSize = COMET_BATCH_SIZE.get(),
- offHeapMode = CometSparkSessionExtensions.isOffHeapEnabled(conf),
+ offHeapMode,
memoryPoolType = COMET_EXEC_MEMORY_POOL_TYPE.get(),
- memoryLimit = CometSparkSessionExtensions.getCometMemoryOverhead(conf),
+ memoryLimit,
memoryLimitPerTask = getMemoryLimitPerTask(conf),
taskAttemptId = TaskContext.get().taskAttemptId,
debug = COMET_DEBUG_ENABLED.get(),
diff --git a/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala b/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala
index a669b36eb..0b423d348 100644
--- a/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala
+++ b/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala
@@ -1334,26 +1334,46 @@ object CometSparkSessionExtensions extends Logging {
CometConf.COMET_NATIVE_SCAN_IMPL.get(conf) == CometConf.SCAN_NATIVE_DATAFUSION
}
- /** Calculates required memory overhead in MB per executor process for Comet. */
+ /**
+ * Whether we should override Spark memory configuration for Comet. This only returns true when
+ * Comet native execution is enabled and/or Comet shuffle is enabled and Comet doesn't use
+ * off-heap mode (unified memory manager).
+ */
+ def shouldOverrideMemoryConf(conf: SparkConf): Boolean = {
+ val cometEnabled = getBooleanConf(conf, CometConf.COMET_ENABLED)
+ val cometShuffleEnabled = getBooleanConf(conf, CometConf.COMET_EXEC_SHUFFLE_ENABLED)
+ val cometExecEnabled = getBooleanConf(conf, CometConf.COMET_EXEC_ENABLED)
+ val offHeapMode = CometSparkSessionExtensions.isOffHeapEnabled(conf)
+ cometEnabled && (cometShuffleEnabled || cometExecEnabled) && !offHeapMode
+ }
+
+ /**
+ * Calculates required memory overhead in MB per executor process for Comet when running in
+ * on-heap mode.
+ *
+ * If `COMET_MEMORY_OVERHEAD` is defined then that value will be used, otherwise the overhead
+ * will be calculated by multiplying executor memory (`spark.executor.memory`) by
+ * `COMET_MEMORY_OVERHEAD_FACTOR`.
+ *
+ * In either case, a minimum value of `COMET_MEMORY_OVERHEAD_MIN_MIB` will be returned.
+ */
def getCometMemoryOverheadInMiB(sparkConf: SparkConf): Long = {
- val baseMemoryMiB = if (isOffHeapEnabled(sparkConf)) {
- ConfigHelpers
- .byteFromString(sparkConf.get("spark.memory.offHeap.size"), ByteUnit.MiB)
- } else {
- // `spark.executor.memory` default value is 1g
- ConfigHelpers
- .byteFromString(sparkConf.get("spark.executor.memory", "1024MB"), ByteUnit.MiB)
+ if (isOffHeapEnabled(sparkConf)) {
+ // when running in off-heap mode we use unified memory management to share
+ // off-heap memory with Spark so do not add overhead
+ return 0
}
- val minimum = ConfigHelpers
- .byteFromString(
- sparkConf.get(
- COMET_MEMORY_OVERHEAD_MIN_MIB.key,
- COMET_MEMORY_OVERHEAD_MIN_MIB.defaultValueString),
- ByteUnit.MiB)
- val overheadFactor = sparkConf.getDouble(
- COMET_MEMORY_OVERHEAD_FACTOR.key,
- COMET_MEMORY_OVERHEAD_FACTOR.defaultValue.get)
+ // `spark.executor.memory` default value is 1g
+ val baseMemoryMiB = ConfigHelpers
+ .byteFromString(sparkConf.get("spark.executor.memory", "1024MB"), ByteUnit.MiB)
+
+ val cometMemoryOverheadMinAsString = sparkConf.get(
+ COMET_MEMORY_OVERHEAD_MIN_MIB.key,
+ COMET_MEMORY_OVERHEAD_MIN_MIB.defaultValueString)
+
+ val minimum = ConfigHelpers.byteFromString(cometMemoryOverheadMinAsString, ByteUnit.MiB)
+ val overheadFactor = getDoubleConf(sparkConf, COMET_MEMORY_OVERHEAD_FACTOR)
val overHeadMemFromConf = sparkConf
.getOption(COMET_MEMORY_OVERHEAD.key)
@@ -1362,7 +1382,16 @@ object CometSparkSessionExtensions extends Logging {
overHeadMemFromConf.getOrElse(math.max((overheadFactor * baseMemoryMiB).toLong, minimum))
}
- /** Calculates required memory overhead in bytes per executor process for Comet. */
+ private def getBooleanConf(conf: SparkConf, entry: ConfigEntry[Boolean]) =
+ conf.getBoolean(entry.key, entry.defaultValue.get)
+
+ private def getDoubleConf(conf: SparkConf, entry: ConfigEntry[Double]) =
+ conf.getDouble(entry.key, entry.defaultValue.get)
+
+ /**
+ * Calculates required memory overhead in bytes per executor process for Comet when running in
+ * on-heap mode.
+ */
def getCometMemoryOverhead(sparkConf: SparkConf): Long = {
ByteUnit.MiB.toBytes(getCometMemoryOverheadInMiB(sparkConf))
}
@@ -1391,11 +1420,6 @@ object CometSparkSessionExtensions extends Logging {
}
}
- /** Calculates Comet shuffle memory size in MB */
- def getCometShuffleMemorySizeInMiB(sparkConf: SparkConf, conf: SQLConf = SQLConf.get): Long = {
- ByteUnit.BYTE.toMiB(getCometShuffleMemorySize(sparkConf, conf))
- }
-
def isOffHeapEnabled(sparkConf: SparkConf): Boolean = {
sparkConf.getBoolean("spark.memory.offHeap.enabled", false)
}
diff --git a/spark/src/main/scala/org/apache/spark/Plugins.scala b/spark/src/main/scala/org/apache/spark/Plugins.scala
index 64426c716..1f3d4bba8 100644
--- a/spark/src/main/scala/org/apache/spark/Plugins.scala
+++ b/spark/src/main/scala/org/apache/spark/Plugins.scala
@@ -28,7 +28,7 @@ import org.apache.spark.internal.Logging
import org.apache.spark.internal.config.{EXECUTOR_MEMORY, EXECUTOR_MEMORY_OVERHEAD, EXECUTOR_MEMORY_OVERHEAD_FACTOR}
import org.apache.spark.sql.internal.StaticSQLConf
-import org.apache.comet.{CometConf, CometSparkSessionExtensions}
+import org.apache.comet.CometSparkSessionExtensions
/**
* Comet driver plugin. This class is loaded by Spark's plugin framework. It will be instantiated
@@ -50,7 +50,7 @@ class CometDriverPlugin extends DriverPlugin with Logging with ShimCometDriverPl
// register CometSparkSessionExtensions if it isn't already registered
CometDriverPlugin.registerCometSessionExtension(sc.conf)
- if (shouldOverrideMemoryConf(sc.getConf)) {
+ if (CometSparkSessionExtensions.shouldOverrideMemoryConf(sc.getConf)) {
val execMemOverhead = if (sc.getConf.contains(EXECUTOR_MEMORY_OVERHEAD.key)) {
sc.getConf.getSizeAsMb(EXECUTOR_MEMORY_OVERHEAD.key)
} else {
@@ -63,9 +63,6 @@ class CometDriverPlugin extends DriverPlugin with Logging with ShimCometDriverPl
Math.max((executorMemory * memoryOverheadFactor).toLong, memoryOverheadMinMib)
}
- // we should never reach this code in off-heap mode due to earlier check
- // in `shouldOverrideMemoryConf`
- assert(!CometSparkSessionExtensions.isOffHeapEnabled(sc.getConf))
val cometMemOverhead = CometSparkSessionExtensions.getCometMemoryOverheadInMiB(sc.getConf)
sc.conf.set(EXECUTOR_MEMORY_OVERHEAD.key, s"${execMemOverhead + cometMemOverhead}M")
val newExecMemOverhead = sc.getConf.getSizeAsMb(EXECUTOR_MEMORY_OVERHEAD.key)
@@ -76,6 +73,8 @@ class CometDriverPlugin extends DriverPlugin with Logging with ShimCometDriverPl
- Comet memory overhead: ${cometMemOverhead}MB
- Updated Spark executor memory overhead: ${newExecMemOverhead}MB
""")
+ } else {
+ logInfo("Comet is running in unified memory mode and sharing off-heap memory with Spark")
}
Collections.emptyMap[String, String]
@@ -92,24 +91,6 @@ class CometDriverPlugin extends DriverPlugin with Logging with ShimCometDriverPl
override def registerMetrics(appId: String, pluginContext: PluginContext): Unit =
super.registerMetrics(appId, pluginContext)
- /**
- * Whether we should override Spark memory configuration for Comet. This only returns true when
- * Comet native execution is enabled and/or Comet shuffle is enabled and Comet doesn't use
- * off-heap mode (unified memory manager).
- */
- private def shouldOverrideMemoryConf(conf: SparkConf): Boolean = {
- val cometEnabled =
- conf.getBoolean(CometConf.COMET_ENABLED.key, CometConf.COMET_ENABLED.defaultValue.get)
- val cometExecShuffle = conf.getBoolean(
- CometConf.COMET_EXEC_SHUFFLE_ENABLED.key,
- CometConf.COMET_EXEC_SHUFFLE_ENABLED.defaultValue.get)
- val cometExec = conf.getBoolean(
- CometConf.COMET_EXEC_ENABLED.key,
- CometConf.COMET_EXEC_ENABLED.defaultValue.get)
- val offHeapMode = CometSparkSessionExtensions.isOffHeapEnabled(conf)
-
- cometEnabled && (cometExecShuffle || cometExec) && !offHeapMode
- }
}
object CometDriverPlugin extends Logging {
diff --git a/spark/src/test/scala/org/apache/comet/CometSparkSessionExtensionsSuite.scala b/spark/src/test/scala/org/apache/comet/CometSparkSessionExtensionsSuite.scala
index 516fb09f8..33552a116 100644
--- a/spark/src/test/scala/org/apache/comet/CometSparkSessionExtensionsSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/CometSparkSessionExtensionsSuite.scala
@@ -24,11 +24,14 @@ import org.apache.spark.sql._
import org.apache.spark.sql.internal.SQLConf
class CometSparkSessionExtensionsSuite extends CometTestBase {
+
+ import CometSparkSessionExtensions._
+
test("isCometEnabled") {
val conf = new SQLConf
conf.setConfString(CometConf.COMET_ENABLED.key, "false")
- assert(!CometSparkSessionExtensions.isCometEnabled(conf))
+ assert(!isCometEnabled(conf))
// Since the native lib is probably already loaded due to previous tests, we reset it here
NativeBase.setLoaded(false)
@@ -36,12 +39,12 @@ class CometSparkSessionExtensionsSuite extends CometTestBase {
conf.setConfString(CometConf.COMET_ENABLED.key, "true")
val oldProperty = System.getProperty("os.name")
System.setProperty("os.name", "foo")
- assert(!CometSparkSessionExtensions.isCometEnabled(conf))
+ assert(!isCometEnabled(conf))
System.setProperty("os.name", oldProperty)
conf.setConf(SQLConf.PARQUET_INT96_TIMESTAMP_CONVERSION, true)
- assert(!CometSparkSessionExtensions.isCometEnabled(conf))
+ assert(!isCometEnabled(conf))
// Restore the original state
NativeBase.setLoaded(true)
@@ -75,9 +78,7 @@ class CometSparkSessionExtensionsSuite extends CometTestBase {
test("Minimum Comet memory overhead") {
val conf = new SparkConf()
- assert(
- CometSparkSessionExtensions
- .getCometMemoryOverhead(conf) == getBytesFromMib(384))
+ assert(getCometMemoryOverhead(conf) == getBytesFromMib(384))
}
test("Comet memory overhead factor with executor memory") {
@@ -85,25 +86,41 @@ class CometSparkSessionExtensionsSuite extends CometTestBase {
sparkConf.set("spark.executor.memory", "16g")
sparkConf.set(CometConf.COMET_MEMORY_OVERHEAD_FACTOR.key, "0.5")
- assert(
- CometSparkSessionExtensions
- .getCometMemoryOverhead(sparkConf) == getBytesFromMib(8 * 1024))
+ assert(getCometMemoryOverhead(sparkConf) == getBytesFromMib(8 * 1024))
}
test("Comet memory overhead factor with default executor memory") {
val sparkConf = new SparkConf()
sparkConf.set(CometConf.COMET_MEMORY_OVERHEAD_FACTOR.key, "0.5")
- assert(
- CometSparkSessionExtensions
- .getCometMemoryOverhead(sparkConf) == getBytesFromMib(512))
+ assert(getCometMemoryOverhead(sparkConf) == getBytesFromMib(512))
}
test("Comet memory overhead") {
val sparkConf = new SparkConf()
sparkConf.set(CometConf.COMET_MEMORY_OVERHEAD.key, "10g")
- assert(
- CometSparkSessionExtensions
- .getCometMemoryOverhead(sparkConf) == getBytesFromMib(1024 * 10))
+ assert(getCometMemoryOverhead(sparkConf) == getBytesFromMib(1024 * 10))
+ assert(shouldOverrideMemoryConf(sparkConf))
+ }
+
+ test("Comet memory overhead (min)") {
+ val sparkConf = new SparkConf()
+ sparkConf.set(CometConf.COMET_MEMORY_OVERHEAD_MIN_MIB.key, "2g")
+ assert(getCometMemoryOverhead(sparkConf) == getBytesFromMib(1024 * 2))
+ }
+
+ test("Comet memory overhead (factor)") {
+ val sparkConf = new SparkConf()
+ sparkConf.set(CometConf.COMET_MEMORY_OVERHEAD_FACTOR.key, "0.5")
+ assert(getCometMemoryOverhead(sparkConf) == getBytesFromMib(512))
+ }
+
+ test("Comet memory overhead (off heap)") {
+ val sparkConf = new SparkConf()
+ sparkConf.set(CometConf.COMET_MEMORY_OVERHEAD.key, "64g")
+ sparkConf.set("spark.memory.offHeap.enabled", "true")
+ sparkConf.set("spark.memory.offHeap.size", "10g")
+ assert(getCometMemoryOverhead(sparkConf) == 0)
+ assert(!shouldOverrideMemoryConf(sparkConf))
}
test("Comet shuffle memory factor") {
@@ -114,12 +131,12 @@ class CometSparkSessionExtensionsSuite extends CometTestBase {
// Minimum Comet memory overhead is 384MB
assert(
- CometSparkSessionExtensions.getCometShuffleMemorySize(conf, sqlConf) ==
+ getCometShuffleMemorySize(conf, sqlConf) ==
getBytesFromMib((384 * 0.2).toLong))
conf.set(CometConf.COMET_MEMORY_OVERHEAD_FACTOR.key, "0.5")
assert(
- CometSparkSessionExtensions.getCometShuffleMemorySize(conf, sqlConf) ==
+ getCometShuffleMemorySize(conf, sqlConf) ==
getBytesFromMib((1024 * 0.5 * 0.2).toLong))
}
@@ -129,9 +146,20 @@ class CometSparkSessionExtensionsSuite extends CometTestBase {
conf.set(CometConf.COMET_MEMORY_OVERHEAD.key, "1g")
sqlConf.setConfString(CometConf.COMET_COLUMNAR_SHUFFLE_MEMORY_SIZE.key, "512m")
- assert(
- CometSparkSessionExtensions
- .getCometShuffleMemorySize(conf, sqlConf) == getBytesFromMib(512))
+ assert(getCometShuffleMemorySize(conf, sqlConf) == getBytesFromMib(512))
+ }
+
+ test("Comet shuffle memory (off-heap)") {
+ val conf = new SparkConf()
+ val sqlConf = new SQLConf
+ conf.set(CometConf.COMET_MEMORY_OVERHEAD.key, "1g")
+ conf.set("spark.memory.offHeap.enabled", "true")
+ conf.set("spark.memory.offHeap.size", "10g")
+ sqlConf.setConfString(CometConf.COMET_COLUMNAR_SHUFFLE_MEMORY_SIZE.key, "512m")
+
+ assertThrows[AssertionError] {
+ getCometShuffleMemorySize(conf, sqlConf)
+ }
}
test("Comet shuffle memory cannot be larger than Comet memory overhead") {
@@ -139,8 +167,6 @@ class CometSparkSessionExtensionsSuite extends CometTestBase {
val sqlConf = new SQLConf
conf.set(CometConf.COMET_MEMORY_OVERHEAD.key, "1g")
sqlConf.setConfString(CometConf.COMET_COLUMNAR_SHUFFLE_MEMORY_SIZE.key, "10g")
- assert(
- CometSparkSessionExtensions
- .getCometShuffleMemorySize(conf, sqlConf) == getBytesFromMib(1024))
+ assert(getCometShuffleMemorySize(conf, sqlConf) == getBytesFromMib(1024))
}
}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org