experiments with Spark conf

Author: comphead

native/core/src/config.rs

@@ -0,0 +1,30 @@
+// 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.
+
+use datafusion::common::extensions_options;
+use datafusion::config::ConfigExtension;
+
+extensions_options! {
+    pub struct CometNativeOptions {
+        // Default File System
+        pub default_fs: Option<String>, default = None
+    }
+}
+
+impl ConfigExtension for CometNativeOptions {
+    const PREFIX: &'static str = "comet";
+}

native/core/src/execution/jni_api.rs

@@ -29,6 +29,7 @@ use crate::{
 use arrow::array::RecordBatch;
 use arrow::datatypes::DataType as ArrowDataType;
 use datafusion::common::ScalarValue;
+use datafusion::config::Extensions;
 use datafusion::execution::memory_pool::{
     FairSpillPool, GreedyMemoryPool, MemoryPool, TrackConsumersPool, UnboundedMemoryPool,
 };
@@ -41,7 +42,8 @@ use datafusion::{
 use datafusion_comet_proto::spark_operator::Operator;
 use futures::poll;
 use futures::stream::StreamExt;
-use jni::objects::{JByteBuffer, JMap};
+use jni::objects::JByteBuffer;
+use jni::objects::JMap;
 use jni::sys::JNI_FALSE;
 use jni::{
     errors::Result as JNIResult,
@@ -63,6 +65,7 @@ use std::time::{Duration, Instant};
 use std::{collections::HashMap, sync::Arc, task::Poll};
 use tokio::runtime::Runtime;
 
+use crate::config::CometNativeOptions;
 use crate::execution::fair_memory_pool::CometFairMemoryPool;
 use crate::execution::operators::ScanExec;
 use crate::execution::shuffle::{read_ipc_compressed, CompressionCodec};
@@ -127,8 +130,6 @@ struct ExecutionContext {
     pub explain_native: bool,
     /// Memory pool config
     pub memory_pool_config: MemoryPoolConfig,
-    /// Apache Spark config
-    pub spark_config: HashMap<String, String>,
 }
 
 #[derive(PartialEq, Eq)]
@@ -249,11 +250,28 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_createPlan(
             local_dirs.push(local_dir.into());
         }
 
-        // We need to keep the session context alive. Some session state like temporary
-        // dictionaries are stored in session context. If it is dropped, the temporary
+        // Read Apache Spark runtime config and memoize them in Comet options
+        let spark_conf_map = JMap::from_env(&mut env, &spark_conf)?;
+        let mut spark_conf_iter = spark_conf_map.iter(&mut env)?;
+        let mut comet_options = CometNativeOptions::default();
+        while let Some((key, value)) = spark_conf_iter.next(&mut env)? {
+            let key: String = env.get_string(&JString::from(key)).unwrap().into();
+            if key == "spark.hadoop.fs.defaultFS" {
+                let value: String = env.get_string(&JString::from(value)).unwrap().into();
+                comet_options.default_fs = Some(value);
+            }
+        }
+
+        // We need to keep the session context alive. Some session states,
+        // like temporary dictionaries, are stored in the
+        // session context. If it is dropped, the temporary
         // dictionaries will be dropped as well.
-        let session =
-            prepare_datafusion_session_context(batch_size as usize, memory_pool, local_dirs)?;
+        let session = prepare_datafusion_session_context(
+            batch_size as usize,
+            memory_pool,
+            local_dirs,
+            Arc::from(comet_options),
+        )?;
 
         let plan_creation_time = start.elapsed();
 
@@ -263,17 +281,6 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_createPlan(
             None
         };
 
-        // Read Apache Spark runtime config
-        let spark_conf_map = JMap::from_env(&mut env, &spark_conf)?;
-        let mut spark_conf_iter = spark_conf_map.iter(&mut env)?;
-        let mut spark_conf = HashMap::new();
-
-        while let Some((key, value)) = spark_conf_iter.next(&mut env)? {
-            let key: String = env.get_string(&JString::from(key)).unwrap().into();
-            let value: String = env.get_string(&JString::from(value)).unwrap().into();
-            spark_conf.insert(key, value);
-        }
-
         let exec_context = Box::new(ExecutionContext {
             id,
             task_attempt_id,
@@ -291,7 +298,6 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_createPlan(
             debug_native: debug_native == 1,
             explain_native: explain_native == 1,
             memory_pool_config,
-            spark_config: spark_conf,
         });
 
         Ok(Box::into_raw(exec_context) as i64)
@@ -303,6 +309,7 @@ fn prepare_datafusion_session_context(
     batch_size: usize,
     memory_pool: Arc<dyn MemoryPool>,
     local_dirs: Vec<String>,
+    comet_opts: Arc<CometNativeOptions>,
 ) -> CometResult<SessionContext> {
     let disk_manager_config =
         DiskManagerConfig::NewSpecified(local_dirs.into_iter().map(PathBuf::from).collect());
@@ -323,7 +330,8 @@ fn prepare_datafusion_session_context(
             // maximum value is 1.0, so we set the threshold a little higher just
             // to be safe
             &ScalarValue::Float64(Some(1.1)),
-        );
+        )
+        .with_extension(comet_opts);
 
     #[allow(deprecated)]
     let runtime = RuntimeEnv::try_new(rt_config)?;

native/core/src/execution/planner.rs

@@ -1113,8 +1113,9 @@ impl PhysicalPlanner {
                     .and_then(|f| f.partitioned_file.first())
                     .map(|f| f.file_path.clone())
                     .ok_or(GeneralError("Failed to locate file".to_string()))?;
+
                 let (object_store_url, _) =
-                    prepare_object_store(self.session_ctx.runtime_env(), one_file)?;
+                    prepare_object_store(Arc::clone(&self.session_ctx), one_file)?;
 
                 // Generate file groups
                 let mut file_groups: Vec<Vec<PartitionedFile>> =
@@ -2501,6 +2502,10 @@ fn create_case_expr(
 
 #[cfg(test)]
 mod tests {
+    use datafusion::config::ConfigExtension;
+    use datafusion::config::ConfigOptions;
+    use datafusion::config::Extensions;
+    use datafusion::prelude::SessionConfig;
     use std::{sync::Arc, task::Poll};
 
     use futures::{poll, StreamExt};
@@ -2953,4 +2958,38 @@ mod tests {
             }
         });
     }
+
+    use datafusion::common::extensions_options;
+
+    extensions_options! {
+        struct CometNativeOptions {
+            // Default File System
+            default_fs: Option<String>, default = Some("123".to_string())
+        }
+    }
+
+    impl ConfigExtension for CometNativeOptions {
+        const PREFIX: &'static str = "comet";
+    }
+
+    #[test]
+    fn test_create_array1() {
+        //let mut extensions = Extensions::new();
+        //extensions.insert(CometNativeOptions::default());
+
+        #[derive(Debug)]
+        struct Ext1 {
+            default_fs: String,
+        }
+
+        let ext1a = Arc::new(Ext1 {
+            default_fs: "123".to_string(),
+        });
+
+        //let config = SessionConfig::from(ConfigOptions::new().with_extensions(extensions));
+        let config = SessionConfig::default().with_extension(Arc::clone(&ext1a));
+        let session_ctx = SessionContext::new_with_config(config);
+
+        dbg!(session_ctx.state().config().get_extension::<Ext1>());
+    }
 }

native/core/src/lib.rs

@@ -46,6 +46,7 @@ use errors::{try_unwrap_or_throw, CometError, CometResult};
 mod errors;
 #[macro_use]
 pub mod common;
+mod config;
 pub mod execution;
 mod jvm_bridge;
 pub mod parquet;

native/core/src/parquet/mod.rs

@@ -26,11 +26,11 @@ pub mod parquet_support;
 pub mod read;
 pub mod schema_adapter;
 
+use crate::errors::{try_unwrap_or_throw, CometError};
+use std::collections::HashMap;
 use std::task::Poll;
 use std::{boxed::Box, ptr::NonNull, sync::Arc};
 
-use crate::errors::{try_unwrap_or_throw, CometError};
-
 use arrow::ffi::FFI_ArrowArray;
 
 /// JNI exposed methods
@@ -59,7 +59,9 @@ use datafusion::execution::SendableRecordBatchStream;
 use datafusion::physical_plan::ExecutionPlan;
 use datafusion::prelude::{SessionConfig, SessionContext};
 use futures::{poll, StreamExt};
-use jni::objects::{JBooleanArray, JByteArray, JLongArray, JPrimitiveArray, JString, ReleaseMode};
+use jni::objects::{
+    JBooleanArray, JByteArray, JLongArray, JMap, JObject, JPrimitiveArray, JString, ReleaseMode,
+};
 use jni::sys::jstring;
 use object_store::path::Path;
 use read::ColumnReader;
@@ -657,6 +659,7 @@ pub unsafe extern "system" fn Java_org_apache_comet_parquet_Native_initRecordBat
     data_schema: jbyteArray,
     session_timezone: jstring,
     batch_size: jint,
+    spark_conf: JObject,
 ) -> jlong {
     try_unwrap_or_throw(&e, |mut env| unsafe {
         let session_config = SessionConfig::new().with_batch_size(batch_size as usize);
@@ -670,7 +673,7 @@ pub unsafe extern "system" fn Java_org_apache_comet_parquet_Native_initRecordBat
             .into();
 
         let (object_store_url, object_store_path) =
-            prepare_object_store(session_ctx.runtime_env(), path.clone())?;
+            prepare_object_store(Arc::clone(session_ctx), path.clone())?;
 
         let required_schema_array = JByteArray::from_raw(required_schema);
         let required_schema_buffer = env.convert_byte_array(&required_schema_array)?;
@@ -706,6 +709,17 @@ pub unsafe extern "system" fn Java_org_apache_comet_parquet_Native_initRecordBat
             .unwrap()
             .into();
 
+        // Read Apache Spark runtime config
+        let spark_conf_map = JMap::from_env(&mut env, &spark_conf)?;
+        let mut spark_conf_iter = spark_conf_map.iter(&mut env)?;
+        let mut spark_conf = HashMap::new();
+
+        while let Some((key, value)) = spark_conf_iter.next(&mut env)? {
+            let key: String = env.get_string(&JString::from(key)).unwrap().into();
+            let value: String = env.get_string(&JString::from(value)).unwrap().into();
+            spark_conf.insert(key, value);
+        }
+
         let scan = init_datasource_exec(
             required_schema,
             Some(data_schema),

native/core/src/parquet/parquet_support.rs

@@ -15,7 +15,9 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use crate::config::CometNativeOptions;
 use crate::execution::operators::ExecutionError;
+use crate::parquet::SessionContext;
 use arrow::array::{new_null_array, DictionaryArray, StructArray};
 use arrow::datatypes::DataType;
 use arrow::{
@@ -25,7 +27,6 @@ use arrow::{
 };
 use datafusion::common::{Result as DataFusionResult, ScalarValue};
 use datafusion::execution::object_store::ObjectStoreUrl;
-use datafusion::execution::runtime_env::RuntimeEnv;
 use datafusion::physical_plan::ColumnarValue;
 use datafusion_comet_spark_expr::utils::array_with_timezone;
 use datafusion_comet_spark_expr::EvalMode;
@@ -213,7 +214,7 @@ fn cast_struct_to_struct(
     }
 }
 
-// Mirrors object_store::parse::parse_url for the hdfs object store
+// Mirrors object_store::parse::parse_hdfs_url for the hdfs object store
 #[cfg(feature = "hdfs")]
 fn parse_hdfs_url(url: &Url) -> Result<(Box<dyn ObjectStore>, Path), object_store::Error> {
     match datafusion_comet_objectstore_hdfs::object_store::hdfs::HadoopFileSystem::new(url.as_ref())
@@ -239,24 +240,36 @@ fn parse_hdfs_url(_url: &Url) -> Result<(Box<dyn ObjectStore>, Path), object_sto
 
 /// Parses the url, registers the object store, and returns a tuple of the object store url and object store path
 pub(crate) fn prepare_object_store(
-    runtime_env: Arc<RuntimeEnv>,
+    session_ctx: Arc<SessionContext>,
     url: String,
 ) -> Result<(ObjectStoreUrl, Path), ExecutionError> {
     let mut url = Url::parse(url.as_str())
         .map_err(|e| ExecutionError::GeneralError(format!("Error parsing URL {url}: {e}")))?;
+
     let mut scheme = url.scheme();
+
     if scheme == "s3a" {
         scheme = "s3";
         url.set_scheme("s3").map_err(|_| {
             ExecutionError::GeneralError("Could not convert scheme from s3a to s3".to_string())
         })?;
     }
+
     let url_key = format!(
         "{}://{}",
         scheme,
         &url[url::Position::BeforeHost..url::Position::AfterPort],
     );
 
+    dbg!(&url_key);
+
+    // let x = match session_ctx.state().config().get_extension::<CometNativeOptions>() {
+    //     Some(opts) if scheme.is_empty() && opts.default_fs.is_some() => {
+    //         Url::parse(format!("{}/{}", opts.default_fs.as_deref().unwrap()))
+    //     },
+    //     _ => url
+    // };
+
     let (object_store, object_store_path): (Box<dyn ObjectStore>, Path) = if scheme == "hdfs" {
         parse_hdfs_url(&url)
     } else {
@@ -265,15 +278,20 @@ pub(crate) fn prepare_object_store(
     .map_err(|e| ExecutionError::GeneralError(e.to_string()))?;
 
     let object_store_url = ObjectStoreUrl::parse(url_key.clone())?;
-    runtime_env.register_object_store(&url, Arc::from(object_store));
+    session_ctx
+        .runtime_env()
+        .register_object_store(&url, Arc::from(object_store));
     Ok((object_store_url, object_store_path))
 }
 
 #[cfg(test)]
 mod tests {
+    use crate::config::CometNativeOptions;
     use crate::parquet::parquet_support::prepare_object_store;
+    use datafusion::common::ScalarValue;
     use datafusion::execution::object_store::ObjectStoreUrl;
     use datafusion::execution::runtime_env::RuntimeEnv;
+    use datafusion::prelude::{SessionConfig, SessionContext};
     use object_store::path::Path;
     use std::sync::Arc;
     use url::Url;
@@ -306,7 +324,7 @@ mod tests {
 
         for (i, url_str) in all_urls.iter().enumerate() {
             let url = &Url::parse(url_str).unwrap();
-            let res = prepare_object_store(Arc::new(RuntimeEnv::default()), url.to_string());
+            let res = prepare_object_store(Arc::new(SessionContext::default()), url.to_string());
 
             let expected = expected.get(i).unwrap();
             match expected {
@@ -327,21 +345,15 @@ mod tests {
     }
 
     #[test]
-    #[cfg(feature = "hdfs")]
-    fn test_prepare_object_store() {
-        // we use a local file system url instead of an hdfs url because the latter requires
-        // a running namenode
-        let hdfs_url = "file:///comet/spark-warehouse/part-00000.snappy.parquet";
-        let expected: (ObjectStoreUrl, Path) = (
-            ObjectStoreUrl::parse("file://").unwrap(),
-            Path::from("/comet/spark-warehouse/part-00000.snappy.parquet"),
-        );
-
-        let url = &Url::parse(hdfs_url).unwrap();
-        let res = prepare_object_store(Arc::new(RuntimeEnv::default()), url.to_string());
-
-        let res = res.unwrap();
-        assert_eq!(res.0, expected.0);
-        assert_eq!(res.1, expected.1);
+    fn test_prepare_object_store_with_default_fs() {
+        let unqualified_url = "file1.parquet";
+        let comet_options = CometNativeOptions {
+            default_fs: Some("hdfs://namenode:9000".to_string()),
+            ..Default::default()
+        };
+        let session_config = SessionConfig::new().with_extension(Arc::from(comet_options));
+        let session = SessionContext::new_with_config(session_config);
+        let object_store = prepare_object_store(Arc::new(session), unqualified_url.to_string());
+        dbg!(object_store);
     }
 }