Integrate Liquid Cache with DataFusion 52.1 for byte-level Parquet caching with comprehensive logging

Signed-off-by: Tanvir Alam <tanvralm@amazon.com>

Author: Tanvir Alam

plugins/engine-datafusion/Cargo.toml

@@ -6,22 +6,22 @@ members = [
 
 [workspace.dependencies]
 # DataFusion dependencies
-datafusion = "51.0.0"
-datafusion-expr = "51.0.0"
-datafusion-datasource = "51.0.0"
-arrow-json = "57.1.0"
-arrow = { version = "57.1.0", features = ["ffi", "ipc_compression"] }
-arrow-array = "57.1.0"
-arrow-schema = "57.1.0"
-arrow-buffer = "57.1.0"
+datafusion = "52.1.0"
+datafusion-expr = "52.1.0"
+datafusion-datasource = "52.1.0"
+arrow-json = "57.3.0"
+arrow = { version = "57.3.0", features = ["ffi", "ipc_compression"] }
+arrow-array = "57.3.0"
+arrow-schema = "57.3.0"
+arrow-buffer = "57.3.0"
 downcast-rs = "1.2"
 
 
 # JNI dependencies
 jni = "0.21"
 
 # Substrait support
-datafusion-substrait = "51.0.0"
+datafusion-substrait = "52.1.0"
 prost = "0.14"
 
 
@@ -42,10 +42,10 @@ thiserror = "1.0"
 # Logging
 log = "0.4"
 # Parquet support
-parquet = "57.1.0"
+parquet = "57.3.0"
 
 # Object store for file access
-object_store = "=0.12.4"
+object_store = "0.12.5"
 url = "2.0"
 
 # Substrait support

plugins/engine-datafusion/jni/Cargo.toml

@@ -61,6 +61,9 @@ num_cpus = "1.16"
 object_store = { workspace = true }
 url = { workspace = true }
 
+# Liquid Cache for byte-level caching
+liquid-cache-datafusion-local = "0.1.12"
+
 # Substrait support
 substrait = { workspace = true }
 

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

@@ -14,7 +14,7 @@ use datafusion::{
     common::tree_node::{Transformed, TreeNode, TreeNodeRecursion},
     config::ConfigOptions,
     datasource::{
-        physical_plan::{FileScanConfig, FileScanConfigBuilder},
+        physical_plan::{FileScanConfig, FileScanConfigBuilder, ParquetSource},
         source::DataSourceExec,
     },
     error::DataFusionError,
@@ -41,7 +41,7 @@ impl AbsoluteRowIdOptimizer {
         datasource_exec_schema: SchemaRef,
     ) -> (SchemaRef, Vec<usize>) {
         // Clone projection and add new field index
-        let mut projected_schema = datasource.projected_schema().clone();
+        let projected_schema = datasource.projected_schema().expect("Failed to get projected schema");
         let file_source_schema = datasource.file_schema().clone();
 
         let mut new_projections = vec![];
@@ -109,9 +109,13 @@ impl AbsoluteRowIdOptimizer {
     ) -> Result<ProjectionExec, DataFusionError> {
         let (new_schema, new_projections) =
             self.build_updated_file_source_schema(datasource, data_source_exec_schema.clone());
+        
+        // Create a new file source with the updated schema
+        let new_file_source = ParquetSource::new(new_schema.clone());
+        
         let file_scan_config = FileScanConfigBuilder::from(datasource.clone())
-            .with_source(datasource.file_source.with_schema(TableSchema::from_file_schema(new_schema.clone())))
-            .with_projection_indices(Some(new_projections))
+            .with_source(Arc::new(new_file_source))
+            .with_projection_indices(Some(new_projections))?
             .build();
 
         let new_datasource = DataSourceExec::from_data_source(file_scan_config);

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

@@ -110,7 +110,7 @@ impl CacheAccessor<ObjectMeta, Arc<dyn datafusion::execution::cache::cache_manag
         }
     }
 
-    fn remove(&mut self, k: &ObjectMeta) -> Option<Arc<dyn datafusion::execution::cache::cache_manager::FileMetadata>> {
+    fn remove(&self, k: &ObjectMeta) -> Option<Arc<dyn datafusion::execution::cache::cache_manager::FileMetadata>> {
         match self.inner.lock() {
             Ok(mut cache) => cache.remove(k),
             Err(e) => {

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

@@ -61,14 +61,9 @@ impl CustomCacheManager {
                 .with_metadata_cache_limit(cache.cache_limit());
         }
 
-        // Add statistics cache if available - use CustomStatisticsCache directly
-        if let Some(stats_cache) = &self.statistics_cache {
-            config = config.with_files_statistics_cache(Some(stats_cache.clone() as FileStatisticsCache));
-        } else {
-            // Default statistics cache if none set
-            let default_stats = Arc::new(DefaultFileStatisticsCache::default());
-            config = config.with_files_statistics_cache(Some(default_stats));
-        }
+        // Add statistics cache if available - use default since CustomStatisticsCache doesn't implement FileStatisticsCache trait
+        let default_stats = Arc::new(DefaultFileStatisticsCache::default());
+        config = config.with_files_statistics_cache(Some(default_stats));
 
         config
     }

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

@@ -80,11 +80,15 @@ use crate::runtime_manager::RuntimeManager;
 
 mod statistics_cache;
 mod eviction_policy;
+mod liquid_cache_runtime;
+
+use liquid_cache_runtime::LiquidOnlyRuntime;
 
 struct DataFusionRuntime {
     runtime_env: RuntimeEnv,
     custom_cache_manager: Option<CustomCacheManager>,
     monitor: Arc<Monitor>,
+    liquid_cache_runtime_env: Arc<RuntimeEnv>,
 }
 
 // TASK monitorint metrics
@@ -291,10 +295,25 @@ pub extern "system" fn Java_org_opensearch_datafusion_jni_NativeBridge_createGlo
         .with_disk_manager_builder(builder)
         .build().unwrap();
 
+    log_info!("[LiquidCache] Initializing global Liquid Cache (1GB max)");
+    let liquid_runtime = match LiquidOnlyRuntime::init(1024 * 1024 * 1024) {
+        Ok(runtime) => {
+            log_info!("[LiquidCache] ✓ Liquid Cache initialized successfully");
+            runtime
+        },
+        Err(e) => {
+            log_error!("[LiquidCache] ERROR: Failed to initialize Liquid Cache: {}", e);
+            panic!("Failed to initialize Liquid Cache: {}", e);
+        }
+    };
+
+    let liquid_cache_runtime_env = liquid_runtime.runtime_env();
+
     let runtime = DataFusionRuntime {
         runtime_env,
         custom_cache_manager,
         monitor,
+        liquid_cache_runtime_env,
     };
 
     Box::into_raw(Box::new(runtime)) as jlong

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

@@ -0,0 +1,99 @@
+/*
+ * 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.
+ */
+
+use std::{fs, path::PathBuf, sync::{Arc, OnceLock}};
+
+use datafusion::{
+    common::DataFusionError,
+    execution::runtime_env::RuntimeEnv,
+    prelude::SessionConfig,
+};
+
+use liquid_cache_datafusion_local::LiquidCacheLocalBuilder;
+
+type LiquidCacheRef = std::sync::Arc<dyn std::any::Any + Send + Sync>;
+use vectorized_exec_spi::{log_info, log_debug};
+
+pub const LIQUID_CACHE_DIR: &str = "/var/lib/opensearch/liquid_cache";
+
+pub struct LiquidOnlyRuntime {
+    runtime_env: Arc<RuntimeEnv>,
+    liquid_cache: LiquidCacheRef,
+}
+
+static LIQUID_ONLY: OnceLock<Result<LiquidOnlyRuntime, String>> = OnceLock::new();
+
+impl LiquidOnlyRuntime {
+    pub fn init(max_cache_bytes: u64) -> Result<&'static LiquidOnlyRuntime, DataFusionError> {
+        log_info!("[LiquidCache] Initializing Liquid Cache with max_cache_bytes={}", max_cache_bytes);
+        
+        let result = LIQUID_ONLY.get_or_init(|| {
+            log_info!("[LiquidCache] Creating new Liquid Cache instance");
+            
+            let cache_dir = PathBuf::from(LIQUID_CACHE_DIR);
+            log_info!("[LiquidCache] Cache directory: {:?}", cache_dir);
+            
+            match fs::create_dir_all(&cache_dir) {
+                Ok(_) => log_info!("[LiquidCache] Cache directory created/verified successfully"),
+                Err(e) => {
+                    log_info!("[LiquidCache] ERROR: Failed to create cache directory: {}", e);
+                    return Err(format!(
+                        "Failed to create liquid cache dir {:?}: {}",
+                        cache_dir, e
+                    ));
+                }
+            }
+
+            let bootstrap_cfg = SessionConfig::new();
+            log_info!("[LiquidCache] Building Liquid Cache with LiquidCacheLocalBuilder");
+
+            let (liquid_ctx, cache_ref) = match LiquidCacheLocalBuilder::new()
+                .with_max_cache_bytes(max_cache_bytes as usize)
+                .with_cache_dir(cache_dir.clone())
+                .build(bootstrap_cfg) {
+                Ok(result) => {
+                    log_info!("[LiquidCache] ✓ Liquid Cache built successfully");
+                    log_info!("[LiquidCache] Cache directory: {:?}", cache_dir);
+                    log_info!("[LiquidCache] Max cache size: {} bytes ({} MB)", 
+                        max_cache_bytes, max_cache_bytes / 1024 / 1024);
+                    result
+                },
+                Err(e) => {
+                    log_info!("[LiquidCache] ERROR: Failed to build Liquid Cache: {}", e);
+                    return Err(format!("Liquid cache build failed: {e}"));
+                }
+            };
+
+            log_info!("[LiquidCache] ✓ Liquid Cache initialization complete");
+            
+            Ok(LiquidOnlyRuntime {
+                runtime_env: liquid_ctx.runtime_env(),
+                liquid_cache: cache_ref,
+            })
+        });
+
+        match result.as_ref() {
+            Ok(_) => {
+                log_info!("[LiquidCache] ✓ Liquid Cache ready for use");
+                Ok(result.as_ref().unwrap())
+            },
+            Err(e) => {
+                log_info!("[LiquidCache] ERROR: Liquid Cache initialization failed: {}", e);
+                Err(DataFusionError::Execution(e.clone()))
+            }
+        }
+    }
+
+    pub fn runtime_env(&self) -> Arc<RuntimeEnv> {
+        self.runtime_env.clone()
+    }
+
+    pub fn cache(&self) -> &LiquidCacheRef {
+        &self.liquid_cache
+    }
+}

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

@@ -948,7 +948,7 @@ pub struct ListingTable {
     schema_source: SchemaSource,
     options: ListingOptions,
     definition: Option<String>,
-    collected_statistics: FileStatisticsCache,
+    collected_statistics: Option<Arc<dyn FileStatisticsCache>>,
     constraints: Constraints,
     column_defaults: HashMap<String, Expr>,
     /// Optional [`SchemaAdapterFactory`] for creating schema adapters
@@ -992,7 +992,7 @@ impl ListingTable {
             schema_source,
             options,
             definition: None,
-            collected_statistics: Arc::new(DefaultFileStatisticsCache::default()),
+            collected_statistics: Some(Arc::new(DefaultFileStatisticsCache::default())),
             constraints: Constraints::default(),
             column_defaults: HashMap::new(),
             schema_adapter_factory: config.schema_adapter_factory,
@@ -1020,9 +1020,9 @@ impl ListingTable {
     /// multiple times in the same session.
     ///
     /// If `None`, creates a new [`DefaultFileStatisticsCache`] scoped to this query.
-    pub fn with_cache(mut self, cache: Option<FileStatisticsCache>) -> Self {
+    pub fn with_cache(mut self, cache: Option<Arc<dyn FileStatisticsCache>>) -> Self {
         self.collected_statistics =
-            cache.unwrap_or_else(|| Arc::new(DefaultFileStatisticsCache::default()));
+            Some(cache.unwrap_or_else(|| Arc::new(DefaultFileStatisticsCache::default())));
         self
     }
 
@@ -1098,7 +1098,7 @@ impl ListingTable {
 
     /// Creates a file source and applies schema adapter factory if available
     fn create_file_source_with_schema_adapter(&self) -> Result<Arc<dyn FileSource>> {
-        let mut source = self.options.format.file_source();
+        let mut source = self.options.format.file_source(self.file_schema.clone().into());
         // Apply schema adapter to source if available
         //
         // The source will use this SchemaAdapter to adapt data batches as they flow up the plan.
@@ -1309,16 +1309,14 @@ impl TableProvider for ListingTable {
                 state,
                 FileScanConfigBuilder::new(
                     object_store_url,
-                    Arc::clone(&self.file_schema),
                     file_source,
                 )
                 .with_file_groups(partitioned_file_lists)
                 .with_constraints(self.constraints.clone())
                 .with_statistics(statistics)
-                .with_projection_indices(projection.cloned())
+                .with_projection_indices(projection.cloned())?
                 .with_limit(limit)
                 .with_output_ordering(output_ordering)
-                .with_table_partition_cols(table_partition_cols)
                 .with_expr_adapter(self.expr_adapter_factory.clone())
                 .build(),
             )
@@ -1496,7 +1494,8 @@ impl ListingTable {
     ) -> Result<Arc<Statistics>> {
         match self
             .collected_statistics
-            .get_with_extra(&part_file.object_meta.location, &part_file.object_meta)
+            .as_ref()
+            .and_then(|cache| cache.get_with_extra(&part_file.object_meta.location, &part_file.object_meta))
         {
             Some(statistics) => Ok(statistics),
             None => {
@@ -1511,11 +1510,13 @@ impl ListingTable {
                     )
                     .await?;
                 let statistics = Arc::new(statistics);
-                self.collected_statistics.put_with_extra(
+                if let Some(cache) = &self.collected_statistics {
+                    cache.put_with_extra(
                     &part_file.object_meta.location,
                     Arc::clone(&statistics),
                     &part_file.object_meta,
                 );
+                }
                 Ok(statistics)
             }
         }