diff --git a/sandbox/libs/dataformat-native/rust/Cargo.lock b/sandbox/libs/dataformat-native/rust/Cargo.lock index e5d42661fca32..3ae84106485da 100644 --- a/sandbox/libs/dataformat-native/rust/Cargo.lock +++ b/sandbox/libs/dataformat-native/rust/Cargo.lock @@ -684,6 +684,15 @@ version = "0.4.32" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cc14f565cf027a105f7a44ccf9e5b424348421a1d8952a8fc9d499d313107789" +[[package]] +name = "congee" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72819224a9c43e0a5d64bdad35a71fde68bf479992fb8989635c72bc547909be" +dependencies = [ + "crossbeam-epoch", +] + [[package]] name = "const-oid" version = "0.10.2" @@ -710,6 +719,12 @@ dependencies = [ "tiny-keccak", ] +[[package]] +name = "const_for" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "988d3bd6bf67b6d7ae2b519c55296fa57411c27c312575e47b2975f8d1d8590b" + [[package]] name = "constant_time_eq" version = "0.4.2" @@ -743,6 +758,12 @@ dependencies = [ "winapi", ] +[[package]] +name = "core_detect" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f8f80099a98041a3d1622845c271458a2d73e688351bf3cb999266764b81d48" + [[package]] name = "cpufeatures" version = "0.2.17" @@ -1666,6 +1687,19 @@ dependencies = [ "web-time", ] +[[package]] +name = "fastlanes" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9272a674b446f53253f66582596614e41552d14968a3d1aaa590576a91c1188" +dependencies = [ + "const_for", + "core_detect", + "num-traits", + "paste", + "seq-macro", +] + [[package]] name = "fastrand" version = "2.4.1" @@ -2689,6 +2723,7 @@ checksum = "1d87ecb2933e8aeadb3e3a02b828fed80a7528047e68b4f424523a0981a3a084" name = "native-bridge-common" version = "0.1.0" dependencies = [ + "log", "native-bridge-macros", "tikv-jemalloc-ctl", "tikv-jemalloc-sys", @@ -2861,6 +2896,8 @@ dependencies = [ "object_store", "once_cell", "opensearch-block-cache", + "opensearch-liquid-cache-core", + "opensearch-liquid-cache-datafusion", "opensearch-tiered-storage", "parking_lot", "parquet", @@ -2881,6 +2918,41 @@ dependencies = [ "url", ] +[[package]] +name = "opensearch-liquid-cache-core" +version = "0.1.0" +dependencies = [ + "ahash", + "arrow", + "arrow-schema", + "congee", + "datafusion-common", + "datafusion-expr-common", + "datafusion-physical-expr", + "fastlanes", + "log", + "num-traits", + "serde", +] + +[[package]] +name = "opensearch-liquid-cache-datafusion" +version = "0.1.0" +dependencies = [ + "ahash", + "arrow", + "arrow-schema", + "bytes", + "datafusion", + "futures", + "log", + "object_store", + "opensearch-liquid-cache-core", + "parquet", + "tempfile", + "tokio", +] + [[package]] name = "opensearch-native-lib" version = "0.1.0" diff --git a/sandbox/libs/dataformat-native/rust/Cargo.toml b/sandbox/libs/dataformat-native/rust/Cargo.toml index b5f7cb0440f92..0af25e842e7b9 100644 --- a/sandbox/libs/dataformat-native/rust/Cargo.toml +++ b/sandbox/libs/dataformat-native/rust/Cargo.toml @@ -4,6 +4,8 @@ members = [ "common", "macros", "lib", + "liquid-cache/core", + "liquid-cache/datafusion", "../../../plugins/analytics-backend-datafusion/rust", "../../../plugins/parquet-data-format/src/main/rust", "../../../plugins/native-repository-s3/src/main/rust", @@ -30,6 +32,8 @@ datafusion-datasource = "=54.0.0" datafusion-common = "=54.0.0" datafusion-execution = "=54.0.0" datafusion-physical-expr = "=54.0.0" +datafusion-physical-expr-common = "=54.0.0" +datafusion-expr-common = "=54.0.0" datafusion-substrait = "=54.0.0" # Async @@ -56,6 +60,12 @@ tikv-jemallocator = { version = "=0.6.1", features = ["disable_initial_exec_tls" tikv-jemalloc-ctl = { version = "=0.6.1", features = ["stats"] } tikv-jemalloc-sys = { version = "=0.6.1", features = ["stats"] } +# Liquid cache (vendored subset — see liquid-cache/README.md) +ahash = "=0.8.12" +congee = "=0.4.1" +fastlanes = "=0.5.2" +num-traits = "=0.2.19" + # Misc dashmap = { version = "=5.5.3", features = ["raw-api", "serde"] } num_cpus = "=1.17.0" @@ -77,6 +87,8 @@ proptest = "=1.4.0" # Internal native-bridge-common = { path = "common" } +opensearch-liquid-cache-core = { path = "liquid-cache/core" } +opensearch-liquid-cache-datafusion = { path = "liquid-cache/datafusion" } opensearch-tiered-storage = { path = "../../tiered-storage/src/main/rust" } opensearch-repository-s3 = { path = "../../../plugins/native-repository-s3/src/main/rust" } opensearch-repository-gcs = { path = "../../../plugins/native-repository-gcs/src/main/rust" } diff --git a/sandbox/libs/dataformat-native/rust/common/Cargo.toml b/sandbox/libs/dataformat-native/rust/common/Cargo.toml index 641b96da148e1..93c42c77e3e75 100644 --- a/sandbox/libs/dataformat-native/rust/common/Cargo.toml +++ b/sandbox/libs/dataformat-native/rust/common/Cargo.toml @@ -13,6 +13,7 @@ native-bridge-macros = { path = "../macros" } tikv-jemalloc-ctl = { workspace = true } tikv-jemalloc-sys = { workspace = true } tokio = { workspace = true } +log = { workspace = true } [dev-dependencies] tikv-jemallocator = { workspace = true } diff --git a/sandbox/libs/dataformat-native/rust/common/src/logger.rs b/sandbox/libs/dataformat-native/rust/common/src/logger.rs index f48ba3259a53a..3e47d3e35bbd6 100644 --- a/sandbox/libs/dataformat-native/rust/common/src/logger.rs +++ b/sandbox/libs/dataformat-native/rust/common/src/logger.rs @@ -49,10 +49,42 @@ pub extern "C" fn native_logger_set_level(level: i32) { MAX_LEVEL.store(clamped, Ordering::Relaxed); } +/// A `log` crate backend that forwards to our native bridge callback. +/// This makes `log::info!()` from any crate (including liquid-cache) +/// appear in OpenSearch's log file via the Java RustLoggerBridge. +struct NativeBridgeLogger; + +impl ::log::Log for NativeBridgeLogger { + fn enabled(&self, metadata: &::log::Metadata) -> bool { + metadata.level() <= ::log::Level::Info + } + + fn log(&self, record: &::log::Record) { + if !self.enabled(record.metadata()) { + return; + } + let level = match record.level() { + ::log::Level::Error => LogLevel::Error, + ::log::Level::Warn | ::log::Level::Info => LogLevel::Info, + _ => LogLevel::Debug, + }; + let msg = format!("{}", record.args()); + log(level, &msg); + } + + fn flush(&self) {} +} + +static NATIVE_BRIDGE_LOGGER: NativeBridgeLogger = NativeBridgeLogger; + /// Called by Java at startup to register the log callback. #[no_mangle] pub unsafe extern "C" fn native_logger_init(callback: LogCallback) { LOG_CALLBACK.store(callback as *mut (), Ordering::Release); + // Initialize the `log` crate facade so that log::info!() from any + // dependency (e.g. liquid-cache) routes through our native bridge. + let _ = ::log::set_logger(&NATIVE_BRIDGE_LOGGER); + ::log::set_max_level(::log::LevelFilter::Info); log(LogLevel::Info, "Native logger initialized successfully"); } diff --git a/sandbox/libs/dataformat-native/rust/liquid-cache/README.md b/sandbox/libs/dataformat-native/rust/liquid-cache/README.md new file mode 100644 index 0000000000000..d5c43b9ae656a --- /dev/null +++ b/sandbox/libs/dataformat-native/rust/liquid-cache/README.md @@ -0,0 +1,43 @@ +# opensearch-liquid-cache (vendored subset) + +In-memory-only subset of [liquid-cache](https://github.com/XiangpengHao/liquid-cache), +vendored so the OpenSearch sandbox takes **no Cargo dependency** on the upstream +package. + +## Provenance + +- Vendored from: https://github.com/cocosz/liquid-cache +- Branch: `lc-opensearch-df54-v2` +- Commit: `8311bccc258756127adbebee3e54140b60fc09ba` +- License: Apache-2.0 (same as upstream) + +## What was removed relative to upstream + +Everything needed only for the disk tier, client/server mode, or data types the +OpenSearch integration gate never caches: + +- **Disk tier**: `t4` store (io-uring), `DiskLiquid`/`DiskArrow` cache entries, + `SqueezeIoHandler`/`DefaultSqueezeIo`, squeeze-to-disk policies, `ipc.rs` + serialization and per-array `to_bytes`/`from_bytes`. The cache is memory-only; + under pressure entries are transcoded (Arrow → Liquid) and then evicted. +- **Client/server**: the whole `liquid-cache-common` crate (arrow-flight/tonic/axum), + `datafusion-client`, `datafusion-server`. +- **String/binary caching**: `byte_view_array` + FSST (`fsst-rs`). The integration + gate only engages LC for numeric/date/timestamp/boolean projections. +- **Variant/JSON shredding**: `variant_array`, `parquet-variant-*` deps, variant UDFs. +- **Date32 squeeze**: `squeezed_date32_array` (a squeeze-to-disk optimization). +- **LineageOptimizer**: excluded by the POC already (planning overhead). +- **Tracing**: `fastrace`, `sysinfo`. + +## Crates + +- `core/` — package `opensearch-liquid-cache-core`, **lib name `liquid_cache`**: + cache storage (index, budget, eviction policies, transcode) + numeric + LiquidArray encodings. +- `datafusion/` — package `opensearch-liquid-cache-datafusion`, **lib name + `liquid_cache_datafusion`**: `LiquidParquetSource` (the ParquetSource + replacement), reader/stream machinery, `LocalModeOptimizer`. + +The lib names intentionally match upstream so vendored code and the +`analytics-backend-datafusion` integration compile with unchanged `use` paths, +which also keeps future re-syncs against upstream reviewable. diff --git a/sandbox/libs/dataformat-native/rust/liquid-cache/core/Cargo.toml b/sandbox/libs/dataformat-native/rust/liquid-cache/core/Cargo.toml new file mode 100644 index 0000000000000..dbcefdcdf512e --- /dev/null +++ b/sandbox/libs/dataformat-native/rust/liquid-cache/core/Cargo.toml @@ -0,0 +1,31 @@ +# SPDX-License-Identifier: Apache-2.0 +# +# Vendored in-memory subset of liquid-cache core. +# Provenance: https://github.com/cocosz/liquid-cache branch lc-opensearch-df54-v2 +# commit 8311bccc258756127adbebee3e54140b60fc09ba. See ../README.md. + +[package] +name = "opensearch-liquid-cache-core" +version = "0.1.0" +edition = "2024" +license = "Apache-2.0" +publish = false + +[lib] +# Keep the upstream crate name so vendored code and the integration keep +# their `use liquid_cache::...` paths, making future upstream re-syncs diffable. +name = "liquid_cache" +path = "src/lib.rs" + +[dependencies] +ahash = { workspace = true } +arrow = { workspace = true } +arrow-schema = { workspace = true } +congee = { workspace = true } +datafusion-common = { workspace = true } +datafusion-expr-common = { workspace = true } +datafusion-physical-expr = { workspace = true } +fastlanes = { workspace = true } +log = { workspace = true } +num-traits = { workspace = true } +serde = { workspace = true } diff --git a/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/cache/budget.rs b/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/cache/budget.rs new file mode 100644 index 0000000000000..4069f23a28417 --- /dev/null +++ b/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/cache/budget.rs @@ -0,0 +1,164 @@ +use crate::sync::atomic::{AtomicUsize, Ordering}; + +#[derive(Debug)] +pub struct BudgetAccounting { + max_memory_bytes: AtomicUsize, + used_memory_bytes: AtomicUsize, +} + +impl BudgetAccounting { + pub(super) fn new(max_memory_bytes: usize) -> Self { + Self { + max_memory_bytes: AtomicUsize::new(max_memory_bytes), + used_memory_bytes: AtomicUsize::new(0), + } + } + + pub(super) fn reset_usage(&self) { + self.used_memory_bytes.store(0, Ordering::Relaxed); + } + + /// Dynamically update the max memory limit. Takes effect for new reservations. + pub fn set_max_memory_bytes(&self, new_limit: usize) { + self.max_memory_bytes.store(new_limit, Ordering::Relaxed); + } + + pub fn max_memory_bytes(&self) -> usize { + self.max_memory_bytes.load(Ordering::Relaxed) + } + + /// Try to reserve memory in the cache. + /// Returns ok if the memory was reserved, err if the memory budget is full. + pub(super) fn try_reserve_memory(&self, request_bytes: usize) -> Result<(), ()> { + let used = self.used_memory_bytes.load(Ordering::Relaxed); + if used + request_bytes > self.max_memory_bytes.load(Ordering::Relaxed) { + return Err(()); + } + + match self.used_memory_bytes.compare_exchange( + used, + used + request_bytes, + Ordering::Relaxed, + Ordering::Relaxed, + ) { + Ok(_) => Ok(()), + Err(_) => self.try_reserve_memory(request_bytes), + } + } + + /// Adjust memory usage after transcoding. + /// Returns ok if the usage was adjusted, err if the memory budget is full when new_size is larger than old_size. + pub(super) fn try_update_memory_usage( + &self, + old_size: usize, + new_size: usize, + ) -> Result<(), ()> { + if old_size < new_size { + let diff = new_size - old_size; + self.try_reserve_memory(diff)?; + Ok(()) + } else { + self.used_memory_bytes + .fetch_sub(old_size - new_size, Ordering::Relaxed); + Ok(()) + } + } + + pub fn memory_usage_bytes(&self) -> usize { + self.used_memory_bytes.load(Ordering::Relaxed) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::sync::{Arc, Barrier, thread}; + + fn test_budget(max_memory_bytes: usize) -> BudgetAccounting { + BudgetAccounting::new(max_memory_bytes) + } + + #[test] + fn test_memory_reservation_and_accounting() { + let config = test_budget(1000); + + assert_eq!(config.memory_usage_bytes(), 0); + + assert!(config.try_reserve_memory(500).is_ok()); + assert_eq!(config.memory_usage_bytes(), 500); + + assert!(config.try_reserve_memory(300).is_ok()); + assert_eq!(config.memory_usage_bytes(), 800); + + assert!(config.try_reserve_memory(300).is_err()); + assert_eq!(config.memory_usage_bytes(), 800); + + config.reset_usage(); + assert_eq!(config.memory_usage_bytes(), 0); + } + + #[test] + fn test_concurrent_memory_operations() { + test_concurrent_memory_budget(); + } + + fn test_concurrent_memory_budget() { + let num_threads = 3; + let max_memory = 10000; + let operations_per_thread = 100; + + let budget = Arc::new(test_budget(max_memory)); + let barrier = Arc::new(Barrier::new(num_threads)); + + let mut thread_handles = vec![]; + + for _ in 0..num_threads { + let budget_clone = budget.clone(); + let barrier_clone = barrier.clone(); + + let handle = thread::spawn(move || { + let mut successful_reservations = Vec::new(); + + barrier_clone.wait(); + + for i in 0..operations_per_thread { + let reserve_size = 10 + (i % 20) * 5; // 10 to 105 bytes + if budget_clone.try_reserve_memory(reserve_size).is_ok() { + successful_reservations.push(reserve_size); + } + + if i.is_multiple_of(5) && !successful_reservations.is_empty() { + let idx = i % successful_reservations.len(); + let old_size = successful_reservations[idx]; + let new_size = if i.is_multiple_of(2) { + old_size + 5 // Grow + } else { + old_size.saturating_sub(5) // Shrink + }; + + if budget_clone + .try_update_memory_usage(old_size, new_size) + .is_ok() + { + successful_reservations[idx] = new_size; + } + } + } + successful_reservations + }); + + thread_handles.push(handle); + } + + let mut expected_memory_usage = 0; + for handle in thread_handles { + let reservations = handle.join().unwrap(); + for size in reservations { + expected_memory_usage += size; + } + } + + assert_eq!(budget.memory_usage_bytes(), expected_memory_usage); + assert!(budget.memory_usage_bytes() <= max_memory); + } +} diff --git a/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/cache/builders.rs b/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/cache/builders.rs new file mode 100644 index 0000000000000..4d0570b879003 --- /dev/null +++ b/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/cache/builders.rs @@ -0,0 +1,391 @@ +use std::future::{IntoFuture, Ready}; + +use arrow::array::{ + Array, ArrayData, ArrayRef, BinaryViewArray, BooleanArray, StringViewArray, make_array, +}; +use arrow::buffer::BooleanBuffer; + +use super::cached_batch::CacheEntry; +use super::core::{CacheFull, LiquidCache}; +use super::policies::{CachePolicy, SqueezePolicy, TranscodeEvict}; +use super::{EntryID, LiquidExpr, LiquidPolicy}; +use crate::sync::Arc; + +/// Builder for [LiquidCache]. +/// +/// Example: +/// ```rust +/// use liquid_cache::cache::LiquidCacheBuilder; +/// use liquid_cache::cache_policies::LiquidPolicy; +/// +/// let _storage = LiquidCacheBuilder::new() +/// .with_batch_size(8192) +/// .with_max_memory_bytes(1024 * 1024 * 1024) +/// .with_cache_policy(Box::new(LiquidPolicy::new())) +/// .build(); +/// ``` +pub struct LiquidCacheBuilder { + batch_size: usize, + max_memory_bytes: usize, + cache_policy: Box, + squeeze_policy: Box, +} + +impl Default for LiquidCacheBuilder { + fn default() -> Self { + Self::new() + } +} + +/// Default max memory when none is provided: 1 GiB. +const DEFAULT_MAX_MEMORY_BYTES: usize = 1 << 30; + +impl LiquidCacheBuilder { + /// Create a new instance of [LiquidCacheBuilder]. + pub fn new() -> Self { + Self { + batch_size: 8192, + max_memory_bytes: DEFAULT_MAX_MEMORY_BYTES, + cache_policy: Box::new(LiquidPolicy::new()), + squeeze_policy: Box::new(TranscodeEvict), + } + } + + /// Set the batch size for the cache. + /// Default is 8192. + pub fn with_batch_size(mut self, batch_size: usize) -> Self { + self.batch_size = batch_size; + self + } + + /// Set the max memory bytes for the cache. + /// Default is 1 GiB. + pub fn with_max_memory_bytes(mut self, max_memory_bytes: usize) -> Self { + self.max_memory_bytes = max_memory_bytes; + self + } + + /// Set the cache policy for the cache. + /// Default is [LiquidPolicy]. + pub fn with_cache_policy(mut self, policy: Box) -> Self { + self.cache_policy = policy; + self + } + + /// Set the squeeze policy for the cache. + /// Default is [TranscodeEvict]. + pub fn with_squeeze_policy(mut self, policy: Box) -> Self { + self.squeeze_policy = policy; + self + } + + /// Build the cache storage. + /// + /// The cache storage is wrapped in an [Arc] to allow for concurrent access. + pub fn build(self) -> Arc { + Arc::new(LiquidCache::new( + self.batch_size, + self.max_memory_bytes, + self.squeeze_policy, + self.cache_policy, + )) + } +} + +/// Builder returned by [`LiquidCache::insert`] for configuring cache writes. +#[derive(Debug)] +pub struct Insert<'a> { + pub(super) storage: &'a LiquidCache, + pub(super) entry_id: EntryID, + pub(super) batch: ArrayRef, + pub(super) skip_gc: bool, +} + +impl<'a> Insert<'a> { + pub(super) fn new(storage: &'a LiquidCache, entry_id: EntryID, batch: ArrayRef) -> Self { + Self { + storage, + entry_id, + batch, + skip_gc: false, + } + } + + /// Skip garbage collection of view arrays. + pub fn with_skip_gc(mut self) -> Self { + self.skip_gc = true; + self + } + + /// Insert the batch into the cache. + pub fn execute(self) -> Result<(), CacheFull> { + let batch = if self.skip_gc { + self.batch.clone() + } else { + maybe_gc_view_arrays(&self.batch).unwrap_or_else(|| self.batch.clone()) + }; + let batch = CacheEntry::memory_arrow(batch); + self.storage.insert_inner(self.entry_id, batch) + } +} + +impl<'a> IntoFuture for Insert<'a> { + type Output = Result<(), CacheFull>; + type IntoFuture = Ready; + + fn into_future(self) -> Self::IntoFuture { + std::future::ready(self.execute()) + } +} + +/// Builder returned by [`LiquidCache::get`] for configuring cache reads. +#[derive(Debug)] +pub struct Get<'a> { + pub(super) storage: &'a LiquidCache, + pub(super) entry_id: &'a EntryID, + pub(super) selection: Option<&'a BooleanBuffer>, +} + +impl<'a> Get<'a> { + pub(super) fn new(storage: &'a LiquidCache, entry_id: &'a EntryID) -> Self { + Self { + storage, + entry_id, + selection: None, + } + } + + /// Attach a selection bitmap used to filter rows prior to materialization. + pub fn with_selection(mut self, selection: &'a BooleanBuffer) -> Self { + self.selection = Some(selection); + self + } + + /// Materialize the cached array as [`ArrayRef`]. + pub fn read(self) -> Option { + self.storage.observer().on_get(self.selection.is_some()); + self.storage.read_arrow_array(self.entry_id, self.selection) + } +} + +impl<'a> IntoFuture for Get<'a> { + type Output = Option; + type IntoFuture = Ready; + + fn into_future(self) -> Self::IntoFuture { + std::future::ready(self.read()) + } +} + +/// Recursively garbage collects view arrays (BinaryView/StringView) within an array tree. +fn maybe_gc_view_arrays(array: &ArrayRef) -> Option { + if let Some(binary_view) = array.as_any().downcast_ref::() { + return Some(Arc::new(binary_view.gc())); + } + if let Some(utf8_view) = array.as_any().downcast_ref::() { + return Some(Arc::new(utf8_view.gc())); + } + + let data = array.to_data(); + if data.child_data().is_empty() { + return None; + } + + let mut changed = false; + let mut children: Vec = Vec::with_capacity(data.child_data().len()); + for child in data.child_data() { + let child_array = make_array(child.clone()); + if let Some(gc_child) = maybe_gc_view_arrays(&child_array) { + changed = true; + children.push(gc_child.to_data()); + } else { + children.push(child.clone()); + } + } + + if !changed { + return None; + } + + let new_data = data.into_builder().child_data(children).build().ok()?; + Some(make_array(new_data)) +} + +/// Builder for predicate evaluation on cached data. +#[derive(Debug)] +pub struct EvaluatePredicate<'a> { + pub(super) storage: &'a LiquidCache, + pub(super) entry_id: &'a EntryID, + pub(super) predicate: &'a LiquidExpr, + pub(super) selection: Option<&'a BooleanBuffer>, +} + +impl<'a> EvaluatePredicate<'a> { + pub(super) fn new( + storage: &'a LiquidCache, + entry_id: &'a EntryID, + predicate: &'a LiquidExpr, + ) -> Self { + Self { + storage, + entry_id, + predicate, + selection: None, + } + } + + /// Attach a selection bitmap used to pre-filter rows before predicate evaluation. + pub fn with_selection(mut self, selection: &'a BooleanBuffer) -> Self { + self.selection = Some(selection); + self + } + + /// Evaluate the predicate against the cached data. + pub fn read(self) -> Option { + self.storage + .eval_predicate_internal(self.entry_id, self.selection, self.predicate) + } +} + +impl<'a> IntoFuture for EvaluatePredicate<'a> { + type Output = Option; + type IntoFuture = Ready; + + fn into_future(self) -> Self::IntoFuture { + std::future::ready(self.read()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{AsArray, StructArray}; + use arrow_schema::{DataType, Field, Fields}; + + #[test] + fn insert_gcs_view_arrays_recursively() { + // Build view arrays then slice to create non-zero offsets (and larger backing buffers). + let bin = Arc::new(BinaryViewArray::from(vec![ + Some(b"long_prefix_m0" as &[u8]), + Some(b"m1"), + ])) as ArrayRef; + let str_view = Arc::new(StringViewArray::from(vec![ + Some("long_prefix_s0"), + Some("s1"), + ])) as ArrayRef; + let nested_metadata = Arc::new(BinaryViewArray::from(vec![ + Some(b"meta0" as &[u8]), + Some(b"meta1"), + ])) as ArrayRef; + let nested_value = Arc::new(BinaryViewArray::from(vec![ + Some(b"value0" as &[u8]), + Some(b"value1"), + ])) as ArrayRef; + + // Slice to keep only the second element so buffers still reference unused bytes. + let bin_slice = bin.slice(1, 1); + let str_slice = str_view.slice(1, 1); + let nested_metadata_slice = nested_metadata.slice(1, 1); + let nested_value_slice = nested_value.slice(1, 1); + + // Nested struct: metadata (BinaryView), value (BinaryView), and a typed string view. + let nested_typed_fields = Fields::from(vec![Arc::new(Field::new( + "typed_str", + DataType::Utf8View, + true, + ))]); + let nested_struct_fields = Fields::from(vec![ + Arc::new(Field::new("metadata", DataType::BinaryView, true)), + Arc::new(Field::new("value", DataType::BinaryView, true)), + Arc::new(Field::new( + "typed_value", + DataType::Struct(nested_typed_fields.clone()), + true, + )), + ]); + let nested_struct = Arc::new(StructArray::new( + nested_struct_fields.clone(), + vec![ + nested_metadata_slice.clone(), + nested_value_slice.clone(), + Arc::new(StructArray::new( + nested_typed_fields.clone(), + vec![str_slice.clone()], + None, + )) as ArrayRef, + ], + None, + )); + + let root_fields = Fields::from(vec![ + Arc::new(Field::new("bin_view", DataType::BinaryView, true)), + Arc::new(Field::new("str_view", DataType::Utf8View, true)), + Arc::new(Field::new( + "nested", + DataType::Struct(nested_struct_fields.clone()), + true, + )), + ]); + let root = Arc::new(StructArray::new( + root_fields, + vec![ + bin_slice.clone(), + str_slice.clone(), + nested_struct.clone() as ArrayRef, + ], + None, + )) as ArrayRef; + + let pre_size = root.get_array_memory_size(); + + let cache = LiquidCacheBuilder::new().build(); + let entry_id = EntryID::from(123usize); + cache.insert(entry_id, root.clone()).execute().unwrap(); + + let stored = cache.get(&entry_id).read().expect("array present"); + let post_size = stored.get_array_memory_size(); + + // GC should have compacted the view arrays, reducing memory footprint. + assert!(post_size < pre_size, "expected gc to reduce memory usage"); + + // Validate values are preserved. + let struct_out = stored + .as_any() + .downcast_ref::() + .expect("struct array"); + + assert_eq!(struct_out.len(), 1); + + let bin_out = struct_out + .column_by_name("bin_view") + .unwrap() + .as_binary_view(); + assert_eq!(bin_out.value(0), b"m1"); + + let str_out = struct_out + .column_by_name("str_view") + .unwrap() + .as_string_view(); + assert_eq!(str_out.value(0), "s1"); + + let nested_out = struct_out.column_by_name("nested").unwrap().as_struct(); + let meta_out = nested_out + .column_by_name("metadata") + .unwrap() + .as_binary_view(); + assert_eq!(meta_out.value(0), b"meta1"); + + let val_out = nested_out.column_by_name("value").unwrap().as_binary_view(); + assert_eq!(val_out.value(0), b"value1"); + + let typed_out = nested_out + .column_by_name("typed_value") + .unwrap() + .as_struct(); + let typed_str_out = typed_out + .column_by_name("typed_str") + .unwrap() + .as_string_view(); + assert_eq!(typed_str_out.value(0), "s1"); + } +} diff --git a/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/cache/cached_batch.rs b/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/cache/cached_batch.rs new file mode 100644 index 0000000000000..053edf0eb7206 --- /dev/null +++ b/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/cache/cached_batch.rs @@ -0,0 +1,71 @@ +//! Cached batch types. + +use std::{fmt::Display, sync::Arc}; + +use arrow::array::ArrayRef; + +use crate::liquid_array::LiquidArrayRef; + +/// A cached entry storing data in various formats. +#[derive(Debug, Clone)] +pub enum CacheEntry { + /// Cached batch in memory as Arrow array. + MemoryArrow(ArrayRef), + /// Cached batch in memory as liquid array. + MemoryLiquid(LiquidArrayRef), +} + +impl CacheEntry { + /// Construct a cached batch stored as an in-memory Arrow array. + pub fn memory_arrow(array: ArrayRef) -> Self { + Self::MemoryArrow(array) + } + + /// Construct a cached batch stored as an in-memory Liquid array. + pub fn memory_liquid(array: LiquidArrayRef) -> Self { + Self::MemoryLiquid(array) + } + + /// Memory usage reported by the underlying representation. + pub fn memory_usage_bytes(&self) -> usize { + match self { + Self::MemoryArrow(array) => array.get_array_memory_size(), + Self::MemoryLiquid(array) => array.get_array_memory_size(), + } + } + + /// Reference count (if any) of the backing storage. + pub fn reference_count(&self) -> usize { + match self { + Self::MemoryArrow(array) => Arc::strong_count(array), + Self::MemoryLiquid(array) => Arc::strong_count(array), + } + } +} + +impl Display for CacheEntry { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::MemoryArrow(_) => write!(f, "MemoryArrow"), + Self::MemoryLiquid(_) => write!(f, "MemoryLiquid"), + } + } +} + +/// The type of the cached batch. +#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize)] +pub enum CachedBatchType { + /// Cached batch in memory as Arrow array. + MemoryArrow, + /// Cached batch in memory as liquid array. + MemoryLiquid, +} + +impl From<&CacheEntry> for CachedBatchType { + fn from(batch: &CacheEntry) -> Self { + match batch { + CacheEntry::MemoryArrow(_) => Self::MemoryArrow, + CacheEntry::MemoryLiquid(_) => Self::MemoryLiquid, + } + } +} diff --git a/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/cache/core.rs b/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/cache/core.rs new file mode 100644 index 0000000000000..9099ea0dd9394 --- /dev/null +++ b/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/cache/core.rs @@ -0,0 +1,681 @@ +use arrow::array::cast::AsArray; +use arrow::array::{ArrayRef, BooleanArray}; +use arrow::buffer::BooleanBuffer; +use arrow::record_batch::RecordBatch; +use arrow_schema::{Field, Schema}; + +use super::{ + budget::BudgetAccounting, + builders::{EvaluatePredicate, Get, Insert}, + cached_batch::{CacheEntry, CachedBatchType}, + observer::Observer, + policies::CachePolicy, + utils::CacheConfig, +}; +use crate::cache::CacheStats; +use crate::cache::policies::{SqueezeOutcome, SqueezePolicy}; +use crate::cache::{LiquidExpr, index::ArtIndex, utils::EntryID}; +use crate::sync::Arc; + +/// The cache could not free enough memory to admit an entry. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct CacheFull; + +/// Cache storage for liquid cache. +/// +/// Example: +/// ```rust +/// use liquid_cache::cache::{LiquidCacheBuilder, EntryID}; +/// use arrow::array::UInt64Array; +/// use std::sync::Arc; +/// +/// let storage = LiquidCacheBuilder::new() +/// .with_max_memory_bytes(1024 * 1024) +/// .build(); +/// +/// let entry_id = EntryID::from(0); +/// let arrow_array = Arc::new(UInt64Array::from_iter_values(0..32)); +/// storage.insert(entry_id, arrow_array.clone()).execute().unwrap(); +/// +/// // Get the arrow array back +/// let retrieved = storage.get(&entry_id).read().unwrap(); +/// assert_eq!(retrieved.as_ref(), arrow_array.as_ref()); +/// ``` +#[derive(Debug)] +pub struct LiquidCache { + index: ArtIndex, + config: CacheConfig, + budget: BudgetAccounting, + cache_policy: Box, + squeeze_policy: Box, + observer: Arc, +} + +impl LiquidCache { + /// Return current cache statistics: counts and resource usage. + pub fn stats(&self) -> CacheStats { + // Count entries by format + let total_entries = self.index.entry_count(); + + let mut memory_arrow_entries = 0usize; + let mut memory_liquid_entries = 0usize; + + let mut memory_arrow_bytes = 0usize; + let mut memory_liquid_bytes = 0usize; + + self.index.for_each(|_, batch| match batch { + CacheEntry::MemoryArrow(array) => { + memory_arrow_entries += 1; + memory_arrow_bytes += array.get_array_memory_size(); + } + CacheEntry::MemoryLiquid(array) => { + memory_liquid_entries += 1; + memory_liquid_bytes += array.get_array_memory_size(); + } + }); + + let memory_usage_bytes = self.budget.memory_usage_bytes(); + let runtime = self.observer.runtime_snapshot(); + + CacheStats { + total_entries, + memory_arrow_entries, + memory_liquid_entries, + memory_arrow_bytes, + memory_liquid_bytes, + memory_usage_bytes, + max_memory_bytes: self.budget.max_memory_bytes(), + runtime, + } + } + + /// Insert a batch into the cache. + pub fn insert<'a>(&'a self, entry_id: EntryID, batch_to_cache: ArrayRef) -> Insert<'a> { + Insert::new(self, entry_id, batch_to_cache) + } + + /// Create a [`Get`] builder for the provided entry. + pub fn get<'a>(&'a self, entry_id: &'a EntryID) -> Get<'a> { + Get::new(self, entry_id) + } + + /// Create an [`EvaluatePredicate`] builder for evaluating predicates on cached data. + pub fn eval_predicate<'a>( + &'a self, + entry_id: &'a EntryID, + predicate: &'a LiquidExpr, + ) -> EvaluatePredicate<'a> { + EvaluatePredicate::new(self, entry_id, predicate) + } + + /// Try to read a liquid array from the cache. + /// Returns None if the cached data is not in liquid format. + pub fn try_read_liquid( + &self, + entry_id: &EntryID, + ) -> Option { + self.observer.on_try_read_liquid(); + let batch = self.index.get(entry_id)?; + self.cache_policy + .notify_access(entry_id, CachedBatchType::from(batch.as_ref())); + + match batch.as_ref() { + CacheEntry::MemoryLiquid(array) => Some(array.clone()), + CacheEntry::MemoryArrow(_) => None, + } + } + + /// Iterate over all entries in the cache. + /// No guarantees are made about the order of the entries. + /// Isolation level: read-committed + pub fn for_each_entry(&self, mut f: impl FnMut(&EntryID, &CacheEntry)) { + self.index.for_each(&mut f); + } + + /// Reset the cache. + pub fn reset(&self) { + self.index.reset(); + self.budget.reset_usage(); + } + + /// Check if a batch is cached. + pub fn is_cached(&self, entry_id: &EntryID) -> bool { + self.index.is_cached(entry_id) + } + + /// Get the config of the cache. + pub fn config(&self) -> &CacheConfig { + &self.config + } + + /// Get the budget of the cache. + pub fn budget(&self) -> &BudgetAccounting { + &self.budget + } + + /// Access the cache observer (runtime stats). + pub fn observer(&self) -> &Observer { + &self.observer + } +} + +impl LiquidCache { + /// Insert a batch into the cache, it will run cache replacement policy until the batch is inserted. + pub(crate) fn insert_inner( + &self, + entry_id: EntryID, + mut batch_to_cache: CacheEntry, + ) -> Result<(), CacheFull> { + loop { + let Err(not_inserted) = self.try_insert(entry_id, batch_to_cache) else { + return Ok(()); + }; + + let victims = self.cache_policy.find_memory_victim(8); + if victims.is_empty() { + // No advice, because the cache is already empty: + // the entry to be inserted does not fit in memory at all. + return Err(CacheFull); + } + self.squeeze_victims(victims); + + batch_to_cache = not_inserted; + } + } + + /// Create a new instance of CacheStorage. + pub(crate) fn new( + batch_size: usize, + max_memory_bytes: usize, + squeeze_policy: Box, + cache_policy: Box, + ) -> Self { + let config = CacheConfig::new(batch_size, max_memory_bytes); + let observer = Arc::new(Observer::new()); + Self { + index: ArtIndex::new(), + budget: BudgetAccounting::new(config.max_memory_bytes()), + config, + cache_policy, + squeeze_policy, + observer, + } + } + + fn try_insert(&self, entry_id: EntryID, to_insert: CacheEntry) -> Result<(), CacheEntry> { + let new_memory_size = to_insert.memory_usage_bytes(); + let cached_batch_type = if let Some(entry) = self.index.get(&entry_id) { + let old_memory_size = entry.memory_usage_bytes(); + if self + .budget + .try_update_memory_usage(old_memory_size, new_memory_size) + .is_err() + { + return Err(to_insert); + } + let batch_type = CachedBatchType::from(&to_insert); + self.index.insert(&entry_id, to_insert); + batch_type + } else { + if self.budget.try_reserve_memory(new_memory_size).is_err() { + return Err(to_insert); + } + let batch_type = CachedBatchType::from(&to_insert); + self.index.insert(&entry_id, to_insert); + batch_type + }; + + self.cache_policy + .notify_insert(&entry_id, cached_batch_type); + + Ok(()) + } + + fn remove_memory_entry(&self, entry_id: EntryID) { + let Some(removed) = self.index.remove(&entry_id) else { + return; + }; + self.budget + .try_update_memory_usage(removed.memory_usage_bytes(), 0) + .expect("memory release cannot fail"); + self.cache_policy.notify_remove(&entry_id); + self.observer.on_memory_eviction(); + } + + /// Get the index of the cache. + #[cfg(test)] + pub(crate) fn index(&self) -> &ArtIndex { + &self.index + } + + fn squeeze_victims(&self, victims: Vec) { + for victim in victims { + self.squeeze_victim_inner(victim); + } + } + + fn squeeze_victim_inner(&self, to_squeeze: EntryID) { + let Some(mut to_squeeze_batch) = self.index.get(&to_squeeze) else { + return; + }; + + loop { + let outcome = self.squeeze_policy.squeeze(to_squeeze_batch.as_ref()); + + match outcome { + SqueezeOutcome::Replace(new_batch) => { + match self.try_insert(to_squeeze, new_batch) { + Ok(()) => { + self.observer.on_transcode(); + break; + } + Err(batch) => { + // The replacement did not fit either; squeeze it further. + // The squeeze policies guarantee this converges to `Remove`. + to_squeeze_batch = Arc::new(batch); + } + } + } + SqueezeOutcome::Remove => { + self.remove_memory_entry(to_squeeze); + break; + } + } + } + } + + pub(crate) fn read_arrow_array( + &self, + entry_id: &EntryID, + selection: Option<&BooleanBuffer>, + ) -> Option { + use arrow::array::BooleanArray; + + let batch = self.index.get(entry_id)?; + self.cache_policy + .notify_access(entry_id, CachedBatchType::from(batch.as_ref())); + + match batch.as_ref() { + CacheEntry::MemoryArrow(array) => match selection { + Some(selection) => { + let selection_array = BooleanArray::new(selection.clone(), None); + arrow::compute::filter(array, &selection_array).ok() + } + None => Some(array.clone()), + }, + CacheEntry::MemoryLiquid(array) => match selection { + Some(selection) => Some(array.filter(selection)), + None => Some(array.to_arrow_array()), + }, + } + } + + pub(crate) fn eval_predicate_internal( + &self, + entry_id: &EntryID, + selection_opt: Option<&BooleanBuffer>, + predicate: &LiquidExpr, + ) -> Option { + use arrow::array::BooleanArray; + + self.observer.on_eval_predicate(); + let batch = self.index.get(entry_id)?; + self.cache_policy + .notify_access(entry_id, CachedBatchType::from(batch.as_ref())); + + match batch.as_ref() { + CacheEntry::MemoryArrow(array) => { + let mut owned = None; + let selection = selection_opt.unwrap_or_else(|| { + owned = Some(BooleanBuffer::new_set(array.len())); + owned.as_ref().unwrap() + }); + let selection_array = BooleanArray::new(selection.clone(), None); + let filtered = arrow::compute::filter(array, &selection_array) + .expect("selection must match array length"); + Some(self.eval_predicate_on_array(filtered, predicate)) + } + CacheEntry::MemoryLiquid(array) => { + let mut owned = None; + let selection = selection_opt.unwrap_or_else(|| { + owned = Some(BooleanBuffer::new_set(array.len())); + owned.as_ref().unwrap() + }); + Some(array.try_eval_predicate(predicate, selection)) + } + } + } + + fn eval_predicate_on_array(&self, array: ArrayRef, predicate: &LiquidExpr) -> BooleanArray { + let schema = Arc::new(Schema::new(vec![Field::new( + "liquid_predicate_col", + array.data_type().clone(), + true, + )])); + let record_batch = + RecordBatch::try_new(schema, vec![array]).expect("single-column predicate batch"); + let result = predicate + .physical_expr() + .evaluate(&record_batch) + .expect("validated LiquidExpr must evaluate"); + let boolean_array = result + .into_array(record_batch.num_rows()) + .expect("predicate output must be an array"); + boolean_array.as_boolean().clone() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::cache::{ + CacheEntry, CachePolicy, LiquidCacheBuilder, LiquidPolicy, transcode_liquid_inner, + utils::{create_cache_store, create_test_array, create_test_arrow_array}, + }; + use crate::sync::thread; + use arrow::array::{Array, ArrayRef, BooleanArray, Int32Array}; + use datafusion_common::ScalarValue; + use datafusion_expr_common::operator::Operator as DFOperator; + use datafusion_physical_expr::PhysicalExpr; + use datafusion_physical_expr::expressions::{BinaryExpr, Column, Literal}; + use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; + + // Unified advice type for more concise testing + #[derive(Debug)] + struct TestPolicy { + target_id: Option, + advice_count: AtomicUsize, + } + + impl TestPolicy { + fn new(target_id: Option) -> Self { + Self { + target_id, + advice_count: AtomicUsize::new(0), + } + } + } + + impl CachePolicy for TestPolicy { + fn find_memory_victim(&self, _cnt: usize) -> Vec { + self.advice_count.fetch_add(1, Ordering::SeqCst); + let id_to_use = self.target_id.unwrap(); + vec![id_to_use] + } + } + + #[test] + fn test_basic_cache_operations() { + // Test basic insert, get, and size tracking in one test + let budget_size = 10 * 1024; + let store = create_cache_store(budget_size, Box::new(LiquidPolicy::new())); + + // 1. Initial budget should be empty + assert_eq!(store.budget.memory_usage_bytes(), 0); + + // 2. Insert and verify first entry + let entry_id1: EntryID = EntryID::from(1); + let array1 = create_test_array(100); + let size1 = array1.memory_usage_bytes(); + store.insert_inner(entry_id1, array1).unwrap(); + + // Verify budget usage and data correctness + assert_eq!(store.budget.memory_usage_bytes(), size1); + let retrieved1 = store.index().get(&entry_id1).unwrap(); + match retrieved1.as_ref() { + CacheEntry::MemoryArrow(arr) => assert_eq!(arr.len(), 100), + _ => panic!("Expected ArrowMemory"), + } + + let entry_id2: EntryID = EntryID::from(2); + let array2 = create_test_array(200); + let size2 = array2.memory_usage_bytes(); + store.insert_inner(entry_id2, array2).unwrap(); + + assert_eq!(store.budget.memory_usage_bytes(), size1 + size2); + + let array3 = create_test_array(150); + let size3 = array3.memory_usage_bytes(); + store.insert_inner(entry_id1, array3).unwrap(); + + assert_eq!(store.budget.memory_usage_bytes(), size3 + size2); + assert!(store.index().get(&EntryID::from(999)).is_none()); + } + + #[test] + fn test_cache_advice_strategies() { + // Under memory pressure, the advised victim is transcoded to liquid. + let entry_id1 = EntryID::from(1); + let entry_id2 = EntryID::from(2); + + let advisor = TestPolicy::new(Some(entry_id1)); + let store = create_cache_store(8000, Box::new(advisor)); // Small budget to force advice + + store + .insert_inner(entry_id1, create_test_array(800)) + .unwrap(); + match store.index().get(&entry_id1).unwrap().as_ref() { + CacheEntry::MemoryArrow(_) => {} + other => panic!("Expected ArrowMemory, got {other:?}"), + } + + store + .insert_inner(entry_id2, create_test_array(800)) + .unwrap(); + match store.index().get(&entry_id1).unwrap().as_ref() { + CacheEntry::MemoryLiquid(_) => {} + other => panic!("Expected LiquidMemory after eviction, got {other:?}"), + } + } + + #[test] + fn test_concurrent_cache_operations() { + concurrent_cache_operations(); + } + + fn concurrent_cache_operations() { + let num_threads = 3; + let ops_per_thread = 50; + + // Large enough that no entry is evicted: with no disk tier, evicted + // entries are gone, so the retrievability invariant below requires + // every entry to fit in memory. + let budget_size = + num_threads * ops_per_thread * create_test_arrow_array(100).get_array_memory_size() * 2; + let store = create_cache_store(budget_size, Box::new(LiquidPolicy::new())); + + let mut handles = vec![]; + for thread_id in 0..num_threads { + let store = store.clone(); + handles.push(thread::spawn(move || { + for i in 0..ops_per_thread { + let unique_id = thread_id * ops_per_thread + i; + let entry_id: EntryID = EntryID::from(unique_id); + let array = create_test_arrow_array(100); + store.insert(entry_id, array).execute().unwrap(); + } + })); + } + for handle in handles { + handle.join().unwrap(); + } + + // Invariant 1: Every previously inserted entry can be retrieved + for thread_id in 0..num_threads { + for i in 0..ops_per_thread { + let unique_id = thread_id * ops_per_thread + i; + let entry_id: EntryID = EntryID::from(unique_id); + assert!(store.index().get(&entry_id).is_some()); + } + } + + // Invariant 2: Number of entries matches number of insertions + assert_eq!(store.index().keys().len(), num_threads * ops_per_thread); + } + + #[test] + fn test_cache_stats_memory_usage() { + let storage = LiquidCacheBuilder::new() + .with_max_memory_bytes(10 * 1024 * 1024) + .build(); + + // Insert two small batches + let arr1: ArrayRef = Arc::new(Int32Array::from_iter_values(0..64)); + let arr2: ArrayRef = Arc::new(Int32Array::from_iter_values(0..128)); + storage + .insert(EntryID::from(1usize), arr1) + .execute() + .unwrap(); + storage + .insert(EntryID::from(2usize), arr2) + .execute() + .unwrap(); + + // Stats after insert: 2 entries, memory usage > 0 + let s = storage.stats(); + assert_eq!(s.total_entries, 2); + assert_eq!(s.memory_arrow_entries, 2); + assert_eq!(s.memory_liquid_entries, 0); + assert!(s.memory_usage_bytes > 0); + assert_eq!(s.max_memory_bytes, 10 * 1024 * 1024); + } + + #[test] + fn insert_returns_cache_full_when_memory_is_saturated() { + let cache = LiquidCacheBuilder::new().with_max_memory_bytes(0).build(); + let array: ArrayRef = Arc::new(Int32Array::from_iter_values(0..16)); + + let err = cache.insert(EntryID::from(900usize), array).execute(); + + assert_eq!(err, Err(CacheFull)); + assert!(!cache.is_cached(&EntryID::from(900usize))); + } + + #[test] + fn eviction_under_memory_pressure_keeps_newest_entries() { + // Budget fits roughly one arrow entry; older entries must be transcoded or evicted. + let first_array: ArrayRef = Arc::new(Int32Array::from_iter_values(0..1024)); + let entry_size = first_array.get_array_memory_size(); + let cache = LiquidCacheBuilder::new() + .with_max_memory_bytes(entry_size + entry_size / 2) + .with_cache_policy(Box::new(LiquidPolicy::new())) + .build(); + + let first = EntryID::from(910usize); + let second = EntryID::from(911usize); + cache.insert(first, first_array).execute().unwrap(); + assert!(cache.is_cached(&first)); + + let second_array: ArrayRef = Arc::new(Int32Array::from_iter_values(1024..2048)); + cache.insert(second, second_array).execute().unwrap(); + assert!(cache.is_cached(&second)); + assert!(cache.budget().memory_usage_bytes() <= cache.budget().max_memory_bytes()); + } + + #[test] + fn eviction_releases_budget() { + let cache = LiquidCacheBuilder::new() + .with_max_memory_bytes(1 << 20) + .build(); + let entry = EntryID::from(914usize); + let array: ArrayRef = Arc::new(Int32Array::from_iter_values(0..16)); + cache.insert(entry, array).execute().unwrap(); + let before = cache.stats().memory_usage_bytes; + assert!(before > 0); + + cache.remove_memory_entry(entry); + + assert_eq!(cache.stats().memory_usage_bytes, 0); + assert!(!cache.is_cached(&entry)); + } + + #[test] + fn get_with_selection_filters_rows() { + let cache = LiquidCacheBuilder::new() + .with_max_memory_bytes(1 << 20) + .build(); + let entry = EntryID::from(915usize); + let array: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3, 4])); + cache.insert(entry, array).execute().unwrap(); + + let selection = arrow::buffer::BooleanBuffer::from(vec![true, false, true, false]); + let result = cache + .get(&entry) + .with_selection(&selection) + .read() + .expect("present"); + let expected = Int32Array::from(vec![1, 3]); + assert_eq!(result.as_ref(), &expected as &dyn Array); + } + + #[test] + fn eval_predicate_on_arrow_and_liquid_entries() { + let cache = LiquidCacheBuilder::new() + .with_max_memory_bytes(1 << 20) + .build(); + let arrow_entry = EntryID::from(916usize); + let liquid_entry = EntryID::from(917usize); + + let array: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3, 4])); + cache.insert(arrow_entry, array.clone()).execute().unwrap(); + + let liquid = transcode_liquid_inner(&array).unwrap(); + cache + .insert_inner(liquid_entry, CacheEntry::memory_liquid(liquid)) + .unwrap(); + + let expr: Arc = Arc::new(BinaryExpr::new( + Arc::new(Column::new("liquid_predicate_col", 0)), + DFOperator::Gt, + Arc::new(Literal::new(ScalarValue::Int32(Some(2)))), + )); + let predicate = LiquidExpr::new_unchecked(expr); + let expected = BooleanArray::from(vec![false, false, true, true]); + + for entry in [arrow_entry, liquid_entry] { + let got = cache + .eval_predicate(&entry, &predicate) + .read() + .expect("entry present"); + assert_eq!(got, expected); + } + } + + #[test] + fn try_read_liquid_returns_only_liquid_entries() { + let cache = LiquidCacheBuilder::new() + .with_max_memory_bytes(1 << 20) + .build(); + let arrow_entry = EntryID::from(918usize); + let liquid_entry = EntryID::from(919usize); + + let array: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3])); + cache.insert(arrow_entry, array.clone()).execute().unwrap(); + let liquid = transcode_liquid_inner(&array).unwrap(); + cache + .insert_inner(liquid_entry, CacheEntry::memory_liquid(liquid)) + .unwrap(); + + assert!(cache.try_read_liquid(&arrow_entry).is_none()); + let read = cache.try_read_liquid(&liquid_entry).expect("liquid entry"); + assert_eq!(read.to_arrow_array().as_ref(), array.as_ref()); + } + + #[test] + fn reset_clears_entries_and_budget() { + let cache = LiquidCacheBuilder::new() + .with_max_memory_bytes(1 << 20) + .build(); + let entry = EntryID::from(920usize); + let array: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3])); + cache.insert(entry, array).execute().unwrap(); + assert!(cache.is_cached(&entry)); + + cache.reset(); + + assert!(!cache.is_cached(&entry)); + assert_eq!(cache.budget().memory_usage_bytes(), 0); + assert_eq!(cache.stats().total_entries, 0); + } +} diff --git a/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/cache/index.rs b/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/cache/index.rs new file mode 100644 index 0000000000000..a25fec751ac6c --- /dev/null +++ b/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/cache/index.rs @@ -0,0 +1,146 @@ +use congee::CongeeArc; +use std::{ + fmt::{Debug, Formatter}, + sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, + }, +}; + +use crate::cache::{cached_batch::CacheEntry, utils::EntryID}; + +pub(crate) struct ArtIndex { + art: CongeeArc, + entry_count: AtomicUsize, +} + +impl Debug for ArtIndex { + fn fmt(&self, _f: &mut Formatter<'_>) -> std::fmt::Result { + Ok(()) + } +} + +impl ArtIndex { + pub(crate) fn new() -> Self { + let art: CongeeArc = CongeeArc::new(); + Self { + art, + entry_count: AtomicUsize::new(0), + } + } + + pub(crate) fn get(&self, entry_id: &EntryID) -> Option> { + let guard = self.art.pin(); + let batch = self.art.get(*entry_id, &guard)?; + Some(batch) + } + + pub(crate) fn is_cached(&self, entry_id: &EntryID) -> bool { + let guard = self.art.pin(); + self.art.get(*entry_id, &guard).is_some() + } + + pub(crate) fn insert(&self, entry_id: &EntryID, batch: CacheEntry) { + let guard = self.art.pin(); + let existing = self + .art + .insert(*entry_id, Arc::new(batch), &guard) + .expect("Insertion failed"); + if existing.is_none() { + self.entry_count.fetch_add(1, Ordering::Relaxed); + } + } + + pub(crate) fn remove(&self, entry_id: &EntryID) -> Option> { + let guard = self.art.pin(); + let removed = self.art.remove(*entry_id, &guard); + if removed.is_some() { + self.entry_count.fetch_sub(1, Ordering::Relaxed); + } + removed + } + + pub(crate) fn reset(&self) { + let guard = self.art.pin(); + self.art.keys().into_iter().for_each(|k| { + _ = self.art.remove(k, &guard).unwrap(); + }); + self.entry_count.store(0, Ordering::Relaxed); + } + + pub(crate) fn for_each(&self, mut f: impl FnMut(&EntryID, &CacheEntry)) { + let guard = self.art.pin(); + for id in self.art.keys().into_iter() { + f( + &id, + &self + .art + .get(id, &guard) + .expect("Failed to get value from ART"), + ); + } + } + + #[cfg(test)] + pub(crate) fn keys(&self) -> Vec { + self.art.keys() + } + + pub(crate) fn entry_count(&self) -> usize { + self.entry_count.load(Ordering::Relaxed) + } +} + +#[cfg(test)] +mod tests { + use crate::cache::cached_batch::CacheEntry; + use crate::cache::utils::create_test_array; + + use super::*; + + #[test] + fn test_get_and_is_cached() { + let store = ArtIndex::new(); + let entry_id1: EntryID = EntryID::from(1); + let entry_id2: EntryID = EntryID::from(2); + let array1 = create_test_array(100); + + // Initially, entries should not be cached + assert!(!store.is_cached(&entry_id1)); + assert!(!store.is_cached(&entry_id2)); + assert!(store.get(&entry_id1).is_none()); + + // Insert an entry and verify it's cached + { + store.insert(&entry_id1, array1.clone()); + } + + assert!(store.is_cached(&entry_id1)); + assert!(!store.is_cached(&entry_id2)); + + // Get should return the cached value + match store.get(&entry_id1) { + Some(batch) => match batch.as_ref() { + CacheEntry::MemoryArrow(arr) => assert_eq!(arr.len(), 100), + _ => panic!("Expected ArrowMemory batch"), + }, + None => panic!("Expected ArrowMemory batch"), + } + } + + #[test] + fn test_reset() { + let store = ArtIndex::new(); + let entry_id: EntryID = EntryID::from(1); + let array = create_test_array(100); + + store.insert(&entry_id, array.clone()); + + let entry_id: EntryID = EntryID::from(1); + assert!(store.is_cached(&entry_id)); + + store.reset(); + let entry_id: EntryID = EntryID::from(1); + assert!(!store.is_cached(&entry_id)); + } +} diff --git a/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/cache/liquid_expr.rs b/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/cache/liquid_expr.rs new file mode 100644 index 0000000000000..0c808c198a33a --- /dev/null +++ b/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/cache/liquid_expr.rs @@ -0,0 +1,196 @@ +use arrow_schema::DataType; +use datafusion_common::ScalarValue; +use datafusion_expr_common::operator::Operator; +use datafusion_physical_expr::expressions::{ + BinaryExpr, CastExpr, Column, DynamicFilterPhysicalExpr, Literal, TryCastExpr, +}; +use datafusion_physical_expr::{PhysicalExpr, ScalarFunctionExpr}; + +use crate::sync::Arc; +use crate::utils::get_bytes_needle; + +/// A predicate expression validated for LiquidCache predicate evaluation. +#[derive(Clone)] +pub struct LiquidExpr { + expr: Arc, +} + +impl std::fmt::Debug for LiquidExpr { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("LiquidExpr") + .field("expr", &self.expr.to_string()) + .finish() + } +} + +impl LiquidExpr { + /// Validate and wrap a physical expression for LiquidCache predicate evaluation. + /// + /// Returns `None` when the expression shape or operator is unsupported for the + /// provided column type. + pub fn try_new(expr: Arc, data_type: &DataType) -> Option { + let normalized = unwrap_dynamic_filter(&expr)?; + if supports_expr(&normalized, data_type) { + Some(Self { expr: normalized }) + } else { + None + } + } + + /// Get the underlying validated physical expression. + pub fn physical_expr(&self) -> &Arc { + &self.expr + } + + #[cfg(test)] + pub(crate) fn new_unchecked(expr: Arc) -> Self { + Self { expr } + } +} + +fn unwrap_dynamic_filter(expr: &Arc) -> Option> { + if let Some(dynamic_filter) = expr.downcast_ref::() { + dynamic_filter.current().ok() + } else { + Some(expr.clone()) + } +} + +fn supports_expr(expr: &Arc, data_type: &DataType) -> bool { + if let Some(binary) = expr.downcast_ref::() { + return supports_binary_expr(binary, data_type); + } + + if let Some(literal) = expr.downcast_ref::() { + return matches!(literal.value(), ScalarValue::Boolean(Some(_))) && is_byte_like(data_type); + } + + false +} + +fn supports_binary_expr(binary: &BinaryExpr, data_type: &DataType) -> bool { + let Some(literal) = binary.right().downcast_ref::() else { + return false; + }; + let op = binary.op(); + if is_byte_like(data_type) { + if !is_column_like(binary.left()) { + return false; + } + + match op { + Operator::Eq + | Operator::NotEq + | Operator::Lt + | Operator::LtEq + | Operator::Gt + | Operator::GtEq => get_bytes_needle(literal.value()).is_some(), + _ => false, + } + } else if is_numeric_like(data_type) { + matches!( + op, + Operator::Eq + | Operator::NotEq + | Operator::Lt + | Operator::LtEq + | Operator::Gt + | Operator::GtEq + ) && (is_column_like(binary.left()) || is_to_timestamp_seconds_column(binary.left())) + } else { + false + } +} + +fn is_column_like(expr: &Arc) -> bool { + if expr.downcast_ref::().is_some() { + return true; + } + if let Some(cast_expr) = expr.downcast_ref::() { + return is_column_like(cast_expr.expr()); + } + if let Some(try_cast_expr) = expr.downcast_ref::() { + return is_column_like(try_cast_expr.expr()); + } + false +} + +fn is_to_timestamp_seconds_column(expr: &Arc) -> bool { + if let Some(func) = expr.downcast_ref::() + && func.name() == "to_timestamp_seconds" + && let [arg] = func.args() + { + return is_column_like(arg); + } + false +} + +fn is_byte_like(data_type: &DataType) -> bool { + match data_type { + DataType::Utf8 | DataType::Utf8View | DataType::Binary | DataType::BinaryView => true, + DataType::Dictionary(_, value_type) => is_byte_like(value_type.as_ref()), + _ => false, + } +} + +fn is_numeric_like(data_type: &DataType) -> bool { + matches!( + data_type, + DataType::Int8 + | DataType::Int16 + | DataType::Int32 + | DataType::Int64 + | DataType::UInt8 + | DataType::UInt16 + | DataType::UInt32 + | DataType::UInt64 + | DataType::Float32 + | DataType::Float64 + | DataType::Date32 + | DataType::Date64 + | DataType::Decimal128(_, _) + | DataType::Decimal256(_, _) + ) || matches!(data_type, DataType::Timestamp(_, None)) +} + +#[cfg(test)] +mod tests { + use super::*; + use datafusion_common::ScalarValue; + use datafusion_physical_expr::expressions::{BinaryExpr, Column}; + + #[test] + fn validates_byte_comparison_with_literal() { + let expr: Arc = Arc::new(BinaryExpr::new( + Arc::new(Column::new("c", 0)), + Operator::Eq, + Arc::new(Literal::new(ScalarValue::Utf8(Some("x".to_string())))), + )); + let liquid_expr = LiquidExpr::try_new(expr, &DataType::Utf8); + assert!(liquid_expr.is_some()); + } + + #[test] + fn validates_numeric_comparison() { + let expr: Arc = Arc::new(BinaryExpr::new( + Arc::new(Column::new("c", 0)), + Operator::Gt, + Arc::new(Literal::new(ScalarValue::Int32(Some(42)))), + )); + let liquid_expr = LiquidExpr::try_new(expr, &DataType::Int32); + assert!(liquid_expr.is_some()); + } + + #[test] + fn rejects_unsupported_like_expression() { + use datafusion_physical_expr::expressions::LikeExpr; + let expr: Arc = Arc::new(LikeExpr::new( + false, + false, + Arc::new(Column::new("c", 0)), + Arc::new(Literal::new(ScalarValue::Utf8(Some("%abc%".to_string())))), + )); + let liquid_expr = LiquidExpr::try_new(expr, &DataType::Utf8); + assert!(liquid_expr.is_none()); + } +} diff --git a/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/cache/mod.rs b/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/cache/mod.rs new file mode 100644 index 0000000000000..4ae208e5bae1e --- /dev/null +++ b/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/cache/mod.rs @@ -0,0 +1,35 @@ +//! Cache layer for liquid cache. + +mod budget; +mod builders; +mod cached_batch; +mod core; +mod index; +mod liquid_expr; +mod observer; +pub mod policies; +mod transcode; +mod utils; + +pub use builders::{EvaluatePredicate, Get, Insert, LiquidCacheBuilder}; +pub use cached_batch::{CacheEntry, CachedBatchType}; +pub use core::{CacheFull, LiquidCache}; +pub use liquid_expr::LiquidExpr; +pub use observer::Observer; +pub use observer::{CacheStats, RuntimeStats, RuntimeStatsSnapshot}; +pub use policies::{ + CachePolicy, Evict, LiquidPolicy, LruPolicy, SqueezeOutcome, SqueezePolicy, TranscodeEvict, +}; +pub use transcode::transcode_liquid_inner; +pub use utils::EntryID; + +// Backwards-compatible module paths for existing imports. +/// Legacy path: re-export cache policy types under `cache::cache_policies`. +pub mod cache_policies { + pub use super::policies::cache::*; +} + +/// Legacy path: re-export squeeze policy types under `cache::squeeze_policies`. +pub mod squeeze_policies { + pub use super::policies::squeeze::*; +} diff --git a/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/cache/observer/mod.rs b/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/cache/observer/mod.rs new file mode 100644 index 0000000000000..07ab4e1c208bd --- /dev/null +++ b/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/cache/observer/mod.rs @@ -0,0 +1,64 @@ +mod stats; + +pub use stats::{CacheStats, RuntimeStats, RuntimeStatsSnapshot}; + +use stats::{RuntimeStats as RuntimeStatsInner, RuntimeStatsSnapshot as RuntimeStatsSnapshotInner}; + +#[derive(Debug)] +/// Cache-side observer for runtime stats. +pub struct Observer { + runtime: RuntimeStatsInner, +} + +impl Default for Observer { + fn default() -> Self { + Self::new() + } +} + +impl Observer { + /// Create a new observer with all counters reset. + pub fn new() -> Self { + Self { + runtime: RuntimeStatsInner::default(), + } + } + + /// Snapshot runtime counters and reset them to zero. + pub fn runtime_snapshot(&self) -> RuntimeStatsSnapshotInner { + self.runtime.consume_snapshot() + } + + #[inline] + pub(crate) fn on_get(&self, selection: bool) { + self.runtime.incr_get(); + if selection { + self.runtime.incr_get_with_selection(); + } + } + + #[inline] + pub(crate) fn on_try_read_liquid(&self) { + self.runtime.incr_try_read_liquid(); + } + + #[inline] + pub(crate) fn on_eval_predicate(&self) { + self.runtime.incr_eval_predicate(); + } + + #[inline] + pub(crate) fn on_memory_eviction(&self) { + self.runtime.incr_memory_evictions(); + } + + #[inline] + pub(crate) fn on_transcode(&self) { + self.runtime.incr_transcodes(); + } + + /// Access the underlying runtime statistics counters. + pub fn runtime_stats(&self) -> &RuntimeStats { + &self.runtime + } +} diff --git a/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/cache/observer/stats.rs b/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/cache/observer/stats.rs new file mode 100644 index 0000000000000..d4b39dfbcbfcf --- /dev/null +++ b/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/cache/observer/stats.rs @@ -0,0 +1,127 @@ +use std::fmt; +use std::sync::atomic::{AtomicU64, Ordering}; + +/// Macro to define runtime statistics metrics. +/// +/// Usage: +/// ```ignore +/// define_runtime_stats! { +/// (field_name, "doc comment", method_name), +/// ... +/// } +/// ``` +/// +/// This generates: +/// - Fields in `RuntimeStats` struct +/// - Fields in `RuntimeStatsSnapshot` struct +/// - Increment methods (`incr_*`) +/// - `consume_snapshot` implementation +/// - `reset` implementation +macro_rules! define_runtime_stats { + ( + $( + ($field:ident, $doc:literal, $method:ident) + ),* $(,)? + ) => { + /// Atomic runtime counters for cache API calls. + #[derive(Debug, Default)] + pub struct RuntimeStats { + $( + #[doc = $doc] + pub(crate) $field: AtomicU64, + )* + } + + /// Immutable snapshot of [`RuntimeStats`]. + #[derive(Debug, Clone, Copy, serde::Serialize, serde::Deserialize)] + pub struct RuntimeStatsSnapshot { + $( + #[doc = concat!("Total ", stringify!($field), ".")] + pub $field: u64, + )* + } + + impl RuntimeStats { + /// Return an immutable snapshot of the current runtime counters and reset the stats to 0. + pub fn consume_snapshot(&self) -> RuntimeStatsSnapshot { + let v = RuntimeStatsSnapshot { + $( + $field: self.$field.load(Ordering::Relaxed), + )* + }; + self.reset(); + v + } + + $( + /// Increment counter. + #[inline] + pub fn $method(&self) { + self.$field.fetch_add(1, Ordering::Relaxed); + } + )* + + /// Reset the runtime stats to 0. + pub fn reset(&self) { + $( + self.$field.store(0, Ordering::Relaxed); + )* + } + } + + impl fmt::Display for RuntimeStats { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + writeln!(f, "RuntimeStats:")?; + $( + writeln!(f, " {}: {}", stringify!($field), self.$field.load(Ordering::Relaxed))?; + )* + Ok(()) + } + } + + impl fmt::Display for RuntimeStatsSnapshot { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + writeln!(f, "RuntimeStatsSnapshot:")?; + $( + writeln!(f, " {}: {}", stringify!($field), self.$field)?; + )* + Ok(()) + } + } + }; +} + +// Define all runtime statistics metrics here. +// To add a new metric, add a line: (field_name, "doc comment", method_name) +define_runtime_stats! { + (get, "Number of `get` calls issued via `CachedData`.", incr_get), + (get_with_selection, "Number of `get_with_selection` calls issued via `CachedData`.", incr_get_with_selection), + (eval_predicate, "Number of `eval_predicate` calls issued via `CachedData`.", incr_eval_predicate), + (cache_hit, "Number of cache hits (data found in cache).", incr_cache_hit), + (cache_miss, "Number of cache misses (data not in cache, fell back to Parquet).", incr_cache_miss), + (try_read_liquid_calls, "Number of `try_read_liquid` calls issued via `CachedData`.", incr_try_read_liquid), + (eval_predicate_on_liquid_failed, "Number of `eval_predicate` calls that failed on Liquid array.", incr_eval_predicate_on_liquid_failed), + (memory_evictions, "Number of cache entries evicted from memory.", incr_memory_evictions), + (transcodes, "Number of Arrow entries transcoded to Liquid under memory pressure.", incr_transcodes), +} + +/// Snapshot of cache statistics. +#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] +pub struct CacheStats { + /// Total number of entries in the cache. + pub total_entries: usize, + /// Number of in-memory Arrow entries. + pub memory_arrow_entries: usize, + /// Number of in-memory Liquid entries. + pub memory_liquid_entries: usize, + /// Total size of in-memory Arrow entries in bytes. + pub memory_arrow_bytes: usize, + /// Total size of in-memory Liquid entries in bytes. + pub memory_liquid_bytes: usize, + /// Total memory usage of the cache. + pub memory_usage_bytes: usize, + /// Maximum memory size. + pub max_memory_bytes: usize, + /// Runtime counters snapshot. + pub runtime: RuntimeStatsSnapshot, +} diff --git a/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/cache/policies/cache/doubly_linked_list.rs b/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/cache/policies/cache/doubly_linked_list.rs new file mode 100644 index 0000000000000..1a48caafdb26b --- /dev/null +++ b/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/cache/policies/cache/doubly_linked_list.rs @@ -0,0 +1,125 @@ +use std::ptr::NonNull; + +/// Intrusive doubly linked list node used by cache policies. +#[derive(Debug)] +pub(crate) struct DoublyLinkedNode { + pub(crate) data: T, + pub(crate) prev: Option>, + pub(crate) next: Option>, +} + +impl DoublyLinkedNode { + pub(crate) fn new(data: T) -> Box { + Box::new(Self { + data, + prev: None, + next: None, + }) + } +} + +/// Intrusive doubly linked list utility shared across cache policies. +#[derive(Debug)] +pub(crate) struct DoublyLinkedList { + head: Option>>, + tail: Option>>, +} + +impl Default for DoublyLinkedList { + fn default() -> Self { + Self::new() + } +} + +impl DoublyLinkedList { + pub(crate) fn new() -> Self { + Self { + head: None, + tail: None, + } + } + + pub(crate) fn head(&self) -> Option>> { + self.head + } + + #[allow(dead_code)] + pub(crate) fn tail(&self) -> Option>> { + self.tail + } + + /// Inserts the node at the front (head) of the list. + pub(crate) unsafe fn push_front(&mut self, mut node_ptr: NonNull>) { + let node = unsafe { node_ptr.as_mut() }; + node.prev = None; + node.next = self.head; + + if let Some(mut head) = self.head { + unsafe { head.as_mut().prev = Some(node_ptr) }; + } else { + self.tail = Some(node_ptr); + } + + self.head = Some(node_ptr); + } + + /// Inserts the node at the back (tail) of the list. + pub(crate) unsafe fn push_back(&mut self, mut node_ptr: NonNull>) { + let node = unsafe { node_ptr.as_mut() }; + node.next = None; + node.prev = self.tail; + + if let Some(mut tail) = self.tail { + unsafe { tail.as_mut().next = Some(node_ptr) }; + } else { + self.head = Some(node_ptr); + } + + self.tail = Some(node_ptr); + } + + /// Moves an existing node to the front of the list. + #[allow(dead_code)] + pub(crate) unsafe fn move_to_front(&mut self, node_ptr: NonNull>) { + unsafe { + self.unlink(node_ptr); + self.push_front(node_ptr); + } + } + + /// Unlinks `node_ptr` from the list without deallocating it. + pub(crate) unsafe fn unlink(&mut self, mut node_ptr: NonNull>) { + let node = unsafe { node_ptr.as_mut() }; + + match node.prev { + Some(mut prev) => unsafe { prev.as_mut().next = node.next }, + None => self.head = node.next, + } + + match node.next { + Some(mut next) => unsafe { next.as_mut().prev = node.prev }, + None => self.tail = node.prev, + } + + node.prev = None; + node.next = None; + } + + /// Drops all nodes currently owned by the list. + pub(crate) unsafe fn drop_all(&mut self) { + let mut current = self.head; + while let Some(node_ptr) = current { + current = unsafe { node_ptr.as_ref().next }; + unsafe { + drop(Box::from_raw(node_ptr.as_ptr())); + } + } + self.head = None; + self.tail = None; + } +} + +/// Drops the boxed node referenced by `ptr`. +pub(crate) unsafe fn drop_boxed_node(ptr: NonNull>) { + unsafe { drop(Box::from_raw(ptr.as_ptr())) } +} diff --git a/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/cache/policies/cache/lru.rs b/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/cache/policies/cache/lru.rs new file mode 100644 index 0000000000000..0cd8df7d45bc0 --- /dev/null +++ b/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/cache/policies/cache/lru.rs @@ -0,0 +1,246 @@ +//! LRU (Least Recently Used) cache eviction policy. +//! +//! Uses 2 queues by entry type (Arrow, Liquid), same as LiquidPolicy. +//! Within each queue, entries are ordered by recency (moved to back on access). +//! Eviction priority: Arrow (largest) first, then Liquid. +//! Within each queue, evicts the LRU entry (front of the queue). + +use crate::cache::cached_batch::CachedBatchType; +use crate::cache::utils::EntryID; +use crate::sync::Mutex; +use ahash::AHashMap; +use std::ptr::NonNull; + +use super::CachePolicy; +use super::doubly_linked_list::{DoublyLinkedList, DoublyLinkedNode, drop_boxed_node}; + +/// Which queue an entry belongs to. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum LruQueueKind { + Arrow, + Liquid, +} + +/// LRU cache eviction policy with type-aware queues. +/// +/// On insert: entry goes to the back of its type queue (most recently used). +/// On access: entry moves to the back of its type queue. +/// On eviction: picks the LRU entry from Arrow queue first, then Liquid. +#[derive(Debug)] +pub struct LruPolicy { + inner: Mutex, +} + +#[derive(Debug)] +struct LruInner { + arrow: DoublyLinkedList, + liquid: DoublyLinkedList, + /// Maps entry_id → (node_ptr, which queue it's in) + map: AHashMap>, LruQueueKind)>, +} + +// Safety: We control access via Mutex and never expose raw pointers outside. +unsafe impl Send for LruInner {} +unsafe impl Sync for LruInner {} + +impl LruInner { + fn queue_mut(&mut self, kind: LruQueueKind) -> &mut DoublyLinkedList { + match kind { + LruQueueKind::Arrow => &mut self.arrow, + LruQueueKind::Liquid => &mut self.liquid, + } + } + + fn queue_kind_for(batch_type: CachedBatchType) -> LruQueueKind { + match batch_type { + CachedBatchType::MemoryArrow => LruQueueKind::Arrow, + CachedBatchType::MemoryLiquid => LruQueueKind::Liquid, + } + } + + /// Pop the LRU entry (front) from a specific queue. + fn pop_lru(&mut self, kind: LruQueueKind) -> Option { + let list = self.queue_mut(kind); + let node_ptr = list.head()?; + let entry_id = unsafe { node_ptr.as_ref().data }; + unsafe { list.unlink(node_ptr) }; + self.map.remove(&entry_id); + unsafe { drop_boxed_node(node_ptr) }; + Some(entry_id) + } +} + +impl LruPolicy { + /// Create a new LRU policy. + pub fn new() -> Self { + Self { + inner: Mutex::new(LruInner { + arrow: DoublyLinkedList::new(), + liquid: DoublyLinkedList::new(), + map: AHashMap::new(), + }), + } + } +} + +impl Default for LruPolicy { + fn default() -> Self { + Self::new() + } +} + +impl CachePolicy for LruPolicy { + fn find_memory_victim(&self, cnt: usize) -> Vec { + let mut inner = self.inner.lock().unwrap(); + let mut victims = Vec::with_capacity(cnt); + + while victims.len() < cnt { + // Priority: evict Arrow (largest) first, then Liquid + if let Some(entry) = inner.pop_lru(LruQueueKind::Arrow) { + victims.push(entry); + continue; + } + if let Some(entry) = inner.pop_lru(LruQueueKind::Liquid) { + victims.push(entry); + continue; + } + break; + } + + victims + } + + fn notify_insert(&self, entry_id: &EntryID, batch_type: CachedBatchType) { + let mut inner = self.inner.lock().unwrap(); + let target = LruInner::queue_kind_for(batch_type); + + // If already present, remove from old position/queue + if let Some((node_ptr, old_kind)) = inner.map.remove(entry_id) { + let old_list = inner.queue_mut(old_kind); + unsafe { old_list.unlink(node_ptr) }; + unsafe { drop_boxed_node(node_ptr) }; + } + + // Insert at the back of the target queue (most recently used) + let node = DoublyLinkedNode::new(*entry_id); + let node_ptr = NonNull::new(Box::into_raw(node)).unwrap(); + let list = inner.queue_mut(target); + unsafe { list.push_back(node_ptr) }; + inner.map.insert(*entry_id, (node_ptr, target)); + } + + fn notify_access(&self, entry_id: &EntryID, _batch_type: CachedBatchType) { + let mut inner = self.inner.lock().unwrap(); + + let Some(&(node_ptr, kind)) = inner.map.get(entry_id) else { + return; + }; + + // Move to back of its current queue (most recently used) + let list = inner.queue_mut(kind); + unsafe { + list.unlink(node_ptr); + list.push_back(node_ptr); + } + } + + fn notify_remove(&self, entry_id: &EntryID) { + let mut inner = self.inner.lock().unwrap(); + + if let Some((node_ptr, kind)) = inner.map.remove(entry_id) { + let list = inner.queue_mut(kind); + unsafe { list.unlink(node_ptr) }; + unsafe { drop_boxed_node(node_ptr) }; + } + } +} + +impl Drop for LruPolicy { + fn drop(&mut self) { + let mut inner = self.inner.lock().unwrap(); + unsafe { + inner.arrow.drop_all(); + inner.liquid.drop_all(); + } + inner.map.clear(); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn entry(id: usize) -> EntryID { + EntryID::from(id) + } + + #[test] + fn test_lru_evicts_arrow_first() { + let policy = LruPolicy::new(); + + // Insert entries: some Arrow, some Liquid + policy.notify_insert(&entry(1), CachedBatchType::MemoryArrow); + policy.notify_insert(&entry(2), CachedBatchType::MemoryLiquid); + policy.notify_insert(&entry(3), CachedBatchType::MemoryArrow); + policy.notify_insert(&entry(4), CachedBatchType::MemoryLiquid); + + // Evict 1 — should pick from Arrow queue first (LRU = entry 1) + let victims = policy.find_memory_victim(1); + assert_eq!(victims, vec![entry(1)]); + + // Evict 1 more — next Arrow LRU = entry 3 + let victims = policy.find_memory_victim(1); + assert_eq!(victims, vec![entry(3)]); + + // Evict 1 more — Arrow empty, now Liquid LRU = entry 2 + let victims = policy.find_memory_victim(1); + assert_eq!(victims, vec![entry(2)]); + } + + #[test] + fn test_lru_access_moves_to_back() { + let policy = LruPolicy::new(); + + policy.notify_insert(&entry(1), CachedBatchType::MemoryArrow); + policy.notify_insert(&entry(2), CachedBatchType::MemoryArrow); + policy.notify_insert(&entry(3), CachedBatchType::MemoryArrow); + + // Access entry 1 — moves it to back + policy.notify_access(&entry(1), CachedBatchType::MemoryArrow); + + // Evict 1 — LRU is now entry 2 (entry 1 was moved to back) + let victims = policy.find_memory_victim(1); + assert_eq!(victims, vec![entry(2)]); + } + + #[test] + fn test_lru_insert_moves_between_queues() { + let policy = LruPolicy::new(); + + // Insert as Arrow + policy.notify_insert(&entry(1), CachedBatchType::MemoryArrow); + + // Re-insert same entry as Liquid (simulates transcode) + policy.notify_insert(&entry(1), CachedBatchType::MemoryLiquid); + + // Arrow queue should be empty now + let victims = policy.find_memory_victim(1); + // Should come from Liquid queue + assert_eq!(victims, vec![entry(1)]); + } + + #[test] + fn test_lru_remove() { + let policy = LruPolicy::new(); + + policy.notify_insert(&entry(1), CachedBatchType::MemoryArrow); + policy.notify_insert(&entry(2), CachedBatchType::MemoryArrow); + + // Remove entry 1 + policy.notify_remove(&entry(1)); + + // Evict — should get entry 2 (entry 1 was removed) + let victims = policy.find_memory_victim(1); + assert_eq!(victims, vec![entry(2)]); + } +} diff --git a/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/cache/policies/cache/mod.rs b/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/cache/policies/cache/mod.rs new file mode 100644 index 0000000000000..5e2e6f483c653 --- /dev/null +++ b/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/cache/policies/cache/mod.rs @@ -0,0 +1,188 @@ +//! Cache policies for liquid cache. + +use crate::cache::cached_batch::CachedBatchType; +use crate::cache::utils::EntryID; + +mod doubly_linked_list; +mod lru; +mod three_queue; + +pub use lru::LruPolicy; +pub use three_queue::LiquidPolicy; + +/// The cache policy that guides the replacement of LiquidCache +pub trait CachePolicy: std::fmt::Debug + Send + Sync { + /// Give cnt amount of entries to evict when cache is full. + fn find_memory_victim(&self, cnt: usize) -> Vec; + + /// Notify the cache policy that an entry was inserted. + fn notify_insert(&self, _entry_id: &EntryID, _batch_type: CachedBatchType) {} + + /// Notify the cache policy that an entry was accessed. + fn notify_access(&self, _entry_id: &EntryID, _batch_type: CachedBatchType) {} + + /// Notify the cache policy that an entry was removed. + fn notify_remove(&self, _entry_id: &EntryID) {} +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::cache::utils::EntryID; + use crate::sync::{Arc, Mutex, thread}; + + fn entry(id: usize) -> EntryID { + id.into() + } + + fn concurrent_invariant_advice_once(policy: Arc) { + let num_threads = 4; + + for i in 0..100 { + policy.notify_insert(&entry(i), CachedBatchType::MemoryArrow); + } + + let advised_entries = Arc::new(Mutex::new(Vec::new())); + + let mut handles = Vec::new(); + for _ in 0..num_threads { + let policy_clone = policy.clone(); + let advised_entries_clone = advised_entries.clone(); + + let handle = thread::spawn(move || { + let advice = policy_clone.find_memory_victim(1); + if let Some(entry_id) = advice.first() { + let mut entries = advised_entries_clone.lock().unwrap(); + entries.push(*entry_id); + } + }); + + handles.push(handle); + } + + for handle in handles { + handle.join().unwrap(); + } + + let entries = advised_entries.lock().unwrap(); + let mut unique_entries = entries.clone(); + unique_entries.sort(); + unique_entries.dedup(); + + assert_eq!( + entries.len(), + unique_entries.len(), + "Some entries were advised for eviction multiple times: {entries:?}" + ); + } + + fn run_concurrent_invariant_tests() { + concurrent_invariant_advice_once(Arc::new(LiquidPolicy::new())); + concurrent_invariant_advice_once(Arc::new(super::LruPolicy::new())); + } + + #[test] + fn test_concurrent_invariant_advice_once() { + run_concurrent_invariant_tests(); + } + + #[test] + fn lru_evicts_least_recently_used() { + let policy = LruPolicy::new(); + + // Insert entries 0, 1, 2, 3 in order + for i in 0..4 { + policy.notify_insert(&entry(i), CachedBatchType::MemoryArrow); + } + + // Evict 1 → should be entry 0 (oldest) + let victims = policy.find_memory_victim(1); + assert_eq!(victims, vec![entry(0)]); + + // Evict 1 more → should be entry 1 + let victims = policy.find_memory_victim(1); + assert_eq!(victims, vec![entry(1)]); + } + + #[test] + fn lru_access_moves_to_back() { + let policy = LruPolicy::new(); + + // Insert 0, 1, 2 + for i in 0..3 { + policy.notify_insert(&entry(i), CachedBatchType::MemoryArrow); + } + + // Access entry 0 → moves it to back + policy.notify_access(&entry(0), CachedBatchType::MemoryArrow); + + // Evict → should be entry 1 now (0 was moved to back) + let victims = policy.find_memory_victim(1); + assert_eq!(victims, vec![entry(1)]); + + // Evict → entry 2 + let victims = policy.find_memory_victim(1); + assert_eq!(victims, vec![entry(2)]); + + // Evict → entry 0 (was moved to back by access) + let victims = policy.find_memory_victim(1); + assert_eq!(victims, vec![entry(0)]); + } + + #[test] + fn lru_remove_removes_entry() { + let policy = LruPolicy::new(); + + for i in 0..3 { + policy.notify_insert(&entry(i), CachedBatchType::MemoryArrow); + } + + // Remove entry 1 + policy.notify_remove(&entry(1)); + + // Evict → entry 0, then entry 2 (entry 1 is gone) + let victims = policy.find_memory_victim(2); + assert_eq!(victims, vec![entry(0), entry(2)]); + } + + #[test] + fn lru_evict_more_than_available() { + let policy = LruPolicy::new(); + + policy.notify_insert(&entry(0), CachedBatchType::MemoryArrow); + policy.notify_insert(&entry(1), CachedBatchType::MemoryArrow); + + // Ask for 5 but only 2 exist + let victims = policy.find_memory_victim(5); + assert_eq!(victims.len(), 2); + assert_eq!(victims, vec![entry(0), entry(1)]); + + // Empty now + let victims = policy.find_memory_victim(1); + assert!(victims.is_empty()); + } + + #[test] + fn lru_reinsert_updates_position() { + let policy = LruPolicy::new(); + + for i in 0..3 { + policy.notify_insert(&entry(i), CachedBatchType::MemoryArrow); + } + + // Re-insert entry 0 → should move to back + policy.notify_insert(&entry(0), CachedBatchType::MemoryLiquid); + + // Evict → entry 1 (oldest now) + let victims = policy.find_memory_victim(1); + assert_eq!(victims, vec![entry(1)]); + + // Evict → entry 2 + let victims = policy.find_memory_victim(1); + assert_eq!(victims, vec![entry(2)]); + + // Evict → entry 0 (re-inserted last) + let victims = policy.find_memory_victim(1); + assert_eq!(victims, vec![entry(0)]); + } +} diff --git a/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/cache/policies/cache/three_queue.rs b/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/cache/policies/cache/three_queue.rs new file mode 100644 index 0000000000000..2be4e2559a04c --- /dev/null +++ b/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/cache/policies/cache/three_queue.rs @@ -0,0 +1,267 @@ +use std::{collections::HashMap, ptr::NonNull}; + +use crate::{ + cache::{CachePolicy, EntryID, cached_batch::CachedBatchType}, + sync::Mutex, +}; + +use super::doubly_linked_list::{DoublyLinkedList, DoublyLinkedNode, drop_boxed_node}; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum QueueKind { + Arrow, + Liquid, +} + +#[derive(Debug)] +struct QueueNode { + entry_id: EntryID, + queue: QueueKind, +} + +type NodePtr = NonNull>; + +#[derive(Default, Debug)] +struct LiquidQueueInternalState { + map: HashMap, + arrow: DoublyLinkedList, + liquid: DoublyLinkedList, +} + +impl LiquidQueueInternalState { + unsafe fn list_mut(&mut self, queue: QueueKind) -> &mut DoublyLinkedList { + match queue { + QueueKind::Arrow => &mut self.arrow, + QueueKind::Liquid => &mut self.liquid, + } + } + + unsafe fn push_back(&mut self, queue: QueueKind, mut node_ptr: NodePtr) { + unsafe { + node_ptr.as_mut().data.queue = queue; + self.list_mut(queue).push_back(node_ptr); + } + } + + unsafe fn detach(&mut self, node_ptr: NodePtr) { + unsafe { + let queue = node_ptr.as_ref().data.queue; + self.list_mut(queue).unlink(node_ptr); + } + } + + fn upsert_into_queue(&mut self, entry_id: EntryID, target: QueueKind) { + if let Some(node_ptr) = self.map.get(&entry_id).copied() { + unsafe { + self.detach(node_ptr); + self.push_back(target, node_ptr); + } + return; + } + + let node = DoublyLinkedNode::new(QueueNode { + entry_id, + queue: target, + }); + let node_ptr = NonNull::from(Box::leak(node)); + + self.map.insert(entry_id, node_ptr); + unsafe { + self.push_back(target, node_ptr); + } + } + + fn pop_front(&mut self, queue: QueueKind) -> Option { + let list = match queue { + QueueKind::Arrow => &mut self.arrow, + QueueKind::Liquid => &mut self.liquid, + }; + + let head_ptr = list.head()?; + let entry_id = unsafe { head_ptr.as_ref().data.entry_id }; + let node_ptr = self + .map + .remove(&entry_id) + .expect("list head must exist in map"); + unsafe { + list.unlink(node_ptr); + drop_boxed_node(node_ptr); + } + Some(entry_id) + } + + fn remove(&mut self, entry_id: &EntryID) -> Option { + let node_ptr = self.map.remove(entry_id)?; + let removed = unsafe { node_ptr.as_ref().data.entry_id }; + unsafe { + self.detach(node_ptr); + drop_boxed_node(node_ptr); + } + Some(removed) + } +} + +impl Drop for LiquidQueueInternalState { + fn drop(&mut self) { + let nodes: Vec<_> = self.map.drain().map(|(_, ptr)| ptr).collect(); + for node_ptr in nodes { + unsafe { + match node_ptr.as_ref().data.queue { + QueueKind::Arrow => self.arrow.unlink(node_ptr), + QueueKind::Liquid => self.liquid.unlink(node_ptr), + } + drop_boxed_node(node_ptr); + } + } + + unsafe { + self.arrow.drop_all(); + self.liquid.drop_all(); + } + } +} + +/// Cache policy that keeps independent FIFO queues per batch type. +#[derive(Debug, Default)] +pub struct LiquidPolicy { + inner: Mutex, +} + +impl LiquidPolicy { + /// Create a new [`LiquidPolicy`]. + pub fn new() -> Self { + Self { + inner: Mutex::new(LiquidQueueInternalState::default()), + } + } +} + +// SAFETY: Access to raw pointers is protected by the internal `Mutex`. +unsafe impl Send for LiquidPolicy {} +unsafe impl Sync for LiquidPolicy {} + +impl CachePolicy for LiquidPolicy { + fn notify_insert(&self, entry_id: &EntryID, batch_type: CachedBatchType) { + let mut inner = self.inner.lock().unwrap(); + let target = match batch_type { + CachedBatchType::MemoryArrow => QueueKind::Arrow, + CachedBatchType::MemoryLiquid => QueueKind::Liquid, + }; + + inner.upsert_into_queue(*entry_id, target); + } + + fn find_memory_victim(&self, cnt: usize) -> Vec { + if cnt == 0 { + return vec![]; + } + + let mut inner = self.inner.lock().unwrap(); + let mut victims = Vec::with_capacity(cnt); + + while victims.len() < cnt { + if let Some(entry) = inner.pop_front(QueueKind::Arrow) { + victims.push(entry); + continue; + } + + if let Some(entry) = inner.pop_front(QueueKind::Liquid) { + victims.push(entry); + continue; + } + + break; + } + + victims + } + + fn notify_access(&self, _entry_id: &EntryID, _batch_type: CachedBatchType) {} + + fn notify_remove(&self, entry_id: &EntryID) { + let mut inner = self.inner.lock().unwrap(); + inner.remove(entry_id); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::cache::utils::EntryID; + + fn entry(id: usize) -> EntryID { + id.into() + } + + #[test] + fn test_fifo_within_each_queue() { + let policy = LiquidPolicy::new(); + + let arrow_a = entry(1); + let arrow_b = entry(2); + let liquid_a = entry(3); + let liquid_b = entry(4); + + policy.notify_insert(&arrow_a, CachedBatchType::MemoryArrow); + policy.notify_insert(&arrow_b, CachedBatchType::MemoryArrow); + policy.notify_insert(&liquid_a, CachedBatchType::MemoryLiquid); + policy.notify_insert(&liquid_b, CachedBatchType::MemoryLiquid); + + assert_eq!(policy.find_memory_victim(1), vec![arrow_a]); + assert_eq!(policy.find_memory_victim(2), vec![arrow_b, liquid_a]); + assert_eq!(policy.find_memory_victim(1), vec![liquid_b]); + } + + #[test] + fn test_queue_priority_order() { + let policy = LiquidPolicy::new(); + + let arrow_entry = entry(1); + let liquid_entry = entry(2); + + policy.notify_insert(&liquid_entry, CachedBatchType::MemoryLiquid); + policy.notify_insert(&arrow_entry, CachedBatchType::MemoryArrow); + + // Request more victims than available to ensure we only get what exists. + let victims = policy.find_memory_victim(5); + assert_eq!(victims, vec![arrow_entry, liquid_entry]); + } + + #[test] + fn test_zero_victim_request_returns_empty() { + let policy = LiquidPolicy::new(); + + policy.notify_insert(&entry(1), CachedBatchType::MemoryArrow); + assert!(policy.find_memory_victim(0).is_empty()); + } + + #[test] + fn test_reinsert_moves_entry_to_back_of_queue() { + let policy = LiquidPolicy::new(); + + let first = entry(1); + let second = entry(2); + + policy.notify_insert(&first, CachedBatchType::MemoryArrow); + policy.notify_insert(&second, CachedBatchType::MemoryArrow); + + // Reinserting should refresh the entry as the newest arrow batch. + policy.notify_insert(&first, CachedBatchType::MemoryArrow); + + assert_eq!(policy.find_memory_victim(1), vec![second]); + assert_eq!(policy.find_memory_victim(1), vec![first]); + } + + #[test] + fn test_reinsert_handles_cross_queue_move() { + let policy = LiquidPolicy::new(); + + let entry_id = entry(42); + + policy.notify_insert(&entry_id, CachedBatchType::MemoryArrow); + policy.notify_insert(&entry_id, CachedBatchType::MemoryLiquid); + + let victims = policy.find_memory_victim(2); + assert_eq!(victims, vec![entry_id]); + } +} diff --git a/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/cache/policies/mod.rs b/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/cache/policies/mod.rs new file mode 100644 index 0000000000000..fd5fae6a955c0 --- /dev/null +++ b/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/cache/policies/mod.rs @@ -0,0 +1,7 @@ +//! Policy modules for cache eviction and squeezing. + +pub mod cache; +pub mod squeeze; + +pub use cache::*; +pub use squeeze::*; diff --git a/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/cache/policies/squeeze.rs b/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/cache/policies/squeeze.rs new file mode 100644 index 0000000000000..c5d43152c4e43 --- /dev/null +++ b/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/cache/policies/squeeze.rs @@ -0,0 +1,106 @@ +//! Squeeze policies for liquid cache. + +use crate::cache::{cached_batch::CacheEntry, transcode_liquid_inner}; + +/// What to do when we need to squeeze an entry? +#[derive(Debug, Clone)] +pub enum SqueezeOutcome { + /// Replace the cache entry with a smaller in-memory representation. + Replace(CacheEntry), + /// Remove the entry entirely. + Remove, +} + +/// Policy that chooses the next representation for an entry under memory pressure. +pub trait SqueezePolicy: std::fmt::Debug + Send + Sync { + /// Squeeze the entry. + fn squeeze(&self, entry: &CacheEntry) -> SqueezeOutcome; +} + +/// Evict the entry from memory. +#[derive(Debug, Default, Clone)] +pub struct Evict; + +impl SqueezePolicy for Evict { + fn squeeze(&self, _entry: &CacheEntry) -> SqueezeOutcome { + SqueezeOutcome::Remove + } +} + +/// Transcode Arrow entries to liquid memory; evict liquid entries. +#[derive(Debug, Default, Clone)] +pub struct TranscodeEvict; + +impl SqueezePolicy for TranscodeEvict { + fn squeeze(&self, entry: &CacheEntry) -> SqueezeOutcome { + match entry { + CacheEntry::MemoryArrow(array) => match transcode_liquid_inner(array) { + Ok(liquid_array) => { + SqueezeOutcome::Replace(CacheEntry::memory_liquid(liquid_array)) + } + Err(_) => SqueezeOutcome::Remove, + }, + CacheEntry::MemoryLiquid(_) => SqueezeOutcome::Remove, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::cache::cached_batch::CacheEntry; + use arrow::array::{ArrayRef, BooleanArray, Int32Array}; + use std::sync::Arc; + + fn int_array(n: i32) -> ArrayRef { + Arc::new(Int32Array::from_iter_values(0..n)) + } + + #[test] + fn test_evict_policy_always_removes() { + let policy = Evict; + + let arrow_entry = CacheEntry::memory_arrow(int_array(8)); + assert!(matches!( + policy.squeeze(&arrow_entry), + SqueezeOutcome::Remove + )); + + let arr = int_array(8); + let liquid = transcode_liquid_inner(&arr).unwrap(); + let liquid_entry = CacheEntry::memory_liquid(liquid); + assert!(matches!( + policy.squeeze(&liquid_entry), + SqueezeOutcome::Remove + )); + } + + #[test] + fn test_transcode_evict_policy() { + let policy = TranscodeEvict; + + // MemoryArrow -> MemoryLiquid + let arr = int_array(8); + let outcome = policy.squeeze(&CacheEntry::memory_arrow(arr.clone())); + match outcome { + SqueezeOutcome::Replace(CacheEntry::MemoryLiquid(liq)) => { + assert_eq!(liq.to_arrow_array().as_ref(), arr.as_ref()); + } + other => panic!("unexpected: {other:?}"), + } + + // MemoryLiquid -> Remove + let liquid = transcode_liquid_inner(&arr).unwrap(); + let outcome = policy.squeeze(&CacheEntry::memory_liquid(liquid)); + assert!(matches!(outcome, SqueezeOutcome::Remove)); + } + + #[test] + fn transcode_evict_untranscodable_removes() { + let policy = TranscodeEvict; + // Boolean arrays are not supported by the transcoder. + let arr: ArrayRef = Arc::new(BooleanArray::from(vec![true, false, true])); + let outcome = policy.squeeze(&CacheEntry::memory_arrow(arr)); + assert!(matches!(outcome, SqueezeOutcome::Remove)); + } +} diff --git a/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/cache/transcode.rs b/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/cache/transcode.rs new file mode 100644 index 0000000000000..284cafa31036f --- /dev/null +++ b/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/cache/transcode.rs @@ -0,0 +1,201 @@ +use std::sync::Arc; + +use arrow::array::types::*; +use arrow::array::{ArrayRef, AsArray}; +use arrow_schema::{DataType, TimeUnit}; + +use crate::liquid_array::{ + LiquidArrayRef, LiquidDecimalArray, LiquidFloatArray, LiquidPrimitiveArray, +}; + +/// This method is used to transcode an arrow array into a liquid array. +/// +/// Returns the transcoded liquid array if successful, otherwise returns the original arrow array. +pub fn transcode_liquid_inner(array: &ArrayRef) -> Result { + let data_type = array.data_type(); + if data_type.is_primitive() { + // For primitive types, perform the transcoding. + let liquid_array: LiquidArrayRef = match data_type { + DataType::Int8 => Arc::new(LiquidPrimitiveArray::::from_arrow_array( + array.as_primitive::().clone(), + )), + DataType::Int16 => Arc::new(LiquidPrimitiveArray::::from_arrow_array( + array.as_primitive::().clone(), + )), + DataType::Int32 => Arc::new(LiquidPrimitiveArray::::from_arrow_array( + array.as_primitive::().clone(), + )), + DataType::Int64 => Arc::new(LiquidPrimitiveArray::::from_arrow_array( + array.as_primitive::().clone(), + )), + DataType::UInt8 => Arc::new(LiquidPrimitiveArray::::from_arrow_array( + array.as_primitive::().clone(), + )), + DataType::UInt16 => Arc::new(LiquidPrimitiveArray::::from_arrow_array( + array.as_primitive::().clone(), + )), + DataType::UInt32 => Arc::new(LiquidPrimitiveArray::::from_arrow_array( + array.as_primitive::().clone(), + )), + DataType::UInt64 => Arc::new(LiquidPrimitiveArray::::from_arrow_array( + array.as_primitive::().clone(), + )), + DataType::Date32 => Arc::new(LiquidPrimitiveArray::::from_arrow_array( + array.as_primitive::().clone(), + )), + DataType::Date64 => Arc::new(LiquidPrimitiveArray::::from_arrow_array( + array.as_primitive::().clone(), + )), + DataType::Timestamp(TimeUnit::Second, None) => Arc::new(LiquidPrimitiveArray::< + TimestampSecondType, + >::from_arrow_array( + array.as_primitive::().clone(), + )), + DataType::Timestamp(TimeUnit::Millisecond, None) => Arc::new(LiquidPrimitiveArray::< + TimestampMillisecondType, + >::from_arrow_array( + array.as_primitive::().clone(), + )), + DataType::Timestamp(TimeUnit::Microsecond, None) => Arc::new(LiquidPrimitiveArray::< + TimestampMicrosecondType, + >::from_arrow_array( + array.as_primitive::().clone(), + )), + DataType::Timestamp(TimeUnit::Nanosecond, None) => Arc::new(LiquidPrimitiveArray::< + TimestampNanosecondType, + >::from_arrow_array( + array.as_primitive::().clone(), + )), + DataType::Timestamp(_, Some(_)) => { + log::warn!("unsupported timestamp type with timezone {data_type:?}"); + return Err(array); + } + DataType::Float32 => Arc::new(LiquidFloatArray::::from_arrow_array( + array.as_primitive::().clone(), + )), + DataType::Float64 => Arc::new(LiquidFloatArray::::from_arrow_array( + array.as_primitive::().clone(), + )), + DataType::Decimal128(_, _) => { + let decimals = array.as_primitive::(); + if LiquidDecimalArray::fits_u64(decimals) { + return Ok(Arc::new(LiquidDecimalArray::from_decimal_array(decimals))); + } + log::debug!("decimal128 does not fit u64, not transcoding"); + return Err(array); + } + DataType::Decimal256(_, _) => { + let decimals = array.as_primitive::(); + if LiquidDecimalArray::fits_u64(decimals) { + return Ok(Arc::new(LiquidDecimalArray::from_decimal_array(decimals))); + } + log::debug!("decimal256 does not fit u64, not transcoding"); + return Err(array); + } + _ => { + // For unsupported primitive types, leave the value unchanged. + log::warn!("unsupported primitive type {data_type:?}"); + return Err(array); + } + }; + return Ok(liquid_array); + } + + log::debug!("unsupported data type {:?}", array.data_type()); + Err(array) +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{ + ArrayRef, BooleanArray, Float32Array, Float64Array, Int32Array, Int64Array, + TimestampMicrosecondArray, + }; + + const TEST_ARRAY_SIZE: usize = 8192; + + fn assert_transcode(original: &ArrayRef, transcoded: &LiquidArrayRef) { + assert!( + transcoded.get_array_memory_size() < original.get_array_memory_size(), + "transcoded size: {}, original size: {}", + transcoded.get_array_memory_size(), + original.get_array_memory_size() + ); + let back_to_arrow = transcoded.to_arrow_array(); + assert_eq!(original, &back_to_arrow); + } + + #[test] + fn test_transcode_int32() { + let array: ArrayRef = Arc::new(Int32Array::from_iter_values(0..TEST_ARRAY_SIZE as i32)); + let transcoded = transcode_liquid_inner(&array).unwrap(); + assert_transcode(&array, &transcoded); + } + + #[test] + fn test_transcode_int64() { + let array: ArrayRef = Arc::new(Int64Array::from_iter_values(0..TEST_ARRAY_SIZE as i64)); + let transcoded = transcode_liquid_inner(&array).unwrap(); + assert_transcode(&array, &transcoded); + } + + #[test] + fn test_transcode_float32() { + let array: ArrayRef = Arc::new(Float32Array::from_iter_values( + (0..TEST_ARRAY_SIZE).map(|i| i as f32), + )); + let transcoded = transcode_liquid_inner(&array).unwrap(); + assert_transcode(&array, &transcoded); + } + + #[test] + fn test_transcode_float64() { + let array: ArrayRef = Arc::new(Float64Array::from_iter_values( + (0..TEST_ARRAY_SIZE).map(|i| i as f64), + )); + + let transcoded = transcode_liquid_inner(&array).unwrap(); + assert_transcode(&array, &transcoded); + } + + #[test] + fn test_transcode_timestamp_microsecond() { + let array: ArrayRef = Arc::new(TimestampMicrosecondArray::from_iter_values( + (0..TEST_ARRAY_SIZE).map(|i| (i as i64) * 1_000), + )); + + let transcoded = transcode_liquid_inner(&array).unwrap(); + assert_transcode(&array, &transcoded); + } + + #[test] + fn test_transcode_decimal128() { + use arrow::array::Decimal128Builder; + let mut builder = Decimal128Builder::new(); + for i in 0..TEST_ARRAY_SIZE { + builder.append_value(i as i128 * 100); + } + let array: ArrayRef = Arc::new(builder.finish().with_precision_and_scale(20, 2).unwrap()); + + let transcoded = transcode_liquid_inner(&array).unwrap(); + let back_to_arrow = transcoded.to_arrow_array(); + assert_eq!(&array, &back_to_arrow); + } + + #[test] + fn test_transcode_unsupported_type() { + // Create a boolean array which is not supported by the transcoder + let values: Vec = (0..TEST_ARRAY_SIZE).map(|i| i.is_multiple_of(2)).collect(); + let array: ArrayRef = Arc::new(BooleanArray::from(values)); + + // Try to transcode and expect an error + let result = transcode_liquid_inner(&array); + + // Verify it returns Err with the original array + assert!(result.is_err()); + if let Err(original) = result { + assert_eq!(&array, original); + } + } +} diff --git a/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/cache/utils.rs b/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/cache/utils.rs new file mode 100644 index 0000000000000..215b16ecc93d5 --- /dev/null +++ b/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/cache/utils.rs @@ -0,0 +1,80 @@ +#[cfg(test)] +use crate::cache::cached_batch::CacheEntry; +#[cfg(test)] +use crate::sync::Arc; +#[cfg(test)] +use arrow::array::ArrayRef; + +#[derive(Debug)] +pub struct CacheConfig { + batch_size: usize, + max_memory_bytes: usize, +} + +impl CacheConfig { + pub(super) fn new(batch_size: usize, max_memory_bytes: usize) -> Self { + Self { + batch_size, + max_memory_bytes, + } + } + + pub fn batch_size(&self) -> usize { + self.batch_size + } + + pub fn max_memory_bytes(&self) -> usize { + self.max_memory_bytes + } +} + +// Helper methods +#[cfg(test)] +pub(crate) fn create_test_array(size: usize) -> CacheEntry { + use arrow::array::Int64Array; + use std::sync::Arc; + + CacheEntry::memory_arrow(Arc::new(Int64Array::from_iter_values(0..size as i64))) +} + +// Helper methods +#[cfg(test)] +pub(crate) fn create_test_arrow_array(size: usize) -> ArrayRef { + use arrow::array::Int64Array; + Arc::new(Int64Array::from_iter_values(0..size as i64)) +} + +#[cfg(test)] +pub(crate) fn create_cache_store( + max_memory_bytes: usize, + policy: Box, +) -> Arc { + use crate::cache::{LiquidCacheBuilder, TranscodeEvict}; + + let batch_size = 128; + + let builder = LiquidCacheBuilder::new() + .with_batch_size(batch_size) + .with_max_memory_bytes(max_memory_bytes) + .with_squeeze_policy(Box::new(TranscodeEvict)) + .with_cache_policy(policy); + builder.build() +} + +/// EntryID is a unique identifier for a batch of rows, i.e., the cache key. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Ord, PartialOrd, serde::Serialize)] +pub struct EntryID { + val: usize, +} + +impl From for EntryID { + fn from(val: usize) -> Self { + Self { val } + } +} + +impl From for usize { + fn from(val: EntryID) -> Self { + val.val + } +} diff --git a/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/lib.rs b/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/lib.rs new file mode 100644 index 0000000000000..a84cbcdc78af6 --- /dev/null +++ b/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/lib.rs @@ -0,0 +1,12 @@ +// SPDX-License-Identifier: Apache-2.0 +// Vendored in-memory subset of liquid-cache core. See ../../README.md for provenance. + +//! In-memory liquid cache: cache storage (index, budget, eviction policies, +//! transcode) plus the numeric LiquidArray encodings. + +pub mod cache; +pub mod liquid_array; +mod sync; +pub mod utils; + +pub use cache::cache_policies; diff --git a/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/liquid_array/decimal_array.rs b/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/liquid_array/decimal_array.rs new file mode 100644 index 0000000000000..10bb692c137cd --- /dev/null +++ b/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/liquid_array/decimal_array.rs @@ -0,0 +1,207 @@ +use std::any::Any; +use std::mem::size_of; +use std::sync::Arc; + +use arrow::array::{Array, ArrayRef, PrimitiveArray}; +use arrow::buffer::ScalarBuffer; +use arrow::datatypes::{Decimal128Type, Decimal256Type, DecimalType, UInt64Type, i256}; +use arrow_schema::DataType; +use num_traits::ToPrimitive; + +use super::{LiquidArray, LiquidDataType}; +use crate::liquid_array::raw::BitPackedArray; +use crate::utils::get_bit_width; + +#[derive(Debug, Clone, Copy)] +struct DecimalMeta { + precision: u8, + scale: i8, + is_256: bool, +} + +impl DecimalMeta { + fn from_data_type(data_type: &DataType) -> Self { + match data_type { + DataType::Decimal128(precision, scale) => Self { + precision: *precision, + scale: *scale, + is_256: false, + }, + DataType::Decimal256(precision, scale) => Self { + precision: *precision, + scale: *scale, + is_256: true, + }, + _ => panic!("unsupported decimal data type: {data_type:?}"), + } + } + + fn data_type(&self) -> DataType { + if self.is_256 { + DataType::Decimal256(self.precision, self.scale) + } else { + DataType::Decimal128(self.precision, self.scale) + } + } +} + +/// Liquid decimal array stored as a compressed u64 primitive. +#[derive(Debug)] +pub struct LiquidDecimalArray { + meta: DecimalMeta, + bit_packed: BitPackedArray, + reference_value: u64, +} + +impl LiquidDecimalArray { + pub(crate) fn fits_u64(array: &PrimitiveArray) -> bool + where + T::Native: ToPrimitive, + { + array.iter().flatten().all(|v| v.to_u64().is_some()) + } + + pub(crate) fn from_decimal_array(array: &PrimitiveArray) -> Self + where + T::Native: ToPrimitive, + { + debug_assert!(Self::fits_u64(array)); + let meta = DecimalMeta::from_data_type(array.data_type()); + if array.null_count() == array.len() { + return Self { + meta, + bit_packed: BitPackedArray::new_null_array(array.len()), + reference_value: 0, + }; + } + + let nulls = array.nulls().cloned(); + let mut min = u64::MAX; + let mut max = 0u64; + let values: Vec = array + .iter() + .map(|v| match v { + Some(v) => { + let value = v.to_u64().expect("decimal fits u64"); + if value < min { + min = value; + } + if value > max { + max = value; + } + value + } + None => 0, + }) + .collect(); + + let bit_width = get_bit_width(max - min); + let offsets = ScalarBuffer::from_iter(values.iter().map(|v| v.saturating_sub(min))); + let unsigned_array = PrimitiveArray::::new(offsets, nulls); + let bit_packed = BitPackedArray::from_primitive(unsigned_array, bit_width); + + Self { + meta, + bit_packed, + reference_value: min, + } + } + + fn to_u64_array(&self) -> PrimitiveArray { + let unsigned_array = self.bit_packed.to_primitive(); + let (_data_type, values, _nulls) = unsigned_array.into_parts(); + let nulls = self.bit_packed.nulls(); + let values = if self.reference_value != 0 { + let reference_value = self.reference_value; + ScalarBuffer::from_iter(values.iter().map(|v| v.wrapping_add(reference_value))) + } else { + values + }; + PrimitiveArray::::new(values, nulls.cloned()) + } +} + +impl LiquidArray for LiquidDecimalArray { + fn as_any(&self) -> &dyn Any { + self + } + + fn get_array_memory_size(&self) -> usize { + self.bit_packed.get_array_memory_size() + size_of::() + size_of::() + } + + fn len(&self) -> usize { + self.bit_packed.len() + } + + fn to_arrow_array(&self) -> ArrayRef { + let u64_array = self.to_u64_array(); + let (_data_type, values, nulls) = u64_array.into_parts(); + let data_type = self.meta.data_type(); + if self.meta.is_256 { + let values_i256 = + ScalarBuffer::from_iter(values.iter().map(|v| i256::from_i128(*v as i128))); + let array = PrimitiveArray::::new(values_i256, nulls); + Arc::new(array.with_data_type(data_type)) + } else { + let values_i128 = ScalarBuffer::from_iter(values.iter().map(|v| *v as i128)); + let array = PrimitiveArray::::new(values_i128, nulls); + Arc::new(array.with_data_type(data_type)) + } + } + + fn original_arrow_data_type(&self) -> DataType { + self.meta.data_type() + } + + fn data_type(&self) -> LiquidDataType { + LiquidDataType::Decimal + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::cache::LiquidExpr; + use arrow::array::{BooleanArray, Decimal128Builder}; + use arrow::buffer::BooleanBuffer; + use datafusion_common::ScalarValue; + use datafusion_expr_common::operator::Operator as DFOperator; + use datafusion_physical_expr::PhysicalExpr; + use datafusion_physical_expr::expressions::{BinaryExpr, Column, Literal}; + use std::sync::Arc; + + #[test] + fn decimal_u64_roundtrip() { + let mut builder = Decimal128Builder::new(); + builder.append_value(100_i128); + builder.append_null(); + builder.append_value(250_i128); + let original = builder.finish().with_precision_and_scale(10, 2).unwrap(); + + let liquid = LiquidDecimalArray::from_decimal_array(&original); + let arrow = liquid.to_arrow_array(); + assert_eq!(arrow.as_ref(), &original); + } + + #[test] + fn decimal_predicate_eval() { + let mut builder = Decimal128Builder::new(); + builder.append_value(100_i128); + builder.append_value(200_i128); + builder.append_null(); + builder.append_value(300_i128); + let original = builder.finish().with_precision_and_scale(10, 2).unwrap(); + + let liquid = LiquidDecimalArray::from_decimal_array(&original); + + let mask = BooleanBuffer::new_set(original.len()); + let lit = Arc::new(Literal::new(ScalarValue::Decimal128(Some(150_i128), 10, 2))); + let col = Arc::new(Column::new("col", 0)); + let expr: Arc = Arc::new(BinaryExpr::new(col, DFOperator::GtEq, lit)); + + let got = liquid.try_eval_predicate(&LiquidExpr::new_unchecked(expr), &mask); + let expected = BooleanArray::from(vec![Some(false), Some(true), None, Some(true)]); + assert_eq!(got, expected); + } +} diff --git a/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/liquid_array/float_array.rs b/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/liquid_array/float_array.rs new file mode 100644 index 0000000000000..932ca04685d05 --- /dev/null +++ b/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/liquid_array/float_array.rs @@ -0,0 +1,590 @@ +/// +/// Acknowledgement: +/// The ALP compression implemented in this file is based on the Rust implementation available at https://github.com/spiraldb/alp +/// +use std::{any::Any, fmt::Debug, ops::Mul, sync::Arc}; + +use arrow::{ + array::{Array, ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType, BooleanArray, PrimitiveArray}, + buffer::{BooleanBuffer, ScalarBuffer}, + datatypes::{ + ArrowNativeType, Float32Type, Float64Type, Int32Type, Int64Type, UInt32Type, UInt64Type, + }, +}; +use arrow_schema::DataType; +use fastlanes::BitPacking; +use num_traits::{AsPrimitive, Float, FromPrimitive}; + +use super::LiquidDataType; +use crate::cache::LiquidExpr; +use crate::liquid_array::LiquidArray; +use crate::liquid_array::eval_predicate_on_array; +use crate::liquid_array::raw::BitPackedArray; +use crate::utils::get_bit_width; + +mod private { + use arrow::{ + array::ArrowNumericType, + datatypes::{Float32Type, Float64Type}, + }; + use num_traits::AsPrimitive; + + pub trait Sealed: ArrowNumericType + AsPrimitive> {} + + impl Sealed for Float32Type {} + impl Sealed for Float64Type {} +} + +const NUM_SAMPLES: usize = 1024; // we use FASTLANES to encode array, the sample size needs to be at least 1024 to get a good estimate of the best exponents + +/// LiquidFloatType is a sealed trait that represents all the float types supported by Liquid. +/// Implementors are Float32Type and Float64Type. TODO(): What about Float16Type, decimal types? +pub trait LiquidFloatType: + ArrowPrimitiveType< + Native: AsPrimitive< + ::Native // Native must be convertible to the Native type of Self::UnSignedType + > + + AsPrimitive<::Native> + + FromPrimitive + + AsPrimitive<::Native> + + Mul<::Native> + + Float // required for decode_single and encode_single_unchecked + > + + private::Sealed + + Debug +{ + type UnsignedIntType: + ArrowPrimitiveType< + Native: BitPacking + + AsPrimitive<::Native> + + AsPrimitive<::Native> + + AsPrimitive + > + + Debug; + type SignedIntType: + ArrowPrimitiveType< + Native: AsPrimitive<::Native> + + AsPrimitive<::Native> + + Ord + + From + > + + Debug + Sync + Send; + + const SWEET: ::Native; + const MAX_EXPONENT: u8; + const FRACTIONAL_BITS: u8; + const F10: &'static [::Native]; + const IF10: &'static [::Native]; + + #[inline] + fn fast_round(val: ::Native) -> ::Native { + ((val + Self::SWEET) - Self::SWEET).as_() + } + + #[inline] + fn encode_single_unchecked(val: &::Native, exp: &Exponents) -> ::Native { + Self::fast_round(*val * Self::F10[exp.e as usize] * Self::IF10[exp.f as usize]) + } + + #[inline] + fn decode_single(val: &::Native, exp: &Exponents) -> ::Native { + let decoded_float: ::Native = (*val).as_(); + decoded_float * Self::F10[exp.f as usize] * Self::IF10[exp.e as usize] + } + +} + +impl LiquidFloatType for Float32Type { + type UnsignedIntType = UInt32Type; + type SignedIntType = Int32Type; + const FRACTIONAL_BITS: u8 = 23; + const MAX_EXPONENT: u8 = 10; + const SWEET: ::Native = (1 << Self::FRACTIONAL_BITS) + as ::Native + + (1 << (Self::FRACTIONAL_BITS - 1)) as ::Native; + const F10: &'static [::Native] = &[ + 1.0, + 10.0, + 100.0, + 1000.0, + 10000.0, + 100000.0, + 1000000.0, + 10000000.0, + 100000000.0, + 1000000000.0, + 10000000000.0, // 10^10 + ]; + const IF10: &'static [::Native] = &[ + 1.0, + 0.1, + 0.01, + 0.001, + 0.0001, + 0.00001, + 0.000001, + 0.0000001, + 0.00000001, + 0.000000001, + 0.0000000001, // 10^-10 + ]; +} + +impl LiquidFloatType for Float64Type { + type UnsignedIntType = UInt64Type; + type SignedIntType = Int64Type; + const FRACTIONAL_BITS: u8 = 52; + const MAX_EXPONENT: u8 = 18; + const SWEET: ::Native = (1u64 << Self::FRACTIONAL_BITS) + as ::Native + + (1u64 << (Self::FRACTIONAL_BITS - 1)) as ::Native; + const F10: &'static [::Native] = &[ + 1.0, + 10.0, + 100.0, + 1000.0, + 10000.0, + 100000.0, + 1000000.0, + 10000000.0, + 100000000.0, + 1000000000.0, + 10000000000.0, + 100000000000.0, + 1000000000000.0, + 10000000000000.0, + 100000000000000.0, + 1000000000000000.0, + 10000000000000000.0, + 100000000000000000.0, + 1000000000000000000.0, + 10000000000000000000.0, + 100000000000000000000.0, + 1000000000000000000000.0, + 10000000000000000000000.0, + 100000000000000000000000.0, // 10^23 + ]; + + const IF10: &'static [::Native] = &[ + 1.0, + 0.1, + 0.01, + 0.001, + 0.0001, + 0.00001, + 0.000001, + 0.0000001, + 0.00000001, + 0.000000001, + 0.0000000001, + 0.00000000001, + 0.000000000001, + 0.0000000000001, + 0.00000000000001, + 0.000000000000001, + 0.0000000000000001, + 0.00000000000000001, + 0.000000000000000001, + 0.0000000000000000001, + 0.00000000000000000001, + 0.000000000000000000001, + 0.0000000000000000000001, + 0.00000000000000000000001, // 10^-23 + ]; +} + +/// Liquid's single-precision floating point array +pub type LiquidFloat32Array = LiquidFloatArray; +/// Liquid's double precision floating point array +pub type LiquidFloat64Array = LiquidFloatArray; + +/// An array that stores floats in ALP +#[derive(Debug, Clone)] +pub struct LiquidFloatArray { + exponent: Exponents, + bit_packed: BitPackedArray, + patch_indices: Vec, + patch_values: Vec, + reference_value: ::Native, +} + +impl LiquidFloatArray +where + T: LiquidFloatType, +{ + /// Check if the Liquid float array is empty. + pub fn is_empty(&self) -> bool { + self.len() == 0 + } + + /// Get the length of the Liquid float array. + pub fn len(&self) -> usize { + self.bit_packed.len() + } + + /// Get the memory size of the Liquid primitive array. + pub fn get_array_memory_size(&self) -> usize { + self.bit_packed.get_array_memory_size() + + size_of::() + + self.patch_indices.capacity() * size_of::() + + self.patch_values.capacity() * size_of::() + + size_of::<::Native>() + } + + /// Create a Liquid primitive array from an Arrow float array. + pub fn from_arrow_array(arrow_array: arrow::array::PrimitiveArray) -> LiquidFloatArray { + let best_exponents = get_best_exponents::(&arrow_array); + encode_arrow_array(&arrow_array, &best_exponents) + } +} + +impl LiquidArray for LiquidFloatArray +where + T: LiquidFloatType, +{ + fn as_any(&self) -> &dyn Any { + self + } + + fn get_array_memory_size(&self) -> usize { + self.get_array_memory_size() + } + + fn len(&self) -> usize { + self.len() + } + + #[inline] + fn to_arrow_array(&self) -> ArrayRef { + let unsigned_array = self.bit_packed.to_primitive(); + let (_data_type, values, _nulls) = unsigned_array.into_parts(); + let nulls = self.bit_packed.nulls(); + // TODO(): Check if we should align vectors to cache line boundary + let mut decoded_values = Vec::from_iter(values.iter().map(|v| { + let mut val: ::Native = (*v).as_(); + val = val.add_wrapping(self.reference_value); + T::decode_single(&val, &self.exponent) + })); + + // Patch values + if !self.patch_indices.is_empty() { + for i in 0..self.patch_indices.len() { + decoded_values[self.patch_indices[i].as_usize()] = self.patch_values[i]; + } + } + + Arc::new(PrimitiveArray::::new( + ScalarBuffer::<::Native>::from(decoded_values), + nulls.cloned(), + )) + } + + fn original_arrow_data_type(&self) -> DataType { + T::DATA_TYPE.clone() + } + + fn data_type(&self) -> LiquidDataType { + LiquidDataType::Float + } + + fn is_empty(&self) -> bool { + self.len() == 0 + } + + fn to_best_arrow_array(&self) -> ArrayRef { + self.to_arrow_array() + } + + fn try_eval_predicate(&self, predicate: &LiquidExpr, filter: &BooleanBuffer) -> BooleanArray { + let filtered = self.filter(filter); + eval_predicate_on_array(filtered, predicate) + } +} + +#[derive(Debug, Copy, Clone, PartialEq, Eq)] +pub struct Exponents { + pub(crate) e: u8, + pub(crate) f: u8, +} + +fn encode_arrow_array( + arrow_array: &PrimitiveArray, + exp: &Exponents, // fill_value: &mut Option<::Native> +) -> LiquidFloatArray { + let mut patch_indices: Vec = Vec::new(); + let mut patch_values: Vec = Vec::new(); + let mut patch_count: usize = 0; + let mut fill_value: Option<::Native> = None; + let values = arrow_array.values(); + let nulls = arrow_array.nulls(); + + // All values are null + if arrow_array.null_count() == arrow_array.len() { + return LiquidFloatArray:: { + bit_packed: BitPackedArray::new_null_array(arrow_array.len()), + exponent: Exponents { e: 0, f: 0 }, + patch_indices: Vec::new(), + patch_values: Vec::new(), + reference_value: ::Native::ZERO, + }; + } + + let mut encoded_values = Vec::with_capacity(arrow_array.len()); + for v in values.iter() { + let encoded = T::encode_single_unchecked(&v.as_(), exp); + let decoded = T::decode_single(&encoded, exp); + // TODO(): Check if this is a bitwise comparison + let neq = !decoded.eq(&v.as_()) as usize; + patch_count += neq; + encoded_values.push(encoded); + } + + if patch_count > 0 { + patch_indices.resize_with(patch_count + 1, Default::default); + patch_values.resize_with(patch_count + 1, Default::default); + let mut patch_index: usize = 0; + + for i in 0..encoded_values.len() { + let decoded = T::decode_single(&encoded_values[i], exp); + patch_indices[patch_index] = i.as_(); + patch_values[patch_index] = arrow_array.value(i).as_(); + patch_index += !(decoded.eq(&values[i].as_())) as usize; + } + assert_eq!(patch_index, patch_count); + unsafe { + patch_indices.set_len(patch_count); + patch_values.set_len(patch_count); + } + } + + // find the first successfully encoded value (i.e., not patched) + // this is our fill value for missing values + if patch_count > 0 && patch_count < arrow_array.len() { + for i in 0..encoded_values.len() { + if i >= patch_indices.len() || patch_indices[i] != i as u64 { + fill_value = encoded_values.get(i).copied(); + break; + } + } + } + + // replace the patched values in the encoded array with the fill value + // for better downstream compression + if let Some(fill_value) = fill_value { + // handle the edge case where the first N >= 1 chunks are all patches + for patch_idx in &patch_indices { + encoded_values[*patch_idx as usize] = fill_value; + } + } + + let min = *encoded_values + .iter() + .min() + .expect("`encoded_values` shouldn't be all nulls"); + let max = *encoded_values + .iter() + .max() + .expect("`encoded_values` shouldn't be all nulls"); + let sub: ::Native = max.sub_wrapping(min).as_(); + + let unsigned_encoded_values = encoded_values + .iter() + .map(|v| { + let k: ::Native = v.sub_wrapping(min).as_(); + k + }) + .collect::>(); + let encoded_output = PrimitiveArray::<::UnsignedIntType>::new( + ScalarBuffer::from(unsigned_encoded_values), + nulls.cloned(), + ); + + let bit_width = get_bit_width(sub.as_()); + let bit_packed_array = BitPackedArray::from_primitive(encoded_output, bit_width); + + LiquidFloatArray:: { + bit_packed: bit_packed_array, + exponent: *exp, + patch_indices, + patch_values, + reference_value: min, + } +} + +fn get_best_exponents(arrow_array: &PrimitiveArray) -> Exponents { + let mut best_exponents = Exponents { e: 0, f: 0 }; + let mut min_encoded_size: usize = usize::MAX; + + let sample_arrow_array: Option> = + (arrow_array.len() > NUM_SAMPLES).then(|| { + arrow_array + .iter() + .step_by(arrow_array.len() / NUM_SAMPLES) + .filter(|s| s.is_some()) + .collect() + }); + + for e in 0..T::MAX_EXPONENT { + for f in 0..e { + let exp = Exponents { e, f }; + let liquid_array = + encode_arrow_array(sample_arrow_array.as_ref().unwrap_or(arrow_array), &exp); + if liquid_array.get_array_memory_size() < min_encoded_size { + best_exponents = exp; + min_encoded_size = liquid_array.get_array_memory_size(); + } + } + } + best_exponents +} + +#[cfg(test)] +mod tests { + use super::*; + + macro_rules! test_roundtrip { + ($test_name: ident, $type:ty, $values: expr) => { + #[test] + fn $test_name() { + let original: Vec::Native>> = $values; + let array = PrimitiveArray::<$type>::from(original.clone()); + + // Convert to Liquid array and back + let liquid_array = LiquidFloatArray::<$type>::from_arrow_array(array.clone()); + let result_array = liquid_array.to_arrow_array(); + + assert_eq!(result_array.as_ref(), &array); + } + }; + } + + // Test cases for Float32 + test_roundtrip!( + test_float32_roundtrip_basic, + Float32Type, + vec![Some(-1.0), Some(1.0), Some(0.0)] + ); + + test_roundtrip!( + test_float32_roundtrip_with_nones, + Float32Type, + vec![Some(-1.0), Some(1.0), Some(0.0), None] + ); + + test_roundtrip!( + test_float32_roundtrip_all_nones, + Float32Type, + vec![None, None, None, None] + ); + + test_roundtrip!(test_float32_roundtrip_empty, Float32Type, vec![]); + + // Test cases for Float64 + test_roundtrip!( + test_float64_roundtrip_basic, + Float64Type, + vec![Some(-1.0), Some(1.0), Some(0.0)] + ); + + test_roundtrip!( + test_float64_roundtrip_with_nones, + Float64Type, + vec![Some(-1.0), Some(1.0), Some(0.0), None] + ); + + test_roundtrip!( + test_float64_roundtrip_all_nones, + Float64Type, + vec![None, None, None, None] + ); + + test_roundtrip!(test_float64_roundtrip_empty, Float64Type, vec![]); + + // Tests with ilters + #[test] + fn test_filter_basic() { + // Create original array with some values + let original = vec![Some(1.0), Some(2.1), Some(3.2), None, Some(5.5)]; + let array = PrimitiveArray::::from(original); + let liquid_array = LiquidFloatArray::::from_arrow_array(array); + + // Create selection mask: keep indices 0, 2, and 4 + let selection = BooleanBuffer::from(vec![true, false, true, false, true]); + + // Apply filter + let result_array = liquid_array.filter(&selection); + + // Expected result after filtering + let expected = PrimitiveArray::::from(vec![Some(1.0), Some(3.2), Some(5.5)]); + + assert_eq!(result_array.as_ref(), &expected); + } + + #[test] + fn test_original_arrow_data_type_returns_float32() { + let array = PrimitiveArray::::from(vec![Some(1.0), Some(2.5)]); + let liquid = LiquidFloatArray::::from_arrow_array(array); + assert_eq!(liquid.original_arrow_data_type(), DataType::Float32); + } + + #[test] + fn test_filter_all_nulls() { + // Create array with all nulls + let original = vec![None, None, None, None]; + let array = PrimitiveArray::::from(original); + let liquid_array = LiquidFloatArray::::from_arrow_array(array); + + // Keep first and last elements + let selection = BooleanBuffer::from(vec![true, false, false, true]); + + let result_array = liquid_array.filter(&selection); + + let expected = PrimitiveArray::::from(vec![None, None]); + + assert_eq!(result_array.as_ref(), &expected); + } + + #[test] + fn test_filter_empty_result() { + let original = vec![Some(1.0), Some(2.1), Some(3.3)]; + let array = PrimitiveArray::::from(original); + let liquid_array = LiquidFloatArray::::from_arrow_array(array); + + // Filter out all elements + let selection = BooleanBuffer::from(vec![false, false, false]); + + let result_array = liquid_array.filter(&selection); + + assert_eq!(result_array.len(), 0); + } + + #[test] + fn test_compression_f32_f64() { + fn run_compression_test( + type_name: &str, + data_fn: impl Fn(usize) -> T::Native, + ) { + let original: Vec = (0..2000).map(data_fn).collect(); + let array = PrimitiveArray::::from_iter_values(original); + let uncompressed_size = array.get_array_memory_size(); + + let liquid_array = LiquidFloatArray::::from_arrow_array(array); + let compressed_size = liquid_array.get_array_memory_size(); + + println!( + "Type: {type_name}, uncompressed_size: {uncompressed_size}, compressed_size: {compressed_size}" + ); + // Assert that compression actually reduced the size + assert!( + compressed_size < uncompressed_size, + "{type_name} compression failed to reduce size" + ); + } + + // Run for f32 + run_compression_test::("f32", |i| i as f32); + + // Run for f64 + run_compression_test::("f64", |i| i as f64); + } +} diff --git a/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/liquid_array/linear_integer_array.rs b/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/liquid_array/linear_integer_array.rs new file mode 100644 index 0000000000000..695e8be37018b --- /dev/null +++ b/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/liquid_array/linear_integer_array.rs @@ -0,0 +1,655 @@ +use std::any::Any; +use std::fmt::Debug; +use std::marker::PhantomData; +use std::sync::Arc; + +use super::PrimitiveKind; +use super::{LiquidArray, LiquidDataType, LiquidPrimitiveType}; +use crate::cache::LiquidExpr; +use crate::liquid_array::LiquidPrimitiveArray; +use crate::liquid_array::eval_predicate_on_array; +use arrow::array::{ + Array, ArrayRef, ArrowPrimitiveType, BooleanArray, PrimitiveArray, + cast::AsArray, + types::{ + Date32Type, Date64Type, Int8Type, Int16Type, Int32Type, Int64Type, UInt8Type, UInt16Type, + UInt32Type, UInt64Type, + }, +}; +use arrow::buffer::{BooleanBuffer, ScalarBuffer}; +use arrow::compute::kernels::filter; +use arrow_schema::DataType; +use num_traits::{AsPrimitive, Bounded, FromPrimitive}; + +/// A linear-model based integer array, **only use it when you know the array is monotonic** and **you don't care about encoding speed!**. +/// +/// Under the hood, it uses a linear model to predict the values and store the residuals: +/// value\[i\] = intercept + round(slope * i) + residual\[i\] +/// +/// Where `intercept` and `slope` are computed using a L-infinity linear fit, **this is time-consuming!**. +/// +/// This array is only recommended if you know the array follows a linear model, e.g., kinetic values, offsets, etc. +/// +/// Examples of not recommended use cases: random values like ids, categorical values, etc. +#[derive(Debug)] +pub struct LiquidLinearArray +where + T::Native: AsPrimitive + FromPrimitive + Bounded, +{ + // Signed residuals, bit-packed as a Liquid primitive array of i64. + residuals: LiquidPrimitiveArray, + // Intercept term stored as f64 for simpler math/IO. + intercept: f64, + // Slope term of the linear model. + slope: f64, + // Keep the logical type parameter. + _phantom: PhantomData, +} + +/// Backward-compatible alias for i32. +pub type LiquidLinearI32Array = LiquidLinearArray; +/// Linear-model array for `i8`. +pub type LiquidLinearI8Array = LiquidLinearArray; +/// Linear-model array for `i16`. +pub type LiquidLinearI16Array = LiquidLinearArray; +/// Linear-model array for `i64`. +pub type LiquidLinearI64Array = LiquidLinearArray; +/// Linear-model array for `u8`. +pub type LiquidLinearU8Array = LiquidLinearArray; +/// Linear-model array for `u16`. +pub type LiquidLinearU16Array = LiquidLinearArray; +/// Linear-model array for `u32`. +pub type LiquidLinearU32Array = LiquidLinearArray; +/// Linear-model array for `u64`. +pub type LiquidLinearU64Array = LiquidLinearArray; +/// Linear-model array for `Date32` (days since epoch). +pub type LiquidLinearDate32Array = LiquidLinearArray; +/// Linear-model array for `Date64` (ms since epoch). +pub type LiquidLinearDate64Array = LiquidLinearArray; + +impl LiquidLinearArray +where + T: LiquidPrimitiveType, + T::Native: AsPrimitive + FromPrimitive + Bounded, +{ + /// Build from an Arrow `PrimitiveArray` by training a linear model + /// using a fast L-infinity fit (Option 3) and storing residuals. + pub fn from_arrow_array(arrow_array: PrimitiveArray) -> Self { + let len = arrow_array.len(); + + // All nulls + if arrow_array.null_count() == len { + // All nulls + let res = PrimitiveArray::::new_null(len); + return Self { + residuals: LiquidPrimitiveArray::::from_arrow_array(res), + intercept: 0.0, // arbitrary, unused since all nulls + slope: 0.0, + _phantom: PhantomData, + }; + } + + // Prepare compact non-null buffers for fast fitting (avoid iterator Option cost). + let (nn_values, nn_indices) = collect_non_null_f64_and_indices::(&arrow_array); + + // Option 3 parameters: L-infinity (Chebyshev) regression. + let (mut intercept, mut slope) = fit_linf(&nn_values, &nn_indices); + + // Compute residuals in one unified loop; also track ranges for fallback decision. + let mut residuals: Vec = Vec::with_capacity(len); + let is_unsigned = ::IS_UNSIGNED; + let vals = arrow_array.values(); + let nulls_opt = arrow_array.nulls(); + + // Original value range + let mut orig_min_u64 = u64::MAX; + let mut orig_max_u64 = 0u64; + let mut orig_min_i64 = i64::MAX; + let mut orig_max_i64 = i64::MIN; + // Residual range + let mut res_min = i64::MAX; + let mut res_max = i64::MIN; + + if is_unsigned { + let max_u64: u64 = ::MAX_U64; + for i in 0..len { + let valid = nulls_opt.as_ref().is_none_or(|n| n.is_valid(i)); + if valid { + type U = + <::UnSignedType as ArrowPrimitiveType>::Native; + let v_u: U = vals[i].as_(); + let v_u64: u64 = v_u.as_(); + if v_u64 < orig_min_u64 { + orig_min_u64 = v_u64; + } + if v_u64 > orig_max_u64 { + orig_max_u64 = v_u64; + } + let pr = slope * (i as f64) + intercept; + let p = predict_u64_saturated(pr, max_u64); + let (pos, mag) = if v_u64 >= p { + (true, v_u64 - p) + } else { + (false, p - v_u64) + }; + let m = (mag & (i64::MAX as u64)) as i64; + let r = if pos { m } else { -m }; + if r < res_min { + res_min = r; + } + if r > res_max { + res_max = r; + } + residuals.push(r); + } else { + residuals.push(0); + } + } + } else { + let (min_i64, max_i64): (i64, i64) = + (::MIN_I64, ::MAX_I64); + for i in 0..len { + let valid = nulls_opt.as_ref().is_none_or(|n| n.is_valid(i)); + if valid { + let v_i64: i64 = vals[i].as_(); + if v_i64 < orig_min_i64 { + orig_min_i64 = v_i64; + } + if v_i64 > orig_max_i64 { + orig_max_i64 = v_i64; + } + let pr = slope * (i as f64) + intercept; + let p = predict_i64_saturated(pr, min_i64, max_i64); + let r = v_i64 - p; + if r < res_min { + res_min = r; + } + if r > res_max { + res_max = r; + } + residuals.push(r); + } else { + residuals.push(0); + } + } + } + + // Fallback: ensure residual range is strictly smaller than original range + let res_width: u128 = (res_max as i128 - res_min as i128) as u128; + let orig_width: u128 = if is_unsigned { + (orig_max_u64 as u128).saturating_sub(orig_min_u64 as u128) + } else { + (orig_max_i64 as i128 - orig_min_i64 as i128) as u128 + }; + if res_width >= orig_width { + // Rebuild residuals with zero model + intercept = 0.0; + slope = 0.0; + residuals.clear(); + if is_unsigned { + for i in 0..len { + let valid = nulls_opt.as_ref().is_none_or(|n| n.is_valid(i)); + if valid { + type U = <::UnSignedType as ArrowPrimitiveType>::Native; + let v_u: U = vals[i].as_(); + let v_u64: u64 = v_u.as_(); + let r = (v_u64 & (i64::MAX as u64)) as i64; + residuals.push(r); + } else { + residuals.push(0); + } + } + } else { + for i in 0..len { + let valid = nulls_opt.as_ref().is_none_or(|n| n.is_valid(i)); + if valid { + let v_i64: i64 = vals[i].as_(); + residuals.push(v_i64); + } else { + residuals.push(0); + } + } + } + } + let residuals_buf: ScalarBuffer = ScalarBuffer::from(residuals); + let nulls = arrow_array.nulls().cloned(); + let res_prim = PrimitiveArray::::new(residuals_buf, nulls); + let residuals = LiquidPrimitiveArray::::from_arrow_array(res_prim); + + Self { + residuals, + intercept, + slope, + _phantom: PhantomData, + } + } + + fn len(&self) -> usize { + self.residuals.len() + } +} + +impl LiquidArray for LiquidLinearArray +where + T: LiquidPrimitiveType, + T::Native: AsPrimitive + FromPrimitive + Bounded, +{ + fn as_any(&self) -> &dyn Any { + self + } + + fn original_arrow_data_type(&self) -> DataType { + T::DATA_TYPE.clone() + } + + fn get_array_memory_size(&self) -> usize { + self.residuals.get_array_memory_size() + + std::mem::size_of::() // intercept + + std::mem::size_of::() // slope + } + + fn len(&self) -> usize { + self.len() + } + + fn to_arrow_array(&self) -> ArrayRef { + let arr = self.residuals.to_arrow_array(); + let (_dt, residuals, nulls) = arr.as_primitive::().clone().into_parts(); + + // Reconstruct final values: predicted(i) +/- |residual_i| + let mut final_values = Vec::::with_capacity(self.len()); + let is_unsigned = ::IS_UNSIGNED; + if is_unsigned { + let max_u64: u64 = ::MAX_U64; + for (i, &e) in residuals.iter().enumerate() { + let pr = self.slope * (i as f64) + self.intercept; + let p = predict_u64_saturated(pr, max_u64); + let mag = e.unsigned_abs(); + let sum = if e >= 0 { + p.saturating_add(mag) + } else { + p.saturating_sub(mag) + }; + final_values.push(T::Native::from_u64(sum).unwrap()); + } + } else { + let (min_i64, max_i64): (i64, i64) = + (::MIN_I64, ::MAX_I64); + for (i, &e) in residuals.iter().enumerate() { + let pr = self.slope * (i as f64) + self.intercept; + let p = predict_i64_saturated(pr, min_i64, max_i64); + let sum = p.saturating_add(e); + final_values.push(T::Native::from_i64(sum).unwrap()); + } + } + + let values_buf: ScalarBuffer = ScalarBuffer::from(final_values); + Arc::new(PrimitiveArray::::new(values_buf, nulls)) + } + + fn filter(&self, selection: &BooleanBuffer) -> ArrayRef { + let arr = self.to_arrow_array(); + let selection = BooleanArray::new(selection.clone(), None); + filter::filter(&arr, &selection).unwrap() + } + + fn try_eval_predicate(&self, predicate: &LiquidExpr, filter: &BooleanBuffer) -> BooleanArray { + let arr = self.filter(filter); + eval_predicate_on_array(arr, predicate) + } + + fn data_type(&self) -> LiquidDataType { + LiquidDataType::LinearInteger + } +} + +#[inline] +fn predict_u64_saturated(pred: f64, max_u64: u64) -> u64 { + if !pred.is_finite() || pred <= 0.0 { + 0 + } else if pred >= max_u64 as f64 { + max_u64 + } else { + pred.round() as u64 + } +} + +#[inline] +fn predict_i64_saturated(pred: f64, min_i64: i64, max_i64: i64) -> i64 { + if !pred.is_finite() { + 0 + } else if pred <= min_i64 as f64 { + min_i64 + } else if pred >= max_i64 as f64 { + max_i64 + } else { + pred.round() as i64 + } +} + +/// L-infinity linear fit for y\[i\] ≈ intercept + slope * i (before rounding), +/// minimizing the maximum absolute error over all non-null points. +/// +/// Approach: minimize R(m) = max_i (y_i - m i) - min_i (y_i - m i), which is +/// convex in m. Given m, the best intercept is b = (max_i s_i + min_i s_i)/2 +/// where s_i = y_i - m i. We find m by a few rounds of bisection using the +/// subgradient sign derived from the argmax/argmin indices. O(n) per iteration. +fn fit_linf(values: &[f64], idxs: &[u32]) -> (f64, f64) { + let n = values.len(); + assert_eq!(values.len(), idxs.len()); + if n == 0 { + return (0.0, 0.0); + } + if n == 1 { + return (values[0], 0.0); + } + + let mut slope_min = f64::INFINITY; + let mut slope_max = f64::NEG_INFINITY; + for k in 1..n { + let di = (idxs[k] - idxs[k - 1]) as f64; + if di > 0.0 { + let dv = values[k] - values[k - 1]; + let s = dv / di; + if s < slope_min { + slope_min = s; + } + if s > slope_max { + slope_max = s; + } + } + } + if !slope_min.is_finite() || !slope_max.is_finite() { + slope_min = 0.0; + slope_max = 0.0; + } + + let mut lo = slope_min.min(slope_max); + let mut hi = slope_min.max(slope_max); + if (hi - lo).abs() < 1e-12 { + let pad = if hi.abs() < 1.0 { 1.0 } else { hi.abs() * 1e-6 }; + lo -= pad; + hi += pad; + } + + #[inline] + fn range_stats(values: &[f64], idxs: &[u32], m: f64) -> (f64, u32, f64, u32) { + let mut min_s = f64::INFINITY; + let mut max_s = f64::NEG_INFINITY; + let mut i_min = 0u32; + let mut i_max = 0u32; + for k in 0..values.len() { + let i = idxs[k] as f64; + let s = values[k] - m * i; + if s < min_s { + min_s = s; + i_min = idxs[k]; + } + if s > max_s { + max_s = s; + i_max = idxs[k]; + } + } + (min_s, i_min, max_s, i_max) + } + + const MAX_ITERS: usize = 8; + for _ in 0..MAX_ITERS { + let m = 0.5 * (lo + hi); + let (_min_s, i_min, _max_s, i_max) = range_stats(values, idxs, m); + let g = (i_min as i64) - (i_max as i64); + if g > 0 { + hi = m; + } else if g < 0 { + lo = m; + } else { + lo = m; + hi = m; + break; + } + if (hi - lo).abs() < 1e-12 { + break; + } + } + + let m = 0.5 * (lo + hi); + let (min_s, _i_min, max_s, _i_max) = range_stats(values, idxs, m); + let b = 0.5 * (max_s + min_s); + (b, m) +} + +#[inline] +fn collect_non_null_f64_and_indices(arr: &PrimitiveArray) -> (Vec, Vec) +where + T: LiquidPrimitiveType, + T::Native: AsPrimitive, +{ + let nn = arr.len() - arr.null_count(); + let mut values = Vec::with_capacity(nn); + let mut idxs = Vec::with_capacity(nn); + let vals = arr.values(); + if arr.null_count() == 0 { + for (i, v) in vals.iter().enumerate() { + values.push(v.as_()); + idxs.push(i as u32); + } + } else { + let nulls = arr.nulls().unwrap(); + for (i, v) in vals.iter().enumerate() { + if nulls.is_valid(i) { + values.push(v.as_()); + idxs.push(i as u32); + } + } + } + (values, idxs) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn roundtrip_eq(values: Vec>) { + let arr = PrimitiveArray::::from(values.clone()); + let linear = LiquidLinearI32Array::from_arrow_array(arr.clone()); + let decoded = linear.to_arrow_array(); + assert_eq!(decoded.as_ref(), &arr); + } + + macro_rules! roundtrip_eq_t { + ($T:ty, $values:expr) => {{ + let arr = PrimitiveArray::<$T>::from(($values).clone()); + let linear = LiquidLinearArray::<$T>::from_arrow_array(arr.clone()); + let decoded = linear.to_arrow_array(); + assert_eq!(decoded.as_ref(), &arr); + }}; + } + + #[test] + fn test_roundtrip_basic() { + // Non-monotonic values to ensure we don't rely on simple increasing sequences + roundtrip_eq(vec![ + Some(10), + Some(15), + Some(14), + Some(20), + Some(18), + Some(25), + Some(24), + ]); + } + + #[test] + fn test_roundtrip_with_nulls() { + roundtrip_eq(vec![Some(10), None, Some(30), None, Some(50), Some(70)]); + } + + #[test] + fn test_all_nulls() { + roundtrip_eq(vec![None, None, None, None]); + } + + #[test] + fn test_single_value() { + roundtrip_eq(vec![Some(42)]); + } + + #[test] + fn test_empty() { + roundtrip_eq(vec![]); + } + + #[test] + fn test_negative_values() { + roundtrip_eq(vec![ + Some(-100), + Some(-50), + Some(0), + Some(50), + Some(25), + None, + Some(-25), + ]); + } + + #[test] + fn test_filter_basic() { + let original: Vec> = vec![Some(1), Some(2), Some(3), None, Some(5), Some(8)]; + let arr = PrimitiveArray::::from(original.clone()); + let linear = LiquidLinearI32Array::from_arrow_array(arr); + let selection = BooleanBuffer::from(vec![true, false, true, false, true, false]); + let result = linear.filter(&selection); + let expected = PrimitiveArray::::from(vec![Some(1), Some(3), Some(5)]); + assert_eq!(result.as_ref(), &expected); + } + + #[test] + fn test_original_arrow_data_type_returns_int32() { + let arr = PrimitiveArray::::from(vec![Some(1), Some(2)]); + let linear = LiquidLinearI32Array::from_arrow_array(arr); + assert_eq!(linear.original_arrow_data_type(), DataType::Int32); + } + + #[test] + fn test_roundtrip_i8() { + roundtrip_eq_t!(Int8Type, vec![Some(-10), Some(0), Some(10), None, Some(20)]); + } + + #[test] + fn test_roundtrip_i16() { + roundtrip_eq_t!( + Int16Type, + vec![Some(-1000), Some(0), Some(1000), None, Some(2000)] + ); + } + + #[test] + fn test_roundtrip_i64() { + roundtrip_eq_t!( + Int64Type, + vec![ + Some(-10_000_000_000), + Some(0), + Some(10_000_000_000), + None, + Some(20_000_000_000), + ] + ); + } + + #[test] + fn test_roundtrip_u8() { + roundtrip_eq_t!( + UInt8Type, + vec![Some(0), Some(10), Some(200), None, Some(255)] + ); + } + + #[test] + fn test_roundtrip_u16() { + roundtrip_eq_t!( + UInt16Type, + vec![Some(0), Some(1000), Some(60000), None, Some(500)] + ); + } + + #[test] + fn test_roundtrip_u32() { + roundtrip_eq_t!( + UInt32Type, + vec![ + Some(0), + Some(1_000_000), + Some(3_000_000_000), + None, + Some(123_456_789), + ] + ); + } + + #[test] + fn test_roundtrip_u64() { + roundtrip_eq_t!( + UInt64Type, + vec![ + Some(0), + Some(10_000_000_000), + Some(9_000_000_000_000_000_000u64), + None, + Some(42), + ] + ); + } + + #[test] + fn test_roundtrip_date32() { + roundtrip_eq_t!( + Date32Type, + vec![Some(-365), Some(0), Some(365), None, Some(18262)] + ); + } + + #[test] + fn test_roundtrip_date64() { + roundtrip_eq_t!( + Date64Type, + vec![ + Some(-86_400_000), + Some(0), + Some(86_400_000), + None, + Some(1_000_000_000_000), + ] + ); + } + + #[test] + fn test_compression() { + let original = (0..1_000_000).step_by(100).collect::>(); + + let original = PrimitiveArray::::from_iter_values(original); + let arrow_size = original.get_array_memory_size(); + + let liquid_linear = LiquidLinearI32Array::from_arrow_array(original.clone()); + let liquid_linear_size = liquid_linear.get_array_memory_size(); + + let liquid_primitive = + LiquidPrimitiveArray::::from_arrow_array(original.clone()); + let liquid_primitive_size = liquid_primitive.get_array_memory_size(); + + println!( + "arrow_size: {arrow_size}, liquid_linear_size: {liquid_linear_size}, liquid_primitive_size: {liquid_primitive_size}", + ); + + assert!(liquid_linear_size < arrow_size); + assert!(liquid_primitive_size < arrow_size); + assert!(liquid_linear_size < liquid_primitive_size); + + let original: ArrayRef = Arc::new(original); + assert_eq!(original.as_ref(), liquid_linear.to_arrow_array().as_ref()); + assert_eq!( + original.as_ref(), + liquid_primitive.to_arrow_array().as_ref() + ); + } +} diff --git a/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/liquid_array/mod.rs b/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/liquid_array/mod.rs new file mode 100644 index 0000000000000..d6b3d607b1982 --- /dev/null +++ b/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/liquid_array/mod.rs @@ -0,0 +1,173 @@ +//! LiquidArray is the core data structure of LiquidCache. +//! You should not use this module directly. +//! Instead, use `liquid_cache_datafusion_server` or `liquid_cache_datafusion_client` to interact with LiquidCache. +mod decimal_array; +mod float_array; +mod linear_integer_array; +mod primitive_array; +pub mod raw; + +use std::{any::Any, sync::Arc}; + +use arrow::{ + array::{ArrayRef, BooleanArray, cast::AsArray}, + buffer::BooleanBuffer, + record_batch::RecordBatch, +}; +use arrow_schema::{DataType, Field, Schema}; +pub use decimal_array::LiquidDecimalArray; +pub use float_array::{LiquidFloat32Array, LiquidFloat64Array, LiquidFloatArray}; +pub use linear_integer_array::{ + LiquidLinearArray, LiquidLinearDate32Array, LiquidLinearDate64Array, LiquidLinearI8Array, + LiquidLinearI16Array, LiquidLinearI32Array, LiquidLinearI64Array, LiquidLinearU8Array, + LiquidLinearU16Array, LiquidLinearU32Array, LiquidLinearU64Array, +}; +pub use primitive_array::{ + LiquidDate32Array, LiquidDate64Array, LiquidI8Array, LiquidI16Array, LiquidI32Array, + LiquidI64Array, LiquidPrimitiveArray, LiquidPrimitiveDeltaArray, LiquidPrimitiveType, + LiquidU8Array, LiquidU16Array, LiquidU32Array, LiquidU64Array, +}; + +use crate::cache::LiquidExpr; + +/// Liquid data type is only logical type +#[derive(Debug, Clone, Copy)] +#[repr(u16)] +pub enum LiquidDataType { + /// An integer. + Integer = 1, + /// A float. + Float = 2, + /// A linear-model based integer (signed residuals + model params). + LinearInteger = 5, + /// A decimal encoded as a primitive u64 array. + Decimal = 6, +} + +/// A Liquid array. +pub trait LiquidArray: std::fmt::Debug + Send + Sync { + /// Get the underlying any type. + fn as_any(&self) -> &dyn Any; + + /// Get the memory size of the Liquid array. + fn get_array_memory_size(&self) -> usize; + + /// Get the length of the Liquid array. + fn len(&self) -> usize; + + /// Check if the Liquid array is empty. + fn is_empty(&self) -> bool { + self.len() == 0 + } + + /// Convert the Liquid array to an Arrow array. + fn to_arrow_array(&self) -> ArrayRef; + + /// Convert the Liquid array to an Arrow array. + /// Except that it will pick the best encoding for the arrow array. + /// Meaning that it may not obey the data type of the original arrow array. + fn to_best_arrow_array(&self) -> ArrayRef { + self.to_arrow_array() + } + + /// Get the logical data type of the Liquid array. + fn data_type(&self) -> LiquidDataType; + + /// Get the original arrow data type of the Liquid array. + fn original_arrow_data_type(&self) -> DataType; + + /// Filter the Liquid array with a boolean array and return an **arrow array**. + fn filter(&self, selection: &BooleanBuffer) -> ArrayRef { + let arrow_array = self.to_arrow_array(); + let selection = BooleanArray::new(selection.clone(), None); + arrow::compute::kernels::filter::filter(&arrow_array, &selection).unwrap() + } + + /// Evaluate a predicate on the Liquid array with a filter. + /// + /// Note that the filter is a boolean buffer, not a boolean array, i.e., filter can't be nullable. + /// The returned boolean mask is nullable if the the original array is nullable. + fn try_eval_predicate(&self, predicate: &LiquidExpr, filter: &BooleanBuffer) -> BooleanArray { + let filtered = self.filter(filter); + eval_predicate_on_array(filtered, predicate) + } +} + +/// A reference to a Liquid array. +pub type LiquidArrayRef = Arc; + +pub(crate) fn eval_predicate_on_array(array: ArrayRef, predicate: &LiquidExpr) -> BooleanArray { + let schema = Arc::new(Schema::new(vec![Field::new( + "liquid_predicate_col", + array.data_type().clone(), + true, + )])); + let record_batch = RecordBatch::try_new(schema, vec![array]).expect("predicate input batch"); + let result = predicate + .physical_expr() + .evaluate(&record_batch) + .expect("validated LiquidExpr must evaluate"); + let boolean_array = result + .into_array(record_batch.num_rows()) + .expect("predicate output must be an array"); + boolean_array.as_boolean().clone() +} + +/// Compile-time info about primitive kind (signed vs unsigned) and bounds. +/// Implemented for all Liquid-supported primitive integer and date types. +pub trait PrimitiveKind { + /// Whether the logical type is unsigned (true for u8/u16/u32/u64). + const IS_UNSIGNED: bool; + /// Maximum representable value as u64 for unsigned types (unused for signed). + const MAX_U64: u64; + /// Minimum representable value as i64 for signed/date types (unused for unsigned). + const MIN_I64: i64; + /// Maximum representable value as i64 for signed/date types (unused for unsigned). + const MAX_I64: i64; +} + +macro_rules! impl_unsigned_kind { + ($t:ty, $max:expr) => { + impl PrimitiveKind for $t { + const IS_UNSIGNED: bool = true; + const MAX_U64: u64 = $max as u64; + const MIN_I64: i64 = 0; // unused + const MAX_I64: i64 = 0; // unused + } + }; +} + +macro_rules! impl_signed_kind { + ($t:ty, $min:expr, $max:expr) => { + impl PrimitiveKind for $t { + const IS_UNSIGNED: bool = false; + const MAX_U64: u64 = 0; // unused + const MIN_I64: i64 = $min as i64; + const MAX_I64: i64 = $max as i64; + } + }; +} + +use arrow::datatypes::{ + Date32Type, Date64Type, Int8Type, Int16Type, Int32Type, Int64Type, TimestampMicrosecondType, + TimestampMillisecondType, TimestampNanosecondType, TimestampSecondType, UInt8Type, UInt16Type, + UInt32Type, UInt64Type, +}; + +impl_unsigned_kind!(UInt8Type, u8::MAX); +impl_unsigned_kind!(UInt16Type, u16::MAX); +impl_unsigned_kind!(UInt32Type, u32::MAX); +impl_unsigned_kind!(UInt64Type, u64::MAX); + +impl_signed_kind!(Int8Type, i8::MIN, i8::MAX); +impl_signed_kind!(Int16Type, i16::MIN, i16::MAX); +impl_signed_kind!(Int32Type, i32::MIN, i32::MAX); +impl_signed_kind!(Int64Type, i64::MIN, i64::MAX); + +// Dates are logically signed in Arrow (Date32: i32 days, Date64: i64 ms) +impl_signed_kind!(Date32Type, i32::MIN, i32::MAX); +impl_signed_kind!(Date64Type, i64::MIN, i64::MAX); +impl_signed_kind!(TimestampSecondType, i64::MIN, i64::MAX); +impl_signed_kind!(TimestampMillisecondType, i64::MIN, i64::MAX); +impl_signed_kind!(TimestampMicrosecondType, i64::MIN, i64::MAX); +impl_signed_kind!(TimestampNanosecondType, i64::MIN, i64::MAX); diff --git a/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/liquid_array/primitive_array.rs b/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/liquid_array/primitive_array.rs new file mode 100644 index 0000000000000..3eab53acb5f6a --- /dev/null +++ b/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/liquid_array/primitive_array.rs @@ -0,0 +1,706 @@ +use std::any::Any; +use std::fmt::{Debug, Display}; +use std::sync::Arc; + +use arrow::array::{ + ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType, BooleanArray, PrimitiveArray, + types::{ + Date32Type, Date64Type, Int8Type, Int16Type, Int32Type, Int64Type, + TimestampMicrosecondType, TimestampMillisecondType, TimestampNanosecondType, + TimestampSecondType, UInt8Type, UInt16Type, UInt32Type, UInt64Type, + }, +}; +use arrow::buffer::{BooleanBuffer, ScalarBuffer}; +use arrow_schema::DataType; +use fastlanes::BitPacking; +use num_traits::{AsPrimitive, FromPrimitive}; + +use super::LiquidDataType; +use crate::cache::LiquidExpr; +use crate::liquid_array::raw::BitPackedArray; +use crate::liquid_array::{LiquidArray, PrimitiveKind, eval_predicate_on_array}; +use crate::utils::get_bit_width; +use arrow::datatypes::ArrowNativeType; + +mod private { + pub trait Sealed {} +} + +/// LiquidPrimitiveType is a sealed trait that represents the primitive types supported by Liquid. +/// Implemented for all supported integer, date, and timestamp Arrow primitive types. +/// +/// I have to admit this trait is super complicated. +/// Luckily users never have to worry about it, they can just use the types that are already implemented. +/// We could have implemented this as a macro, but macro is ugly. +/// Type is spec, code is proof. +pub trait LiquidPrimitiveType: + ArrowPrimitiveType< + Native: AsPrimitive<::Native> + + AsPrimitive + + FromPrimitive + + Display, + > + Debug + + Send + + Sync + + private::Sealed + + PrimitiveKind +{ + /// The unsigned type that can be used to represent the signed type. + type UnSignedType: ArrowPrimitiveType + AsPrimitive + BitPacking> + + Debug; +} + +macro_rules! impl_has_unsigned_type { + ($($signed:ty => $unsigned:ty),*) => { + $( + impl private::Sealed for $signed {} + impl LiquidPrimitiveType for $signed { + type UnSignedType = $unsigned; + } + )* + } +} + +impl_has_unsigned_type! { + Int32Type => UInt32Type, + Int64Type => UInt64Type, + Int16Type => UInt16Type, + Int8Type => UInt8Type, + UInt32Type => UInt32Type, + UInt64Type => UInt64Type, + UInt16Type => UInt16Type, + UInt8Type => UInt8Type, + Date64Type => UInt64Type, + Date32Type => UInt32Type, + TimestampSecondType => UInt64Type, + TimestampMillisecondType => UInt64Type, + TimestampMicrosecondType => UInt64Type, + TimestampNanosecondType => UInt64Type +} + +/// Liquid's unsigned 8-bit integer array. +pub type LiquidU8Array = LiquidPrimitiveArray; +/// Liquid's unsigned 16-bit integer array. +pub type LiquidU16Array = LiquidPrimitiveArray; +/// Liquid's unsigned 32-bit integer array. +pub type LiquidU32Array = LiquidPrimitiveArray; +/// Liquid's unsigned 64-bit integer array. +pub type LiquidU64Array = LiquidPrimitiveArray; +/// Liquid's signed 8-bit integer array. +pub type LiquidI8Array = LiquidPrimitiveArray; +/// Liquid's signed 16-bit integer array. +pub type LiquidI16Array = LiquidPrimitiveArray; +/// Liquid's signed 32-bit integer array. +pub type LiquidI32Array = LiquidPrimitiveArray; +/// Liquid's signed 64-bit integer array. +pub type LiquidI64Array = LiquidPrimitiveArray; +/// Liquid's 32-bit date array. +pub type LiquidDate32Array = LiquidPrimitiveArray; +/// Liquid's 64-bit date array. +pub type LiquidDate64Array = LiquidPrimitiveArray; + +/// Liquid's primitive array +#[derive(Debug)] +pub struct LiquidPrimitiveArray { + bit_packed: BitPackedArray, + reference_value: T::Native, +} + +/// Liquid's primitive array which uses delta encoding for compression +#[derive(Debug, Clone)] +pub struct LiquidPrimitiveDeltaArray { + bit_packed: BitPackedArray, + reference_value: T::Native, +} + +impl LiquidPrimitiveArray +where + T: LiquidPrimitiveType, +{ + /// Get the memory size of the Liquid primitive array. + pub fn get_array_memory_size(&self) -> usize { + self.bit_packed.get_array_memory_size() + std::mem::size_of::() + } + + /// Get the length of the Liquid primitive array. + pub fn len(&self) -> usize { + self.bit_packed.len() + } + + /// Check if the Liquid primitive array is empty. + pub fn is_empty(&self) -> bool { + self.len() == 0 + } + + /// Create a Liquid primitive array from an Arrow primitive array. + pub fn from_arrow_array(arrow_array: PrimitiveArray) -> LiquidPrimitiveArray { + let min = match arrow::compute::kernels::aggregate::min(&arrow_array) { + Some(v) => v, + None => { + // entire array is null + return Self { + bit_packed: BitPackedArray::new_null_array(arrow_array.len()), + reference_value: T::Native::ZERO, + }; + } + }; + let max = arrow::compute::kernels::aggregate::max(&arrow_array).unwrap(); + + // be careful of overflow: + // Want: 127i8 - (-128i8) -> 255u64, + // but we get -1i8 + // (-1i8) as u8 as u64 -> 255u64 + let sub = max.sub_wrapping(min) as ::Native; + let sub: <::UnSignedType as ArrowPrimitiveType>::Native = + sub.as_(); + let bit_width = get_bit_width(sub.as_()); + + let (_data_type, values, nulls) = arrow_array.clone().into_parts(); + let values = if min != T::Native::ZERO { + ScalarBuffer::from_iter(values.iter().map(|v| { + let k: <::UnSignedType as ArrowPrimitiveType>::Native = + v.sub_wrapping(min).as_(); + k + })) + } else { + #[allow(clippy::missing_transmute_annotations)] + unsafe { + std::mem::transmute(values) + } + }; + + let unsigned_array = + PrimitiveArray::<::UnSignedType>::new(values, nulls); + + let bit_packed_array = BitPackedArray::from_primitive(unsigned_array, bit_width); + + Self { + bit_packed: bit_packed_array, + reference_value: min, + } + } +} + +impl LiquidPrimitiveDeltaArray +where + T: LiquidPrimitiveType, +{ + /// Get the memory size of the Liquid primitive delta array. + pub fn get_array_memory_size(&self) -> usize { + self.bit_packed.get_array_memory_size() + std::mem::size_of::() + } + + /// Get the length of the Liquid primitive delta array. + pub fn len(&self) -> usize { + self.bit_packed.len() + } + + /// Check if the Liquid primitive delta array is empty. + pub fn is_empty(&self) -> bool { + self.len() == 0 + } + + /// Create a Liquid primitive delta array from an Arrow primitive array. + pub fn from_arrow_array(arrow_array: PrimitiveArray) -> LiquidPrimitiveDeltaArray { + use arrow::array::Array; + + let len = arrow_array.len(); + // check if entire array is already null + if arrow_array.null_count() == len { + return Self { + bit_packed: BitPackedArray::new_null_array(len), + reference_value: T::Native::ZERO, + }; + } + + let (_dt, values, nulls) = arrow_array.clone().into_parts(); + let vals: Vec = values.to_vec(); + + type UnsignedNative = + <::UnSignedType as ArrowPrimitiveType>::Native; + let mut out: Vec> = Vec::with_capacity(len); + let mut max_value: UnsignedNative = UnsignedNative::::ZERO; + let mut anchor: T::Native = T::Native::ZERO; + + if let Some(_nb) = &nulls { + // Nulls present: write 0 for nulls; prev will the last prev non-null value + let nb = nulls.as_ref().unwrap(); + let mut have_prev = false; + let mut prev: T::Native = T::Native::ZERO; + + for (i, &cur) in vals.iter().enumerate() { + if !nb.is_valid(i) { + out.push(UnsignedNative::::ZERO); + continue; + } + if !have_prev { + anchor = cur; + prev = cur; + have_prev = true; + out.push(UnsignedNative::::ZERO); + continue; + } + let delta: T::Native = cur.sub_wrapping(prev); + // zig zag encoding + let delta_i64: i64 = delta.as_(); + let zigzag: u64 = ((delta_i64 << 1) ^ (delta_i64 >> 63)) as u64; + let delta_unsigned: UnsignedNative = + UnsignedNative::::usize_as(zigzag as usize); + if delta_unsigned > max_value { + max_value = delta_unsigned; + } + out.push(delta_unsigned); + prev = cur; + } + } else { + // No nulls: first value is anchor, remainder are deltas with their previous values + anchor = vals[0]; + let mut prev: T::Native = anchor; + out.push(UnsignedNative::::ZERO); // anchor will have a difference of 0 + for &cur in vals.iter().skip(1) { + let delta: T::Native = cur.sub_wrapping(prev); + // zig zag encoding + let delta_i64: i64 = delta.as_(); + let zigzag: u64 = ((delta_i64 << 1) ^ (delta_i64 >> 63)) as u64; + let delta_unsigned: UnsignedNative = + UnsignedNative::::usize_as(zigzag as usize); + if delta_unsigned > max_value { + max_value = delta_unsigned; + } + out.push(delta_unsigned); + prev = cur; + } + } + + let bit_width = get_bit_width(max_value.as_()); + let values = ScalarBuffer::from_iter(out); + let unsigned_array = + PrimitiveArray::<::UnSignedType>::new(values, nulls); + let bit_packed_array = BitPackedArray::from_primitive(unsigned_array, bit_width); + + Self { + bit_packed: bit_packed_array, + reference_value: anchor, + } + } +} + +impl LiquidArray for LiquidPrimitiveArray +where + T: LiquidPrimitiveType + super::PrimitiveKind, +{ + fn get_array_memory_size(&self) -> usize { + self.get_array_memory_size() + } + + fn len(&self) -> usize { + self.len() + } + + fn original_arrow_data_type(&self) -> DataType { + T::DATA_TYPE.clone() + } + + fn as_any(&self) -> &dyn Any { + self + } + + #[inline] + fn to_arrow_array(&self) -> ArrayRef { + let unsigned_array = self.bit_packed.to_primitive(); + let (_data_type, values, _nulls) = unsigned_array.into_parts(); + let nulls = self.bit_packed.nulls(); + let values = if self.reference_value != T::Native::ZERO { + let reference_v = self.reference_value.as_(); + ScalarBuffer::from_iter(values.iter().map(|v| { + let k: ::Native = (*v).add_wrapping(reference_v).as_(); + k + })) + } else { + #[allow(clippy::missing_transmute_annotations)] + unsafe { + std::mem::transmute(values) + } + }; + + Arc::new(PrimitiveArray::::new(values, nulls.cloned())) + } + + fn filter(&self, selection: &BooleanBuffer) -> ArrayRef { + let arrow_array = self.to_arrow_array(); + let selection = BooleanArray::new(selection.clone(), None); + arrow::compute::kernels::filter::filter(&arrow_array, &selection).unwrap() + } + + fn try_eval_predicate(&self, predicate: &LiquidExpr, filter: &BooleanBuffer) -> BooleanArray { + let filtered = self.filter(filter); + eval_predicate_on_array(filtered, predicate) + } + + fn data_type(&self) -> LiquidDataType { + LiquidDataType::Integer + } +} + +impl LiquidArray for LiquidPrimitiveDeltaArray +where + T: LiquidPrimitiveType + super::PrimitiveKind, +{ + fn get_array_memory_size(&self) -> usize { + self.get_array_memory_size() + } + + fn len(&self) -> usize { + self.len() + } + + fn original_arrow_data_type(&self) -> DataType { + T::DATA_TYPE.clone() + } + + fn as_any(&self) -> &dyn Any { + self + } + + #[inline] + fn to_arrow_array(&self) -> ArrayRef { + // Reconstruct original values from deltas + let unsigned_array = self.bit_packed.to_primitive(); + let (_data_type, delta_values, _nulls) = unsigned_array.into_parts(); + let nulls = self.bit_packed.nulls(); + + // Reconstruct original values by applying deltas + let mut reconstructed = Vec::with_capacity(delta_values.len()); + let mut current_value = self.reference_value; // anchor + + if let Some(nulls) = nulls { + let mut have_prev = false; + for (i, &delta_unsigned) in delta_values.iter().enumerate() { + if !nulls.is_valid(i) { + reconstructed.push(T::Native::ZERO); // Will be masked out by nulls + continue; + } + if !have_prev { + // First non-null value is the anchor + reconstructed.push(current_value); + have_prev = true; + } else { + // Apply delta to get next value + let zigzag: u64 = delta_unsigned.as_(); + let delta_i64 = (zigzag >> 1) as i64 ^ -((zigzag & 1) as i64); + let delta: T::Native = T::Native::from_i64(delta_i64).unwrap(); + current_value = current_value.add_wrapping(delta); + reconstructed.push(current_value); + } + } + } else { + // No nulls case + reconstructed.push(current_value); // First value is anchor + for &delta_unsigned in delta_values.iter().skip(1) { + let zigzag: u64 = delta_unsigned.as_(); + let delta_i64 = (zigzag >> 1) as i64 ^ -((zigzag & 1) as i64); + let delta: T::Native = T::Native::from_i64(delta_i64).unwrap(); + current_value = current_value.add_wrapping(delta); + reconstructed.push(current_value); + } + } + + let values = ScalarBuffer::from_iter(reconstructed); + Arc::new(PrimitiveArray::::new(values, nulls.cloned())) + } + + fn filter(&self, selection: &BooleanBuffer) -> ArrayRef { + let arrow_array = self.to_arrow_array(); + let selection = BooleanArray::new(selection.clone(), None); + arrow::compute::kernels::filter::filter(&arrow_array, &selection).unwrap() + } + + fn try_eval_predicate(&self, predicate: &LiquidExpr, filter: &BooleanBuffer) -> BooleanArray { + let filtered = self.filter(filter); + eval_predicate_on_array(filtered, predicate) + } + + fn data_type(&self) -> LiquidDataType { + LiquidDataType::Integer + } +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::Array; + + macro_rules! test_roundtrip { + ($test_name:ident, $type:ty, $values:expr) => { + #[test] + fn $test_name() { + // Create the original array + let original: Vec::Native>> = $values; + let array = PrimitiveArray::<$type>::from(original.clone()); + + // Convert to Liquid array and back + let liquid_array = LiquidPrimitiveArray::<$type>::from_arrow_array(array.clone()); + let result_array = liquid_array.to_arrow_array(); + + assert_eq!(result_array.as_ref(), &array); + } + }; + } + + // Test cases for Int8Type + test_roundtrip!( + test_int8_roundtrip_basic, + Int8Type, + vec![Some(1), Some(2), Some(3), None, Some(5)] + ); + test_roundtrip!( + test_int8_roundtrip_negative, + Int8Type, + vec![Some(-128), Some(-64), Some(0), Some(63), Some(127)] + ); + + // Test cases for Int16Type + test_roundtrip!( + test_int16_roundtrip_basic, + Int16Type, + vec![Some(1), Some(2), Some(3), None, Some(5)] + ); + test_roundtrip!( + test_int16_roundtrip_negative, + Int16Type, + vec![ + Some(-32768), + Some(-16384), + Some(0), + Some(16383), + Some(32767) + ] + ); + + // Test cases for Int32Type + test_roundtrip!( + test_int32_roundtrip_basic, + Int32Type, + vec![Some(1), Some(2), Some(3), None, Some(5)] + ); + test_roundtrip!( + test_int32_roundtrip_negative, + Int32Type, + vec![ + Some(-2147483648), + Some(-1073741824), + Some(0), + Some(1073741823), + Some(2147483647) + ] + ); + + // Test cases for Int64Type + test_roundtrip!( + test_int64_roundtrip_basic, + Int64Type, + vec![Some(1), Some(2), Some(3), None, Some(5)] + ); + test_roundtrip!( + test_int64_roundtrip_negative, + Int64Type, + vec![ + Some(-9223372036854775808), + Some(-4611686018427387904), + Some(0), + Some(4611686018427387903), + Some(9223372036854775807) + ] + ); + + // Test cases for unsigned types + test_roundtrip!( + test_uint8_roundtrip, + UInt8Type, + vec![Some(0), Some(128), Some(255), None, Some(64)] + ); + test_roundtrip!( + test_uint16_roundtrip, + UInt16Type, + vec![Some(0), Some(32768), Some(65535), None, Some(16384)] + ); + test_roundtrip!( + test_uint32_roundtrip, + UInt32Type, + vec![ + Some(0), + Some(2147483648), + Some(4294967295), + None, + Some(1073741824) + ] + ); + test_roundtrip!( + test_uint64_roundtrip, + UInt64Type, + vec![ + Some(0), + Some(9223372036854775808), + Some(18446744073709551615), + None, + Some(4611686018427387904) + ] + ); + + test_roundtrip!( + test_date32_roundtrip, + Date32Type, + vec![Some(-365), Some(0), Some(365), None, Some(18262)] + ); + + test_roundtrip!( + test_date64_roundtrip, + Date64Type, + vec![Some(-365), Some(0), Some(365), None, Some(18262)] + ); + + // Edge cases + #[test] + fn test_all_nulls() { + let original: Vec> = vec![None, None, None]; + let array = PrimitiveArray::::from(original.clone()); + let liquid_array = LiquidPrimitiveArray::::from_arrow_array(array); + let result_array = liquid_array.to_arrow_array(); + + assert_eq!(result_array.len(), original.len()); + assert_eq!(result_array.null_count(), original.len()); + } + + #[test] + fn test_all_nulls_filter() { + let original: Vec> = vec![None, None, None]; + let array = PrimitiveArray::::from(original.clone()); + let liquid_array = LiquidPrimitiveArray::::from_arrow_array(array); + let result_array = liquid_array.filter(&BooleanBuffer::from(vec![true, false, true])); + + assert_eq!(result_array.len(), 2); + assert_eq!(result_array.null_count(), 2); + } + + #[test] + fn test_zero_reference_value() { + let original: Vec> = vec![Some(0), Some(1), Some(2), None, Some(4)]; + let array = PrimitiveArray::::from(original.clone()); + let liquid_array = LiquidPrimitiveArray::::from_arrow_array(array.clone()); + let result_array = liquid_array.to_arrow_array(); + + assert_eq!(liquid_array.reference_value, 0); + assert_eq!(result_array.as_ref(), &array); + } + + #[test] + fn test_single_value() { + let original: Vec> = vec![Some(42)]; + let array = PrimitiveArray::::from(original.clone()); + let liquid_array = LiquidPrimitiveArray::::from_arrow_array(array.clone()); + let result_array = liquid_array.to_arrow_array(); + + assert_eq!(result_array.as_ref(), &array); + } + + #[test] + fn test_filter_basic() { + // Create original array with some values + let original = vec![Some(1), Some(2), Some(3), None, Some(5)]; + let array = PrimitiveArray::::from(original); + let liquid_array = LiquidPrimitiveArray::::from_arrow_array(array); + + // Create selection mask: keep indices 0, 2, and 4 + let selection = BooleanBuffer::from(vec![true, false, true, false, true]); + + // Apply filter + let result_array = liquid_array.filter(&selection); + + // Expected result after filtering + let expected = PrimitiveArray::::from(vec![Some(1), Some(3), Some(5)]); + + assert_eq!(result_array.as_ref(), &expected); + } + + #[test] + fn test_original_arrow_data_type_returns_int32() { + let array = PrimitiveArray::::from(vec![Some(1), Some(2)]); + let liquid = LiquidPrimitiveArray::::from_arrow_array(array); + assert_eq!(liquid.original_arrow_data_type(), DataType::Int32); + } + + #[test] + fn test_filter_all_nulls() { + // Create array with all nulls + let original = vec![None, None, None, None]; + let array = PrimitiveArray::::from(original); + let liquid_array = LiquidPrimitiveArray::::from_arrow_array(array); + + // Keep first and last elements + let selection = BooleanBuffer::from(vec![true, false, false, true]); + + let result_array = liquid_array.filter(&selection); + + let expected = PrimitiveArray::::from(vec![None, None]); + + assert_eq!(result_array.as_ref(), &expected); + } + + #[test] + fn test_filter_empty_result() { + let original = vec![Some(1), Some(2), Some(3)]; + let array = PrimitiveArray::::from(original); + let liquid_array = LiquidPrimitiveArray::::from_arrow_array(array); + + // Filter out all elements + let selection = BooleanBuffer::from(vec![false, false, false]); + + let result_array = liquid_array.filter(&selection); + + assert_eq!(result_array.len(), 0); + } + + #[test] + fn test_delta_encoding_basic_roundtrip() { + let original = vec![Some(1), Some(3), Some(6), Some(10), Some(15)]; + let array = PrimitiveArray::::from(original.clone()); + + let liquid_delta = LiquidPrimitiveDeltaArray::::from_arrow_array(array.clone()); + let result_array = liquid_delta.to_arrow_array(); + + assert_eq!(result_array.as_ref(), &array); + } + + #[test] + fn test_delta_encoding_with_nulls() { + let original = vec![Some(1), None, Some(4), Some(7), None, Some(12)]; + let array = PrimitiveArray::::from(original.clone()); + + let liquid_delta = LiquidPrimitiveDeltaArray::::from_arrow_array(array.clone()); + let result_array = liquid_delta.to_arrow_array(); + + assert_eq!(result_array.as_ref(), &array); + } + + #[test] + fn test_memory_comparison_sequential_data() { + // Sequential data: delta encoding performs better + let sequential_data: Vec> = (0..1000).map(Some).collect(); + let array = PrimitiveArray::::from(sequential_data); + + let liquid_regular = LiquidPrimitiveArray::::from_arrow_array(array.clone()); + let liquid_delta = LiquidPrimitiveDeltaArray::::from_arrow_array(array); + + let regular_size = liquid_regular.get_array_memory_size(); + let delta_size = liquid_delta.get_array_memory_size(); + + println!( + "Sequential data - Regular: {} bytes, Delta: {} bytes", + regular_size, delta_size + ); + assert!( + delta_size <= regular_size, + "Delta encoding should be more efficient for sequential data" + ); + } +} diff --git a/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/liquid_array/raw/bit_pack_array.rs b/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/liquid_array/raw/bit_pack_array.rs new file mode 100644 index 0000000000000..9c693616db3e3 --- /dev/null +++ b/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/liquid_array/raw/bit_pack_array.rs @@ -0,0 +1,347 @@ +use std::mem::size_of; +use std::num::NonZero; + +use arrow::array::{ArrowPrimitiveType, PrimitiveArray}; +use arrow::buffer::{Buffer, NullBuffer, ScalarBuffer}; +use arrow::datatypes::ArrowNativeType; +use fastlanes::BitPacking; + +/// A bit-packed array. +#[derive(Debug)] +pub struct BitPackedArray +where + T::Native: BitPacking, +{ + packed_values: ScalarBuffer, + nulls: Option, + bit_width: Option>, // if None, the array is entirely null + original_len: usize, +} + +/// Implement Clone for any T that implements ArrowPrimitiveType and BitPacking +/// This allows us to clone it without requiring T to implement Clone +impl Clone for BitPackedArray +where + T::Native: BitPacking, +{ + fn clone(&self) -> Self { + Self { + packed_values: self.packed_values.clone(), + nulls: self.nulls.clone(), + bit_width: self.bit_width, + original_len: self.original_len, + } + } +} + +impl BitPackedArray +where + T::Native: BitPacking, +{ + /// Creates a new null array with the given length. + pub fn new_null_array(len: usize) -> Self { + Self { + packed_values: vec![T::Native::usize_as(0); len].into(), + nulls: Some(NullBuffer::new_null(len)), + bit_width: None, + original_len: len, + } + } + + pub(crate) fn len(&self) -> usize { + self.original_len + } + + pub(crate) fn nulls(&self) -> Option<&NullBuffer> { + self.nulls.as_ref() + } + + /// Returns true if the array is nullable. + #[cfg(test)] + fn is_nullable(&self) -> bool { + self.nulls.is_some() + } + + /// Creates a new bit-packed array from a primitive array and a bit width. + pub fn from_primitive(array: PrimitiveArray, bit_width: NonZero) -> Self { + let original_len = array.len(); + let (_data_type, values, nulls) = array.into_parts(); + + let bit_width_usize = bit_width.get() as usize; + let num_chunks = original_len.div_ceil(1024); + let num_full_chunks = original_len / 1024; + let packed_len = (1024 * bit_width_usize).div_ceil(size_of::() * 8); + + let mut output = Vec::::with_capacity(num_chunks * packed_len); + + (0..num_full_chunks).for_each(|i| { + let start_elem = i * 1024; + + output.reserve(packed_len); + let output_len = output.len(); + unsafe { + output.set_len(output_len + packed_len); + BitPacking::unchecked_pack( + bit_width_usize, + &values[start_elem..][..1024], + &mut output[output_len..][..packed_len], + ); + } + }); + + if num_chunks != num_full_chunks { + let last_chunk_size = values.len() % 1024; + let mut last_chunk = vec![T::Native::default(); 1024]; + last_chunk[..last_chunk_size] + .copy_from_slice(&values[values.len() - last_chunk_size..]); + + output.reserve(packed_len); + let output_len = output.len(); + unsafe { + output.set_len(output_len + packed_len); + BitPacking::unchecked_pack( + bit_width_usize, + &last_chunk, + &mut output[output_len..][..packed_len], + ); + } + } + + let buffer = Buffer::from(output); + let scalar_buffer = ScalarBuffer::new(buffer, 0, num_chunks * packed_len); + + Self { + packed_values: scalar_buffer, + nulls, + bit_width: Some(bit_width), + original_len, + } + } + + /// Converts the bit-packed array to a primitive array. + pub fn to_primitive(&self) -> PrimitiveArray { + // Special case for all nulls, don't unpack + let bit_width = if let Some(bit_width) = self.bit_width { + bit_width.get() as usize + } else { + return PrimitiveArray::::new_null(self.original_len); + }; + let packed = self.packed_values.as_ref(); + let length = self.original_len; + let offset = 0; + + let num_chunks = (offset + length).div_ceil(1024); + let elements_per_chunk = (1024 * bit_width).div_ceil(size_of::() * 8); + + let mut output = Vec::::with_capacity(num_chunks * 1024 - offset); + + let first_full_chunk = if offset != 0 { + let chunk: &[T::Native] = &packed[0..elements_per_chunk]; + let mut decoded = vec![T::Native::default(); 1024]; + unsafe { BitPacking::unchecked_unpack(bit_width, chunk, &mut decoded) }; + output.extend_from_slice(&decoded[offset..]); + 1 + } else { + 0 + }; + + (first_full_chunk..num_chunks).for_each(|i| { + let chunk: &[T::Native] = &packed[i * elements_per_chunk..][0..elements_per_chunk]; + unsafe { + let output_len = output.len(); + output.set_len(output_len + 1024); + BitPacking::unchecked_unpack(bit_width, chunk, &mut output[output_len..][..1024]); + } + }); + + output.truncate(length); + if output.len() < 1024 { + output.shrink_to_fit(); + } + + let nulls = self.nulls.clone(); + PrimitiveArray::::new(ScalarBuffer::from(output), nulls) + } + + /// Returns the memory size of the bit-packed array. + pub fn get_array_memory_size(&self) -> usize { + std::mem::size_of::() + + self.packed_values.inner().capacity() + + self + .nulls + .as_ref() + .map_or(0, |nulls| nulls.buffer().capacity()) + } +} + +#[allow(dead_code)] +fn best_arrow_primitive_width(bit_width: NonZero) -> usize { + match bit_width.get() { + 0..=8 => 8, + 9..=16 => 16, + 17..=32 => 32, + 33..=64 => 64, + _ => panic!("Unsupported bit width: {}", bit_width.get()), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::{ + array::Array, + datatypes::{UInt16Type, UInt32Type}, + }; + + #[test] + fn test_bit_pack_roundtrip() { + // Test with a full chunk (1024 elements) + let values: Vec = (0..1024).collect(); + + let array = PrimitiveArray::::from(values); + let before_size = array.get_array_memory_size(); + let bit_packed = BitPackedArray::from_primitive(array, NonZero::new(10).unwrap()); + let after_size = bit_packed.get_array_memory_size(); + println!("before: {before_size}, after: {after_size}"); + let unpacked = bit_packed.to_primitive(); + + assert_eq!(unpacked.len(), 1024); + for i in 0..1024 { + assert_eq!(unpacked.value(i), i as u32); + } + } + + #[test] + fn test_bit_pack_partial_chunk() { + // Test with a partial chunk (500 elements) + let values: Vec = (0..500).collect(); + let array = PrimitiveArray::::from(values); + let bit_packed = BitPackedArray::from_primitive(array, NonZero::new(10).unwrap()); + let unpacked = bit_packed.to_primitive(); + + assert_eq!(unpacked.len(), 500); + for i in 0..500 { + assert_eq!(unpacked.value(i), i as u32); + } + } + + #[test] + fn test_bit_pack_multiple_chunks() { + // Test with multiple chunks (2048 elements = 2 full chunks) + let values: Vec = (0..2048).collect(); + let array = PrimitiveArray::::from(values); + let bit_packed = BitPackedArray::from_primitive(array, NonZero::new(11).unwrap()); + let unpacked = bit_packed.to_primitive(); + + assert_eq!(unpacked.len(), 2048); + for i in 0..2048 { + assert_eq!(unpacked.value(i), i as u32); + } + } + + #[test] + fn test_bit_pack_with_nulls() { + let values: Vec> = (0..1000) + .map(|i| if i % 2 == 0 { Some(i as u32) } else { None }) + .collect(); + let array = PrimitiveArray::::from(values); + let bit_packed = BitPackedArray::from_primitive(array, NonZero::new(10).unwrap()); + let unpacked = bit_packed.to_primitive(); + + assert_eq!(unpacked.len(), 1000); + for i in 0..1000_usize { + if i.is_multiple_of(2) { + assert_eq!(unpacked.value(i), i as u32); + } else { + assert!(unpacked.is_null(i)); + } + } + } + + #[test] + fn test_different_bit_widths() { + // Test with different bit widths + let values: Vec = (0..100).map(|i| i * 2).collect(); + let array = PrimitiveArray::::from(values); + + for bit_width in [8, 16, 24, 32] { + let bit_packed = + BitPackedArray::from_primitive(array.clone(), NonZero::new(bit_width).unwrap()); + let unpacked = bit_packed.to_primitive(); + + assert_eq!(unpacked.len(), 100); + for i in 0..100 { + assert_eq!(unpacked.value(i), i as u32 * 2); + } + } + } + + #[test] + fn test_roundtrip_with_nulls_and_offset() { + let values: Vec> = (0..32) + .map(|i| if i % 3 == 0 { None } else { Some(i as u16) }) + .collect(); + let array = PrimitiveArray::::from(values); + + // Slice to create a non-zero offset (and therefore a non-zero null bitmap bit offset). + let sliced = array.slice(1, 23); + + let bit_width = NonZero::new(16).unwrap(); + let original = BitPackedArray::from_primitive(sliced.clone(), bit_width); + assert!(original.is_nullable()); + + let roundtripped = original.to_primitive(); + assert_eq!(roundtripped, sliced); + } + + #[test] + fn test_memory_size_calculation() { + use super::*; + use arrow::buffer::{Buffer, NullBuffer, ScalarBuffer}; + use arrow::datatypes::UInt32Type; + + let scalar_buffer = ScalarBuffer::::new(Buffer::from(vec![0; 1024]), 0, 1024); + + // --- Test without nulls --- + let bit_packed_no_nulls = BitPackedArray:: { + packed_values: scalar_buffer.clone(), + nulls: None, + bit_width: Some(NonZero::new(10).unwrap()), + original_len: 1024, + }; + + let expected_size_no_nulls = + size_of::>() + scalar_buffer.inner().capacity(); + assert_eq!( + bit_packed_no_nulls.get_array_memory_size(), + expected_size_no_nulls, + "Memory size mismatch without nulls" + ); + + // --- Test with nulls --- + // Create dummy null buffer + let null_buffer = NullBuffer::new_null(1024); + let nulls = Some(null_buffer); + + let bit_packed_with_nulls = BitPackedArray:: { + packed_values: scalar_buffer.clone(), + nulls: nulls.clone(), // Clone the Option + bit_width: Some(NonZero::new(10).unwrap()), + original_len: 1024, + }; + + // Calculate expected size including null buffer + // Note: Arrow's Buffer might allocate slightly more than null_bitmap_len_bytes + // We use the actual buffer capacity for a more precise comparison + let actual_null_buffer_size = nulls.as_ref().map_or(0, |nb| nb.buffer().capacity()); + let expected_size_with_nulls = size_of::>() + + scalar_buffer.inner().capacity() + + actual_null_buffer_size; + + assert_eq!( + bit_packed_with_nulls.get_array_memory_size(), + expected_size_with_nulls, + "Memory size mismatch with nulls" + ); + } +} diff --git a/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/liquid_array/raw/mod.rs b/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/liquid_array/raw/mod.rs new file mode 100644 index 0000000000000..c7414ae8bc691 --- /dev/null +++ b/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/liquid_array/raw/mod.rs @@ -0,0 +1,5 @@ +//! Low level array primitives. +//! You should not use this module directly. +//! Instead, use `liquid_cache_datafusion_server` or `liquid_cache_datafusion_client` to interact with LiquidCache. +pub(super) mod bit_pack_array; +pub use bit_pack_array::BitPackedArray; diff --git a/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/sync.rs b/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/sync.rs new file mode 100644 index 0000000000000..4e2a741253538 --- /dev/null +++ b/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/sync.rs @@ -0,0 +1,2 @@ +#[allow(unused_imports)] +pub use std::{sync::*, thread}; diff --git a/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/utils/mod.rs b/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/utils/mod.rs new file mode 100644 index 0000000000000..3e8c51a66771b --- /dev/null +++ b/sandbox/libs/dataformat-native/rust/liquid-cache/core/src/utils/mod.rs @@ -0,0 +1,32 @@ +//! Utility functions for the storage module. + +use std::num::NonZero; + +use datafusion_common::ScalarValue; + +/// Get the bit width for a given max value. +/// Returns 1 if the max value is 0. +/// Returns 64 - max_value.leading_zeros() as u8 otherwise. +pub(crate) fn get_bit_width(max_value: u64) -> NonZero { + if max_value == 0 { + // todo: here we actually should return 0, as we should just use constant encoding. + // but that's not implemented yet. + NonZero::new(1).unwrap() + } else { + NonZero::new(64 - max_value.leading_zeros() as u8).unwrap() + } +} + +pub(crate) fn get_bytes_needle(value: &ScalarValue) -> Option> { + match value { + ScalarValue::Utf8(Some(v)) => Some(v.as_bytes().to_vec()), + ScalarValue::Utf8View(Some(v)) => Some(v.as_bytes().to_vec()), + ScalarValue::LargeUtf8(Some(v)) => Some(v.as_bytes().to_vec()), + ScalarValue::Binary(Some(v)) => Some(v.clone()), + ScalarValue::BinaryView(Some(v)) => Some(v.clone()), + ScalarValue::FixedSizeBinary(_, Some(v)) => Some(v.clone()), + ScalarValue::LargeBinary(Some(v)) => Some(v.clone()), + ScalarValue::Dictionary(_, value) => get_bytes_needle(value.as_ref()), + _ => None, + } +} diff --git a/sandbox/libs/dataformat-native/rust/liquid-cache/datafusion/Cargo.toml b/sandbox/libs/dataformat-native/rust/liquid-cache/datafusion/Cargo.toml new file mode 100644 index 0000000000000..80950a688c28d --- /dev/null +++ b/sandbox/libs/dataformat-native/rust/liquid-cache/datafusion/Cargo.toml @@ -0,0 +1,34 @@ +# SPDX-License-Identifier: Apache-2.0 +# +# Vendored in-memory subset of liquid-cache DataFusion integration. +# Provenance: https://github.com/cocosz/liquid-cache branch lc-opensearch-df54-v2 +# commit 8311bccc258756127adbebee3e54140b60fc09ba. See ../README.md. + +[package] +name = "opensearch-liquid-cache-datafusion" +version = "0.1.0" +edition = "2024" +license = "Apache-2.0" +publish = false + +[lib] +# Keep the upstream crate name so the integration keeps its +# `use liquid_cache_datafusion::...` paths. See ../README.md. +name = "liquid_cache_datafusion" +path = "src/lib.rs" + +[dependencies] +ahash = { workspace = true } +arrow = { workspace = true } +arrow-schema = { workspace = true } +bytes = { workspace = true } +datafusion = { workspace = true } +futures = { workspace = true } +log = { workspace = true } +object_store = { workspace = true } +parquet = { workspace = true } +tokio = { workspace = true } +opensearch-liquid-cache-core = { workspace = true } + +[dev-dependencies] +tempfile = { workspace = true } diff --git a/sandbox/libs/dataformat-native/rust/liquid-cache/datafusion/src/cache/column.rs b/sandbox/libs/dataformat-native/rust/liquid-cache/datafusion/src/cache/column.rs new file mode 100644 index 0000000000000..ea029cfb61e13 --- /dev/null +++ b/sandbox/libs/dataformat-native/rust/liquid-cache/datafusion/src/cache/column.rs @@ -0,0 +1,196 @@ +use arrow::{ + array::{Array, ArrayRef, BooleanArray}, + buffer::BooleanBuffer, + compute::prep_null_mask_filter, + record_batch::RecordBatch, +}; +use arrow_schema::{ArrowError, DataType, Field, Schema}; +use liquid_cache::cache::{CacheFull, LiquidCache, LiquidExpr}; +use parquet::arrow::arrow_reader::ArrowPredicate; + +use crate::{ + LiquidPredicate, + cache::{BatchID, ColumnAccessPath, ParquetArrayID}, +}; +use std::sync::Arc; + +/// A column in the cache. +#[derive(Debug)] +pub struct CachedColumn { + cache_store: Arc, + field: Arc, + column_path: ColumnAccessPath, + /// Whether this column is used in a predicate (WHERE clause). + /// In predicate-only mode, only predicate columns are cached. + is_predicate_column: bool, +} + +/// A reference to a cached column. +pub type CachedColumnRef = Arc; + +/// Error type for inserting an arrow array into the cache. +#[derive(Debug)] +pub enum InsertArrowArrayError { + /// The array is already cached. + AlreadyCached, + /// The cache does not have enough memory budget to accept the array. + CacheFull, +} + +impl From for InsertArrowArrayError { + fn from(_: CacheFull) -> Self { + Self::CacheFull + } +} + +impl CachedColumn { + pub(crate) fn new( + field: Arc, + cache_store: Arc, + column_access_path: ColumnAccessPath, + is_predicate_column: bool, + ) -> Self { + Self { + field, + cache_store, + column_path: column_access_path, + is_predicate_column, + } + } + + /// row_id must be on a batch boundary. + pub(crate) fn entry_id(&self, batch_id: BatchID) -> ParquetArrayID { + self.column_path.entry_id(batch_id) + } + + pub(crate) fn is_cached(&self, batch_id: BatchID) -> bool { + self.cache_store.is_cached(&self.entry_id(batch_id).into()) + } + + /// Returns the Arrow field metadata for this cached column. + pub fn field(&self) -> Arc { + self.field.clone() + } + + /// Returns whether this column is a predicate column (used in WHERE clause). + pub fn is_predicate_column(&self) -> bool { + self.is_predicate_column + } + + fn array_to_record_batch(&self, array: ArrayRef) -> RecordBatch { + let schema = Arc::new(Schema::new(vec![self.field.clone()])); + RecordBatch::try_new(schema, vec![array]).unwrap() + } + + /// Evaluates a predicate on a cached column. + pub fn eval_predicate_with_filter( + &self, + batch_id: BatchID, + filter: &BooleanBuffer, + predicate: &mut LiquidPredicate, + ) -> Option> { + let entry_id = self.entry_id(batch_id).into(); + let liquid_expr = LiquidExpr::try_new( + Arc::clone(predicate.physical_expr()), + self.field.data_type(), + ); + + if let Some(liquid_expr) = liquid_expr + && let Some(boolean_array) = self + .cache_store + .eval_predicate(&entry_id, &liquid_expr) + .with_selection(filter) + .read() + { + let predicate_filter = match boolean_array.null_count() { + 0 => boolean_array, + _ => prep_null_mask_filter(&boolean_array), + }; + return Some(Ok(predicate_filter)); + } + + let array = self.get_arrow_array_with_filter(batch_id, filter)?; + let record_batch = self.array_to_record_batch(array); + let boolean_array = match predicate.evaluate(record_batch) { + Ok(arr) => arr, + Err(err) => return Some(Err(err)), + }; + let predicate_filter = match boolean_array.null_count() { + 0 => boolean_array, + _ => prep_null_mask_filter(&boolean_array), + }; + Some(Ok(predicate_filter)) + } + + pub(crate) fn liquid_expr_for_predicate( + &self, + expr: Arc, + ) -> Option { + LiquidExpr::try_new(expr, self.field.data_type()) + } + + /// Get an arrow array with a filter applied. + /// Returns None for non-predicate columns or string predicate columns + /// (only numeric predicate columns are cached and served from cache). + pub fn get_arrow_array_with_filter( + &self, + batch_id: BatchID, + filter: &BooleanBuffer, + ) -> Option { + if !self.is_predicate_column || is_string_type(self.field.data_type()) { + return None; + } + let entry_id = self.entry_id(batch_id).into(); + let result = self + .cache_store + .get(&entry_id) + .with_selection(filter) + .read(); + if result.is_some() { + self.cache_store.observer().runtime_stats().incr_cache_hit(); + } else { + self.cache_store + .observer() + .runtime_stats() + .incr_cache_miss(); + } + result + } + + #[cfg(test)] + pub(crate) fn get_arrow_array_test_only(&self, batch_id: BatchID) -> Option { + let entry_id = self.entry_id(batch_id).into(); + self.cache_store.get(&entry_id).read() + } + + /// Insert an array into the cache. + /// Only numeric predicate columns are cached; string predicates and + /// non-predicate columns return CacheFull. + pub fn insert( + self: &Arc, + batch_id: BatchID, + array: ArrayRef, + ) -> Result<(), InsertArrowArrayError> { + if !self.is_predicate_column || is_string_type(self.field.data_type()) { + return Err(InsertArrowArrayError::CacheFull); + } + + if self.is_cached(batch_id) { + return Err(InsertArrowArrayError::AlreadyCached); + } + + self.cache_store + .insert(self.entry_id(batch_id).into(), array) + .execute()?; + Ok(()) + } +} + +fn is_string_type(data_type: &DataType) -> bool { + match data_type { + DataType::Utf8 | DataType::Utf8View | DataType::LargeUtf8 => true, + DataType::Binary | DataType::BinaryView | DataType::LargeBinary => true, + DataType::Dictionary(_, value_type) => is_string_type(value_type.as_ref()), + _ => false, + } +} diff --git a/sandbox/libs/dataformat-native/rust/liquid-cache/datafusion/src/cache/id.rs b/sandbox/libs/dataformat-native/rust/liquid-cache/datafusion/src/cache/id.rs new file mode 100644 index 0000000000000..b52379dea05be --- /dev/null +++ b/sandbox/libs/dataformat-native/rust/liquid-cache/datafusion/src/cache/id.rs @@ -0,0 +1,310 @@ +use std::ops::Deref; + +use liquid_cache::cache::EntryID; + +/// This is a unique identifier for a row in a parquet file. +#[repr(C, align(8))] +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Ord, PartialOrd)] +pub struct ParquetArrayID { + file_id: u16, + rg_id: u16, + col_id: u16, + batch_id: BatchID, +} + +impl From for usize { + fn from(id: ParquetArrayID) -> Self { + (id.file_id as usize) << 48 + | (id.rg_id as usize) << 32 + | (id.col_id as usize) << 16 + | (id.batch_id.v as usize) + } +} + +impl From for ParquetArrayID { + fn from(value: usize) -> Self { + Self { + file_id: (value >> 48) as u16, + rg_id: ((value >> 32) & 0xFFFF) as u16, + col_id: ((value >> 16) & 0xFFFF) as u16, + batch_id: BatchID::from_raw((value & 0xFFFF) as u16), + } + } +} + +impl ParquetArrayID {} + +impl From for EntryID { + fn from(id: ParquetArrayID) -> Self { + EntryID::from(usize::from(id)) + } +} + +impl From for ParquetArrayID { + fn from(id: EntryID) -> Self { + ParquetArrayID::from(usize::from(id)) + } +} + +const _: () = assert!(std::mem::size_of::() == 8); +const _: () = assert!(std::mem::align_of::() == 8); + +impl ParquetArrayID { + /// Creates a new CacheEntryID. + pub fn new(file_id: u64, row_group_id: u64, column_id: u64, batch_id: BatchID) -> Self { + debug_assert!(file_id <= u16::MAX as u64); + debug_assert!(row_group_id <= u16::MAX as u64); + debug_assert!(column_id <= u16::MAX as u64); + Self { + file_id: file_id as u16, + rg_id: row_group_id as u16, + col_id: column_id as u16, + batch_id, + } + } + + /// Get the batch id. + pub fn batch_id_inner(&self) -> u64 { + self.batch_id.v as u64 + } + + /// Get the file id. + pub fn file_id_inner(&self) -> u64 { + self.file_id as u64 + } + + /// Get the row group id. + pub fn row_group_id_inner(&self) -> u64 { + self.rg_id as u64 + } + + /// Get the column id. + pub fn column_id_inner(&self) -> u64 { + self.col_id as u64 + } + + /// Returns a human-readable string representation of this entry + /// (e.g. `file_1/rg_2/col_3/batch_4`). + pub fn display_path(&self) -> String { + format!( + "file_{}/rg_{}/col_{}/batch_{}", + self.file_id_inner(), + self.row_group_id_inner(), + self.column_id_inner(), + self.batch_id_inner() + ) + } +} + +/// BatchID is a unique identifier for a batch of rows, +/// it is row id divided by the batch size. +/// +// It's very easy to misinterpret this as row id, so we use new type idiom to avoid confusion: +// https://doc.rust-lang.org/rust-by-example/generics/new_types.html +#[repr(C, align(2))] +#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash, Ord, PartialOrd)] +pub struct BatchID { + v: u16, +} + +impl BatchID { + /// Creates a new BatchID from a row id and a batch size. + /// The row id is at the boundary of the batch. + pub fn from_row_id(row_id: usize, batch_size: usize) -> Self { + Self { + v: (row_id / batch_size) as u16, + } + } + + /// Creates a new BatchID from a raw value. + pub fn from_raw(v: u16) -> Self { + Self { v } + } + + /// Increment the batch id. + pub fn inc(&mut self) { + debug_assert!(self.v < u16::MAX); + self.v += 1; + } +} + +impl Deref for BatchID { + type Target = u16; + + fn deref(&self) -> &Self::Target { + &self.v + } +} + +/// Column access path. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Ord, PartialOrd)] +pub struct ColumnAccessPath { + file_id: u16, + rg_id: u16, + col_id: u16, +} + +impl ColumnAccessPath { + /// Create a new instance of ColumnAccessPath. + pub fn new(file_id: u64, row_group_id: u64, column_id: u64) -> Self { + debug_assert!(file_id <= u16::MAX as u64); + debug_assert!(row_group_id <= u16::MAX as u64); + debug_assert!(column_id <= u16::MAX as u64); + Self { + file_id: file_id as u16, + rg_id: row_group_id as u16, + col_id: column_id as u16, + } + } + + /// Get the file id. + fn file_id_inner(&self) -> u64 { + self.file_id as u64 + } + + /// Get the row group id. + fn row_group_id_inner(&self) -> u64 { + self.rg_id as u64 + } + + /// Get the column id. + fn column_id_inner(&self) -> u64 { + self.col_id as u64 + } + + /// Get the entry id. + pub fn entry_id(&self, batch_id: BatchID) -> ParquetArrayID { + ParquetArrayID::new( + self.file_id_inner(), + self.row_group_id_inner(), + self.column_id_inner(), + batch_id, + ) + } +} + +impl From for ColumnAccessPath { + fn from(value: ParquetArrayID) -> Self { + Self { + file_id: value.file_id_inner() as u16, + rg_id: value.row_group_id_inner() as u16, + col_id: value.column_id_inner() as u16, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_cache_entry_id_new_and_getters() { + let file_id = 10u64; + let row_group_id = 20u64; + let column_id = 30u64; + let batch_id = BatchID::from_raw(40); + let entry_id = ParquetArrayID::new(file_id, row_group_id, column_id, batch_id); + + assert_eq!(entry_id.file_id_inner(), file_id); + assert_eq!(entry_id.row_group_id_inner(), row_group_id); + assert_eq!(entry_id.column_id_inner(), column_id); + assert_eq!(entry_id.batch_id_inner(), *batch_id as u64); + } + + #[test] + fn test_cache_entry_id_boundaries() { + let file_id = u16::MAX as u64; + let row_group_id = 0u64; + let column_id = u16::MAX as u64; + let batch_id = BatchID::from_raw(0); + let entry_id = ParquetArrayID::new(file_id, row_group_id, column_id, batch_id); + + assert_eq!(entry_id.file_id_inner(), file_id); + assert_eq!(entry_id.row_group_id_inner(), row_group_id); + assert_eq!(entry_id.column_id_inner(), column_id); + assert_eq!(entry_id.batch_id_inner(), *batch_id as u64); + } + + #[test] + #[should_panic] + fn test_cache_entry_id_new_panic_file_id() { + ParquetArrayID::new((u16::MAX as u64) + 1, 0, 0, BatchID::from_raw(0)); + } + + #[test] + #[should_panic] + fn test_cache_entry_id_new_panic_row_group_id() { + ParquetArrayID::new(0, (u16::MAX as u64) + 1, 0, BatchID::from_raw(0)); + } + + #[test] + #[should_panic] + fn test_cache_entry_id_new_panic_column_id() { + ParquetArrayID::new(0, 0, (u16::MAX as u64) + 1, BatchID::from_raw(0)); + } + + #[test] + fn test_cache_entry_id_display_path() { + let entry_id = ParquetArrayID::new(1, 2, 3, BatchID::from_raw(4)); + assert_eq!(entry_id.display_path(), "file_1/rg_2/col_3/batch_4"); + } + + #[test] + fn test_batch_id_from_row_id() { + let batch_id = BatchID::from_row_id(256, 128); + assert_eq!(batch_id.v, 2); + } + + #[test] + fn test_batch_id_from_raw() { + let batch_id = BatchID::from_raw(5); + assert_eq!(batch_id.v, 5); + } + + #[test] + fn test_batch_id_inc() { + let mut batch_id = BatchID::from_raw(10); + batch_id.inc(); + assert_eq!(batch_id.v, 11); + } + + #[test] + #[should_panic] + fn test_batch_id_inc_overflow() { + let mut batch_id = BatchID::from_raw(u16::MAX); + // Should panic because incrementing exceeds u16::MAX + batch_id.inc(); + } + + #[test] + fn test_batch_id_deref() { + let batch_id = BatchID::from_raw(15); + assert_eq!(*batch_id, 15); + } + + #[test] + fn test_column_path_from_cache_entry_id() { + let entry_id = ParquetArrayID::new(1, 2, 3, BatchID::from_raw(4)); + let column_path: ColumnAccessPath = entry_id.into(); + + assert_eq!(column_path.file_id, 1); + assert_eq!(column_path.rg_id, 2); + assert_eq!(column_path.col_id, 3); + } + + #[test] + fn test_column_path_entry_id() { + let file_id = 5u64; + let row_group_id = 6u64; + let column_id = 7u64; + let column_path = ColumnAccessPath::new(file_id, row_group_id, column_id); + + let batch_id = BatchID::from_raw(8); + let entry_id = column_path.entry_id(batch_id); + + assert_eq!(entry_id.file_id_inner(), file_id); + assert_eq!(entry_id.row_group_id_inner(), row_group_id); + assert_eq!(entry_id.column_id_inner(), column_id); + assert_eq!(entry_id.batch_id_inner(), *batch_id as u64); + } +} diff --git a/sandbox/libs/dataformat-native/rust/liquid-cache/datafusion/src/cache/mod.rs b/sandbox/libs/dataformat-native/rust/liquid-cache/datafusion/src/cache/mod.rs new file mode 100644 index 0000000000000..04f6f3161c885 --- /dev/null +++ b/sandbox/libs/dataformat-native/rust/liquid-cache/datafusion/src/cache/mod.rs @@ -0,0 +1,480 @@ +//! This module contains the cache implementation for the Parquet reader. +//! + +use crate::reader::{LiquidPredicate, extract_multi_column_or}; +use crate::sync::Mutex; +use ahash::AHashMap; +use arrow::array::{BooleanArray, RecordBatch}; +use arrow::buffer::BooleanBuffer; +use arrow_schema::{ArrowError, Field, Schema, SchemaRef}; +use liquid_cache::cache::{ + CachePolicy, CacheStats, LiquidCache, LiquidCacheBuilder, SqueezePolicy, +}; +use parquet::arrow::arrow_reader::ArrowPredicate; +use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering}; + +mod column; +mod id; +mod stats; + +pub use column::{CachedColumn, CachedColumnRef, InsertArrowArrayError}; +pub(crate) use id::ColumnAccessPath; +pub use id::{BatchID, ParquetArrayID}; + +#[derive(Default, Debug)] +struct ColumnMaps { + // invariant: Arc::ptr_eq(map[field.name()], map[field.id()]) + by_id: AHashMap, + by_name: AHashMap, +} + +/// A row group in the cache. +#[derive(Debug)] +pub struct CachedRowGroup { + columns: ColumnMaps, + cache_store: Arc, +} + +impl CachedRowGroup { + /// Create a new row group. + /// The column_ids are the indices of the columns in the file schema. + /// So they may not start from 0. + fn new( + cache_store: Arc, + row_group_idx: u64, + file_idx: u64, + columns: &[(u64, Arc, bool)], + ) -> Self { + let mut column_maps = ColumnMaps::default(); + for (column_id, field, is_predicate_column) in columns { + let column_access_path = ColumnAccessPath::new(file_idx, row_group_idx, *column_id); + let column = Arc::new(CachedColumn::new( + Arc::clone(field), + Arc::clone(&cache_store), + column_access_path, + *is_predicate_column, + )); + column_maps.by_id.insert(*column_id, column.clone()); + column_maps.by_name.insert(field.name().to_string(), column); + } + + Self { + columns: column_maps, + cache_store, + } + } + + /// Returns the batch size configured for this cached row group. + pub fn batch_size(&self) -> usize { + self.cache_store.config().batch_size() + } + + /// Get a column from the row group. + pub fn get_column(&self, column_id: u64) -> Option { + self.columns.by_id.get(&column_id).cloned() + } + + /// Get a column from the row group by its field name. + pub fn get_column_by_name(&self, column_name: &str) -> Option { + if let Some(column) = self.columns.by_name.get(column_name) { + return Some(column.clone()); + } + + // DataFusion may carry qualified names in physical expressions + // (e.g. "table.col"), while cache fields are keyed by file schema names. + let unqualified = column_name.rsplit('.').next().unwrap_or(column_name); + self.columns.by_name.get(unqualified).cloned() + } + + /// Evaluate a predicate on a row group. + pub fn evaluate_selection_with_predicate( + &self, + batch_id: BatchID, + selection: &BooleanBuffer, + predicate: &mut LiquidPredicate, + ) -> Option> { + let column_ids = predicate.predicate_column_ids(); + + if column_ids.len() == 1 { + // If we only have one column, we can short-circuit and try to evaluate the predicate on encoded data. + let column_id = column_ids[0]; + let cache = self.get_column(column_id as u64)?; + return cache.eval_predicate_with_filter(batch_id, selection, predicate); + } else if column_ids.len() >= 2 { + // Try to extract multiple column-literal expressions from OR structure + if let Some(column_exprs) = + extract_multi_column_or(predicate.physical_expr_physical_column_index()) + { + let mut combined_buffer: Option = None; + + for (col_name, expr) in column_exprs { + let column = self.get_column_by_name(col_name)?; + let liquid_expr = column.liquid_expr_for_predicate(Arc::clone(&expr)); + let liquid_expr = match liquid_expr { + Some(expr) => expr, + None => { + combined_buffer = None; + break; + } + }; + let entry_id = column.entry_id(batch_id).into(); + let liquid_array = self.cache_store.try_read_liquid(&entry_id); + let liquid_array = match liquid_array { + None => { + combined_buffer = None; + break; + } + Some(array) => array, + }; + let buffer = liquid_array.try_eval_predicate(&liquid_expr, selection); + + combined_buffer = Some(match combined_buffer { + None => buffer, + Some(existing) => { + arrow::compute::kernels::boolean::or_kleene(&existing, &buffer).ok()? + } + }); + } + + if let Some(result) = combined_buffer { + return Some(Ok(result)); + } + } + } + // Otherwise, we need to first convert the data into arrow arrays. + let mut arrays = Vec::new(); + let mut fields = Vec::new(); + for column_id in column_ids { + let column = self.get_column(column_id as u64)?; + let array = column.get_arrow_array_with_filter(batch_id, selection)?; + arrays.push(array); + fields.push(column.field()); + } + let schema = Arc::new(Schema::new(fields)); + let record_batch = RecordBatch::try_new(schema, arrays).unwrap(); + let boolean_array = predicate.evaluate(record_batch).unwrap(); + Some(Ok(boolean_array)) + } +} + +pub(crate) type CachedRowGroupRef = Arc; + +/// A file in the cache. +#[derive(Debug)] +pub struct CachedFile { + cache_store: Arc, + file_id: u64, + file_schema: SchemaRef, +} + +impl CachedFile { + fn new(cache_store: Arc, file_id: u64, file_schema: SchemaRef) -> Self { + Self { + cache_store, + file_id, + file_schema, + } + } + + /// Create a row group handle scoped to the current query context. + pub fn create_row_group( + &self, + row_group_id: u64, + predicate_column_ids: Vec, + ) -> CachedRowGroupRef { + let columns: Vec<(u64, Arc, bool)> = self + .file_schema + .fields() + .iter() + .enumerate() + .map(|(idx, field)| { + let is_predicate_column = predicate_column_ids.contains(&idx); + (idx as u64, Arc::clone(field), is_predicate_column) + }) + .collect(); + + Arc::new(CachedRowGroup::new( + self.cache_store.clone(), + row_group_id, + self.file_id, + &columns, + )) + } + + /// Return the configured cache batch size. + pub fn batch_size(&self) -> usize { + self.cache_store.config().batch_size() + } + + /// Return the full file schema tracked by the cache entry. + pub fn schema(&self) -> SchemaRef { + Arc::clone(&self.file_schema) + } +} + +/// A reference to a cached file. +pub(crate) type CachedFileRef = Arc; + +/// The main cache structure. +#[derive(Debug)] +pub struct LiquidCacheParquet { + /// Map file path to file id. + files: Mutex>, + + cache_store: Arc, + + current_file_id: AtomicU64, +} + +/// A reference to the main cache structure. +pub type LiquidCacheParquetRef = Arc; + +impl LiquidCacheParquet { + /// Create a new in-memory cache for parquet files. + pub fn new( + batch_size: usize, + max_memory_bytes: usize, + cache_policy: Box, + squeeze_policy: Box, + ) -> Self { + assert!(batch_size.is_power_of_two()); + let cache_storage = LiquidCacheBuilder::new() + .with_batch_size(batch_size) + .with_max_memory_bytes(max_memory_bytes) + .with_squeeze_policy(squeeze_policy) + .with_cache_policy(cache_policy) + .build(); + + LiquidCacheParquet { + files: Mutex::new(AHashMap::new()), + cache_store: cache_storage, + current_file_id: AtomicU64::new(0), + } + } + + /// Register a file in the cache. + pub fn register_or_get_file( + &self, + file_path: String, + full_file_schema: SchemaRef, + ) -> CachedFileRef { + let mut files = self.files.lock().unwrap(); + let file_id = *files + .entry(file_path.clone()) + .or_insert_with(|| self.current_file_id.fetch_add(1, Ordering::Relaxed)); + drop(files); + + Arc::new(CachedFile::new( + self.cache_store.clone(), + file_id, + full_file_schema, + )) + } + + /// Get the batch size of the cache. + pub fn batch_size(&self) -> usize { + self.cache_store.config().batch_size() + } + + /// Get the max memory bytes of the cache. + pub fn max_memory_bytes(&self) -> usize { + self.cache_store.config().max_memory_bytes() + } + + /// Get the memory usage of the cache in bytes. + pub fn memory_usage_bytes(&self) -> usize { + self.cache_store.budget().memory_usage_bytes() + } + + /// Get a snapshot of the cache statistics. + pub fn stats(&self) -> CacheStats { + self.cache_store.stats() + } + + /// Reset the cache. + /// + /// # Safety + /// This is unsafe because resetting the cache while other threads are using the cache may cause undefined behavior. + /// You should only call this when no one else is using the cache. + pub unsafe fn reset(&self) { + let mut files = self.files.lock().unwrap(); + files.clear(); + self.cache_store.reset(); + } + + /// Get the storage of the cache. + pub fn storage(&self) -> &Arc { + &self.cache_store + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::cache::{CachedRowGroupRef, LiquidCacheParquet}; + use crate::reader::FilterCandidateBuilder; + use arrow::array::Int32Array; + use arrow::buffer::BooleanBuffer; + use arrow::datatypes::{DataType, Field, Schema}; + use arrow::record_batch::RecordBatch; + use datafusion::common::ScalarValue; + use datafusion::logical_expr::Operator; + use datafusion::physical_expr::PhysicalExpr; + use datafusion::physical_expr::expressions::{BinaryExpr, Literal}; + use datafusion::physical_plan::expressions::Column; + use liquid_cache::cache::TranscodeEvict; + use liquid_cache::cache_policies::LiquidPolicy; + use parquet::arrow::ArrowWriter; + use parquet::arrow::arrow_reader::{ArrowReaderMetadata, ArrowReaderOptions}; + use std::sync::Arc; + + fn setup_cache(batch_size: usize, schema: SchemaRef) -> CachedRowGroupRef { + let cache = LiquidCacheParquet::new( + batch_size, + usize::MAX, + Box::new(LiquidPolicy::new()), + Box::new(TranscodeEvict), + ); + let file = cache.register_or_get_file("test".to_string(), schema); + // Mark all columns as predicate columns so they are cacheable in tests. + let all_columns: Vec = (0..file.schema().fields().len()).collect(); + file.create_row_group(0, all_columns) + } + + #[test] + fn evaluate_or_on_cached_columns() { + let batch_size = 4; + + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + ])); + let row_group = setup_cache(batch_size, schema.clone()); + + let col_a = row_group.get_column(0).unwrap(); + let col_b = row_group.get_column(1).unwrap(); + + let batch_id = BatchID::from_row_id(0, batch_size); + + let array_a = Arc::new(Int32Array::from(vec![1, 2, 3, 4])); + let array_b = Arc::new(Int32Array::from(vec![10, 20, 30, 40])); + + assert!(col_a.insert(batch_id, array_a.clone()).is_ok()); + assert!(col_b.insert(batch_id, array_b.clone()).is_ok()); + + // build parquet metadata for predicate construction + let tmp_meta = tempfile::NamedTempFile::new().unwrap(); + let mut writer = + ArrowWriter::try_new(tmp_meta.reopen().unwrap(), Arc::clone(&schema), None).unwrap(); + let batch = RecordBatch::try_new(Arc::clone(&schema), vec![array_a, array_b]).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + let file_reader = std::fs::File::open(tmp_meta.path()).unwrap(); + let metadata = ArrowReaderMetadata::load(&file_reader, ArrowReaderOptions::new()).unwrap(); + + // expression a = 3 OR b = 20 + let expr_a: Arc = Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::Eq, + Arc::new(Literal::new(ScalarValue::Int32(Some(3)))), + )); + let expr_b: Arc = Arc::new(BinaryExpr::new( + Arc::new(Column::new("b", 1)), + Operator::Eq, + Arc::new(Literal::new(ScalarValue::Int32(Some(20)))), + )); + let expr: Arc = Arc::new(BinaryExpr::new(expr_a, Operator::Or, expr_b)); + + let builder = FilterCandidateBuilder::new(expr, Arc::clone(&schema)); + let candidate = builder.build(metadata.metadata()).unwrap().unwrap(); + let projection = candidate.projection(metadata.metadata()); + let mut predicate = LiquidPredicate::try_new(candidate, projection).unwrap(); + + let selection = BooleanBuffer::new_set(batch_size); + let result = row_group + .evaluate_selection_with_predicate(batch_id, &selection, &mut predicate) + .unwrap() + .unwrap(); + + let expected = BooleanBuffer::collect_bool(batch_size, |i| i == 1 || i == 2).into(); + assert_eq!(result, expected); + } + + #[test] + fn evaluate_three_column_or() { + let batch_size = 8; + + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + Field::new("c", DataType::Int32, false), + ])); + + let row_group = setup_cache(batch_size, schema.clone()); + + let col_a = row_group.get_column(0).unwrap(); + let col_b = row_group.get_column(1).unwrap(); + let col_c = row_group.get_column(2).unwrap(); + + let batch_id = BatchID::from_row_id(0, batch_size); + + let array_a = Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8])); + let array_b = Arc::new(Int32Array::from(vec![10, 20, 30, 40, 50, 60, 70, 80])); + let array_c = Arc::new(Int32Array::from(vec![ + 100, 200, 300, 400, 500, 600, 700, 800, + ])); + + assert!(col_a.insert(batch_id, array_a.clone()).is_ok()); + assert!(col_b.insert(batch_id, array_b.clone()).is_ok()); + assert!(col_c.insert(batch_id, array_c.clone()).is_ok()); + + // build parquet metadata for predicate construction + let tmp_meta = tempfile::NamedTempFile::new().unwrap(); + let mut writer = + ArrowWriter::try_new(tmp_meta.reopen().unwrap(), Arc::clone(&schema), None).unwrap(); + let batch = + RecordBatch::try_new(Arc::clone(&schema), vec![array_a, array_b, array_c]).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + let file_reader = std::fs::File::open(tmp_meta.path()).unwrap(); + let metadata = ArrowReaderMetadata::load(&file_reader, ArrowReaderOptions::new()).unwrap(); + + // expression: a = 2 OR b = 40 OR c = 600 + let expr_a: Arc = Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::Eq, + Arc::new(Literal::new(ScalarValue::Int32(Some(2)))), + )); + let expr_b: Arc = Arc::new(BinaryExpr::new( + Arc::new(Column::new("b", 1)), + Operator::Eq, + Arc::new(Literal::new(ScalarValue::Int32(Some(40)))), + )); + let expr_c: Arc = Arc::new(BinaryExpr::new( + Arc::new(Column::new("c", 2)), + Operator::Eq, + Arc::new(Literal::new(ScalarValue::Int32(Some(600)))), + )); + + // Build nested OR: (a = 2 OR b = 40) OR c = 600 + let expr_ab = Arc::new(BinaryExpr::new(expr_a, Operator::Or, expr_b)); + let expr: Arc = Arc::new(BinaryExpr::new(expr_ab, Operator::Or, expr_c)); + + let builder = FilterCandidateBuilder::new(expr, Arc::clone(&schema)); + let candidate = builder.build(metadata.metadata()).unwrap().unwrap(); + let projection = candidate.projection(metadata.metadata()); + let mut predicate = LiquidPredicate::try_new(candidate, projection).unwrap(); + + let selection = BooleanBuffer::new_set(batch_size); + let result = row_group + .evaluate_selection_with_predicate(batch_id, &selection, &mut predicate) + .unwrap() + .unwrap(); + + // Expected: row 1 (a=2), row 3 (b=40), row 5 (c=600) -> indices 1, 3, 5 + let expected = + BooleanBuffer::collect_bool(batch_size, |i| i == 1 || i == 3 || i == 5).into(); + assert_eq!(result, expected); + } +} diff --git a/sandbox/libs/dataformat-native/rust/liquid-cache/datafusion/src/cache/stats.rs b/sandbox/libs/dataformat-native/rust/liquid-cache/datafusion/src/cache/stats.rs new file mode 100644 index 0000000000000..59191edee813a --- /dev/null +++ b/sandbox/libs/dataformat-native/rust/liquid-cache/datafusion/src/cache/stats.rs @@ -0,0 +1,268 @@ +use super::LiquidCacheParquet; +use crate::{cache::id::ParquetArrayID, sync::Arc}; +use arrow::array::{ArrayBuilder, RecordBatch, StringBuilder, UInt64Builder}; +use arrow_schema::{DataType, Field, Schema, SchemaRef}; +use liquid_cache::cache::CacheEntry; +use parquet::{ + arrow::ArrowWriter, basic::Compression, errors::ParquetError, + file::properties::WriterProperties, +}; +use std::{fs::File, path::Path}; + +struct StatsWriter { + writer: ArrowWriter, + schema: SchemaRef, + file_path_builder: StringBuilder, + row_group_id_builder: UInt64Builder, + column_id_builder: UInt64Builder, + row_start_id_builder: UInt64Builder, + row_count_builder: UInt64Builder, + memory_size_builder: UInt64Builder, + cache_type_builder: StringBuilder, + reference_count_builder: UInt64Builder, +} + +impl StatsWriter { + fn new(file_path: impl AsRef) -> Result { + let schema = Arc::new(Schema::new(vec![ + Field::new("row_group_id", DataType::UInt64, false), + Field::new("column_id", DataType::UInt64, false), + Field::new("row_start_id", DataType::UInt64, false), + Field::new("row_count", DataType::UInt64, true), + Field::new("memory_size", DataType::UInt64, false), + Field::new("cache_type", DataType::Utf8, false), + Field::new("file_path", DataType::Utf8, false), + Field::new("reference_count", DataType::UInt64, false), + ])); + + let file = File::create(file_path)?; + let write_props = WriterProperties::builder() + .set_compression(Compression::LZ4) + .set_created_by("liquid-cache-stats".to_string()) + .build(); + let writer = ArrowWriter::try_new(file, schema.clone(), Some(write_props))?; + Ok(Self { + writer, + schema, + file_path_builder: StringBuilder::with_capacity(8192, 8192), + row_group_id_builder: UInt64Builder::new(), + column_id_builder: UInt64Builder::new(), + row_start_id_builder: UInt64Builder::new(), + row_count_builder: UInt64Builder::new(), + memory_size_builder: UInt64Builder::new(), + cache_type_builder: StringBuilder::with_capacity(8192, 8192), + reference_count_builder: UInt64Builder::new(), + }) + } + + fn build_batch(&mut self) -> Result { + let row_group_id_array = self.row_group_id_builder.finish(); + let column_id_array = self.column_id_builder.finish(); + let row_start_id_array = self.row_start_id_builder.finish(); + let row_count_array = self.row_count_builder.finish(); + let memory_size_array = self.memory_size_builder.finish(); + let cache_type_array = self.cache_type_builder.finish(); + let file_path_array = self.file_path_builder.finish(); + let reference_count_array = self.reference_count_builder.finish(); + Ok(RecordBatch::try_new( + self.schema.clone(), + vec![ + Arc::new(row_group_id_array), + Arc::new(column_id_array), + Arc::new(row_start_id_array), + Arc::new(row_count_array), + Arc::new(memory_size_array), + Arc::new(cache_type_array), + Arc::new(file_path_array), + Arc::new(reference_count_array), + ], + )?) + } + + #[allow(clippy::too_many_arguments)] + fn append_entry( + &mut self, + file_path: &str, + row_group_id: u64, + column_id: u64, + row_start_id: u64, + row_count: Option, + memory_size: u64, + cache_type: &str, + reference_count: u64, + ) -> Result<(), ParquetError> { + self.row_group_id_builder.append_value(row_group_id); + self.column_id_builder.append_value(column_id); + self.row_start_id_builder.append_value(row_start_id); + self.row_count_builder.append_option(row_count); + self.memory_size_builder.append_value(memory_size); + self.cache_type_builder.append_value(cache_type); + self.reference_count_builder.append_value(reference_count); + self.file_path_builder.append_value(file_path); + if self.row_start_id_builder.len() >= 8192 { + let batch = self.build_batch()?; + self.writer.write(&batch)?; + } + Ok(()) + } + + fn finish(mut self) -> Result<(), ParquetError> { + let batch = self.build_batch()?; + self.writer.write(&batch)?; + self.writer.close()?; + Ok(()) + } +} + +impl LiquidCacheParquet { + /// Get the memory usage of the cache in bytes. + pub fn compute_memory_usage_bytes(&self) -> u64 { + self.cache_store.budget().memory_usage_bytes() as u64 + } + + /// Write the stats of the cache to a parquet file. + pub fn write_stats(&self, parquet_file_path: impl AsRef) -> Result<(), ParquetError> { + let mut writer = StatsWriter::new(parquet_file_path)?; + self.cache_store.for_each_entry(|entry_id, cached_batch| { + let memory_size = cached_batch.memory_usage_bytes(); + let row_count = match cached_batch { + CacheEntry::MemoryArrow(array) => Some(array.len() as u64), + CacheEntry::MemoryLiquid(array) => Some(array.len() as u64), + }; + let cache_type = match cached_batch { + CacheEntry::MemoryArrow(_) => "InMemory", + CacheEntry::MemoryLiquid(_) => "LiquidMemory", + }; + let reference_count = cached_batch.reference_count(); + let entry_id = ParquetArrayID::from(*entry_id); + writer + .append_entry( + &entry_id.display_path(), + entry_id.row_group_id_inner(), + entry_id.column_id_inner(), + entry_id.batch_id_inner() * self.batch_size() as u64, + row_count, + memory_size as u64, + cache_type, + reference_count as u64, + ) + .unwrap(); + }); + + writer.finish()?; + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use std::io::Read; + + use crate::cache::id::BatchID; + + use super::*; + use arrow::{ + array::{Array, AsArray}, + datatypes::UInt64Type, + }; + use bytes::Bytes; + use liquid_cache::{cache::Evict, cache_policies::LiquidPolicy}; + use parquet::arrow::arrow_reader::ParquetRecordBatchReader; + use tempfile::NamedTempFile; + + #[test] + fn test_stats_writer() -> Result<(), ParquetError> { + let cache = LiquidCacheParquet::new( + 1024, + usize::MAX, + Box::new(LiquidPolicy::new()), + Box::new(Evict), + ); + let fields: Vec = (0..8) + .map(|i| Field::new(format!("test_{i}"), DataType::Int32, false)) + .collect(); + let schema = Arc::new(Schema::new(fields)); + let array = Arc::new(arrow::array::Int32Array::from(vec![1, 2, 3])); + let num_rows = 8 * 8 * 8 * 8; + + let mut row_group_id_sum = 0; + let mut column_id_sum = 0; + let mut row_start_id_sum = 0; + let mut row_count_sum = 0; + let mut memory_size_sum = 0; + for file_no in 0..8 { + let file_name = format!("test_{file_no}.parquet"); + let file = cache.register_or_get_file(file_name, schema.clone()); + for rg in 0..8 { + // Mark all columns as predicate columns so inserts are accepted. + let row_group = file.create_row_group(rg, (0..8).collect()); + for col in 0..8 { + let column = row_group.get_column(col).unwrap(); + for batch in 0..8 { + let batch_id = BatchID::from_raw(batch); + assert!(column.insert(batch_id, array.clone()).is_ok()); + row_group_id_sum += rg; + column_id_sum += col; + row_start_id_sum += *batch_id as u64 * cache.batch_size() as u64; + row_count_sum += array.len() as u64; + memory_size_sum += array.get_array_memory_size(); + + if batch.is_multiple_of(2) { + _ = column.get_arrow_array_test_only(batch_id).unwrap(); + } + } + } + } + } + + let mut tmp_file = NamedTempFile::new()?; + cache.write_stats(tmp_file.path())?; + + // Read and verify stats + let mut bytes = Vec::new(); + tmp_file.read_to_end(&mut bytes)?; + let bytes = Bytes::from(bytes); + let reader = ParquetRecordBatchReader::try_new(bytes, 8192)?; + + let batch = reader.into_iter().next().unwrap()?; + assert_eq!(batch.num_rows(), num_rows); + + macro_rules! uint64_col { + ($batch:expr, $col_idx:expr) => { + $batch + .column_by_name($col_idx) + .unwrap() + .as_primitive::() + }; + } + + let row_group_id_array = uint64_col!(batch, "row_group_id"); + let column_id_array = uint64_col!(batch, "column_id"); + let row_start_id_array = uint64_col!(batch, "row_start_id"); + let row_count_array = uint64_col!(batch, "row_count"); + let memory_size_array = uint64_col!(batch, "memory_size"); + + assert_eq!( + row_group_id_array.iter().map(|v| v.unwrap()).sum::(), + row_group_id_sum + ); + assert_eq!( + column_id_array.iter().map(|v| v.unwrap()).sum::(), + column_id_sum + ); + assert_eq!( + row_start_id_array.iter().map(|v| v.unwrap()).sum::(), + row_start_id_sum + ); + assert_eq!( + row_count_array.iter().map(|v| v.unwrap()).sum::(), + row_count_sum + ); + assert_eq!( + memory_size_array.iter().map(|v| v.unwrap()).sum::(), + memory_size_sum as u64 + ); + + Ok(()) + } +} diff --git a/sandbox/libs/dataformat-native/rust/liquid-cache/datafusion/src/lib.rs b/sandbox/libs/dataformat-native/rust/liquid-cache/datafusion/src/lib.rs new file mode 100644 index 0000000000000..2c7356525d901 --- /dev/null +++ b/sandbox/libs/dataformat-native/rust/liquid-cache/datafusion/src/lib.rs @@ -0,0 +1,24 @@ +// SPDX-License-Identifier: Apache-2.0 +// Vendored in-memory subset of liquid-cache DataFusion integration. See ../README.md for provenance. + +//! DataFusion integration for the vendored in-memory liquid cache: +//! a Parquet [`FileSource`](datafusion::datasource::physical_plan::FileSource) +//! replacement that serves decoded batches from the cache, plus a physical +//! optimizer rule that rewrites `ParquetSource` scans to use it. +#![warn(missing_docs)] + +pub mod optimizers; +mod reader; +mod sync; +pub(crate) mod utils; + +pub mod cache; +pub use cache::{LiquidCacheParquet, LiquidCacheParquetRef}; +pub use optimizers::{LocalModeOptimizer, rewrite_data_source_plan}; +pub use reader::plantime::engagement_policy::{ + AlwaysEngagePolicy, CacheEngagementPolicy, DEFAULT_SELECTIVITY_THRESHOLD, EngagementContext, + EngagementDecision, NeverEngagePolicy, SelectivityThresholdPolicy, default_engagement_policy, +}; +pub use reader::plantime::source::pre_seed_metadata_cache; +pub use reader::{FilterCandidateBuilder, LiquidParquetSource, LiquidPredicate, LiquidRowFilter}; +pub use utils::boolean_buffer_and_then; diff --git a/sandbox/libs/dataformat-native/rust/liquid-cache/datafusion/src/optimizers/mod.rs b/sandbox/libs/dataformat-native/rust/liquid-cache/datafusion/src/optimizers/mod.rs new file mode 100644 index 0000000000000..95896c6dc7993 --- /dev/null +++ b/sandbox/libs/dataformat-native/rust/liquid-cache/datafusion/src/optimizers/mod.rs @@ -0,0 +1,314 @@ +//! Optimizers for the Parquet module + +use std::sync::Arc; + +use datafusion::{ + catalog::memory::DataSourceExec, + common::tree_node::{Transformed, TreeNode, TreeNodeRecursion}, + config::ConfigOptions, + datasource::{ + physical_plan::{FileSource, ParquetSource}, + source::DataSource, + }, + physical_optimizer::PhysicalOptimizerRule, + physical_plan::ExecutionPlan, +}; + +use crate::{LiquidCacheParquetRef, LiquidParquetSource}; + +/// Physical optimizer rule for local mode liquid cache +/// +/// This optimizer rewrites DataSourceExec nodes that read Parquet files +/// to use LiquidParquetSource instead of the default ParquetSource +#[derive(Debug)] +pub struct LocalModeOptimizer { + cache: LiquidCacheParquetRef, +} + +impl LocalModeOptimizer { + /// Create an optimizer with an existing cache instance + pub fn new(cache: LiquidCacheParquetRef) -> Self { + Self { cache } + } + + /// Create an optimizer with an existing cache instance + pub fn with_cache(cache: LiquidCacheParquetRef) -> Self { + Self { cache } + } +} + +impl PhysicalOptimizerRule for LocalModeOptimizer { + fn optimize( + &self, + plan: Arc, + _config: &ConfigOptions, + ) -> Result, datafusion::error::DataFusionError> { + Ok(rewrite_data_source_plan(plan, &self.cache)) + } + + fn name(&self) -> &str { + "LocalModeLiquidCacheOptimizer" + } + + fn schema_check(&self) -> bool { + true + } +} + +/// Rewrite the data source plan to use liquid cache. +pub fn rewrite_data_source_plan( + plan: Arc, + cache: &LiquidCacheParquetRef, +) -> Arc { + let rewritten = plan + .transform_up(|node| try_optimize_parquet_source(node, cache)) + .unwrap(); + rewritten.data +} + +/// Returns true if a data type is uncacheable by LC (string/binary). +fn is_uncacheable_type(dt: &arrow_schema::DataType) -> bool { + use arrow_schema::DataType; + matches!( + dt, + DataType::Utf8 + | DataType::Utf8View + | DataType::LargeUtf8 + | DataType::Binary + | DataType::BinaryView + | DataType::LargeBinary + ) || matches!(dt, DataType::Dictionary(_, v) if is_uncacheable_type(v)) +} + +/// Max number of output columns for which LC wrapping is worthwhile. +/// Per-column cache overhead exceeds decode savings for wide projections. +const MAX_LC_COLUMNS: usize = 4; + +fn try_optimize_parquet_source( + plan: Arc, + cache: &LiquidCacheParquetRef, +) -> Result>, datafusion::error::DataFusionError> { + if let Some(data_source_exec) = plan.downcast_ref::() + && let Some((file_scan_config, parquet_source)) = + data_source_exec.downcast_to_file_source::() + { + // Skip LC wrapping if: + // - Output has zero columns (COUNT(*) — just needs row count from metadata) + // - Too many output columns (cache overhead exceeds decode savings) + // - ANY output column is string/binary (LC can't cache, fallback negates hits) + // - Predicate references a string column + let output_schema = plan.schema(); + if output_schema.fields().is_empty() { + log::debug!("[LC-Optimizer] SKIP: empty projection (COUNT(*))"); + return Ok(Transformed::no(plan)); + } + + if output_schema.fields().len() > MAX_LC_COLUMNS { + log::debug!( + "[LC-Optimizer] SKIP: too many columns ({} > {})", + output_schema.fields().len(), + MAX_LC_COLUMNS + ); + return Ok(Transformed::no(plan)); + } + + let has_string_output = output_schema + .fields() + .iter() + .any(|f| is_uncacheable_type(f.data_type())); + + let predicate_has_string = parquet_source.filter().is_some_and(|pred| { + use datafusion::physical_expr::utils::collect_columns; + let file_schema = file_scan_config.file_schema(); + let cols = collect_columns(&pred); + cols.iter().any(|col| { + file_schema + .fields() + .get(col.index()) + .is_some_and(|f| is_uncacheable_type(f.data_type())) + }) + }); + + if has_string_output || predicate_has_string { + log::debug!( + "[LC-Optimizer] SKIP: string_in_output={}, string_in_predicate={}, output_cols={}", + has_string_output, + predicate_has_string, + output_schema.fields().len() + ); + return Ok(Transformed::no(plan)); + } + + let num_fields = output_schema.fields().len(); + let has_predicate = parquet_source.filter().is_some(); + log::debug!( + "[LC-Optimizer] WRAP: all {} output columns cacheable, predicate={}", + num_fields, + has_predicate + ); + + let mut new_config = file_scan_config.clone(); + let new_source = + LiquidParquetSource::from_parquet_source(parquet_source.clone(), cache.clone()); + + new_config.file_source = Arc::new(new_source); + let new_file_source: Arc = Arc::new(new_config); + let new_plan = Arc::new(DataSourceExec::new(new_file_source)); + + return Ok(Transformed::new( + new_plan, + true, + TreeNodeRecursion::Continue, + )); + } + Ok(Transformed::no(plan)) +} + +#[cfg(test)] +mod tests { + use arrow::array::{Int32Array, RecordBatch, StringArray}; + use arrow_schema::{DataType, Field, Schema}; + use datafusion::{datasource::physical_plan::FileScanConfig, prelude::SessionContext}; + use liquid_cache::{cache::TranscodeEvict, cache_policies::LiquidPolicy}; + use parquet::arrow::ArrowWriter; + + use crate::LiquidCacheParquet; + + use super::*; + + fn test_cache() -> LiquidCacheParquetRef { + Arc::new(LiquidCacheParquet::new( + 8192, + 1_000_000, + Box::new(LiquidPolicy::new()), + Box::new(TranscodeEvict), + )) + } + + /// Write a small parquet file with numeric and string columns. + fn write_test_parquet(dir: &std::path::Path) -> String { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + Field::new("c", DataType::Int32, false), + Field::new("d", DataType::Int32, false), + Field::new("e", DataType::Int32, false), + Field::new("s", DataType::Utf8, false), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3, 4])), + Arc::new(Int32Array::from(vec![10, 20, 30, 40])), + Arc::new(Int32Array::from(vec![100, 200, 300, 400])), + Arc::new(Int32Array::from(vec![5, 6, 7, 8])), + Arc::new(Int32Array::from(vec![50, 60, 70, 80])), + Arc::new(StringArray::from(vec!["w", "x", "y", "z"])), + ], + ) + .unwrap(); + let path = dir.join("data.parquet"); + let file = std::fs::File::create(&path).unwrap(); + let mut writer = ArrowWriter::try_new(file, schema, None).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + path.to_string_lossy().into_owned() + } + + fn count_sources(plan: &Arc) -> (usize, usize) { + let mut liquid = 0; + let mut parquet = 0; + plan.apply(|node| { + if let Some(exec) = node.downcast_ref::() { + let data_source = exec.data_source(); + if let Some(config) = data_source.downcast_ref::() { + let file_source = config.file_source(); + if file_source.downcast_ref::().is_some() { + liquid += 1; + } else if file_source.downcast_ref::().is_some() { + parquet += 1; + } + } + } + Ok(TreeNodeRecursion::Continue) + }) + .unwrap(); + (liquid, parquet) + } + + async fn plan_for_sql(sql: &str, path: &str) -> Arc { + let ctx = SessionContext::new(); + ctx.register_parquet("t", path, Default::default()) + .await + .unwrap(); + let df = ctx.sql(sql).await.unwrap(); + df.create_physical_plan().await.unwrap() + } + + #[tokio::test] + async fn test_plan_rewrite_wraps_numeric_scan() { + let tmp_dir = tempfile::tempdir().unwrap(); + let path = write_test_parquet(tmp_dir.path()); + let plan = plan_for_sql("SELECT a, b FROM t WHERE a > 1", &path).await; + let expected_schema = plan.schema(); + + let rewritten = rewrite_data_source_plan(plan, &test_cache()); + + let (liquid, parquet) = count_sources(&rewritten); + assert_eq!(liquid, 1); + assert_eq!(parquet, 0); + assert_eq!(rewritten.schema(), expected_schema); + } + + #[tokio::test] + async fn test_plan_rewrite_skips_string_output() { + let tmp_dir = tempfile::tempdir().unwrap(); + let path = write_test_parquet(tmp_dir.path()); + let plan = plan_for_sql("SELECT a, s FROM t WHERE a > 1", &path).await; + + let rewritten = rewrite_data_source_plan(plan, &test_cache()); + + let (liquid, parquet) = count_sources(&rewritten); + assert_eq!(liquid, 0); + assert_eq!(parquet, 1); + } + + #[tokio::test] + async fn test_plan_rewrite_skips_string_predicate() { + let tmp_dir = tempfile::tempdir().unwrap(); + let path = write_test_parquet(tmp_dir.path()); + let plan = plan_for_sql("SELECT a, b FROM t WHERE s = 'x'", &path).await; + + let rewritten = rewrite_data_source_plan(plan, &test_cache()); + + let (liquid, parquet) = count_sources(&rewritten); + assert_eq!(liquid, 0); + assert_eq!(parquet, 1); + } + + #[tokio::test] + async fn test_plan_rewrite_skips_wide_projection() { + let tmp_dir = tempfile::tempdir().unwrap(); + let path = write_test_parquet(tmp_dir.path()); + // 5 numeric columns > MAX_LC_COLUMNS (4) + let plan = plan_for_sql("SELECT a, b, c, d, e FROM t", &path).await; + + let rewritten = rewrite_data_source_plan(plan, &test_cache()); + + let (liquid, parquet) = count_sources(&rewritten); + assert_eq!(liquid, 0); + assert_eq!(parquet, 1); + } + + #[tokio::test] + async fn test_plan_rewrite_skips_empty_projection() { + let tmp_dir = tempfile::tempdir().unwrap(); + let path = write_test_parquet(tmp_dir.path()); + let plan = plan_for_sql("SELECT COUNT(*) FROM t", &path).await; + + let rewritten = rewrite_data_source_plan(plan, &test_cache()); + + let (liquid, _parquet) = count_sources(&rewritten); + assert_eq!(liquid, 0); + } +} diff --git a/sandbox/libs/dataformat-native/rust/liquid-cache/datafusion/src/reader/mod.rs b/sandbox/libs/dataformat-native/rust/liquid-cache/datafusion/src/reader/mod.rs new file mode 100644 index 0000000000000..ac32eac7ebdd9 --- /dev/null +++ b/sandbox/libs/dataformat-native/rust/liquid-cache/datafusion/src/reader/mod.rs @@ -0,0 +1,6 @@ +pub mod plantime; +mod runtime; +mod utils; + +pub use plantime::{FilterCandidateBuilder, LiquidParquetSource, LiquidPredicate, LiquidRowFilter}; +pub(crate) use runtime::extract_multi_column_or; diff --git a/sandbox/libs/dataformat-native/rust/liquid-cache/datafusion/src/reader/plantime/engagement_policy.rs b/sandbox/libs/dataformat-native/rust/liquid-cache/datafusion/src/reader/plantime/engagement_policy.rs new file mode 100644 index 0000000000000..e8b2428777431 --- /dev/null +++ b/sandbox/libs/dataformat-native/rust/liquid-cache/datafusion/src/reader/plantime/engagement_policy.rs @@ -0,0 +1,95 @@ +//! Cache engagement policy — decides per-file whether LC should serve +//! decoded batches from cache or delegate to plain parquet. + +use std::fmt::Debug; +use std::sync::Arc; + +/// Context passed to the policy at file-open time. +#[derive(Debug, Clone)] +pub struct EngagementContext { + /// Fraction of rows surviving RG + page pruning (0.0–1.0). + pub estimated_selectivity: f64, + /// Whether the query has a pushdown predicate on this file. + pub has_predicate: bool, + /// Total rows across selected row groups. + pub total_rows: usize, + /// Rows that survived pruning. + pub selected_rows: usize, + /// File path for logging. + pub file_path: String, +} + +/// Decision returned by the policy. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum EngagementDecision { + /// Serve decoded batches from cache, fill on miss. + UseLiquidCache, + /// Bypass LC, delegate to plain parquet reader. + DelegateToParquet, +} + +/// Decides per-file whether Liquid Cache should engage or delegate. +pub trait CacheEngagementPolicy: Send + Sync + Debug { + /// Decide whether to use LC or delegate to plain parquet for this file. + fn decide(&self, ctx: &EngagementContext) -> EngagementDecision; +} + +/// Default selectivity threshold below which LC delegates to parquet. +pub const DEFAULT_SELECTIVITY_THRESHOLD: f64 = 0.5; + +/// Delegate when selectivity < threshold AND a predicate exists. +#[derive(Debug, Clone)] +pub struct SelectivityThresholdPolicy { + /// Selectivity below which LC delegates to plain parquet. + pub threshold: f64, +} + +impl Default for SelectivityThresholdPolicy { + fn default() -> Self { + Self { + threshold: DEFAULT_SELECTIVITY_THRESHOLD, + } + } +} + +impl SelectivityThresholdPolicy { + /// Create a policy with the given selectivity threshold. + pub fn new(threshold: f64) -> Self { + Self { threshold } + } +} + +impl CacheEngagementPolicy for SelectivityThresholdPolicy { + fn decide(&self, ctx: &EngagementContext) -> EngagementDecision { + if ctx.estimated_selectivity < self.threshold && ctx.has_predicate { + EngagementDecision::DelegateToParquet + } else { + EngagementDecision::UseLiquidCache + } + } +} + +/// Always use LC regardless of selectivity. +#[derive(Debug, Clone)] +pub struct AlwaysEngagePolicy; + +impl CacheEngagementPolicy for AlwaysEngagePolicy { + fn decide(&self, _ctx: &EngagementContext) -> EngagementDecision { + EngagementDecision::UseLiquidCache + } +} + +/// Never use LC — always delegate to plain parquet. +#[derive(Debug, Clone)] +pub struct NeverEngagePolicy; + +impl CacheEngagementPolicy for NeverEngagePolicy { + fn decide(&self, _ctx: &EngagementContext) -> EngagementDecision { + EngagementDecision::DelegateToParquet + } +} + +/// Returns the default engagement policy. +pub fn default_engagement_policy() -> Arc { + Arc::new(SelectivityThresholdPolicy::default()) +} diff --git a/sandbox/libs/dataformat-native/rust/liquid-cache/datafusion/src/reader/plantime/mod.rs b/sandbox/libs/dataformat-native/rust/liquid-cache/datafusion/src/reader/plantime/mod.rs new file mode 100644 index 0000000000000..0f22e9ae5be9a --- /dev/null +++ b/sandbox/libs/dataformat-native/rust/liquid-cache/datafusion/src/reader/plantime/mod.rs @@ -0,0 +1,12 @@ +#[cfg(test)] +pub(crate) use source::CachedMetaReaderFactory; +pub use source::LiquidParquetSource; +pub(crate) use source::ParquetMetadataCacheReader; + +pub mod engagement_policy; +mod opener; +mod row_filter; +mod row_group_filter; +pub mod source; + +pub use row_filter::{FilterCandidateBuilder, LiquidPredicate, LiquidRowFilter}; diff --git a/sandbox/libs/dataformat-native/rust/liquid-cache/datafusion/src/reader/plantime/opener.rs b/sandbox/libs/dataformat-native/rust/liquid-cache/datafusion/src/reader/plantime/opener.rs new file mode 100644 index 0000000000000..af0fac5393811 --- /dev/null +++ b/sandbox/libs/dataformat-native/rust/liquid-cache/datafusion/src/reader/plantime/opener.rs @@ -0,0 +1,509 @@ +use std::sync::Arc; + +use crate::{ + cache::LiquidCacheParquetRef, + reader::{ + plantime::{ + engagement_policy::{CacheEngagementPolicy, EngagementContext, EngagementDecision}, + row_filter::build_row_filter, + row_group_filter::RowGroupAccessPlanFilter, + }, + runtime::LiquidStreamBuilder, + }, +}; +use arrow::array::{RecordBatch, RecordBatchOptions}; +use arrow_schema::SchemaRef; +use datafusion::{ + common::exec_err, + datasource::{ + listing::PartitionedFile, + physical_plan::{ + FileOpenFuture, FileOpener, ParquetFileMetrics, ParquetFileReaderFactory, + parquet::{PagePruningAccessPlanFilter, ParquetAccessPlan}, + }, + table_schema::TableSchema, + }, + error::DataFusionError, + physical_expr::PhysicalExprSimplifier, + physical_expr::projection::ProjectionExprs, + physical_expr::utils::reassign_expr_columns, + physical_expr_adapter::{PhysicalExprAdapterFactory, replace_columns_with_literals}, + physical_expr_common::physical_expr::is_dynamic_physical_expr, + physical_optimizer::pruning::{FilePruner, PruningPredicate, build_pruning_predicate}, + physical_plan::{ + PhysicalExpr, + metrics::{Count, ExecutionPlanMetricsSet, MetricBuilder}, + }, +}; +use futures::StreamExt; +use futures::TryStreamExt; +use log::debug; +use parquet::arrow::{ + ParquetRecordBatchStreamBuilder, ProjectionMask, + arrow_reader::{ArrowReaderMetadata, ArrowReaderOptions}, +}; +use parquet::file::metadata::ParquetMetaData; + +use super::source::CachedMetaReaderFactory; + +pub struct LiquidParquetOpener { + partition_index: usize, + projection: ProjectionExprs, + batch_size: usize, + limit: Option, + predicate: Option>, + table_schema: TableSchema, + metrics: ExecutionPlanMetricsSet, + parquet_file_reader_factory: Arc, + reorder_filters: bool, + liquid_cache: LiquidCacheParquetRef, + expr_adapter_factory: Arc, + /// Optional caller-provided reader factory with pre-loaded metadata. + caller_reader_factory: Option>, + /// Policy that decides whether to use LC stream or delegate to parquet. + engagement_policy: Arc, +} + +impl LiquidParquetOpener { + #[allow(clippy::too_many_arguments)] + pub fn new( + partition_index: usize, + projection: ProjectionExprs, + batch_size: usize, + limit: Option, + predicate: Option>, + table_schema: TableSchema, + metrics: ExecutionPlanMetricsSet, + liquid_cache: LiquidCacheParquetRef, + parquet_file_reader_factory: Arc, + reorder_filters: bool, + expr_adapter_factory: Arc, + caller_reader_factory: Option>, + engagement_policy: Arc, + ) -> Self { + Self { + partition_index, + projection, + batch_size, + limit, + predicate, + table_schema, + metrics, + liquid_cache, + parquet_file_reader_factory, + reorder_filters, + expr_adapter_factory, + caller_reader_factory, + engagement_policy, + } + } +} + +impl FileOpener for LiquidParquetOpener { + fn open(&self, partitioned_file: PartitionedFile) -> Result { + let file_range = partitioned_file.range.clone(); + let access_plan_ext = partitioned_file.extensions.get_arc::(); + let file_name = partitioned_file.object_meta.location.to_string(); + let file_metrics = ParquetFileMetrics::new(self.partition_index, &file_name, &self.metrics); + + let metadata_size_hint = partitioned_file.metadata_size_hint; + let has_predicate = self.predicate.is_some(); + log::debug!( + "[LC-Opener] open file={}, predicate={}, batch_size={}, limit={:?}", + file_name, + has_predicate, + self.batch_size, + self.limit + ); + + let lc = self.liquid_cache.clone(); + let file_loc = partitioned_file.object_meta.location.to_string(); + + // If caller provided a reader factory with pre-loaded metadata, create + // a reader from it. We'll use this for get_metadata() to avoid refetching. + let caller_metadata_reader = self.caller_reader_factory.as_ref().and_then(|factory| { + factory + .create_reader( + self.partition_index, + partitioned_file.clone(), + metadata_size_hint, + &self.metrics, + ) + .ok() + }); + + let mut async_file_reader = self.parquet_file_reader_factory.create_liquid_reader( + self.partition_index, + partitioned_file.clone(), + metadata_size_hint, + &self.metrics, + ); + + let batch_size = self.batch_size; + let logical_file_schema = Arc::clone(self.table_schema.file_schema()); + let output_schema = Arc::new( + self.projection + .project_schema(self.table_schema.table_schema())?, + ); + let mut projection = self.projection.clone(); + let mut predicate = self.predicate.clone(); + let mut literal_columns = std::collections::HashMap::new(); + for (field, value) in self + .table_schema + .table_partition_cols() + .iter() + .zip(partitioned_file.partition_values.iter()) + { + literal_columns.insert(field.name().clone(), value.clone()); + } + if !literal_columns.is_empty() { + projection = projection.try_map_exprs(|expr| { + replace_columns_with_literals(Arc::clone(&expr), &literal_columns) + })?; + predicate = predicate + .map(|p| replace_columns_with_literals(p, &literal_columns)) + .transpose()?; + } + let reorder_predicates = self.reorder_filters; + let limit = self.limit; + + let predicate_creation_errors = + MetricBuilder::new(&self.metrics).global_counter("num_predicate_creation_errors"); + + let expr_adapter_factory = Arc::clone(&self.expr_adapter_factory); + let engagement_policy = Arc::clone(&self.engagement_policy); + Ok(Box::pin(async move { + // Prune this file using the file level statistics and partition values. + // Since dynamic filters may have been updated since planning it is possible that we are able + // to prune files now that we couldn't prune at planning time. + // It is assumed that there is no point in doing pruning here if the predicate is not dynamic, + // as it would have been done at planning time. + // We'll also check this after every record batch we read, + // and if at some point we are able to prove we can prune the file using just the file level statistics + // we can end the stream early. + let mut file_pruner = predicate + .as_ref() + .filter(|p| is_dynamic_physical_expr(p) || partitioned_file.has_statistics()) + .and_then(|p| { + FilePruner::try_new( + Arc::clone(p), + &logical_file_schema, + &partitioned_file, + predicate_creation_errors.clone(), + ) + }); + + if let Some(file_pruner) = &mut file_pruner + && file_pruner.should_prune()? + { + file_metrics.files_ranges_pruned_statistics.add_pruned(1); + return Ok(futures::stream::empty().boxed()); + } + + file_metrics.files_ranges_pruned_statistics.add_matched(1); + + let options = ArrowReaderOptions::new() + .with_page_index_policy(parquet::file::metadata::PageIndexPolicy::Required); + let mut metadata_timer = file_metrics.metadata_load_time.timer(); + + // Try to get metadata from caller's pre-loaded reader first (instant). + // Fall back to loading from the object store if not available. + let reader_metadata = if let Some(mut caller_reader) = caller_metadata_reader { + match caller_reader.get_metadata(Some(&options)).await { + Ok(meta) => { + log::debug!("[LC-Meta] REUSE from caller factory: {}", file_name); + ArrowReaderMetadata::try_new(meta, options.clone())? + } + Err(_) => { + log::debug!( + "[LC-Meta] caller factory failed, loading from store: {}", + file_name + ); + ArrowReaderMetadata::load_async(&mut async_file_reader, options.clone()) + .await? + } + } + } else { + ArrowReaderMetadata::load_async(&mut async_file_reader, options.clone()).await? + }; + + // Note about schemas: we are actually dealing with **3 different schemas** here: + // - The table schema as defined by the TableProvider. + // This is what the user sees, what they get when they `SELECT * FROM table`, etc. + // - The logical file schema: this is the table schema minus any hive partition columns and projections. + // This is what the physical file schema is coerced to. + // - The physical file schema: this is the schema as defined by the parquet file. This is what the parquet file actually contains. + let physical_file_schema = Arc::clone(reader_metadata.schema()); + let cache_full_schema = Arc::clone(&physical_file_schema); + + let rewriter = expr_adapter_factory.create( + Arc::clone(&logical_file_schema), + Arc::clone(&physical_file_schema), + )?; + let simplifier = PhysicalExprSimplifier::new(&physical_file_schema); + predicate = predicate + .map(|p| simplifier.simplify(rewriter.rewrite(p)?)) + .transpose()?; + projection = projection.try_map_exprs(|p| simplifier.simplify(rewriter.rewrite(p)?))?; + + let (pruning_predicate, page_pruning_predicate) = build_pruning_predicates( + predicate.as_ref(), + &physical_file_schema, + &predicate_creation_errors, + ); + + metadata_timer.stop(); + + let mut builder = ParquetRecordBatchStreamBuilder::new_with_metadata( + async_file_reader.clone(), + reader_metadata.clone(), + ); + let indices = projection.column_indices(); + let mask = ProjectionMask::roots(builder.parquet_schema(), indices); + + // Filter pushdown: evaluate predicates during scan + let row_filter = predicate.as_ref().and_then(|p| { + let row_filter = build_row_filter( + p, + &physical_file_schema, + reader_metadata.metadata(), + reorder_predicates, + &file_metrics, + ); + + match row_filter { + Ok(Some(filter)) => Some(filter), + Ok(None) => None, + Err(e) => { + debug!("Ignoring error building row filter for '{predicate:?}': {e:?}"); + None + } + } + }); + + // Determine which row groups to actually read. The idea is to skip + // as many row groups as possible based on the metadata and query + let file_metadata: Arc = Arc::clone(builder.metadata()); + let predicate = pruning_predicate.as_ref().map(|p| p.as_ref()); + let rg_metadata = file_metadata.row_groups(); + // track which row groups to actually read + let access_plan = create_initial_plan(&file_name, access_plan_ext, rg_metadata.len())?; + let mut row_groups = RowGroupAccessPlanFilter::new(access_plan); + // if there is a range restricting what parts of the file to read + if let Some(range) = file_range.as_ref() { + row_groups.prune_by_range(rg_metadata, range); + } + // If there is a predicate that can be evaluated against the metadata + if let Some(predicate) = predicate.as_ref() { + row_groups.prune_by_statistics( + &physical_file_schema, + builder.parquet_schema(), + rg_metadata, + predicate, + &file_metrics, + ); + + if !row_groups.is_empty() { + row_groups + .prune_by_bloom_filters( + &physical_file_schema, + &mut builder, + predicate, + &file_metrics, + ) + .await; + } + } + + let mut access_plan = row_groups.build(); + + // page index pruning: if all data on individual pages can + // be ruled using page metadata, rows from other columns + // with that range can be skipped as well + if !access_plan.is_empty() + && let Some(p) = page_pruning_predicate + { + access_plan = p.prune_plan_with_page_index( + access_plan, + &physical_file_schema, + builder.parquet_schema(), + file_metadata.as_ref(), + &file_metrics, + ); + } + + let row_group_indexes = access_plan.row_group_indexes(); + + // Early exit: if all row groups were pruned, return empty stream + if row_group_indexes.is_empty() { + log::debug!("[LC-Opener] EMPTY: all RGs pruned, file={}", file_name); + return Ok(futures::stream::empty().boxed()); + } + + let row_selection = access_plan.into_overall_row_selection(rg_metadata)?; + + // Estimate selectivity from row_selection: how many rows survived + // RG pruning + page index pruning vs total rows in selected RGs. + let total_rows: usize = row_group_indexes + .iter() + .map(|&idx| rg_metadata[idx].num_rows() as usize) + .sum(); + let selected_rows = row_selection + .as_ref() + .map(|sel| sel.row_count()) + .unwrap_or(total_rows); + let estimated_selectivity = if total_rows > 0 { + selected_rows as f64 / total_rows as f64 + } else { + 1.0 + }; + + // If selectivity is low (few rows match), the decode cost is already + // minimal — LC cache overhead would dominate for negligible savings. + // Delegate to plain parquet for fast pass-through. + // For high selectivity (most rows match = lots of decode), LC cache + // saves significant decode work on warm iterations. + let engagement_ctx = EngagementContext { + estimated_selectivity, + has_predicate: predicate.is_some(), + total_rows, + selected_rows, + file_path: file_name.clone(), + }; + + match engagement_policy.decide(&engagement_ctx) { + EngagementDecision::DelegateToParquet => { + log::debug!( + "[LC-Opener] DELEGATE to plain parquet: selectivity={:.3}, file={}", + estimated_selectivity, + file_name + ); + let mut plain_builder = ParquetRecordBatchStreamBuilder::new_with_metadata( + async_file_reader, + reader_metadata, + ) + .with_batch_size(batch_size) + .with_projection(mask) + .with_row_groups(row_group_indexes); + + if let Some(sel) = row_selection { + plain_builder = plain_builder.with_row_selection(sel); + } + if let Some(lim) = limit { + plain_builder = plain_builder.with_limit(lim); + } + + let stream = plain_builder.build()?; + let adapted = stream.map_err(|e| DataFusionError::External(Box::new(e))); + return Ok(adapted.boxed()); + } + EngagementDecision::UseLiquidCache => { + log::debug!( + "[LC-Opener] LC STREAM: selectivity={:.3}, predicate={}, file={}", + estimated_selectivity, + predicate.is_some(), + file_name + ); + } + } + + let mut liquid_builder = + LiquidStreamBuilder::new(async_file_reader, Arc::clone(reader_metadata.metadata())) + .with_batch_size(batch_size) + .with_row_groups(row_group_indexes) + .with_projection(mask) + .with_selection(row_selection) + .with_limit(limit); + + if let Some(row_filter) = row_filter { + liquid_builder = liquid_builder.with_row_filter(row_filter); + } + + let liquid_cache = lc.register_or_get_file(file_loc, Arc::clone(&cache_full_schema)); + + let stream = liquid_builder.build(liquid_cache)?; + + let stream_schema = Arc::clone(stream.schema()); + let replace_schema = !stream_schema.eq(&output_schema); + let projection = + projection.try_map_exprs(|expr| reassign_expr_columns(expr, &stream_schema))?; + let projector = projection.make_projector(&stream_schema)?; + + let adapted = stream + .map_err(|e| DataFusionError::External(Box::new(e))) + .map(move |batch| { + batch.and_then(|batch| { + let batch = projector.project_batch(&batch)?; + if replace_schema { + let (_schema, arrays, num_rows) = batch.into_parts(); + let options = RecordBatchOptions::new().with_row_count(Some(num_rows)); + RecordBatch::try_new_with_options( + Arc::clone(&output_schema), + arrays, + &options, + ) + .map_err(Into::into) + } else { + Ok(batch) + } + }) + }); + + Ok(adapted.boxed()) + })) + } +} + +fn create_initial_plan( + file_name: &str, + access_plan: Option>, + row_group_count: usize, +) -> Result { + if let Some(access_plan) = access_plan { + let plan_len = access_plan.len(); + if plan_len != row_group_count { + return exec_err!( + "Invalid ParquetAccessPlan for {file_name}. Specified {plan_len} row groups, but file has {row_group_count}" + ); + } + + // check row group count matches the plan + return Ok(access_plan.as_ref().clone()); + } + + // default to scanning all row groups + Ok(ParquetAccessPlan::new_all(row_group_count)) +} + +pub(crate) fn build_pruning_predicates( + predicate: Option<&Arc>, + file_schema: &SchemaRef, + predicate_creation_errors: &Count, +) -> ( + Option>, + Option>, +) { + let Some(predicate) = predicate.as_ref() else { + return (None, None); + }; + let pruning_predicate = build_pruning_predicate( + Arc::clone(predicate), + file_schema, + predicate_creation_errors, + ); + let page_pruning_predicate = build_page_pruning_predicate(predicate, file_schema); + (pruning_predicate, Some(page_pruning_predicate)) +} + +/// Build a page pruning predicate from an optional predicate expression. +/// If the predicate is None or the predicate cannot be converted to a page pruning +/// predicate, return None. +pub(crate) fn build_page_pruning_predicate( + predicate: &Arc, + file_schema: &SchemaRef, +) -> Arc { + Arc::new(PagePruningAccessPlanFilter::new( + predicate, + Arc::clone(file_schema), + )) +} diff --git a/sandbox/libs/dataformat-native/rust/liquid-cache/datafusion/src/reader/plantime/row_filter.rs b/sandbox/libs/dataformat-native/rust/liquid-cache/datafusion/src/reader/plantime/row_filter.rs new file mode 100644 index 0000000000000..fe9c5ea16246c --- /dev/null +++ b/sandbox/libs/dataformat-native/rust/liquid-cache/datafusion/src/reader/plantime/row_filter.rs @@ -0,0 +1,515 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Utilities to push down of DataFusion filter predicates (any DataFusion +//! `PhysicalExpr` that evaluates to a [`BooleanArray`]) to the parquet decoder +//! level in `arrow-rs`. +//! +//! DataFusion will use a `ParquetRecordBatchStream` to read data from parquet +//! into [`RecordBatch`]es. +//! +//! The `ParquetRecordBatchStream` takes an optional `RowFilter` which is itself +//! a Vec of `Box`. During decoding, the predicates are +//! evaluated in order, to generate a mask which is used to avoid decoding rows +//! in projected columns which do not pass the filter which can significantly +//! reduce the amount of compute required for decoding and thus improve query +//! performance. +//! +//! Since the predicates are applied serially in the order defined in the +//! `RowFilter`, the optimal ordering depends on the exact filters. The best +//! filters to execute first have two properties: +//! +//! 1. They are relatively inexpensive to evaluate (e.g. they read +//! column chunks which are relatively small) +//! +//! 2. They filter many (contiguous) rows, reducing the amount of decoding +//! required for subsequent filters and projected columns +//! +//! If requested, this code will reorder the filters based on heuristics try and +//! reduce the evaluation cost. +//! +//! The basic algorithm for constructing the `RowFilter` is as follows +//! +//! 1. Break conjunctions into separate predicates. An expression +//! like `a = 1 AND (b = 2 AND c = 3)` would be +//! separated into the expressions `a = 1`, `b = 2`, and `c = 3`. +//! 2. Determine whether each predicate can be evaluated as an `ArrowPredicate`. +//! 3. Determine, for each predicate, the total compressed size of all +//! columns required to evaluate the predicate. +//! 4. Determine, for each predicate, whether all columns required to +//! evaluate the expression are sorted. +//! 5. Re-order the predicate by total size (from step 3). +//! 6. Partition the predicates according to whether they are sorted (from step 4) +//! 7. "Compile" each predicate `Expr` to a `DatafusionArrowPredicate`. +//! 8. Build the `RowFilter` with the sorted predicates followed by +//! the unsorted predicates. Within each partition, predicates are +//! still be sorted by size. + +use std::cmp::Ordering; +use std::collections::BTreeSet; +use std::sync::Arc; + +use arrow::array::BooleanArray; +use arrow::datatypes::{DataType, Schema}; +use arrow::error::{ArrowError, Result as ArrowResult}; +use arrow::record_batch::RecordBatch; +use arrow_schema::SchemaRef; +use datafusion::datasource::physical_plan::ParquetFileMetrics; +use datafusion::logical_expr::Operator; +use datafusion::physical_expr::utils::reassign_expr_columns; +use datafusion::physical_plan::expressions::{BinaryExpr, LikeExpr}; +use datafusion::physical_plan::metrics; +use parquet::arrow::ProjectionMask; +use parquet::arrow::arrow_reader::ArrowPredicate; +use parquet::file::metadata::ParquetMetaData; + +use datafusion::common::Result; +use datafusion::common::cast::as_boolean_array; +use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion, TreeNodeVisitor}; +use datafusion::physical_expr::expressions::Column; +use datafusion::physical_expr::{PhysicalExpr, split_conjunction}; + +/// A row filter that can be used to filter rows from a parquet file. +pub struct LiquidRowFilter { + predicates: Vec, +} + +impl LiquidRowFilter { + /// Create a new `LiquidRowFilter` from a vector of `LiquidPredicate`s. + pub fn new(predicates: Vec) -> Self { + Self { predicates } + } + + /// Get the predicates of the `LiquidRowFilter`. + pub fn predicates(&self) -> &[LiquidPredicate] { + &self.predicates + } + + /// Get the predicates of the `LiquidRowFilter` as mutable. + pub fn predicates_mut(&mut self) -> &mut [LiquidPredicate] { + &mut self.predicates + } +} + +pub(crate) fn get_predicate_column_id(projection: &parquet::arrow::ProjectionMask) -> Vec { + #[derive(Debug, Clone)] + struct ProjectionMaskLiquid { + mask: Option>, + } + let project_inner: &ProjectionMaskLiquid = unsafe { std::mem::transmute(projection) }; + project_inner + .mask + .as_ref() + .map(|m| { + m.iter() + .enumerate() + .filter_map(|(pos, &x)| if x { Some(pos) } else { None }) + .collect::>() + }) + .unwrap_or_default() +} + +/// A "compiled" predicate passed to `ParquetRecordBatchStream` to perform +/// row-level filtering during parquet decoding. +/// +/// See the module level documentation for more information. +/// +/// Implements the `ArrowPredicate` trait used by the parquet decoder +/// +/// An expression can be evaluated as a `DatafusionArrowPredicate` if it: +/// * Does not reference any projected columns +/// * Does not reference columns with non-primitive types (e.g. structs / lists) +#[derive(Debug, Clone)] +pub struct LiquidPredicate { + /// the filter expression + physical_expr: Arc, + /// the filter expression without reassigned index + physical_expr_physical_column_index: Arc, + /// Path to the columns in the parquet schema required to evaluate the + /// expression + projection_mask: ProjectionMask, + /// how many rows were filtered out by this predicate + rows_pruned: metrics::Count, + /// how many rows passed this predicate + rows_matched: metrics::Count, + /// how long was spent evaluating this predicate + time: metrics::Time, +} + +impl std::fmt::Display for LiquidPredicate { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}", self.physical_expr.as_ref()) + } +} + +impl LiquidPredicate { + /// Create a new `LiquidPredicate` from a `FilterCandidate` + pub fn try_new_with_metrics( + candidate: FilterCandidate, + projection: ProjectionMask, + rows_pruned: metrics::Count, + rows_matched: metrics::Count, + time: metrics::Time, + ) -> Result { + let physical_expr = + reassign_expr_columns(candidate.expr.clone(), &candidate.filter_schema)?; + + Ok(Self { + physical_expr, + physical_expr_physical_column_index: candidate.expr, + projection_mask: projection, + rows_pruned, + rows_matched, + time, + }) + } + + /// Create a new `LiquidPredicate` from a `FilterCandidate` + pub fn try_new(candidate: FilterCandidate, projection: ProjectionMask) -> Result { + Self::try_new_with_metrics( + candidate, + projection, + metrics::Count::new(), + metrics::Count::new(), + metrics::Time::new(), + ) + } + + /// Get the physical expression with physical column index. + pub fn physical_expr_physical_column_index(&self) -> &Arc { + &self.physical_expr_physical_column_index + } + + /// Get the physical expression rewritten to the projected batch schema. + pub fn physical_expr(&self) -> &Arc { + &self.physical_expr + } + + /// Get the column ids of the predicate. + pub fn predicate_column_ids(&self) -> Vec { + let projection = self.projection(); + get_predicate_column_id(projection) + } +} + +impl ArrowPredicate for LiquidPredicate { + fn projection(&self) -> &ProjectionMask { + &self.projection_mask + } + + fn evaluate(&mut self, batch: RecordBatch) -> ArrowResult { + // scoped timer updates on drop + let mut timer = self.time.timer(); + + self.physical_expr + .evaluate(&batch) + .and_then(|v| v.into_array(batch.num_rows())) + .and_then(|array| { + let bool_arr = as_boolean_array(&array)?.clone(); + let num_matched = bool_arr.true_count(); + let num_pruned = bool_arr.len() - num_matched; + self.rows_pruned.add(num_pruned); + self.rows_matched.add(num_matched); + timer.stop(); + Ok(bool_arr) + }) + .map_err(|e| { + ArrowError::ComputeError(format!("Error evaluating filter predicate: {e:?}")) + }) + } +} + +/// A candidate expression for creating a `RowFilter`. +/// +/// Each candidate contains the expression as well as data to estimate the cost +/// of evaluating the resulting expression. +/// +/// See the module level documentation for more information. +pub struct FilterCandidate { + expr: Arc, + required_bytes: usize, + can_use_index: bool, + projection: Vec, + /// The projected file schema that this filter references + filter_schema: SchemaRef, +} + +impl FilterCandidate { + pub fn projection(&self, metadata: &ParquetMetaData) -> ProjectionMask { + ProjectionMask::roots( + metadata.file_metadata().schema_descr(), + self.projection.iter().copied(), + ) + } +} + +/// Helper to build a `FilterCandidate`. +/// +/// This will do several things +/// 1. Determine the columns required to evaluate the expression +/// 2. Calculate data required to estimate the cost of evaluating the filter +pub struct FilterCandidateBuilder { + expr: Arc, + /// The schema of this parquet file. + /// Expressions are already adapted to this schema before row-filter construction. + file_schema: SchemaRef, +} + +impl FilterCandidateBuilder { + /// Create a new `FilterCandidateBuilder` + pub fn new(expr: Arc, file_schema: SchemaRef) -> Self { + Self { expr, file_schema } + } + + /// Attempt to build a `FilterCandidate` from the expression + /// + /// # Return values + /// + /// * `Ok(Some(candidate))` if the expression can be used as an ArrowFilter + /// * `Ok(None)` if the expression cannot be used as an ArrowFilter + /// * `Err(e)` if an error occurs while building the candidate + pub fn build(self, metadata: &ParquetMetaData) -> Result> { + let Some(required_indices_into_file_schema) = + pushdown_columns(&self.expr, &self.file_schema)? + else { + return Ok(None); + }; + + if required_indices_into_file_schema.is_empty() { + return Ok(None); + } + + let projected_file_schema = Arc::new( + self.file_schema + .project(&required_indices_into_file_schema)?, + ); + + let required_bytes = size_of_columns(&required_indices_into_file_schema, metadata)?; + let can_use_index = columns_sorted(&required_indices_into_file_schema, metadata)?; + + Ok(Some(FilterCandidate { + expr: self.expr, + required_bytes, + can_use_index, + projection: required_indices_into_file_schema, + filter_schema: Arc::clone(&projected_file_schema), + })) + } +} + +// a struct that implements TreeNodeRewriter to traverse a PhysicalExpr tree structure to determine +// if any column references in the expression would prevent it from being predicate-pushed-down. +// if non_primitive_columns || projected_columns, it can't be pushed down. +// can't be reused between calls to `rewrite`; each construction must be used only once. +struct PushdownChecker<'schema> { + /// Does the expression require any non-primitive columns (like structs)? + non_primitive_columns: bool, + /// Does the expression reference any columns that are not in the file schema? + projected_columns: bool, + // Indices into the file schema of the columns required to evaluate the expression + required_columns: BTreeSet, + file_schema: &'schema Schema, +} + +impl<'schema> PushdownChecker<'schema> { + fn new(file_schema: &'schema Schema) -> Self { + Self { + non_primitive_columns: false, + projected_columns: false, + required_columns: BTreeSet::default(), + file_schema, + } + } + + fn check_single_column(&mut self, column_name: &str) -> Option { + if let Ok(idx) = self.file_schema.index_of(column_name) { + self.required_columns.insert(idx); + if DataType::is_nested(self.file_schema.field(idx).data_type()) { + self.non_primitive_columns = true; + return Some(TreeNodeRecursion::Jump); + } + } else { + // If the column does not exist in the file schema then it cannot be pushed down. + self.projected_columns = true; + return Some(TreeNodeRecursion::Jump); + } + + None + } + + #[inline] + fn prevents_pushdown(&self) -> bool { + self.non_primitive_columns || self.projected_columns + } +} + +impl TreeNodeVisitor<'_> for PushdownChecker<'_> { + type Node = Arc; + + fn f_down(&mut self, node: &Self::Node) -> Result { + if let Some(column) = node.downcast_ref::() + && let Some(recursion) = self.check_single_column(column.name()) + { + return Ok(recursion); + } + + Ok(TreeNodeRecursion::Continue) + } +} + +// Checks if a given expression can be pushed down into `DataSourceExec` as opposed to being evaluated +// post-parquet-scan in a `FilterExec`. If it can be pushed down, this returns all the +// columns in the given expression so that they can be used in the parquet scanning, along with the +// expression rewritten as defined in [`PushdownChecker::f_up`] +fn pushdown_columns( + expr: &Arc, + file_schema: &Schema, +) -> Result>> { + let mut checker = PushdownChecker::new(file_schema); + expr.visit(&mut checker)?; + Ok((!checker.prevents_pushdown()).then_some(checker.required_columns.into_iter().collect())) +} + +/// Calculate the total compressed size of all `Column`'s required for +/// predicate `Expr`. +/// +/// This value represents the total amount of IO required to evaluate the +/// predicate. +fn size_of_columns(columns: &[usize], metadata: &ParquetMetaData) -> Result { + let mut total_size = 0; + let row_groups = metadata.row_groups(); + for idx in columns { + for rg in row_groups.iter() { + total_size += rg.column(*idx).compressed_size() as usize; + } + } + + Ok(total_size) +} + +/// For a given set of `Column`s required for predicate `Expr` determine whether +/// all columns are sorted. +/// +/// Sorted columns may be queried more efficiently in the presence of +/// a PageIndex. +fn columns_sorted(_columns: &[usize], _metadata: &ParquetMetaData) -> Result { + // TODO How do we know this? + Ok(false) +} + +/// Build a [`LiquidRowFilter`] from the given predicate `Expr` if possible +/// +/// # returns +/// * `Ok(Some(row_filter))` if the expression can be used as RowFilter +/// * `Ok(None)` if the expression cannot be used as an RowFilter +/// * `Err(e)` if an error occurs while building the filter +/// +/// Note that the returned `RowFilter` may not contains all conjuncts in the +/// original expression. This is because some conjuncts may not be able to be +/// evaluated as an `ArrowPredicate` and will be ignored. +/// +/// For example, if the expression is `a = 1 AND b = 2 AND c = 3` and `b = 2` +/// can not be evaluated for some reason, the returned `RowFilter` will contain +/// `a = 1` and `c = 3`. +pub fn build_row_filter( + expr: &Arc, + physical_file_schema: &SchemaRef, + metadata: &ParquetMetaData, + reorder_predicates: bool, + file_metrics: &ParquetFileMetrics, +) -> Result> { + let rows_pruned = &file_metrics.pushdown_rows_pruned; + let rows_matched = &file_metrics.pushdown_rows_matched; + let time = &file_metrics.row_pushdown_eval_time; + + // Split into conjuncts: + // `a = 1 AND b = 2 AND c = 3` -> [`a = 1`, `b = 2`, `c = 3`] + let predicates = split_conjunction(expr); + + // Determine which conjuncts can be evaluated as ArrowPredicates, if any + let mut candidates: Vec = predicates + .into_iter() + .map(|expr| { + FilterCandidateBuilder::new(Arc::clone(expr), physical_file_schema.clone()) + .build(metadata) + }) + .collect::, _>>()? + .into_iter() + .flatten() + .collect(); + + // no candidates + if candidates.is_empty() { + return Ok(None); + } + + if reorder_predicates { + candidates.sort_unstable_by(|c1, c2| match c1.can_use_index.cmp(&c2.can_use_index) { + Ordering::Equal => c1.required_bytes.cmp(&c2.required_bytes), + ord => ord, + }); + } + + candidates.sort_unstable_by(|c1, c2| { + let p1 = get_priority(&c1.expr); + let p2 = get_priority(&c2.expr); + p1.cmp(&p2) + }); + + log::debug!( + "Predicate eval order: {}", + candidates + .iter() + .map(|c| c.expr.as_ref().to_string()) + .collect::>() + .join(", ") + ); + + candidates + .into_iter() + .map(|candidate| { + let projection = ProjectionMask::roots( + metadata.file_metadata().schema_descr(), + candidate.projection.iter().copied(), + ); + LiquidPredicate::try_new_with_metrics( + candidate, + projection, + rows_pruned.clone(), + rows_matched.clone(), + time.clone(), + ) + }) + .collect::, _>>() + .map(|filters| Some(LiquidRowFilter::new(filters))) +} + +fn get_priority(expr: &Arc) -> u8 { + if let Some(binary) = expr.downcast_ref::() { + match binary.op() { + Operator::Eq | Operator::NotEq => 0, // Highest priority + Operator::LikeMatch | Operator::ILikeMatch => 1, + Operator::NotLikeMatch | Operator::NotILikeMatch => 2, + Operator::Lt | Operator::LtEq | Operator::Gt | Operator::GtEq => 3, + _ => 4, + } + } else if expr.is::() { + 1 // LIKE expressions + } else { + 5 // All other expression types + } +} diff --git a/sandbox/libs/dataformat-native/rust/liquid-cache/datafusion/src/reader/plantime/row_group_filter.rs b/sandbox/libs/dataformat-native/rust/liquid-cache/datafusion/src/reader/plantime/row_group_filter.rs new file mode 100644 index 0000000000000..eb26b34b052ae --- /dev/null +++ b/sandbox/libs/dataformat-native/rust/liquid-cache/datafusion/src/reader/plantime/row_group_filter.rs @@ -0,0 +1,427 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use arrow::{array::ArrayRef, array::BooleanArray, array::UInt64Array, datatypes::Schema}; +use datafusion::common::{Column, Result, ScalarValue}; +use datafusion::datasource::listing::FileRange; +use datafusion::datasource::physical_plan::ParquetFileMetrics; +use datafusion::datasource::physical_plan::parquet::ParquetAccessPlan; +use datafusion::physical_optimizer::pruning::{PruningPredicate, PruningStatistics}; +use parquet::arrow::arrow_reader::statistics::StatisticsConverter; +use parquet::arrow::parquet_column; +use parquet::basic::Type; +use parquet::data_type::Decimal; +use parquet::schema::types::SchemaDescriptor; +use parquet::{ + arrow::{ParquetRecordBatchStreamBuilder, async_reader::AsyncFileReader}, + bloom_filter::Sbbf, + file::metadata::RowGroupMetaData, +}; +use std::collections::{HashMap, HashSet}; +use std::sync::Arc; + +/// Reduces the [`ParquetAccessPlan`] based on row group level metadata. +/// +/// This struct implements the various types of pruning that are applied to a +/// set of row groups within a parquet file, progressively narrowing down the +/// set of row groups (and ranges/selections within those row groups) that +/// should be scanned, based on the available metadata. +#[derive(Debug, Clone, PartialEq)] +pub struct RowGroupAccessPlanFilter { + /// which row groups should be accessed + access_plan: ParquetAccessPlan, +} + +impl RowGroupAccessPlanFilter { + /// Create a new `RowGroupPlanBuilder` for pruning out the groups to scan + /// based on metadata and statistics + pub fn new(access_plan: ParquetAccessPlan) -> Self { + Self { access_plan } + } + + /// Return true if there are no row groups + pub fn is_empty(&self) -> bool { + self.access_plan.is_empty() + } + + /// Returns the inner access plan + pub fn build(self) -> ParquetAccessPlan { + self.access_plan + } + + /// Prune remaining row groups to only those within the specified range. + /// + /// Updates this set to mark row groups that should not be scanned + /// + /// # Panics + /// if `groups.len() != self.len()` + pub fn prune_by_range(&mut self, groups: &[RowGroupMetaData], range: &FileRange) { + assert_eq!(groups.len(), self.access_plan.len()); + for (idx, metadata) in groups.iter().enumerate() { + if !self.access_plan.should_scan(idx) { + continue; + } + + // Skip the row group if the first dictionary/data page are not + // within the range. + // + // note don't use the location of metadata + // + let col = metadata.column(0); + let offset = col + .dictionary_page_offset() + .unwrap_or_else(|| col.data_page_offset()); + if !range.contains(offset) { + self.access_plan.skip(idx); + } + } + } + /// Prune remaining row groups using min/max/null_count statistics and + /// the [`PruningPredicate`] to determine if the predicate can not be true. + /// + /// Updates this set to mark row groups that should not be scanned + /// + /// Note: This method currently ignores ColumnOrder + /// + /// + /// # Panics + /// if `groups.len() != self.len()` + pub fn prune_by_statistics( + &mut self, + arrow_schema: &Schema, + parquet_schema: &SchemaDescriptor, + groups: &[RowGroupMetaData], + predicate: &PruningPredicate, + metrics: &ParquetFileMetrics, + ) { + // scoped timer updates on drop + let _timer_guard = metrics.statistics_eval_time.timer(); + + assert_eq!(groups.len(), self.access_plan.len()); + // Indexes of row groups still to scan + let row_group_indexes = self.access_plan.row_group_indexes(); + let row_group_metadatas = row_group_indexes + .iter() + .map(|&i| &groups[i]) + .collect::>(); + + let pruning_stats = RowGroupPruningStatistics { + parquet_schema, + row_group_metadatas, + arrow_schema, + }; + + // try to prune the row groups in a single call + match predicate.prune(&pruning_stats) { + Ok(values) => { + // values[i] is false means the predicate could not be true for row group i + for (idx, &value) in row_group_indexes.iter().zip(values.iter()) { + if !value { + self.access_plan.skip(*idx); + metrics.row_groups_pruned_statistics.add_pruned(1); + } else { + metrics.row_groups_pruned_statistics.add_matched(1); + } + } + } + // stats filter array could not be built, so we can't prune + Err(e) => { + log::debug!("Error evaluating row group predicate values {e}"); + metrics.predicate_evaluation_errors.add(1); + } + } + } + + /// Prune remaining row groups using available bloom filters and the + /// [`PruningPredicate`]. + /// + /// Updates this set with row groups that should not be scanned + /// + /// # Panics + /// if the builder does not have the same number of row groups as this set + pub async fn prune_by_bloom_filters( + &mut self, + arrow_schema: &Schema, + builder: &mut ParquetRecordBatchStreamBuilder, + predicate: &PruningPredicate, + metrics: &ParquetFileMetrics, + ) { + // scoped timer updates on drop + let _timer_guard = metrics.bloom_filter_eval_time.timer(); + + assert_eq!(builder.metadata().num_row_groups(), self.access_plan.len()); + for idx in 0..self.access_plan.len() { + if !self.access_plan.should_scan(idx) { + continue; + } + + // Attempt to find bloom filters for filtering this row group + let literal_columns = predicate.literal_columns(); + let mut column_sbbf = HashMap::with_capacity(literal_columns.len()); + + for column_name in literal_columns { + let Some((column_idx, _field)) = + parquet_column(builder.parquet_schema(), arrow_schema, &column_name) + else { + continue; + }; + + let bf = match builder + .get_row_group_column_bloom_filter(idx, column_idx) + .await + { + Ok(Some(bf)) => bf, + Ok(None) => continue, // no bloom filter for this column + Err(e) => { + log::debug!("Ignoring error reading bloom filter: {e}"); + metrics.predicate_evaluation_errors.add(1); + continue; + } + }; + let physical_type = builder.parquet_schema().column(column_idx).physical_type(); + + column_sbbf.insert(column_name.to_string(), (bf, physical_type)); + } + + let stats = BloomFilterStatistics { column_sbbf }; + + // Can this group be pruned? + let prune_group = match predicate.prune(&stats) { + Ok(values) => !values[0], + Err(e) => { + log::debug!("Error evaluating row group predicate on bloom filter: {e}"); + metrics.predicate_evaluation_errors.add(1); + false + } + }; + + if prune_group { + metrics.row_groups_pruned_bloom_filter.add_pruned(1); + self.access_plan.skip(idx) + } else if !stats.column_sbbf.is_empty() { + metrics.row_groups_pruned_bloom_filter.add_matched(1); + } + } + } +} +/// Implements [`PruningStatistics`] for Parquet Split Block Bloom Filters (SBBF) +struct BloomFilterStatistics { + /// Maps column name to the parquet bloom filter and parquet physical type + column_sbbf: HashMap, +} + +impl BloomFilterStatistics { + /// Helper function for checking if [`Sbbf`] filter contains [`ScalarValue`]. + /// + /// In case the type of scalar is not supported, returns `true`, assuming that the + /// value may be present. + fn check_scalar(sbbf: &Sbbf, value: &ScalarValue, parquet_type: &Type) -> bool { + match value { + ScalarValue::Utf8(Some(v)) + | ScalarValue::Utf8View(Some(v)) + | ScalarValue::LargeUtf8(Some(v)) => sbbf.check(&v.as_str()), + ScalarValue::Binary(Some(v)) + | ScalarValue::BinaryView(Some(v)) + | ScalarValue::LargeBinary(Some(v)) => sbbf.check(v), + ScalarValue::FixedSizeBinary(_size, Some(v)) => sbbf.check(v), + ScalarValue::Boolean(Some(v)) => sbbf.check(v), + ScalarValue::Float64(Some(v)) => sbbf.check(v), + ScalarValue::Float32(Some(v)) => sbbf.check(v), + ScalarValue::Int64(Some(v)) => sbbf.check(v), + ScalarValue::Int32(Some(v)) => sbbf.check(v), + ScalarValue::UInt64(Some(v)) => sbbf.check(v), + ScalarValue::UInt32(Some(v)) => sbbf.check(v), + ScalarValue::Decimal128(Some(v), p, s) => match parquet_type { + Type::INT32 => { + //https://github.com/apache/parquet-format/blob/eb4b31c1d64a01088d02a2f9aefc6c17c54cc6fc/Encodings.md?plain=1#L35-L42 + // All physical type are little-endian + if *p > 9 { + //DECIMAL can be used to annotate the following types: + // + // int32: for 1 <= precision <= 9 + // int64: for 1 <= precision <= 18 + return true; + } + let b = (*v as i32).to_le_bytes(); + // Use Decimal constructor after https://github.com/apache/arrow-rs/issues/5325 + let decimal = Decimal::Int32 { + value: b, + precision: *p as i32, + scale: *s as i32, + }; + sbbf.check(&decimal) + } + Type::INT64 => { + if *p > 18 { + return true; + } + let b = (*v as i64).to_le_bytes(); + let decimal = Decimal::Int64 { + value: b, + precision: *p as i32, + scale: *s as i32, + }; + sbbf.check(&decimal) + } + Type::FIXED_LEN_BYTE_ARRAY => { + // keep with from_bytes_to_i128 + let b = v.to_be_bytes().to_vec(); + // Use Decimal constructor after https://github.com/apache/arrow-rs/issues/5325 + let decimal = Decimal::Bytes { + value: b.into(), + precision: *p as i32, + scale: *s as i32, + }; + sbbf.check(&decimal) + } + _ => true, + }, + // One more pattern matching since not all data types are supported + // inside of a Dictionary + ScalarValue::Dictionary(_, inner) => match inner.as_ref() { + ScalarValue::Int32(_) + | ScalarValue::Int64(_) + | ScalarValue::UInt32(_) + | ScalarValue::UInt64(_) + | ScalarValue::Float32(_) + | ScalarValue::Float64(_) + | ScalarValue::Utf8(_) + | ScalarValue::LargeUtf8(_) + | ScalarValue::Binary(_) + | ScalarValue::LargeBinary(_) => { + BloomFilterStatistics::check_scalar(sbbf, inner, parquet_type) + } + _ => true, + }, + _ => true, + } + } +} + +impl PruningStatistics for BloomFilterStatistics { + fn min_values(&self, _column: &Column) -> Option { + None + } + + fn max_values(&self, _column: &Column) -> Option { + None + } + + fn num_containers(&self) -> usize { + 1 + } + + fn null_counts(&self, _column: &Column) -> Option { + None + } + + fn row_counts(&self) -> Option { + None + } + + /// Use bloom filters to determine if we are sure this column can not + /// possibly contain `values` + /// + /// The `contained` API returns false if the bloom filters knows that *ALL* + /// of the values in a column are not present. + fn contained(&self, column: &Column, values: &HashSet) -> Option { + let (sbbf, parquet_type) = self.column_sbbf.get(column.name.as_str())?; + + // Bloom filters are probabilistic data structures that can return false + // positives (i.e. it might return true even if the value is not + // present) however, the bloom filter will return `false` if the value is + // definitely not present. + + let known_not_present = values + .iter() + .map(|value| BloomFilterStatistics::check_scalar(sbbf, value, parquet_type)) + // The row group doesn't contain any of the values if + // all the checks are false + .all(|v| !v); + + let contains = if known_not_present { + Some(false) + } else { + // Given the bloom filter is probabilistic, we can't be sure that + // the row group actually contains the values. Return `None` to + // indicate this uncertainty + None + }; + + Some(BooleanArray::from(vec![contains])) + } +} + +/// Wraps a slice of [`RowGroupMetaData`] in a way that implements [`PruningStatistics`] +struct RowGroupPruningStatistics<'a> { + parquet_schema: &'a SchemaDescriptor, + row_group_metadatas: Vec<&'a RowGroupMetaData>, + arrow_schema: &'a Schema, +} + +impl<'a> RowGroupPruningStatistics<'a> { + /// Return an iterator over the row group metadata + fn metadata_iter(&'a self) -> impl Iterator + 'a { + self.row_group_metadatas.iter().copied() + } + + fn statistics_converter<'b>(&'a self, column: &'b Column) -> Result> { + Ok(StatisticsConverter::try_new( + &column.name, + self.arrow_schema, + self.parquet_schema, + )?) + } +} + +impl PruningStatistics for RowGroupPruningStatistics<'_> { + fn min_values(&self, column: &Column) -> Option { + self.statistics_converter(column) + .and_then(|c| Ok(c.row_group_mins(self.metadata_iter())?)) + .ok() + } + + fn max_values(&self, column: &Column) -> Option { + self.statistics_converter(column) + .and_then(|c| Ok(c.row_group_maxes(self.metadata_iter())?)) + .ok() + } + + fn num_containers(&self) -> usize { + self.row_group_metadatas.len() + } + + fn null_counts(&self, column: &Column) -> Option { + self.statistics_converter(column) + .and_then(|c| Ok(c.row_group_null_counts(self.metadata_iter())?)) + .ok() + .map(|counts| Arc::new(counts) as ArrayRef) + } + + fn row_counts(&self) -> Option { + // Row counts are container-level — read directly from row group metadata. + let counts: UInt64Array = self + .metadata_iter() + .map(|rg| Some(rg.num_rows() as u64)) + .collect(); + Some(Arc::new(counts) as ArrayRef) + } + + fn contained(&self, _column: &Column, _values: &HashSet) -> Option { + None + } +} diff --git a/sandbox/libs/dataformat-native/rust/liquid-cache/datafusion/src/reader/plantime/source.rs b/sandbox/libs/dataformat-native/rust/liquid-cache/datafusion/src/reader/plantime/source.rs new file mode 100644 index 0000000000000..fa9e3c7b3590a --- /dev/null +++ b/sandbox/libs/dataformat-native/rust/liquid-cache/datafusion/src/reader/plantime/source.rs @@ -0,0 +1,399 @@ +use super::opener::LiquidParquetOpener; +use crate::cache::LiquidCacheParquetRef; +use crate::reader::plantime::engagement_policy::{ + CacheEngagementPolicy, default_engagement_policy, +}; +use ahash::{HashMap, HashMapExt}; +use arrow_schema::Schema; +use bytes::Bytes; +use datafusion::{ + config::TableParquetOptions, + datasource::{ + listing::PartitionedFile, + physical_plan::{ + FileScanConfig, FileSource, ParquetFileMetrics, ParquetFileReaderFactory, + ParquetSource, parquet::PagePruningAccessPlanFilter, + }, + table_schema::TableSchema, + }, + error::Result, + physical_expr::projection::ProjectionExprs, + physical_expr_adapter::DefaultPhysicalExprAdapterFactory, + physical_optimizer::pruning::PruningPredicate, + physical_plan::{ + PhysicalExpr, + metrics::{ExecutionPlanMetricsSet, MetricBuilder}, + }, +}; +use futures::{FutureExt, future::BoxFuture}; +use object_store::{ObjectStore, path::Path}; +use parquet::{ + arrow::{ + arrow_reader::ArrowReaderOptions, + async_reader::{AsyncFileReader, ParquetObjectReader}, + }, + file::metadata::{PageIndexPolicy, ParquetMetaData, ParquetMetaDataReader}, +}; +use std::{ + ops::Range, + sync::{Arc, LazyLock}, +}; +use tokio::sync::RwLock; + +static META_CACHE: LazyLock = LazyLock::new(MetadataCache::new); + +/// Pre-seed the metadata cache with already-loaded metadata. +/// Callers that have pre-loaded parquet metadata (e.g., from a custom +/// ParquetFileReaderFactory) can inject it here so that LC's opener +/// skips the expensive `ArrowReaderMetadata::load_async()` call. +pub async fn pre_seed_metadata_cache(path: &Path, metadata: Arc) { + let mut cache = META_CACHE.val.write().await; + cache.entry(path.clone()).or_insert(metadata); +} + +#[derive(Debug)] +pub(crate) struct CachedMetaReaderFactory { + store: Arc, +} + +impl CachedMetaReaderFactory { + pub(crate) fn new(store: Arc) -> Self { + Self { store } + } + + pub(crate) fn create_liquid_reader( + &self, + partition_index: usize, + partitioned_file: PartitionedFile, + metadata_size_hint: Option, + metrics: &ExecutionPlanMetricsSet, + ) -> ParquetMetadataCacheReader { + let path = partitioned_file.object_meta.location.clone(); + let store = Arc::clone(&self.store); + let mut inner = ParquetObjectReader::new(store, path.clone()) + .with_file_size(partitioned_file.object_meta.size); + + if let Some(hint) = metadata_size_hint { + inner = inner.with_footer_size_hint(hint); + } + + ParquetMetadataCacheReader { + file_metrics: ParquetFileMetrics::new(partition_index, path.as_ref(), metrics), + inner, + path, + } + } +} + +impl ParquetFileReaderFactory for CachedMetaReaderFactory { + fn create_reader( + &self, + partition_index: usize, + partitioned_file: PartitionedFile, + metadata_size_hint: Option, + metrics: &ExecutionPlanMetricsSet, + ) -> Result> { + let reader = self.create_liquid_reader( + partition_index, + partitioned_file, + metadata_size_hint, + metrics, + ); + Ok(Box::new(reader)) + } +} + +struct MetadataCache { + val: RwLock>>, +} + +impl MetadataCache { + fn new() -> Self { + Self { + val: RwLock::new(HashMap::new()), + } + } +} + +#[derive(Clone)] +pub struct ParquetMetadataCacheReader { + file_metrics: ParquetFileMetrics, + inner: ParquetObjectReader, + path: Path, +} + +impl AsyncFileReader for ParquetMetadataCacheReader { + fn get_byte_ranges( + &mut self, + ranges: Vec>, + ) -> BoxFuture<'_, parquet::errors::Result>> { + let total: u64 = ranges.iter().map(|r| r.end - r.start).sum(); + self.file_metrics.bytes_scanned.add(total as usize); + self.inner.get_byte_ranges(ranges) + } + + fn get_bytes(&mut self, range: Range) -> BoxFuture<'_, parquet::errors::Result> { + self.file_metrics + .bytes_scanned + .add((range.end - range.start) as usize); + self.inner.get_bytes(range) + } + + fn get_metadata( + &mut self, + options: Option<&ArrowReaderOptions>, + ) -> BoxFuture<'_, parquet::errors::Result>> { + let path = self.path.clone(); + let options = options.cloned(); + async move { + // First check with read lock + { + let cache = META_CACHE.val.read().await; + if let Some(meta) = cache.get(&path) { + log::debug!("[LC-Meta] HIT path={}", path); + return Ok(meta.clone()); + } + } + + // Upgrade to write lock and double-check + let mut cache = META_CACHE.val.write().await; + match cache.entry(path.clone()) { + std::collections::hash_map::Entry::Occupied(entry) => { + log::debug!("[LC-Meta] HIT (race) path={}", path); + Ok(entry.get().clone()) + } + std::collections::hash_map::Entry::Vacant(entry) => { + log::debug!("[LC-Meta] MISS (loading from store) path={}", path); + let meta = self.inner.get_metadata(options.as_ref()).await?; + let meta = Arc::try_unwrap(meta).unwrap_or_else(|e| e.as_ref().clone()); + let mut reader = ParquetMetaDataReader::new_with_metadata(meta.clone()) + .with_page_index_policy(PageIndexPolicy::Optional); + reader.load_page_index(&mut self.inner).await?; + let meta = Arc::new(reader.finish()?); + entry.insert(meta.clone()); + Ok(meta) + } + } + } + .boxed() + } +} + +/// The data source for LiquidCache +#[derive(Clone)] +pub struct LiquidParquetSource { + metrics: ExecutionPlanMetricsSet, + predicate: Option>, + pruning_predicate: Option>, + page_pruning_predicate: Option>, + table_parquet_options: TableParquetOptions, + liquid_cache: LiquidCacheParquetRef, + batch_size: Option, + projection: ProjectionExprs, + table_schema: TableSchema, + /// Optional caller-provided reader factory with pre-loaded metadata. + /// When set, LC's opener uses this to get metadata instantly instead of + /// fetching from the object store. + parquet_file_reader_factory: Option>, + /// Policy that decides per-file whether to use LC stream or delegate to parquet. + engagement_policy: Arc, +} + +impl LiquidParquetSource { + fn reorder_filters(&self) -> bool { + self.table_parquet_options.global.reorder_filters + } + + /// Set the table schema for the LiquidParquetSource + pub fn with_table_schema(&self, table_schema: TableSchema) -> Self { + Self { + table_schema, + ..self.clone() + } + } + + /// Set predicate information, also sets pruning_predicate and page_pruning_predicate attributes + pub fn with_predicate( + mut self, + file_schema: Arc, + predicate: Arc, + ) -> Self { + let metrics = ExecutionPlanMetricsSet::new(); + let predicate_creation_errors = + MetricBuilder::new(&metrics).global_counter("num_predicate_creation_errors"); + + self.metrics = metrics; + self.predicate = Some(Arc::clone(&predicate)); + + match PruningPredicate::try_new(Arc::clone(&predicate), Arc::clone(&file_schema)) { + Ok(pruning_predicate) => { + if !pruning_predicate.always_true() { + self.pruning_predicate = Some(Arc::new(pruning_predicate)); + } + } + Err(e) => { + log::debug!("Could not create pruning predicate for: {e}"); + predicate_creation_errors.add(1); + } + }; + + let page_pruning_predicate = Arc::new(PagePruningAccessPlanFilter::new( + &predicate, + Arc::clone(&file_schema), + )); + self.page_pruning_predicate = Some(page_pruning_predicate); + + self + } + + /// Set predicate for row_filter only — no page-index pruning. + /// Used by the indexed path where the BoolNode's RowSelection is authoritative + /// and page-level statistics must not override it. + pub fn with_predicate_no_page_pruning( + mut self, + file_schema: Arc, + predicate: Arc, + ) -> Self { + let metrics = ExecutionPlanMetricsSet::new(); + let predicate_creation_errors = + MetricBuilder::new(&metrics).global_counter("num_predicate_creation_errors"); + self.metrics = metrics; + self.predicate = Some(Arc::clone(&predicate)); + + match PruningPredicate::try_new(Arc::clone(&predicate), Arc::clone(&file_schema)) { + Ok(pruning_predicate) => { + if !pruning_predicate.always_true() { + self.pruning_predicate = Some(Arc::new(pruning_predicate)); + } + } + Err(e) => { + log::debug!("Could not create pruning predicate for: {e}"); + predicate_creation_errors.add(1); + } + }; + + // Deliberately skip page_pruning_predicate — page-index stats must not + // prune pages that the caller's RowSelection already validated. + self + } + + /// Create a new LiquidParquetSource from a ParquetSource + pub fn from_parquet_source(source: ParquetSource, liquid_cache: LiquidCacheParquetRef) -> Self { + let predicate = source.filter(); + let reader_factory = source.parquet_file_reader_factory().cloned(); + + let table_schema = source.table_schema().clone(); + let file_schema = table_schema.file_schema().clone(); + let projection = source.projection().cloned().unwrap_or_else(|| { + let table_schema = table_schema.table_schema(); + ProjectionExprs::from_indices( + &(0..table_schema.fields().len()).collect::>(), + table_schema, + ) + }); + let mut v = Self { + table_schema, + table_parquet_options: source.table_parquet_options().clone(), + batch_size: Some(liquid_cache.batch_size()), + liquid_cache, + projection, + metrics: source.metrics().clone(), + predicate: None, + pruning_predicate: None, + page_pruning_predicate: None, + parquet_file_reader_factory: reader_factory, + engagement_policy: default_engagement_policy(), + }; + + if let Some(predicate) = predicate { + v = v.with_predicate(file_schema, predicate); + } + + v + } + + /// Set a custom cache engagement policy. + pub fn with_engagement_policy(mut self, policy: Arc) -> Self { + self.engagement_policy = policy; + self + } + + /// Get the predicate for the LiquidParquetSource + pub fn predicate(&self) -> Option> { + self.predicate.clone() + } +} + +impl FileSource for LiquidParquetSource { + fn create_file_opener( + &self, + object_store: Arc, + base_config: &FileScanConfig, + partition: usize, + ) -> Result> { + let expr_adapter_factory = base_config + .expr_adapter_factory + .clone() + .unwrap_or_else(|| Arc::new(DefaultPhysicalExprAdapterFactory) as _); + + let reader_factory = Arc::new(CachedMetaReaderFactory::new(object_store)); + + // If the caller provided a ParquetFileReaderFactory (with pre-loaded + // metadata), pre-seed LC's metadata cache for all files in this partition. + // This avoids expensive ArrowReaderMetadata::load_async() calls in the + // opener for files whose metadata we already have. + // Pass caller's reader factory to the opener so it can use pre-loaded + // metadata (avoids re-fetching from object store). + let caller_reader_factory = self.parquet_file_reader_factory.clone(); + + let opener = LiquidParquetOpener::new( + partition, + self.projection.clone(), + self.batch_size + .expect("Batch size must be set before creating LiquidParquetOpener"), + base_config.limit, + self.predicate.clone(), + self.table_schema.clone(), + self.metrics.clone(), + self.liquid_cache.clone(), + reader_factory, + self.reorder_filters(), + expr_adapter_factory, + caller_reader_factory, + Arc::clone(&self.engagement_policy), + ); + + Ok(Arc::new(opener)) + } + + fn with_batch_size(&self, batch_size: usize) -> Arc { + let mut conf = self.clone(); + conf.batch_size = Some(batch_size); + Arc::new(conf) + } + + fn table_schema(&self) -> &TableSchema { + &self.table_schema + } + + fn try_pushdown_projection( + &self, + projection: &ProjectionExprs, + ) -> Result>> { + let mut source = self.clone(); + source.projection = self.projection.try_merge(projection)?; + Ok(Some(Arc::new(source))) + } + + fn projection(&self) -> Option<&ProjectionExprs> { + Some(&self.projection) + } + + fn metrics(&self) -> &ExecutionPlanMetricsSet { + &self.metrics + } + + fn file_type(&self) -> &str { + "liquid_parquet" + } +} diff --git a/sandbox/libs/dataformat-native/rust/liquid-cache/datafusion/src/reader/runtime/liquid_cache_reader.rs b/sandbox/libs/dataformat-native/rust/liquid-cache/datafusion/src/reader/runtime/liquid_cache_reader.rs new file mode 100644 index 0000000000000..22dffbc7ff725 --- /dev/null +++ b/sandbox/libs/dataformat-native/rust/liquid-cache/datafusion/src/reader/runtime/liquid_cache_reader.rs @@ -0,0 +1,900 @@ +use std::collections::VecDeque; +use std::pin::Pin; +use std::sync::Arc; +use std::task::{Context, Poll}; + +use arrow::array::{Array, ArrayRef, BooleanArray, RecordBatch}; +use arrow::buffer::BooleanBuffer; +use arrow::compute::prep_null_mask_filter; +use arrow::record_batch::RecordBatchOptions; +use arrow_schema::{ArrowError, Schema, SchemaRef}; +use futures::{Stream, StreamExt, future::BoxFuture, stream::BoxStream}; +use parquet::arrow::arrow_reader::{ + ArrowPredicate, ArrowReaderMetadata, ArrowReaderOptions, RowSelection, RowSelector, +}; +use parquet::arrow::{ParquetRecordBatchStreamBuilder, ProjectionMask}; +use parquet::errors::ParquetError; +use parquet::file::metadata::ParquetMetaData; + +use crate::cache::{BatchID, CachedRowGroupRef}; +use crate::reader::plantime::{LiquidRowFilter, ParquetMetadataCacheReader}; +use crate::reader::runtime::utils::take_next_batch; +use crate::utils::{boolean_buffer_and_then, row_selector_to_boolean_buffer}; + +pub(crate) struct LiquidCacheReader { + state: ReaderState, + row_filter: Option, +} + +enum ReaderState { + Ready(Box), + Processing( + BoxFuture< + 'static, + ( + LiquidCacheReaderInner, + Option, + ProcessResult, + ), + >, + ), + Finished, +} + +enum ProcessResult { + Emit(Result), + Skip, +} + +struct LiquidCacheReaderInner { + cached_row_group: CachedRowGroupRef, + current_batch_id: BatchID, + selection: VecDeque, + schema: SchemaRef, + batch_size: usize, + projection_columns: Vec, + parquet_fallback: ParquetFallback, + last_pull: Option<(BatchID, RecordBatch)>, +} + +pub(crate) struct LiquidCacheReaderConfig { + pub(crate) batch_size: usize, + pub(crate) selection: RowSelection, + pub(crate) row_filter: Option, + pub(crate) cached_row_group: CachedRowGroupRef, + pub(crate) projection_columns: Vec, + pub(crate) schema: SchemaRef, + pub(crate) parquet_fallback: ParquetFallbackConfig, +} + +#[derive(Clone)] +pub(crate) struct ParquetFallbackConfig { + pub(crate) row_group_idx: usize, + pub(crate) metadata: Arc, + pub(crate) input: ParquetMetadataCacheReader, + pub(crate) cache_projection: ProjectionMask, + pub(crate) cache_column_ids: Vec, + pub(crate) cache_batch_size: usize, + pub(crate) row_count: usize, +} + +struct ParquetFallback { + row_group_idx: usize, + metadata: Arc, + input: ParquetMetadataCacheReader, + cache_projection: ProjectionMask, + cache_column_ids: Vec, + cache_batch_size: usize, + row_count: usize, + stream: Option>>, + next_batch_id: BatchID, +} + +impl LiquidCacheReader { + pub(crate) fn new(config: LiquidCacheReaderConfig) -> Self { + let inner = LiquidCacheReaderInner::new( + config.batch_size, + config.selection, + config.cached_row_group, + config.projection_columns, + Arc::clone(&config.schema), + ParquetFallback::new(config.parquet_fallback), + ); + Self { + state: ReaderState::Ready(Box::new(inner)), + row_filter: config.row_filter, + } + } + + pub(crate) fn into_filter(self) -> Option { + debug_assert!( + matches!(self.state, ReaderState::Finished), + "cannot extract filter before reader completes" + ); + self.row_filter + } +} + +impl Stream for LiquidCacheReader { + type Item = Result; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + loop { + let state = std::mem::replace(&mut self.state, ReaderState::Finished); + + match state { + ReaderState::Processing(mut fut) => match fut.as_mut().poll(cx) { + Poll::Pending => { + self.state = ReaderState::Processing(fut); + return Poll::Pending; + } + Poll::Ready((inner, row_filter, result)) => { + self.row_filter = row_filter; + self.state = ReaderState::Ready(Box::new(inner)); + match result { + ProcessResult::Emit(item) => return Poll::Ready(Some(item)), + ProcessResult::Skip => continue, + } + } + }, + ReaderState::Ready(mut inner) => { + match take_next_batch(&mut inner.selection, inner.batch_size) { + Some(selection) => { + let inner = *inner; + let future = inner.next_batch(self.row_filter.take(), selection); + self.state = ReaderState::Processing(future); + continue; + } + None => { + self.state = ReaderState::Finished; + return Poll::Ready(None); + } + } + } + ReaderState::Finished => { + self.state = ReaderState::Finished; + return Poll::Ready(None); + } + } + } + } +} + +impl ParquetFallback { + fn new(config: ParquetFallbackConfig) -> Self { + Self { + row_group_idx: config.row_group_idx, + metadata: config.metadata, + input: config.input, + cache_projection: config.cache_projection, + cache_column_ids: config.cache_column_ids, + cache_batch_size: config.cache_batch_size, + row_count: config.row_count, + stream: None, + next_batch_id: BatchID::from_raw(0), + } + } + + async fn fetch_batch(&mut self, batch_id: BatchID) -> Result { + if self.stream.is_none() || batch_id != self.next_batch_id { + self.rebuild_stream(batch_id)?; + } + + let stream = self.stream.as_mut().expect("fallback stream is present"); + let record_batch = stream.next().await.transpose()?.ok_or_else(|| { + ParquetError::General(format!( + "parquet fallback ended before batch {}", + *batch_id as usize + )) + })?; + + self.next_batch_id = batch_id; + self.next_batch_id.inc(); + Ok(record_batch) + } + + fn rebuild_stream(&mut self, batch_id: BatchID) -> Result<(), ParquetError> { + let reader_metadata = + ArrowReaderMetadata::try_new(Arc::clone(&self.metadata), ArrowReaderOptions::new())?; + let row_selection = + build_row_selection_from(batch_id, self.cache_batch_size, self.row_count); + + let stream = + ParquetRecordBatchStreamBuilder::new_with_metadata(self.input.clone(), reader_metadata) + .with_projection(self.cache_projection.clone()) + .with_row_groups(vec![self.row_group_idx]) + .with_batch_size(self.cache_batch_size) + .with_row_selection(row_selection) + .build()? + .boxed(); + + self.stream = Some(stream); + self.next_batch_id = batch_id; + Ok(()) + } +} + +fn build_row_selection_from( + batch_id: BatchID, + batch_size: usize, + row_count: usize, +) -> RowSelection { + let start = usize::from(*batch_id) * batch_size; + let mut selectors = Vec::new(); + + if start > 0 { + selectors.push(RowSelector::skip(start.min(row_count))); + } + + if start >= row_count { + return RowSelection::from(selectors); + } + + let mut remaining = row_count - start; + while remaining > 0 { + let selected = remaining.min(batch_size); + selectors.push(RowSelector::select(selected)); + remaining -= selected; + } + + RowSelection::from(selectors) +} + +impl LiquidCacheReaderInner { + fn new( + batch_size: usize, + selection: RowSelection, + cached_row_group: CachedRowGroupRef, + projection_columns: Vec, + schema: SchemaRef, + parquet_fallback: ParquetFallback, + ) -> Self { + Self { + cached_row_group, + current_batch_id: BatchID::from_raw(0), + selection: selection.into(), + schema, + batch_size, + projection_columns, + parquet_fallback, + last_pull: None, + } + } + + fn next_batch( + self, + row_filter: Option, + selection: Vec, + ) -> BoxFuture<'static, (Self, Option, ProcessResult)> { + Box::pin(async move { + let mut inner = self; + let mut row_filter = row_filter; + inner.last_pull = None; + + let result = match inner + .build_predicate_filter(&mut row_filter, selection) + .await + { + Ok(selection) => match inner.read_from_cache(&selection).await { + // read_from_cache is async because a cache miss falls back + // to reading the batch from parquet. + Ok(Some(batch)) => { + inner.current_batch_id.inc(); + ProcessResult::Emit(Ok(batch)) + } + Ok(None) => { + inner.current_batch_id.inc(); + ProcessResult::Skip + } + Err(e) => ProcessResult::Emit(Err(e)), + }, + Err(e) => ProcessResult::Emit(Err(e)), + }; + + (inner, row_filter, result) + }) + } + + async fn build_predicate_filter( + &mut self, + row_filter: &mut Option, + selection: Vec, + ) -> Result { + let mut input_selection = row_selector_to_boolean_buffer(&selection); + + let Some(filter) = row_filter.as_mut() else { + return Ok(input_selection); + }; + + for predicate in filter.predicates_mut() { + if input_selection.count_set_bits() == 0 { + break; + } + + let boolean_array = match self.cached_row_group.evaluate_selection_with_predicate( + self.current_batch_id, + &input_selection, + predicate, + ) { + Some(result) => result?, + None => { + self.evaluate_predicate_after_materialize(&input_selection, predicate) + .await? + } + }; + + let boolean_mask = if boolean_array.null_count() == 0 { + boolean_array.into_parts().0 + } else { + prep_null_mask_filter(&boolean_array).into_parts().0 + }; + + input_selection = boolean_buffer_and_then(&input_selection, &boolean_mask); + } + + Ok(input_selection) + } + + async fn read_from_cache( + &mut self, + selection: &BooleanBuffer, + ) -> Result, ArrowError> { + let selected_rows = selection.count_set_bits(); + if selected_rows == 0 { + return Ok(None); + } + + if self.projection_columns.is_empty() { + let options = RecordBatchOptions::new().with_row_count(Some(selected_rows)); + let batch = + RecordBatch::try_new_with_options(self.schema.clone(), Vec::new(), &options) + .unwrap(); + return Ok(Some(batch)); + } + + // Phase 1: Try cache for all columns, collect hits and track misses. + let mut arrays: Vec> = vec![None; self.projection_columns.len()]; + let mut has_miss = false; + let mut hit_count = 0usize; + let mut miss_count = 0usize; + + for (i, &column_idx) in self.projection_columns.iter().enumerate() { + let column = self + .cached_row_group + .get_column(column_idx as u64) + .ok_or_else(|| { + ArrowError::ComputeError(format!( + "column {column_idx} not present in liquid cache" + )) + })?; + + let array = column.get_arrow_array_with_filter(self.current_batch_id, selection); + + if let Some(arr) = array { + arrays[i] = Some(arr); + hit_count += 1; + } else { + has_miss = true; + miss_count += 1; + } + } + + if *self.current_batch_id == 0 { + log::debug!( + "[LC-Reader] batch_id={}, selected_rows={}, cols={}, hits={}, misses={}, fallback={}", + *self.current_batch_id, + selected_rows, + self.projection_columns.len(), + hit_count, + miss_count, + has_miss, + ); + } + + // Phase 2: If any columns missed, read from parquet ONCE for all misses. + if has_miss { + let record_batch = self + .read_parquet_batch_and_fill_cache(self.current_batch_id) + .await?; + + for (i, &column_idx) in self.projection_columns.iter().enumerate() { + if arrays[i].is_none() { + let array = self.parquet_array(&record_batch, column_idx)?; + arrays[i] = Some(filter_array(array, selection)?); + } + } + } + + let final_arrays: Vec = arrays.into_iter().map(|a| a.unwrap()).collect(); + Ok(Some( + RecordBatch::try_new(self.schema.clone(), final_arrays).unwrap(), + )) + } + + async fn read_parquet_batch_and_fill_cache( + &mut self, + batch_id: BatchID, + ) -> Result { + if let Some((pulled_batch_id, record_batch)) = &self.last_pull + && *pulled_batch_id == batch_id + { + return Ok(record_batch.clone()); + } + + let record_batch = self + .parquet_fallback + .fetch_batch(batch_id) + .await + .map_err(|e| ArrowError::ComputeError(format!("parquet fallback read failed: {e}")))?; + + // Spawn cache fill asynchronously — transcoding Arrow→Liquid should not + // block the query hot path. The batch is returned immediately; cache + // population happens in the background. + let cached_row_group = self.cached_row_group.clone(); + let cache_column_ids = self.parquet_fallback.cache_column_ids.clone(); + let batch_for_cache = record_batch.clone(); + tokio::spawn(async move { + for (col_idx, file_column_id) in cache_column_ids.iter().copied().enumerate() { + let Some(column) = cached_row_group.get_column(file_column_id as u64) else { + continue; + }; + let array = Arc::clone(batch_for_cache.column(col_idx)); + let _ = column.insert(batch_id, array); + } + }); + + self.last_pull = Some((batch_id, record_batch.clone())); + Ok(record_batch) + } + + async fn evaluate_predicate_after_materialize( + &mut self, + selection: &BooleanBuffer, + predicate: &mut crate::reader::LiquidPredicate, + ) -> Result { + let record_batch = self + .read_parquet_batch_and_fill_cache(self.current_batch_id) + .await?; + + if let Some(result) = self.cached_row_group.evaluate_selection_with_predicate( + self.current_batch_id, + selection, + predicate, + ) { + return result; + } + + let column_ids = predicate.predicate_column_ids(); + let mut arrays = Vec::with_capacity(column_ids.len()); + let mut fields = Vec::with_capacity(column_ids.len()); + + for column_id in column_ids { + let array = self.parquet_array(&record_batch, column_id)?; + arrays.push(filter_array(array, selection)?); + + let field = self + .cached_row_group + .get_column(column_id as u64) + .ok_or_else(|| { + ArrowError::ComputeError(format!( + "column {column_id} not present in liquid cache" + )) + })? + .field() + .as_ref() + .clone(); + fields.push(field); + } + + let schema = Arc::new(Schema::new(fields)); + let predicate_batch = if arrays.is_empty() { + let options = + RecordBatchOptions::new().with_row_count(Some(selection.count_set_bits())); + RecordBatch::try_new_with_options(schema, arrays, &options)? + } else { + RecordBatch::try_new(schema, arrays)? + }; + + predicate.evaluate(predicate_batch) + } + + fn parquet_array( + &self, + record_batch: &RecordBatch, + file_column_id: usize, + ) -> Result { + let position = self + .parquet_fallback + .cache_column_ids + .iter() + .position(|column_id| *column_id == file_column_id) + .ok_or_else(|| { + ArrowError::ComputeError(format!( + "column {file_column_id} not present in parquet fallback projection" + )) + })?; + + Ok(Arc::clone(record_batch.column(position))) + } +} + +fn filter_array(array: ArrayRef, selection: &BooleanBuffer) -> Result { + let selection_array = BooleanArray::new(selection.clone(), None); + arrow::compute::filter(array.as_ref(), &selection_array) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{ + cache::LiquidCacheParquet, + reader::plantime::CachedMetaReaderFactory, + reader::{FilterCandidateBuilder, LiquidPredicate, LiquidRowFilter}, + }; + use arrow::array::{ArrayRef, Int32Array}; + use arrow::record_batch::RecordBatch; + use arrow_schema::{DataType, Field, Schema, SchemaRef}; + use datafusion::datasource::listing::PartitionedFile; + use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet; + use datafusion::{ + logical_expr::Operator, + physical_expr::PhysicalExpr, + physical_expr::expressions::{BinaryExpr, Column, Literal}, + scalar::ScalarValue, + }; + use futures::{StreamExt, pin_mut}; + use liquid_cache::cache::TranscodeEvict; + use liquid_cache::cache_policies::LiquidPolicy; + use object_store::local::LocalFileSystem; + use parquet::arrow::{ + ArrowWriter, ProjectionMask, + arrow_reader::{ArrowReaderMetadata, ArrowReaderOptions, RowSelection, RowSelector}, + }; + use std::fs::File; + use std::sync::Arc; + + struct TestRowGroup { + batch_size: usize, + row_group: CachedRowGroupRef, + schema: SchemaRef, + fallback: ParquetFallbackConfig, + _tmp_dir: tempfile::TempDir, + } + + struct ReaderRequest { + selection: RowSelection, + row_filter: Option, + projection_columns: Vec, + schema: SchemaRef, + } + + impl TestRowGroup { + fn reader(&self, request: ReaderRequest) -> LiquidCacheReader { + LiquidCacheReader::new(LiquidCacheReaderConfig { + batch_size: self.batch_size, + selection: request.selection, + row_filter: request.row_filter, + cached_row_group: Arc::clone(&self.row_group), + projection_columns: request.projection_columns, + schema: request.schema, + parquet_fallback: self.fallback.clone(), + }) + } + } + + async fn make_row_group(batch_size: usize, batches: &[Vec]) -> TestRowGroup { + let tmp_dir = tempfile::tempdir().unwrap(); + let field = Arc::new(Field::new("col0", DataType::Int32, false)); + let schema = Arc::new(Schema::new(vec![field.clone()])); + let parquet_path = tmp_dir.path().join("data.parquet"); + let file = File::create(&parquet_path).unwrap(); + let mut writer = ArrowWriter::try_new(file, Arc::clone(&schema), None).unwrap(); + for values in batches { + let array: ArrayRef = Arc::new(Int32Array::from(values.clone())); + let batch = RecordBatch::try_new(Arc::clone(&schema), vec![array]).unwrap(); + writer.write(&batch).unwrap(); + } + writer.close().unwrap(); + + let metadata_file = File::open(&parquet_path).unwrap(); + let reader_metadata = + ArrowReaderMetadata::load(&metadata_file, ArrowReaderOptions::new()).unwrap(); + let object_store = Arc::new(LocalFileSystem::new_with_prefix(tmp_dir.path()).unwrap()); + let partitioned_file = PartitionedFile::new( + "data.parquet", + std::fs::metadata(&parquet_path).unwrap().len(), + ); + let metrics = ExecutionPlanMetricsSet::new(); + let input = CachedMetaReaderFactory::new(object_store).create_liquid_reader( + 0, + partitioned_file, + None, + &metrics, + ); + let projection = ProjectionMask::roots( + reader_metadata.metadata().file_metadata().schema_descr(), + [0], + ); + + let cache = LiquidCacheParquet::new( + batch_size, + usize::MAX, + Box::new(LiquidPolicy::new()), + Box::new(TranscodeEvict), + ); + let file = cache.register_or_get_file("test".to_string(), schema.clone()); + // Mark col0 as a predicate column so it is cacheable in tests. + let row_group = file.create_row_group(0, vec![0]); + let column = row_group.get_column(0).unwrap(); + + for (idx, values) in batches.iter().enumerate() { + let array: ArrayRef = Arc::new(Int32Array::from(values.clone())); + column + .insert(BatchID::from_raw(idx as u16), array) + .expect("cache insert"); + } + + TestRowGroup { + batch_size, + row_group, + schema, + fallback: ParquetFallbackConfig { + row_group_idx: 0, + metadata: Arc::clone(reader_metadata.metadata()), + input, + cache_projection: projection, + cache_column_ids: vec![0], + cache_batch_size: batch_size, + row_count: flatten_batches(batches).len(), + }, + _tmp_dir: tmp_dir, + } + } + + fn flatten_batches(batches: &[Vec]) -> Vec { + batches.iter().flat_map(|b| b.iter().copied()).collect() + } + + fn collect_batches(reader: LiquidCacheReader) -> Vec { + futures::executor::block_on( + reader + .map(|batch| batch.expect("valid record batch")) + .collect::>(), + ) + } + + fn as_i32_values(batch: &RecordBatch) -> Vec { + let array = batch + .column(0) + .as_any() + .downcast_ref::() + .expect("int column"); + array.iter().map(|v| v.expect("non-null")).collect() + } + + fn build_filter( + schema: SchemaRef, + values: &[i32], + expr: Arc, + ) -> LiquidRowFilter { + let tmp_meta = tempfile::NamedTempFile::new().unwrap(); + let array: ArrayRef = Arc::new(Int32Array::from(values.to_vec())); + let batch = RecordBatch::try_new(Arc::clone(&schema), vec![array.clone()]).unwrap(); + let mut writer = + ArrowWriter::try_new(tmp_meta.reopen().unwrap(), Arc::clone(&schema), None).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + + let file = std::fs::File::open(tmp_meta.path()).unwrap(); + let metadata = ArrowReaderMetadata::load(&file, ArrowReaderOptions::new()).unwrap(); + + let builder = FilterCandidateBuilder::new(expr, Arc::clone(&schema)); + let candidate = builder.build(metadata.metadata()).unwrap().unwrap(); + let projection = candidate.projection(metadata.metadata()); + let predicate = LiquidPredicate::try_new(candidate, projection).unwrap(); + + LiquidRowFilter::new(vec![predicate]) + } + + fn make_gt_filter(schema: SchemaRef, values: &[i32], literal: i32) -> LiquidRowFilter { + let expr: Arc = Arc::new(BinaryExpr::new( + Arc::new(Column::new("col0", 0)), + Operator::Gt, + Arc::new(Literal::new(ScalarValue::Int32(Some(literal)))), + )); + build_filter(schema, values, expr) + } + + #[tokio::test] + async fn reads_batches_in_order() { + let batch_size = 2; + let test = make_row_group(batch_size, &[vec![1, 2], vec![3, 4]]).await; + let selection = RowSelection::from(vec![RowSelector::select(4)]); + + let reader = test.reader(ReaderRequest { + selection, + row_filter: None, + projection_columns: vec![0], + schema: Arc::clone(&test.schema), + }); + + let batches = collect_batches(reader); + assert_eq!(batches.len(), 2); + assert_eq!(as_i32_values(&batches[0]), vec![1, 2]); + assert_eq!(as_i32_values(&batches[1]), vec![3, 4]); + } + + #[tokio::test] + async fn skips_unselected_batches() { + let batch_size = 2; + let test = make_row_group(batch_size, &[vec![1, 2], vec![3, 4]]).await; + let selection = RowSelection::from(vec![RowSelector::skip(2), RowSelector::select(2)]); + + let reader = test.reader(ReaderRequest { + selection, + row_filter: None, + projection_columns: vec![0], + schema: Arc::clone(&test.schema), + }); + + let batches = collect_batches(reader); + assert_eq!(batches.len(), 1); + assert_eq!(as_i32_values(&batches[0]), vec![3, 4]); + } + + #[tokio::test] + async fn empty_projection_emits_schema_only_batches() { + let batch_size = 2; + let test = make_row_group(batch_size, &[vec![10, 11]]).await; + let selection = RowSelection::from(vec![RowSelector::select(2)]); + + let reader = test.reader(ReaderRequest { + selection, + row_filter: None, + projection_columns: Vec::new(), + schema: Arc::new(Schema::new(Vec::::new())), + }); + + let batches = collect_batches(reader); + assert_eq!(batches.len(), 1); + let batch = &batches[0]; + assert_eq!(batch.num_columns(), 0); + assert_eq!(batch.num_rows(), 2); + } + + #[tokio::test] + async fn into_filter_returns_stored_filter_after_completion() { + let batch_size = 2; + let test = make_row_group(batch_size, &[vec![1, 2]]).await; + let selection = RowSelection::from(Vec::::new()); + let filter = LiquidRowFilter::new(Vec::new()); + + let mut reader = test.reader(ReaderRequest { + selection, + row_filter: Some(filter), + projection_columns: vec![0], + schema: Arc::clone(&test.schema), + }); + + let waker = futures::task::noop_waker(); + let mut cx = Context::from_waker(&waker); + assert!(matches!( + Pin::new(&mut reader).poll_next(&mut cx), + Poll::Ready(None) + )); + + assert!(reader.into_filter().is_some()); + } + + #[tokio::test] + async fn predicate_filters_rows_across_batches() { + let batches = vec![vec![1, 2], vec![3, 4]]; + let batch_size = 2; + let all_values = flatten_batches(&batches); + let test = make_row_group(batch_size, &batches).await; + let filter = make_gt_filter(Arc::clone(&test.schema), &all_values, 2); + let selection = RowSelection::from(vec![RowSelector::select(4)]); + + let reader = test.reader(ReaderRequest { + selection, + row_filter: Some(filter), + projection_columns: vec![0], + schema: Arc::clone(&test.schema), + }); + + let batches = collect_batches(reader); + assert_eq!(batches.len(), 1); + assert_eq!(as_i32_values(&batches[0]), vec![3, 4]); + } + + fn make_or_filter( + schema: SchemaRef, + values: &[i32], + gt_literal: i32, + lt_literal: i32, + ) -> LiquidRowFilter { + let cond1: Arc = Arc::new(BinaryExpr::new( + Arc::new(Column::new("col0", 0)), + Operator::Gt, + Arc::new(Literal::new(ScalarValue::Int32(Some(gt_literal)))), + )); + let cond2: Arc = Arc::new(BinaryExpr::new( + Arc::new(Column::new("col0", 0)), + Operator::Lt, + Arc::new(Literal::new(ScalarValue::Int32(Some(lt_literal)))), + )); + let expr: Arc = Arc::new(BinaryExpr::new(cond1, Operator::Or, cond2)); + build_filter(schema, values, expr) + } + + #[tokio::test] + async fn predicate_filters_or_rows() { + let batches = vec![vec![1, 2], vec![3, 4], vec![5, 6]]; + let batch_size = 2; + let all_values = flatten_batches(&batches); + let test = make_row_group(batch_size, &batches).await; + let filter = make_or_filter(Arc::clone(&test.schema), &all_values, 4, 2); + let selection = RowSelection::from(vec![RowSelector::select(6)]); + + let reader = test.reader(ReaderRequest { + selection, + row_filter: Some(filter), + projection_columns: vec![0], + schema: Arc::clone(&test.schema), + }); + + let batches = collect_batches(reader); + assert_eq!(batches.len(), 2); + assert_eq!(as_i32_values(&batches[0]), vec![1]); + assert_eq!(as_i32_values(&batches[1]), vec![5, 6]); + } + + #[tokio::test] + async fn predicate_combines_with_selection() { + let batches = vec![vec![1, 2, 3, 4]]; + let batch_size = 4; + let all_values = flatten_batches(&batches); + let test = make_row_group(batch_size, &batches).await; + let filter = make_gt_filter(Arc::clone(&test.schema), &all_values, 2); + let selection = RowSelection::from(vec![ + RowSelector::skip(1), + RowSelector::select(2), + RowSelector::skip(1), + ]); + + let reader = test.reader(ReaderRequest { + selection, + row_filter: Some(filter), + projection_columns: vec![0], + schema: Arc::clone(&test.schema), + }); + + let mut batches = collect_batches(reader); + assert_eq!(batches.len(), 1); + assert_eq!(as_i32_values(&batches.pop().unwrap()), vec![3]); + } + + #[tokio::test] + async fn predicate_can_filter_all_rows() { + let batches = vec![vec![1, 2]]; + let batch_size = 2; + let all_values = flatten_batches(&batches); + let test = make_row_group(batch_size, &batches).await; + let filter = make_gt_filter(Arc::clone(&test.schema), &all_values, 10); + let selection = RowSelection::from(vec![RowSelector::select(2)]); + + let reader = test.reader(ReaderRequest { + selection, + row_filter: Some(filter), + projection_columns: vec![0], + schema: Arc::clone(&test.schema), + }); + + let next_batch = futures::executor::block_on(async { + pin_mut!(reader); + reader.next().await + }); + + assert!(next_batch.is_none()); + } +} diff --git a/sandbox/libs/dataformat-native/rust/liquid-cache/datafusion/src/reader/runtime/liquid_predicate.rs b/sandbox/libs/dataformat-native/rust/liquid-cache/datafusion/src/reader/runtime/liquid_predicate.rs new file mode 100644 index 0000000000000..2cae21e29421d --- /dev/null +++ b/sandbox/libs/dataformat-native/rust/liquid-cache/datafusion/src/reader/runtime/liquid_predicate.rs @@ -0,0 +1,163 @@ +use std::sync::Arc; + +use datafusion::logical_expr::Operator; +use datafusion::physical_expr::PhysicalExpr; +use datafusion::physical_expr::expressions::{ + BinaryExpr, CastExpr, Column, LikeExpr, Literal, TryCastExpr, +}; + +/// Extract multiple column-literal expressions from a nested OR structure. +/// Returns a vector of (column_name, expression) pairs if all leaf expressions +/// are column-literal expressions connected by OR operators. +pub(crate) fn extract_multi_column_or( + expr: &Arc, +) -> Option)>> { + let mut result = Vec::new(); + + fn collect_or_expressions<'a>( + expr: &'a Arc, + result: &mut Vec<(&'a str, Arc)>, + ) -> bool { + if let Some(binary) = expr.downcast_ref::() + && binary.op() == &Operator::Or + { + // Recursively collect from left and right + return collect_or_expressions(binary.left(), result) + && collect_or_expressions(binary.right(), result); + } + + // Try to extract column-literal from this expression + if let Some(column_literal) = extract_column_literal(expr) { + result.push(column_literal); + true + } else { + false + } + } + + if collect_or_expressions(expr, &mut result) && result.len() >= 2 { + Some(result) + } else { + None + } +} + +fn extract_column_literal(expr: &Arc) -> Option<(&str, Arc)> { + if let Some(binary) = expr.downcast_ref::() + && binary.right().is::() + { + return extract_column_literal(binary.left()); + } else if let Some(like_expr) = expr.downcast_ref::() + && like_expr.pattern().is::() + { + return extract_column_literal(like_expr.expr()); + } else if let Some(cast_expr) = expr.downcast_ref::() { + return extract_column_literal(cast_expr.expr()); + } else if let Some(try_cast_expr) = expr.downcast_ref::() { + return extract_column_literal(try_cast_expr.expr()); + } else if let Some(column) = expr.downcast_ref::() { + return Some((column.name(), Arc::clone(expr))); + } + None +} + +#[cfg(test)] +mod tests { + use super::*; + use datafusion::common::ScalarValue; + use datafusion::logical_expr::Operator; + use datafusion::physical_expr::expressions::{BinaryExpr, Literal}; + use datafusion::physical_plan::expressions::Column; + + #[test] + fn test_extract_multi_column_or_valid_three_columns() { + // Test case: a = 1 OR b = 2 OR c = 3 + // This should extract 3 column-literal pairs + + let expr_a: Arc = Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::Eq, + Arc::new(Literal::new(ScalarValue::Int32(Some(1)))), + )); + + let expr_b: Arc = Arc::new(BinaryExpr::new( + Arc::new(Column::new("b", 1)), + Operator::Eq, + Arc::new(Literal::new(ScalarValue::Int32(Some(2)))), + )); + + let expr_c: Arc = Arc::new(BinaryExpr::new( + Arc::new(Column::new("c", 2)), + Operator::Eq, + Arc::new(Literal::new(ScalarValue::Int32(Some(3)))), + )); + + // Build nested OR: (a = 1 OR b = 2) OR c = 3 + let expr_ab = Arc::new(BinaryExpr::new(expr_a, Operator::Or, expr_b)); + let expr_final: Arc = + Arc::new(BinaryExpr::new(expr_ab, Operator::Or, expr_c)); + + let result = extract_multi_column_or(&expr_final); + + assert!(result.is_some()); + let column_exprs = result.unwrap(); + assert_eq!(column_exprs.len(), 3); + + // Verify we got the correct column names + let mut column_names: Vec<&str> = column_exprs.iter().map(|(name, _)| *name).collect(); + column_names.sort(); + assert_eq!(column_names, vec!["a", "b", "c"]); + } + + #[test] + fn test_extract_multi_column_or_invalid_expression() { + // Test case: a + b = 5 (not a column-literal OR expression) + // This should return None + + let add_expr: Arc = Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::Plus, + Arc::new(Column::new("b", 1)), + )); + + let expr: Arc = Arc::new(BinaryExpr::new( + add_expr, + Operator::Eq, + Arc::new(Literal::new(ScalarValue::Int32(Some(5)))), + )); + + let result = extract_multi_column_or(&expr); + assert!(result.is_none()); + + // Test case: Single column expression (a = 1) + // This should return None because we need >= 2 columns + let single_expr: Arc = Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::Eq, + Arc::new(Literal::new(ScalarValue::Int32(Some(1)))), + )); + + let result = extract_multi_column_or(&single_expr); + assert!(result.is_none()); + + // Test case: Mixed valid and invalid OR (a = 1 OR (b + c)) + // This should return None because one branch is not column-literal + let valid_expr: Arc = Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::Eq, + Arc::new(Literal::new(ScalarValue::Int32(Some(1)))), + )); + + let invalid_expr: Arc = Arc::new(BinaryExpr::new( + Arc::new(Column::new("b", 1)), + Operator::Plus, + Arc::new(Column::new("c", 2)), + )); + + let mixed_expr: Arc = + Arc::new(BinaryExpr::new(valid_expr, Operator::Or, invalid_expr)); + + let result = extract_multi_column_or(&mixed_expr); + assert!(result.is_none()); + } +} diff --git a/sandbox/libs/dataformat-native/rust/liquid-cache/datafusion/src/reader/runtime/liquid_stream.rs b/sandbox/libs/dataformat-native/rust/liquid-cache/datafusion/src/reader/runtime/liquid_stream.rs new file mode 100644 index 0000000000000..30669ae6fac76 --- /dev/null +++ b/sandbox/libs/dataformat-native/rust/liquid-cache/datafusion/src/reader/runtime/liquid_stream.rs @@ -0,0 +1,712 @@ +use crate::cache::{CachedFileRef, CachedRowGroupRef}; +use crate::reader::plantime::{LiquidRowFilter, ParquetMetadataCacheReader}; +use arrow::array::RecordBatch; +use arrow_schema::{Schema, SchemaRef}; +use futures::Stream; +use parquet::{ + arrow::{ + ProjectionMask, + arrow_reader::{ArrowPredicate, RowSelection, RowSelector}, + }, + errors::ParquetError, + file::metadata::ParquetMetaData, +}; +use std::{ + collections::VecDeque, + fmt::Formatter, + pin::Pin, + sync::Arc, + task::{Context, Poll}, +}; + +use super::liquid_cache_reader::{ + LiquidCacheReader, LiquidCacheReaderConfig, ParquetFallbackConfig, +}; +use super::utils::{get_root_column_ids, limit_row_selection, offset_row_selection}; + +type PlanResult = Option; + +struct ReaderFactory { + metadata: Arc, + + input: ParquetMetadataCacheReader, + + filter: Option, + + limit: Option, + + offset: Option, + + cached_file: CachedFileRef, +} + +impl ReaderFactory { + /// Plans what to read from cache vs parquet for the next row group + fn plan_row_group( + &mut self, + row_group_idx: usize, + selection: Option, + projection: ProjectionMask, + batch_size: usize, + ) -> PlanResult { + let meta = self.metadata.row_group(row_group_idx); + + let mut predicate_projection: Option = None; + if let Some(filter) = self.filter.as_mut() { + for predicate in filter.predicates_mut() { + let p_projection = predicate.projection(); + if let Some(ref mut p) = predicate_projection { + p.union(p_projection); + } else { + predicate_projection = Some(p_projection.clone()); + } + } + } + + let mut selection = + selection.unwrap_or_else(|| vec![RowSelector::select(meta.num_rows() as usize)].into()); + + let rows_before = selection.row_count(); + + if rows_before == 0 { + return None; + } + + if let Some(offset) = self.offset { + selection = offset_row_selection(selection, offset); + } + + if let Some(limit) = self.limit { + selection = limit_row_selection(selection, limit); + } + + let rows_after = selection.row_count(); + + // Update offset if necessary + if let Some(offset) = &mut self.offset { + // Reduction is either because of offset or limit, as limit is applied + // after offset has been "exhausted" can just use saturating sub here + *offset = offset.saturating_sub(rows_before - rows_after) + } + + if rows_after == 0 { + return None; + } + + if let Some(limit) = &mut self.limit { + *limit -= rows_after; + } + + let mut cache_projection = projection.clone(); + if let Some(ref predicate_projection) = predicate_projection { + cache_projection.union(predicate_projection); + } + + let schema_descr = self.metadata.file_metadata().schema_descr(); + let cache_column_ids = get_root_column_ids(schema_descr, &cache_projection); + // When no predicate is present, all projected columns are cacheable. + // Without this, `is_predicate_column` is false for every column and + // the cache is effectively disabled (get returns None, insert rejects). + let predicate_column_ids = if let Some(ref predicate_projection) = predicate_projection { + get_root_column_ids(schema_descr, predicate_projection) + } else { + cache_column_ids.clone() + }; + + if row_group_idx == 0 { + log::debug!( + "[LC-Stream] plan_row_group: rg={}, cache_cols={:?}, predicate_cols={:?}, \ + has_filter={}, rows={}", + row_group_idx, + &cache_column_ids, + &predicate_column_ids, + self.filter.is_some(), + self.metadata.row_group(row_group_idx).num_rows(), + ); + } + + let cached_row_group = self + .cached_file + .create_row_group(row_group_idx as u64, predicate_column_ids.clone()); + + let projection_column_ids = get_root_column_ids(schema_descr, &projection); + + let context = PlanningContext { + row_group_idx, + selection, + batch_size, + cached_row_group, + cache_projection, + projection_column_ids, + cache_column_ids, + }; + + Some(context) + } +} + +fn build_projection_schema(file_schema: &SchemaRef, projection_column_ids: &[usize]) -> SchemaRef { + let fields: Vec<_> = projection_column_ids + .iter() + .filter_map(|column_id| file_schema.fields().get(*column_id)) + .map(|field_ref| field_ref.as_ref().clone()) + .collect(); + Arc::new(Schema::new(fields)) +} + +/// Context for planning what to read from cache vs parquet +struct PlanningContext { + row_group_idx: usize, + selection: RowSelection, + batch_size: usize, + cached_row_group: CachedRowGroupRef, + cache_projection: ProjectionMask, + projection_column_ids: Vec, + cache_column_ids: Vec, +} + +fn build_liquid_cache_reader( + reader_factory: &mut ReaderFactory, + context: PlanningContext, + schema: SchemaRef, +) -> LiquidCacheReader { + let row_count = reader_factory + .metadata + .row_group(context.row_group_idx) + .num_rows() as usize; + let cache_batch_size = context.cached_row_group.batch_size(); + LiquidCacheReader::new(LiquidCacheReaderConfig { + batch_size: context.batch_size, + selection: context.selection, + row_filter: reader_factory.filter.take(), + cached_row_group: context.cached_row_group, + projection_columns: context.projection_column_ids, + schema, + parquet_fallback: ParquetFallbackConfig { + row_group_idx: context.row_group_idx, + metadata: Arc::clone(&reader_factory.metadata), + input: reader_factory.input.clone(), + cache_projection: context.cache_projection, + cache_column_ids: context.cache_column_ids, + cache_batch_size, + row_count, + }, + }) +} + +enum StreamState { + /// At the start of a new row group, or the end of the parquet stream + Init, + /// Decoding a batch from cache + ReadFromCache(Box), +} + +impl std::fmt::Debug for StreamState { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + match self { + StreamState::Init => write!(f, "StreamState::Init"), + StreamState::ReadFromCache(_) => write!(f, "StreamState::Decoding"), + } + } +} + +pub struct LiquidStreamBuilder { + pub(crate) input: ParquetMetadataCacheReader, + + pub(crate) metadata: Arc, + + pub(crate) batch_size: usize, + + pub(crate) row_groups: Option>, + + pub(crate) projection: ProjectionMask, + + pub(crate) filter: Option, + + pub(crate) selection: Option, + + pub(crate) limit: Option, + + pub(crate) offset: Option, +} + +impl LiquidStreamBuilder { + pub fn new(input: ParquetMetadataCacheReader, metadata: Arc) -> Self { + Self { + input, + metadata, + batch_size: 1024, + row_groups: None, + projection: ProjectionMask::all(), + filter: None, + selection: None, + limit: None, + offset: None, + } + } + + pub fn with_batch_size(mut self, batch_size: usize) -> Self { + self.batch_size = batch_size; + self + } + + pub fn with_row_groups(mut self, row_groups: Vec) -> Self { + self.row_groups = Some(row_groups); + self + } + + pub fn with_projection(mut self, projection: ProjectionMask) -> Self { + self.projection = projection; + self + } + + pub fn with_selection(mut self, selection: Option) -> Self { + self.selection = selection; + self + } + + pub fn with_limit(mut self, limit: Option) -> Self { + self.limit = limit; + self + } + + pub fn with_row_filter(mut self, filter: LiquidRowFilter) -> Self { + self.filter = Some(filter); + self + } + + pub fn build(self, liquid_cache: CachedFileRef) -> Result { + let num_row_groups = self.metadata.row_groups().len(); + + let row_groups: VecDeque = match self.row_groups { + Some(row_groups) => { + if let Some(col) = row_groups.iter().find(|x| **x >= num_row_groups) { + return Err(ParquetError::ArrowError(format!( + "row group {col} out of bounds 0..{num_row_groups}" + ))); + } + row_groups.into() + } + None => (0..self.metadata.row_groups().len()).collect(), + }; + + let batch_size = self + .batch_size + .min(self.metadata.file_metadata().num_rows() as usize); + + let schema_descr = self.metadata.file_metadata().schema_descr(); + let projection_column_ids = get_root_column_ids(schema_descr, &self.projection); + let file_schema = liquid_cache.schema(); + let schema = build_projection_schema(&file_schema, &projection_column_ids); + + let reader = ReaderFactory { + metadata: Arc::clone(&self.metadata), + input: self.input, + filter: self.filter, + limit: self.limit, + offset: self.offset, + cached_file: liquid_cache, + }; + + Ok(LiquidStream { + metadata: self.metadata, + schema, + row_groups, + projection: self.projection, + batch_size, + selection: self.selection, + reader: Some(reader), + state: StreamState::Init, + }) + } +} + +pub struct LiquidStream { + metadata: Arc, + + schema: SchemaRef, + + row_groups: VecDeque, + + projection: ProjectionMask, + + batch_size: usize, + + selection: Option, + + /// This is an option so it can be moved into a future + reader: Option, + + state: StreamState, +} + +impl std::fmt::Debug for LiquidStream { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ParquetRecordBatchStream") + .field("metadata", &self.metadata) + .field("schema", &self.schema) + .field("batch_size", &self.batch_size) + .field("projection", &self.projection) + .field("state", &self.state) + .finish() + } +} + +impl LiquidStream { + pub fn schema(&self) -> &SchemaRef { + &self.schema + } +} + +impl Stream for LiquidStream { + type Item = Result; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + loop { + let state = std::mem::replace(&mut self.state, StreamState::Init); + + match state { + StreamState::ReadFromCache(mut batch_reader) => { + match Pin::new(&mut *batch_reader).poll_next(cx) { + Poll::Ready(Some(Ok(batch))) => { + self.state = StreamState::ReadFromCache(batch_reader); + return Poll::Ready(Some(Ok(batch))); + } + Poll::Ready(Some(Err(e))) => { + panic!("Decoding next batch error: {e:?}"); + } + Poll::Ready(None) => { + let batch_reader = *batch_reader; + let filter = batch_reader.into_filter(); + self.reader.as_mut().unwrap().filter = filter; + // state left as Init, continue loop to plan next row group + } + Poll::Pending => { + self.state = StreamState::ReadFromCache(batch_reader); + return Poll::Pending; + } + } + } + StreamState::Init => { + let row_group_idx = match self.row_groups.pop_front() { + Some(idx) => idx, + None => return Poll::Ready(None), + }; + + let row_count = self.metadata.row_group(row_group_idx).num_rows() as usize; + + let selection = self.selection.as_mut().map(|s| s.split_off(row_count)); + + let projection = self.projection.clone(); + let batch_size = self.batch_size; + let maybe_context = self.reader.as_mut().expect("lost reader").plan_row_group( + row_group_idx, + selection, + projection, + batch_size, + ); + match maybe_context { + Some(context) => { + let schema = Arc::clone(&self.schema); + let reader_factory = self.reader.as_mut().unwrap(); + let batch_reader = + build_liquid_cache_reader(reader_factory, context, schema); + self.state = StreamState::ReadFromCache(Box::new(batch_reader)); + } + None => { + self.state = StreamState::Init; + } + } + } + } + } + } +} +#[cfg(test)] +mod tests { + use super::*; + use crate::cache::{BatchID, CachedFileRef, LiquidCacheParquet}; + use crate::reader::plantime::{ + CachedMetaReaderFactory, FilterCandidateBuilder, LiquidPredicate, LiquidRowFilter, + }; + use arrow::array::{ArrayRef, Int32Array}; + use arrow_schema::{DataType, Field, Schema}; + use datafusion::common::ScalarValue; + use datafusion::datasource::listing::PartitionedFile; + use datafusion::logical_expr::Operator; + use datafusion::physical_expr::PhysicalExpr; + use datafusion::physical_expr::expressions::{BinaryExpr, Column, Literal}; + use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet; + use futures::StreamExt; + use liquid_cache::cache::TranscodeEvict; + use liquid_cache::cache_policies::LiquidPolicy; + use object_store::local::LocalFileSystem; + use parquet::arrow::ArrowWriter; + use parquet::arrow::arrow_reader::{ArrowReaderMetadata, ArrowReaderOptions}; + use std::fs::File; + use std::sync::Arc; + + fn write_two_row_group_file(path: &std::path::Path, schema: SchemaRef) { + let file = File::create(path).unwrap(); + let mut writer = ArrowWriter::try_new(file, schema.clone(), None).unwrap(); + let batch0 = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![0, 1, 2, 3])), + Arc::new(Int32Array::from(vec![10, 11, 12, 13])), + ], + ) + .unwrap(); + let batch1 = RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from(vec![4, 5, 6, 7])), + Arc::new(Int32Array::from(vec![14, 15, 16, 17])), + ], + ) + .unwrap(); + writer.write(&batch0).unwrap(); + writer.flush().unwrap(); + writer.write(&batch1).unwrap(); + writer.close().unwrap(); + } + + fn make_liquid_stream( + max_memory_bytes: usize, + row_filter: Option, + projection_columns: Vec, + ) -> ( + LiquidStream, + Arc, + CachedFileRef, + tempfile::TempDir, + ) { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + ])); + let tmp_dir = tempfile::tempdir().unwrap(); + let parquet_path = tmp_dir.path().join("data.parquet"); + write_two_row_group_file(&parquet_path, schema.clone()); + let metadata_file = File::open(&parquet_path).unwrap(); + let reader_metadata = + ArrowReaderMetadata::load(&metadata_file, ArrowReaderOptions::new()).unwrap(); + let object_store = Arc::new(LocalFileSystem::new_with_prefix(tmp_dir.path()).unwrap()); + let partitioned_file = PartitionedFile::new( + "data.parquet", + std::fs::metadata(&parquet_path).unwrap().len(), + ); + let metrics = ExecutionPlanMetricsSet::new(); + let input = CachedMetaReaderFactory::new(object_store).create_liquid_reader( + 0, + partitioned_file, + None, + &metrics, + ); + + let cache = Arc::new(LiquidCacheParquet::new( + 4, + max_memory_bytes, + Box::new(LiquidPolicy::new()), + Box::new(TranscodeEvict), + )); + let cached_file = cache.register_or_get_file("data.parquet".to_string(), schema); + let projection = ProjectionMask::roots( + reader_metadata.metadata().file_metadata().schema_descr(), + projection_columns, + ); + let mut builder = LiquidStreamBuilder::new(input, Arc::clone(reader_metadata.metadata())) + .with_batch_size(4) + .with_row_groups(vec![0, 1]) + .with_projection(projection); + if let Some(row_filter) = row_filter { + builder = builder.with_row_filter(row_filter); + } + let stream = builder.build(cached_file.clone()).unwrap(); + (stream, cache, cached_file, tmp_dir) + } + + async fn collect_liquid_values(stream: LiquidStream) -> (Vec, Vec) { + let batches = stream + .map(|batch| batch.expect("valid liquid stream batch")) + .collect::>() + .await; + let mut a = Vec::new(); + let mut b = Vec::new(); + for batch in batches { + let a_array = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let b_array = batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + a.extend(a_array.iter().map(|value| value.unwrap())); + b.extend(b_array.iter().map(|value| value.unwrap())); + } + (a, b) + } + + async fn collect_projected_a(stream: LiquidStream) -> Vec { + let batches = stream + .map(|batch| batch.expect("valid liquid stream batch")) + .collect::>() + .await; + let mut a = Vec::new(); + for batch in batches { + let a_array = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + a.extend(a_array.iter().map(|value| value.unwrap())); + } + a + } + + fn gt_filter_on( + schema: SchemaRef, + col_name: &str, + col_idx: usize, + literal: i32, + ) -> LiquidRowFilter { + let expr: Arc = Arc::new(BinaryExpr::new( + Arc::new(Column::new(col_name, col_idx)), + Operator::Gt, + Arc::new(Literal::new(ScalarValue::Int32(Some(literal)))), + )); + let tmp_meta = tempfile::NamedTempFile::new().unwrap(); + write_two_row_group_file(tmp_meta.path(), schema.clone()); + let file = File::open(tmp_meta.path()).unwrap(); + let metadata = ArrowReaderMetadata::load(&file, ArrowReaderOptions::new()).unwrap(); + let builder = FilterCandidateBuilder::new(expr, schema); + let candidate = builder.build(metadata.metadata()).unwrap().unwrap(); + let projection = candidate.projection(metadata.metadata()); + let predicate = LiquidPredicate::try_new(candidate, projection).unwrap(); + LiquidRowFilter::new(vec![predicate]) + } + + /// Let the background cache-fill tasks (spawned via tokio::spawn) run. + async fn drain_background_tasks() { + for _ in 0..64 { + tokio::task::yield_now().await; + } + } + + fn is_cached(row_group: &CachedRowGroupRef, column_id: usize, batch_idx: u16) -> bool { + row_group + .get_column(column_id as u64) + .unwrap() + .get_arrow_array_test_only(BatchID::from_raw(batch_idx)) + .is_some() + } + + #[tokio::test] + async fn reads_two_row_groups_without_filter() { + let (stream, _cache, cached_file, _tmp_dir) = + make_liquid_stream(usize::MAX, None, vec![0, 1]); + + let (a, b) = collect_liquid_values(stream).await; + + assert_eq!(a, vec![0, 1, 2, 3, 4, 5, 6, 7]); + assert_eq!(b, vec![10, 11, 12, 13, 14, 15, 16, 17]); + + drain_background_tasks().await; + // Without a predicate, all projected columns are cacheable. + let row_group0 = cached_file.create_row_group(0, vec![0, 1]); + let row_group1 = cached_file.create_row_group(1, vec![0, 1]); + assert!(is_cached(&row_group0, 0, 0)); + assert!(is_cached(&row_group0, 1, 0)); + assert!(is_cached(&row_group1, 0, 0)); + assert!(is_cached(&row_group1, 1, 0)); + } + + #[tokio::test] + async fn row_filter_filters_rows() { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + ])); + let filter = gt_filter_on(schema, "a", 0, 2); + let (stream, _cache, _cached_file, _tmp_dir) = + make_liquid_stream(usize::MAX, Some(filter), vec![0, 1]); + + let (a, b) = collect_liquid_values(stream).await; + + assert_eq!(a, vec![3, 4, 5, 6, 7]); + assert_eq!(b, vec![13, 14, 15, 16, 17]); + } + + #[tokio::test] + async fn predicate_fallback_uses_predicate_projection() { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + ])); + let filter = gt_filter_on(schema, "b", 1, 13); + let (stream, _cache, cached_file, _tmp_dir) = + make_liquid_stream(usize::MAX, Some(filter), vec![0]); + + let a_values = collect_projected_a(stream).await; + + assert_eq!(a_values, vec![4, 5, 6, 7]); + + drain_background_tasks().await; + // With a predicate on `b`, only `b` is a predicate column and cacheable. + let row_group0 = cached_file.create_row_group(0, vec![1]); + let row_group1 = cached_file.create_row_group(1, vec![1]); + assert!(is_cached(&row_group0, 1, 0)); + assert!(is_cached(&row_group1, 1, 0)); + } + + #[tokio::test] + async fn zero_memory_budget_recovers() { + let (stream, _cache, cached_file, _tmp_dir) = make_liquid_stream(0, None, vec![0, 1]); + + let (a, b) = collect_liquid_values(stream).await; + + assert_eq!(a, vec![0, 1, 2, 3, 4, 5, 6, 7]); + assert_eq!(b, vec![10, 11, 12, 13, 14, 15, 16, 17]); + + drain_background_tasks().await; + let row_group0 = cached_file.create_row_group(0, vec![0, 1]); + let row_group1 = cached_file.create_row_group(1, vec![0, 1]); + assert!(!is_cached(&row_group0, 0, 0)); + assert!(!is_cached(&row_group0, 1, 0)); + assert!(!is_cached(&row_group1, 0, 0)); + assert!(!is_cached(&row_group1, 1, 0)); + } + + #[tokio::test] + async fn pre_populated_cache_serves_reads() { + let (stream, _cache, cached_file, _tmp_dir) = + make_liquid_stream(usize::MAX, None, vec![0, 1]); + let row_group0 = cached_file.create_row_group(0, vec![0, 1]); + let row_group1 = cached_file.create_row_group(1, vec![0, 1]); + let a0: ArrayRef = Arc::new(Int32Array::from(vec![0, 1, 2, 3])); + let a1: ArrayRef = Arc::new(Int32Array::from(vec![4, 5, 6, 7])); + row_group0 + .get_column(0) + .unwrap() + .insert(BatchID::from_raw(0), a0) + .unwrap(); + row_group1 + .get_column(0) + .unwrap() + .insert(BatchID::from_raw(0), a1) + .unwrap(); + + let (a, b) = collect_liquid_values(stream).await; + + assert_eq!(a, vec![0, 1, 2, 3, 4, 5, 6, 7]); + assert_eq!(b, vec![10, 11, 12, 13, 14, 15, 16, 17]); + + drain_background_tasks().await; + // Missing column `b` was pulled from parquet and back-filled. + assert!(is_cached(&row_group0, 1, 0)); + assert!(is_cached(&row_group1, 1, 0)); + } +} diff --git a/sandbox/libs/dataformat-native/rust/liquid-cache/datafusion/src/reader/runtime/mod.rs b/sandbox/libs/dataformat-native/rust/liquid-cache/datafusion/src/reader/runtime/mod.rs new file mode 100644 index 0000000000000..bfde28dcc9382 --- /dev/null +++ b/sandbox/libs/dataformat-native/rust/liquid-cache/datafusion/src/reader/runtime/mod.rs @@ -0,0 +1,7 @@ +pub(crate) use liquid_predicate::extract_multi_column_or; +pub(crate) use liquid_stream::LiquidStreamBuilder; + +mod liquid_cache_reader; +mod liquid_predicate; +mod liquid_stream; +mod utils; diff --git a/sandbox/libs/dataformat-native/rust/liquid-cache/datafusion/src/reader/runtime/utils.rs b/sandbox/libs/dataformat-native/rust/liquid-cache/datafusion/src/reader/runtime/utils.rs new file mode 100644 index 0000000000000..a05b79942b1c2 --- /dev/null +++ b/sandbox/libs/dataformat-native/rust/liquid-cache/datafusion/src/reader/runtime/utils.rs @@ -0,0 +1,212 @@ +use std::collections::VecDeque; + +use parquet::{ + arrow::{ + ProjectionMask, + arrow_reader::{RowSelection, RowSelector}, + }, + schema::types::SchemaDescriptor, +}; + +pub(crate) fn get_root_column_ids( + schema: &SchemaDescriptor, + projection: &ProjectionMask, +) -> Vec { + let mut root_mask = vec![false; schema.root_schema().get_fields().len()]; + + for leaf_idx in 0..schema.num_columns() { + if projection.leaf_included(leaf_idx) { + let root_idx = schema.get_column_root_idx(leaf_idx); + root_mask[root_idx] = true; + } + } + + root_mask + .into_iter() + .enumerate() + .filter_map(|(idx, included)| included.then_some(idx)) + .collect() +} + +pub(crate) fn offset_row_selection(selection: RowSelection, offset: usize) -> RowSelection { + if offset == 0 { + return selection; + } + + let mut selected_count = 0; + let mut skipped_count = 0; + + let mut selectors: Vec = selection.into(); + + let find = selectors.iter().position(|selector| match selector.skip { + true => { + skipped_count += selector.row_count; + false + } + false => { + selected_count += selector.row_count; + selected_count > offset + } + }); + + let split_idx = match find { + Some(idx) => idx, + None => { + selectors.clear(); + return RowSelection::from(selectors); + } + }; + + let mut new_selectors = Vec::with_capacity(selectors.len() - split_idx + 1); + new_selectors.push(RowSelector::skip(skipped_count + offset)); + new_selectors.push(RowSelector::select(selected_count - offset)); + new_selectors.extend_from_slice(&selectors[split_idx + 1..]); + + RowSelection::from(new_selectors) +} + +pub(crate) fn limit_row_selection(selection: RowSelection, mut limit: usize) -> RowSelection { + let mut selectors: Vec = selection.into(); + + if limit == 0 { + selectors.clear(); + } + + for (idx, selection) in selectors.iter_mut().enumerate() { + if !selection.skip { + if selection.row_count >= limit { + selection.row_count = limit; + selectors.truncate(idx + 1); + break; + } else { + limit -= selection.row_count; + } + } + } + RowSelection::from(selectors) +} + +/// Take the next batch from the selection queue. +/// The returning selection will have exactly the batch size, or less if the selection is exhausted. +pub(super) fn take_next_batch( + selection: &mut VecDeque, + batch_size: usize, +) -> Option> { + let mut current_selected = 0; + let mut rt = Vec::new(); + while let Some(mut front) = selection.pop_front() { + if front.row_count + current_selected > batch_size { + let to_select = batch_size - current_selected; + if to_select > 0 { + let mut sub_front = front; + sub_front.row_count = to_select; + rt.push(sub_front); + } + let remaining = front.row_count - to_select; + front.row_count = remaining; + selection.push_front(front); + current_selected += to_select; + break; + } else { + rt.push(front); + current_selected += front.row_count; + } + } + if current_selected == 0 { + return None; + } + Some(rt) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_take_next_batch() { + { + let mut queue = VecDeque::new(); + let selection = take_next_batch(&mut queue, 8); + assert!(selection.is_none()); + assert!(queue.is_empty()); + } + + { + let mut queue = VecDeque::from(vec![RowSelector::select(8)]); + let selection = take_next_batch(&mut queue, 8).unwrap(); + assert_eq!(selection, vec![RowSelector::select(8)]); + assert!(queue.is_empty()); + } + + { + let mut queue = VecDeque::from(vec![RowSelector::select(10)]); + let selection = take_next_batch(&mut queue, 8).unwrap(); + assert_eq!(selection, vec![RowSelector::select(8)]); + assert_eq!(queue, vec![RowSelector::select(2)]); + } + + { + let mut queue = VecDeque::from(vec![ + RowSelector::select(2), + RowSelector::skip(2), + RowSelector::select(2), + RowSelector::skip(2), + RowSelector::select(2), + ]); + let selection = take_next_batch(&mut queue, 8).unwrap(); + assert_eq!( + selection, + vec![ + RowSelector::select(2), + RowSelector::skip(2), + RowSelector::select(2), + RowSelector::skip(2), + ] + ); + assert_eq!(queue, vec![RowSelector::select(2)]); + } + + { + let mut queue = VecDeque::from(vec![RowSelector::select(3), RowSelector::skip(2)]); + let selection = take_next_batch(&mut queue, 8).unwrap(); + assert_eq!( + selection, + vec![RowSelector::select(3), RowSelector::skip(2)] + ); + assert!(queue.is_empty()); + } + + { + let mut queue = VecDeque::from(vec![ + RowSelector::select(2), + RowSelector::skip(4), + RowSelector::select(6), + ]); + let selection = take_next_batch(&mut queue, 8).unwrap(); + assert_eq!( + selection, + vec![ + RowSelector::select(2), + RowSelector::skip(4), + RowSelector::select(2), + ] + ); + assert_eq!(queue, vec![RowSelector::select(4)]); + } + + { + let mut queue = VecDeque::from(vec![ + RowSelector::skip(5), + RowSelector::select(3), + RowSelector::skip(2), + RowSelector::select(7), + ]); + let selection = take_next_batch(&mut queue, 8).unwrap(); + assert_eq!( + selection, + vec![RowSelector::skip(5), RowSelector::select(3),] + ); + assert_eq!(queue, vec![RowSelector::skip(2), RowSelector::select(7)]); + } + } +} diff --git a/sandbox/libs/dataformat-native/rust/liquid-cache/datafusion/src/reader/utils/boolean_selection.rs b/sandbox/libs/dataformat-native/rust/liquid-cache/datafusion/src/reader/utils/boolean_selection.rs new file mode 100644 index 0000000000000..c031d45e4c1f6 --- /dev/null +++ b/sandbox/libs/dataformat-native/rust/liquid-cache/datafusion/src/reader/utils/boolean_selection.rs @@ -0,0 +1,273 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::ops::Range; + +use arrow::array::{Array, BooleanArray, BooleanBufferBuilder}; +use arrow::buffer::{BooleanBuffer, MutableBuffer}; +use arrow::util::bit_iterator::BitIndexIterator; +use parquet::arrow::arrow_reader::{RowSelection, RowSelector}; + +/// A selection of rows using a boolean array. +#[derive(Debug, Clone, PartialEq)] +pub struct BooleanSelection { + selectors: BooleanBuffer, +} + +impl BooleanSelection { + /// Create a new BooleanSelection from a list of BooleanArray. + pub fn from_filters(filters: &[BooleanArray]) -> Self { + let arrays: Vec<&dyn Array> = filters.iter().map(|x| x as &dyn Array).collect(); + let result = arrow::compute::concat(&arrays).unwrap().into_data(); + let (boolean_array, _null) = BooleanArray::from(result).into_parts(); + BooleanSelection { + selectors: boolean_array, + } + } + + /// Create a new BooleanSelection with all rows unselected + pub fn new_unselected(row_count: usize) -> Self { + let buffer = BooleanBuffer::new_unset(row_count); + + BooleanSelection { selectors: buffer } + } + + /// Create a new BooleanSelection with all rows selected + pub fn new_selected(row_count: usize) -> Self { + let buffer = BooleanBuffer::new_set(row_count); + + BooleanSelection { selectors: buffer } + } + + /// Returns a new BooleanSelection that selects the inverse of this BooleanSelection. + pub fn as_inverted(&self) -> Self { + let buffer = !&self.selectors; + BooleanSelection { selectors: buffer } + } + + /// Returns the number of rows in this BooleanSelection. + pub fn len(&self) -> usize { + self.selectors.len() + } + + /// Check if the BooleanSelection is empty. + pub fn is_empty(&self) -> bool { + self.len() == 0 + } + + /// Returns the number of rows selected by this BooleanSelection. + pub fn row_count(&self) -> usize { + self.selectors.count_set_bits() + } + + /// Create a new BooleanSelection from a list of consecutive ranges. + pub fn from_consecutive_ranges( + ranges: impl Iterator>, + total_rows: usize, + ) -> Self { + let mut buffer = BooleanBufferBuilder::new(total_rows); + let mut last_end = 0; + + for range in ranges { + let len = range.end - range.start; + if len == 0 { + continue; + } + + if range.start > last_end { + buffer.append_n(range.start - last_end, false); + } + buffer.append_n(len, true); + last_end = range.end; + } + + if last_end != total_rows { + buffer.append_n(total_rows - last_end, false); + } + + BooleanSelection { + selectors: buffer.finish(), + } + } + + /// Compute the union of two [`RowSelection`] + /// For example: + /// self: NNYYYYNNYYNYN + /// other: NYNNNNNNN + /// + /// returned: NYYYYYNNYYNYN + #[must_use] + pub fn union(&self, other: &Self) -> Self { + // use arrow::compute::kernels::boolean::or; + + let union_selectors = &self.selectors | &other.selectors; + + BooleanSelection { + selectors: union_selectors, + } + } + + /// Compute the intersection of two [`RowSelection`] + /// For example: + /// self: NNYYYYNNYYNYN + /// other: NYNNNNNNY + /// + /// returned: NNNNNNNNYYNYN + #[must_use] + pub fn intersection(&self, other: &Self) -> Self { + let intersection_selectors = &self.selectors & &other.selectors; + + BooleanSelection { + selectors: intersection_selectors, + } + } + + /// Combines this `BooleanSelection` with another using logical AND on the selected bits. + /// + /// The `other` `BooleanSelection` must have exactly as many set bits as `self`. + /// This method will keep only the bits in `self` that are also set in `other` + /// at the positions corresponding to `self`'s set bits. + pub fn and_then(&self, other: &Self) -> Self { + // Ensure that 'other' has exactly as many set bits as 'self' + debug_assert_eq!( + self.row_count(), + other.len(), + "The 'other' selection must have exactly as many set bits as 'self'." + ); + + if self.len() == other.len() { + // fast path if the two selections are the same length + // common if this is the first predicate + debug_assert_eq!(self.row_count(), self.len()); + return self.intersection(other); + } + + let mut buffer = MutableBuffer::from_len_zeroed(self.selectors.inner().len()); + buffer.copy_from_slice(self.selectors.values()); + let mut builder = BooleanBufferBuilder::new_from_buffer(buffer, self.len()); + + // Create iterators for 'self' and 'other' bits + let mut other_bits = other.selectors.iter(); + + for bit_idx in self.positive_iter() { + let predicate = other_bits + .next() + .expect("Mismatch in set bits between self and other"); + if !predicate { + builder.set_bit(bit_idx, false); + } + } + + BooleanSelection { + selectors: builder.finish(), + } + } + + /// Returns an iterator over the indices of the set bits in this [`BooleanSelection`] + pub fn positive_iter(&self) -> BitIndexIterator<'_> { + self.selectors.set_indices() + } + + /// Returns `true` if this [BooleanSelection] selects any rows + pub fn selects_any(&self) -> bool { + self.row_count() > 0 + } + + /// Returns a new BooleanSelection that selects the rows in this BooleanSelection from `offset` to `offset + len` + pub fn slice(&self, offset: usize, len: usize) -> BooleanArray { + BooleanArray::new(self.selectors.slice(offset, len), None) + } +} + +impl From> for BooleanSelection { + fn from(selection: Vec) -> Self { + let selection = RowSelection::from(selection); + RowSelection::into(selection) + } +} + +impl From for BooleanSelection { + fn from(selection: RowSelection) -> Self { + let total_rows = selection.row_count(); + let mut builder = BooleanBufferBuilder::new(total_rows); + + for selector in selection.iter() { + if selector.skip { + builder.append_n(selector.row_count, false); + } else { + builder.append_n(selector.row_count, true); + } + } + + BooleanSelection { + selectors: builder.finish(), + } + } +} + +impl From<&BooleanSelection> for RowSelection { + fn from(selection: &BooleanSelection) -> Self { + let array = BooleanArray::new(selection.selectors.clone(), None); + RowSelection::from_filters(&[array]) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_boolean_selection_and_then() { + // Initial mask: 001011010101 + let self_filters = vec![BooleanArray::from(vec![ + false, false, true, false, true, true, false, true, false, true, false, true, + ])]; + let self_selection = BooleanSelection::from_filters(&self_filters); + + // Predicate mask (only for selected bits): 001101 + let other_filters = vec![BooleanArray::from(vec![ + false, false, true, true, false, true, + ])]; + let other_selection = BooleanSelection::from_filters(&other_filters); + + let result = self_selection.and_then(&other_selection); + + // Expected result: 000001010001 + let expected_filters = vec![BooleanArray::from(vec![ + false, false, false, false, false, true, false, true, false, false, false, true, + ])]; + let expected_selection = BooleanSelection::from_filters(&expected_filters); + + assert_eq!(result, expected_selection); + } + + #[test] + #[should_panic( + expected = "The 'other' selection must have exactly as many set bits as 'self'." + )] + fn test_and_then_mismatched_set_bits() { + let self_filters = vec![BooleanArray::from(vec![true, true, false])]; + let self_selection = BooleanSelection::from_filters(&self_filters); + + // 'other' has only one set bit, but 'self' has two + let other_filters = vec![BooleanArray::from(vec![true, false, false])]; + let other_selection = BooleanSelection::from_filters(&other_filters); + + // This should panic + let _ = self_selection.and_then(&other_selection); + } +} diff --git a/sandbox/libs/dataformat-native/rust/liquid-cache/datafusion/src/reader/utils/mod.rs b/sandbox/libs/dataformat-native/rust/liquid-cache/datafusion/src/reader/utils/mod.rs new file mode 100644 index 0000000000000..237e831a8d5d0 --- /dev/null +++ b/sandbox/libs/dataformat-native/rust/liquid-cache/datafusion/src/reader/utils/mod.rs @@ -0,0 +1 @@ +pub(crate) mod boolean_selection; diff --git a/sandbox/libs/dataformat-native/rust/liquid-cache/datafusion/src/sync.rs b/sandbox/libs/dataformat-native/rust/liquid-cache/datafusion/src/sync.rs new file mode 100644 index 0000000000000..4e2a741253538 --- /dev/null +++ b/sandbox/libs/dataformat-native/rust/liquid-cache/datafusion/src/sync.rs @@ -0,0 +1,2 @@ +#[allow(unused_imports)] +pub use std::{sync::*, thread}; diff --git a/sandbox/libs/dataformat-native/rust/liquid-cache/datafusion/src/utils.rs b/sandbox/libs/dataformat-native/rust/liquid-cache/datafusion/src/utils.rs new file mode 100644 index 0000000000000..076cc2b91035f --- /dev/null +++ b/sandbox/libs/dataformat-native/rust/liquid-cache/datafusion/src/utils.rs @@ -0,0 +1,335 @@ +use arrow::{ + array::BooleanBufferBuilder, + buffer::{BooleanBuffer, MutableBuffer}, +}; +use parquet::arrow::arrow_reader::RowSelector; + +fn boolean_buffer_and_then_fallback(left: &BooleanBuffer, right: &BooleanBuffer) -> BooleanBuffer { + debug_assert_eq!( + left.count_set_bits(), + right.len(), + "the right selection must have the same number of set bits as the left selection" + ); + + if left.len() == right.len() { + debug_assert_eq!(left.count_set_bits(), left.len()); + return right.clone(); + } + + let mut buffer = MutableBuffer::from_len_zeroed(left.values().len()); + buffer.copy_from_slice(left.values()); + let mut builder = BooleanBufferBuilder::new_from_buffer(buffer, left.len()); + + let mut other_bits = right.iter(); + + for bit_idx in left.set_indices() { + let predicate = other_bits + .next() + .expect("Mismatch in set bits between self and other"); + if !predicate { + builder.set_bit(bit_idx, false); + } + } + + builder.finish() +} + +/// Combines this [`BooleanBuffer`] with another using logical AND on the selected bits. +/// +/// Unlike intersection, the `other` [`BooleanBuffer`] must have exactly as many **set bits** as `self`, +/// i.e., self.count_set_bits() == other.len(). +/// +/// This method will keep only the bits in `self` that are also set in `other` +/// at the positions corresponding to `self`'s set bits. +/// For example: +/// left: NNYYYNNYYNYN +/// right: YNY NY N +/// result: NNYNYNNNYNNN +/// +/// Optimized version of `boolean_buffer_and_then` using BMI2 PDEP instructions. +/// This function performs the same operation but uses bit manipulation instructions +/// for better performance on supported x86_64 CPUs. +pub fn boolean_buffer_and_then(left: &BooleanBuffer, right: &BooleanBuffer) -> BooleanBuffer { + debug_assert_eq!( + left.count_set_bits(), + right.len(), + "the right selection must have the same number of set bits as the left selection" + ); + + if left.len() == right.len() { + debug_assert_eq!(left.count_set_bits(), left.len()); + return right.clone(); + } + + // Fast path for BMI2 support on x86_64 + #[cfg(target_arch = "x86_64")] + { + if is_x86_feature_detected!("bmi2") { + return unsafe { boolean_buffer_and_then_bmi2(left, right) }; + } + } + + boolean_buffer_and_then_fallback(left, right) +} + +#[cfg(target_arch = "x86_64")] +#[inline(always)] +fn load_u64_zero_padded(base_ptr: *const u8, total_bytes: usize, offset: usize) -> u64 { + let remaining = total_bytes.saturating_sub(offset); + if remaining >= 8 { + unsafe { core::ptr::read_unaligned(base_ptr.add(offset) as *const u64) } + } else if remaining > 0 { + let mut tmp = [0u8; 8]; + unsafe { + core::ptr::copy_nonoverlapping(base_ptr.add(offset), tmp.as_mut_ptr(), remaining); + } + u64::from_le_bytes(tmp) + } else { + 0 + } +} + +#[cfg(target_arch = "x86_64")] +#[target_feature(enable = "bmi2")] +unsafe fn boolean_buffer_and_then_bmi2( + left: &BooleanBuffer, + right: &BooleanBuffer, +) -> BooleanBuffer { + use core::arch::x86_64::_pdep_u64; + + debug_assert_eq!(left.count_set_bits(), right.len()); + + let bit_len = left.len(); + let byte_len = bit_len.div_ceil(8); + let left_ptr = left.values().as_ptr(); + let right_bytes = right.len().div_ceil(8); + let right_ptr = right.values().as_ptr(); + + let mut out = MutableBuffer::from_len_zeroed(byte_len); + let out_ptr = out.as_mut_ptr(); + + let full_words = bit_len / 64; + let mut right_bit_idx = 0; // how many bits we have processed from right + + for word_idx in 0..full_words { + let left_word = + unsafe { core::ptr::read_unaligned(left_ptr.add(word_idx * 8) as *const u64) }; + + if left_word == 0 { + continue; + } + + let need = left_word.count_ones(); + + // Absolute byte & bit offset of the first needed bit inside `right`. + let rb_byte = right_bit_idx / 8; + let rb_bit = (right_bit_idx & 7) as u32; + + let need_high = rb_bit != 0; + let safe16 = right_bytes.saturating_sub(16); + let mut r_bits; + if rb_byte <= safe16 { + let low = unsafe { core::ptr::read_unaligned(right_ptr.add(rb_byte) as *const u64) }; + if need_high { + let high = + unsafe { core::ptr::read_unaligned(right_ptr.add(rb_byte + 8) as *const u64) }; + r_bits = (low >> rb_bit) | (high << (64 - rb_bit)); + } else { + r_bits = low >> rb_bit; + } + } else { + let low = load_u64_zero_padded(right_ptr, right_bytes, rb_byte); + r_bits = low >> rb_bit; + if need_high { + let high = load_u64_zero_padded(right_ptr, right_bytes, rb_byte + 8); + r_bits |= high << (64 - rb_bit); + } + } + + // Mask off the high garbage. + r_bits &= 1u64.unbounded_shl(need).wrapping_sub(1); + + // The PDEP instruction: https://www.felixcloutier.com/x86/pdep + // It takes left_word as the mask, and deposit the packed bits into the sparse positions of `left_word`. + let result = _pdep_u64(r_bits, left_word); + + unsafe { + core::ptr::write_unaligned(out_ptr.add(word_idx * 8) as *mut u64, result); + } + + right_bit_idx += need as usize; + } + + // Handle remaining bits that are less than 64 bits + let tail_bits = bit_len & 63; + if tail_bits != 0 { + // Build the mask from the remaining bytes in one bounded copy + let tail_bytes = tail_bits.div_ceil(8); + let base = unsafe { left_ptr.add(full_words * 8) }; + let mut mask: u64; + if tail_bytes == 8 { + mask = unsafe { core::ptr::read_unaligned(base as *const u64) }; + } else { + let mut buf = [0u8; 8]; + unsafe { + core::ptr::copy_nonoverlapping(base, buf.as_mut_ptr(), tail_bytes); + } + mask = u64::from_le_bytes(buf); + } + // Clear any high bits beyond the actual tail length + mask &= (1u64 << tail_bits) - 1; + + if mask != 0 { + let need = mask.count_ones(); + + let rb_byte = right_bit_idx / 8; + let rb_bit = (right_bit_idx & 7) as u32; + + let need_high = rb_bit != 0; + let safe16 = right_bytes.saturating_sub(16); + let mut r_bits; + if rb_byte <= safe16 { + let low = + unsafe { core::ptr::read_unaligned(right_ptr.add(rb_byte) as *const u64) }; + if need_high { + let high = unsafe { + core::ptr::read_unaligned(right_ptr.add(rb_byte + 8) as *const u64) + }; + r_bits = (low >> rb_bit) | (high << (64 - rb_bit)); + } else { + r_bits = low >> rb_bit; + } + } else { + let low = load_u64_zero_padded(right_ptr, right_bytes, rb_byte); + r_bits = low >> rb_bit; + if need_high { + let high = load_u64_zero_padded(right_ptr, right_bytes, rb_byte + 8); + r_bits |= high << (64 - rb_bit); + } + } + + r_bits &= 1u64.unbounded_shl(need).wrapping_sub(1); + + let result = _pdep_u64(r_bits, mask); + + let tail_bytes = tail_bits.div_ceil(8); + let result_bytes = result.to_le_bytes(); + let dst_off = full_words * 8; + unsafe { + let dst = core::slice::from_raw_parts_mut(out_ptr, byte_len); + dst[dst_off..dst_off + tail_bytes].copy_from_slice(&result_bytes[..tail_bytes]); + } + } + } + + BooleanBuffer::new(out.into(), 0, bit_len) +} + +pub(super) fn row_selector_to_boolean_buffer(selection: &[RowSelector]) -> BooleanBuffer { + let mut buffer = BooleanBufferBuilder::new(8192); + for selector in selection.iter() { + if selector.skip { + buffer.append_n(selector.row_count, false); + } else { + buffer.append_n(selector.row_count, true); + } + } + buffer.finish() +} + +#[cfg(all(test, target_arch = "x86_64"))] +mod tests { + use super::*; + + #[test] + fn test_boolean_buffer_and_then_bmi2_large() { + use super::boolean_buffer_and_then_bmi2; + + // Test with larger buffer (more than 64 bits) + let size = 128; + let mut left_builder = BooleanBufferBuilder::new(size); + let mut right_bits = Vec::new(); + + // Create a pattern where every 3rd bit is set in left + for i in 0..size { + let is_set = i.is_multiple_of(3); + left_builder.append(is_set); + if is_set { + // For right buffer, alternate between true/false + right_bits.push(right_bits.len().is_multiple_of(2)); + } + } + let left = left_builder.finish(); + + let mut right_builder = BooleanBufferBuilder::new(right_bits.len()); + for bit in right_bits { + right_builder.append(bit); + } + let right = right_builder.finish(); + + let result_bmi2 = unsafe { boolean_buffer_and_then_bmi2(&left, &right) }; + let result_orig = boolean_buffer_and_then_fallback(&left, &right); + + assert_eq!(result_bmi2.len(), result_orig.len()); + assert_eq!(result_bmi2.len(), size); + + // Verify they produce the same result + for i in 0..size { + assert_eq!( + result_bmi2.value(i), + result_orig.value(i), + "Mismatch at position {i}" + ); + } + } + + #[test] + fn test_boolean_buffer_and_then_bmi2_edge_cases() { + use super::boolean_buffer_and_then_bmi2; + + // Test case: all bits set in left, alternating pattern in right + let mut left_builder = BooleanBufferBuilder::new(16); + for _ in 0..16 { + left_builder.append(true); + } + let left = left_builder.finish(); + + let mut right_builder = BooleanBufferBuilder::new(16); + for i in 0..16 { + right_builder.append(i % 2 == 0); + } + let right = right_builder.finish(); + + let result_bmi2 = unsafe { boolean_buffer_and_then_bmi2(&left, &right) }; + let result_orig = boolean_buffer_and_then_fallback(&left, &right); + + assert_eq!(result_bmi2.len(), result_orig.len()); + for i in 0..16 { + assert_eq!( + result_bmi2.value(i), + result_orig.value(i), + "Mismatch at position {i}" + ); + // Should be true for even indices, false for odd + assert_eq!(result_bmi2.value(i), i.is_multiple_of(2)); + } + + // Test case: no bits set in left + let mut left_empty_builder = BooleanBufferBuilder::new(8); + for _ in 0..8 { + left_empty_builder.append(false); + } + let left_empty = left_empty_builder.finish(); + let right_empty = BooleanBufferBuilder::new(0).finish(); + + let result_bmi2_empty = unsafe { boolean_buffer_and_then_bmi2(&left_empty, &right_empty) }; + let result_orig_empty = boolean_buffer_and_then_fallback(&left_empty, &right_empty); + + assert_eq!(result_bmi2_empty.len(), result_orig_empty.len()); + assert_eq!(result_bmi2_empty.len(), 8); + for i in 0..8 { + assert!(!result_bmi2_empty.value(i)); + assert!(!result_orig_empty.value(i)); + } + } +} diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/Cargo.toml b/sandbox/plugins/analytics-backend-datafusion/rust/Cargo.toml index c12a4cb13c5c2..25df4167b9ef9 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/Cargo.toml +++ b/sandbox/plugins/analytics-backend-datafusion/rust/Cargo.toml @@ -95,6 +95,11 @@ crc32fast = { workspace = true } # rewrite. STANDARD alphabet matches the OpenSearch `binary` field wire contract. base64 = "0.22" +# Liquid Cache for decoded-batch Parquet caching (vendored in-memory subset, +# see sandbox/libs/dataformat-native/rust/liquid-cache/README.md). +opensearch-liquid-cache-core = { workspace = true } +opensearch-liquid-cache-datafusion = { workspace = true } + [dev-dependencies] criterion = { workspace = true } tempfile = { workspace = true } diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/benches/row_id_bench.rs b/sandbox/plugins/analytics-backend-datafusion/rust/benches/row_id_bench.rs index 7726542d92dff..86f940895c9c3 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/benches/row_id_bench.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/benches/row_id_bench.rs @@ -56,6 +56,7 @@ fn setup() -> (RuntimeManager, DataFusionRuntime) { runtime_env, custom_cache_manager: None, dynamic_limit_handle: handle, + liquid_cache_optimizer: None, }; (mgr, df_runtime) } diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/api.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/api.rs index c884828e71791..3ba0ce7f11e19 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/api.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/api.rs @@ -313,6 +313,8 @@ pub struct DataFusionRuntime { pub runtime_env: datafusion::execution::runtime_env::RuntimeEnv, pub custom_cache_manager: Option, pub dynamic_limit_handle: DynamicLimitHandle, + pub liquid_cache_optimizer: + Option>, } /// Per-file metadata passed from Java at shard view creation time. @@ -417,8 +419,23 @@ impl DataFusionRuntime { runtime_env, custom_cache_manager: None, dynamic_limit_handle: handle, + liquid_cache_optimizer: None, } } + + pub fn apply_liquid_cache_optimizers( + &self, + mut builder: SessionStateBuilder, + ) -> SessionStateBuilder { + if let Some(ref optimizer) = self.liquid_cache_optimizer { + builder = builder.with_physical_optimizer_rule(optimizer.clone()); + } + builder + } + + pub fn has_liquid_cache(&self) -> bool { + self.liquid_cache_optimizer.is_some() + } } /// Opaque shard view handle returned to the caller. @@ -511,6 +528,9 @@ pub fn create_global_runtime( cache_manager_ptr: i64, spill_dir: &str, spill_limit: i64, + liquid_cache_enabled: bool, + liquid_cache_size: i64, + liquid_cache_eviction_policy: &str, ) -> Result { if memory_pool_limit < 0 { return Err(DataFusionError::Configuration(format!( @@ -712,10 +732,21 @@ pub fn create_global_runtime( .with_cache_manager(cache_manager_config) .build()?; + let liquid_cache_optimizer = if liquid_cache_enabled { + let liquid_runtime = crate::liquid_cache::LiquidOnlyRuntime::init( + liquid_cache_size as u64, + liquid_cache_eviction_policy, + )?; + Some(liquid_runtime.optimizer()) + } else { + None + }; + let runtime = DataFusionRuntime { runtime_env, custom_cache_manager, dynamic_limit_handle, + liquid_cache_optimizer, }; Ok(Box::into_raw(Box::new(runtime)) as i64) } @@ -1511,6 +1542,23 @@ pub fn cancel_query(context_id: i64) { query_tracker::cancel_query(context_id); } +/// Clears all caching layers: Liquid Cache (in-memory) and DataFusion +/// metadata caches (parquet footers + column statistics). +/// +/// # Safety +/// `runtime_ptr` must be 0 or a valid pointer from `create_global_runtime`. +pub unsafe fn clear_liquid_cache(runtime_ptr: i64) { + crate::liquid_cache::LiquidOnlyRuntime::reset_cache_if_initialized(); + + if runtime_ptr == 0 { + return; + } + let runtime = &*(runtime_ptr as *const DataFusionRuntime); + if let Some(ref cache_manager) = runtime.custom_cache_manager { + cache_manager.clear_all(); + } +} + /// Converts SQL to Substrait plan bytes (test only). /// /// # Safety @@ -2174,7 +2222,8 @@ mod tests { // memory_guard SPILL_ENABLED flag off so per_query_spill_budget returns // Disabled (not Critical) — preventing the 1-partition clamp. let _guard = SPILL_GLOBALS_LOCK.lock().unwrap_or_else(|e| e.into_inner()); - let ptr = create_global_runtime(64 * 1024 * 1024, 0, "", 0).expect("runtime build"); + let ptr = create_global_runtime(64 * 1024 * 1024, 0, "", 0, false, 0, "lru") + .expect("runtime build"); assert!(ptr > 0); let runtime = unsafe { &*(ptr as *const DataFusionRuntime) }; assert!( @@ -2211,8 +2260,16 @@ mod tests { "sentinel must exist before runtime build" ); - let ptr = create_global_runtime(64 * 1024 * 1024, 0, spill_path, 1024 * 1024 * 1024) - .expect("runtime build"); + let ptr = create_global_runtime( + 64 * 1024 * 1024, + 0, + spill_path, + 1024 * 1024 * 1024, + false, + 0, + "lru", + ) + .expect("runtime build"); assert!(ptr > 0); // Phase 1 renames the sentinel file to leaked_from_prior_run.tmp.stale @@ -2275,8 +2332,16 @@ mod tests { assert!(top_file.exists()); assert!(nested_file.exists()); - let ptr = create_global_runtime(64 * 1024 * 1024, 0, spill_path, 1024 * 1024 * 1024) - .expect("runtime build"); + let ptr = create_global_runtime( + 64 * 1024 * 1024, + 0, + spill_path, + 1024 * 1024 * 1024, + false, + 0, + "lru", + ) + .expect("runtime build"); assert!(ptr > 0); // Phase 1: original names gone (renamed to *.stale). @@ -2320,7 +2385,8 @@ mod tests { // accidental fs::remove_dir_all("") would error and break boot. This test // guards against future refactors that hoist the cleanup out of the else-branch. let _guard = SPILL_GLOBALS_LOCK.lock().unwrap_or_else(|e| e.into_inner()); - let ptr = create_global_runtime(64 * 1024 * 1024, 0, "", 0).expect("runtime build"); + let ptr = create_global_runtime(64 * 1024 * 1024, 0, "", 0, false, 0, "lru") + .expect("runtime build"); assert!(ptr > 0); unsafe { close_global_runtime(ptr) }; } @@ -2338,7 +2404,7 @@ mod tests { fs::write(&bad_path, b"not a directory").expect("seed regular file"); let bad_path_str = bad_path.to_str().expect("utf-8 path"); - let err = create_global_runtime(64 * 1024 * 1024, 0, bad_path_str, 0) + let err = create_global_runtime(64 * 1024 * 1024, 0, bad_path_str, 0, false, 0, "lru") .expect_err("create_global_runtime must fail when cleanup fails"); // Operator-facing message must include the offending path + io kind so the @@ -2412,7 +2478,7 @@ mod tests { fs::set_permissions(parent.path(), locked).expect("chmod parent 555"); let spill_str = spill_path.to_str().expect("utf-8 path"); - let result = create_global_runtime(64 * 1024 * 1024, 0, spill_str, 0); + let result = create_global_runtime(64 * 1024 * 1024, 0, spill_str, 0, false, 0, "lru"); // Restore parent perms via RAII so tempdir cleanup runs even on assertion failure. struct RestorePerms<'a> { @@ -2491,7 +2557,8 @@ mod tests { assert!(link.is_symlink(), "precondition: link is a symlink"); let spill_str = spill_path.to_str().expect("utf-8 path"); - let ptr = create_global_runtime(64 * 1024 * 1024, 0, spill_str, 0).expect("runtime build"); + let ptr = create_global_runtime(64 * 1024 * 1024, 0, spill_str, 0, false, 0, "lru") + .expect("runtime build"); assert!(ptr > 0); // Phase 1: symlink is renamed inline (fs::rename does not follow symlinks). @@ -2541,7 +2608,8 @@ mod tests { fs::create_dir(&leaked_b).expect("create leaked b"); fs::write(leaked_b.join("tmp_002.arrow"), b"data b").expect("write b"); - let ptr = create_global_runtime(64 * 1024 * 1024, 0, spill_path, 0).expect("runtime build"); + let ptr = create_global_runtime(64 * 1024 * 1024, 0, spill_path, 0, false, 0, "lru") + .expect("runtime build"); assert!(ptr > 0); // Originals were renamed inline — gone immediately by the original name. @@ -2573,7 +2641,8 @@ mod tests { fs::create_dir(&leftover).expect("create leftover"); fs::write(leftover.join("residue.arrow"), b"prior boot data").expect("write residue"); - let ptr = create_global_runtime(64 * 1024 * 1024, 0, spill_path, 0).expect("runtime build"); + let ptr = create_global_runtime(64 * 1024 * 1024, 0, spill_path, 0, false, 0, "lru") + .expect("runtime build"); assert!(ptr > 0); let cleaned = wait_until(2000, || !leftover.exists()); diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs index 60864b563014b..270f094a22010 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs @@ -156,12 +156,34 @@ pub unsafe extern "C" fn df_create_global_runtime( spill_dir_ptr: *const u8, spill_dir_len: i64, spill_limit: i64, + liquid_cache_enabled: i64, + liquid_cache_size: i64, + liquid_cache_eviction_policy_ptr: *const u8, + liquid_cache_eviction_policy_len: i64, ) -> i64 { crate::memory_guard::set_pool_limit_for_guard(memory_pool_limit); let spill_dir = str_from_raw(spill_dir_ptr, spill_dir_len) .map_err(|e| format!("df_create_global_runtime: {}", e))?; - api::create_global_runtime(memory_pool_limit, cache_manager_ptr, spill_dir, spill_limit) - .map_err(|e| e.to_string()) + let liquid_cache_eviction_policy = str_from_raw( + liquid_cache_eviction_policy_ptr, + liquid_cache_eviction_policy_len, + ) + .map_err(|e| { + format!( + "df_create_global_runtime: liquid_cache_eviction_policy: {}", + e + ) + })?; + api::create_global_runtime( + memory_pool_limit, + cache_manager_ptr, + spill_dir, + spill_limit, + liquid_cache_enabled != 0, + liquid_cache_size, + liquid_cache_eviction_policy, + ) + .map_err(|e| e.to_string()) } #[no_mangle] @@ -169,6 +191,39 @@ pub unsafe extern "C" fn df_close_global_runtime(ptr: i64) { api::close_global_runtime(ptr); } +// ---- Liquid Cache FFM entry points ---- + +#[no_mangle] +pub unsafe extern "C" fn df_clear_liquid_cache(runtime_ptr: i64) { + api::clear_liquid_cache(runtime_ptr); +} + +#[no_mangle] +pub extern "C" fn df_set_liquid_cache_enabled(enabled: i64) { + crate::liquid_cache::LiquidOnlyRuntime::set_enabled_globally(enabled != 0); +} + +#[no_mangle] +pub extern "C" fn df_set_liquid_cache_memory_limit(bytes: i64) { + if bytes >= 0 { + crate::liquid_cache::LiquidOnlyRuntime::set_max_memory_bytes_globally(bytes as usize); + } +} + +#[no_mangle] +pub extern "C" fn df_set_liquid_cache_selectivity_threshold(permille: i64) { + if (0..=1000).contains(&permille) { + crate::liquid_cache::set_lc_selectivity_threshold(permille as f64 / 1000.0); + } +} + +#[no_mangle] +pub extern "C" fn df_set_liquid_cache_max_columns(count: i64) { + if count > 0 { + crate::liquid_cache::set_lc_max_columns(count as usize); + } +} + // ---- Memory pool observability and dynamic limit ---- /// Returns current memory pool usage in bytes. @@ -489,6 +544,7 @@ pub unsafe extern "C" fn df_stream_next(stream_ptr: i64) -> i64 { #[no_mangle] pub unsafe extern "C" fn df_stream_close(stream_ptr: i64) { api::stream_close(stream_ptr); + crate::liquid_cache::LiquidOnlyRuntime::log_stats_if_initialized(); } /// Returns execution metrics as JSON bytes for the given stream. diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/helper.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/helper.rs index c35b6ea8d9f99..66bee6b67c35f 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/helper.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/helper.rs @@ -178,6 +178,15 @@ pub fn build_query_session_context( // combine-partial-final physical optimizer pass. builder = builder.with_physical_optimizer_rules(physical_optimizer_rules_without_combine()); } + // Only apply the physical optimizer (LocalModeLiquidCacheOptimizer) — it + // wraps ParquetSource with LiquidParquetSource for filter pushdown + + // decoded-batch caching. The LineageOptimizer (logical) is excluded because + // it adds planning overhead that causes regression. + if crate::liquid_cache::LiquidOnlyRuntime::is_enabled_globally() { + if let Some(optimizer) = crate::liquid_cache::LiquidOnlyRuntime::optimizer_globally() { + builder = builder.with_physical_optimizer_rule(optimizer); + } + } let state = builder.build(); let ctx = SessionContext::new_with_state(state); 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..200e0f5fc3674 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 @@ -46,6 +46,7 @@ use datafusion_datasource::source::DataSourceExec; use datafusion_datasource::PartitionedFile; use futures::future::BoxFuture; use futures::FutureExt; +use native_bridge_common::log_debug; use object_store::{ObjectStore, ObjectStoreExt}; use prost::bytes::Bytes; @@ -177,16 +178,22 @@ pub struct RowGroupStreamConfig { /// /// Predicate pushdown IS safe here — `RowSelection` is applied during decode, /// so the predicate sees only selected rows and indices stay aligned. +/// +/// `selectivity` is the fraction of rows in this RG that are candidates +/// (0.0 = no rows, 1.0 = all rows). Used to gate LC: highly selective +/// queries (low selectivity) benefit from LC's column-by-column decode +/// with filter pushdown. pub fn create_row_selection_stream( config: &RowGroupStreamConfig, rg_index: usize, selection: RowSelection, push_predicate: bool, + selectivity: f64, ) -> Result<(SendableRecordBatchStream, Arc)> { let num_rgs = config.metadata.num_row_groups(); let mut access_plan = ParquetAccessPlan::new_none(num_rgs); access_plan.set(rg_index, RowGroupAccess::Selection(selection)); - create_stream_with_access_plan(config, access_plan, push_predicate) + create_stream_with_access_plan(config, access_plan, push_predicate, selectivity) } /// Create a stream that reads a single row group with full scan. @@ -218,13 +225,15 @@ pub fn create_full_scan_stream( // rows with gaps?) so the caller's post-decode mask alignment stays // correct. Documented in `pr-reviews/EVALUATOR_HANDOFF.md`. access_plan.set(rg_index, RowGroupAccess::Scan); - create_stream_with_access_plan(config, access_plan, false) + // Full scan = selectivity 1.0 (all rows). + create_stream_with_access_plan(config, access_plan, false, 1.0) } fn create_stream_with_access_plan( config: &RowGroupStreamConfig, access_plan: ParquetAccessPlan, push_predicate: bool, + selectivity: f64, ) -> Result<(SendableRecordBatchStream, Arc)> { let partitioned_file = PartitionedFile::new(config.file_path.clone(), config.file_size) .with_extensions(Arc::new(access_plan)); @@ -250,9 +259,76 @@ fn create_stream_with_access_plan( } } - let mut config_builder = + // Liquid Cache engagement gate: wrap the ParquetSource with + // LiquidParquetSource when ALL projected columns are cacheable + // (numeric/date/timestamp/boolean) and no predicate column is a string. + // The opener decides per-file whether to STREAM or DELEGATE. + let use_lc = { + let lc_globally_enabled = crate::liquid_cache::LiquidOnlyRuntime::is_enabled_globally(); + let max_cols = crate::liquid_cache::lc_max_columns(); + let all_numeric_projection = config.projection.as_ref().map_or(false, |proj| { + !proj.is_empty() + && proj.len() <= max_cols + && proj.iter().all(|&idx| { + config.full_schema.fields().get(idx).map_or(false, |f| { + f.data_type().is_numeric() + || matches!( + f.data_type(), + datafusion::arrow::datatypes::DataType::Date32 + | datafusion::arrow::datatypes::DataType::Date64 + | datafusion::arrow::datatypes::DataType::Timestamp(_, _) + | datafusion::arrow::datatypes::DataType::Boolean + ) + }) + }) + }); + let predicate_has_string = config.predicate.as_ref().map_or(false, |pred| { + let referenced = datafusion::physical_expr::utils::collect_columns(pred); + referenced.iter().any(|col| { + config + .full_schema + .fields() + .get(col.index()) + .map_or(false, |f| { + matches!( + f.data_type(), + datafusion::arrow::datatypes::DataType::Utf8 + | datafusion::arrow::datatypes::DataType::Utf8View + | datafusion::arrow::datatypes::DataType::LargeUtf8 + | datafusion::arrow::datatypes::DataType::Binary + | datafusion::arrow::datatypes::DataType::BinaryView + | datafusion::arrow::datatypes::DataType::LargeBinary + ) + }) + }) + }); + let result = lc_globally_enabled && all_numeric_projection && !predicate_has_string; + log_debug!( + "[parquet_bridge] gate: selectivity={:.3}, all_numeric_proj={}, pred_has_string={}, use_lc={}", + selectivity, + all_numeric_projection, + predicate_has_string, + result, + ); + result + }; + + let mut config_builder = if use_lc { + if let Some(cache_ref) = crate::liquid_cache::LiquidOnlyRuntime::cache_ref_globally() { + let liquid_source = liquid_cache_datafusion::LiquidParquetSource::from_parquet_source( + parquet_source, + cache_ref, + ); + FileScanConfigBuilder::new(config.store_url.clone(), Arc::new(liquid_source)) + .with_file(partitioned_file) + } else { + FileScanConfigBuilder::new(config.store_url.clone(), Arc::new(parquet_source)) + .with_file(partitioned_file) + } + } else { FileScanConfigBuilder::new(config.store_url.clone(), Arc::new(parquet_source)) - .with_file(partitioned_file); + .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/indexed_table/stream.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/stream.rs index 791a3ec9db04b..ce4ea0fc09878 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/stream.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/stream.rs @@ -690,12 +690,14 @@ impl IndexedStream { rg: &RowGroupInfo, selection: RowSelection, push_predicate: bool, + selectivity: f64, ) -> Result<(SendableRecordBatchStream, Arc)> { parquet_bridge::create_row_selection_stream( &self.bridge_config(), rg.index, selection, push_predicate, + selectivity, ) } @@ -1079,6 +1081,15 @@ impl IndexedStream { let min_skip_run = self.pick_min_skip_run(candidates.len() as usize, rg.num_rows as usize); + // Candidate selectivity for this RG — forwarded to + // parquet_bridge so the Liquid Cache gate can log/decide + // per-RG engagement. + let selectivity = if rg.num_rows > 0 { + candidates.len() as f64 / rg.num_rows as f64 + } else { + 1.0 + }; + // Metrics: track which regime we landed in, using the // same counters as before so `EXPLAIN ANALYZE` output // stays comparable. @@ -1158,7 +1169,7 @@ impl IndexedStream { && !alignment_risk && !self.evaluator.forbid_parquet_pushdown(); - match self.create_row_selection_stream(&rg, selection, push) { + match self.create_row_selection_stream(&rg, selection, push, selectivity) { Ok((stream, plan)) => { if let Some(ref timer) = self.metrics.parquet_time { timer.add_duration(t_plan.elapsed()); 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..73faa9945ef30 --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/liquid_cache.rs @@ -0,0 +1,191 @@ +/* SPDX-License-Identifier: Apache-2.0 */ + +//! Liquid Cache integration — in-memory decoded-batch cache for Parquet scans. +//! +//! Backed by the vendored in-memory liquid-cache subset +//! (`sandbox/libs/dataformat-native/rust/liquid-cache`). There is no disk +//! tier: entries are transcoded to the Liquid format under memory pressure +//! and evicted when the budget is exhausted (`TranscodeEvict`). + +use std::sync::{ + atomic::{AtomicBool, AtomicU32, Ordering}, + Arc, OnceLock, +}; + +use datafusion::{common::DataFusionError, physical_optimizer::PhysicalOptimizerRule}; + +use liquid_cache::cache::{CachePolicy, LiquidCache, TranscodeEvict}; +use liquid_cache::cache_policies::{LiquidPolicy, LruPolicy}; +use liquid_cache_datafusion::{LiquidCacheParquet, LiquidCacheParquetRef, LocalModeOptimizer}; +use native_bridge_common::log_debug; + +const EVICTION_POLICY_LRU: &str = "lru"; + +/// Liquid cache batch size — must be a power of two (upstream default). +const LIQUID_CACHE_BATCH_SIZE: usize = 8192; + +static INSTANCE: OnceLock> = OnceLock::new(); + +// Dynamic tuning knobs — updated via cluster settings without restart. +// Selectivity threshold stored as permille (800 = 0.800) to avoid floating-point atomics. +static LC_SELECTIVITY_THRESHOLD_PERMILLE: AtomicU32 = AtomicU32::new(800); +static LC_MAX_COLUMNS: AtomicU32 = AtomicU32::new(10); + +pub fn lc_selectivity_threshold() -> f64 { + LC_SELECTIVITY_THRESHOLD_PERMILLE.load(Ordering::Relaxed) as f64 / 1000.0 +} + +pub fn lc_max_columns() -> usize { + LC_MAX_COLUMNS.load(Ordering::Relaxed) as usize +} + +pub fn set_lc_selectivity_threshold(value: f64) { + let permille = (value * 1000.0) as u32; + LC_SELECTIVITY_THRESHOLD_PERMILLE.store(permille, Ordering::Relaxed); +} + +pub fn set_lc_max_columns(value: usize) { + LC_MAX_COLUMNS.store(value as u32, Ordering::Relaxed); +} + +pub struct LiquidOnlyRuntime { + optimizer: Arc, + cache_ref: LiquidCacheParquetRef, + storage: Arc, + enabled: AtomicBool, +} + +impl LiquidOnlyRuntime { + pub fn init( + max_cache_bytes: u64, + eviction_policy: &str, + ) -> Result<&'static Self, DataFusionError> { + INSTANCE + .get_or_init(|| Self::build(max_cache_bytes, eviction_policy)) + .as_ref() + .map_err(|e| DataFusionError::Execution(e.clone())) + } + + fn build(max_cache_bytes: u64, eviction_policy: &str) -> Result { + let policy: Box = match eviction_policy { + EVICTION_POLICY_LRU => Box::new(LruPolicy::new()), + _ => Box::new(LiquidPolicy::new()), + }; + + let cache_ref: LiquidCacheParquetRef = Arc::new(LiquidCacheParquet::new( + LIQUID_CACHE_BATCH_SIZE, + max_cache_bytes as usize, + policy, + Box::new(TranscodeEvict), + )); + let optimizer = Arc::new(LocalModeOptimizer::new(cache_ref.clone())); + + Ok(Self { + optimizer, + storage: cache_ref.storage().clone(), + cache_ref, + enabled: AtomicBool::new(true), + }) + } + + pub fn optimizer(&self) -> Arc { + self.optimizer.clone() + } + + pub fn cache_ref(&self) -> &LiquidCacheParquetRef { + &self.cache_ref + } + + pub fn cache_ref_globally() -> Option { + Self::get().map(|rt| rt.cache_ref.clone()) + } + + pub fn optimizer_globally() -> Option> { + Self::get().map(|rt| rt.optimizer()) + } + + pub fn is_enabled(&self) -> bool { + self.enabled.load(Ordering::Relaxed) + } + + pub fn set_enabled(&self, enabled: bool) { + self.enabled.store(enabled, Ordering::Relaxed); + } + + pub fn set_max_memory_bytes(&self, bytes: usize) { + self.storage.budget().set_max_memory_bytes(bytes); + } + + pub fn reset_cache(&self) { + // Safety: callers only reset via the explicit REST clear action; + // concurrent queries may observe a cold cache, which is benign for + // a read-through cache (they fall back to Parquet decode). + unsafe { self.cache_ref.reset() }; + let stats = self.storage.stats(); + log_debug!( + "[LiquidCache] cache cleared: entries={}, mem_usage={} bytes", + stats.total_entries, + stats.memory_usage_bytes + ); + } + + pub fn log_stats(&self) { + let s = self.storage.stats(); + log_debug!( + "[LiquidCache] entries={}, mem={}/{}, arrow={}({} B), liquid={}({} B)", + s.total_entries, + s.memory_usage_bytes, + s.max_memory_bytes, + s.memory_arrow_entries, + s.memory_arrow_bytes, + s.memory_liquid_entries, + s.memory_liquid_bytes, + ); + let mem_pct = if s.max_memory_bytes > 0 { + (s.memory_usage_bytes as f64 / s.max_memory_bytes as f64 * 100.0) as u64 + } else { + 0 + }; + log_debug!( + "[LiquidCache] hits={}, misses={}, predicate_evals={}, mem_evictions={}, transcodes={}, mem_pressure={}%", + s.runtime.cache_hit, + s.runtime.cache_miss, + s.runtime.eval_predicate, + s.runtime.memory_evictions, + s.runtime.transcodes, + mem_pct, + ); + } + + fn get() -> Option<&'static Self> { + INSTANCE.get().and_then(|r| r.as_ref().ok()) + } + + pub fn is_enabled_globally() -> bool { + Self::get().map(|rt| rt.is_enabled()).unwrap_or(false) + } + + pub fn set_enabled_globally(enabled: bool) { + if let Some(rt) = Self::get() { + rt.set_enabled(enabled); + } + } + + pub fn set_max_memory_bytes_globally(bytes: usize) { + if let Some(rt) = Self::get() { + rt.set_max_memory_bytes(bytes); + } + } + + pub fn log_stats_if_initialized() { + if let Some(rt) = Self::get() { + rt.log_stats(); + } + } + + pub fn reset_cache_if_initialized() { + if let Some(rt) = Self::get() { + rt.reset_cache(); + } + } +} diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/session_context.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/session_context.rs index 58c37458a4332..af4bd4254a102 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/session_context.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/session_context.rs @@ -243,6 +243,8 @@ pub async unsafe fn create_session_context( .execution .split_file_groups_by_statistics = true; } + // LC session config is applied only when LC is actually engaged (inside + // parquet_bridge.rs), not globally. let mut state_builder = SessionStateBuilder::new() .with_config(config) @@ -263,6 +265,15 @@ pub async unsafe fn create_session_context( runtime.runtime_env.cache_manager.get_file_metadata_cache(), ))); } + // Only the physical optimizer (LocalModeLiquidCacheOptimizer) is applied. + // It wraps ParquetSource with LiquidParquetSource for decoded-batch caching. + // The LineageOptimizer (logical) is excluded — it adds O(nodes*exprs) + // planning overhead that causes 2x regression on string-heavy queries. + if crate::liquid_cache::LiquidOnlyRuntime::is_enabled_globally() { + if let Some(ref optimizer) = runtime.liquid_cache_optimizer { + state_builder = state_builder.with_physical_optimizer_rule(optimizer.clone()); + } + } let state = state_builder.build(); @@ -394,7 +405,7 @@ pub async unsafe fn create_session_context( shard_view.sort_fields.len() ); - error!( + log_debug!( "create_session_context: successfully registered table '{}', table_name_len={}", table_name, table_name.len() diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/tests/local_exec_test.rs b/sandbox/plugins/analytics-backend-datafusion/rust/tests/local_exec_test.rs index 0587b92b05b53..272002e2e3dd6 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/tests/local_exec_test.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/tests/local_exec_test.rs @@ -82,6 +82,10 @@ impl RuntimeGuard { spill_bytes.as_ptr(), spill_bytes.len() as i64, 64 * 1024 * 1024, + 0, // liquid_cache_enabled — off for this test + 0, // liquid_cache_size + b"lru".as_ptr(), + 3, // liquid_cache_eviction_policy ) }; assert!(rc > 0, "df_create_global_runtime returned {}", rc); diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/tests/stringview_gc_test.rs b/sandbox/plugins/analytics-backend-datafusion/rust/tests/stringview_gc_test.rs index 0a23451ec6b65..2448278b9c0a7 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/tests/stringview_gc_test.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/tests/stringview_gc_test.rs @@ -65,6 +65,10 @@ impl RuntimeGuard { spill_bytes.as_ptr(), spill_bytes.len() as i64, 64 * 1024 * 1024, + 0, // liquid_cache_enabled — off for this test + 0, // liquid_cache_size + b"lru".as_ptr(), + 3, // liquid_cache_eviction_policy ) }; assert!(rc > 0, "df_create_global_runtime returned {}", rc); diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionPlugin.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionPlugin.java index e05d6127c9c0d..3ce6e6dd79f14 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionPlugin.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionPlugin.java @@ -16,6 +16,7 @@ import org.opensearch.arrow.allocator.ArrowNativeAllocator; import org.opensearch.arrow.spi.NativeAllocator; import org.opensearch.arrow.spi.PoolGroup; +import org.opensearch.be.datafusion.action.LiquidCacheClearAction; import org.opensearch.be.datafusion.action.stats.DataFusionStatsActionType; import org.opensearch.be.datafusion.action.stats.RestDataFusionStatsAction; import org.opensearch.be.datafusion.action.stats.TransportDataFusionStatsAction; @@ -33,6 +34,7 @@ import org.opensearch.common.settings.Setting; import org.opensearch.common.settings.Settings; import org.opensearch.common.settings.SettingsFilter; +import org.opensearch.common.util.FeatureFlags; import org.opensearch.core.action.ActionResponse; import org.opensearch.core.common.breaker.CircuitBreaker; import org.opensearch.core.common.io.stream.NamedWriteableRegistry; @@ -560,6 +562,9 @@ public Collection createComponents( .spillDirectory(spillDir) .datanodeMultiplier(DatafusionSettings.CONCURRENCY_DATANODE_MULTIPLIER.get(settings)) .clusterSettings(clusterService.getClusterSettings()) + .liquidCacheEnabled(FeatureFlags.isEnabled(FeatureFlags.LIQUID_CACHE_EXPERIMENTAL_SETTING)) + .liquidCacheSize(DatafusionSettings.LIQUID_CACHE_SIZE.get(settings)) + .liquidCacheEvictionPolicy(DatafusionSettings.LIQUID_CACHE_EVICTION_POLICY.get(settings)) .build(); dataFusionService.start(); logger.debug("DataFusion plugin initialized — memory pool {}B, spill limit {}B", memoryPoolLimit, spillMemoryLimit); @@ -651,6 +656,31 @@ public Collection createComponents( DATAFUSION_MEMORY_GUARD_EXECUTION_CRITICAL_THRESHOLD.get(settings) ); + // Wire Liquid Cache dynamic settings only when the experimental flag is enabled. + if (FeatureFlags.isEnabled(FeatureFlags.LIQUID_CACHE_EXPERIMENTAL_SETTING)) { + clusterService.getClusterSettings() + .addSettingsUpdateConsumer(DatafusionSettings.LIQUID_CACHE_ENABLED, enabled -> NativeBridge.setLiquidCacheEnabled(enabled)); + clusterService.getClusterSettings() + .addSettingsUpdateConsumer(DatafusionSettings.LIQUID_CACHE_SIZE, bytes -> NativeBridge.setLiquidCacheMemoryLimit(bytes)); + clusterService.getClusterSettings() + .addSettingsUpdateConsumer( + DatafusionSettings.LIQUID_CACHE_SELECTIVITY_THRESHOLD, + threshold -> NativeBridge.setLiquidCacheSelectivityThreshold((long) (threshold * 1000)) + ); + clusterService.getClusterSettings() + .addSettingsUpdateConsumer( + DatafusionSettings.LIQUID_CACHE_MAX_COLUMNS, + count -> NativeBridge.setLiquidCacheMaxColumns((long) count) + ); + + NativeBridge.setLiquidCacheEnabled(DatafusionSettings.LIQUID_CACHE_ENABLED.get(settings)); + NativeBridge.setLiquidCacheMemoryLimit(DatafusionSettings.LIQUID_CACHE_SIZE.get(settings)); + NativeBridge.setLiquidCacheSelectivityThreshold( + (long) (DatafusionSettings.LIQUID_CACHE_SELECTIVITY_THRESHOLD.get(settings) * 1000) + ); + NativeBridge.setLiquidCacheMaxColumns((long) DatafusionSettings.LIQUID_CACHE_MAX_COLUMNS.get(settings)); + } + this.datafusionSettings = new DatafusionSettings(clusterService); // Expose per-task native-memory usage to search backpressure. @@ -993,6 +1023,13 @@ public List getRestHandlers( if (dataFusionService == null) { return Collections.emptyList(); } + if (FeatureFlags.isEnabled(FeatureFlags.LIQUID_CACHE_EXPERIMENTAL_SETTING)) { + return List.of( + new RestDataFusionStatsAction(), + new org.opensearch.be.datafusion.action.stats.RestClearCacheAction(), + new LiquidCacheClearAction(dataFusionService) + ); + } return List.of(new RestDataFusionStatsAction(), new org.opensearch.be.datafusion.action.stats.RestClearCacheAction()); } diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionService.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionService.java index 31a0838b18ee8..420a60a1f8f80 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionService.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionService.java @@ -49,6 +49,9 @@ public class DataFusionService extends AbstractLifecycleComponent { private final double datanodeMultiplier; private final double coordinatorMultiplier; private final ClusterSettings clusterSettings; + private final boolean liquidCacheEnabled; + private final long liquidCacheSize; + private final String liquidCacheEvictionPolicy; /** Handle to the native DataFusion global runtime (memory pool + cache). */ private volatile NativeRuntimeHandle runtimeHandle; @@ -64,6 +67,9 @@ private DataFusionService(Builder builder) { this.datanodeMultiplier = builder.datanodeMultiplier; this.coordinatorMultiplier = builder.coordinatorMultiplier; this.clusterSettings = builder.clusterSettings; + this.liquidCacheEnabled = builder.liquidCacheEnabled; + this.liquidCacheSize = builder.liquidCacheSize; + this.liquidCacheEvictionPolicy = builder.liquidCacheEvictionPolicy; } /** Creates a new builder. */ @@ -98,7 +104,15 @@ protected void doStart() { } try { - long ptr = NativeBridge.createGlobalRuntime(memoryPoolLimit, cacheManagerPtr, spillDirectory, spillMemoryLimit); + long ptr = NativeBridge.createGlobalRuntime( + memoryPoolLimit, + cacheManagerPtr, + spillDirectory, + spillMemoryLimit, + liquidCacheEnabled, + liquidCacheSize, + liquidCacheEvictionPolicy + ); if (cacheHandle != null) { cacheHandle.markConsumed(); } @@ -296,6 +310,13 @@ public void onFilesDeleted(Collection filePaths) { } } + /** + * Clears the Liquid Cache and DataFusion internal caches. + */ + public void clearLiquidCache() { + NativeBridge.clearLiquidCache(getNativeRuntime().get()); + } + private void releaseRuntime() { NativeRuntimeHandle handle = runtimeHandle; if (handle != null) { @@ -316,6 +337,9 @@ public static class Builder { private double datanodeMultiplier = 1.0; private double coordinatorMultiplier = 1.0; private ClusterSettings clusterSettings; + private boolean liquidCacheEnabled = false; + private long liquidCacheSize = 1L * 1024 * 1024 * 1024; // 1GB default + private String liquidCacheEvictionPolicy = "lru"; private Builder() {} @@ -376,6 +400,33 @@ public Builder clusterSettings(ClusterSettings clusterSettings) { return this; } + /** + * Enables or disables Liquid Cache for byte-level Parquet caching. + * @param enabled whether to enable liquid cache + */ + public Builder liquidCacheEnabled(boolean enabled) { + this.liquidCacheEnabled = enabled; + return this; + } + + /** + * Sets the Liquid Cache size in bytes. + * @param bytes cache size limit + */ + public Builder liquidCacheSize(long bytes) { + this.liquidCacheSize = bytes; + return this; + } + + /** + * Sets the Liquid Cache eviction policy (liquid, lru). + * @param policy the eviction policy + */ + public Builder liquidCacheEvictionPolicy(String policy) { + this.liquidCacheEvictionPolicy = policy; + return this; + } + /** Builds the {@link DataFusionService}. */ public DataFusionService build() { return new DataFusionService(this); diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionSettings.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionSettings.java index 5f3fbb818e865..4828db8a206d3 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionSettings.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionSettings.java @@ -188,6 +188,69 @@ static String derivePoolMinDefault(Settings settings, int percent) { // ── All settings registered by the plugin ── + /** + * Dynamically enables or disables Liquid Cache for new queries. + * When false, optimizers are not injected and queries bypass the cache. + * The cache remains initialized (for fast re-enable) but idle. + */ + public static final Setting LIQUID_CACHE_ENABLED = Setting.boolSetting( + "datafusion.liquid_cache.enabled", + true, + Setting.Property.NodeScope, + Setting.Property.Dynamic + ); + + /** + * Controls the Liquid Cache max memory size in bytes for byte-level Parquet caching. + * Only used when liquid cache is enabled via the experimental feature flag. + */ + public static final Setting LIQUID_CACHE_SIZE = Setting.longSetting( + "datafusion.liquid_cache.size_bytes", + 1L * 1024 * 1024 * 1024, // 1GB default + 0L, + Setting.Property.NodeScope, + Setting.Property.Dynamic + ); + + /** + * Liquid Cache eviction policy. + *
    + *
  • {@code liquid} (default) — Independent FIFO queues per batch type (LiquidPolicy)
    • + *
  • {@code lru} — Standard LRU eviction across all entry types
    • + *
+ */ + public static final Setting LIQUID_CACHE_EVICTION_POLICY = Setting.simpleString( + "datafusion.liquid_cache.eviction_policy", + "lru", + Setting.Property.NodeScope, + Setting.Property.Final + ); + + /** + * Selectivity threshold for LC STREAM vs DELEGATE (0.0 to 1.0). + * Files with selectivity below this threshold are delegated to plain parquet. + */ + public static final Setting LIQUID_CACHE_SELECTIVITY_THRESHOLD = Setting.floatSetting( + "datafusion.liquid_cache.selectivity_threshold", + 0.8f, + 0.0f, + 1.0f, + Setting.Property.NodeScope, + Setting.Property.Dynamic + ); + + /** + * Maximum number of output columns for LC engagement. + * Queries projecting more columns than this are skipped by the optimizer. + */ + public static final Setting LIQUID_CACHE_MAX_COLUMNS = Setting.intSetting( + "datafusion.liquid_cache.max_columns", + 10, + 1, + Setting.Property.NodeScope, + Setting.Property.Dynamic + ); + public static final List> ALL_SETTINGS = List.of( // Runtime settings — memory pool, spill, reduce input mode, and budget tuning @@ -229,7 +292,13 @@ static String derivePoolMinDefault(Settings settings, int percent) { INDEXED_PUSHDOWN_FILTERS, INDEXED_MIN_SKIP_RUN_DEFAULT, INDEXED_MIN_SKIP_RUN_SELECTIVITY_THRESHOLD, - INDEXED_FORCE_STRATEGY + INDEXED_FORCE_STRATEGY, + + LIQUID_CACHE_ENABLED, + LIQUID_CACHE_SIZE, + LIQUID_CACHE_EVICTION_POLICY, + LIQUID_CACHE_SELECTIVITY_THRESHOLD, + LIQUID_CACHE_MAX_COLUMNS ); // ── Snapshot management ── diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/action/LiquidCacheClearAction.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/action/LiquidCacheClearAction.java new file mode 100644 index 0000000000000..e372a5aa84ed1 --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/action/LiquidCacheClearAction.java @@ -0,0 +1,64 @@ +/* + * 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.be.datafusion.action; + +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.opensearch.be.datafusion.DataFusionService; +import org.opensearch.core.rest.RestStatus; +import org.opensearch.core.xcontent.XContentBuilder; +import org.opensearch.rest.BaseRestHandler; +import org.opensearch.rest.BytesRestResponse; +import org.opensearch.rest.RestRequest; +import org.opensearch.transport.client.node.NodeClient; + +import java.util.List; + +/** + * REST handler for {@code POST _plugins/analytics_backend_datafusion/liquid_cache/clear}. + *

+ * Clears all Liquid Cache entries (in-memory) and DataFusion internal caches. + */ +public class LiquidCacheClearAction extends BaseRestHandler { + + private static final Logger logger = LogManager.getLogger(LiquidCacheClearAction.class); + private final DataFusionService dataFusionService; + + public LiquidCacheClearAction(DataFusionService dataFusionService) { + this.dataFusionService = dataFusionService; + } + + @Override + public String getName() { + return "liquid_cache_clear_action"; + } + + @Override + public List routes() { + return List.of(new Route(RestRequest.Method.POST, "_plugins/analytics_backend_datafusion/liquid_cache/clear")); + } + + @Override + protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) { + return channel -> { + try { + logger.info("Clearing Liquid Cache via REST endpoint"); + dataFusionService.clearLiquidCache(); + logger.info("Liquid Cache cleared successfully"); + XContentBuilder builder = channel.newBuilder(); + builder.startObject(); + builder.field("acknowledged", true); + builder.endObject(); + channel.sendResponse(new BytesRestResponse(RestStatus.OK, builder)); + } catch (Exception e) { + channel.sendResponse(new BytesRestResponse(channel, e)); + } + }; + } +} diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/nativelib/NativeBridge.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/nativelib/NativeBridge.java index bf24dcb0330f4..91cfe2fb36444 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/nativelib/NativeBridge.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/nativelib/NativeBridge.java @@ -143,6 +143,11 @@ private static RuntimeException rethrowConverted(RuntimeException e) { private static final MethodHandle SET_SCOPED_PAGE_INDEX_ENABLED; private static final MethodHandle CANCEL_QUERY; private static final MethodHandle SET_CANCEL_STATS_THRESHOLD_MS; + private static final MethodHandle CLEAR_LIQUID_CACHE; + private static final MethodHandle SET_LIQUID_CACHE_ENABLED; + private static final MethodHandle SET_LIQUID_CACHE_MEMORY_LIMIT; + private static final MethodHandle SET_LIQUID_CACHE_SELECTIVITY_THRESHOLD; + private static final MethodHandle SET_LIQUID_CACHE_MAX_COLUMNS; private static final MethodHandle STATS; private static final MethodHandle QUERY_REGISTRY_TOP_N_BY_CURRENT; private static final MethodHandle DF_NATIVE_NODE_STATS; @@ -174,6 +179,10 @@ private static RuntimeException rethrowConverted(RuntimeException e) { ValueLayout.JAVA_LONG, ValueLayout.ADDRESS, ValueLayout.JAVA_LONG, + ValueLayout.JAVA_LONG, + ValueLayout.JAVA_LONG, + ValueLayout.JAVA_LONG, + ValueLayout.ADDRESS, ValueLayout.JAVA_LONG ) ); @@ -183,6 +192,31 @@ private static RuntimeException rethrowConverted(RuntimeException e) { FunctionDescriptor.ofVoid(ValueLayout.JAVA_LONG) ); + CLEAR_LIQUID_CACHE = linker.downcallHandle( + lib.find("df_clear_liquid_cache").orElseThrow(), + FunctionDescriptor.ofVoid(ValueLayout.JAVA_LONG) + ); + + SET_LIQUID_CACHE_ENABLED = linker.downcallHandle( + lib.find("df_set_liquid_cache_enabled").orElseThrow(), + FunctionDescriptor.ofVoid(ValueLayout.JAVA_LONG) + ); + + SET_LIQUID_CACHE_MEMORY_LIMIT = linker.downcallHandle( + lib.find("df_set_liquid_cache_memory_limit").orElseThrow(), + FunctionDescriptor.ofVoid(ValueLayout.JAVA_LONG) + ); + + SET_LIQUID_CACHE_SELECTIVITY_THRESHOLD = linker.downcallHandle( + lib.find("df_set_liquid_cache_selectivity_threshold").orElseThrow(), + FunctionDescriptor.ofVoid(ValueLayout.JAVA_LONG) + ); + + SET_LIQUID_CACHE_MAX_COLUMNS = linker.downcallHandle( + lib.find("df_set_liquid_cache_max_columns").orElseThrow(), + FunctionDescriptor.ofVoid(ValueLayout.JAVA_LONG) + ); + GET_MEMORY_POOL_USAGE = linker.downcallHandle( lib.find("df_get_memory_pool_usage").orElseThrow(), FunctionDescriptor.of(ValueLayout.JAVA_LONG, ValueLayout.JAVA_LONG) @@ -768,10 +802,30 @@ public static void updateConcurrencyGate(String gateName, int newMaxPermits) { * This pointer is not a MemorySegment — it's a Rust heap address that lives * until {@link #closeGlobalRuntime} is called. */ - public static long createGlobalRuntime(long memoryLimit, long cacheManagerPtr, String spillDir, long spillLimit) { + public static long createGlobalRuntime( + long memoryLimit, + long cacheManagerPtr, + String spillDir, + long spillLimit, + boolean liquidCacheEnabled, + long liquidCacheSize, + String liquidCacheEvictionPolicy + ) { try (var call = new NativeCall()) { var dir = call.str(spillDir); - return call.invoke(CREATE_GLOBAL_RUNTIME, memoryLimit, cacheManagerPtr, dir.segment(), dir.len(), spillLimit); + var eviction = call.str(liquidCacheEvictionPolicy); + return call.invoke( + CREATE_GLOBAL_RUNTIME, + memoryLimit, + cacheManagerPtr, + dir.segment(), + dir.len(), + spillLimit, + liquidCacheEnabled ? 1L : 0L, + liquidCacheSize, + eviction.segment(), + eviction.len() + ); } } @@ -780,6 +834,31 @@ public static void closeGlobalRuntime(long ptr) { NativeCall.invokeVoid(CLOSE_GLOBAL_RUNTIME, ptr); } + /** Clears all Liquid Cache entries and DataFusion internal caches. */ + public static void clearLiquidCache(long runtimePtr) { + NativeCall.invokeVoid(CLEAR_LIQUID_CACHE, runtimePtr); + } + + /** Dynamically enable or disable Liquid Cache for new queries. */ + public static void setLiquidCacheEnabled(boolean enabled) { + NativeCall.invokeVoid(SET_LIQUID_CACHE_ENABLED, enabled ? 1L : 0L); + } + + /** Dynamically update the Liquid Cache memory limit in bytes. */ + public static void setLiquidCacheMemoryLimit(long bytes) { + NativeCall.invokeVoid(SET_LIQUID_CACHE_MEMORY_LIMIT, bytes); + } + + /** Dynamically update the LC selectivity threshold (permille: 800 = 0.8). */ + public static void setLiquidCacheSelectivityThreshold(long permille) { + NativeCall.invokeVoid(SET_LIQUID_CACHE_SELECTIVITY_THRESHOLD, permille); + } + + /** Dynamically update the max columns for LC engagement. */ + public static void setLiquidCacheMaxColumns(long count) { + NativeCall.invokeVoid(SET_LIQUID_CACHE_MAX_COLUMNS, count); + } + // ---- Memory pool observability and dynamic limit ---- /** Returns current memory pool usage in bytes. */ diff --git a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DataFusionNativeBridgeTests.java b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DataFusionNativeBridgeTests.java index 7f93b4d9b9a81..9c333025862b4 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DataFusionNativeBridgeTests.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DataFusionNativeBridgeTests.java @@ -51,7 +51,10 @@ public void testRuntimeLifecycle() { 64 * 1024 * 1024, // 64MB 0L, spillDir.toString(), - 32 * 1024 * 1024 // 32MB spill + 32 * 1024 * 1024, // 32MB spill + false, + 0L, + "lru" ); assertTrue("Runtime pointer should be non-zero", runtimePtr != 0); @@ -62,7 +65,7 @@ public void testRuntimeLifecycle() { public void testReaderLifecycle() throws Exception { NativeBridge.initTokioRuntimeManager(2); Path spillDir = createTempDir("datafusion-spill"); - long runtimePtr = NativeBridge.createGlobalRuntime(64 * 1024 * 1024, 0L, spillDir.toString(), 32 * 1024 * 1024); + long runtimePtr = NativeBridge.createGlobalRuntime(64 * 1024 * 1024, 0L, spillDir.toString(), 32 * 1024 * 1024, false, 0L, "lru"); // Copy test parquet to a temp dir Path dataDir = createTempDir("datafusion-data"); @@ -88,7 +91,7 @@ public void testReaderLifecycle() throws Exception { public void testSessionContextCreationAndTableRegistration() throws Exception { NativeBridge.initTokioRuntimeManager(2); Path spillDir = createTempDir("datafusion-spill"); - long runtimePtr = NativeBridge.createGlobalRuntime(64 * 1024 * 1024, 0L, spillDir.toString(), 32 * 1024 * 1024); + long runtimePtr = NativeBridge.createGlobalRuntime(64 * 1024 * 1024, 0L, spillDir.toString(), 32 * 1024 * 1024, false, 0L, "lru"); NativeRuntimeHandle runtimeHandle = new NativeRuntimeHandle(runtimePtr); Path dataDir = createTempDir("datafusion-data"); @@ -166,7 +169,7 @@ public void onFailure(Exception exception) { public void testReaderWithRealNativeStoreHandle() throws Exception { NativeBridge.initTokioRuntimeManager(2); Path spillDir = createTempDir("datafusion-spill"); - long runtimePtr = NativeBridge.createGlobalRuntime(64 * 1024 * 1024, 0L, spillDir.toString(), 32 * 1024 * 1024); + long runtimePtr = NativeBridge.createGlobalRuntime(64 * 1024 * 1024, 0L, spillDir.toString(), 32 * 1024 * 1024, false, 0L, "lru"); // Copy test parquet to a temp dir Path dataDir = createTempDir("datafusion-data"); diff --git a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DataFusionPluginSettingsTests.java b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DataFusionPluginSettingsTests.java index be7142952b5de..24148b5dc9f91 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DataFusionPluginSettingsTests.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DataFusionPluginSettingsTests.java @@ -111,7 +111,7 @@ public void testGetSettingsReturnsAllIndexedSettings() { public void testGetSettingsReturnsTotalExpectedCount() { try (DataFusionPlugin plugin = new DataFusionPlugin()) { List> settings = plugin.getSettings(); - assertEquals(31, settings.size()); + assertEquals(36, settings.size()); } catch (Exception e) { throw new AssertionError(e); } diff --git a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DataFusionQueryExecutionTests.java b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DataFusionQueryExecutionTests.java index 307d0e5e15421..bb06c83361b77 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DataFusionQueryExecutionTests.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DataFusionQueryExecutionTests.java @@ -53,7 +53,7 @@ public void setUp() throws Exception { NativeBridge.initTokioRuntimeManager(2); Path spillDir = createTempDir("datafusion-spill"); runtimeHandle = new NativeRuntimeHandle( - NativeBridge.createGlobalRuntime(128 * 1024 * 1024, 0L, spillDir.toString(), 64 * 1024 * 1024) + NativeBridge.createGlobalRuntime(128 * 1024 * 1024, 0L, spillDir.toString(), 64 * 1024 * 1024, false, 0L, "lru") ); // Create a real TieredObjectStore (local-only) and wrap in NativeStoreHandle diff --git a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DataFusionServiceTests.java b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DataFusionServiceTests.java index 1cf221e4e9f57..a0a5738fca301 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DataFusionServiceTests.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DataFusionServiceTests.java @@ -82,7 +82,7 @@ public void testGetNativeRuntimeBeforeStartThrows() { public void testNativeRuntimeHandleCloseIsIdempotent() { ensureTokioInit(); Path spillDir = createTempDir("spill"); - long ptr = NativeBridge.createGlobalRuntime(64 * 1024 * 1024, 0L, spillDir.toString(), 32 * 1024 * 1024); + long ptr = NativeBridge.createGlobalRuntime(64 * 1024 * 1024, 0L, spillDir.toString(), 32 * 1024 * 1024, false, 0L, "lru"); NativeRuntimeHandle handle = new NativeRuntimeHandle(ptr); assertTrue(handle.isOpen()); @@ -96,7 +96,7 @@ public void testNativeRuntimeHandleCloseIsIdempotent() { public void testNativeRuntimeHandleGetAfterCloseThrows() { ensureTokioInit(); Path spillDir = createTempDir("spill"); - long ptr = NativeBridge.createGlobalRuntime(64 * 1024 * 1024, 0L, spillDir.toString(), 32 * 1024 * 1024); + long ptr = NativeBridge.createGlobalRuntime(64 * 1024 * 1024, 0L, spillDir.toString(), 32 * 1024 * 1024, false, 0L, "lru"); NativeRuntimeHandle handle = new NativeRuntimeHandle(ptr); handle.close(); expectThrows(IllegalStateException.class, handle::get); @@ -202,7 +202,15 @@ public void testRuntimeWithCacheManagerPointer() { NativeBridge.createCache(cachePtr, "STATISTICS", 100 * 1024 * 1024, "LRU"); Path spillDir = createTempDir("spill"); - long runtimePtr = NativeBridge.createGlobalRuntime(64 * 1024 * 1024, cachePtr, spillDir.toString(), 32 * 1024 * 1024); + long runtimePtr = NativeBridge.createGlobalRuntime( + 64 * 1024 * 1024, + cachePtr, + spillDir.toString(), + 32 * 1024 * 1024, + false, + 0L, + "lru" + ); assertTrue(runtimePtr != 0); NativeBridge.closeGlobalRuntime(runtimePtr); @@ -217,7 +225,15 @@ public void testCacheManagerHandleConsumedAfterRuntimeCreation() { assertTrue(org.opensearch.analytics.backend.jni.NativeHandle.isLivePointer(ptrBefore)); Path spillDir = createTempDir("spill"); - long runtimePtr = NativeBridge.createGlobalRuntime(64 * 1024 * 1024, handle.getPointer(), spillDir.toString(), 32 * 1024 * 1024); + long runtimePtr = NativeBridge.createGlobalRuntime( + 64 * 1024 * 1024, + handle.getPointer(), + spillDir.toString(), + 32 * 1024 * 1024, + false, + 0L, + "lru" + ); handle.markConsumed(); assertFalse(org.opensearch.analytics.backend.jni.NativeHandle.isLivePointer(ptrBefore)); diff --git a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionMemtableReduceSinkTests.java b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionMemtableReduceSinkTests.java index fa41ad2d42fe7..02dccfa8a06ee 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionMemtableReduceSinkTests.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionMemtableReduceSinkTests.java @@ -59,7 +59,7 @@ public void testInputIdConstantMatchesDesign() { public void testFeedDrainsSumToDownstream() throws Exception { NativeBridge.initTokioRuntimeManager(2); Path spillDir = createTempDir("datafusion-spill"); - long runtimePtr = NativeBridge.createGlobalRuntime(64 * 1024 * 1024, 0L, spillDir.toString(), 32 * 1024 * 1024); + long runtimePtr = NativeBridge.createGlobalRuntime(64 * 1024 * 1024, 0L, spillDir.toString(), 32 * 1024 * 1024, false, 0L, "lru"); assertTrue("runtime ptr non-zero", runtimePtr != 0); NativeRuntimeHandle runtimeHandle = new NativeRuntimeHandle(runtimePtr); diff --git a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionReduceSinkTests.java b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionReduceSinkTests.java index b2c8800558d20..4a310b96155aa 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionReduceSinkTests.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionReduceSinkTests.java @@ -107,7 +107,7 @@ public void testReceiverDroppedSentinelIsPositiveAndDistinctFromSuccess() { public void testFeedDrainsSumToDownstream() throws Exception { NativeBridge.initTokioRuntimeManager(2); Path spillDir = createTempDir("datafusion-spill"); - long runtimePtr = NativeBridge.createGlobalRuntime(64 * 1024 * 1024, 0L, spillDir.toString(), 32 * 1024 * 1024); + long runtimePtr = NativeBridge.createGlobalRuntime(64 * 1024 * 1024, 0L, spillDir.toString(), 32 * 1024 * 1024, false, 0L, "lru"); assertTrue("runtime ptr non-zero", runtimePtr != 0); NativeRuntimeHandle runtimeHandle = new NativeRuntimeHandle(runtimePtr); @@ -169,7 +169,7 @@ public void testFeedDrainsSumToDownstream() throws Exception { public void testReduceWithLimitProducesLimitedOutput() throws Exception { NativeBridge.initTokioRuntimeManager(2); Path spillDir = createTempDir("datafusion-spill"); - long runtimePtr = NativeBridge.createGlobalRuntime(64 * 1024 * 1024, 0L, spillDir.toString(), 32 * 1024 * 1024); + long runtimePtr = NativeBridge.createGlobalRuntime(64 * 1024 * 1024, 0L, spillDir.toString(), 32 * 1024 * 1024, false, 0L, "lru"); NativeRuntimeHandle runtimeHandle = new NativeRuntimeHandle(runtimePtr); try (RootAllocator alloc = new RootAllocator(Long.MAX_VALUE)) { @@ -214,7 +214,7 @@ public void testReduceWithLimitProducesLimitedOutput() throws Exception { public void testDrainTaskKeepsUpWithProducer() throws Exception { NativeBridge.initTokioRuntimeManager(2); Path spillDir = createTempDir("datafusion-spill"); - long runtimePtr = NativeBridge.createGlobalRuntime(64 * 1024 * 1024, 0L, spillDir.toString(), 32 * 1024 * 1024); + long runtimePtr = NativeBridge.createGlobalRuntime(64 * 1024 * 1024, 0L, spillDir.toString(), 32 * 1024 * 1024, false, 0L, "lru"); NativeRuntimeHandle runtimeHandle = new NativeRuntimeHandle(runtimePtr); try (RootAllocator alloc = new RootAllocator(Long.MAX_VALUE)) { @@ -268,7 +268,7 @@ public void testDrainTaskKeepsUpWithProducer() throws Exception { public void testReduceProducesOutputIncrementallyForPipelinedPlan() throws Exception { NativeBridge.initTokioRuntimeManager(2); Path spillDir = createTempDir("datafusion-spill"); - long runtimePtr = NativeBridge.createGlobalRuntime(64 * 1024 * 1024, 0L, spillDir.toString(), 32 * 1024 * 1024); + long runtimePtr = NativeBridge.createGlobalRuntime(64 * 1024 * 1024, 0L, spillDir.toString(), 32 * 1024 * 1024, false, 0L, "lru"); NativeRuntimeHandle runtimeHandle = new NativeRuntimeHandle(runtimePtr); try (RootAllocator alloc = new RootAllocator(Long.MAX_VALUE)) { @@ -344,7 +344,7 @@ public void testReduceProducesOutputIncrementallyForPipelinedPlan() throws Excep public void testCloseWhileFeederParkedOnFullChannelDoesNotDeadlock() throws Exception { NativeBridge.initTokioRuntimeManager(2); Path spillDir = createTempDir("datafusion-spill"); - long runtimePtr = NativeBridge.createGlobalRuntime(64 * 1024 * 1024, 0L, spillDir.toString(), 32 * 1024 * 1024); + long runtimePtr = NativeBridge.createGlobalRuntime(64 * 1024 * 1024, 0L, spillDir.toString(), 32 * 1024 * 1024, false, 0L, "lru"); NativeRuntimeHandle runtimeHandle = new NativeRuntimeHandle(runtimePtr); try (RootAllocator alloc = new RootAllocator(Long.MAX_VALUE)) { @@ -421,7 +421,7 @@ public void testCloseWhileFeederParkedOnFullChannelDoesNotDeadlock() throws Exce public void testCancelBeforeFirstBatchUnwindsDrain() throws Exception { NativeBridge.initTokioRuntimeManager(2); Path spillDir = createTempDir("datafusion-spill"); - long runtimePtr = NativeBridge.createGlobalRuntime(64 * 1024 * 1024, 0L, spillDir.toString(), 32 * 1024 * 1024); + long runtimePtr = NativeBridge.createGlobalRuntime(64 * 1024 * 1024, 0L, spillDir.toString(), 32 * 1024 * 1024, false, 0L, "lru"); NativeRuntimeHandle runtimeHandle = new NativeRuntimeHandle(runtimePtr); try (RootAllocator alloc = new RootAllocator(Long.MAX_VALUE)) { @@ -473,7 +473,7 @@ public void testCancelBeforeFirstBatchUnwindsDrain() throws Exception { public void testCancelAfterFirstBatchUnwindsDrain() throws Exception { NativeBridge.initTokioRuntimeManager(2); Path spillDir = createTempDir("datafusion-spill"); - long runtimePtr = NativeBridge.createGlobalRuntime(64 * 1024 * 1024, 0L, spillDir.toString(), 32 * 1024 * 1024); + long runtimePtr = NativeBridge.createGlobalRuntime(64 * 1024 * 1024, 0L, spillDir.toString(), 32 * 1024 * 1024, false, 0L, "lru"); NativeRuntimeHandle runtimeHandle = new NativeRuntimeHandle(runtimePtr); try (RootAllocator alloc = new RootAllocator(Long.MAX_VALUE)) { @@ -519,7 +519,7 @@ public void testCancelAfterFirstBatchUnwindsDrain() throws Exception { public void testDoubleCloseIsIdempotent() throws Exception { NativeBridge.initTokioRuntimeManager(2); Path spillDir = createTempDir("datafusion-spill"); - long runtimePtr = NativeBridge.createGlobalRuntime(64 * 1024 * 1024, 0L, spillDir.toString(), 32 * 1024 * 1024); + long runtimePtr = NativeBridge.createGlobalRuntime(64 * 1024 * 1024, 0L, spillDir.toString(), 32 * 1024 * 1024, false, 0L, "lru"); NativeRuntimeHandle runtimeHandle = new NativeRuntimeHandle(runtimePtr); try (RootAllocator alloc = new RootAllocator(Long.MAX_VALUE)) { diff --git a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionResultStreamTests.java b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionResultStreamTests.java index 28db10793bffc..6dd6010a48e5d 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionResultStreamTests.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionResultStreamTests.java @@ -47,7 +47,7 @@ public void setUp() throws Exception { super.setUp(); NativeBridge.initTokioRuntimeManager(2); Path spillDir = createTempDir("spill"); - long ptr = NativeBridge.createGlobalRuntime(128 * 1024 * 1024, 0L, spillDir.toString(), 64 * 1024 * 1024); + long ptr = NativeBridge.createGlobalRuntime(128 * 1024 * 1024, 0L, spillDir.toString(), 64 * 1024 * 1024, false, 0L, "lru"); runtimeHandle = new NativeRuntimeHandle(ptr); testRootAllocator = new RootAllocator(Long.MAX_VALUE); @@ -237,7 +237,7 @@ public void testCloseAfterNativeStreamNextFailure() throws Exception { // Create a valid stream, close the runtime handle to force streamNext failure, // then verify the stream still closes cleanly Path spillDir2 = createTempDir("spill2"); - long ptr2 = NativeBridge.createGlobalRuntime(128 * 1024 * 1024, 0L, spillDir2.toString(), 64 * 1024 * 1024); + long ptr2 = NativeBridge.createGlobalRuntime(128 * 1024 * 1024, 0L, spillDir2.toString(), 64 * 1024 * 1024, false, 0L, "lru"); NativeRuntimeHandle tempRuntime = new NativeRuntimeHandle(ptr2); byte[] substrait = NativeBridge.sqlToSubstrait( diff --git a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionSearchExecEngineTests.java b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionSearchExecEngineTests.java index 48b380ea44056..6892251783fca 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionSearchExecEngineTests.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionSearchExecEngineTests.java @@ -44,7 +44,7 @@ public void setUp() throws Exception { super.setUp(); NativeBridge.initTokioRuntimeManager(2); Path spillDir = createTempDir("datafusion-spill"); - long ptr = NativeBridge.createGlobalRuntime(128 * 1024 * 1024, 0L, spillDir.toString(), 64 * 1024 * 1024); + long ptr = NativeBridge.createGlobalRuntime(128 * 1024 * 1024, 0L, spillDir.toString(), 64 * 1024 * 1024, false, 0L, "lru"); runtimeHandle = new NativeRuntimeHandle(ptr); // Create a real TieredObjectStore (local-only) and wrap in NativeStoreHandle diff --git a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionSettingsTests.java b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionSettingsTests.java index e7cde67cd052d..fc8181c59109c 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionSettingsTests.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionSettingsTests.java @@ -55,7 +55,7 @@ public void testMinSkipRunSelectivityThresholdSettingDefinition() { } public void testAllSettingsContainsAllExpectedSettings() { - assertEquals(31, DatafusionSettings.ALL_SETTINGS.size()); + assertEquals(36, DatafusionSettings.ALL_SETTINGS.size()); assertTrue(DatafusionSettings.ALL_SETTINGS.contains(DataFusionPlugin.DATAFUSION_REDUCE_TARGET_PARTITIONS)); assertTrue(DatafusionSettings.ALL_SETTINGS.contains(DataFusionPlugin.DATAFUSION_MEMORY_GUARD_SPILL_EXEMPT_CAP)); assertTrue(DatafusionSettings.ALL_SETTINGS.contains(DataFusionPlugin.DATAFUSION_SPILL_DIRECTORY)); diff --git a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/NativeBridgeLocalSessionTests.java b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/NativeBridgeLocalSessionTests.java index c80e2434c5908..fec13a83b8ca5 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/NativeBridgeLocalSessionTests.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/NativeBridgeLocalSessionTests.java @@ -53,7 +53,7 @@ public class NativeBridgeLocalSessionTests extends OpenSearchTestCase { private NativeRuntimeHandle createRuntime() { NativeBridge.initTokioRuntimeManager(2); Path spillDir = createTempDir("datafusion-spill"); - long runtimePtr = NativeBridge.createGlobalRuntime(64 * 1024 * 1024, 0L, spillDir.toString(), 32 * 1024 * 1024); + long runtimePtr = NativeBridge.createGlobalRuntime(64 * 1024 * 1024, 0L, spillDir.toString(), 32 * 1024 * 1024, false, 0L, "lru"); assertTrue("runtime ptr non-zero", runtimePtr != 0); return new NativeRuntimeHandle(runtimePtr); } diff --git a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/NativeBridgePreparedPlanTests.java b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/NativeBridgePreparedPlanTests.java index b605fb3bc6527..f542ed012cd72 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/NativeBridgePreparedPlanTests.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/NativeBridgePreparedPlanTests.java @@ -42,7 +42,7 @@ public void testExecuteLocalPreparedPlanRejectsNullPointer() { public void testPrepareFinalPlanWithInvalidBytesThrowsDecodeError() { NativeBridge.initTokioRuntimeManager(2); Path spillDir = createTempDir("datafusion-spill"); - long runtimePtr = NativeBridge.createGlobalRuntime(64 * 1024 * 1024, 0L, spillDir.toString(), 32 * 1024 * 1024); + long runtimePtr = NativeBridge.createGlobalRuntime(64 * 1024 * 1024, 0L, spillDir.toString(), 32 * 1024 * 1024, false, 0L, "lru"); NativeRuntimeHandle runtimeHandle = new NativeRuntimeHandle(runtimePtr); DatafusionLocalSession session = new DatafusionLocalSession(runtimeHandle.get()); try { diff --git a/sandbox/qa/analytics-engine-rest/build.gradle b/sandbox/qa/analytics-engine-rest/build.gradle index cb822cdf1b0b6..bed91a032f488 100644 --- a/sandbox/qa/analytics-engine-rest/build.gradle +++ b/sandbox/qa/analytics-engine-rest/build.gradle @@ -105,6 +105,9 @@ def configureAnalyticsCluster = { cluster -> // wildcards, or comma-lists). cluster.setting 'cluster.pluggable.dataformat', 'composite' + // Enable liquid cache feature flag (in-memory decoded-batch cache) + cluster.systemProperty 'opensearch.experimental.feature.liquid_cache.enabled', 'true' + // analytics-engine requires the streaming transport — fragment dispatch is streaming-only. cluster.systemProperty 'opensearch.experimental.feature.transport.stream.enabled', 'true' } diff --git a/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/LiquidCacheBenchmarkIT.java b/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/LiquidCacheBenchmarkIT.java new file mode 100644 index 0000000000000..7509c44829f57 --- /dev/null +++ b/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/LiquidCacheBenchmarkIT.java @@ -0,0 +1,228 @@ +/* + * 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.analytics.qa; + +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.opensearch.client.Request; +import org.opensearch.client.Response; + +import java.util.ArrayList; +import java.util.List; +import java.util.Random; +import java.util.concurrent.TimeUnit; + +/** + * Benchmark comparing query latency with and without Liquid Cache on a + * scaled-down ClickBench-like dataset (~100K rows). + * + * Run with: + * ./gradlew :sandbox:qa:analytics-engine-rest:integTest \ + * --tests "org.opensearch.analytics.qa.LiquidCacheBenchmarkIT" \ + * -Dsandbox.enabled=true + * + * Results are printed to stdout — compare "LC OFF" vs "LC ON" latencies. + */ +@SuppressWarnings("unchecked") +public class LiquidCacheBenchmarkIT extends AnalyticsRestTestCase { + + private static final Logger logger = LogManager.getLogger(LiquidCacheBenchmarkIT.class); + private static final String INDEX_NAME = "liquid_cache_bench"; + private static final int NUM_DOCS = 100_000; + private static final int BULK_BATCH_SIZE = 5000; + private static final int WARMUP_ITERATIONS = 3; + private static final int MEASURE_ITERATIONS = 10; + + private static final String[] QUERIES = { + // Q1: COUNT with equality filter (keyword) + "source=" + INDEX_NAME + " | where status = '200' | stats count() as cnt", + // Q2: Aggregation with range filter (numeric) + "source=" + INDEX_NAME + " | where response_time > 500 | stats avg(response_time) as avg_rt", + // Q3: GROUP BY low-cardinality keyword + "source=" + INDEX_NAME + " | stats count() as cnt by status", + // Q4: GROUP BY + range filter (time-range pattern) + "source=" + INDEX_NAME + " | where event_time > 1700000000000 | stats count() as cnt by region", + // Q5: Multi-column aggregation + "source=" + INDEX_NAME + " | where status = '500' | stats avg(response_time) as avg_rt, max(response_time) as max_rt", + // Q6: High selectivity filter + count + "source=" + INDEX_NAME + " | where user_id = 'user_42' | stats count() as cnt", }; + + public void testLiquidCacheBenchmark() throws Exception { + setupIndex(); + + logger.info("=== Liquid Cache Benchmark: {} docs, {} queries ===", NUM_DOCS, QUERIES.length); + + // Phase 1: LC OFF — baseline + updateSetting("datafusion.liquid_cache.enabled", "false"); + long[] baselineLatencies = runBenchmark("LC OFF"); + + // Phase 2: LC ON — first run (cold cache, populates cache) + updateSetting("datafusion.liquid_cache.enabled", "true"); + long[] coldCacheLatencies = runBenchmark("LC ON (cold)"); + + // Phase 3: LC ON — second run (warm cache, should be faster) + long[] warmCacheLatencies = runBenchmark("LC ON (warm)"); + + // Report + logger.info("=== RESULTS (median of {} iterations, ms) ===", MEASURE_ITERATIONS); + logger.info(String.format("%-60s %10s %10s %10s %10s", "Query", "Baseline", "Cold", "Warm", "Speedup")); + for (int i = 0; i < QUERIES.length; i++) { + String q = QUERIES[i].length() > 58 ? QUERIES[i].substring(0, 58) + ".." : QUERIES[i]; + double speedup = baselineLatencies[i] > 0 ? (double) baselineLatencies[i] / warmCacheLatencies[i] : 0; + logger.info( + String.format( + "%-60s %8dms %8dms %8dms %8.1fx", + q, + baselineLatencies[i], + coldCacheLatencies[i], + warmCacheLatencies[i], + speedup + ) + ); + } + } + + private long[] runBenchmark(String label) throws Exception { + long[] medians = new long[QUERIES.length]; + for (int q = 0; q < QUERIES.length; q++) { + // Warmup + for (int i = 0; i < WARMUP_ITERATIONS; i++) { + executePplRaw(QUERIES[q]); + } + // Measure + List latencies = new ArrayList<>(); + for (int i = 0; i < MEASURE_ITERATIONS; i++) { + long start = System.nanoTime(); + executePplRaw(QUERIES[q]); + long elapsed = TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - start); + latencies.add(elapsed); + } + latencies.sort(Long::compareTo); + medians[q] = latencies.get(latencies.size() / 2); + logger.info( + "[{}] Q{}: median={}ms, min={}ms, max={}ms", + label, + q + 1, + medians[q], + latencies.get(0), + latencies.get(latencies.size() - 1) + ); + } + return medians; + } + + private void setupIndex() throws Exception { + deleteIndexIfExists(INDEX_NAME); + createIndex(); + ingestData(); + flush(); + logger.info("Index created with {} docs", NUM_DOCS); + } + + private void createIndex() throws Exception { + Request request = new Request("PUT", "/" + INDEX_NAME); + request.setJsonEntity( + "{" + + "\"settings\": {" + + " \"number_of_shards\": 1," + + " \"number_of_replicas\": 0," + + " \"refresh_interval\": \"-1\"," + + " \"index.pluggable.dataformat.enabled\": true," + + " \"index.pluggable.dataformat\": \"composite\"," + + " \"index.composite.primary_data_format\": \"parquet\"" + + "}," + + "\"mappings\": {" + + " \"properties\": {" + + " \"event_time\": {\"type\": \"long\"}," + + " \"status\": {\"type\": \"keyword\"}," + + " \"region\": {\"type\": \"keyword\"}," + + " \"user_id\": {\"type\": \"keyword\"}," + + " \"response_time\": {\"type\": \"integer\"}," + + " \"bytes_sent\": {\"type\": \"long\"}," + + " \"url\": {\"type\": \"keyword\"}" + + " }" + + "}" + + "}" + ); + assertEquals(200, client().performRequest(request).getStatusLine().getStatusCode()); + } + + private void ingestData() throws Exception { + Random rng = new Random(42); // deterministic seed + String[] statuses = { "200", "301", "404", "500" }; + String[] regions = { "us-east-1", "us-west-2", "eu-west-1", "ap-southeast-1" }; + + int ingested = 0; + while (ingested < NUM_DOCS) { + int batchSize = Math.min(BULK_BATCH_SIZE, NUM_DOCS - ingested); + StringBuilder bulk = new StringBuilder(); + for (int i = 0; i < batchSize; i++) { + bulk.append("{\"index\":{}}\n"); + bulk.append("{") + .append("\"event_time\":") + .append(1700000000000L + rng.nextInt(86400000)) + .append(",") + .append("\"status\":\"") + .append(statuses[rng.nextInt(statuses.length)]) + .append("\",") + .append("\"region\":\"") + .append(regions[rng.nextInt(regions.length)]) + .append("\",") + .append("\"user_id\":\"user_") + .append(rng.nextInt(10000)) + .append("\",") + .append("\"response_time\":") + .append(50 + rng.nextInt(2000)) + .append(",") + .append("\"bytes_sent\":") + .append(100 + rng.nextInt(50000)) + .append(",") + .append("\"url\":\"/api/v") + .append(rng.nextInt(3) + 1) + .append("/resource/") + .append(rng.nextInt(1000)) + .append("\"") + .append("}\n"); + } + Request request = new Request("POST", "/" + INDEX_NAME + "/_bulk"); + request.setJsonEntity(bulk.toString()); + Response response = client().performRequest(request); + assertEquals(200, response.getStatusLine().getStatusCode()); + ingested += batchSize; + if (ingested % 20000 == 0) { + logger.info("Ingested {}/{} docs", ingested, NUM_DOCS); + } + } + } + + private void flush() throws Exception { + client().performRequest(new Request("POST", "/" + INDEX_NAME + "/_refresh")); + client().performRequest(new Request("POST", "/" + INDEX_NAME + "/_flush?force=true")); + } + + private void executePplRaw(String query) throws Exception { + Request request = new Request("POST", "/_plugins/_ppl"); + request.setJsonEntity("{\"query\":\"" + query + "\"}"); + client().performRequest(request); + } + + private void updateSetting(String key, String value) throws Exception { + Request request = new Request("PUT", "/_cluster/settings"); + request.setJsonEntity("{\"transient\":{\"" + key + "\":\"" + value + "\"}}"); + client().performRequest(request); + } + + private void deleteIndexIfExists(String name) throws Exception { + try { + client().performRequest(new Request("DELETE", "/" + name)); + } catch (Exception e) { + // ignore — index may not exist + } + } +} diff --git a/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/LiquidCacheIT.java b/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/LiquidCacheIT.java new file mode 100644 index 0000000000000..c6d41d90089a1 --- /dev/null +++ b/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/LiquidCacheIT.java @@ -0,0 +1,190 @@ +/* + * 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.analytics.qa; + +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.opensearch.client.Request; +import org.opensearch.client.Response; + +import java.util.List; +import java.util.Map; + +/** + * Integration tests for Liquid Cache functionality within the analytics engine. + *

+ * Validates: + *

    + *
  • Composite parquet index creation and data ingestion
    • + *
  • PPL query execution through DataFusion with numeric predicates
    • + *
  • Dynamic enable/disable of liquid cache via cluster settings
    • + *
  • Dynamic resize of memory budget at runtime
    • + *
  • Query correctness across all cache states
    • + *
+ *

+ * Requires feature flags: + * {@code opensearch.experimental.feature.pluggable.dataformat.enabled=true}, + * {@code opensearch.experimental.feature.liquid_cache.enabled=true} + */ +@SuppressWarnings("unchecked") +public class LiquidCacheIT extends AnalyticsRestTestCase { + + private static final Logger logger = LogManager.getLogger(LiquidCacheIT.class); + private static final String INDEX_NAME = "liquid_cache_integ"; + private static final String PPL_ENDPOINT = "/_analytics/ppl"; + private static final long EXPECTED_SUM_AGE_GT_25 = 300000L; + + /** + * End-to-end test: index lifecycle, query execution, dynamic toggle, and budget resize. + */ + public void testLiquidCacheEndToEnd() throws Exception { + setupIndex(); + + verifyQueryReturnsExpectedResult(); + verifyDynamicDisableAndReenable(); + verifyDynamicBudgetResize(); + } + + private void verifyQueryReturnsExpectedResult() throws Exception { + logger.info("Executing PPL query with numeric predicate (age > 25)"); + long latency = executePplAndAssert( + "source=" + INDEX_NAME + " | where age > 25 | stats sum(salary) as total", + EXPECTED_SUM_AGE_GT_25 + ); + logger.info("Query latency: {}ms", latency); + } + + private void verifyDynamicDisableAndReenable() throws Exception { + updateSetting("datafusion.liquid_cache.enabled", "false"); + logger.info("Liquid cache disabled via cluster settings"); + + long disabledLatency = executePplAndAssert( + "source=" + INDEX_NAME + " | where age > 25 | stats sum(salary) as total", + EXPECTED_SUM_AGE_GT_25 + ); + logger.info("Query latency (LC disabled): {}ms", disabledLatency); + + updateSetting("datafusion.liquid_cache.enabled", "true"); + logger.info("Liquid cache re-enabled via cluster settings"); + + long reenabledLatency = executePplAndAssert( + "source=" + INDEX_NAME + " | where age > 25 | stats sum(salary) as total", + EXPECTED_SUM_AGE_GT_25 + ); + logger.info("Query latency (LC re-enabled): {}ms", reenabledLatency); + } + + private void verifyDynamicBudgetResize() throws Exception { + long newMemory = 512L * 1024 * 1024; + + updateSetting("datafusion.liquid_cache.size_bytes", String.valueOf(newMemory)); + + Response response = client().performRequest(new Request("GET", "/_cluster/settings?flat_settings=true&include_defaults=false")); + Map settings = entityAsMap(response); + Map transient_ = (Map) settings.get("transient"); + + assertEquals("Memory budget not updated", String.valueOf(newMemory), transient_.get("datafusion.liquid_cache.size_bytes")); + logger.info("Budget resize verified: memory={}MB", newMemory / (1024 * 1024)); + } + + // ── Helpers ────────────────────────────────────────────────────────────── + + private void setupIndex() throws Exception { + deleteIndexIfExists(INDEX_NAME); + createCompositeParquetIndex(); + bulkIngestTestData(); + flushAndForceMerge(); + verifyParquetFormat(); + } + + private void createCompositeParquetIndex() throws Exception { + Request request = new Request("PUT", "/" + INDEX_NAME); + request.setJsonEntity( + "{" + + "\"settings\": {" + + " \"number_of_shards\": 1," + + " \"number_of_replicas\": 0," + + " \"index.pluggable.dataformat.enabled\": true," + + " \"index.pluggable.dataformat\": \"composite\"," + + " \"index.composite.primary_data_format\": \"parquet\"" + + "}," + + "\"mappings\": {" + + " \"properties\": {" + + " \"name\": {\"type\": \"keyword\"}," + + " \"age\": {\"type\": \"integer\"}," + + " \"salary\": {\"type\": \"long\"}" + + " }" + + "}" + + "}" + ); + assertEquals(200, client().performRequest(request).getStatusLine().getStatusCode()); + } + + private void bulkIngestTestData() throws Exception { + Request request = new Request("POST", "/" + INDEX_NAME + "/_bulk"); + request.addParameter("refresh", "true"); + request.setJsonEntity( + "{\"index\":{}}\n{\"name\":\"Alice\",\"age\":30,\"salary\":75000}\n" + + "{\"index\":{}}\n{\"name\":\"Bob\",\"age\":25,\"salary\":60000}\n" + + "{\"index\":{}}\n{\"name\":\"Charlie\",\"age\":35,\"salary\":90000}\n" + + "{\"index\":{}}\n{\"name\":\"Diana\",\"age\":28,\"salary\":70000}\n" + + "{\"index\":{}}\n{\"name\":\"Eve\",\"age\":35,\"salary\":65000}\n" + ); + assertEquals(200, client().performRequest(request).getStatusLine().getStatusCode()); + } + + private void flushAndForceMerge() throws Exception { + client().performRequest(new Request("POST", "/" + INDEX_NAME + "/_flush?force=true")); + Request merge = new Request("POST", "/" + INDEX_NAME + "/_forcemerge"); + merge.addParameter("max_num_segments", "1"); + client().performRequest(merge); + Thread.sleep(5000); + } + + private void verifyParquetFormat() throws Exception { + Response response = client().performRequest(new Request("GET", "/" + INDEX_NAME + "/_settings?flat_settings=true")); + Map settings = entityAsMap(response); + Map indexSettings = (Map) ((Map) settings.get(INDEX_NAME)).get("settings"); + assertEquals("parquet", indexSettings.get("index.composite.primary_data_format")); + } + + private long executePplAndAssert(String pplQuery, long expectedValue) throws Exception { + long start = System.currentTimeMillis(); + Request request = new Request("POST", PPL_ENDPOINT); + request.setJsonEntity("{\"query\": \"" + pplQuery + "\"}"); + Response response = client().performRequest(request); + long elapsed = System.currentTimeMillis() - start; + + assertEquals(200, response.getStatusLine().getStatusCode()); + Map body = entityAsMap(response); + logger.info("PPL response: {}", body); + assertNotNull("Response body should not be null", body); + + @SuppressWarnings("unchecked") + List> rows = (List>) body.get("rows"); + assertNotNull("Response should contain rows", rows); + assertFalse("Rows should not be empty", rows.isEmpty()); + Number actual = (Number) rows.get(0).get(0); + assertEquals("Query result mismatch", expectedValue, actual.longValue()); + + return elapsed; + } + + private void deleteIndexIfExists(String index) { + try { + client().performRequest(new Request("DELETE", "/" + index)); + } catch (Exception ignored) {} + } + + private void updateSetting(String key, String value) throws Exception { + Request request = new Request("PUT", "/_cluster/settings"); + request.setJsonEntity("{\"transient\":{\"" + key + "\":\"" + value + "\"}}"); + assertEquals(200, client().performRequest(request).getStatusLine().getStatusCode()); + } +} diff --git a/server/src/main/java/org/opensearch/common/util/FeatureFlags.java b/server/src/main/java/org/opensearch/common/util/FeatureFlags.java index 2a53dc05f6559..a37814f2f848d 100644 --- a/server/src/main/java/org/opensearch/common/util/FeatureFlags.java +++ b/server/src/main/java/org/opensearch/common/util/FeatureFlags.java @@ -127,6 +127,16 @@ public class FeatureFlags { public static final String STREAM_TRANSPORT = FEATURE_FLAG_PREFIX + "transport.stream.enabled"; public static final Setting STREAM_TRANSPORT_SETTING = Setting.boolSetting(STREAM_TRANSPORT, false, Property.NodeScope); + /** + * Gates the functionality of Liquid Cache for byte-level Parquet caching. + */ + public static final String LIQUID_CACHE_EXPERIMENTAL_FLAG = FEATURE_FLAG_PREFIX + "liquid_cache.enabled"; + public static final Setting LIQUID_CACHE_EXPERIMENTAL_SETTING = Setting.boolSetting( + LIQUID_CACHE_EXPERIMENTAL_FLAG, + false, + Property.NodeScope + ); + /** * Underlying implementation for feature flags. * All settable feature flags are tracked here in FeatureFlagsImpl.featureFlags. @@ -153,6 +163,7 @@ static class FeatureFlagsImpl { put(STREAM_TRANSPORT_SETTING, STREAM_TRANSPORT_SETTING.getDefault(Settings.EMPTY)); put(CONTEXT_AWARE_MIGRATION_EXPERIMENTAL_SETTING, CONTEXT_AWARE_MIGRATION_EXPERIMENTAL_SETTING.getDefault(Settings.EMPTY)); put(PLUGGABLE_DATAFORMAT_EXPERIMENTAL_SETTING, PLUGGABLE_DATAFORMAT_EXPERIMENTAL_SETTING.getDefault(Settings.EMPTY)); + put(LIQUID_CACHE_EXPERIMENTAL_SETTING, LIQUID_CACHE_EXPERIMENTAL_SETTING.getDefault(Settings.EMPTY)); } };