diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index 2e98bf9aa7a..04ab7f3ebb6 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -114,6 +114,7 @@ use crate::execution::spark_config::{ COMET_TRACING_ENABLED, SPARK_EXECUTOR_CORES, }; use crate::parquet::encryption_support::{CometEncryptionFactory, ENCRYPTION_FACTORY_ID}; +use crate::parquet::parquet_support::CometObjectStoreRegistry; use datafusion_comet_proto::spark_operator::operator::OpStruct; use log::info; use std::sync::OnceLock; @@ -569,7 +570,9 @@ fn prepare_datafusion_session_context( let disk_manager = DiskManagerBuilder::default() .with_mode(DiskManagerMode::Directories(paths)) .with_max_temp_directory_size(max_temp_directory_size); - let mut rt_config = RuntimeEnvBuilder::new().with_disk_manager_builder(disk_manager); + let mut rt_config = RuntimeEnvBuilder::new() + .with_disk_manager_builder(disk_manager) + .with_object_store_registry(Arc::new(CometObjectStoreRegistry::default())); rt_config = rt_config.with_memory_pool(memory_pool); let mut session_config = SessionConfig::new() diff --git a/native/core/src/parquet/parquet_support.rs b/native/core/src/parquet/parquet_support.rs index 2ee1230ed87..476f3a802a9 100644 --- a/native/core/src/parquet/parquet_support.rs +++ b/native/core/src/parquet/parquet_support.rs @@ -31,7 +31,9 @@ use arrow::{ }; use datafusion::common::{Result as DataFusionResult, ScalarValue}; use datafusion::error::DataFusionError; -use datafusion::execution::object_store::ObjectStoreUrl; +use datafusion::execution::object_store::{ + DefaultObjectStoreRegistry, ObjectStoreRegistry, ObjectStoreUrl, +}; use datafusion::execution::runtime_env::RuntimeEnv; use datafusion::physical_plan::ColumnarValue; use datafusion_comet_spark_expr::EvalMode; @@ -439,6 +441,51 @@ fn is_azure_scheme(scheme: &str) -> bool { matches!(scheme, "abfs" | "abfss") } +fn object_store_url_key(url: &Url) -> String { + let authority_start = if is_azure_scheme(url.scheme()) { + // ABFS URLs encode the container in the userinfo + url::Position::BeforeUsername + } else { + url::Position::BeforeHost + }; + format!( + "{}://{}", + url.scheme(), + &url[authority_start..url::Position::AfterPort], + ) +} + +#[derive(Debug, Default)] +pub(crate) struct CometObjectStoreRegistry { + default: DefaultObjectStoreRegistry, + azure_stores: parking_lot::RwLock>>, +} + +impl ObjectStoreRegistry for CometObjectStoreRegistry { + fn register_store( + &self, + url: &Url, + store: Arc, + ) -> Option> { + if is_azure_scheme(url.scheme()) { + self.azure_stores + .write() + .insert(object_store_url_key(url), store) + } else { + self.default.register_store(url, store) + } + } + + fn get_store(&self, url: &Url) -> DataFusionResult> { + if is_azure_scheme(url.scheme()) { + if let Some(store) = self.azure_stores.read().get(&object_store_url_key(url)) { + return Ok(Arc::clone(store)); + } + } + self.default.get_store(url) + } +} + // Creates an OpenDAL HDFS Operator from a URL with optional configuration #[cfg(feature = "hdfs-opendal")] pub(crate) fn create_hdfs_operator(url: &Url) -> Result { @@ -498,7 +545,8 @@ fn create_hdfs_object_store( type ObjectStoreCache = RwLock>>; -/// Process-wide cache of object stores, keyed by `(scheme://host:port, config_hash)`. +/// Process-wide cache of object stores, keyed by +/// `(scheme://[container@]host:port, config_hash)`. /// /// ## Why static / process lifetime? /// @@ -510,12 +558,15 @@ type ObjectStoreCache = RwLock>>; /// deployment model each executor process is dedicated to a single Spark application, so /// process lifetime and application lifetime are equivalent; the cache is reclaimed when /// the executor pod terminates. +/// Per-container isolation in shared native DataFusion runtimes also depends on +/// `CometObjectStoreRegistry` using the same ABFS-aware key. /// /// ## Unbounded size /// -/// Cache entries are indexed by `(scheme://host:port, hash-of-configs)`. A typical Spark -/// job accesses a small, fixed set of buckets with a stable configuration, so the number of -/// distinct keys is O(buckets × credential-configs) and remains small throughout the job. +/// Cache entries are indexed by `(scheme://[container@]host:port, hash-of-configs)`. A typical +/// Spark job accesses a small, fixed set of buckets or containers with a stable configuration, +/// so the number of distinct keys is O(buckets/containers × credential-configs) and remains small +/// throughout the job. /// Entries are cheap relative to the cost of creating a new object store (new HTTP /// connection pool + DNS resolution), and there is no meaningful benefit from eviction, so /// no eviction policy is applied. @@ -563,11 +614,7 @@ pub(crate) fn prepare_object_store_with_configs( ExecutionError::GeneralError("Could not convert scheme from s3a to s3".to_string()) })?; } - let url_key = format!( - "{}://{}", - scheme, - &url[url::Position::BeforeHost..url::Position::AfterPort], - ); + let url_key = object_store_url_key(&url); let config_hash = hash_object_store_configs(object_store_configs); let cache_key = (url_key.clone(), config_hash); @@ -617,18 +664,15 @@ pub(crate) fn prepare_object_store_with_configs( mod tests { #[cfg(not(feature = "hdfs-opendal"))] use datafusion::execution::object_store::ObjectStoreUrl; - #[cfg(not(feature = "hdfs-opendal"))] use datafusion::execution::runtime_env::RuntimeEnv; #[cfg(not(feature = "hdfs-opendal"))] use object_store::path::Path; - #[cfg(not(feature = "hdfs-opendal"))] use std::sync::Arc; #[cfg(not(feature = "hdfs-opendal"))] use url::Url; #[cfg(not(feature = "hdfs-opendal"))] use crate::execution::operators::ExecutionError; - #[cfg(not(feature = "hdfs-opendal"))] use std::collections::HashMap; /// Parses the url, registers the object store, and returns a tuple of the object store url and object store path @@ -683,4 +727,75 @@ mod tests { } } } + + #[test] + fn test_azure_containers_use_distinct_stores() { + let configs = HashMap::from([("fs.azure.account.key".into(), "c2VjcmV0".into())]); + let object_store = |container| { + let runtime_env = Arc::new(RuntimeEnv::default()); + let (object_store_url, _) = super::prepare_object_store_with_configs( + Arc::clone(&runtime_env), + format!("abfss://{container}@myacct.dfs.core.windows.net/path/file.parquet"), + &configs, + ) + .unwrap(); + runtime_env.object_store(&object_store_url).unwrap() + }; + + let container_a = object_store("container-a"); + assert!(Arc::ptr_eq(&container_a, &object_store("container-a"))); + assert!(!Arc::ptr_eq(&container_a, &object_store("container-b"))); + } + + #[test] + fn test_shared_runtime_env_uses_distinct_azure_container_stores() { + use datafusion::execution::runtime_env::RuntimeEnvBuilder; + + let configs = HashMap::from([("fs.azure.account.key".into(), "c2VjcmV0".into())]); + let runtime_env = Arc::new( + RuntimeEnvBuilder::new() + .with_object_store_registry(Arc::new(super::CometObjectStoreRegistry::default())) + .build() + .unwrap(), + ); + let register = |container| { + let (url, _) = super::prepare_object_store_with_configs( + Arc::clone(&runtime_env), + format!("abfss://{container}@myacct.dfs.core.windows.net/path/file.parquet"), + &configs, + ) + .unwrap(); + runtime_env.object_store(&url).unwrap() + }; + + let container_a = register("container-a"); + let container_b = register("container-b"); + assert!(!Arc::ptr_eq(&container_a, &container_b)); + } + + #[test] + fn test_s3_store_cache_keys_by_host() { + let configs = HashMap::from([ + ( + "fs.s3a.aws.credentials.provider".into(), + "org.apache.hadoop.fs.s3a.AnonymousAWSCredentialsProvider".into(), + ), + ("fs.s3a.endpoint.region".into(), "us-east-1".into()), + ]); + let object_store = |bucket| { + let runtime_env = Arc::new(RuntimeEnv::default()); + let (object_store_url, _) = super::prepare_object_store_with_configs( + Arc::clone(&runtime_env), + format!("s3://{bucket}@shared-host/path/file.parquet"), + &configs, + ) + .unwrap(); + runtime_env.object_store(&object_store_url).unwrap() + }; + + assert!(Arc::ptr_eq( + &object_store("bucket-a"), + &object_store("bucket-b") + )); + } }