This is an automated email from the ASF dual-hosted git repository.
viirya 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 380f03d0 feat: Add native thread configs (#828)
380f03d0 is described below
commit 380f03d067b7703adefd3dce1849d9709bad3da9
Author: Liang-Chi Hsieh <[email protected]>
AuthorDate: Wed Aug 14 22:52:28 2024 -0700
feat: Add native thread configs (#828)
---
common/src/main/scala/org/apache/comet/CometConf.scala | 16 ++++++++++++++++
native/core/src/execution/jni_api.rs | 13 +++++++++++++
.../main/scala/org/apache/comet/CometExecIterator.scala | 4 +++-
3 files changed, 32 insertions(+), 1 deletion(-)
diff --git a/common/src/main/scala/org/apache/comet/CometConf.scala
b/common/src/main/scala/org/apache/comet/CometConf.scala
index 1f4874eb..1db80298 100644
--- a/common/src/main/scala/org/apache/comet/CometConf.scala
+++ b/common/src/main/scala/org/apache/comet/CometConf.scala
@@ -358,6 +358,22 @@ object CometConf extends ShimCometConf {
.booleanConf
.createWithDefault(false)
+ val COMET_WORKER_THREADS: ConfigEntry[Int] =
+ conf("spark.comet.workerThreads")
+ .internal()
+ .doc("The number of worker threads used for Comet native execution. " +
+ "By default, this config is 4.")
+ .intConf
+ .createWithDefault(4)
+
+ val COMET_BLOCKING_THREADS: ConfigEntry[Int] =
+ conf("spark.comet.blockingThreads")
+ .internal()
+ .doc("The number of blocking threads used for Comet native execution. " +
+ "By default, this config is 10.")
+ .intConf
+ .createWithDefault(10)
+
val COMET_BATCH_SIZE: ConfigEntry[Int] = conf("spark.comet.batchSize")
.doc("The columnar batch size, i.e., the maximum number of rows that a
batch can contain.")
.intConf
diff --git a/native/core/src/execution/jni_api.rs
b/native/core/src/execution/jni_api.rs
index 6d61f4a4..587defaa 100644
--- a/native/core/src/execution/jni_api.rs
+++ b/native/core/src/execution/jni_api.rs
@@ -133,8 +133,21 @@ pub unsafe extern "system" fn
Java_org_apache_comet_Native_createPlan(
let debug_native = parse_bool(&configs, "debug_native")?;
let explain_native = parse_bool(&configs, "explain_native")?;
+ let worker_threads = configs
+ .get("worker_threads")
+ .map(String::as_str)
+ .unwrap_or("4")
+ .parse::<usize>()?;
+ let blocking_threads = configs
+ .get("blocking_threads")
+ .map(String::as_str)
+ .unwrap_or("10")
+ .parse::<usize>()?;
+
// Use multi-threaded tokio runtime to prevent blocking spawned tasks
if any
let runtime = tokio::runtime::Builder::new_multi_thread()
+ .worker_threads(worker_threads)
+ .max_blocking_threads(blocking_threads)
.enable_all()
.build()?;
diff --git a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala
b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala
index ef18850b..29eb2f0c 100644
--- a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala
+++ b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala
@@ -23,7 +23,7 @@ import org.apache.spark._
import org.apache.spark.sql.comet.CometMetricNode
import org.apache.spark.sql.vectorized._
-import org.apache.comet.CometConf.{COMET_BATCH_SIZE, COMET_DEBUG_ENABLED,
COMET_EXEC_MEMORY_FRACTION, COMET_EXPLAIN_NATIVE_ENABLED}
+import org.apache.comet.CometConf.{COMET_BATCH_SIZE, COMET_BLOCKING_THREADS,
COMET_DEBUG_ENABLED, COMET_EXEC_MEMORY_FRACTION, COMET_EXPLAIN_NATIVE_ENABLED,
COMET_WORKER_THREADS}
import org.apache.comet.vector.NativeUtil
/**
@@ -86,6 +86,8 @@ class CometExecIterator(
result.put("batch_size", String.valueOf(COMET_BATCH_SIZE.get()))
result.put("debug_native", String.valueOf(COMET_DEBUG_ENABLED.get()))
result.put("explain_native",
String.valueOf(COMET_EXPLAIN_NATIVE_ENABLED.get()))
+ result.put("worker_threads", String.valueOf(COMET_WORKER_THREADS.get()))
+ result.put("blocking_threads",
String.valueOf(COMET_BLOCKING_THREADS.get()))
// Strip mandatory prefix spark. which is not required for DataFusion
session params
conf.getAll.foreach {
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
