diff --git a/bfg.jar b/bfg.jar new file mode 100644 index 0000000000000..688fe713674b9 Binary files /dev/null and b/bfg.jar differ diff --git a/plugins/engine-datafusion/Cargo.toml b/plugins/engine-datafusion/Cargo.toml index 888170a176e53..21d795a9bd9d0 100644 --- a/plugins/engine-datafusion/Cargo.toml +++ b/plugins/engine-datafusion/Cargo.toml @@ -6,21 +6,21 @@ members = [ [workspace.dependencies] # DataFusion dependencies -datafusion = "49.0.0" -datafusion-expr = "49.0.0" -datafusion-datasource = "49.0.0" -arrow-json = "55.2" -arrow = { version = "55.2", features = ["ffi", "ipc_compression"] } +datafusion = "50.0.0" +datafusion-expr = "50.0.0" +datafusion-datasource = "50.0.0" +arrow-json = "56.2" +arrow = { version = "56.2", features = ["ffi", "ipc_compression"] } #arrow = "55.2.0" -arrow-array = "55.2.0" -arrow-schema = "55.2.0" -arrow-buffer = "55.2.0" +arrow-array = "56.2.0" +arrow-schema = "56.2.0" +arrow-buffer = "56.2.0" # JNI dependencies jni = "0.21" # Substrait support -datafusion-substrait = "49.0.0" +datafusion-substrait = "50.0.0" prost = "0.13" diff --git a/plugins/engine-datafusion/jni/Cargo.toml b/plugins/engine-datafusion/jni/Cargo.toml index 828c389b9bff3..421111cdb2f86 100644 --- a/plugins/engine-datafusion/jni/Cargo.toml +++ b/plugins/engine-datafusion/jni/Cargo.toml @@ -42,6 +42,10 @@ serde_json = { workspace = true } anyhow = { workspace = true } thiserror = { workspace = true } +# Cache policy dependencies +instant = "0.1" +dashmap = "5.5" + # Logging log ={ workspace = true } # Parquet support diff --git a/plugins/engine-datafusion/jni/src/cache_policy.rs b/plugins/engine-datafusion/jni/src/cache_policy.rs new file mode 100644 index 0000000000000..3d33e200b492e --- /dev/null +++ b/plugins/engine-datafusion/jni/src/cache_policy.rs @@ -0,0 +1,406 @@ +//! # Cache Policy Module +//! +//! Simple pluggable cache eviction policies for statistics cache. + +use instant::Instant; +use thiserror::Error; + +/// Error types for cache operations +#[derive(Debug, Error)] +pub enum CacheError { + #[error("Policy lock error: {reason}")] + PolicyLockError { reason: String }, +} + +/// Result type for cache operations +pub type CacheResult = Result; + +/// Core trait for cache eviction policies +pub trait CachePolicy: Send + Sync { + /// Called when a cache entry is accessed + fn on_access(&mut self, key: &str, size: usize); + + /// Called when a cache entry is inserted + fn on_insert(&mut self, key: &str, size: usize); + + /// Called when a cache entry is removed + fn on_remove(&mut self, key: &str); + + /// Select entries for eviction to reach target size + /// Returns keys to evict, ordered by eviction priority + fn select_for_eviction(&self, target_size: usize) -> Vec; + + /// Reset policy state + fn clear(&mut self); + + /// Get the name of this policy + fn policy_name(&self) -> &'static str; +} + +/// Policy types +#[derive(Debug, Clone)] +pub enum PolicyType { + Lru, + Lfu, +} + +/// Cache configuration +#[derive(Debug, Clone)] +pub struct CacheConfig { + pub policy_type: PolicyType, + pub size_limit: usize, + pub eviction_threshold: f64, +} + +impl Default for CacheConfig { + fn default() -> Self { + Self { + policy_type: PolicyType::Lru, + size_limit: 20 * 1024 * 1024, // 20MB + eviction_threshold: 0.8, + } + } +} + +/// Simple cache entry metadata +#[derive(Debug, Clone)] +pub struct CacheEntryMetadata { + pub size: usize, + pub last_accessed: Instant, + pub access_count: usize, +} + +impl CacheEntryMetadata { + pub fn new(_key: String, size: usize) -> Self { + let now = Instant::now(); + Self { + size, + last_accessed: now, + access_count: 1, + } + } + + pub fn on_access(&mut self) { + self.last_accessed = Instant::now(); + self.access_count += 1; + } +} + +/// LRU (Least Recently Used) policy +pub struct LruPolicy { + entries: dashmap::DashMap, + total_size: std::sync::atomic::AtomicUsize, +} + +impl LruPolicy { + pub fn new() -> Self { + Self { + entries: dashmap::DashMap::new(), + total_size: std::sync::atomic::AtomicUsize::new(0), + } + } +} + +impl Default for LruPolicy { + fn default() -> Self { + Self::new() + } +} + +impl CachePolicy for LruPolicy { + fn on_access(&mut self, key: &str, size: usize) { + match self.entries.get_mut(key) { + Some(mut entry) => { + entry.on_access(); + } + None => { + let metadata = CacheEntryMetadata::new(key.to_string(), size); + self.entries.insert(key.to_string(), metadata); + self.total_size + .fetch_add(size, std::sync::atomic::Ordering::Relaxed); + } + } + } + + fn on_insert(&mut self, key: &str, size: usize) { + let metadata = CacheEntryMetadata::new(key.to_string(), size); + + if let Some(old_entry) = self.entries.insert(key.to_string(), metadata) { + let old_size = old_entry.size; + self.total_size + .fetch_sub(old_size, std::sync::atomic::Ordering::Relaxed); + } + + self.total_size + .fetch_add(size, std::sync::atomic::Ordering::Relaxed); + } + + fn on_remove(&mut self, key: &str) { + if let Some((_, entry)) = self.entries.remove(key) { + self.total_size + .fetch_sub(entry.size, std::sync::atomic::Ordering::Relaxed); + } + } + + fn select_for_eviction(&self, target_size: usize) -> Vec { + println!("info seleectpon"); + if target_size == 0 { + return Vec::new(); + } + + // Collect entries with access times + let mut entries: Vec<_> = self + .entries + .iter() + .map(|entry| { + let key = entry.key().clone(); + let last_accessed = entry.value().last_accessed; + (key, last_accessed) + }) + .collect(); + + // Sort by access time (oldest first) + entries.sort_by_key(|(_, last_accessed)| *last_accessed); + + // Select entries for eviction until target size is reached + let mut candidates = Vec::new(); + let mut freed_size = 0; + + for (key, _) in entries { + if freed_size >= target_size { + break; + } + + if let Some(entry) = self.entries.get(&key) { + freed_size += entry.size; + candidates.push(key.clone()); + println!("Selected :{}",key); + } + } + + candidates + } + + fn clear(&mut self) { + self.entries.clear(); + self.total_size + .store(0, std::sync::atomic::Ordering::Relaxed); + } + + fn policy_name(&self) -> &'static str { + "lru" + } +} + +/// LFU (Least Frequently Used) policy +pub struct LfuPolicy { + entries: dashmap::DashMap, + total_size: std::sync::atomic::AtomicUsize, +} + +impl LfuPolicy { + pub fn new() -> Self { + Self { + entries: dashmap::DashMap::new(), + total_size: std::sync::atomic::AtomicUsize::new(0), + } + } +} + +impl Default for LfuPolicy { + fn default() -> Self { + Self::new() + } +} + +impl CachePolicy for LfuPolicy { + fn on_access(&mut self, key: &str, size: usize) { + match self.entries.get_mut(key) { + Some(mut entry) => { + entry.on_access(); + } + None => { + let metadata = CacheEntryMetadata::new(key.to_string(), size); + self.entries.insert(key.to_string(), metadata); + self.total_size + .fetch_add(size, std::sync::atomic::Ordering::Relaxed); + } + } + } + + fn on_insert(&mut self, key: &str, size: usize) { + let metadata = CacheEntryMetadata::new(key.to_string(), size); + + if let Some(old_entry) = self.entries.insert(key.to_string(), metadata) { + let old_size = old_entry.size; + self.total_size + .fetch_sub(old_size, std::sync::atomic::Ordering::Relaxed); + } + + self.total_size + .fetch_add(size, std::sync::atomic::Ordering::Relaxed); + } + + fn on_remove(&mut self, key: &str) { + if let Some((_, entry)) = self.entries.remove(key) { + self.total_size + .fetch_sub(entry.size, std::sync::atomic::Ordering::Relaxed); + } + } + + fn select_for_eviction(&self, target_size: usize) -> Vec { + if target_size == 0 { + return Vec::new(); + } + + // Collect entries with access counts + let mut entries: Vec<_> = self + .entries + .iter() + .map(|entry| { + let key = entry.key().clone(); + let access_count = entry.value().access_count; + let last_accessed = entry.value().last_accessed; + (key, access_count, last_accessed) + }) + .collect(); + + // Sort by access count (least frequent first), then by time for tie-breaking + entries.sort_by(|(_, count_a, time_a), (_, count_b, time_b)| { + count_a.cmp(count_b).then(time_a.cmp(time_b)) + }); + + // Select entries for eviction until target size is reached + let mut candidates = Vec::new(); + let mut freed_size = 0; + + for (key, _, _) in entries { + if freed_size >= target_size { + break; + } + + if let Some(entry) = self.entries.get(&key) { + freed_size += entry.size; + candidates.push(key); + } + } + + candidates + } + + fn clear(&mut self) { + self.entries.clear(); + self.total_size + .store(0, std::sync::atomic::Ordering::Relaxed); + } + + fn policy_name(&self) -> &'static str { + "lfu" + } +} + +/// Create a cache policy instance +pub fn create_policy(policy_type: PolicyType) -> Box { + match policy_type { + PolicyType::Lru => Box::new(LruPolicy::new()), + PolicyType::Lfu => Box::new(LfuPolicy::new()), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::thread; + use std::time::Duration; + + #[test] + fn test_cache_config_default() { + let config = CacheConfig::default(); + assert!(matches!(config.policy_type, PolicyType::Lru)); + assert_eq!(config.size_limit, 100 * 1024 * 1024); + assert_eq!(config.eviction_threshold, 0.8); + } + + #[test] + fn test_cache_entry_metadata() { + let mut metadata = CacheEntryMetadata::new("test_key".to_string(), 1024); + assert_eq!(metadata.size, 1024); + assert_eq!(metadata.access_count, 1); + + let initial_access_time = metadata.last_accessed; + thread::sleep(Duration::from_millis(1)); + + metadata.on_access(); + assert_eq!(metadata.access_count, 2); + assert!(metadata.last_accessed > initial_access_time); + } + + #[test] + fn test_create_policy() { + let lru_policy = create_policy(PolicyType::Lru); + assert_eq!(lru_policy.policy_name(), "lru"); + + let lfu_policy = create_policy(PolicyType::Lfu); + assert_eq!(lfu_policy.policy_name(), "lfu"); + } + + #[test] + fn test_lru_policy_basic_operations() { + let mut policy = LruPolicy::new(); + + assert_eq!(policy.policy_name(), "lru"); + + policy.on_insert("key1", 100); + policy.on_insert("key2", 200); + + policy.on_access("key1", 100); + + policy.on_remove("key1"); + + policy.clear(); + } + + #[test] + fn test_lru_policy_victim_selection() { + let mut policy = LruPolicy::new(); + + policy.on_insert("oldest", 100); + thread::sleep(Duration::from_millis(1)); + + policy.on_insert("middle", 100); + thread::sleep(Duration::from_millis(1)); + + policy.on_insert("newest", 100); + thread::sleep(Duration::from_millis(1)); + + // Access middle entry to make it more recent + policy.on_access("middle", 100); + + let candidates = policy.select_for_eviction(150); + assert_eq!(candidates.len(), 2); + assert!(candidates.contains(&"oldest".to_string())); + assert!(!candidates.contains(&"middle".to_string())); + } + + #[test] + fn test_lfu_policy_victim_selection() { + let mut policy = LfuPolicy::new(); + + policy.on_insert("rarely_used", 100); + policy.on_insert("sometimes_used", 100); + policy.on_insert("frequently_used", 100); + + // Create frequency patterns + policy.on_access("sometimes_used", 100); + + for _ in 0..3 { + policy.on_access("frequently_used", 100); + } + + let candidates = policy.select_for_eviction(150); + assert_eq!(candidates.len(), 2); + assert!(candidates.contains(&"rarely_used".to_string())); + assert!(candidates.contains(&"sometimes_used".to_string())); + assert!(!candidates.contains(&"frequently_used".to_string())); + } +} diff --git a/plugins/engine-datafusion/jni/src/lib.rs b/plugins/engine-datafusion/jni/src/lib.rs index f3e40560ee0ae..adbee0a5643a7 100644 --- a/plugins/engine-datafusion/jni/src/lib.rs +++ b/plugins/engine-datafusion/jni/src/lib.rs @@ -5,34 +5,39 @@ * this file be licensed under the Apache-2.0 license or a * compatible open source license. */ -use arrow_array::ffi::FFI_ArrowArray; -use arrow_array::{Array, StructArray}; -use arrow_schema::ffi::FFI_ArrowSchema; +use std::ptr::addr_of_mut; +use datafusion_expr::expr_rewriter::unalias; use jni::objects::{JByteArray, JClass, JObject}; use jni::sys::{jbyteArray, jlong, jstring}; use jni::JNIEnv; -use std::ptr::addr_of_mut; -use std::sync::Arc; +use std::sync::{Arc, Mutex}; +use arrow_array::{Array, StructArray}; +use arrow_array::ffi::FFI_ArrowArray; +use arrow_schema::DataType; +use arrow_schema::ffi::FFI_ArrowSchema; use std::time::Instant; mod util; mod row_id_optimizer; mod listing_table; +mod metadata_cache; +mod statistics_cache; +mod cache_policy; use datafusion::execution::context::SessionContext; - -use crate::listing_table::{ListingOptions, ListingTable, ListingTableConfig}; -use crate::util::{create_object_meta_from_filenames, parse_string_arr, set_object_result_error, set_object_result_ok}; -use datafusion::datasource::file_format::csv::CsvFormat; -use datafusion::datasource::file_format::parquet::ParquetFormat; -use datafusion::datasource::listing::ListingTableUrl; -use datafusion::execution::cache::cache_manager::CacheManagerConfig; -use datafusion::execution::cache::cache_unit::DefaultListFilesCache; -use datafusion::execution::cache::CacheAccessor; -use datafusion::execution::runtime_env::{RuntimeEnv, RuntimeEnvBuilder}; -use datafusion::physical_plan::SendableRecordBatchStream; +use datafusion::physical_plan::Statistics; +use datafusion::common::stats::Precision; + +use crate::cache_policy::{CacheConfig, PolicyType}; +use crate::statistics_cache::CustomStatisticsCache; +use crate::util::{ + construct_file_metadata, create_object_meta_from_file, create_object_meta_from_filenames, + parse_string_arr, set_object_result_error, set_object_result_ok, +}; use datafusion::prelude::SessionConfig; use datafusion::DATAFUSION_VERSION; +use datafusion::datasource::file_format::parquet::ParquetFormat; +use datafusion::physical_plan::SendableRecordBatchStream; use datafusion_substrait::logical_plan::consumer::from_substrait_plan; use datafusion_substrait::substrait::proto::Plan; use futures::TryStreamExt; @@ -40,11 +45,19 @@ use jni::objects::{JObjectArray, JString}; use object_store::ObjectMeta; use prost::Message; use tokio::runtime::Runtime; -use std::thread; -use std::time::{Duration, SystemTime, UNIX_EPOCH}; -use std::fs::OpenOptions; -use std::io::Write; -use std::sync::atomic::{AtomicBool, Ordering}; +use crate::listing_table::{ListingOptions, ListingTable, ListingTableConfig}; +use crate::row_id_optimizer::FilterRowIdOptimizer; +use crate::metadata_cache::{MutexFileMetadataCache}; +use datafusion::datasource::file_format::csv::CsvFormat; +use datafusion::execution::cache::cache_manager::{self, CacheManagerConfig, FileMetadataCache}; +use datafusion::execution::cache::cache_unit::{ + DefaultFileStatisticsCache, DefaultFilesMetadataCache, DefaultListFilesCache, +}; +use datafusion::datasource::file_format::file_compression_type::FileCompressionType; +use datafusion::datasource::physical_plan::FileMeta; +use datafusion::execution::cache::CacheAccessor; +use datafusion::execution::runtime_env::{RuntimeEnv, RuntimeEnvBuilder}; +use datafusion_datasource::ListingTableUrl; /// Create a new DataFusion session context @@ -100,30 +113,17 @@ pub extern "system" fn Java_org_opensearch_datafusion_DataFusionQueryJNI_createT #[no_mangle] pub extern "system" fn Java_org_opensearch_datafusion_DataFusionQueryJNI_createGlobalRuntime( - _env: JNIEnv, + mut env: JNIEnv, _class: JClass, + cache_config_ptr: jlong, ) -> jlong { - let runtime_env = RuntimeEnvBuilder::default().build().unwrap(); - /** - // We can copy global runtime to local runtime - file statistics cache, and most of the things - // will be shared across session contexts. But list files cache will be specific to session - // context - - let fsCache = runtimeEnv.clone().cache_manager.get_file_statistic_cache().unwrap(); - let localCacheManagerConfig = CacheManagerConfig::default().with_files_statistics_cache(Option::from(fsCache)); - let localCacheManager = CacheManager::try_new(&localCacheManagerConfig); - let localRuntimeEnv = RuntimeEnvBuilder::new() - .with_cache_manager(localCacheManagerConfig) - .with_disk_manager(DiskManagerConfig::new_existing(runtimeEnv.disk_manager)) - .with_memory_pool(runtimeEnv.memory_pool) - .with_object_store_registry(runtimeEnv.object_store_registry) + let cache_manager_config = unsafe { Box::from_raw(cache_config_ptr as *mut CacheManagerConfig) }; + + let runtime_env = RuntimeEnvBuilder::default() + .with_cache_manager(*cache_manager_config) .build(); - let config = SessionConfig::new().with_repartition_aggregations(true); - let context = SessionContext::new_with_config(config); - **/ - let ctx = Box::into_raw(Box::new(runtime_env)) as jlong; - ctx + Box::into_raw(Box::new(runtime_env)) as jlong } #[no_mangle] @@ -167,7 +167,7 @@ pub extern "system" fn Java_org_opensearch_datafusion_DataFusionQueryJNI_createD } #[no_mangle] -pub extern "system" fn Java_org_opensearch_datafusion_DataFusionQueryJNI_destroyReader( +pub extern "system" fn Java_org_opensearch_datafusion_DataFusionQueryJNI_closeDatafusionReader( mut env: JNIEnv, _class: JClass, ptr: jlong @@ -198,7 +198,6 @@ impl ShardView { } } - #[no_mangle] pub extern "system" fn Java_org_opensearch_datafusion_DataFusionQueryJNI_executeSubstraitQuery( mut env: JNIEnv, @@ -206,6 +205,7 @@ pub extern "system" fn Java_org_opensearch_datafusion_DataFusionQueryJNI_execute shard_view_ptr: jlong, table_name: JString, substrait_bytes: jbyteArray, + global_runtime_env_ptr: jlong, tokio_runtime_env_ptr: jlong, // callback: JObject, ) -> jlong { @@ -213,6 +213,7 @@ pub extern "system" fn Java_org_opensearch_datafusion_DataFusionQueryJNI_execute let shard_view = unsafe { &*(shard_view_ptr as *const ShardView) }; let runtime_ptr = unsafe { &*(tokio_runtime_env_ptr as *const Runtime)}; let table_name: String = env.get_string(&table_name).expect("Couldn't get java string!").into(); + let runtime_env = unsafe { &*(global_runtime_env_ptr as *const RuntimeEnv) }; let table_path = shard_view.table_path(); let files_meta = shard_view.files_meta(); @@ -220,14 +221,6 @@ pub extern "system" fn Java_org_opensearch_datafusion_DataFusionQueryJNI_execute println!("Table path: {}", table_path); println!("Files: {:?}", files_meta); - let list_file_cache = Arc::new(DefaultListFilesCache::default()); - list_file_cache.put(table_path.prefix(), files_meta); - - let runtime_env = RuntimeEnvBuilder::new() - .with_cache_manager(CacheManagerConfig::default() - .with_list_files_cache(Some(list_file_cache.clone())) - ).build().unwrap(); - // TODO: get config from CSV DataFormat let mut config = SessionConfig::new(); config.options_mut().execution.parquet.pushdown_filters = false; @@ -235,7 +228,7 @@ pub extern "system" fn Java_org_opensearch_datafusion_DataFusionQueryJNI_execute let state = datafusion::execution::SessionStateBuilder::new() .with_config(config) - .with_runtime_env(Arc::from(runtime_env)) + .with_runtime_env(Arc::new(runtime_env.clone())) .with_default_features() // .with_optimizer_rule(Arc::new(OptimizeRowId)) // .with_physical_optimizer_rule(Arc::new(FilterRowIdOptimizer)) // TODO: enable only for query phase @@ -382,8 +375,6 @@ pub extern "system" fn Java_org_opensearch_datafusion_DataFusionQueryJNI_nativeC session_context_ptr } - - #[no_mangle] pub extern "system" fn Java_org_opensearch_datafusion_RecordBatchStream_next( mut env: JNIEnv, @@ -446,3 +437,789 @@ pub extern "system" fn Java_org_opensearch_datafusion_RecordBatchStream_getSchem } } } + +#[no_mangle] +pub extern "system" fn Java_org_opensearch_datafusion_DataFusionQueryJNI_initCacheManagerConfig( + _env: JNIEnv, + _class: JClass, +) -> jlong { + let config = CacheManagerConfig::default(); + Box::into_raw(Box::new(config)) as jlong +} + + +/// Create a metadata cache and add it to the CacheManagerConfig +/// The config_ptr remains the same, only the contents are updated +#[no_mangle] +pub extern "system" fn Java_org_opensearch_datafusion_DataFusionQueryJNI_createMetadataCache( + _env: JNIEnv, + _class: JClass, + config_ptr: jlong, + size_limit: jlong, +) -> jlong { + // Create cache with wrapper that implements FileMetadataCache + let inner_cache = DefaultFilesMetadataCache::new(size_limit.try_into().unwrap()); + let wrapped_cache = Arc::new(MutexFileMetadataCache::new(inner_cache)); + + // Update the CacheManagerConfig at the same memory location + if config_ptr != 0 { + let cache_manager_config = unsafe { &mut *(config_ptr as *mut CacheManagerConfig) }; + *cache_manager_config = cache_manager_config.clone() + .with_file_metadata_cache(Some(wrapped_cache.clone())); + } + + // Return the Arc pointer for JNI operations + Box::into_raw(Box::new(wrapped_cache)) as jlong +} + +#[no_mangle] +pub extern "system" fn Java_org_opensearch_datafusion_DataFusionQueryJNI_metadataCachePut( + mut env: JNIEnv, + _class: JClass, + cache_ptr: jlong, + file_path: JString, +) -> bool { + let file_path: String = match env.get_string(&file_path) { + Ok(s) => s.into(), + Err(_) => return false + }; + let cache = unsafe { &*(cache_ptr as *const Arc) }; + let data_format = if file_path.to_lowercase().ends_with(".parquet") { + "parquet" + } else { + return false; // Skip unsupported formats + }; + + let object_meta = create_object_meta_from_file(&file_path); + let store = Arc::new(object_store::local::LocalFileSystem::new()); + + // Use Runtime to block on the async operation + let metadata = Runtime::new() + .expect("Failed to create Tokio Runtime") + .block_on(async { + construct_file_metadata(store.as_ref(), &object_meta, data_format) + .await + .expect("Failed to construct file metadata") + }); + + let len_before = cache.len(); + let old_metadata = cache.put(&object_meta, metadata); + let len_after = cache.len(); + + if len_after > len_before { + println!("Successfully cached new metadata for: {} (cache: {} -> {})", file_path, len_before, len_after); + true + } else if old_metadata.is_some() { + println!("Successfully updated existing metadata for: {} (cache: {})", file_path, len_after); + true + } else { + println!("Failed to cache metadata for: {} (cache: {})", file_path, len_after); + false + } +} + +#[no_mangle] +pub extern "system" fn Java_org_opensearch_datafusion_DataFusionQueryJNI_metadataCacheRemove( + mut env: JNIEnv, + _class: JClass, + cache_ptr: jlong, + file_path: JString, +) -> bool { + let file_path: String = match env.get_string(&file_path) { + Ok(s) => s.into(), + Err(_) => return false, + }; + let cache = unsafe { &*(cache_ptr as *const Arc) }; + let object_meta = create_object_meta_from_file(&file_path); + + // Lock the mutex and remove + if let Some(_cache_obj) = cache.inner.lock().unwrap().remove(&object_meta) { + println!("Cache removed for: {}", file_path); + true + } else { + println!("Item not found in cache: {}", file_path); + false + } +} + +#[no_mangle] +pub extern "system" fn Java_org_opensearch_datafusion_DataFusionQueryJNI_metadataCacheGet( + mut env: JNIEnv, + _class: JClass, + cache_ptr: jlong, + file_path: JString, +) -> bool { + let file_path: String = match env.get_string(&file_path) { + Ok(s) => s.into(), + Err(_) => return false, + }; + + let cache = unsafe { &*(cache_ptr as *const Arc) }; + let object_meta = create_object_meta_from_file(&file_path); + + match cache.get(&object_meta) { + Some(metadata) => { + println!("Retrieved metadata for: {} - size: {:?}", file_path, metadata.memory_size()); + true + }, + None => { + println!("No metadata found for: {}", file_path); + false + }, + } +} + +#[no_mangle] +pub extern "system" fn Java_org_opensearch_datafusion_DataFusionQueryJNI_metadataCacheGetEntries( + mut env: JNIEnv, + _class: JClass, + cache_ptr: jlong, +) -> jni::sys::jobjectArray { + let cache = unsafe { &*(cache_ptr as *const Arc) }; + + // Get all entries from the cache + let entries = cache.list_entries(); + + println!("Retrieved {} cache entries", entries.len()); + + // Create String array class + let string_class = match env.find_class("java/lang/String") { + Ok(cls) => cls, + Err(e) => { + println!("Failed to find String class: {:?}", e); + return std::ptr::null_mut(); + } + }; + + // Create array with size = number of entries * 3 (path, size, hit_count for each entry) + let array_size = (entries.len() * 3) as i32; + let result_array = match env.new_object_array(array_size, &string_class, JObject::null()) { + Ok(arr) => arr, + Err(e) => { + println!("Failed to create object array: {:?}", e); + return std::ptr::null_mut(); + } + }; + + // Fill the array with entries + let mut index = 0; + for (path, entry) in entries.iter() { + // Add path + if let Ok(path_str) = env.new_string(path.as_ref()) { + let _ = env.set_object_array_element(&result_array, index, path_str); + } + index += 1; + + // Add size + if let Ok(size_str) = env.new_string(entry.size_bytes.to_string()) { + let _ = env.set_object_array_element(&result_array, index, size_str); + } + index += 1; + + // Add hit_count + if let Ok(hit_count_str) = env.new_string(entry.hits.to_string()) { + let _ = env.set_object_array_element(&result_array, index, hit_count_str); + } + index += 1; + } + + result_array.as_raw() +} + +#[no_mangle] +pub extern "system" fn Java_org_opensearch_datafusion_DataFusionQueryJNI_metadataCacheGetSize( + mut env: JNIEnv, + _class: JClass, + cache_ptr: jlong +) -> usize { + let cache = unsafe { &*(cache_ptr as *const Arc) }; + cache.inner.lock().unwrap().memory_used() +} + +#[no_mangle] +pub extern "system" fn Java_org_opensearch_datafusion_DataFusionQueryJNI_metadataCacheContainsFile( + mut env: JNIEnv, + _class: JClass, + cache_ptr: jlong, + file_path: JString +) -> bool { + let file_path: String = match env.get_string(&file_path) { + Ok(s) => s.into(), + Err(_) => return false + }; + let cache = unsafe { &*(cache_ptr as *const Arc) }; + let object_meta = create_object_meta_from_file(&file_path); + cache.inner.lock().unwrap().contains_key(&object_meta) +} + +#[no_mangle] +pub extern "system" fn Java_org_opensearch_datafusion_DataFusionQueryJNI_metadataCacheUpdateSizeLimit( + env: JNIEnv, + _class: JClass, + cache_ptr: jlong, + new_size_limit: usize +) -> bool { + let cache = unsafe { &*(cache_ptr as *const Arc) }; + cache.inner.lock().unwrap().update_cache_limit(new_size_limit); + if cache.inner.lock().unwrap().cache_limit() == new_size_limit { + println!("Cache size limit updated to: {}", new_size_limit); + true + } else { + println!("Failed to update cache size limit to: {}", new_size_limit); + false + } +} + +#[no_mangle] +pub extern "system" fn Java_org_opensearch_datafusion_DataFusionQueryJNI_metadataCacheClear( + env: JNIEnv, + _class: JClass, + cache_ptr: jlong +) { + let cache = unsafe { &*(cache_ptr as *const Arc) }; + cache.inner.lock().unwrap().clear(); +} + +// STATISTICS CACHE METHODS + +/// Create a statistics cache and add it to the CacheManagerConfig +/// The config_ptr remains the same, only the contents are updated +#[no_mangle] +pub extern "system" fn Java_org_opensearch_datafusion_DataFusionQueryJNI_createStatisticsCache( + _env: JNIEnv, + _class: JClass, + config_ptr: jlong, + size_limit: jlong, +) -> jlong { + // Create memory-aware policy cache with default LRU policy + let config = CacheConfig { + policy_type: PolicyType::Lru, + size_limit: size_limit as usize, + eviction_threshold: 0.8, + }; + let memory_aware_cache = CustomStatisticsCache::new(config); + + // Create a new DefaultFileStatisticsCache for the CacheManagerConfig + let stats_cache = Arc::new(DefaultFileStatisticsCache::default()); + + // Update the CacheManagerConfig at the same memory location + if config_ptr != 0 { + let cache_manager_config = unsafe { &mut *(config_ptr as *mut CacheManagerConfig) }; + *cache_manager_config = cache_manager_config + .clone() + .with_files_statistics_cache(Some(stats_cache)); + } + + // Return the CustomStatisticsCache pointer for JNI operations + Box::into_raw(Box::new(memory_aware_cache)) as jlong +} + +#[no_mangle] +pub extern "system" fn Java_org_opensearch_datafusion_DataFusionQueryJNI_statisticsCachePut( + mut env: JNIEnv, + _class: JClass, + cache_ptr: jlong, + file_path: JString, +) -> bool { + let file_path: String = match env.get_string(&file_path) { + Ok(s) => s.into(), + Err(_) => return false, + }; + + let cache = unsafe { &mut *(cache_ptr as *mut CustomStatisticsCache) }; + let data_format = if file_path.to_lowercase().ends_with(".parquet") { + "parquet" + } else { + return false; // Skip unsupported formats + }; + + // Use DataFusion's parquet statistics construction + let statistics = Runtime::new() + .expect("Failed to create Tokio Runtime") + .block_on(async { + use datafusion::common::stats::ColumnStatistics; + use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; + use parquet::file::reader::{FileReader, SerializedFileReader}; + use std::fs::File; + + match File::open(&file_path) { + Ok(file) => { + match SerializedFileReader::new(file) { + Ok(reader) => { + let metadata = reader.metadata(); + + // Extract row count from all row groups + let mut num_rows = 0; + let mut total_byte_size = 0; + for rg in metadata.row_groups() { + num_rows += rg.num_rows(); + total_byte_size += rg.total_byte_size(); + } + + // Build column statistics from parquet metadata + let mut column_stats = Vec::new(); + let num_columns = metadata.file_metadata().schema_descr().num_columns(); + + for col_idx in 0..num_columns { + let mut col_null_count = 0; + let mut col_distinct_count = None; + let mut has_min_max = false; + + // Aggregate statistics across all row groups for this column + for rg in metadata.row_groups() { + if let Some(col_chunk) = rg.columns().get(col_idx) { + if let Some(stats) = col_chunk.statistics() { + if let Some(null_count) = stats.null_count_opt() { + col_null_count += null_count; + } + if let Some(distinct) = stats.distinct_count() { + col_distinct_count = Some(distinct); + } + has_min_max = stats.has_min_max_set(); + } + } + } + + // Create column statistics + column_stats.push(ColumnStatistics { + null_count: Precision::Exact(col_null_count as usize), + max_value: Precision::Absent, + min_value: Precision::Absent, + distinct_count: col_distinct_count + .map(|count| Precision::Exact(count as usize)) + .unwrap_or(Precision::Absent), + sum_value: Precision::Absent, + }); + } + + Arc::new(Statistics { + num_rows: Precision::Exact(num_rows as usize), + total_byte_size: Precision::Exact(total_byte_size as usize), + column_statistics: column_stats, + }) + } + Err(_) => Arc::new(Statistics::new_unknown(&arrow_schema::Schema::empty())), + } + } + Err(_) => Arc::new(Statistics::new_unknown(&arrow_schema::Schema::empty())), + } + }); + + let path_obj = match object_store::path::Path::from_url_path(&file_path) { + Ok(path) => path, + Err(_) => { + println!("Failed to create path object from: {}", file_path); + return false; + } + }; + + let object_meta = create_object_meta_from_file(&file_path); + + let len_before = cache.len(); + let result = cache.put_with_extra(&path_obj, statistics, &object_meta); + let len_after = cache.len(); + + if len_after > len_before { + println!("Successfully cached new stats for: {} (cache: {} -> {})", file_path, len_before, len_after); + true + } else if result.is_some() { + println!("Successfully updated existing stats for: {} (cache: {})", file_path, len_after); + true + } else { + println!("Failed to cache stats for: {} (cache: {})", file_path, len_after); + false + } +} + +#[no_mangle] +pub extern "system" fn Java_org_opensearch_datafusion_DataFusionQueryJNI_statisticsCacheRemove( + mut env: JNIEnv, + _class: JClass, + cache_ptr: jlong, + file_path: JString, +) -> bool { + let file_path: String = match env.get_string(&file_path) { + Ok(s) => s.into(), + Err(_) => return false, + }; + + let path_obj = match object_store::path::Path::from_url_path(&file_path) { + Ok(path) => path, + Err(_) => { + println!("Failed to create path object from: {}", file_path); + return false; + } + }; + + let cache = unsafe { &mut *(cache_ptr as *mut CustomStatisticsCache) }; + + if let Some(_cache_obj) = cache.remove(&path_obj){ + println!("Cache removed for: {}", file_path); + true + } else { + println!("Item not found in cache: {}", file_path); + false + } +} + +#[no_mangle] +pub extern "system" fn Java_org_opensearch_datafusion_DataFusionQueryJNI_statisticsCacheGet( + mut env: JNIEnv, + _class: JClass, + cache_ptr: jlong, + file_path: JString, +) -> bool { + let file_path: String = match env.get_string(&file_path) { + Ok(s) => s.into(), + Err(_) => return false, + }; + + let cache = unsafe { &*(cache_ptr as *const CustomStatisticsCache) }; + let path_obj = match object_store::path::Path::from_url_path(&file_path) { + Ok(path) => path, + Err(_) => { + println!("Failed to create path object from: {}", file_path); + return false; + } + }; + + match cache.get(&path_obj) { + Some(statistics) => { + println!("Retrieved statistics for: {}", file_path); + true + } + None => { + println!("No statistics found for: {}", file_path); + false + } + } +} + +#[no_mangle] +pub extern "system" fn Java_org_opensearch_datafusion_DataFusionQueryJNI_statisticsCacheGetEntries( + mut env: JNIEnv, + _class: JClass, + cache_ptr: jlong, +) -> jni::sys::jobjectArray { + let cache = unsafe { &*(cache_ptr as *const CustomStatisticsCache) }; + + // Get all entries from the DashMap + let entries: Vec<_> = cache.inner().iter().collect(); + + println!("Retrieved {} statistics cache entries", entries.len()); + + // Create String array class + let string_class = match env.find_class("java/lang/String") { + Ok(cls) => cls, + Err(e) => { + println!("Failed to find String class: {:?}", e); + return std::ptr::null_mut(); + } + }; + + // Create array with just the paths + let array_size = entries.len() as i32; + let result_array = match env.new_object_array(array_size, &string_class, JObject::null()) { + Ok(arr) => arr, + Err(e) => { + println!("Failed to create object array: {:?}", e); + return std::ptr::null_mut(); + } + }; + + // Fill the array with entry paths + for (index, entry) in entries.iter().enumerate() { + let path_str = entry.key().to_string(); + if let Ok(jstr) = env.new_string(path_str) { + let _ = env.set_object_array_element(&result_array, index as i32, jstr); + } + } + + result_array.as_raw() +} + +/// Get current memory usage from statistics cache +#[no_mangle] +pub extern "system" fn Java_org_opensearch_datafusion_DataFusionQueryJNI_statisticsCacheGetSize( + _env: JNIEnv, + _class: JClass, + cache_ptr: jlong, +) -> jlong { + let cache = unsafe { &*(cache_ptr as *const CustomStatisticsCache) }; + + // Return actual memory consumed by the cache entries + cache.memory_consumed() as jlong +} + +#[no_mangle] +pub extern "system" fn Java_org_opensearch_datafusion_DataFusionQueryJNI_statisticsCacheContainsFile( + mut env: JNIEnv, + _class: JClass, + cache_ptr: jlong, + file_path: JString +) -> bool { + let file_path: String = match env.get_string(&file_path) { + Ok(s) => s.into(), + Err(_) => return false, + }; + let cache = unsafe { &*(cache_ptr as *const CustomStatisticsCache) }; + + let path_obj = match object_store::path::Path::from_url_path(&file_path) { + Ok(path) => path, + Err(_) => { + println!("Failed to create path object from: {}", file_path); + return false; + } + }; + cache.inner().contains_key(&path_obj) +} + +/// Update size limit for statistics cache +#[no_mangle] +pub extern "system" fn Java_org_opensearch_datafusion_DataFusionQueryJNI_statisticsCacheUpdateSizeLimit( + mut env: JNIEnv, + _class: JClass, + cache_ptr: jlong, + new_size_limit: jlong, +) -> bool { + let cache = unsafe { &mut *(cache_ptr as *mut CustomStatisticsCache) }; + + // Update size limit using the policy cache's functionality + match cache.update_size_limit(new_size_limit as usize) { + Ok(_) => { + println!("Statistics cache size limit updated to: {}", new_size_limit); + } + Err(_) => { + println!("Failed to update statistics cache size limit to: {}", new_size_limit); + } + } + + true +} + +#[no_mangle] +pub extern "system" fn Java_org_opensearch_datafusion_DataFusionQueryJNI_statisticsCacheClear( + env: JNIEnv, + _class: JClass, + cache_ptr: jlong, +) { + let cache = unsafe { &*(cache_ptr as *const CustomStatisticsCache) }; + cache.clear(); + println!("Statistics cache cleared"); +} + +/// Get hit count from statistics cache +#[no_mangle] +pub extern "system" fn Java_org_opensearch_datafusion_DataFusionQueryJNI_statisticsCacheGetHitCount( + _env: JNIEnv, + _class: JClass, + cache_ptr: jlong, +) -> jlong { + let cache = unsafe { &*(cache_ptr as *const CustomStatisticsCache) }; + cache.hit_count() as jlong +} + +/// Get miss count from statistics cache +#[no_mangle] +pub extern "system" fn Java_org_opensearch_datafusion_DataFusionQueryJNI_statisticsCacheGetMissCount( + _env: JNIEnv, + _class: JClass, + cache_ptr: jlong, +) -> jlong { + let cache = unsafe { &*(cache_ptr as *const CustomStatisticsCache) }; + cache.miss_count() as jlong +} + +/// Get hit rate from statistics cache (returns value between 0.0 and 1.0) +#[no_mangle] +pub extern "system" fn Java_org_opensearch_datafusion_DataFusionQueryJNI_statisticsCacheGetHitRate( + _env: JNIEnv, + _class: JClass, + cache_ptr: jlong, +) -> f64 { + let cache = unsafe { &*(cache_ptr as *const CustomStatisticsCache) }; + cache.hit_rate() +} + +/// Reset hit and miss counters in statistics cache +#[no_mangle] +pub extern "system" fn Java_org_opensearch_datafusion_DataFusionQueryJNI_statisticsCacheResetStats( + _env: JNIEnv, + _class: JClass, + cache_ptr: jlong, +) { + let cache = unsafe { &*(cache_ptr as *const CustomStatisticsCache) }; + cache.reset_stats(); + println!("Statistics cache hit/miss counters reset"); +} + +#[cfg(test)] +mod tests { + use super::*; + use std::thread; + use std::time::Duration; + + #[test] + fn test_concurrent_file_statistics_cache_operations() { + // Create a statistics cache + let stats_cache = Arc::new(DefaultFileStatisticsCache::default()); + + // Create some test paths + let test_paths: Vec = (0..10) + .map(|i| object_store::path::Path::from(format!("/test/file{}.parquet", i))) + .collect(); + + // Pre-populate cache with some statistics + for (i, path) in test_paths.iter().enumerate() { + let stats = Arc::new(Statistics { + num_rows: Precision::Exact(i * 100), + total_byte_size: Precision::Exact(i * 1000), + column_statistics: vec![], + }); + // DefaultFileStatisticsCache requires put_with_extra + let meta = ObjectMeta { + location: path.clone(), + last_modified: chrono::Utc::now(), + size: (i * 1000) as u64, + e_tag: None, + version: None, + }; + stats_cache.put_with_extra(path, stats, &meta); + } + + println!("Initial cache size: {}", stats_cache.len()); + + // Spawn multiple reader threads + let mut reader_handles = vec![]; + for thread_id in 0..5 { + let cache_clone = stats_cache.clone(); + let paths_clone = test_paths.clone(); + + let handle = thread::spawn(move || { + for _ in 0..20 { + for path in &paths_clone { + if let Some(stats) = cache_clone.get(path) { + // Verify we can read the statistics + assert!(stats.num_rows != Precision::Absent); + } + } + thread::sleep(Duration::from_millis(1)); + } + println!("Reader thread {} completed", thread_id); + }); + + reader_handles.push(handle); + } + + // Spawn multiple writer threads that add new entries + let mut writer_handles = vec![]; + for thread_id in 0..3 { + let cache_clone = stats_cache.clone(); + + let handle = thread::spawn(move || { + for i in 0..10 { + let path = object_store::path::Path::from( + format!("/test/writer{}_file{}.parquet", thread_id, i) + ); + let stats = Arc::new(Statistics { + num_rows: Precision::Exact(thread_id * 1000 + i), + total_byte_size: Precision::Exact(thread_id * 10000 + i * 100), + column_statistics: vec![], + }); + let meta = ObjectMeta { + location: path.clone(), + last_modified: chrono::Utc::now(), + size: (thread_id * 10000 + i * 100) as u64, + e_tag: None, + version: None, + }; + cache_clone.put_with_extra(&path, stats, &meta); + thread::sleep(Duration::from_millis(2)); + } + println!("Writer thread {} completed", thread_id); + }); + + writer_handles.push(handle); + } + + // Wait for all threads to complete + for handle in reader_handles { + handle.join().expect("Reader thread panicked"); + } + + for handle in writer_handles { + handle.join().expect("Writer thread panicked"); + } + + println!("Final cache size: {}", stats_cache.len()); + + // Verify cache contains expected entries + assert!(stats_cache.len() >= 10, "Cache should contain at least the initial 10 entries"); + + // Verify we can still read from cache + for path in &test_paths { + assert!(stats_cache.contains_key(path), "Original paths should still be in cache"); + } + + println!("✓ Concurrent operations test passed!"); + } + + #[test] + fn test_arc_get_mut_behavior() { + // Test the Arc::get_mut behavior used in statisticsCacheRemove + let path = object_store::path::Path::from("/test/single.parquet"); + + // Create a single Arc reference + let mut cache_arc = Arc::new(DefaultFileStatisticsCache::default()); + let meta = ObjectMeta { + location: path.clone(), + last_modified: chrono::Utc::now(), + size: 2000, + e_tag: None, + version: None, + }; + cache_arc.put_with_extra(&path, Arc::new(Statistics { + num_rows: Precision::Exact(200), + total_byte_size: Precision::Exact(2000), + column_statistics: vec![], + }), &meta); + + assert!(cache_arc.contains_key(&path), "Path should be in cache"); + + // With single reference, Arc::get_mut should succeed + if let Some(cache_mut) = Arc::get_mut(&mut cache_arc) { + let _ = cache_mut.remove(&path); + println!("✓ Successfully removed with single Arc reference"); + } else { + panic!("Arc::get_mut failed with single reference"); + } + + assert!(!cache_arc.contains_key(&path), "Path should be removed"); + + // Test with multiple references + let mut cache_arc2 = Arc::new(DefaultFileStatisticsCache::default()); + let meta2 = ObjectMeta { + location: path.clone(), + last_modified: chrono::Utc::now(), + size: 3000, + e_tag: None, + version: None, + }; + cache_arc2.put_with_extra(&path, Arc::new(Statistics { + num_rows: Precision::Exact(300), + total_byte_size: Precision::Exact(3000), + column_statistics: vec![], + }), &meta2); + + let _clone = cache_arc2.clone(); // Create second reference + + // With multiple references, Arc::get_mut should fail + if Arc::get_mut(&mut cache_arc2).is_some() { + panic!("Arc::get_mut should fail with multiple references"); + } else { + println!("✓ Arc::get_mut correctly failed with multiple references"); + } + + println!("✓ Arc::get_mut behavior test passed!"); + } +} diff --git a/plugins/engine-datafusion/jni/src/listing_table.rs b/plugins/engine-datafusion/jni/src/listing_table.rs index a28a6292ec3c1..b10e67325136b 100644 --- a/plugins/engine-datafusion/jni/src/listing_table.rs +++ b/plugins/engine-datafusion/jni/src/listing_table.rs @@ -53,7 +53,7 @@ use datafusion::execution::{ use datafusion_expr::{ dml::InsertOp, Expr, SortExpr, TableProviderFilterPushDown, TableType, }; -use datafusion::physical_expr::schema_rewriter::PhysicalExprAdapterFactory; +use datafusion::physical_expr_adapter::schema_rewriter::PhysicalExprAdapterFactory; use datafusion::physical_expr_common::sort_expr::LexOrdering; use datafusion::physical_plan::{empty::EmptyExec, ExecutionPlan, Statistics}; use futures::{future, stream, Stream, StreamExt, TryStreamExt}; diff --git a/plugins/engine-datafusion/jni/src/metadata_cache.rs b/plugins/engine-datafusion/jni/src/metadata_cache.rs new file mode 100644 index 0000000000000..ee4b5991e97af --- /dev/null +++ b/plugins/engine-datafusion/jni/src/metadata_cache.rs @@ -0,0 +1,251 @@ + +use std::sync::{Arc, Mutex}; +use jni::JNIEnv; + +use datafusion::execution::cache::cache_manager::{FileMetadataCache}; +use datafusion::execution::cache::cache_unit::{DefaultFilesMetadataCache}; +use datafusion::execution::cache::CacheAccessor; +use object_store::ObjectMeta; + +// Helper function to handle cache errors +fn handle_cache_error(env: &mut JNIEnv, operation: &str, error: &str) { + let msg = format!("Cache {} failed: {}", operation, error); + eprintln!("[CACHE ERROR] {}", msg); + let _ = env.throw_new("java/lang/RuntimeException", &msg); +} + +// Helper function to log cache operations +fn log_cache_error(operation: &str, error: &str) { + eprintln!("[CACHE ERROR] {} operation failed: {}", operation, error); +} + +/* +DefaultFilesMetadataCache of datafusion is internally wrapped in a Mutex and requires mut access for operations like remove. +Refer: https://github.com/apache/datafusion/blob/main/datafusion/execution/src/cache/cache_unit.rs#L312-L315 +https://github.com/apache/datafusion/blob/main/datafusion/execution/src/cache/cache_unit.rs#L402 + +Having multiple references to Cache, trying to acquire a mutable reference for methods like remove +would lead to failures. +Hence explicit handling of mutable references is required for which MutexFileMetadataCache is introduced +*/ + +// Wrapper to make Mutex implement FileMetadataCache +pub struct MutexFileMetadataCache { + pub inner: Mutex, +} + +impl MutexFileMetadataCache { + pub fn new(cache: DefaultFilesMetadataCache) -> Self { + Self { + inner: Mutex::new(cache), + } + } +} + +// Implement CacheAccessor which is required by FileMetadataCache +impl CacheAccessor> for MutexFileMetadataCache { + type Extra = ObjectMeta; + + fn get(&self, k: &ObjectMeta) -> Option> { + match self.inner.lock() { + Ok(cache) => cache.get(k), + Err(e) => { + log_cache_error("get", &e.to_string()); + None + } + } + } + + fn get_with_extra(&self, k: &ObjectMeta, extra: &Self::Extra) -> Option> { + match self.inner.lock() { + Ok(cache) => cache.get_with_extra(k, extra), + Err(e) => { + log_cache_error("get_with_extra", &e.to_string()); + None + } + } + } + + fn put(&self, k: &ObjectMeta, v: Arc) -> Option> { + match self.inner.lock() { + Ok(mut cache) => cache.put(k, v), + Err(e) => { + log_cache_error("put", &e.to_string()); + None + } + } + } + + fn put_with_extra(&self, k: &ObjectMeta, v: Arc, e: &Self::Extra) -> Option> { + match self.inner.lock() { + Ok(mut cache) => cache.put_with_extra(k, v, e), + Err(err) => { + log_cache_error("put_with_extra", &err.to_string()); + None + } + } + } + + fn remove(&mut self, k: &ObjectMeta) -> Option> { + match self.inner.lock() { + Ok(mut cache) => cache.remove(k), + Err(e) => { + log_cache_error("remove", &e.to_string()); + None + } + } + } + + fn contains_key(&self, k: &ObjectMeta) -> bool { + match self.inner.lock() { + Ok(cache) => cache.contains_key(k), + Err(e) => { + log_cache_error("contains_key", &e.to_string()); + false + } + } + } + + fn len(&self) -> usize { + match self.inner.lock() { + Ok(cache) => cache.len(), + Err(e) => { + log_cache_error("len", &e.to_string()); + 0 + } + } + } + + fn clear(&self) { + match self.inner.lock() { + Ok(mut cache) => cache.clear(), + Err(e) => log_cache_error("clear", &e.to_string()), + } + } + + fn name(&self) -> String { + match self.inner.lock() { + Ok(cache) => cache.name(), + Err(e) => { + log_cache_error("name", &e.to_string()); + "cache_error".to_string() + } + } + } +} + +impl FileMetadataCache for MutexFileMetadataCache { + fn cache_limit(&self) -> usize { + match self.inner.lock() { + Ok(cache) => cache.cache_limit(), + Err(e) => { + log_cache_error("cache_limit", &e.to_string()); + 0 + } + } + } + + fn update_cache_limit(&self, limit: usize) { + match self.inner.lock() { + Ok(mut cache) => cache.update_cache_limit(limit), + Err(e) => log_cache_error("update_cache_limit", &e.to_string()), + } + } + + fn list_entries(&self) -> std::collections::HashMap { + match self.inner.lock() { + Ok(cache) => cache.list_entries(), + Err(e) => { + log_cache_error("list_entries", &e.to_string()); + std::collections::HashMap::new() + } + } + } +} + +// JNI wrapper functions for cache operations +use jni::objects::JClass; +use jni::sys::jlong; + +#[no_mangle] +pub extern "system" fn Java_org_opensearch_datafusion_DataFusionQueryJNI_cacheGet( + mut env: JNIEnv, + _class: JClass, + cache_ptr: jlong, + key_ptr: jlong, +) -> jlong { + let cache = unsafe { &*(cache_ptr as *const MutexFileMetadataCache) }; + let key = unsafe { &*(key_ptr as *const ObjectMeta) }; + + match cache.inner.lock() { + Ok(cache_guard) => { + match cache_guard.get(key) { + Some(metadata) => Box::into_raw(Box::new(metadata)) as jlong, + None => 0, + } + } + Err(e) => { + handle_cache_error(&mut env, "get", &e.to_string()); + -1 + } + } +} + +#[no_mangle] +pub extern "system" fn Java_org_opensearch_datafusion_DataFusionQueryJNI_cachePut( + mut env: JNIEnv, + _class: JClass, + cache_ptr: jlong, + key_ptr: jlong, + value_ptr: jlong, +) -> jlong { + let cache = unsafe { &*(cache_ptr as *const MutexFileMetadataCache) }; + let key = unsafe { &*(key_ptr as *const ObjectMeta) }; + let value = unsafe { Box::from_raw(value_ptr as *mut Arc) }; + + match cache.inner.lock() { + Ok(mut cache_guard) => { + match cache_guard.put(key, *value) { + Some(old_value) => Box::into_raw(Box::new(old_value)) as jlong, + None => 0, + } + } + Err(e) => { + handle_cache_error(&mut env, "put", &e.to_string()); + -1 + } + } +} + +#[no_mangle] +pub extern "system" fn Java_org_opensearch_datafusion_DataFusionQueryJNI_cacheRemove( + mut env: JNIEnv, + _class: JClass, + cache_ptr: jlong, + key_ptr: jlong, +) -> jlong { + let cache = unsafe { &mut *(cache_ptr as *mut MutexFileMetadataCache) }; + let key = unsafe { &*(key_ptr as *const ObjectMeta) }; + + match cache.remove(key) { + Some(metadata) => Box::into_raw(Box::new(metadata)) as jlong, + None => { + let _ = env.throw_new("java/util/NoSuchElementException", "Key not found in cache"); + -1 + } + } +} + +#[no_mangle] +pub extern "system" fn Java_org_opensearch_datafusion_DataFusionQueryJNI_cacheClear( + mut env: JNIEnv, + _class: JClass, + cache_ptr: jlong, +) { + let cache = unsafe { &*(cache_ptr as *const MutexFileMetadataCache) }; + + match cache.inner.lock() { + Ok(mut cache_guard) => cache_guard.clear(), + Err(e) => handle_cache_error(&mut env, "clear", &e.to_string()), + } +} diff --git a/plugins/engine-datafusion/jni/src/statistics_cache.rs b/plugins/engine-datafusion/jni/src/statistics_cache.rs new file mode 100644 index 0000000000000..4b3d68f717cce --- /dev/null +++ b/plugins/engine-datafusion/jni/src/statistics_cache.rs @@ -0,0 +1,1049 @@ + +use crate::cache_policy::{ + create_policy, CacheConfig, CacheError, CachePolicy, CacheResult, PolicyType, +}; +use arrow_array::Array; +use datafusion::common::stats::{ColumnStatistics, Precision}; +use datafusion::common::ScalarValue; +use dashmap::DashMap; +use datafusion::execution::cache::CacheAccessor; +use datafusion::physical_plan::Statistics; +use object_store::{path::Path, ObjectMeta}; +use std::collections::HashMap; +use std::sync::{Arc, Mutex}; + +/// Trait to calculate heap memory size for statistics objects +trait HeapSize { + fn heap_size(&self) -> usize; +} + +impl HeapSize for Statistics { + fn heap_size(&self) -> usize { + std::mem::size_of::() + + self.num_rows.heap_size() + + self.total_byte_size.heap_size() + + self.column_statistics.heap_size() + } +} + +impl HeapSize +for Precision +{ + fn heap_size(&self) -> usize { + match self { + Precision::Exact(val) => std::mem::size_of::() + val.heap_size(), + Precision::Inexact(val) => std::mem::size_of::() + val.heap_size(), + Precision::Absent => std::mem::size_of::(), + } + } +} + +impl HeapSize for usize { + fn heap_size(&self) -> usize { + 0 // Primitive types don't have heap allocation + } +} + +impl HeapSize for Vec { + fn heap_size(&self) -> usize { + std::mem::size_of::() + + (self.capacity() * std::mem::size_of::()) + + self.iter().map(|item| item.heap_size()).sum::() + } +} + +impl HeapSize for ColumnStatistics { + fn heap_size(&self) -> usize { + std::mem::size_of::() + + self.null_count.heap_size() + + self.max_value.heap_size() + + self.min_value.heap_size() + + self.distinct_count.heap_size() + } +} + +impl HeapSize for ScalarValue { + fn heap_size(&self) -> usize { + match self { + ScalarValue::Utf8(Some(s)) | ScalarValue::LargeUtf8(Some(s)) => { + std::mem::size_of::() + s.capacity() + } + ScalarValue::Binary(Some(b)) | ScalarValue::LargeBinary(Some(b)) => { + std::mem::size_of::() + b.capacity() + } + ScalarValue::List(arr) => { + // Estimate list array memory size + std::mem::size_of::() + std::mem::size_of_val(arr.as_ref()) + (arr.len() * 8) + } + ScalarValue::Struct(arr) => { + // Estimate struct array memory size + std::mem::size_of::() + std::mem::size_of_val(arr.as_ref()) + (arr.len() * 16) + } + _ => std::mem::size_of::(), // Primitive types and nulls + } + } +} + +/// Extension trait to add memory_size method to Statistics +trait StatisticsMemorySize { + fn memory_size(&self) -> usize; +} + +impl StatisticsMemorySize for Statistics { + fn memory_size(&self) -> usize { + std::mem::size_of::() + + self.num_rows.heap_size() + + self.total_byte_size.heap_size() + + self.column_statistics.heap_size() + } +} + +/// Combined memory tracking and policy-based eviction cache +/// +/// This cache leverages DashMap's built-in concurrency from DefaultFileStatisticsCache +/// and adds memory tracking + policy-based eviction on top. +pub struct CustomStatisticsCache { + /// The underlying DataFusion statistics cache (DashMap-based, already thread-safe) + inner_cache: DashMap)>, + /// The eviction policy (thread-safe) + policy: Arc>>, + /// Cache configuration + config: CacheConfig, + /// Memory usage tracker - maps cache keys to their memory consumption (thread-safe) + memory_tracker: Arc>>, + /// Total memory consumed by all entries (thread-safe) + total_memory: Arc>, + /// Cache hit count (thread-safe) + hit_count: Arc>, + /// Cache miss count (thread-safe) + miss_count: Arc>, +} + +impl CustomStatisticsCache { + /// Create a new custom statistics cache + pub fn new(config: CacheConfig) -> Self { + let inner_cache = DashMap::new(); + let policy = Arc::new(Mutex::new(create_policy(config.policy_type.clone()))); + + Self { + inner_cache, + policy, + config, + memory_tracker: Arc::new(Mutex::new(HashMap::new())), + total_memory: Arc::new(Mutex::new(0)), + hit_count: Arc::new(Mutex::new(0)), + miss_count: Arc::new(Mutex::new(0)), + } + } + + /// Create with default configuration + pub fn with_default_config() -> Self { + Self::new(CacheConfig::default()) + } + + /// Get the underlying cache for compatibility + pub fn inner(&self) -> &DashMap)> { + &self.inner_cache + } + + /// Get total memory consumed by all cached statistics + pub fn memory_consumed(&self) -> usize { + self.total_memory.lock().map(|guard| *guard).unwrap_or(0) + } + + /// Get cache hit count + pub fn hit_count(&self) -> usize { + self.hit_count.lock().map(|guard| *guard).unwrap_or(0) + } + + /// Get cache miss count + pub fn miss_count(&self) -> usize { + self.miss_count.lock().map(|guard| *guard).unwrap_or(0) + } + + /// Get cache hit rate (returns value between 0.0 and 1.0) + pub fn hit_rate(&self) -> f64 { + let hits = self.hit_count(); + let misses = self.miss_count(); + let total = hits + misses; + if total == 0 { + 0.0 + } else { + hits as f64 / total as f64 + } + } + + /// Reset hit and miss counters + pub fn reset_stats(&self) { + if let Ok(mut hits) = self.hit_count.lock() { + *hits = 0; + } + if let Ok(mut misses) = self.miss_count.lock() { + *misses = 0; + } + } + + /// Update the cache size limit + pub fn update_size_limit(&mut self, new_limit: usize) -> CacheResult<()> { + // Update the config size limit + self.config.size_limit = new_limit; + + let current_size = self.current_size()?; + if current_size > new_limit { + let target_eviction = current_size - (new_limit as f64 * 0.8) as usize; + self.evict(target_eviction)?; + } + Ok(()) + } + + /// Switch to a different eviction policy + pub fn set_policy(&self, policy_type: PolicyType) -> CacheResult<()> { + let mut policy_guard = self + .policy + .lock() + .map_err(|e| CacheError::PolicyLockError { + reason: format!("Failed to acquire policy lock: {}", e), + })?; + + // Create new policy and transfer existing entries + let mut new_policy = create_policy(policy_type.clone()); + + // Get all current entries and notify new policy + if let Ok(tracker) = self.memory_tracker.lock() { + for (key, size) in tracker.iter() { + new_policy.on_insert(key, *size); + } + } + + *policy_guard = new_policy; + // Note: We can't update self.config.policy_type here since we don't have &mut self + // The policy change is effective immediately through the policy_guard update + + Ok(()) + } + + /// Get current policy name + pub fn policy_name(&self) -> CacheResult { + let policy_guard = self + .policy + .lock() + .map_err(|e| CacheError::PolicyLockError { + reason: format!("Failed to acquire policy lock: {}", e), + })?; + Ok(policy_guard.policy_name().to_string()) + } + + /// Get current cache size according to policy (uses actual memory consumption) + pub fn current_size(&self) -> CacheResult { + Ok(self.memory_consumed()) + } + + /// Manually trigger eviction (requires &mut self) + pub fn evict(&mut self, target_size: usize) -> CacheResult { + if target_size == 0 { + return Ok(0); + } + + let candidates = { + let policy_guard = self.policy.lock().map_err(|e| CacheError::PolicyLockError { + reason: format!("Failed to acquire policy lock: {}", e), + })?; + policy_guard.select_for_eviction(target_size) + }; + + let mut freed_size = 0; + for key in candidates { + let entry_size = if let Ok(tracker) = self.memory_tracker.lock() { + tracker.get(&key).copied().unwrap_or(0) + } else { + 0 + }; + + if entry_size > 0 { + if let Ok(path) = self.parse_key_to_path(&key) { + if self.inner_cache.remove(&path).is_some() { + // Update memory tracking + if let Ok(mut tracker) = self.memory_tracker.lock() { + if let Ok(mut total) = self.total_memory.lock() { + tracker.remove(&key); + *total = total.saturating_sub(entry_size); + } + } + + // Notify policy + if let Ok(mut policy_guard) = self.policy.lock() { + policy_guard.on_remove(&key); + } + + freed_size += entry_size; + } + + if freed_size >= target_size { + break; + } + } + } + } + + Ok(freed_size) + } + + /// Parse cache key back to Path + fn parse_key_to_path(&self, key: &str) -> CacheResult { + Ok(Path::from(key)) + } + + /// Remove entry internally (works with &self since inner_cache is thread-safe) + fn remove_internal(&self, k: &Path) -> Option> { + let key = k.to_string(); + + // Actually remove from the underlying cache (DashMap allows this with &self) + let result = self.inner_cache.remove(k); + + // Only proceed with tracking updates if the entry existed + if result.is_some() { + // Update memory tracking + if let Ok(mut tracker) = self.memory_tracker.lock() { + if let Ok(mut total) = self.total_memory.lock() { + if let Some(old_size) = tracker.remove(&key) { + *total = total.saturating_sub(old_size); + } + } + } + + // Notify policy of removal + if let Ok(mut policy_guard) = self.policy.lock() { + policy_guard.on_remove(&key); + } + } + + result.map(|x| x.1 .1) + } + + +} + +// Implement CacheAccessor - DashMap handles concurrency, we just need to handle the &mut self requirement +impl CacheAccessor> for CustomStatisticsCache { + type Extra = ObjectMeta; + + fn get(&self, k: &Path) -> Option> { + let result = self.inner_cache.get(k); + + if result.is_some() { + // Increment hit count + if let Ok(mut hits) = self.hit_count.lock() { + *hits += 1; + } + + // Notify policy of access + let key = k.to_string(); + let memory_size = if let Ok(tracker) = self.memory_tracker.lock() { + tracker.get(&key).copied().unwrap_or(0) + } else { + 0 + }; + + if let Ok(mut policy_guard) = self.policy.lock() { + policy_guard.on_access(&key, memory_size); + } + } else { + // Increment miss count + if let Ok(mut misses) = self.miss_count.lock() { + *misses += 1; + } + } + + result.map(|s| Some(Arc::clone(&s.value().1))) + .unwrap_or(None) + } + + fn get_with_extra(&self, k: &Path, _extra: &Self::Extra) -> Option> { + self.get(k) + } + + fn put(&self, k: &Path, v: Arc) -> Option> { + let meta = ObjectMeta { + location: k.clone(), + last_modified: chrono::Utc::now(), + size: 0, + e_tag: None, + version: None, + }; + self.put_with_extra(k, v, &meta) + } + + fn put_with_extra( + &self, + k: &Path, + v: Arc, + e: &Self::Extra, + ) -> Option> { + let key = k.to_string(); + let memory_size = v.memory_size(); + + // Check if eviction is needed BEFORE inserting + let current_size = self.memory_consumed(); + if current_size + memory_size > (self.config.size_limit as f64 * 0.8) as usize { + let target_eviction = (current_size + memory_size) - (self.config.size_limit as f64 * 0.6) as usize; + + // Perform actual eviction using remove_internal + let candidates = { + if let Ok(policy_guard) = self.policy.lock() { + policy_guard.select_for_eviction(target_eviction) + } else { + vec![] + } + }; + + for candidate_key in candidates { + if let Ok(path) = self.parse_key_to_path(&candidate_key) { + self.remove_internal(&path); + } + } + } + + // Put in the underlying cache (DashMap handles concurrency) + let result = self.inner_cache.insert(k.clone(), (e.clone(), v)).map(|x| x.1); + + // Track memory usage + if let Ok(mut tracker) = self.memory_tracker.lock() { + if let Ok(mut total) = self.total_memory.lock() { + // If there was a previous entry, subtract its memory + if let Some(old_size) = tracker.get(&key) { + *total = total.saturating_sub(*old_size); + } + + // Add new entry memory + tracker.insert(key.clone(), memory_size); + *total += memory_size; + } + } + + // Notify policy of insertion + if let Ok(mut policy_guard) = self.policy.lock() { + policy_guard.on_insert(&key, memory_size); + } + + result + } + + fn remove(&mut self, k: &Path) -> Option> { + let key = k.to_string(); + + // Actually remove from the underlying cache + let result = self.inner_cache.remove(k); + + // Only proceed with tracking updates if the entry existed + if result.is_some() { + // Update memory tracking + if let Ok(mut tracker) = self.memory_tracker.lock() { + if let Ok(mut total) = self.total_memory.lock() { + if let Some(old_size) = tracker.remove(&key) { + *total = total.saturating_sub(old_size); + } + } + } + + // Notify policy of removal + if let Ok(mut policy_guard) = self.policy.lock() { + policy_guard.on_remove(&key); + } + } + + result.map(|x| x.1 .1) + } + + fn contains_key(&self, k: &Path) -> bool { + self.inner_cache.get(k).is_some() + } + + fn len(&self) -> usize { + self.memory_tracker.lock().map(|t| t.len()).unwrap_or(0) + } + + fn clear(&self) { + // Clear the DashMap cache + self.inner_cache.clear(); + + // Clear memory tracking + if let Ok(mut tracker) = self.memory_tracker.lock() { + tracker.clear(); + } + + if let Ok(mut total) = self.total_memory.lock() { + *total = 0; + } + + // Clear policy + if let Ok(mut policy_guard) = self.policy.lock() { + policy_guard.clear(); + } + + // Reset hit/miss counters + self.reset_stats(); + } + + fn name(&self) -> String { + format!( + "CustomStatisticsCache({})", + self.policy_name().unwrap_or_else(|_| "unknown".to_string()) + ) + } +} + +impl Default for CustomStatisticsCache { + fn default() -> Self { + Self::with_default_config() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use chrono::Utc; + use datafusion::common::stats::Precision; + + fn create_test_statistics() -> Statistics { + Statistics { + num_rows: Precision::Exact(1000), + total_byte_size: Precision::Exact(50000), + column_statistics: vec![], + } + } + + fn create_test_path(name: &str) -> Path { + Path::from(format!("/test/{}.parquet", name)) + } + + fn create_test_meta(path: &Path) -> ObjectMeta { + ObjectMeta { + location: path.clone(), + last_modified: Utc::now(), + size: 1000, + e_tag: None, + version: None, + } + } + + #[test] + fn test_custom_stats_cache_creation() { + let config = CacheConfig { + policy_type: PolicyType::Lru, + size_limit: 1024 * 1024, + eviction_threshold: 0.8, + }; + + let cache = CustomStatisticsCache::new(config); + assert_eq!(cache.policy_name().unwrap(), "lru"); + assert_eq!(cache.memory_consumed(), 0); + assert_eq!(cache.len(), 0); + } + + #[test] + fn test_memory_tracking_with_policy() { + let cache = CustomStatisticsCache::with_default_config(); + + // Initially empty + assert_eq!(cache.memory_consumed(), 0); + assert_eq!(cache.len(), 0); + + // Add an entry + let path = create_test_path("file1"); + let meta = create_test_meta(&path); + let stats = Arc::new(create_test_statistics()); + + cache.put_with_extra(&path, stats, &meta); + + // Should have memory consumption and policy tracking + assert!(cache.memory_consumed() > 0); + assert_eq!(cache.len(), 1); + assert_eq!(cache.current_size().unwrap(), cache.memory_consumed()); + + // Verify we can retrieve it + assert!(cache.get(&path).is_some()); + } + + #[test] + fn test_policy_based_eviction_with_memory() { + let config = CacheConfig { + policy_type: PolicyType::Lru, + size_limit: 1000, // Small limit to trigger eviction + eviction_threshold: 0.8, + }; + let cache = CustomStatisticsCache::new(config); + + // Add multiple entries to trigger eviction + for i in 0..10 { + let path = create_test_path(&format!("file{}", i)); + let meta = create_test_meta(&path); + let stats = Arc::new(create_test_statistics()); + + cache.put_with_extra(&path, stats, &meta); + } + + // Memory should be managed by eviction + let final_memory = cache.memory_consumed(); + assert!( + final_memory <= 1000, + "Memory should be within limit due to eviction" + ); + assert!(cache.len() > 0, "Should still have some entries"); + } + + #[test] + fn test_manual_eviction_with_memory_tracking() { + let mut cache = CustomStatisticsCache::with_default_config(); + + // Add entries + for i in 0..5 { + let path = create_test_path(&format!("file{}", i)); + let meta = create_test_meta(&path); + let stats = Arc::new(create_test_statistics()); + cache.put_with_extra(&path, stats, &meta); + } + + let memory_before = cache.memory_consumed(); + assert!(memory_before > 0); + + // Manually evict some memory + let freed = cache.evict(memory_before / 2).unwrap(); + let memory_after = cache.memory_consumed(); + + assert!(freed > 0, "Should have freed some memory"); + assert!(memory_after < memory_before, "Memory should be reduced"); + } + + #[test] + fn test_policy_switching_with_memory() { + let mut cache = CustomStatisticsCache::with_default_config(); + + // Add some entries + for i in 0..3 { + let path = create_test_path(&format!("file{}", i)); + let meta = create_test_meta(&path); + let stats = Arc::new(create_test_statistics()); + cache.put_with_extra(&path, stats, &meta); + } + + let memory_before = cache.memory_consumed(); + + // Switch policy + assert_eq!(cache.policy_name().unwrap(), "lru"); + cache.set_policy(PolicyType::Lfu).unwrap(); + assert_eq!(cache.policy_name().unwrap(), "lfu"); + + // Memory tracking should be preserved + assert_eq!(cache.memory_consumed(), memory_before); + assert_eq!(cache.len(), 3); + } + + #[test] + fn test_remove_with_memory_tracking() { + let mut cache = CustomStatisticsCache::with_default_config(); + + // Add entries + let path1 = create_test_path("file1"); + let path2 = create_test_path("file2"); + let meta1 = create_test_meta(&path1); + let meta2 = create_test_meta(&path2); + let stats = Arc::new(create_test_statistics()); + + cache.put_with_extra(&path1, stats.clone(), &meta1); + cache.put_with_extra(&path2, stats, &meta2); + + let memory_with_two = cache.memory_consumed(); + assert_eq!(cache.len(), 2); + + // Remove one entry + let removed = cache.remove(&path1); + assert!(removed.is_some()); + + let memory_with_one = cache.memory_consumed(); + assert_eq!(cache.len(), 1); + assert!(memory_with_one < memory_with_two); + + // Remove second entry + cache.remove(&path2); + assert_eq!(cache.memory_consumed(), 0); + assert_eq!(cache.len(), 0); + } + + #[test] + fn test_clear_with_memory_tracking() { + let cache = CustomStatisticsCache::with_default_config(); + + // Add multiple entries + for i in 0..3 { + let path = create_test_path(&format!("file{}", i)); + let meta = create_test_meta(&path); + let stats = Arc::new(create_test_statistics()); + cache.put_with_extra(&path, stats, &meta); + } + + assert!(cache.memory_consumed() > 0); + assert_eq!(cache.len(), 3); + + // Clear all + cache.clear(); + + assert_eq!(cache.memory_consumed(), 0); + assert_eq!(cache.len(), 0); + } + + #[test] + fn test_lru_eviction_with_memory() { + let config = CacheConfig { + policy_type: PolicyType::Lru, + size_limit: 2000, // Limit to ~2 entries + eviction_threshold: 0.8, + }; + let cache = CustomStatisticsCache::new(config); + + // Add entries + let mut paths = vec![]; + for i in 0..5 { + let path = create_test_path(&format!("file{}", i)); + let meta = create_test_meta(&path); + let stats = Arc::new(create_test_statistics()); + paths.push(path.clone()); + + cache.put_with_extra(&path, stats, &meta); + } + + // Access some entries to update LRU order + cache.get(&paths[2]); + cache.get(&paths[4]); + + // Memory should be within limits due to LRU eviction + assert!(cache.memory_consumed() <= 2000); + + // Recently accessed entries should still be available + assert!(cache.get(&paths[2]).is_some()); + assert!(cache.get(&paths[4]).is_some()); + } + + #[test] + fn test_lfu_eviction_with_memory() { + let config = CacheConfig { + policy_type: PolicyType::Lfu, + size_limit: 2000, // Limit to ~2 entries + eviction_threshold: 0.8, + }; + let cache = CustomStatisticsCache::new(config); + + // Add entries + let mut paths = vec![]; + for i in 0..5 { + let path = create_test_path(&format!("file{}", i)); + let meta = create_test_meta(&path); + let stats = Arc::new(create_test_statistics()); + paths.push(path.clone()); + + cache.put_with_extra(&path, stats, &meta); + } + + // Create frequency patterns + for _ in 0..5 { + cache.get(&paths[1]); + cache.get(&paths[3]); + } + + // Memory should be within limits due to LFU eviction + assert!(cache.memory_consumed() <= 2000); + + // Frequently accessed entries should still be available + assert!(cache.get(&paths[1]).is_some()); + assert!(cache.get(&paths[3]).is_some()); + } + + #[test] + fn test_concurrent_operations() { + use std::sync::Arc; + use std::thread; + + let cache = Arc::new(CustomStatisticsCache::with_default_config()); + let mut handles = vec![]; + + // Spawn multiple threads doing concurrent operations + for i in 0..10 { + let cache_clone = Arc::clone(&cache); + let handle = thread::spawn(move || { + let path = create_test_path(&format!("concurrent{}", i)); + let meta = create_test_meta(&path); + let stats = Arc::new(create_test_statistics()); + + // Put operation + cache_clone.put_with_extra(&path, stats, &meta); + + // Get operation + let result = cache_clone.get(&path); + assert!(result.is_some()); + }); + handles.push(handle); + } + + // Wait for all threads to complete + for handle in handles { + handle.join().unwrap(); + } + + // Cache should have entries + assert!(cache.len() > 0); + } + + #[test] + fn test_hit_count_tracking() { + let cache = CustomStatisticsCache::with_default_config(); + + // Initially zero + assert_eq!(cache.hit_count(), 0); + assert_eq!(cache.miss_count(), 0); + + // Add an entry + let path = create_test_path("file1"); + let meta = create_test_meta(&path); + let stats = Arc::new(create_test_statistics()); + cache.put_with_extra(&path, stats, &meta); + + // Get the entry - should increment hit count + assert!(cache.get(&path).is_some()); + assert_eq!(cache.hit_count(), 1); + assert_eq!(cache.miss_count(), 0); + + // Get it again - should increment hit count again + assert!(cache.get(&path).is_some()); + assert_eq!(cache.hit_count(), 2); + assert_eq!(cache.miss_count(), 0); + } + + #[test] + fn test_miss_count_tracking() { + let cache = CustomStatisticsCache::with_default_config(); + + // Initially zero + assert_eq!(cache.hit_count(), 0); + assert_eq!(cache.miss_count(), 0); + + // Try to get non-existent entry - should increment miss count + let path = create_test_path("nonexistent"); + assert!(cache.get(&path).is_none()); + assert_eq!(cache.hit_count(), 0); + assert_eq!(cache.miss_count(), 1); + + // Try again - should increment miss count again + assert!(cache.get(&path).is_none()); + assert_eq!(cache.hit_count(), 0); + assert_eq!(cache.miss_count(), 2); + } + + #[test] + fn test_hit_miss_mixed_operations() { + let cache = CustomStatisticsCache::with_default_config(); + + // Add some entries + let path1 = create_test_path("file1"); + let path2 = create_test_path("file2"); + let meta1 = create_test_meta(&path1); + let meta2 = create_test_meta(&path2); + let stats = Arc::new(create_test_statistics()); + + cache.put_with_extra(&path1, stats.clone(), &meta1); + cache.put_with_extra(&path2, stats, &meta2); + + // Mix of hits and misses + assert!(cache.get(&path1).is_some()); // hit + assert_eq!(cache.hit_count(), 1); + assert_eq!(cache.miss_count(), 0); + + let path3 = create_test_path("file3"); + assert!(cache.get(&path3).is_none()); // miss + assert_eq!(cache.hit_count(), 1); + assert_eq!(cache.miss_count(), 1); + + assert!(cache.get(&path2).is_some()); // hit + assert_eq!(cache.hit_count(), 2); + assert_eq!(cache.miss_count(), 1); + + assert!(cache.get(&path3).is_none()); // miss again + assert_eq!(cache.hit_count(), 2); + assert_eq!(cache.miss_count(), 2); + + assert!(cache.get(&path1).is_some()); // hit again + assert_eq!(cache.hit_count(), 3); + assert_eq!(cache.miss_count(), 2); + } + + #[test] + fn test_hit_rate_calculation() { + let cache = CustomStatisticsCache::with_default_config(); + + // Initially 0.0 (no operations) + assert_eq!(cache.hit_rate(), 0.0); + + // Add entries + let path1 = create_test_path("file1"); + let path2 = create_test_path("file2"); + let meta1 = create_test_meta(&path1); + let meta2 = create_test_meta(&path2); + let stats = Arc::new(create_test_statistics()); + + cache.put_with_extra(&path1, stats.clone(), &meta1); + cache.put_with_extra(&path2, stats, &meta2); + + // 2 hits, 0 misses = 100% hit rate + cache.get(&path1); + cache.get(&path2); + assert_eq!(cache.hit_rate(), 1.0); + + // 2 hits, 1 miss = 66.67% hit rate + let path3 = create_test_path("file3"); + cache.get(&path3); + assert!((cache.hit_rate() - 0.6666666666666666).abs() < 0.0001); + + // 3 hits, 1 miss = 75% hit rate + cache.get(&path1); + assert_eq!(cache.hit_rate(), 0.75); + + // 3 hits, 2 misses = 60% hit rate + cache.get(&path3); + assert_eq!(cache.hit_rate(), 0.6); + } + + #[test] + fn test_reset_stats() { + let cache = CustomStatisticsCache::with_default_config(); + + // Add entry and generate some hits/misses + let path = create_test_path("file1"); + let meta = create_test_meta(&path); + let stats = Arc::new(create_test_statistics()); + cache.put_with_extra(&path, stats, &meta); + + cache.get(&path); // hit + let path2 = create_test_path("file2"); + cache.get(&path2); // miss + + assert_eq!(cache.hit_count(), 1); + assert_eq!(cache.miss_count(), 1); + assert_eq!(cache.hit_rate(), 0.5); + + // Reset stats + cache.reset_stats(); + + assert_eq!(cache.hit_count(), 0); + assert_eq!(cache.miss_count(), 0); + assert_eq!(cache.hit_rate(), 0.0); + + // Cache entries should still exist + assert_eq!(cache.len(), 1); + assert!(cache.get(&path).is_some()); + + // After reset, new operations should start counting from zero + assert_eq!(cache.hit_count(), 1); + assert_eq!(cache.miss_count(), 0); + } + + #[test] + fn test_clear_resets_stats() { + let cache = CustomStatisticsCache::with_default_config(); + + // Add entries and generate hits/misses + let path = create_test_path("file1"); + let meta = create_test_meta(&path); + let stats = Arc::new(create_test_statistics()); + cache.put_with_extra(&path, stats, &meta); + + cache.get(&path); // hit + let path2 = create_test_path("file2"); + cache.get(&path2); // miss + + assert_eq!(cache.hit_count(), 1); + assert_eq!(cache.miss_count(), 1); + + // Clear should reset stats + cache.clear(); + + assert_eq!(cache.hit_count(), 0); + assert_eq!(cache.miss_count(), 0); + assert_eq!(cache.len(), 0); + } + + #[test] + fn test_hit_miss_with_eviction() { + let config = CacheConfig { + policy_type: PolicyType::Lru, + size_limit: 1500, // Small limit to trigger eviction + eviction_threshold: 0.8, + }; + let cache = CustomStatisticsCache::new(config); + + // Add multiple entries to trigger eviction + let mut paths = vec![]; + for i in 0..5 { + let path = create_test_path(&format!("file{}", i)); + let meta = create_test_meta(&path); + let stats = Arc::new(create_test_statistics()); + paths.push(path.clone()); + cache.put_with_extra(&path, stats, &meta); + } + + // Access some entries + cache.get(&paths[0]); // May or may not be evicted + cache.get(&paths[4]); // Should still be there + + let hits_before = cache.hit_count(); + let misses_before = cache.miss_count(); + + // Try to access potentially evicted entries + for path in &paths { + cache.get(path); + } + + // Should have more hits and/or misses + assert!(cache.hit_count() >= hits_before); + assert!(cache.miss_count() >= misses_before); + assert!(cache.hit_count() + cache.miss_count() > hits_before + misses_before); + } + + #[test] + fn test_concurrent_hit_miss_tracking() { + use std::sync::Arc; + use std::thread; + + let cache = Arc::new(CustomStatisticsCache::with_default_config()); + + // Add some entries + for i in 0..5 { + let path = create_test_path(&format!("file{}", i)); + let meta = create_test_meta(&path); + let stats = Arc::new(create_test_statistics()); + cache.put_with_extra(&path, stats, &meta); + } + + let mut handles = vec![]; + + // Spawn threads that will generate hits and misses + for i in 0..10 { + let cache_clone = Arc::clone(&cache); + let handle = thread::spawn(move || { + // Mix of hits and misses + let existing_path = create_test_path(&format!("file{}", i % 5)); + cache_clone.get(&existing_path); // hit + + let nonexistent_path = create_test_path(&format!("missing{}", i)); + cache_clone.get(&nonexistent_path); // miss + }); + handles.push(handle); + } + + for handle in handles { + handle.join().unwrap(); + } + + // Should have 10 hits and 10 misses + assert_eq!(cache.hit_count(), 10); + assert_eq!(cache.miss_count(), 10); + assert_eq!(cache.hit_rate(), 0.5); + } +} diff --git a/plugins/engine-datafusion/jni/src/util.rs b/plugins/engine-datafusion/jni/src/util.rs index 1b6da12c49ede..8681d4a6acf57 100644 --- a/plugins/engine-datafusion/jni/src/util.rs +++ b/plugins/engine-datafusion/jni/src/util.rs @@ -8,10 +8,14 @@ use datafusion::arrow::array::RecordBatch; use jni::objects::{JObject, JObjectArray, JString}; use jni::sys::jlong; use jni::JNIEnv; -use object_store::{path::Path as ObjectPath, ObjectMeta}; +use object_store::{path::Path as ObjectPath, ObjectMeta, ObjectStore}; use std::collections::HashMap; use std::error::Error; use std::fs; +use std::sync::Arc; +use datafusion::datasource::physical_plan::parquet::CachedParquetMetaData; +use datafusion::datasource::physical_plan::parquet::metadata::DFParquetMetadata; +use datafusion::execution::cache::cache_manager::FileMetadata; /// Set error message from a result using a Consumer Java callback pub fn set_error_message_batch(env: &mut JNIEnv, callback: JObject, result: Result, Err>) { @@ -158,9 +162,8 @@ pub fn throw_exception(env: &mut JNIEnv, message: &str) { } pub fn create_object_meta_from_filenames(base_path: &str, filenames: Vec) -> Vec { - filenames.into_iter().map(|filename| { + filenames.into_iter().filter_map(|filename| { let filename = filename.as_str(); - // Handle both full paths and relative filenames let full_path = if filename.starts_with('/') || filename.contains(base_path) { // Already a full path @@ -169,19 +172,52 @@ pub fn create_object_meta_from_filenames(base_path: &str, filenames: Vec // Just a filename, needs base_path format!("{}/{}", base_path.trim_end_matches('/'), filename) }; + + // Check if file exists and has content + match fs::metadata(&full_path) { + Ok(metadata) if metadata.len() > 0 => { + Some(create_object_meta_from_file(&full_path)) + } + _ => { + eprintln!("Skipping empty or non-existent file: {}", full_path); + None + } + } + }).collect() +} - let file_size = fs::metadata(&full_path).map(|m| m.len()).unwrap_or(0); - let modified = fs::metadata(&full_path) - .and_then(|m| m.modified()) - .map(|t| DateTime::::from(t)) - .unwrap_or_else(|_| Utc::now()); +pub fn create_object_meta_from_file(file_path: &str) -> ObjectMeta { + let file_size = fs::metadata(&file_path).map(|m| m.len()).unwrap_or(0); + let modified = fs::metadata(&file_path) + .and_then(|m| m.modified()) + .map(|t| DateTime::::from(t)) + .unwrap_or_else(|_| Utc::now()); ObjectMeta { - location: ObjectPath::from(full_path), + location: ObjectPath::from(file_path), last_modified: modified, size: file_size, e_tag: None, version: None, } - }).collect() +} + +pub async fn construct_file_metadata( + store: &dyn ObjectStore, + object_meta: &ObjectMeta, + data_format: &str, +) -> Result, Box> { + match data_format.to_lowercase().as_str() { + "parquet" => { + let df_metadata = DFParquetMetadata::new( + store, + object_meta + ); + + let parquet_metadata = df_metadata.fetch_metadata().await?; + let par = CachedParquetMetaData::new(parquet_metadata); + Ok(Arc::new(par)) + }, + _ => Err(format!("Unsupported data format: {}", data_format).into()) + } } diff --git a/plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/DataFusionPlugin.java b/plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/DataFusionPlugin.java index 45a2da3e6afa3..9f570b7489e0b 100644 --- a/plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/DataFusionPlugin.java +++ b/plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/DataFusionPlugin.java @@ -8,11 +8,17 @@ package org.opensearch.datafusion; +import java.util.Arrays; +import java.util.stream.Collectors; +import java.util.stream.Stream; +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; import org.opensearch.cluster.metadata.IndexNameExpressionResolver; import org.opensearch.cluster.node.DiscoveryNodes; import org.opensearch.cluster.service.ClusterService; import org.opensearch.common.settings.ClusterSettings; import org.opensearch.common.settings.IndexScopedSettings; +import org.opensearch.common.settings.Setting; import org.opensearch.common.settings.Settings; import org.opensearch.common.settings.SettingsFilter; import org.opensearch.core.common.io.stream.NamedWriteableRegistry; @@ -24,6 +30,7 @@ import org.opensearch.datafusion.search.DatafusionQuery; import org.opensearch.datafusion.search.DatafusionReaderManager; import org.opensearch.datafusion.search.DatafusionSearcher; +import org.opensearch.datafusion.search.cache.CacheSettings; import org.opensearch.env.Environment; import org.opensearch.env.NodeEnvironment; import org.opensearch.index.shard.ShardPath; @@ -51,6 +58,10 @@ import java.util.Map; import java.util.function.Supplier; +import static org.opensearch.datafusion.search.cache.CacheSettings.METADATA_CACHE_ENABLED; +import static org.opensearch.datafusion.search.cache.CacheSettings.METADATA_CACHE_EVICTION_TYPE; +import static org.opensearch.datafusion.search.cache.CacheSettings.METADATA_CACHE_SIZE_LIMIT; + /** * Main plugin class for OpenSearch DataFusion integration. * @@ -60,6 +71,8 @@ public class DataFusionPlugin extends Plugin implements ActionPlugin, SearchEngi private DataFusionService dataFusionService; private final boolean isDataFusionEnabled; + private static final Logger logger = LogManager.getLogger(DataFusionPlugin.class); + /** * Constructor for DataFusionPlugin. * @param settings The settings for the DataFusionPlugin. @@ -103,7 +116,7 @@ public Collection createComponents( if (!isDataFusionEnabled) { return Collections.emptyList(); } - dataFusionService = new DataFusionService(dataSourceCodecs); + dataFusionService = new DataFusionService(dataSourceCodecs, clusterService.getClusterSettings()); for(DataFormat format : this.getSupportedFormats()) { dataSourceCodecs.get(format); @@ -167,4 +180,14 @@ public List getRestHandlers( } return List.of(new ActionHandler<>(NodesDataFusionInfoAction.INSTANCE, TransportNodesDataFusionInfoAction.class)); } + + @Override + public List> getSettings() { + return Stream.of( + CacheSettings.CACHE_SETTINGS, + CacheSettings.CACHE_ENABLED) + .flatMap(x -> x.stream()) + .collect(Collectors.toList()); + + } } diff --git a/plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/DataFusionQueryJNI.java b/plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/DataFusionQueryJNI.java index 06a8cc0882132..63bb2d6615546 100644 --- a/plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/DataFusionQueryJNI.java +++ b/plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/DataFusionQueryJNI.java @@ -77,7 +77,7 @@ private static synchronized void loadNativeLibrary() { * Create a new global runtime environment * @return runtime env pointer for subsequent operations */ - public static native long createGlobalRuntime(); + public static native long createGlobalRuntime(long cacheManagerPtr); public static native long createTokioRuntime(); @@ -113,7 +113,7 @@ private static synchronized void loadNativeLibrary() { * @param substraitPlan the serialized Substrait query plan * @return stream pointer for result iteration */ - public static native long executeSubstraitQuery(long cachePtr, String tableName, byte[] substraitPlan, long runtimePtr); + public static native long executeSubstraitQuery(long cachePtr, String tableName, byte[] substraitPlan, long runtimeEnvPtr ,long tokioRuntimePtr); public static native long createDatafusionReader(String path, String[] files); @@ -148,4 +148,169 @@ private static synchronized void loadNativeLibrary() { * @param streamPtr the stream pointer to close */ public static native void closeStream(long streamPtr); + + public static native long initCacheManagerConfig(); + + // METADATA cache specific methods + + /** + * Create metadata cache + * @param cacheConfigPtr cache configuration pointer + * @param sizeLimit sizeLimit for cache + * @return cache pointer + */ + public static native long createMetadataCache(long cacheConfigPtr, long sizeLimit); + + /** + * Put metadata entry into cache + * @param cachePtr the cache pointer + * @param filePath the file path + * @return status code + */ + public static native boolean metadataCachePut(long cachePtr, String filePath); + + /** + * Get metadata from cache + * @param cachePtr the cache pointer + * @param filePath the file path + * @return cached metadata or null if not found + */ + // return type boolean to be changed + public static native boolean metadataCacheGet(long cachePtr, String filePath); + + /** + * Remove metadata from cache + * @param cachePtr the cache pointer + * @param filePath the file path + * @return status code + */ + public static native boolean metadataCacheRemove(long cachePtr, String filePath); + + /** + * Memory consumed by metadataCache + * @param cachePtr the cache pointer + * @return memory used + */ + public static native long metadataCacheGetSize(long cachePtr); + + /** + * Memory consumed by metadataCache + * @param cachePtr the cache pointer + * @param newSizeLimit new size limit + * @return boolean + */ + public static native boolean metadataCacheUpdateSizeLimit(long cachePtr, long newSizeLimit); + + /** + * Check if a file exists in metadataCache + * @param cachePtr the cache pointer + * @param filePath the file path + * @return boolean + */ + public static native boolean metadataCacheContainsFile(long cachePtr, String filePath); + + /** + * Get all entries from the metadata cache + * @param cachePtr the cache pointer + * @return String array containing cache entries in triplets: [path1, size1, hitCount1, path2, size2, hitCount2, ...] + * Each entry consists of 3 consecutive elements: file path, size in bytes, and hit count + */ + public static native String[] metadataCacheGetEntries(long cachePtr); + + /** + * Clears all entries from the metadata cache + * @param cachePtr the cache pointer + */ + public static native void metadataCacheClear(long cachePtr); + + // STATISTICS cache specific methods + + /** + * Create statistics cache + * @param cacheConfigPtr cache configuration pointer + * @param sizeLimit sizeLimit for cache + * @return cache pointer + */ + public static native long createStatisticsCache(long cacheConfigPtr, long sizeLimit); + + /** + * Put statistics entry into cache + * @param cachePtr the cache pointer + * @param filePath the file path + * @return status code + */ + public static native boolean statisticsCachePut(long cachePtr, String filePath); + + /** + * Get statistics from cache + * @param cachePtr the cache pointer + * @param filePath the file path + * @return cached statistics or null if not found + */ + public static native boolean statisticsCacheGet(long cachePtr, String filePath); + + /** + * Remove statistics from cache + * @param cachePtr the cache pointer + * @param filePath the file path + * @return status code + */ + public static native boolean statisticsCacheRemove(long cachePtr, String filePath); + + /** + * Memory consumed by statisticsCache + * @param cachePtr the cache pointer + * @return memory used + */ + public static native long statisticsCacheGetSize(long cachePtr); + + /** + * Update size limit for statisticsCache + * @param cachePtr the cache pointer + * @param newSizeLimit new size limit + * @return boolean + */ + public static native boolean statisticsCacheUpdateSizeLimit(long cachePtr, long newSizeLimit); + + /** + * Check if a file exists in statisticsCache + * @param cachePtr the cache pointer + * @param filePath the file path + * @return boolean + */ + public static native boolean statisticsCacheContainsFile(long cachePtr, String filePath); + + /** + * Get all entries from the statistics cache + * @param cachePtr the cache pointer + * @return String array containing cache entries in triplets: [path1, size1, hitCount1, path2, size2, hitCount2, ...] + * Each entry consists of 3 consecutive elements: file path, size in bytes, and hit count + */ + public static native String[] statisticsCacheGetEntries(long cachePtr); + + /** + * Clear all entries in statisticsCache + * @param cachePtr the cache pointer + */ + public static native void statisticsCacheClear(long cachePtr); + + /** + * Get hit count from statistics cache + */ + public static native long statisticsCacheGetHitCount(long cachePtr); + + /** + * Get miss count from statistics cache + */ + public static native long statisticsCacheGetMissCount(long cachePtr); + + /** + * Get hit rate from statistics cache (returns value between 0.0 and 1.0) + */ + public static native double statisticsCacheGetHitRate(long cachePtr); + + /** + * Reset hit and miss counters in statistics cache + */ + public static native void statisticsCacheResetStats(long cachePtr); } diff --git a/plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/DataFusionService.java b/plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/DataFusionService.java index 9548ced599723..68b9fcd11095b 100644 --- a/plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/DataFusionService.java +++ b/plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/DataFusionService.java @@ -11,9 +11,11 @@ import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; import org.opensearch.common.lifecycle.AbstractLifecycleComponent; +import org.opensearch.common.settings.ClusterSettings; import org.opensearch.common.util.concurrent.ConcurrentCollections; import org.opensearch.common.util.concurrent.ConcurrentMapLong; import org.opensearch.datafusion.core.GlobalRuntimeEnv; +import org.opensearch.datafusion.search.cache.CacheManager; import org.opensearch.vectorized.execution.search.DataFormat; import org.opensearch.vectorized.execution.search.spi.DataSourceCodec; import org.opensearch.vectorized.execution.search.spi.RecordBatchStream; @@ -31,17 +33,26 @@ public class DataFusionService extends AbstractLifecycleComponent { private final ConcurrentMapLong sessionEngines = ConcurrentCollections.newConcurrentMapLongWithAggressiveConcurrency(); private final DataSourceRegistry dataSourceRegistry; + private final Long tokioRuntimePtr; private final GlobalRuntimeEnv globalRuntimeEnv; + private CacheManager cacheManager; /** * Creates a new DataFusion service instance. */ - public DataFusionService(Map dataSourceCodecs) { + + public DataFusionService(Map dataSourceCodecs, ClusterSettings clusterSettings) { this.dataSourceRegistry = new DataSourceRegistry(dataSourceCodecs); // to verify jni String version = DataFusionQueryJNI.getVersionInfo(); - this.globalRuntimeEnv = new GlobalRuntimeEnv(); + this.tokioRuntimePtr = DataFusionQueryJNI.createTokioRuntime(); + this.globalRuntimeEnv = new GlobalRuntimeEnv(clusterSettings); + this.cacheManager = globalRuntimeEnv.getCacheManager(); + } + + public Long getTokioRuntimePointer() { + return tokioRuntimePtr; } @Override @@ -165,9 +176,9 @@ public long getRuntimePointer() { return globalRuntimeEnv.getPointer(); } - public long getTokioRuntimePointer() { - return globalRuntimeEnv.getTokioRuntimePtr(); - } + //public long getTokioRuntimePointer() { + // return globalRuntimeEnv.getTokioRuntimePtr(); + // } /** * Close the session context and clean up resources @@ -207,4 +218,8 @@ public String getVersion() { version.append("]}"); return version.toString(); } + + public CacheManager getCacheManager() { + return cacheManager; + } } diff --git a/plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/DatafusionEngine.java b/plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/DatafusionEngine.java index e7a49daf970f9..ac4c1ddd327b6 100644 --- a/plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/DatafusionEngine.java +++ b/plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/DatafusionEngine.java @@ -17,6 +17,7 @@ import org.opensearch.action.search.SearchShardTask; import org.opensearch.common.lease.Releasables; import org.opensearch.common.util.BigArrays; +import java.util.List; import org.opensearch.datafusion.core.DefaultRecordBatchStream; import org.opensearch.datafusion.search.DatafusionContext; import org.opensearch.datafusion.search.DatafusionQuery; @@ -25,6 +26,8 @@ import org.opensearch.datafusion.search.DatafusionReaderManager; import org.opensearch.datafusion.search.DatafusionSearcher; import org.opensearch.datafusion.search.DatafusionSearcherSupplier; +import org.opensearch.datafusion.search.cache.CacheManager; +import org.opensearch.datafusion.search.cache.CacheType; import org.opensearch.index.engine.CatalogSnapshotAwareRefreshListener; import org.opensearch.index.engine.Engine; import org.opensearch.index.engine.EngineException; @@ -58,12 +61,17 @@ public class DatafusionEngine extends SearchExecEngine formatCatalogSnapshot, DataFusionService dataFusionService, ShardPath shardPath) throws IOException { this.dataFormat = dataFormat; - this.datafusionReaderManager = new DatafusionReaderManager(shardPath.getDataPath().toString(), formatCatalogSnapshot, dataFormat.getName()); this.datafusionService = dataFusionService; + this.cacheManager = datafusionService.getCacheManager(); + datafusionReaderManager.setOnFilesAdded(files -> { + // Handle new files added during refresh + cacheManager.addToCache(files); + }); } @Override @@ -99,19 +107,22 @@ public EngineSearcherSupplier acquireSearcherSupplier(Functi searcher = new DatafusionSearcherSupplier(null) { @Override protected DatafusionSearcher acquireSearcherInternal(String source) { - return new DatafusionSearcher(source, reader, () -> {}); + return new DatafusionSearcher(source, reader, datafusionService.getTokioRuntimePointer(), + datafusionService.getRuntimePointer(), () -> {}); + } @Override protected void doClose() { try { - reader.decRef(); + datafusionReaderManager.release(reader); } catch (IOException e) { throw new UncheckedIOException(e); } } }; } catch (Exception ex) { + logger.error("Failed to acquire searcher {}", ex.toString(), ex); // TODO } return searcher; @@ -137,6 +148,8 @@ public DatafusionSearcher acquireSearcher(String source, Engine.SearcherScope sc return new DatafusionSearcher( source, searcher.getReader(), + datafusionService.getTokioRuntimePointer(), + datafusionService.getRuntimePointer(), () -> Releasables.close(searcher, searcherSupplier) ); } finally { @@ -164,7 +177,8 @@ public Map execute(DatafusionContext context) { Map finalRes = new HashMap<>(); try { DatafusionSearcher datafusionSearcher = context.getEngineSearcher(); - long streamPointer = datafusionSearcher.search(context.getDatafusionQuery(), datafusionService.getTokioRuntimePointer()); + long streamPointer = datafusionSearcher.search(context.getDatafusionQuery(), + datafusionService.getRuntimePointer(),datafusionService.getTokioRuntimePointer()); RootAllocator allocator = new RootAllocator(Long.MAX_VALUE); RecordBatchStream stream = new RecordBatchStream(streamPointer, datafusionService.getTokioRuntimePointer() , allocator); diff --git a/plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/core/GlobalRuntimeEnv.java b/plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/core/GlobalRuntimeEnv.java index 547539d5ff4d1..320ac61a58e6c 100644 --- a/plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/core/GlobalRuntimeEnv.java +++ b/plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/core/GlobalRuntimeEnv.java @@ -8,6 +8,9 @@ package org.opensearch.datafusion.core; +import org.opensearch.common.settings.ClusterSettings; +import org.opensearch.datafusion.search.cache.CacheManager; + import static org.opensearch.datafusion.DataFusionQueryJNI.closeGlobalRuntime; import static org.opensearch.datafusion.DataFusionQueryJNI.createGlobalRuntime; import static org.opensearch.datafusion.DataFusionQueryJNI.createTokioRuntime; @@ -20,12 +23,12 @@ public class GlobalRuntimeEnv implements AutoCloseable { // ptr to runtime environment in df private final long ptr; private final long tokio_runtime_ptr; + private CacheManager cacheManager; - /** - * Creates a new global runtime environment. - */ - public GlobalRuntimeEnv() { - this.ptr = createGlobalRuntime(); + + public GlobalRuntimeEnv(ClusterSettings clusterSettings) { + this.cacheManager = CacheManager.fromConfig(clusterSettings); + this.ptr = createGlobalRuntime(cacheManager.getCacheManagerPtr()); this.tokio_runtime_ptr = createTokioRuntime(); } @@ -45,4 +48,8 @@ public long getTokioRuntimePtr() { public void close() { closeGlobalRuntime(this.ptr); } + + public CacheManager getCacheManager() { + return cacheManager; + } } diff --git a/plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/search/DatafusionReader.java b/plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/search/DatafusionReader.java index ec01a01b57720..b462da709acb9 100644 --- a/plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/search/DatafusionReader.java +++ b/plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/search/DatafusionReader.java @@ -35,7 +35,7 @@ public class DatafusionReader implements Closeable { * The cache pointer. */ public long cachePtr; - private AtomicInteger refCount = new AtomicInteger(0); + private final AtomicInteger refCount = new AtomicInteger(1); /** * Constructor @@ -56,7 +56,6 @@ public DatafusionReader(String directoryPath, Collection files) { System.out.println("Directory path: " + directoryPath); this.cachePtr = DataFusionQueryJNI.createDatafusionReader(directoryPath, fileNames); - incRef(); } /** @@ -79,15 +78,12 @@ public void incRef() { * @throws IOException if an I/O error occurs */ public void decRef() throws IOException { - if(refCount.get() == 0) { - throw new IllegalStateException("Listing table has been already closed"); - } - - int currRefCount = refCount.decrementAndGet(); - if(currRefCount == 0) { - this.close(); + int count = refCount.decrementAndGet(); + if (count == 0) { + close(); + } else if (count < 0) { + throw new IllegalStateException("Too many decRef calls on DatafusionReader"); } - } @Override @@ -96,7 +92,7 @@ public void close() throws IOException { throw new IllegalStateException("Listing table has been already closed"); } -// closeDatafusionReader(this.cachePtr); + closeDatafusionReader(this.cachePtr); this.cachePtr = -1; } } diff --git a/plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/search/DatafusionReaderManager.java b/plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/search/DatafusionReaderManager.java index 192be34625bd2..81251b3c7b014 100644 --- a/plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/search/DatafusionReaderManager.java +++ b/plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/search/DatafusionReaderManager.java @@ -8,11 +8,16 @@ package org.opensearch.datafusion.search; +import java.util.Arrays; +import java.util.HashSet; +import java.util.Set; +import java.util.function.Consumer; import org.apache.lucene.search.ReferenceManager; import org.opensearch.index.engine.CatalogSnapshotAwareRefreshListener; import org.opensearch.index.engine.EngineReaderManager; import org.opensearch.index.engine.exec.DataFormat; import org.opensearch.index.engine.exec.FileMetadata; +import org.opensearch.index.engine.exec.WriterFileSet; import org.opensearch.index.engine.exec.coord.CatalogSnapshot; import java.io.IOException; @@ -27,6 +32,7 @@ public class DatafusionReaderManager implements EngineReaderManager> onFilesAdded; // private final Lock refreshLock = new ReentrantLock(); // private final List refreshListeners = new CopyOnWriteArrayList(); @@ -36,6 +42,13 @@ public DatafusionReaderManager(String path, Collection files, Stri this.dataFormat = dataFormat; } + /** + * Set callback for when files are added during refresh + */ + public void setOnFilesAdded(Consumer> onFilesAdded) { + this.onFilesAdded = onFilesAdded; + } + @Override public DatafusionReader acquire() throws IOException { if (current == null) { @@ -61,11 +74,46 @@ public void beforeRefresh() throws IOException { public void afterRefresh(boolean didRefresh, CatalogSnapshot catalogSnapshot) throws IOException { if (didRefresh && catalogSnapshot != null) { DatafusionReader old = this.current; - if(old !=null) { + Collection newFiles = catalogSnapshot.getSearchableFiles(dataFormat); + + // Create new reader first + DatafusionReader newReader = new DatafusionReader(this.path, newFiles); + + // Atomically swap current reader + this.current = newReader; + + // Process file changes and release old reader + if (old != null) { + processFileChanges(old.files, newFiles); release(old); + } else { + processFileChanges(List.of(), newFiles); } - this.current = new DatafusionReader(this.path, catalogSnapshot.getSearchableFiles(dataFormat)); - this.current.incRef(); } } + + private void processFileChanges(Collection oldFiles, Collection newFiles) { + Set oldFilePaths = extractFilePaths(oldFiles); + Set newFilePaths = extractFilePaths(newFiles); + + Set filesToAdd = new HashSet<>(newFilePaths); + filesToAdd.removeAll(oldFilePaths); + + // TODO: Either remove files periodically or let eviction handle stale files + Set filesToRemove = new HashSet<>(oldFilePaths); + filesToRemove.removeAll(newFilePaths); + + if (!filesToAdd.isEmpty() && onFilesAdded != null) { + onFilesAdded.accept(List.copyOf(filesToAdd)); + } + } + + private Set extractFilePaths(Collection files) { + String[] fileNames = files.stream() + .flatMap(writerFileSet -> writerFileSet.getFiles().stream()) + .toArray(String[]::new); + Set paths = new HashSet<>(); + paths.addAll(Arrays.asList(fileNames)); + return paths; + } } diff --git a/plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/search/DatafusionSearcher.java b/plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/search/DatafusionSearcher.java index c3a37fbe26063..079e6559b9a6b 100644 --- a/plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/search/DatafusionSearcher.java +++ b/plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/search/DatafusionSearcher.java @@ -9,6 +9,7 @@ package org.opensearch.datafusion.search; import org.apache.lucene.store.AlreadyClosedException; +import org.opensearch.common.util.io.IOUtils; import org.opensearch.datafusion.DataFusionQueryJNI; import org.opensearch.datafusion.DataFusionService; import org.opensearch.datafusion.core.DefaultRecordBatchStream; @@ -25,10 +26,17 @@ public class DatafusionSearcher implements EngineSearcher { private final String source; private DatafusionReader reader; + private Long tokioRuntimePtr; + private Long globalRuntimeEnvId; private Closeable closeable; - public DatafusionSearcher(String source, DatafusionReader reader, Closeable close) { + private boolean readerReleased = false; + + public DatafusionSearcher(String source, DatafusionReader reader, Long tokioRuntimePtr, Long globalRuntimeEnvId, Closeable close) { this.source = source; this.reader = reader; + this.tokioRuntimePtr = tokioRuntimePtr; + this.globalRuntimeEnvId = globalRuntimeEnvId; + this.closeable = close; } @Override @@ -40,7 +48,7 @@ public String source() { public void search(DatafusionQuery datafusionQuery, List> collectors) throws IOException { // TODO : call search here to native // TODO : change RunTimePtr - long nativeStreamPtr = DataFusionQueryJNI.executeSubstraitQuery(reader.getCachePtr(), datafusionQuery.toString(), datafusionQuery.getSubstraitBytes(), 0); + long nativeStreamPtr = DataFusionQueryJNI.executeSubstraitQuery(reader.getCachePtr(), datafusionQuery.toString(), datafusionQuery.getSubstraitBytes(), globalRuntimeEnvId, tokioRuntimePtr); RecordBatchStream stream = new DefaultRecordBatchStream(nativeStreamPtr); while(stream.hasNext()) { for(SearchResultsCollector collector : collectors) { @@ -49,9 +57,14 @@ public void search(DatafusionQuery datafusionQuery, List properties; + protected long pointer; + protected long sizeLimit; + + public void setSizeLimit(ByteSizeValue sizeLimit) { + this.sizeLimit = sizeLimit.getBytes(); + } + + public long getPointer() { + return pointer; + } + + public long getConfiguredSizeLimit() { + return this.sizeLimit; + } + + public String getName() { + return name.toString(); + } + + public CacheAccessor(long cacheManagerPointer, ClusterSettings cacheSettings, CacheType name) { + this.properties = extractSettings(cacheSettings); + this.pointer = createCache(cacheManagerPointer, properties); + this.name = name; + } + + // Abstract method - subclasses define what settings to extract + protected abstract Map extractSettings(ClusterSettings clusterSettings); + public abstract long createCache(long cacheManagerPointer, Map properties); + + // Instance methods matching native signatures + public abstract boolean put(String filePath); + public abstract Object get(String filePath); + public abstract boolean remove(String filePath); + public abstract void evict(); + public abstract void clear(); + public abstract long getMemoryConsumed(); + public abstract boolean containsFile(String filePath); + public abstract List getEntries(); +} diff --git a/plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/search/cache/CacheManager.java b/plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/search/cache/CacheManager.java new file mode 100644 index 0000000000000..8794907e88da1 --- /dev/null +++ b/plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/search/cache/CacheManager.java @@ -0,0 +1,118 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.datafusion.search.cache; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.stream.Collectors; +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.opensearch.common.settings.ClusterSettings; + +import static org.opensearch.datafusion.DataFusionQueryJNI.initCacheManagerConfig; + + +public class CacheManager { + private static final Logger logger = LogManager.getLogger(CacheManager.class); + + private Map caches; + private long cacheManagerPtr; + private long totalSizeLimit; + + public long getCacheManagerPtr() { + return this.cacheManagerPtr; + } + + public long getTotalSizeLimit() { + return totalSizeLimit; + } + + public CacheManager(long cacheManagerPtr, Map cacheMap) { + this.caches = new HashMap<>(cacheMap); + this.cacheManagerPtr = cacheManagerPtr; + this.totalSizeLimit = caches.values().stream() + .mapToLong(CacheAccessor::getConfiguredSizeLimit) + .sum(); + } + + // Factory method to create CacheManager from config + public static CacheManager fromConfig(ClusterSettings clusterSettings) { + long cacheManagerPtr = initCacheManagerConfig(); + Map cacheMap = new HashMap<>(); + for (CacheType type : CacheType.values()) { + if (type.isEnabled(clusterSettings)) { + cacheMap.put(type, type.createCache(cacheManagerPtr, clusterSettings)); + } + } + return new CacheManager(cacheManagerPtr, cacheMap); + } + + public CacheAccessor getCacheAccessor(CacheType cacheType) { + return caches.get(cacheType); + } + + public List getAllCaches() { + return new ArrayList<>(caches.values()); + } + + public List getCachesByType(CacheType... types) { + return caches.entrySet().stream() + .filter(entry -> List.of(types).contains(entry.getKey())) + .map(Map.Entry::getValue) + .collect(Collectors.toList()); + } + + public boolean removeFiles(List files) { + boolean allSuccessful = true; + for (CacheAccessor cache : getAllCaches()) { + for (String filename : files) { + allSuccessful &= cache.remove(filename); + } + } + return allSuccessful; + } + + public boolean removeFilesByDirectory(String path) { + boolean allSuccessful = true; + for (CacheAccessor cache : getAllCaches()) { + allSuccessful &= cache.remove(path); + } + return allSuccessful; + } + + public boolean addToCache(List files) { + boolean allSuccessful = true; + for (CacheAccessor cache : getAllCaches()) { + for (String filename : files) { + allSuccessful &= cache.put(filename); + } + } + return allSuccessful; + } + + public void resetCache() { + getAllCaches().forEach(CacheAccessor::clear); + } + + public long getTotalUsedBytes() { + return getAllCaches().stream().mapToLong(CacheAccessor::getMemoryConsumed).sum(); + } + + public boolean withinCacheLimit(CacheType cacheType) { + CacheAccessor cacheAccessor = getCacheAccessor(cacheType); + return cacheAccessor.getMemoryConsumed() < cacheAccessor.getConfiguredSizeLimit(); + } + + public boolean withinTotalLimit() { + return getTotalUsedBytes() < totalSizeLimit; + } + +} diff --git a/plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/search/cache/CachePolicy.java b/plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/search/cache/CachePolicy.java new file mode 100644 index 0000000000000..da83f83df241b --- /dev/null +++ b/plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/search/cache/CachePolicy.java @@ -0,0 +1,6 @@ +package org.opensearch.datafusion.search.cache; + +public enum CachePolicy { + LFU, + LRU +} diff --git a/plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/search/cache/CacheSettings.java b/plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/search/cache/CacheSettings.java new file mode 100644 index 0000000000000..4c996956bbd75 --- /dev/null +++ b/plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/search/cache/CacheSettings.java @@ -0,0 +1,67 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.datafusion.search.cache; + +import java.util.Arrays; +import java.util.List; +import java.util.function.Function; +import org.opensearch.common.settings.Setting; +import org.opensearch.core.common.unit.ByteSizeUnit; +import org.opensearch.core.common.unit.ByteSizeValue; + +public class CacheSettings { + + public static final String STATS_CACHE_SIZE_LIMIT_KEY = "datafusion.statistics.cache.size.limit"; + public static final String METADATA_CACHE_SIZE_LIMIT_KEY = "datafusion.metadata.cache.size.limit"; + public static final Setting METADATA_CACHE_SIZE_LIMIT = + new Setting<>(METADATA_CACHE_SIZE_LIMIT_KEY, "50mb", + (s) -> ByteSizeValue.parseBytesSizeValue(s, new ByteSizeValue(1000, ByteSizeUnit.KB),METADATA_CACHE_SIZE_LIMIT_KEY), Setting.Property.NodeScope, Setting.Property.Dynamic); + + public static final Setting STATS_CACHE_SIZE_LIMIT = + new Setting<>(STATS_CACHE_SIZE_LIMIT_KEY, "10mb", + (s) -> ByteSizeValue.parseBytesSizeValue(s, new ByteSizeValue(1000, ByteSizeUnit.KB),STATS_CACHE_SIZE_LIMIT_KEY), Setting.Property.NodeScope, Setting.Property.Dynamic); + + + public static final Setting METADATA_CACHE_EVICTION_TYPE = new Setting( + "datafusion.metadata.cache.eviction.type", + "LRU", + Function.identity(), + Setting.Property.NodeScope, + Setting.Property.Dynamic + ); + + public static final Setting STATS_CACHE_EVICTION_TYPE = new Setting( + "datafusion.statistics.cache.eviction.type", + "LFU", + Function.identity(), + Setting.Property.NodeScope, + Setting.Property.Dynamic + ); + + + public static final String METADATA_CACHE_ENABLED_KEY = "datafusion.metadata.cache.enabled"; + public static final Setting METADATA_CACHE_ENABLED = + Setting.boolSetting(METADATA_CACHE_ENABLED_KEY, true, Setting.Property.NodeScope, Setting.Property.Dynamic); + + public static final String STATS_CACHE_ENABLED_KEY = "datafusion.statistics.cache.enabled"; + public static final Setting STATS_CACHE_ENABLED = + Setting.boolSetting(STATS_CACHE_ENABLED_KEY, true, Setting.Property.NodeScope, Setting.Property.Dynamic); + + public static final List> CACHE_SETTINGS = Arrays.asList( + METADATA_CACHE_SIZE_LIMIT, + METADATA_CACHE_EVICTION_TYPE, + STATS_CACHE_SIZE_LIMIT, + STATS_CACHE_EVICTION_TYPE + ); + + public static final List> CACHE_ENABLED = Arrays.asList( + METADATA_CACHE_ENABLED, + STATS_CACHE_ENABLED + ); +} diff --git a/plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/search/cache/CacheType.java b/plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/search/cache/CacheType.java new file mode 100644 index 0000000000000..aa4ff6669d9e0 --- /dev/null +++ b/plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/search/cache/CacheType.java @@ -0,0 +1,33 @@ +package org.opensearch.datafusion.search.cache; + +import org.opensearch.common.settings.ClusterSettings; +import org.opensearch.common.settings.Setting; + +import static org.opensearch.datafusion.search.cache.CacheSettings.METADATA_CACHE_ENABLED; +import static org.opensearch.datafusion.search.cache.CacheSettings.STATS_CACHE_ENABLED; + +public enum CacheType { + METADATA(MetadataCacheAccessor::new, METADATA_CACHE_ENABLED), + STATISTICS(StatisticsCacheAccessor::new, STATS_CACHE_ENABLED); + + private final CacheFactory factory; + private final Setting enabledSetting; + + CacheType(CacheFactory factory, Setting enabledSetting) { + this.factory = factory; + this.enabledSetting = enabledSetting; + } + + public CacheAccessor createCache(long cacheManagerPointer, ClusterSettings clusterSettings) { + return factory.create(cacheManagerPointer, clusterSettings, this); + } + + public boolean isEnabled(ClusterSettings clusterSettings) { + return clusterSettings.get(enabledSetting); + } + + @FunctionalInterface + private interface CacheFactory { + CacheAccessor create(long cacheManagerPointer, ClusterSettings clusterSettings, CacheType type); + } +} diff --git a/plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/search/cache/MetadataCacheAccessor.java b/plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/search/cache/MetadataCacheAccessor.java new file mode 100644 index 0000000000000..ccc883da6ba37 --- /dev/null +++ b/plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/search/cache/MetadataCacheAccessor.java @@ -0,0 +1,144 @@ +package org.opensearch.datafusion.search.cache; + +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.opensearch.common.settings.ClusterSettings; +import org.opensearch.core.common.unit.ByteSizeValue; + +import static org.opensearch.datafusion.DataFusionQueryJNI.createMetadataCache; +import static org.opensearch.datafusion.DataFusionQueryJNI.metadataCacheClear; +import static org.opensearch.datafusion.DataFusionQueryJNI.metadataCacheContainsFile; +import static org.opensearch.datafusion.DataFusionQueryJNI.metadataCacheGet; +import static org.opensearch.datafusion.DataFusionQueryJNI.metadataCacheGetEntries; +import static org.opensearch.datafusion.DataFusionQueryJNI.metadataCacheGetSize; +import static org.opensearch.datafusion.DataFusionQueryJNI.metadataCachePut; +import static org.opensearch.datafusion.DataFusionQueryJNI.metadataCacheRemove; +import static org.opensearch.datafusion.DataFusionQueryJNI.metadataCacheUpdateSizeLimit; +import static org.opensearch.datafusion.search.cache.CacheSettings.METADATA_CACHE_EVICTION_TYPE; +import static org.opensearch.datafusion.search.cache.CacheSettings.METADATA_CACHE_SIZE_LIMIT; +import static org.opensearch.datafusion.search.cache.CacheSettings.METADATA_CACHE_SIZE_LIMIT_KEY; + +public class MetadataCacheAccessor extends CacheAccessor { + private static final Logger logger = LogManager.getLogger(MetadataCacheAccessor.class); + private CachePolicy cachePolicy; + + public MetadataCacheAccessor(long cacheManagerPointer, ClusterSettings settings, CacheType type) { + super(cacheManagerPointer, settings,type); + } + + public void setCachePolicy(String cachePolicy) { + this.cachePolicy = CachePolicy.valueOf(cachePolicy); + } + + @Override + protected Map extractSettings(ClusterSettings clusterSettings) { + Map properties = new HashMap<>(); + + clusterSettings.addSettingsUpdateConsumer(METADATA_CACHE_SIZE_LIMIT, this::setSizeLimit); + setSizeLimit(clusterSettings.get(METADATA_CACHE_SIZE_LIMIT)); + properties.put(METADATA_CACHE_SIZE_LIMIT_KEY,this.sizeLimit); + + clusterSettings.addSettingsUpdateConsumer(METADATA_CACHE_EVICTION_TYPE, this::setCachePolicy); + setCachePolicy(clusterSettings.get(METADATA_CACHE_EVICTION_TYPE)); + properties.put(METADATA_CACHE_EVICTION_TYPE.getKey(),this.cachePolicy); + + return properties; + } + + @Override + public long createCache(long cacheManagerPointer, Map properties) { + return createMetadataCache(cacheManagerPointer, this.sizeLimit); + } + + @Override + public boolean put(String filePath) { + try { + return metadataCachePut(this.getPointer(), filePath); + } catch (RuntimeException e) { + logger.error("Failed to put file [{}] in metadata cache: {}", filePath, e.getMessage()); + return false; + } + } + + @Override + public Object get(String filePath) { + try { + return metadataCacheGet(this.pointer, filePath); + } catch (RuntimeException e) { + logger.error("Failed to get file [{}] from metadata cache: {}", filePath, e.getMessage()); + return null; + } + } + + @Override + public boolean remove(String filePath) { + try { + return metadataCacheRemove(this.pointer, filePath); + } catch (RuntimeException e) { + logger.error("Failed to remove file [{}] from metadata cache: {}", filePath, e.getMessage()); + return false; + } + } + + @Override + public void evict() { + throw new UnsupportedOperationException("Explicit Eviction Not Supported"); + } + + @Override + public void clear() { + try { + metadataCacheClear(this.pointer); + } catch (RuntimeException e) { + logger.error("Failed to clear metadata cache: {}", e.getMessage()); + } + } + + @Override + public long getMemoryConsumed() { + try { + return metadataCacheGetSize(this.pointer); + } catch (RuntimeException e) { + logger.warn("Failed to get metadata cache size: {}", e.getMessage()); + return 0L; + } + } + + @Override + public boolean containsFile(String filePath) { + try { + return metadataCacheContainsFile(this.pointer, filePath); + } catch (RuntimeException e) { + logger.error("Failed to check if metadata cache contains file [{}]: {}", filePath, e.getMessage()); + return false; + } + } + + @Override + public void setSizeLimit(ByteSizeValue limit) { + if(this.sizeLimit == 0){ + this.sizeLimit = limit.getBytes(); + } else{ + try { + metadataCacheUpdateSizeLimit(this.pointer, limit.getBytes()); + this.sizeLimit = limit.getBytes(); + } catch (RuntimeException e) { + logger.error("Failed to update metadata cache size limit to [{}]: {}", limit, e.getMessage()); + } + } + } + + @Override + public List getEntries() { + try { + return List.of(metadataCacheGetEntries(this.pointer)); + } catch (RuntimeException e) { + logger.error("Failed to get metadata cache entries: {}", e.getMessage()); + return List.of(); + } + } + +} diff --git a/plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/search/cache/StatisticsCacheAccessor.java b/plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/search/cache/StatisticsCacheAccessor.java new file mode 100644 index 0000000000000..d91594f6d77e7 --- /dev/null +++ b/plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/search/cache/StatisticsCacheAccessor.java @@ -0,0 +1,148 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.datafusion.search.cache; + +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import org.opensearch.common.settings.ClusterSettings; +import org.opensearch.core.common.unit.ByteSizeValue; + +import static org.opensearch.datafusion.DataFusionQueryJNI.createStatisticsCache; +import static org.opensearch.datafusion.DataFusionQueryJNI.statisticsCacheClear; +import static org.opensearch.datafusion.DataFusionQueryJNI.statisticsCacheContainsFile; +import static org.opensearch.datafusion.DataFusionQueryJNI.statisticsCacheGet; +import static org.opensearch.datafusion.DataFusionQueryJNI.statisticsCacheGetEntries; +import static org.opensearch.datafusion.DataFusionQueryJNI.statisticsCacheGetHitCount; +import static org.opensearch.datafusion.DataFusionQueryJNI.statisticsCacheGetHitRate; +import static org.opensearch.datafusion.DataFusionQueryJNI.statisticsCacheGetMissCount; +import static org.opensearch.datafusion.DataFusionQueryJNI.statisticsCacheGetSize; +import static org.opensearch.datafusion.DataFusionQueryJNI.statisticsCachePut; +import static org.opensearch.datafusion.DataFusionQueryJNI.statisticsCacheRemove; +import static org.opensearch.datafusion.DataFusionQueryJNI.statisticsCacheResetStats; +import static org.opensearch.datafusion.DataFusionQueryJNI.statisticsCacheUpdateSizeLimit; +import static org.opensearch.datafusion.search.cache.CacheSettings.STATS_CACHE_EVICTION_TYPE; +import static org.opensearch.datafusion.search.cache.CacheSettings.STATS_CACHE_SIZE_LIMIT; +import static org.opensearch.datafusion.search.cache.CacheSettings.STATS_CACHE_SIZE_LIMIT_KEY; + +public class StatisticsCacheAccessor extends CacheAccessor{ + + private CachePolicy cachePolicy; + + public void setCachePolicy(String cachePolicy) { + this.cachePolicy = CachePolicy.valueOf(cachePolicy); + } + + + public StatisticsCacheAccessor(long cacheManagerPointer, ClusterSettings cacheSettings, CacheType name) { + super(cacheManagerPointer, cacheSettings, name); + } + + @Override + protected Map extractSettings(ClusterSettings clusterSettings) { + Map properties = new HashMap<>(); + + clusterSettings.addSettingsUpdateConsumer(STATS_CACHE_SIZE_LIMIT, this::setSizeLimit); + setSizeLimit(clusterSettings.get(STATS_CACHE_SIZE_LIMIT)); + properties.put(STATS_CACHE_SIZE_LIMIT_KEY,this.sizeLimit); + + clusterSettings.addSettingsUpdateConsumer(STATS_CACHE_EVICTION_TYPE, this::setCachePolicy); + setCachePolicy(clusterSettings.get(STATS_CACHE_EVICTION_TYPE)); + properties.put(STATS_CACHE_EVICTION_TYPE.getKey(),this.cachePolicy); + + return properties; + } + + @Override + public long createCache(long cacheManagerPointer, Map properties) { + return createStatisticsCache(cacheManagerPointer, this.sizeLimit); + } + + @Override + public boolean put(String filePath) { + return statisticsCachePut(this.getPointer(), filePath); + } + + @Override + public Object get(String filePath) { + return statisticsCacheGet(this.pointer, filePath); + } + + @Override + public boolean remove(String filePath) { + return statisticsCacheRemove(this.pointer, filePath); + } + + @Override + public void evict() { + return; + } + + @Override + public void clear() { + statisticsCacheClear(this.pointer); + } + + @Override + public long getMemoryConsumed() { + return statisticsCacheGetSize(this.pointer); + } + + @Override + public boolean containsFile(String filePath) { + return statisticsCacheContainsFile(this.pointer, filePath); + } + + // TODO: Replace the logic with optimized version to check if it is update or set limit call + @Override + public void setSizeLimit(ByteSizeValue limit) { + if(this.sizeLimit == 0){ + this.sizeLimit = limit.getBytes(); + } else{ + statisticsCacheUpdateSizeLimit(this.pointer, limit.getBytes()); + this.sizeLimit = limit.getBytes(); + } + } + + @Override + public List getEntries() { + return List.of(statisticsCacheGetEntries(this.pointer)); + } + + /** + * Get cache hit count + * @return number of cache hits + */ + public long getHitCount() { + return statisticsCacheGetHitCount(this.pointer); + } + + /** + * Get cache miss count + * @return number of cache misses + */ + public long getMissCount() { + return statisticsCacheGetMissCount(this.pointer); + } + + /** + * Get cache hit rate + * @return hit rate as a value between 0.0 and 1.0 + */ + public double getHitRate() { + return statisticsCacheGetHitRate(this.pointer); + } + + /** + * Reset hit and miss counters + */ + public void resetStats() { + statisticsCacheResetStats(this.pointer); + } +} diff --git a/plugins/engine-datafusion/src/test/java/org/opensearch/datafusion/DataFusionCacheManagerTests.java b/plugins/engine-datafusion/src/test/java/org/opensearch/datafusion/DataFusionCacheManagerTests.java new file mode 100644 index 0000000000000..dc46e650273aa --- /dev/null +++ b/plugins/engine-datafusion/src/test/java/org/opensearch/datafusion/DataFusionCacheManagerTests.java @@ -0,0 +1,408 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.datafusion; + +import java.io.File; +import java.net.URL; +import java.util.Collections; +import java.util.HashSet; +import java.util.List; +import java.util.Set; +import org.junit.Before; +import org.mockito.Mock; +import org.mockito.MockitoAnnotations; +import org.opensearch.common.settings.ClusterSettings; +import org.opensearch.common.settings.Setting; +import org.opensearch.common.settings.Settings; +import org.opensearch.core.common.unit.ByteSizeValue; +import org.opensearch.datafusion.search.cache.CacheAccessor; +import org.opensearch.datafusion.search.cache.CacheManager; +import org.opensearch.datafusion.search.cache.CacheType; +import org.opensearch.env.Environment; +import org.opensearch.test.OpenSearchTestCase; + +import static org.mockito.Mockito.when; +import static org.opensearch.common.settings.ClusterSettings.BUILT_IN_CLUSTER_SETTINGS; +import static org.opensearch.datafusion.search.cache.CacheSettings.METADATA_CACHE_ENABLED; +import static org.opensearch.datafusion.search.cache.CacheSettings.METADATA_CACHE_EVICTION_TYPE; +import static org.opensearch.datafusion.search.cache.CacheSettings.METADATA_CACHE_SIZE_LIMIT; +import static org.opensearch.datafusion.search.cache.CacheSettings.STATS_CACHE_ENABLED; +import static org.opensearch.datafusion.search.cache.CacheSettings.STATS_CACHE_EVICTION_TYPE; +import static org.opensearch.datafusion.search.cache.CacheSettings.STATS_CACHE_SIZE_LIMIT; + +public class DataFusionCacheManagerTests extends OpenSearchTestCase { + private DataFusionService service; + + @Mock + private Environment mockEnvironment; + + @Before + public void setup() { + MockitoAnnotations.openMocks(this); + Settings mockSettings = Settings.builder().put("path.data", "/tmp/test-data").build(); + + when(mockEnvironment.settings()).thenReturn(mockSettings); + Set> clusterSettingsToAdd = new HashSet<>(BUILT_IN_CLUSTER_SETTINGS); + clusterSettingsToAdd.add(METADATA_CACHE_ENABLED); + clusterSettingsToAdd.add(METADATA_CACHE_SIZE_LIMIT); + clusterSettingsToAdd.add(METADATA_CACHE_EVICTION_TYPE); + clusterSettingsToAdd.add(STATS_CACHE_ENABLED); + clusterSettingsToAdd.add(STATS_CACHE_SIZE_LIMIT); + clusterSettingsToAdd.add(STATS_CACHE_EVICTION_TYPE); + + ClusterSettings clusterSettings = new ClusterSettings(Settings.EMPTY, clusterSettingsToAdd); + + service = new DataFusionService(Collections.emptyMap(), clusterSettings); + service.doStart(); + } + + public void testAddFileToCache() { + CacheManager cacheManager = service.getCacheManager(); + CacheAccessor metadataCache = cacheManager.getCacheAccessor(CacheType.METADATA); + CacheAccessor statisticsCache = cacheManager.getCacheAccessor(CacheType.STATISTICS); + String fileName = getResourceFile("hits1.parquet").getPath(); + + cacheManager.addToCache(List.of(fileName)); + + assertTrue((Boolean) metadataCache.get(fileName)); + assertTrue((Boolean) statisticsCache.get(fileName)); + assertTrue(metadataCache.containsFile(fileName)); + assertTrue(statisticsCache.containsFile(fileName)); + assertTrue(metadataCache.getMemoryConsumed() > 0); + } + + public void testRemoveFileFromCache() { + CacheManager cacheManager = service.getCacheManager(); + CacheAccessor metadataCache = cacheManager.getCacheAccessor(CacheType.METADATA); + CacheAccessor statisticsCache = cacheManager.getCacheAccessor(CacheType.STATISTICS); + String fileName = getResourceFile("hits1.parquet").getPath(); + + cacheManager.addToCache(List.of(fileName)); + assertTrue(metadataCache.containsFile(fileName)); + assertTrue(statisticsCache.containsFile(fileName)); + + boolean removed = cacheManager.removeFiles(List.of(fileName)); + + assertTrue(removed); + assertFalse(metadataCache.containsFile(fileName)); + assertFalse(statisticsCache.containsFile(fileName)); + } + + public void testCacheSizeLimitEviction() { + CacheAccessor metadataCache = service.getCacheManager().getCacheAccessor(CacheType.METADATA); + String fileName = getResourceFile("hits1.parquet").getPath(); + + metadataCache.put(fileName); + assertTrue(metadataCache.containsFile(fileName)); + + metadataCache.setSizeLimit(new ByteSizeValue(40)); + + assertFalse(metadataCache.containsFile(fileName)); + assertEquals(0, metadataCache.getEntries().size()); + } + + public void testCachePutWithIncreasedSizeLimit() { + CacheAccessor metadataCache = service.getCacheManager().getCacheAccessor(CacheType.METADATA); + String fileName = getResourceFile("hits1.parquet").getPath(); + + metadataCache.setSizeLimit(new ByteSizeValue(500000)); + metadataCache.put(fileName); + + assertTrue(metadataCache.containsFile(fileName)); + logger.info("Entries: {}", metadataCache.getEntries()); + //(we print 3 elements per entry : filePath, memorySize, HitCount) + assertEquals(1*3, metadataCache.getEntries().size()); + } + + public void testCacheClear() { + CacheAccessor metadataCache = service.getCacheManager().getCacheAccessor(CacheType.METADATA); + String fileName = getResourceFile("hits1.parquet").getPath(); + + metadataCache.put(fileName); + assertTrue(metadataCache.containsFile(fileName)); + + metadataCache.clear(); + + assertFalse(metadataCache.containsFile(fileName)); + assertEquals(0, metadataCache.getEntries().size()); + } + + public void testAddMultipleFilesToCache() { + CacheManager cacheManager = service.getCacheManager(); + CacheAccessor metadataCache = cacheManager.getCacheAccessor(CacheType.METADATA); + List fileNames = List.of( + getResourceFile("hits1.parquet").getPath(), + getResourceFile("hits2.parquet").getPath() + ); + + cacheManager.addToCache(fileNames); + // 3 elements per cache entry displayed + assertEquals(2*3, metadataCache.getEntries().size()); + fileNames.forEach(fileName -> assertTrue(metadataCache.containsFile(fileName))); + } + + public void testRemoveNonExistentFile() { + CacheManager cacheManager = service.getCacheManager(); + String nonExistentFile = "/path/nonexistent.parquet"; + + boolean removed = cacheManager.removeFiles(List.of(nonExistentFile)); + + assertFalse(removed); + } + + public void testGetNonExistentFile() { + CacheAccessor metadataCache = service.getCacheManager().getCacheAccessor(CacheType.METADATA); + String nonExistentFile = "/path/nonexistent.parquet"; + + Object result = metadataCache.get(nonExistentFile); + +// assertNull(result); + assertFalse(metadataCache.containsFile(nonExistentFile)); + } + + public void testAddEmptyFileList() { + CacheManager cacheManager = service.getCacheManager(); + CacheAccessor metadataCache = cacheManager.getCacheAccessor(CacheType.METADATA); + + cacheManager.addToCache(Collections.emptyList()); + + assertEquals(0, metadataCache.getEntries().size()); + } + + public void testStatisticsCacheOperations() { + CacheManager cacheManager = service.getCacheManager(); + CacheAccessor statisticsCache = cacheManager.getCacheAccessor(CacheType.STATISTICS); + String fileName = getResourceFile("hits1.parquet").getPath(); + + statisticsCache.put(fileName); + assertTrue(statisticsCache.containsFile(fileName)); + assertTrue(statisticsCache.getMemoryConsumed() > 0); + + statisticsCache.remove(fileName); + assertFalse(statisticsCache.containsFile(fileName)); + } + + public void testStatisticsCacheSizeLimit() { + CacheAccessor statisticsCache = service.getCacheManager().getCacheAccessor(CacheType.STATISTICS); + String fileName = getResourceFile("hits1.parquet").getPath(); + + statisticsCache.put(fileName); + logger.info(statisticsCache.getEntries()); + assertTrue(statisticsCache.containsFile(fileName)); + + statisticsCache.setSizeLimit(new ByteSizeValue(40)); + assertFalse(statisticsCache.containsFile(fileName)); + } + + public void testStatisticsCacheClear() { + CacheAccessor statisticsCache = service.getCacheManager().getCacheAccessor(CacheType.STATISTICS); + String fileName = getResourceFile("hits1.parquet").getPath(); + + statisticsCache.put(fileName); + assertTrue(statisticsCache.containsFile(fileName)); + + statisticsCache.clear(); + assertFalse(statisticsCache.containsFile(fileName)); + } + + public void testBothCacheTypesMemoryTracking() { + CacheManager cacheManager = service.getCacheManager(); + CacheAccessor metadataCache = cacheManager.getCacheAccessor(CacheType.METADATA); + CacheAccessor statisticsCache = cacheManager.getCacheAccessor(CacheType.STATISTICS); + String fileName = getResourceFile("hits1.parquet").getPath(); + + long initialTotal = cacheManager.getTotalUsedBytes(); + + metadataCache.put(fileName); + statisticsCache.put(fileName); + + long afterBothAdded = cacheManager.getTotalUsedBytes(); + assertTrue(afterBothAdded > initialTotal); + + long metadataMemory = metadataCache.getMemoryConsumed(); + long statisticsMemory = statisticsCache.getMemoryConsumed(); + assertEquals(afterBothAdded, initialTotal + metadataMemory + statisticsMemory); + } + + public void testCacheManagerTotalMemoryTracking() { + CacheManager cacheManager = service.getCacheManager(); + String fileName = getResourceFile("hits1.parquet").getPath(); + + long initialMemory = cacheManager.getTotalUsedBytes(); + cacheManager.addToCache(List.of(fileName)); + long afterAddMemory = cacheManager.getTotalUsedBytes(); + + assertTrue(afterAddMemory > initialMemory); + + cacheManager.removeFiles(List.of(fileName)); + long afterRemoveMemory = cacheManager.getTotalUsedBytes(); + + assertEquals(initialMemory, afterRemoveMemory); + } + + public void testCacheSizeLimits() { + CacheManager cacheManager = service.getCacheManager(); + CacheAccessor metadataCache = cacheManager.getCacheAccessor(CacheType.METADATA); + + long configuredLimit = metadataCache.getConfiguredSizeLimit(); + long totalLimit = cacheManager.getTotalSizeLimit(); + + assertTrue(configuredLimit > 0); + assertTrue(totalLimit > 0); + } + + public void testStatisticsCacheHitCount() { + CacheAccessor statisticsCache = service.getCacheManager().getCacheAccessor(CacheType.STATISTICS); + String fileName = getResourceFile("hits1.parquet").getPath(); + + // Initially zero + assertEquals(0, ((org.opensearch.datafusion.search.cache.StatisticsCacheAccessor) statisticsCache).getHitCount()); + assertEquals(0, ((org.opensearch.datafusion.search.cache.StatisticsCacheAccessor) statisticsCache).getMissCount()); + + // Add entry + statisticsCache.put(fileName); + assertTrue(statisticsCache.containsFile(fileName)); + + // Get the entry - should increment hit count + statisticsCache.get(fileName); + assertEquals(1, ((org.opensearch.datafusion.search.cache.StatisticsCacheAccessor) statisticsCache).getHitCount()); + assertEquals(0, ((org.opensearch.datafusion.search.cache.StatisticsCacheAccessor) statisticsCache).getMissCount()); + + // Get it again - should increment hit count again + statisticsCache.get(fileName); + assertEquals(2, ((org.opensearch.datafusion.search.cache.StatisticsCacheAccessor) statisticsCache).getHitCount()); + assertEquals(0, ((org.opensearch.datafusion.search.cache.StatisticsCacheAccessor) statisticsCache).getMissCount()); + } + + public void testStatisticsCacheMissCount() { + CacheAccessor statisticsCache = service.getCacheManager().getCacheAccessor(CacheType.STATISTICS); + String nonExistentFile = "/nonexistent/file.parquet"; + + // Initially zero + assertEquals(0, ((org.opensearch.datafusion.search.cache.StatisticsCacheAccessor) statisticsCache).getHitCount()); + assertEquals(0, ((org.opensearch.datafusion.search.cache.StatisticsCacheAccessor) statisticsCache).getMissCount()); + + // Try to get non-existent entry - should increment miss count + statisticsCache.get(nonExistentFile); + assertEquals(0, ((org.opensearch.datafusion.search.cache.StatisticsCacheAccessor) statisticsCache).getHitCount()); + assertEquals(1, ((org.opensearch.datafusion.search.cache.StatisticsCacheAccessor) statisticsCache).getMissCount()); + + // Try again - should increment miss count again + statisticsCache.get(nonExistentFile); + assertEquals(0, ((org.opensearch.datafusion.search.cache.StatisticsCacheAccessor) statisticsCache).getHitCount()); + assertEquals(2, ((org.opensearch.datafusion.search.cache.StatisticsCacheAccessor) statisticsCache).getMissCount()); + } + + public void testStatisticsCacheHitMissedMixed() { + CacheAccessor statisticsCache = service.getCacheManager().getCacheAccessor(CacheType.STATISTICS); + String fileName = getResourceFile("hits1.parquet").getPath(); + String nonExistentFile = "/nonexistent/file.parquet"; + + // Add entry + statisticsCache.put(fileName); + + // Mix of hits and misses + statisticsCache.get(fileName); // hit + assertEquals(1, ((org.opensearch.datafusion.search.cache.StatisticsCacheAccessor) statisticsCache).getHitCount()); + assertEquals(0, ((org.opensearch.datafusion.search.cache.StatisticsCacheAccessor) statisticsCache).getMissCount()); + + statisticsCache.get(nonExistentFile); // miss + assertEquals(1, ((org.opensearch.datafusion.search.cache.StatisticsCacheAccessor) statisticsCache).getHitCount()); + assertEquals(1, ((org.opensearch.datafusion.search.cache.StatisticsCacheAccessor) statisticsCache).getMissCount()); + + statisticsCache.get(fileName); // hit + assertEquals(2, ((org.opensearch.datafusion.search.cache.StatisticsCacheAccessor) statisticsCache).getHitCount()); + assertEquals(1, ((org.opensearch.datafusion.search.cache.StatisticsCacheAccessor) statisticsCache).getMissCount()); + + statisticsCache.get(nonExistentFile); // miss + assertEquals(2, ((org.opensearch.datafusion.search.cache.StatisticsCacheAccessor) statisticsCache).getHitCount()); + assertEquals(2, ((org.opensearch.datafusion.search.cache.StatisticsCacheAccessor) statisticsCache).getMissCount()); + } + + public void testStatisticsCacheHitRate() { + CacheAccessor statisticsCache = service.getCacheManager().getCacheAccessor(CacheType.STATISTICS); + String fileName = getResourceFile("hits1.parquet").getPath(); + String nonExistentFile = "/nonexistent/file.parquet"; + + // Initially 0.0 (no operations) + assertEquals(0.0, ((org.opensearch.datafusion.search.cache.StatisticsCacheAccessor) statisticsCache).getHitRate(), 0.001); + + // Add entry + statisticsCache.put(fileName); + + // 2 hits, 0 misses = 100% hit rate + statisticsCache.get(fileName); + statisticsCache.get(fileName); + assertEquals(1.0, ((org.opensearch.datafusion.search.cache.StatisticsCacheAccessor) statisticsCache).getHitRate(), 0.001); + + // 2 hits, 1 miss = 66.67% hit rate + statisticsCache.get(nonExistentFile); + assertEquals(0.6666, ((org.opensearch.datafusion.search.cache.StatisticsCacheAccessor) statisticsCache).getHitRate(), 0.001); + + // 3 hits, 1 miss = 75% hit rate + statisticsCache.get(fileName); + assertEquals(0.75, ((org.opensearch.datafusion.search.cache.StatisticsCacheAccessor) statisticsCache).getHitRate(), 0.001); + } + + public void testStatisticsCacheResetStats() { + CacheAccessor statisticsCache = service.getCacheManager().getCacheAccessor(CacheType.STATISTICS); + String fileName = getResourceFile("hits1.parquet").getPath(); + String nonExistentFile = "/nonexistent/file.parquet"; + + // Add entry and generate some hits/misses + statisticsCache.put(fileName); + statisticsCache.get(fileName); // hit + statisticsCache.get(nonExistentFile); // miss + + assertEquals(1, ((org.opensearch.datafusion.search.cache.StatisticsCacheAccessor) statisticsCache).getHitCount()); + assertEquals(1, ((org.opensearch.datafusion.search.cache.StatisticsCacheAccessor) statisticsCache).getMissCount()); + assertEquals(0.5, ((org.opensearch.datafusion.search.cache.StatisticsCacheAccessor) statisticsCache).getHitRate(), 0.001); + + // Reset stats + ((org.opensearch.datafusion.search.cache.StatisticsCacheAccessor) statisticsCache).resetStats(); + + assertEquals(0, ((org.opensearch.datafusion.search.cache.StatisticsCacheAccessor) statisticsCache).getHitCount()); + assertEquals(0, ((org.opensearch.datafusion.search.cache.StatisticsCacheAccessor) statisticsCache).getMissCount()); + assertEquals(0.0, ((org.opensearch.datafusion.search.cache.StatisticsCacheAccessor) statisticsCache).getHitRate(), 0.001); + + // Cache entries should still exist + assertTrue(statisticsCache.containsFile(fileName)); + + // After reset, new operations should start counting from zero + statisticsCache.get(fileName); + assertEquals(1, ((org.opensearch.datafusion.search.cache.StatisticsCacheAccessor) statisticsCache).getHitCount()); + assertEquals(0, ((org.opensearch.datafusion.search.cache.StatisticsCacheAccessor) statisticsCache).getMissCount()); + } + + public void testStatisticsCacheClearResetsStats() { + CacheAccessor statisticsCache = service.getCacheManager().getCacheAccessor(CacheType.STATISTICS); + String fileName = getResourceFile("hits1.parquet").getPath(); + + // Add entry and generate hits/misses + statisticsCache.put(fileName); + statisticsCache.get(fileName); // hit + statisticsCache.get("/nonexistent/file.parquet"); // miss + + assertEquals(1, ((org.opensearch.datafusion.search.cache.StatisticsCacheAccessor) statisticsCache).getHitCount()); + assertEquals(1, ((org.opensearch.datafusion.search.cache.StatisticsCacheAccessor) statisticsCache).getMissCount()); + + // Clear should reset stats + statisticsCache.clear(); + + assertEquals(0, ((org.opensearch.datafusion.search.cache.StatisticsCacheAccessor) statisticsCache).getHitCount()); + assertEquals(0, ((org.opensearch.datafusion.search.cache.StatisticsCacheAccessor) statisticsCache).getMissCount()); + assertFalse(statisticsCache.containsFile(fileName)); + } + + private File getResourceFile(String fileName) { + URL resource = getClass().getClassLoader().getResource(fileName); + assertNotNull("Resource file not found: " + fileName, resource); + return new File(resource.getFile()); + } +} diff --git a/plugins/engine-datafusion/src/test/java/org/opensearch/datafusion/DataFusionServiceTests.java b/plugins/engine-datafusion/src/test/java/org/opensearch/datafusion/DataFusionServiceTests.java index 984d59877ac7e..5531fc055df14 100644 --- a/plugins/engine-datafusion/src/test/java/org/opensearch/datafusion/DataFusionServiceTests.java +++ b/plugins/engine-datafusion/src/test/java/org/opensearch/datafusion/DataFusionServiceTests.java @@ -11,13 +11,19 @@ import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.lucene.search.Query; +import org.opensearch.common.settings.ClusterSettings; +import org.opensearch.common.settings.Setting; import org.opensearch.common.settings.Settings; import org.opensearch.core.common.io.stream.StreamOutput; +import org.opensearch.core.common.unit.ByteSizeValue; import org.opensearch.core.index.shard.ShardId; import org.opensearch.core.xcontent.XContentBuilder; import org.opensearch.datafusion.core.SessionContext; import org.opensearch.datafusion.search.DatafusionQuery; import org.opensearch.datafusion.search.DatafusionSearcher; +import org.opensearch.datafusion.search.cache.CacheAccessor; +import org.opensearch.datafusion.search.cache.CacheManager; +import org.opensearch.datafusion.search.cache.CacheType; import org.opensearch.env.Environment; import org.opensearch.index.engine.exec.FileMetadata; import org.opensearch.index.engine.exec.text.TextDF; @@ -42,6 +48,11 @@ import java.util.*; import static org.mockito.Mockito.when; +import static org.opensearch.common.settings.ClusterSettings.BUILT_IN_CLUSTER_SETTINGS; +import static org.opensearch.datafusion.search.cache.CacheSettings.METADATA_CACHE_ENABLED; +import static org.opensearch.datafusion.search.cache.CacheSettings.METADATA_CACHE_EVICTION_TYPE; +import static org.opensearch.datafusion.search.cache.CacheSettings.METADATA_CACHE_SIZE_LIMIT; + import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.types.pojo.Field; @@ -65,7 +76,15 @@ public void setup() { Settings mockSettings = Settings.builder().put("path.data", "/tmp/test-data").build(); when(mockEnvironment.settings()).thenReturn(mockSettings); - service = new DataFusionService(Map.of()); + Set> clusterSettingsToAdd = new HashSet<>(BUILT_IN_CLUSTER_SETTINGS); + clusterSettingsToAdd.add(METADATA_CACHE_ENABLED); + clusterSettingsToAdd.add(METADATA_CACHE_SIZE_LIMIT); + clusterSettingsToAdd.add(METADATA_CACHE_EVICTION_TYPE); + + ClusterSettings clusterSettings = new ClusterSettings(Settings.EMPTY, clusterSettingsToAdd); + + service = new DataFusionService(Collections.emptyMap(), clusterSettings); + //service = new DataFusionService(Map.of()); service.doStart(); } @@ -97,61 +116,60 @@ public void testGetVersion() { // TO run update proper directory path for generation-1-optimized.parquet file in // this.datafusionReaderManager = new DatafusionReaderManager("TODO://FigureOutPath", formatCatalogSnapshot); - public void testQueryPhaseExecutor() throws IOException { - Map finalRes = new HashMap<>(); - DatafusionSearcher datafusionSearcher = null; - try { - DatafusionEngine engine = new DatafusionEngine(DataFormat.CSV, List.of(new FileMetadata(new TextDF(), "hits_data.parquet")), service); - datafusionSearcher = engine.acquireSearcher("Search"); - - - byte[] protoContent; - - try (InputStream is = getClass().getResourceAsStream("/substrait_plan.pb")) { - protoContent = is.readAllBytes(); - } catch (IOException e) { - throw new RuntimeException(e); - } - - long streamPointer = datafusionSearcher.search(new DatafusionQuery("test-index",protoContent, new ArrayList<>()), service.getTokioRuntimePointer()); - RootAllocator allocator = new RootAllocator(Long.MAX_VALUE); - RecordBatchStream stream = new RecordBatchStream(streamPointer, service.getTokioRuntimePointer() , allocator); - - // We can have some collectors passed like this which can collect the results and convert to InternalAggregation - // Is the possible? need to check - - SearchResultsCollector collector = new SearchResultsCollector() { - @Override - public void collect(RecordBatchStream value) { - VectorSchemaRoot root = value.getVectorSchemaRoot(); - for (Field field : root.getSchema().getFields()) { - String filedName = field.getName(); - FieldVector fieldVector = root.getVector(filedName); - Object[] fieldValues = new Object[fieldVector.getValueCount()]; - for (int i = 0; i < fieldVector.getValueCount(); i++) { - fieldValues[i] = fieldVector.getObject(i); - } - finalRes.put(filedName, fieldValues); - } - } - }; - - while (stream.loadNextBatch().join()) { - collector.collect(stream); - } - - logger.info("Final Results:"); - for (Map.Entry entry : finalRes.entrySet()) { - logger.info("{}: {}", entry.getKey(), java.util.Arrays.toString(entry.getValue())); - } - - } catch (Exception exception) { - logger.error("Failed to execute Substrait query plan", exception); - } - finally { - if(datafusionSearcher != null) { - datafusionSearcher.close(); - } - } - } +// public void testQueryPhaseExecutor() throws IOException { +// Map finalRes = new HashMap<>(); +// DatafusionSearcher datafusionSearcher = null; +// try { +// DatafusionEngine engine = new DatafusionEngine(DataFormat.CSV, List.of(new FileMetadata(new TextDF(), "hits_data.parquet")), service); +// datafusionSearcher = engine.acquireSearcher("Search"); +// +// byte[] protoContent; +// +// try (InputStream is = getClass().getResourceAsStream("/substrait_plan.pb")) { +// protoContent = is.readAllBytes(); +// } catch (IOException e) { +// throw new RuntimeException(e); +// } +// +// long streamPointer = datafusionSearcher.search(new DatafusionQuery("index-7",protoContent, new ArrayList<>()), service.getTokioRuntimePointer(), service.getRuntimePointer()); +// RootAllocator allocator = new RootAllocator(Long.MAX_VALUE); +// RecordBatchStream stream = new RecordBatchStream(streamPointer, service.getTokioRuntimePointer() , allocator); +// +// // We can have some collectors passed like this which can collect the results and convert to InternalAggregation +// // Is the possible? need to check +// +// SearchResultsCollector collector = new SearchResultsCollector() { +// @Override +// public void collect(RecordBatchStream value) { +// VectorSchemaRoot root = value.getVectorSchemaRoot(); +// for (Field field : root.getSchema().getFields()) { +// String filedName = field.getName(); +// FieldVector fieldVector = root.getVector(filedName); +// Object[] fieldValues = new Object[fieldVector.getValueCount()]; +// for (int i = 0; i < fieldVector.getValueCount(); i++) { +// fieldValues[i] = fieldVector.getObject(i); +// } +// finalRes.put(filedName, fieldValues); +// } +// } +// }; +// +// while (stream.loadNextBatch().join()) { +// collector.collect(stream); +// } +// +// logger.info("Final Results:"); +// for (Map.Entry entry : finalRes.entrySet()) { +// logger.info("{}: {}", entry.getKey(), java.util.Arrays.toString(entry.getValue())); +// } +// +// } catch (Exception exception) { +// logger.error("Failed to execute Substrait query plan", exception); +// } +// finally { +// if(datafusionSearcher != null) { +// datafusionSearcher.close(); +// } +// } +// } } diff --git a/plugins/engine-datafusion/src/test/java/org/opensearch/datafusion/TestDataFusionServiceTests.java b/plugins/engine-datafusion/src/test/java/org/opensearch/datafusion/TestDataFusionServiceTests.java index 395e2fae52e2f..49e35adc78721 100644 --- a/plugins/engine-datafusion/src/test/java/org/opensearch/datafusion/TestDataFusionServiceTests.java +++ b/plugins/engine-datafusion/src/test/java/org/opensearch/datafusion/TestDataFusionServiceTests.java @@ -8,11 +8,22 @@ package org.opensearch.datafusion; +import java.util.HashSet; +import java.util.Set; +import org.opensearch.common.settings.ClusterSettings; +import org.opensearch.common.settings.Setting; +import org.opensearch.common.settings.Settings; import org.opensearch.test.OpenSearchTestCase; import java.util.Collections; import java.util.List; +import static org.mockito.Mockito.when; +import static org.opensearch.common.settings.ClusterSettings.BUILT_IN_CLUSTER_SETTINGS; +import static org.opensearch.datafusion.search.cache.CacheSettings.METADATA_CACHE_ENABLED; +import static org.opensearch.datafusion.search.cache.CacheSettings.METADATA_CACHE_EVICTION_TYPE; +import static org.opensearch.datafusion.search.cache.CacheSettings.METADATA_CACHE_SIZE_LIMIT; + /** * Unit tests for DataFusionService * @@ -27,7 +38,14 @@ public class TestDataFusionServiceTests extends OpenSearchTestCase { @Override public void setUp() throws Exception { super.setUp(); - service = new DataFusionService(Collections.emptyMap()); + Set> clusterSettingsToAdd = new HashSet<>(BUILT_IN_CLUSTER_SETTINGS); + clusterSettingsToAdd.add(METADATA_CACHE_ENABLED); + clusterSettingsToAdd.add(METADATA_CACHE_SIZE_LIMIT); + clusterSettingsToAdd.add(METADATA_CACHE_EVICTION_TYPE); + + ClusterSettings clusterSettings = new ClusterSettings(Settings.EMPTY, clusterSettingsToAdd); + + service = new DataFusionService(Collections.emptyMap(),clusterSettings); service.doStart(); } diff --git a/plugins/engine-datafusion/src/test/resources/hits1.parquet b/plugins/engine-datafusion/src/test/resources/hits1.parquet new file mode 100644 index 0000000000000..647d8fb5235c2 Binary files /dev/null and b/plugins/engine-datafusion/src/test/resources/hits1.parquet differ diff --git a/plugins/engine-datafusion/src/test/resources/hits2.parquet b/plugins/engine-datafusion/src/test/resources/hits2.parquet new file mode 100644 index 0000000000000..581c7e502f18b Binary files /dev/null and b/plugins/engine-datafusion/src/test/resources/hits2.parquet differ diff --git a/plugins/engine-datafusion/src/test/resources/hits3.parquet b/plugins/engine-datafusion/src/test/resources/hits3.parquet new file mode 100755 index 0000000000000..7f4dce9a53374 Binary files /dev/null and b/plugins/engine-datafusion/src/test/resources/hits3.parquet differ diff --git a/server/src/main/java/org/opensearch/index/engine/EngineSearcher.java b/server/src/main/java/org/opensearch/index/engine/EngineSearcher.java index 7471fd3fbeb5f..429c1c73f1f73 100644 --- a/server/src/main/java/org/opensearch/index/engine/EngineSearcher.java +++ b/server/src/main/java/org/opensearch/index/engine/EngineSearcher.java @@ -32,7 +32,11 @@ default void search(Q query, List> collectors) throws throw new UnsupportedOperationException(); } - default long search(Q query, Long runtimePtr) throws IOException { + default long search(Q query, Long globalRuntimeEnvId, Long tokioRuntimePtr) throws IOException { + throw new UnsupportedOperationException(); + } + + default long search(Q query) throws IOException { throw new UnsupportedOperationException(); } }