From 5e4f91e70529eb245b0be879b6961db927dd37ac Mon Sep 17 00:00:00 2001 From: Abhita Lakkabathini Date: Thu, 25 Sep 2025 16:39:00 +0530 Subject: [PATCH 1/2] DataFusion Cache Support --- plugins/engine-datafusion/jni/Cargo.toml | 6 +- plugins/engine-datafusion/jni/src/lib.rs | 189 +++++++++++++++++- plugins/engine-datafusion/jni/src/util.rs | 62 ++++-- .../datafusion/DataFusionPlugin.java | 19 +- .../datafusion/DataFusionQueryJNI.java | 60 ++++++ .../datafusion/DataFusionService.java | 17 +- .../datafusion/DatafusionEngine.java | 11 +- .../datafusion/core/GlobalRuntimeEnv.java | 8 + .../search/DatafusionReaderManager.java | 50 ++++- .../search/cache/CacheAccessor.java | 50 +++++ .../datafusion/search/cache/CacheManager.java | 127 ++++++++++++ .../datafusion/search/cache/CachePolicy.java | 6 + .../search/cache/CacheSettings.java | 48 +++++ .../datafusion/search/cache/CacheType.java | 32 +++ .../search/cache/MetadataCacheAccessor.java | 91 +++++++++ .../TestDataFusionServiceTests.java | 55 ++++- .../java/org/opensearch/node/MockNode.java | 7 +- .../opensearch/search/MockSearchService.java | 3 +- 18 files changed, 797 insertions(+), 44 deletions(-) create mode 100644 plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/search/cache/CacheAccessor.java create mode 100644 plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/search/cache/CacheManager.java create mode 100644 plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/search/cache/CachePolicy.java create mode 100644 plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/search/cache/CacheSettings.java create mode 100644 plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/search/cache/CacheType.java create mode 100644 plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/search/cache/MetadataCacheAccessor.java diff --git a/plugins/engine-datafusion/jni/Cargo.toml b/plugins/engine-datafusion/jni/Cargo.toml index 5c062a24d891d..ba1ec24214083 100644 --- a/plugins/engine-datafusion/jni/Cargo.toml +++ b/plugins/engine-datafusion/jni/Cargo.toml @@ -9,9 +9,9 @@ crate-type = ["cdylib"] [dependencies] # DataFusion dependencies -datafusion = "49.0.0" -datafusion-substrait = "49.0.0" -arrow = "55.2.0" +datafusion = "50.0.0" +datafusion-substrait = "50.0.0" +arrow = "56.2.0" arrow-array = "55.2.0" arrow-schema = "55.2.0" arrow-buffer = "55.2.0" diff --git a/plugins/engine-datafusion/jni/src/lib.rs b/plugins/engine-datafusion/jni/src/lib.rs index 1806a18963a37..64904516cbb95 100644 --- a/plugins/engine-datafusion/jni/src/lib.rs +++ b/plugins/engine-datafusion/jni/src/lib.rs @@ -6,7 +6,7 @@ * compatible open source license. */ -use jni::objects::{JByteArray, JClass}; +use jni::objects::{JByteArray, JClass, JObject}; use jni::sys::{jbyteArray, jlong, jstring}; use jni::JNIEnv; use std::sync::Arc; @@ -17,12 +17,12 @@ use datafusion::execution::context::SessionContext; use datafusion::DATAFUSION_VERSION; use datafusion::datasource::file_format::csv::CsvFormat; -use datafusion::datasource::file_format::parquet::ParquetFormat; -use datafusion::execution::cache::cache_manager::{CacheManager, CacheManagerConfig, FileStatisticsCache}; +use datafusion::execution::cache::cache_manager::{CacheManager, CacheManagerConfig, FileMetadataCache, FileStatisticsCache}; +use datafusion::execution::cache::cache_unit::DefaultFilesMetadataCache; use datafusion::execution::disk_manager::DiskManagerConfig; use datafusion::execution::runtime_env::{RuntimeEnv, RuntimeEnvBuilder}; use datafusion::prelude::SessionConfig; -use crate::util::{create_object_meta_from_filenames, parse_string_arr}; +use crate::util::{create_object_meta_from_filenames, create_object_meta_from_file, parse_string_arr, construct_file_metadata}; use datafusion::datasource::listing::{ListingOptions, ListingTable, ListingTableConfig, ListingTableUrl}; use datafusion::execution::cache::cache_unit::DefaultListFilesCache; use datafusion::execution::cache::CacheAccessor; @@ -33,6 +33,11 @@ use jni::objects::{JObjectArray, JString}; use prost::Message; use tokio::runtime::Runtime; use object_store::ObjectMeta; +use chrono::Utc; +use std::collections::HashMap; +use std::sync::Mutex; +use std::ops::Deref; + /// Create a new DataFusion session context #[no_mangle] @@ -103,7 +108,7 @@ pub extern "system" fn Java_org_opensearch_datafusion_DataFusionQueryJNI_createG } #[no_mangle] -pub extern "system" fn Java_org_opensearch_datafusion_DataFusionQueryJNI_createSessionContext( +pub extern "system" fn Java_org_opensearch_datafusion_DataFusionQueryJNI_createSessionContextv1( _env: JNIEnv, _class: JClass, runtime_id: jlong, @@ -261,7 +266,7 @@ pub extern "system" fn Java_org_opensearch_datafusion_DataFusionQueryJNI_nativeE }, Err(e) => { println!("SUBSTRAIT Rust: Failed to convert Substrait plan: {}", e); - return; + return 0; } }; @@ -269,22 +274,20 @@ pub extern "system" fn Java_org_opensearch_datafusion_DataFusionQueryJNI_nativeE .await.expect("Failed to run Logical Plan"); // TODO : check if this works - return match dataframe.execute_stream() { + match dataframe.execute_stream().await { Ok(stream) => { let boxed_stream = Box::new(stream); let stream_ptr = Box::into_raw(boxed_stream); stream_ptr as jlong }, Err(e) => { + println!("SUBSTRAIT Rust: Failed to execute stream: {}", e); 0 } } }) - // Create DataFrame from the converted logical plan - - } // If we need to create session context separately @@ -353,4 +356,170 @@ pub extern "system" fn Java_org_opensearch_datafusion_DataFusionQueryJNI_nativeC +#[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 RuntimeEnv using the configured CacheManagerConfig +#[no_mangle] +pub extern "system" fn Java_org_opensearch_datafusion_DataFusionQueryJNI_createGlobalRuntimeWithConfig( + _env: JNIEnv, + _class: JClass, + config_ptr: jlong, +) -> jlong { + // Take ownership of the CacheManagerConfig + let cache_manager_config = unsafe { Box::from_raw(config_ptr as *mut CacheManagerConfig) }; + + // Create RuntimeEnv with the configured cache manager + let runtime_env = RuntimeEnvBuilder::default() + .with_cache_manager(*cache_manager_config) + .build() + .unwrap(); + + Box::into_raw(Box::new(runtime_env)) 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 the cache + let cache = Arc::new(DefaultFilesMetadataCache::new(size_limit.try_into().unwrap())); + + // Update the CacheManagerConfig at the same memory location + if config_ptr != 0 { + let cache_manager_config = unsafe { &mut *(config_ptr as *mut CacheManagerConfig) }; + // This replaces the contents at the same pointer location + *cache_manager_config = cache_manager_config.clone() + .with_file_metadata_cache(Some(cache.clone() as Arc)); + } + + // Return the cache pointer + Box::into_raw(Box::new(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, +) -> i32 { + let file_path: String = match env.get_string(&file_path) { + Ok(s) => s.into(), + Err(_) => return -1, + }; + let cache = unsafe { &mut *(cache_ptr as *mut Arc) }; + let data_format = if file_path.to_lowercase().ends_with(".parquet") { + "parquet" + } else { + return 0; // 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") + }); + + cache.put(&object_meta, metadata); + + println!("Cached metadata for: {}", file_path); + 1 +} + +#[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 { &mut *(cache_ptr as *mut Arc) }; + let object_meta = create_object_meta_from_file(&file_path); + + // Try to get mutable access if there's only one reference + if let Some(cache_mut) = Arc::get_mut(cache) { + cache_mut.remove(&object_meta); + println!("Cache removed for: {}", file_path); + true + } else { + // If there are multiple references, we can't remove the item + // This is a limitation of the current cache design + println!("Cannot remove from cache (multiple references exist): {}", 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, +) -> jlong { + let file_path: String = match env.get_string(&file_path) { + Ok(s) => s.into(), + Err(_) => return 0, + }; + + 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()); + Box::into_raw(Box::new(metadata)) as jlong + }, + None => { + println!("No metadata found for: {}", file_path); + 0 + }, + } +} + +#[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.contains_key(&object_meta) +} +#[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.memory_used() +} diff --git a/plugins/engine-datafusion/jni/src/util.rs b/plugins/engine-datafusion/jni/src/util.rs index a584df7489cdd..23ced069ba46d 100644 --- a/plugins/engine-datafusion/jni/src/util.rs +++ b/plugins/engine-datafusion/jni/src/util.rs @@ -5,13 +5,18 @@ use anyhow::Result; use chrono::{DateTime, Utc}; use datafusion::arrow::array::RecordBatch; +use datafusion::datasource::physical_plan::parquet::metadata::DFParquetMetadata; +use datafusion::datasource::physical_plan::parquet::CachedParquetMetaData; +use datafusion::execution::cache::cache_manager::FileMetadata; 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 object_store::{local::LocalFileSystem}; /// 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>) { @@ -90,7 +95,6 @@ pub fn set_object_result_error(env: &mut JNIEnv, callback: JObject, er .expect("Failed to call object result callback with error"); } - /// Parse a string map from JNI arrays pub fn parse_string_map( env: &mut JNIEnv, @@ -159,20 +163,44 @@ 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| { - let filename = filename.as_str(); let full_path = format!("{}/{}", base_path.trim_end_matches('/'), filename); - 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()); - - ObjectMeta { - location: ObjectPath::from(filename), - last_modified: modified, - size: file_size, - e_tag: None, - version: None, - } + create_object_meta_from_file(&full_path) }).collect() -} \ No newline at end of file +} + +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(file_path), + last_modified: modified, + size: file_size, + e_tag: None, + version: None, + } +} + +// Utility method to construct file metadata using DataFusion's DFParquetMetadata +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 955eca8c97362..76efeac893e81 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,14 @@ package org.opensearch.datafusion; +import java.util.stream.Collectors; +import java.util.stream.Stream; 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,9 +27,9 @@ 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.search.ContextEngineSearcher; import org.opensearch.index.engine.SearchExecEngine; import org.opensearch.index.engine.exec.FileMetadata; import org.opensearch.plugins.ActionPlugin; @@ -40,7 +43,6 @@ import org.opensearch.transport.client.Client; import org.opensearch.vectorized.execution.search.DataFormat; import org.opensearch.vectorized.execution.search.spi.DataSourceCodec; -import org.opensearch.vectorized.execution.search.spi.RecordBatchStream; import org.opensearch.watcher.ResourceWatcherService; import java.io.IOException; @@ -50,6 +52,8 @@ import java.util.Map; import java.util.function.Supplier; +import static org.opensearch.datafusion.search.cache.CacheType.METADATA; + /** * Main plugin class for OpenSearch DataFusion integration. * @@ -102,7 +106,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); @@ -125,7 +129,7 @@ public List getSupportedFormats() { public SearchExecEngine createEngine(DataFormat dataFormat,Collection formatCatalogSnapshot) throws IOException { - return new DatafusionEngine(dataFormat, formatCatalogSnapshot); + return new DatafusionEngine(dataFormat, dataFusionService.getCacheManager(),formatCatalogSnapshot); } /** @@ -166,4 +170,11 @@ 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 a64ca2da182d6..287e1f076d431 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 @@ -99,6 +99,8 @@ private static synchronized void loadNativeLibrary() { */ public static native long createSessionContext(long runtimeId); + public static native long createSessionContextv1(long runtimeId); + /** * Close and cleanup a DataFusion context * @param contextId the context ID to close @@ -146,4 +148,62 @@ private static synchronized void loadNativeLibrary() { * @param streamPtr the stream pointer to close */ public static native void closeStream(long streamPtr); + + public static native long initCacheManagerConfig(); + + public static native long createMetadataCache(long cacheConfigPtr, long sizeLimit); + + public static native long createGlobalRuntimeWithConfig(long cacheConfigPtr); + + public static native long create(java.util.Map properties, String cacheType); + +// // METADATA cache specific methods +// +// /** +// * Create metadata cache +// * @param properties cache configuration properties +// * @return cache pointer +// */ +// public static native long createMetadataCache(java.util.Map properties); + + /** + * 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); + + /** + * 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); + } 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 7b03b584d3444..87eb208989df9 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; @@ -32,16 +34,25 @@ public class DataFusionService extends AbstractLifecycleComponent { private final DataSourceRegistry dataSourceRegistry; private final GlobalRuntimeEnv globalRuntimeEnv; + private final CacheManager cacheManager; + + /** + * Get the cache manager instance + * @return CacheManager instance + */ + public CacheManager getCacheManager() { + return cacheManager; + } /** * Creates a new DataFusion service instance. */ - public DataFusionService(Map dataSourceCodecs) { + public DataFusionService(Map dataSourceCodecs, ClusterSettings settings) { this.dataSourceRegistry = new DataSourceRegistry(dataSourceCodecs); - // to verify jni String version = DataFusionQueryJNI.getVersionInfo(); - this.globalRuntimeEnv = new GlobalRuntimeEnv(); + this.cacheManager = CacheManager.fromConfig(settings); + this.globalRuntimeEnv = new GlobalRuntimeEnv(cacheManager.getCacheManagerPtr()); } @Override 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 20e32899c1f8a..bacb2fadb10ea 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.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.index.engine.CatalogSnapshotAwareRefreshListener; import org.opensearch.index.engine.Engine; import org.opensearch.index.engine.EngineException; @@ -40,10 +41,17 @@ public class DatafusionEngine extends SearchExecEngine formatCatalogSnapshot) throws IOException { + public DatafusionEngine(DataFormat dataFormat, CacheManager cacheManager, Collection formatCatalogSnapshot) throws IOException { this.dataFormat = dataFormat; this.datafusionReaderManager = new DatafusionReaderManager("TODO://FigureOutPath", formatCatalogSnapshot); + this.cacheManager = cacheManager; + // Set up callbacks when creating the reader manager + datafusionReaderManager.setOnFilesAdded(files -> { + // Handle new files added during refresh + cacheManager.addToCache("TODO://FigureOutPath",files); + }); } @Override @@ -85,6 +93,7 @@ protected DatafusionSearcher acquireSearcherInternal(String source) { @Override protected void doClose() { try { + cacheManager.removeFilesByDirectory("TODO://FigureOutPath"); reader.decRef(); } catch (IOException e) { throw new UncheckedIOException(e); 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 1867028fcb945..d09cf46d9d137 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,8 +8,11 @@ package org.opensearch.datafusion.core; +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.createGlobalRuntimeWithConfig; /** * Global runtime environment for DataFusion operations. @@ -26,6 +29,10 @@ public GlobalRuntimeEnv() { this.ptr = createGlobalRuntime(); } + public GlobalRuntimeEnv(long cacheManagerPtr) { + this.ptr = createGlobalRuntimeWithConfig(cacheManagerPtr); + } + /** * Gets the native pointer to the runtime environment. * @return the native pointer @@ -37,5 +44,6 @@ public long getPointer() { @Override public void close() { closeGlobalRuntime(this.ptr); + } } 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 a48e697d6fd16..a2a40d4c9b16e 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,24 +8,36 @@ package org.opensearch.datafusion.search; +import org.apache.lucene.index.IndexReader; import org.apache.lucene.search.ReferenceManager; +import org.opensearch.datafusion.search.cache.CacheManager; import org.opensearch.index.engine.CatalogSnapshotAwareRefreshListener; import org.opensearch.index.engine.EngineReaderManager; import org.opensearch.index.engine.exec.FileMetadata; import org.opensearch.index.engine.exec.coord.CatalogSnapshot; +import org.opensearch.index.shard.IndexEventListener; +import org.opensearch.index.shard.IndexShard; +import org.opensearch.core.index.shard.ShardId; +import org.opensearch.common.settings.Settings; import java.io.IOException; import java.nio.file.Path; import java.util.Collection; import java.util.List; +import java.util.Set; +import java.util.HashSet; import java.util.concurrent.CopyOnWriteArrayList; import java.util.concurrent.locks.Lock; import java.util.concurrent.locks.ReentrantLock; +import java.util.function.Consumer; public class DatafusionReaderManager implements EngineReaderManager, CatalogSnapshotAwareRefreshListener { private DatafusionReader current; private String path; private String dataFormat; + private Consumer> onFilesAdded; + private Consumer> onFilesRemoved; + // private final Lock refreshLock = new ReentrantLock(); // private final List refreshListeners = new CopyOnWriteArrayList(); @@ -35,6 +47,13 @@ public DatafusionReaderManager(String path, Collection files) thro 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) { @@ -50,7 +69,6 @@ public void release(DatafusionReader reference) throws IOException { reference.decRef(); } - @Override public void beforeRefresh() throws IOException { // no op @@ -60,9 +78,37 @@ public void beforeRefresh() throws IOException { public void afterRefresh(boolean didRefresh, CatalogSnapshot catalogSnapshot) throws IOException { if (didRefresh && catalogSnapshot != null) { DatafusionReader old = this.current; + Collection newFiles = catalogSnapshot.getSearchableFiles(dataFormat); + + processFileChanges(old.files, newFiles); + release(old); - this.current = new DatafusionReader(this.path, catalogSnapshot.getSearchableFiles(dataFormat)); + this.current = new DatafusionReader(this.path, newFiles); 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); + + Set filesToRemove = new HashSet<>(oldFilePaths); + filesToRemove.removeAll(newFilePaths); + + if (!filesToAdd.isEmpty() && onFilesAdded != null) { + onFilesAdded.accept(List.copyOf(filesToAdd)); + } + } + + private Set extractFilePaths(Collection files) { + Set paths = new HashSet<>(); + for (FileMetadata file : files) { + paths.add(this.path.concat(file.fileName())); + } + return paths; + } + } diff --git a/plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/search/cache/CacheAccessor.java b/plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/search/cache/CacheAccessor.java new file mode 100644 index 0000000000000..6cf3e698c5809 --- /dev/null +++ b/plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/search/cache/CacheAccessor.java @@ -0,0 +1,50 @@ +package org.opensearch.datafusion.search.cache; + +import java.util.Map; +import org.opensearch.common.settings.ClusterSettings; +import org.opensearch.core.common.unit.ByteSizeValue; + +public abstract class CacheAccessor { + + protected CacheType name; + protected Map 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 void setSizeLimit(int limit); + public abstract int 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..135f814071364 --- /dev/null +++ b/plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/search/cache/CacheManager.java @@ -0,0 +1,127 @@ +/* + * 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.io.IOException; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.stream.Collectors; +import org.opensearch.common.lucene.index.OpenSearchDirectoryReader; +import org.opensearch.common.settings.ClusterSettings; +import org.opensearch.common.settings.Setting; +import org.opensearch.index.engine.CatalogSnapshotAwareRefreshListener; +import org.opensearch.index.engine.exec.coord.CatalogSnapshot; + +import static org.opensearch.datafusion.DataFusionQueryJNI.initCacheManagerConfig; + + +public class CacheManager implements AutoCloseable { + + 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(); + // this.totalSizeLimit = 500000; + } + + // 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(String path, List files) { + boolean allSuccessful = true; + for (CacheAccessor cache : getAllCaches()) { + for (String filename : files) { + allSuccessful &= cache.remove(path.concat(filename)); + } + } + + + return allSuccessful; + } + + public boolean removeFilesByDirectory(String path) { + boolean allSuccessful = true; + for (CacheAccessor cache : getAllCaches()) { + allSuccessful &= cache.remove(path); + } + return allSuccessful; + } + + public boolean addToCache(String path, List files) { + boolean allSuccessful = true; + for (CacheAccessor cache : getAllCaches()) { + for (String filename : files) { + allSuccessful &= cache.put(path.concat("/" + 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; + } + + @Override + public void close() throws Exception { + + } +} 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..f3b9549c0f86c --- /dev/null +++ b/plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/search/cache/CacheSettings.java @@ -0,0 +1,48 @@ +/* + * 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 METADATA_CACHE_SIZE_LIMIT_KEY = "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 METADATA_CACHE_EVICTION_TYPE = new Setting( + "metadata.cache.eviction.type", + "LFU", + Function.identity(), + Setting.Property.NodeScope, + Setting.Property.Dynamic + ); + + + public static final String METADATA_CACHE_ENABLED_KEY = "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 List> CACHE_SETTINGS = Arrays.asList( + METADATA_CACHE_SIZE_LIMIT, + METADATA_CACHE_EVICTION_TYPE + ); + + public static final List> CACHE_ENABLED = Arrays.asList( + METADATA_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..4afbb133d1c43 --- /dev/null +++ b/plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/search/cache/CacheType.java @@ -0,0 +1,32 @@ +package org.opensearch.datafusion.search.cache; + +import org.opensearch.common.settings.*; +import java.util.List; +import java.util.function.Supplier; +import static org.opensearch.datafusion.search.cache.CacheSettings.*; + +public enum CacheType { + METADATA(MetadataCacheAccessor::new, METADATA_CACHE_ENABLED); + // STATS(StatsCacheAccessor::new); + + 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..bea1c84308763 --- /dev/null +++ b/plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/search/cache/MetadataCacheAccessor.java @@ -0,0 +1,91 @@ +package org.opensearch.datafusion.search.cache; + +import java.util.HashMap; +import java.util.Map; +import org.opensearch.common.settings.ClusterSettings; + +import static org.opensearch.datafusion.DataFusionQueryJNI.createMetadataCache; +import static org.opensearch.datafusion.DataFusionQueryJNI.metadataCacheContainsFile; +import static org.opensearch.datafusion.DataFusionQueryJNI.metadataCacheGet; +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.search.cache.CacheSettings.*; + +public class MetadataCacheAccessor extends CacheAccessor { + 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) { + return metadataCachePut(this.getPointer(), filePath); + } + + @Override + public Object get(String filePath) { + return metadataCacheGet(this.pointer,filePath); + } + + @Override + public boolean remove(String filePath) { + return metadataCacheRemove(this.pointer, filePath); + } + + @Override + public void evict() { + return; + } + + @Override + public void clear() { + return; + } + + @Override + public long getMemoryConsumed() { + return metadataCacheGetSize(this.pointer); + } + + @Override + public boolean containsFile(String filePath) { + return metadataCacheContainsFile(this.pointer, filePath); + } + + @Override + public void setSizeLimit(int limit) { + return; + } + + @Override + public int getEntries() { + return 1 ; + } + +} 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..dbf1899452e98 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,31 @@ package org.opensearch.datafusion; +import java.util.HashSet; +import java.util.Set; +import org.opensearch.cluster.routing.allocation.DiskThresholdSettings; +import org.opensearch.common.settings.ClusterSettings; +import org.opensearch.common.settings.Setting; +import org.opensearch.common.settings.Settings; +import org.opensearch.datafusion.search.cache.CacheAccessor; +import org.opensearch.datafusion.search.cache.CacheManager; +import org.opensearch.datafusion.search.cache.CacheType; +import org.opensearch.index.IndexModule; +import org.opensearch.index.store.remote.filecache.FileCacheSettings; import org.opensearch.test.OpenSearchTestCase; import java.util.Collections; import java.util.List; +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_ENABLED_KEY; +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; +import static org.opensearch.index.IndexModule.INDEX_STORE_LOCALITY_SETTING; +import static org.opensearch.index.IndexModule.IS_WARM_INDEX_SETTING; + /** * Unit tests for DataFusionService * @@ -27,7 +47,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(); } @@ -48,6 +75,32 @@ public void testCreateAndCloseContext() { service.getVersion(); } + public void testCacheOperations() { + CacheAccessor metadata = service.getCacheManager().getCacheAccessor(CacheType.METADATA); + + CacheManager cacheManager= service.getCacheManager(); + String dirPath = "/Users/abhital/dev/src/forkedrepo/OpenSearch/plugins/engine-datafusion/src"; + String fileName = "hits9_v1.parquet"; + + String filePath = dirPath + "/" + fileName; + // add File using CacheManager + cacheManager.addToCache(dirPath,List.of(fileName)); + + // Get file using individual Cache Accessor Methods -> Prints Cache content size + assertTrue((Boolean) metadata.get(filePath)); + + logger.info("Memory Consumed by MetadataCache : {}",metadata.getMemoryConsumed()); + logger.info("Memory Consumed by CacheManager : {}",cacheManager.getTotalUsedBytes()); + + logger.info("Total Configured Size Limit for MetadataCache : {}",metadata.getConfiguredSizeLimit()); + logger.info("Total Configured Size Limit for CacheManager : {}",cacheManager.getTotalSizeLimit()); + + boolean removed = cacheManager.removeFiles(dirPath,List.of(fileName)); + logger.info("Is file removed: {}. Contains File Check: {} Ideally remove will not work as we have multiple references",removed, metadata.containsFile(filePath)); + logger.info("Memory Consumed by MetadataCache : {}",metadata.getMemoryConsumed()); + logger.info("Memory Consumed by CacheManager : {}",cacheManager.getTotalUsedBytes()); + } + public void testCodecDiscovery() { // Test that the CSV codec can be discovered via SPI // TODO : test with dummy plugin and dummy codec diff --git a/test/framework/src/main/java/org/opensearch/node/MockNode.java b/test/framework/src/main/java/org/opensearch/node/MockNode.java index b83e77891b92f..32aaecce9f14f 100644 --- a/test/framework/src/main/java/org/opensearch/node/MockNode.java +++ b/test/framework/src/main/java/org/opensearch/node/MockNode.java @@ -50,6 +50,7 @@ import org.opensearch.env.Environment; import org.opensearch.http.HttpServerTransport; import org.opensearch.indices.IndicesService; +import org.opensearch.plugins.DataSourcePlugin; import org.opensearch.plugins.Plugin; import org.opensearch.plugins.PluginInfo; import org.opensearch.plugins.SearchPlugin; @@ -174,7 +175,8 @@ protected SearchService newSearchService( Executor indexSearcherExecutor, TaskResourceTrackingService taskResourceTrackingService, Collection concurrentSearchDeciderFactories, - List pluginProfilers + List pluginProfilers, + List dataSourcePluginList ) { if (getPluginsService().filterPlugins(MockSearchService.TestPlugin.class).isEmpty()) { return super.newSearchService( @@ -190,7 +192,8 @@ protected SearchService newSearchService( indexSearcherExecutor, taskResourceTrackingService, concurrentSearchDeciderFactories, - pluginProfilers + pluginProfilers, + null ); } return new MockSearchService( diff --git a/test/framework/src/main/java/org/opensearch/search/MockSearchService.java b/test/framework/src/main/java/org/opensearch/search/MockSearchService.java index e3bc166e56d6b..f35e3a7e5925b 100644 --- a/test/framework/src/main/java/org/opensearch/search/MockSearchService.java +++ b/test/framework/src/main/java/org/opensearch/search/MockSearchService.java @@ -114,7 +114,8 @@ public MockSearchService( indexSearcherExecutor, taskResourceTrackingService, Collections.emptyList(), - Collections.emptyList() + Collections.emptyList(), + null ); } From cbba1339c425b4f73403263e2128260d27b8554a Mon Sep 17 00:00:00 2001 From: Abhita Lakkabathini Date: Fri, 26 Sep 2025 11:13:45 +0530 Subject: [PATCH 2/2] path fix --- .../opensearch/datafusion/search/cache/CacheManager.java | 2 +- .../opensearch/datafusion/TestDataFusionServiceTests.java | 7 ------- 2 files changed, 1 insertion(+), 8 deletions(-) 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 index 135f814071364..2d723d9c470c3 100644 --- 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 @@ -77,7 +77,7 @@ public boolean removeFiles(String path, List files) { boolean allSuccessful = true; for (CacheAccessor cache : getAllCaches()) { for (String filename : files) { - allSuccessful &= cache.remove(path.concat(filename)); + allSuccessful &= cache.remove(path.concat("/"+filename)); } } 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 dbf1899452e98..17a28e6865bb0 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 @@ -10,15 +10,12 @@ import java.util.HashSet; import java.util.Set; -import org.opensearch.cluster.routing.allocation.DiskThresholdSettings; import org.opensearch.common.settings.ClusterSettings; import org.opensearch.common.settings.Setting; import org.opensearch.common.settings.Settings; import org.opensearch.datafusion.search.cache.CacheAccessor; import org.opensearch.datafusion.search.cache.CacheManager; import org.opensearch.datafusion.search.cache.CacheType; -import org.opensearch.index.IndexModule; -import org.opensearch.index.store.remote.filecache.FileCacheSettings; import org.opensearch.test.OpenSearchTestCase; import java.util.Collections; @@ -26,12 +23,8 @@ 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_ENABLED_KEY; 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; -import static org.opensearch.index.IndexModule.INDEX_STORE_LOCALITY_SETTING; -import static org.opensearch.index.IndexModule.IS_WARM_INDEX_SETTING; /** * Unit tests for DataFusionService