From c51e11c5c5e9de9509f41e272b934df63c439884 Mon Sep 17 00:00:00 2001 From: Arpit Bandejiya Date: Tue, 28 Jul 2026 15:10:06 +0530 Subject: [PATCH] Add optional liquid cache for indexed parquet scans MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds an opt-in, in-memory decoded-batch cache (liquid cache) for the DataFusion analytics backend, wired behind a cargo feature so it is absent from the default build. When the `liquid_cache` feature is off (the default), the integration is an inlined no-op: eligible scans use the plain ParquetSource unchanged. When on, an eligible indexed row-group scan — all projected columns numeric/date/timestamp/boolean, within a column budget, and no string or binary predicate column — is wrapped with LiquidParquetSource so decoded batches are served from and populated into the process-global cache. Enabling is a single-file edit plus a build flag: the liquid cache dependency is a commented placeholder in the analytics-backend-datafusion crate (uncomment it, set the git repo/branch, and uncomment the matching `dep:` feature entries), then build with -PliquidCache (or LIQUID_CACHE=1). The default build declares no liquid cache source, so it resolves and compiles exactly as before. Signed-off-by: Arpit Bandejiya --- sandbox/libs/dataformat-native/build.gradle | 11 ++ .../libs/dataformat-native/rust/Cargo.toml | 4 + .../rust/Cargo.toml | 24 +++ .../rust/src/indexed_table/parquet_bridge.rs | 12 +- .../rust/src/lib.rs | 1 + .../rust/src/liquid_cache.rs | 150 ++++++++++++++++++ 6 files changed, 199 insertions(+), 3 deletions(-) create mode 100644 sandbox/plugins/analytics-backend-datafusion/rust/src/liquid_cache.rs diff --git a/sandbox/libs/dataformat-native/build.gradle b/sandbox/libs/dataformat-native/build.gradle index 529ea63f2376b..f1f936f655e37 100644 --- a/sandbox/libs/dataformat-native/build.gradle +++ b/sandbox/libs/dataformat-native/build.gradle @@ -77,6 +77,17 @@ task buildRustLibrary(type: Exec) { def cargoArgs = [cargoExecutable, 'build', '-p', 'opensearch-native-lib'] if (buildType == 'release') { cargoArgs.add('--release') } + + // Optional Liquid Cache: pass -PliquidCache (or LIQUID_CACHE=1) to turn on the + // `liquid_cache` cargo feature. Off by default, so the plain build has no liquid + // cache dependency. NOTE: the liquid cache deps in Cargo.toml are commented-out + // placeholders by default — fill them in and add the `dep:` entries to the + // feature (see the notes in analytics-backend-datafusion/rust/Cargo.toml) before + // this flag actually compiles liquid cache in. + if (project.hasProperty('liquidCache') || System.getenv('LIQUID_CACHE') == '1') { + logger.lifecycle('Liquid Cache: feature enabled') + cargoArgs.addAll(['--features', 'opensearch-datafusion/liquid_cache']) + } commandLine cargoArgs inputs.files fileTree("${rustWorkspaceDir}/common/src") diff --git a/sandbox/libs/dataformat-native/rust/Cargo.toml b/sandbox/libs/dataformat-native/rust/Cargo.toml index b5f7cb0440f92..f5b51db012aa3 100644 --- a/sandbox/libs/dataformat-native/rust/Cargo.toml +++ b/sandbox/libs/dataformat-native/rust/Cargo.toml @@ -83,6 +83,10 @@ opensearch-repository-gcs = { path = "../../../plugins/native-repository-gcs/src opensearch-repository-azure = { path = "../../../plugins/native-repository-azure/src/main/rust" } opensearch-repository-fs = { path = "../../../plugins/native-repository-fs/src/main/rust" } +# Liquid Cache (optional) is declared directly in the enabling crate, +# analytics-backend-datafusion/rust/Cargo.toml — not here — so that turning it on +# is a single-file edit. See the placeholder block there. + [profile.release] lto = true codegen-units = 1 diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/Cargo.toml b/sandbox/plugins/analytics-backend-datafusion/rust/Cargo.toml index c12a4cb13c5c2..333c65a51a3a2 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/Cargo.toml +++ b/sandbox/plugins/analytics-backend-datafusion/rust/Cargo.toml @@ -95,6 +95,30 @@ crc32fast = { workspace = true } # rewrite. STANDARD alphabet matches the OpenSearch `binary` field wire contract. base64 = "0.22" +# ── Liquid Cache (optional) ────────────────────────────────────────────────── +# In-memory decoded-batch cache for indexed parquet scans. OFF by default and +# self-contained to THIS file: enabling it is a single-file edit here plus the +# Gradle flag. +# +# The dependency lines below are commented out on purpose. Cargo resolves (and +# fetches) every *declared* optional dep even when its feature is inactive, so a +# committed placeholder URL would break the default build. Leaving them commented +# keeps the default build free of any liquid cache dependency. +# +# To ENABLE liquid cache: +# 1. Uncomment the two dependency lines below and set the git repo + branch to +# your liquid cache source (must be API-compatible with this DataFusion / +# arrow version). +# 2. Uncomment the `dep:` entries in the `liquid_cache` feature line below. +# 3. Build with `-PliquidCache` (or LIQUID_CACHE=1); see build.gradle. +# liquid-cache = { git = "https://github.com/OWNER/REPO", branch = "BRANCH", optional = true } +# liquid-cache-datafusion = { git = "https://github.com/OWNER/REPO", branch = "BRANCH", optional = true } + +[features] +# Off by default. Bare flag so the default build resolves no liquid cache dep; +# add the `dep:` entries (see step 2 above) when enabling. +liquid_cache = [] # ["dep:liquid-cache", "dep:liquid-cache-datafusion"] + [dev-dependencies] criterion = { workspace = true } tempfile = { workspace = true } diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/parquet_bridge.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/parquet_bridge.rs index 957b0d15fa0dc..2a93a3995571c 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/parquet_bridge.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/parquet_bridge.rs @@ -250,9 +250,15 @@ fn create_stream_with_access_plan( } } - let mut config_builder = - FileScanConfigBuilder::new(config.store_url.clone(), Arc::new(parquet_source)) - .with_file(partitioned_file); + // Liquid Cache (optional): when the `liquid_cache` feature is compiled in and + // an eligible scan engages, wrap the ParquetSource with LiquidParquetSource so + // decoded batches are served from / populated into the in-memory cache. When + // the feature is off this is a no-op and the plain ParquetSource is used. + let file_source: Arc = + crate::liquid_cache::maybe_wrap_parquet_source(parquet_source, config); + + let mut config_builder = FileScanConfigBuilder::new(config.store_url.clone(), file_source) + .with_file(partitioned_file); if let Some(ref proj) = config.projection { // Empty projection (e.g. COUNT(*)) is honoured as "read no diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/lib.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/lib.rs index 8708627abe3a4..7f2e9eda7e53a 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/lib.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/lib.rs @@ -31,6 +31,7 @@ pub mod ffm; pub mod helper; pub mod indexed_executor; pub mod indexed_table; +pub mod liquid_cache; pub mod local_executor; pub mod memory; pub mod memory_guard; diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/liquid_cache.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/liquid_cache.rs new file mode 100644 index 0000000000000..8e68eec04d94f --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/liquid_cache.rs @@ -0,0 +1,150 @@ +/* + * 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. + */ + +//! Optional Liquid Cache integration — in-memory decoded-batch cache for +//! Parquet scans. +//! +//! This whole module is gated behind the `liquid_cache` cargo feature. When the +//! feature is off (the default), [`maybe_wrap_parquet_source`] is an inlined +//! no-op that hands the plain `ParquetSource` back, and none of the liquid cache +//! crates are compiled or even resolved. When the feature is on, an eligible +//! indexed scan (all projected columns numeric/date/timestamp/boolean, within a +//! column budget, and no string/binary predicate column) is wrapped with +//! `LiquidParquetSource`, which serves decoded batches from / populates them into +//! the process-global cache. + +use std::sync::Arc; + +use datafusion::datasource::physical_plan::{FileSource, ParquetSource}; + +use crate::indexed_table::parquet_bridge::RowGroupStreamConfig; + +/// Wrap `parquet_source` with liquid cache when the feature is enabled and the +/// scan is eligible; otherwise return the source unchanged. The return type is +/// the erased `Arc` accepted by `FileScanConfigBuilder`, so both +/// branches are interchangeable at the call site. +#[cfg(not(feature = "liquid_cache"))] +#[inline] +pub(crate) fn maybe_wrap_parquet_source( + parquet_source: ParquetSource, + _config: &RowGroupStreamConfig, +) -> Arc { + Arc::new(parquet_source) +} + +#[cfg(feature = "liquid_cache")] +pub(crate) fn maybe_wrap_parquet_source( + parquet_source: ParquetSource, + config: &RowGroupStreamConfig, +) -> Arc { + match runtime::LiquidRuntime::cache_ref_if_enabled() { + Some(cache_ref) if runtime::scan_is_eligible(config) => { + let liquid_source = liquid_cache_datafusion::LiquidParquetSource::from_parquet_source( + parquet_source, + cache_ref, + ); + Arc::new(liquid_source) + } + _ => Arc::new(parquet_source), + } +} + +#[cfg(feature = "liquid_cache")] +mod runtime { + use std::sync::atomic::{AtomicU32, Ordering}; + use std::sync::OnceLock; + + use datafusion::arrow::datatypes::DataType; + use liquid_cache_datafusion::LiquidCacheParquetRef; + + use crate::indexed_table::parquet_bridge::RowGroupStreamConfig; + + /// Liquid cache batch size — must be a power of two (upstream default). + const LIQUID_CACHE_BATCH_SIZE: usize = 8192; + + static INSTANCE: OnceLock = OnceLock::new(); + static MAX_COLUMNS: AtomicU32 = AtomicU32::new(10); + + /// Column budget above which a scan is not delegated to liquid cache. + pub(super) fn max_columns() -> usize { + MAX_COLUMNS.load(Ordering::Relaxed) as usize + } + + pub(super) fn set_max_columns(value: usize) { + MAX_COLUMNS.store(value as u32, Ordering::Relaxed); + } + + /// The process-global cache reference, if liquid cache has been initialized. + pub(super) fn cache_ref_if_enabled() -> Option { + INSTANCE.get().cloned() + } + + /// Install the process-global cache reference. First call wins. + pub(super) fn set_cache_ref(cache_ref: LiquidCacheParquetRef) { + let _ = INSTANCE.set(cache_ref); + } + + /// A scan is eligible when every projected column is numeric/date/timestamp/ + /// boolean, the projection is non-empty and within the column budget, and no + /// predicate column is string/binary (parity with the original engagement + /// gate — strings were never cacheable in-process). + pub(super) fn scan_is_eligible(config: &RowGroupStreamConfig) -> bool { + let all_numeric_projection = config.projection.as_ref().is_some_and(|proj| { + !proj.is_empty() + && proj.len() <= max_columns() + && proj.iter().all(|&idx| { + config + .full_schema + .fields() + .get(idx) + .is_some_and(|f| is_cacheable(f.data_type())) + }) + }); + if all_numeric_projection == false { + return false; + } + let predicate_has_string = config.predicate.as_ref().is_some_and(|pred| { + datafusion::physical_expr::utils::collect_columns(pred) + .iter() + .any(|col| { + config + .full_schema + .fields() + .get(col.index()) + .is_some_and(|f| is_string_or_binary(f.data_type())) + }) + }); + predicate_has_string == false + } + + fn is_cacheable(dt: &DataType) -> bool { + dt.is_numeric() + || matches!( + dt, + DataType::Date32 | DataType::Date64 | DataType::Timestamp(_, _) | DataType::Boolean + ) + } + + fn is_string_or_binary(dt: &DataType) -> bool { + matches!( + dt, + DataType::Utf8 + | DataType::Utf8View + | DataType::LargeUtf8 + | DataType::Binary + | DataType::BinaryView + | DataType::LargeBinary + ) + } +} + +/// Re-exports for the initialization path (used by the runtime bootstrap when +/// the feature is enabled). No-ops are not provided for the feature-off build +/// because the initializer is itself feature-gated. +#[cfg(feature = "liquid_cache")] +pub(crate) use runtime::{set_cache_ref, set_max_columns};