chore: Pass Spark configs to native createPlan (#2180)

Author: andygrove

native/core/src/execution/jni_api.rs

@@ -60,6 +60,7 @@ use jni::{
     objects::GlobalRef,
     sys::{jboolean, jdouble, jintArray, jobjectArray, jstring},
 };
+use std::collections::HashMap;
 use std::path::PathBuf;
 use std::time::{Duration, Instant};
 use std::{sync::Arc, task::Poll};
@@ -151,6 +152,7 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_createPlan(
     id: jlong,
     iterators: jobjectArray,
     serialized_query: jbyteArray,
+    serialized_spark_configs: jbyteArray,
     partition_count: jint,
     metrics_node: JObject,
     metrics_update_interval: jlong,
@@ -173,12 +175,20 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_createPlan(
 
             let start = Instant::now();
 
+            // Deserialize query plan
             let array = unsafe { JPrimitiveArray::from_raw(serialized_query) };
             let bytes = env.convert_byte_array(array)?;
-
-            // Deserialize query plan
             let spark_plan = serde::deserialize_op(bytes.as_slice())?;
 
+            // Deserialize Spark configs
+            let array = unsafe { JPrimitiveArray::from_raw(serialized_spark_configs) };
+            let bytes = env.convert_byte_array(array)?;
+            let spark_configs = serde::deserialize_config(bytes.as_slice())?;
+
+            // Convert Spark configs to HashMap
+            let _spark_config_map: HashMap<String, String> =
+                spark_configs.entries.into_iter().collect();
+
             let metrics = Arc::new(jni_new_global_ref!(env, metrics_node)?);
 
             // Get the global references of input sources

native/core/src/execution/serde.rs

@@ -22,7 +22,7 @@ use crate::errors::ExpressionError;
 use arrow::datatypes::{DataType as ArrowDataType, TimeUnit};
 use arrow::datatypes::{Field, Fields};
 use datafusion_comet_proto::{
-    spark_expression,
+    spark_config, spark_expression,
     spark_expression::data_type::{
         data_type_info::DatatypeStruct,
         DataTypeId,
@@ -74,6 +74,14 @@ pub fn deserialize_op(buf: &[u8]) -> Result<spark_operator::Operator, ExecutionE
     }
 }
 
+/// Deserialize bytes to protobuf type of data type
+pub fn deserialize_config(buf: &[u8]) -> Result<spark_config::ConfigMap, ExecutionError> {
+    match spark_config::ConfigMap::decode(&mut Cursor::new(buf)) {
+        Ok(e) => Ok(e),
+        Err(err) => Err(ExecutionError::from(err)),
+    }
+}
+
 /// Deserialize bytes to protobuf type of data type
 pub fn deserialize_data_type(buf: &[u8]) -> Result<spark_expression::DataType, ExecutionError> {
     match spark_expression::DataType::decode(&mut Cursor::new(buf)) {

native/proto/build.rs

@@ -33,6 +33,7 @@ fn main() -> Result<()> {
             "src/proto/metric.proto",
             "src/proto/partitioning.proto",
             "src/proto/operator.proto",
+            "src/proto/config.proto",
         ],
         &["src/proto"],
     )?;

native/proto/src/lib.rs

@@ -43,3 +43,9 @@ pub mod spark_operator {
 pub mod spark_metric {
     include!(concat!("generated", "/spark.spark_metric.rs"));
 }
+
+// Include generated modules from .proto files.
+#[allow(missing_docs)]
+pub mod spark_config {
+    include!(concat!("generated", "/spark.spark_config.rs"));
+}

native/proto/src/proto/config.proto

@@ -0,0 +1,26 @@
+// 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.
+
+syntax = "proto3";
+
+package spark.spark_config;
+
+option java_package = "org.apache.comet.serde";
+
+message ConfigMap {
+  map<string, string> entries = 1;
+}

spark/src/main/scala/org/apache/comet/CometExecIterator.scala

@@ -32,6 +32,7 @@ import org.apache.spark.sql.vectorized._
 
 import org.apache.comet.CometConf.{COMET_BATCH_SIZE, COMET_DEBUG_ENABLED, COMET_EXEC_MEMORY_POOL_TYPE, COMET_EXPLAIN_NATIVE_ENABLED, COMET_METRICS_UPDATE_INTERVAL}
 import org.apache.comet.Tracing.withTrace
+import org.apache.comet.serde.Config.ConfigMap
 import org.apache.comet.vector.NativeUtil
 
 /**
@@ -84,10 +85,19 @@ class CometExecIterator(
       // and `memory_fraction` below.
       CometSparkSessionExtensions.getCometMemoryOverhead(conf)
     }
+
+    // serialize Spark conf in protobuf format
+    val builder = ConfigMap.newBuilder()
+    conf.getAll.foreach { case (k, v) =>
+      builder.putEntries(k, v)
+    }
+    val protobufSparkConfigs = builder.build().toByteArray
+
     nativeLib.createPlan(
       id,
       cometBatchIterators,
       protobufQueryPlan,
+      protobufSparkConfigs,
       numParts,
       nativeMetrics,
       metricsUpdateInterval = COMET_METRICS_UPDATE_INTERVAL.get(),

spark/src/main/scala/org/apache/comet/Native.scala

@@ -54,6 +54,7 @@ class Native extends NativeBase {
       id: Long,
       iterators: Array[CometBatchIterator],
       plan: Array[Byte],
+      configMapProto: Array[Byte],
       partitionCount: Int,
       metrics: CometMetricNode,
       metricsUpdateInterval: Long,