Enable Metadata Cache Complete Warm-up (#20117)

* Update MetadataCache Warm-up to include page-indexes

* disable dynamic metadtaa cache limit updates

Author: abhita

plugins/engine-datafusion/jni/src/custom_cache_manager.rs

@@ -2,9 +2,10 @@ use std::sync::{Arc, Mutex};
 use datafusion::execution::cache::cache_manager::{FileMetadataCache, CacheManagerConfig};
 use datafusion::execution::cache::cache_unit::{DefaultFileStatisticsCache, DefaultFilesMetadataCache, DefaultListFilesCache};
 use datafusion::execution::cache::CacheAccessor;
+use datafusion::datasource::physical_plan::parquet::metadata::DFParquetMetadata;
 use tokio::runtime::Runtime;
 use crate::cache::MutexFileMetadataCache;
-use crate::util::{create_object_meta_from_file, construct_file_metadata};
+use crate::util::{create_object_meta_from_file};
 
 /// Custom CacheManager that holds cache references directly
 pub struct CustomCacheManager {
@@ -39,7 +40,8 @@ impl CustomCacheManager {
         let file_static_cache = Arc::new(DefaultFileStatisticsCache::default());
         // Add file metadata cache if available
         if let Some(cache) = self.get_file_metadata_cache_for_datafusion() {
-            config = config.with_file_metadata_cache(Some(cache));
+            config = config.with_file_metadata_cache(Some(cache.clone()))
+                .with_metadata_cache_limit(cache.cache_limit());
         }
         config = config.with_files_statistics_cache(Some(file_static_cache.clone()));
         // Future: Add stats cache when implemented
@@ -256,31 +258,40 @@ impl CustomCacheManager {
 
         let store = Arc::new(object_store::local::LocalFileSystem::new());
 
-        //TODO: Use TokioRuntimePtr to block on the async operation
-        let metadata = Runtime::new()
+        // Get cache reference for DataFusion metadata loading
+        let cache_ref = self.file_metadata_cache.as_ref()
+            .ok_or_else(|| "No file metadata cache configured".to_string())?;
+
+        let metadata_cache = cache_ref.clone() as Arc<dyn FileMetadataCache>;
+
+        // Use DataFusion's metadata loading by passing reference to file_metadata_cache to get complete metadata
+        // IMPORTANT: When a cache is provided to DFParquetMetadata, fetch_metadata() will:
+        // 1. Enable page index loading (with_page_indexes(true))
+        // 2. Load the complete metadata including column and offset indexes
+        // 3. Automatically put the metadata into the cache (lines 155-160 in datafusion's metadata.rs)
+        // This ensures we cache exactly what DataFusion would cache during query execution
+        let _parquet_metadata = Runtime::new()
             .map_err(|e| format!("Failed to create Tokio Runtime: {}", e))?
             .block_on(async {
-                construct_file_metadata(store.as_ref(), object_meta, data_format)
-                    .await
-                    .map_err(|e| format!("Failed to construct file metadata: {}", e))
-            })?;
+                let df_metadata = DFParquetMetadata::new(store.as_ref(), object_meta)
+                    .with_file_metadata_cache(Some(metadata_cache));
 
-        // Get the cache directly from our stored reference
-        let cache = self.file_metadata_cache.as_ref()
-            .ok_or_else(|| "No file metadata cache configured".to_string())?;
-
-        match cache.inner.lock() {
-            Ok(mut cache_guard) => {
-                cache_guard.put(object_meta, metadata);
+                // fetch_metadata() performs the cache put operation internally
+                df_metadata.fetch_metadata().await
+                    .map_err(|e| format!("Failed to fetch metadata: {}", e))
+            })?;
 
+        // Verify the metadata was cached properly
+        match cache_ref.inner.lock() {
+            Ok(cache_guard) => {
                 if cache_guard.contains_key(object_meta) {
                     Ok(true)
                 } else {
-                    println!("Failed to cache metadata for: {}", file_path);
+                    println!("[CACHE ERROR] Failed to cache metadata for: {}", file_path);
                     Ok(false)
                 }
             }
-            Err(e) => Err(format!("Cache put failed: {}", e))
+            Err(e) => Err(format!("Failed to verify cache: {}", e))
         }
     }
 }

plugins/engine-datafusion/jni/src/lib.rs

@@ -235,7 +235,6 @@ pub extern "system" fn Java_org_opensearch_datafusion_jni_NativeBridge_createGlo
 
         let runtime_env = RuntimeEnvBuilder::new().with_cache_manager(cache_manager_config)
             .with_memory_pool(memory_pool.clone())
-            .with_metadata_cache_limit(250 * 1024 * 1024) // 250 MB
             .build().unwrap();
 
         let runtime = DataFusionRuntime {

plugins/engine-datafusion/jni/src/query_executor.rs

@@ -65,15 +65,16 @@ pub async fn execute_query_with_cross_rt_stream(
 
     let runtimeEnv = &runtime.runtime_env;
 
+    let file_metadata_cache = runtime.runtime_env.cache_manager.get_file_metadata_cache();
+
     let runtime_env = match RuntimeEnvBuilder::from_runtime_env(runtimeEnv)
         .with_cache_manager(
             CacheManagerConfig::default()
                 .with_list_files_cache(Some(list_file_cache.clone()))
-                .with_file_metadata_cache(Some(runtimeEnv.cache_manager.get_file_metadata_cache()))
+                .with_file_metadata_cache(Some(file_metadata_cache.clone()))
+                .with_metadata_cache_limit(file_metadata_cache.cache_limit())
                 .with_files_statistics_cache(runtimeEnv.cache_manager.get_file_statistic_cache()),
-        )
-        .with_metadata_cache_limit(250 * 1024 * 1024) // 250 MB
-        .build() {
+        ).build() {
         Ok(env) => env,
         Err(e) => {
             error!("Failed to build runtime env: {}", e);
@@ -211,13 +212,15 @@ pub async fn execute_fetch_phase(
     let list_file_cache = Arc::new(DefaultListFilesCache::default());
     list_file_cache.put(table_path.prefix(), object_meta);
 
+    let file_metadata_cache = runtime.runtime_env.cache_manager.get_file_metadata_cache();
+
     let runtime_env = RuntimeEnvBuilder::new()
         .with_cache_manager(
             CacheManagerConfig::default().with_list_files_cache(Some(list_file_cache))
-                .with_file_metadata_cache(Some(runtime.runtime_env.cache_manager.get_file_metadata_cache()))
-                .with_files_statistics_cache(runtime.runtime_env.cache_manager.get_file_statistic_cache()),
+                .with_file_metadata_cache(Some(file_metadata_cache.clone()))
+                .with_files_statistics_cache(runtime.runtime_env.cache_manager.get_file_statistic_cache())
+                .with_metadata_cache_limit(file_metadata_cache.cache_limit()),
         )
-        .with_metadata_cache_limit(250 * 1024 * 1024) // 250 MB
         .build()?;
     let ctx = SessionContext::new_with_config_rt(SessionConfig::new(), Arc::new(runtime_env));
 

plugins/engine-datafusion/jni/src/util.rs

@@ -209,25 +209,6 @@ pub fn create_object_meta_from_file(file_path: &str) -> Result<Vec<ObjectMeta>,
     Ok(vec![object_meta])
 }
 
-pub async fn construct_file_metadata(
-    store: &dyn ObjectStore,
-    object_meta: &ObjectMeta,
-    data_format: &str,
-) -> Result<Arc<dyn FileMetadata>, Box<dyn std::error::Error>> {
-    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())
-    }
-}
 /// Set success result by calling an ActionListener
 pub fn set_action_listener_ok(env: &mut JNIEnv, listener: JObject, value: jlong) {
     let long_obj = env.new_object("java/lang/Long", "(J)V", &[value.into()])

plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/search/cache/CacheUtils.java

@@ -105,7 +105,7 @@ public static long createCacheConfig(ClusterSettings clusterSettings) {
                     type.getEvictionType(clusterSettings));
 
                 NativeBridge.createCache(cacheManagerPtr, type.cacheTypeName, type.getSizeLimit(clusterSettings).getBytes(), type.getEvictionType(clusterSettings));
-                clusterSettings.addSettingsUpdateConsumer(type.sizeLimitSetting,(v) -> NativeBridge.cacheManagerUpdateSizeLimitForCacheType(cacheManagerPtr, CacheType.METADATA.getCacheTypeName(),v.getBytes()));
+           //     clusterSettings.addSettingsUpdateConsumer(type.sizeLimitSetting,(v) -> NativeBridge.cacheManagerUpdateSizeLimitForCacheType(cacheManagerPtr, CacheType.METADATA.getCacheTypeName(),v.getBytes()));
             } else {
                 logger.debug("Cache type {} is disabled", type.getCacheTypeName());
             }