diff --git a/src/core/src/cache/budget.rs b/src/core/src/cache/budget.rs index 097cb417..2f54385c 100644 --- a/src/core/src/cache/budget.rs +++ b/src/core/src/cache/budget.rs @@ -3,14 +3,22 @@ use crate::sync::atomic::{AtomicUsize, Ordering}; #[derive(Debug)] pub struct BudgetAccounting { max_memory_bytes: usize, + max_disk_bytes: usize, + watermark_ratio: f64, used_memory_bytes: AtomicUsize, used_disk_bytes: AtomicUsize, } impl BudgetAccounting { - pub(super) fn new(max_memory_bytes: usize) -> Self { + pub(super) fn new( + max_memory_bytes: usize, + max_disk_bytes: usize, + watermark_ratio: f64, + ) -> Self { Self { max_memory_bytes, + max_disk_bytes, + watermark_ratio, used_memory_bytes: AtomicUsize::new(0), used_disk_bytes: AtomicUsize::new(0), } @@ -69,6 +77,20 @@ impl BudgetAccounting { pub fn add_used_disk_bytes(&self, bytes: usize) { self.used_disk_bytes.fetch_add(bytes, Ordering::Relaxed); } + + pub fn sub_used_disk_bytes(&self, bytes: usize) { + self.used_disk_bytes.fetch_sub(bytes, Ordering::Relaxed); + } + + pub fn disk_budget_exceeded(&self) -> bool { + self.max_disk_bytes != usize::MAX + && self.used_disk_bytes.load(Ordering::Relaxed) + > (self.max_disk_bytes as f64 * self.watermark_ratio) as usize + } + + pub fn max_disk_bytes(&self) -> usize { + self.max_disk_bytes + } } #[cfg(test)] @@ -78,7 +100,7 @@ mod tests { #[test] fn test_memory_reservation_and_accounting() { - let config = BudgetAccounting::new(1000); + let config = BudgetAccounting::new(1000, usize::MAX, 0.9); assert_eq!(config.memory_usage_bytes(), 0); @@ -95,6 +117,34 @@ mod tests { assert_eq!(config.memory_usage_bytes(), 0); } + #[test] + fn test_disk_budget_tracking() { + let budget = BudgetAccounting::new(1000, 500, 0.9); + + assert_eq!(budget.disk_usage_bytes(), 0); + assert!(!budget.disk_budget_exceeded()); + + budget.add_used_disk_bytes(400); + assert_eq!(budget.disk_usage_bytes(), 400); + assert!(!budget.disk_budget_exceeded()); + + budget.add_used_disk_bytes(200); + assert_eq!(budget.disk_usage_bytes(), 600); + assert!(budget.disk_budget_exceeded()); + + budget.sub_used_disk_bytes(300); + assert_eq!(budget.disk_usage_bytes(), 300); + assert!(!budget.disk_budget_exceeded()); + } + + #[test] + fn test_disk_budget_unlimited() { + let budget = BudgetAccounting::new(1000, usize::MAX, 0.9); + + budget.add_used_disk_bytes(usize::MAX / 2); + assert!(!budget.disk_budget_exceeded()); + } + #[test] fn test_concurrent_memory_operations() { test_concurrent_memory_budget(); @@ -111,7 +161,7 @@ mod tests { let max_memory = 10000; let operations_per_thread = 100; - let budget = Arc::new(BudgetAccounting::new(max_memory)); + let budget = Arc::new(BudgetAccounting::new(max_memory, usize::MAX, 0.9)); let barrier = Arc::new(Barrier::new(num_threads)); let mut thread_handles = vec![]; diff --git a/src/core/src/cache/builders.rs b/src/core/src/cache/builders.rs index 0555a87d..877aae74 100644 --- a/src/core/src/cache/builders.rs +++ b/src/core/src/cache/builders.rs @@ -32,6 +32,8 @@ use crate::sync::Arc; pub struct LiquidCacheBuilder { batch_size: usize, max_memory_bytes: usize, + max_disk_bytes: usize, + disk_watermark: f64, cache_policy: Box, hydration_policy: Box, squeeze_policy: Box, @@ -52,6 +54,8 @@ impl LiquidCacheBuilder { Self { batch_size: 8192, max_memory_bytes: 1024 * 1024 * 1024, + max_disk_bytes: usize::MAX, + disk_watermark: 0.9, cache_policy: Box::new(LiquidPolicy::new()), hydration_policy: Box::new(super::AlwaysHydrate::new()), squeeze_policy: Box::new(TranscodeSqueezeEvict), @@ -75,6 +79,20 @@ impl LiquidCacheBuilder { self } + /// Set the max disk bytes for the cache. + /// Default is unlimited — the cache will use available disk space without a cap. + pub fn with_max_disk_bytes(mut self, max_disk_bytes: usize) -> Self { + self.max_disk_bytes = max_disk_bytes; + self + } + + /// Set the disk watermark ratio (0.0–1.0). Default is 0.9. + /// Eviction triggers when disk usage exceeds this fraction of max_disk_bytes. + pub fn with_disk_watermark(mut self, ratio: f64) -> Self { + self.disk_watermark = ratio; + self + } + /// Set the cache policy for the cache. /// Default is [LiquidPolicy]. pub fn with_cache_policy(mut self, policy: Box) -> Self { @@ -137,6 +155,8 @@ impl LiquidCacheBuilder { Arc::new(LiquidCache::new( self.batch_size, self.max_memory_bytes, + self.max_disk_bytes, + self.disk_watermark, self.squeeze_policy, self.cache_policy, self.hydration_policy, diff --git a/src/core/src/cache/core.rs b/src/core/src/cache/core.rs index 8963a7c7..b41044fd 100644 --- a/src/core/src/cache/core.rs +++ b/src/core/src/cache/core.rs @@ -113,6 +113,7 @@ impl LiquidCache { memory_usage_bytes, disk_usage_bytes, max_memory_bytes: self.config.max_memory_bytes(), + max_disk_bytes: self.config.max_disk_bytes(), runtime, } } @@ -338,6 +339,8 @@ impl LiquidCache { pub(crate) fn new( batch_size: usize, max_memory_bytes: usize, + max_disk_bytes: usize, + disk_watermark: f64, squeeze_policy: Box, cache_policy: Box, hydration_policy: Box, @@ -345,10 +348,14 @@ impl LiquidCache { store: t4::Store, squeeze_victims_concurrently: bool, ) -> Self { - let config = CacheConfig::new(batch_size, max_memory_bytes); + let config = CacheConfig::new(batch_size, max_memory_bytes, max_disk_bytes, disk_watermark); Self { index: ArtIndex::new(), - budget: BudgetAccounting::new(config.max_memory_bytes()), + budget: BudgetAccounting::new( + config.max_memory_bytes(), + config.max_disk_bytes(), + config.disk_watermark(), + ), config, cache_policy, hydration_policy, @@ -693,6 +700,10 @@ impl LiquidCache { } async fn write_batch_to_disk(&self, entry_id: EntryID, batch: &CacheEntry, bytes: Bytes) { + if self.budget.disk_budget_exceeded() { + self.evict_disk_entries().await; + } + self.trace(InternalEvent::IoWrite { entry: entry_id, kind: CachedBatchType::from(batch), @@ -706,6 +717,29 @@ impl LiquidCache { self.budget.add_used_disk_bytes(len); } + async fn evict_disk_entries(&self) { + while self.budget.disk_budget_exceeded() { + let victims = self.cache_policy.find_disk_victims(8); + if victims.is_empty() { + break; + } + for victim in victims { + if !self.budget.disk_budget_exceeded() { + break; + } + if let Some(_removed) = self.index.remove(&victim) { + let key = crate::cache::io_context::entry_id_to_key(&victim); + let _ = self.store.remove(&key).await; + self.budget.sub_used_disk_bytes( + self.budget + .disk_usage_bytes() + .min(self.config.batch_size() * 8), + ); + } + } + } + } + async fn read_disk_arrow_array(&self, entry_id: &EntryID) -> ArrayRef { let bytes = self .store diff --git a/src/core/src/cache/index.rs b/src/core/src/cache/index.rs index b2a84edd..55b9ee37 100644 --- a/src/core/src/cache/index.rs +++ b/src/core/src/cache/index.rs @@ -80,6 +80,15 @@ impl ArtIndex { pub(crate) fn entry_count(&self) -> usize { self.entry_count.load(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 + } } #[cfg(test)] @@ -134,4 +143,26 @@ mod tests { let entry_id: EntryID = EntryID::from(1); assert!(!store.is_cached(&entry_id)); } + + #[test] + fn test_remove() { + let store = ArtIndex::new(); + let entry_id = EntryID::from(1); + let entry_id2 = EntryID::from(2); + let array = create_test_array(50); + + store.insert(&entry_id, array.clone()); + store.insert(&entry_id2, array.clone()); + assert_eq!(store.entry_count(), 2); + + let removed = store.remove(&entry_id); + assert!(removed.is_some()); + assert!(!store.is_cached(&entry_id)); + assert!(store.is_cached(&entry_id2)); + assert_eq!(store.entry_count(), 1); + + let removed = store.remove(&EntryID::from(99)); + assert!(removed.is_none()); + assert_eq!(store.entry_count(), 1); + } } diff --git a/src/core/src/cache/observer/stats.rs b/src/core/src/cache/observer/stats.rs index cfb74e8d..8b68a11b 100644 --- a/src/core/src/cache/observer/stats.rs +++ b/src/core/src/cache/observer/stats.rs @@ -146,6 +146,8 @@ pub struct CacheStats { pub disk_usage_bytes: usize, /// Maximum memory size. pub max_memory_bytes: usize, + /// Maximum disk size. + pub max_disk_bytes: usize, /// Runtime counters snapshot. pub runtime: RuntimeStatsSnapshot, } diff --git a/src/core/src/cache/policies/cache/mod.rs b/src/core/src/cache/policies/cache/mod.rs index a253d1e3..6cc6c92d 100644 --- a/src/core/src/cache/policies/cache/mod.rs +++ b/src/core/src/cache/policies/cache/mod.rs @@ -24,6 +24,12 @@ pub trait CachePolicy: std::fmt::Debug + Send + Sync { /// Give cnt amount of entries to evict when cache is full. fn find_victim(&self, cnt: usize) -> Vec; + /// Give cnt amount of disk entries to evict when disk is full. + /// Default returns empty — policies that track disk entries can override. + fn find_disk_victims(&self, _cnt: usize) -> Vec { + vec![] + } + /// Notify the cache policy that an entry was inserted. fn notify_insert(&self, _entry_id: &EntryID, _batch_type: CachedBatchType) {} diff --git a/src/core/src/cache/policies/cache/three_queue.rs b/src/core/src/cache/policies/cache/three_queue.rs index 57ff6a4b..9b6cf8e0 100644 --- a/src/core/src/cache/policies/cache/three_queue.rs +++ b/src/core/src/cache/policies/cache/three_queue.rs @@ -185,6 +185,24 @@ impl CachePolicy for LiquidPolicy { victims } + fn find_disk_victims(&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 { + match inner.pop_front(QueueKind::Disk) { + Some(entry) => victims.push(entry), + None => break, + } + } + + victims + } + fn notify_access(&self, _entry_id: &EntryID, _batch_type: CachedBatchType) {} } @@ -289,4 +307,49 @@ mod tests { let victims = policy.find_victim(2); assert_eq!(victims, vec![entry_id]); } + + #[test] + fn test_find_disk_victims_returns_disk_entries_fifo() { + let policy = LiquidPolicy::new(); + + let d1 = entry(10); + let d2 = entry(11); + let d3 = entry(12); + + policy.notify_insert(&d1, CachedBatchType::DiskLiquid); + policy.notify_insert(&d2, CachedBatchType::DiskArrow); + policy.notify_insert(&d3, CachedBatchType::DiskLiquid); + + let victims = policy.find_disk_victims(2); + assert_eq!(victims, vec![d1, d2]); + + let victims = policy.find_disk_victims(5); + assert_eq!(victims, vec![d3]); + + assert!(policy.find_disk_victims(1).is_empty()); + } + + #[test] + fn test_find_disk_victims_zero_returns_empty() { + let policy = LiquidPolicy::new(); + policy.notify_insert(&entry(1), CachedBatchType::DiskLiquid); + assert!(policy.find_disk_victims(0).is_empty()); + } + + #[test] + fn test_find_disk_victims_does_not_affect_memory_eviction() { + let policy = LiquidPolicy::new(); + + let mem = entry(1); + let disk = entry(2); + + policy.notify_insert(&mem, CachedBatchType::MemoryArrow); + policy.notify_insert(&disk, CachedBatchType::DiskLiquid); + + let disk_victims = policy.find_disk_victims(5); + assert_eq!(disk_victims, vec![disk]); + + let mem_victims = policy.find_victim(5); + assert_eq!(mem_victims, vec![mem]); + } } diff --git a/src/core/src/cache/utils.rs b/src/core/src/cache/utils.rs index f22f398d..a54643ab 100644 --- a/src/core/src/cache/utils.rs +++ b/src/core/src/cache/utils.rs @@ -9,13 +9,23 @@ use bytes::Bytes; pub struct CacheConfig { batch_size: usize, max_memory_bytes: usize, + max_disk_bytes: usize, + + disk_watermark: f64, } impl CacheConfig { - pub(super) fn new(batch_size: usize, max_memory_bytes: usize) -> Self { + pub(super) fn new( + batch_size: usize, + max_memory_bytes: usize, + max_disk_bytes: usize, + disk_watermark: f64, + ) -> Self { Self { batch_size, max_memory_bytes, + max_disk_bytes, + disk_watermark, } } @@ -26,6 +36,14 @@ impl CacheConfig { pub fn max_memory_bytes(&self) -> usize { self.max_memory_bytes } + + pub fn max_disk_bytes(&self) -> usize { + self.max_disk_bytes + } + + pub fn disk_watermark(&self) -> f64 { + self.disk_watermark + } } // Helper methods diff --git a/src/datafusion-local/src/lib.rs b/src/datafusion-local/src/lib.rs index 9f23d654..9056f27a 100644 --- a/src/datafusion-local/src/lib.rs +++ b/src/datafusion-local/src/lib.rs @@ -59,6 +59,8 @@ pub struct LiquidCacheLocalBuilder { batch_size: usize, /// Maximum memory size in bytes max_memory_bytes: usize, + /// Maximum disk size in bytes + max_disk_bytes: usize, /// Directory for disk cache cache_dir: PathBuf, /// Cache policy @@ -75,6 +77,7 @@ impl Default for LiquidCacheLocalBuilder { Self { batch_size: 8192, max_memory_bytes: 1024 * 1024 * 1024, // 1GB + max_disk_bytes: usize::MAX, cache_dir: std::env::temp_dir(), cache_policy: Box::new(LiquidPolicy::new()), squeeze_policy: Box::new(TranscodeSqueezeEvict), @@ -102,6 +105,12 @@ impl LiquidCacheLocalBuilder { self } + /// Set maximum disk size in bytes + pub fn with_max_disk_bytes(mut self, max_disk_bytes: usize) -> Self { + self.max_disk_bytes = max_disk_bytes; + self + } + /// Set cache directory pub fn with_cache_dir(mut self, cache_dir: PathBuf) -> Self { self.cache_dir = cache_dir; @@ -152,20 +161,23 @@ impl LiquidCacheLocalBuilder { .await .map_err(|e| datafusion::error::DataFusionError::External(Box::new(e)))?; #[cfg(not(test))] - let cache = LiquidCacheParquet::new( + let cache = LiquidCacheParquet::new_with_options( self.batch_size, self.max_memory_bytes, + self.max_disk_bytes, store, self.cache_policy, self.squeeze_policy, self.hydration_policy, + true, ) .await; #[cfg(test)] - let cache = LiquidCacheParquet::new_with_squeeze_victim_concurrency( + let cache = LiquidCacheParquet::new_with_options( self.batch_size, self.max_memory_bytes, + self.max_disk_bytes, store, self.cache_policy, self.squeeze_policy, diff --git a/src/datafusion/src/cache/mod.rs b/src/datafusion/src/cache/mod.rs index daeb1967..71b5c6f6 100644 --- a/src/datafusion/src/cache/mod.rs +++ b/src/datafusion/src/cache/mod.rs @@ -249,9 +249,10 @@ impl LiquidCacheParquet { squeeze_policy: Box, hydration_policy: Box, ) -> Self { - Self::new_with_squeeze_victim_concurrency( + Self::new_with_options( batch_size, max_memory_bytes, + usize::MAX, store, cache_policy, squeeze_policy, @@ -271,12 +272,37 @@ impl LiquidCacheParquet { squeeze_policy: Box, hydration_policy: Box, squeeze_victims_concurrently: bool, + ) -> Self { + Self::new_with_options( + batch_size, + max_memory_bytes, + usize::MAX, + store, + cache_policy, + squeeze_policy, + hydration_policy, + squeeze_victims_concurrently, + ) + .await + } + + /// Create a new cache with all options including disk budget. + pub async fn new_with_options( + batch_size: usize, + max_memory_bytes: usize, + max_disk_bytes: usize, + store: t4::Store, + cache_policy: Box, + squeeze_policy: Box, + hydration_policy: Box, + squeeze_victims_concurrently: bool, ) -> Self { assert!(batch_size.is_power_of_two()); let metadata = Arc::new(ParquetCacheMetadata::new()); let cache_storage = LiquidCacheBuilder::new() .with_batch_size(batch_size) .with_max_memory_bytes(max_memory_bytes) + .with_max_disk_bytes(max_disk_bytes) .with_squeeze_policy(squeeze_policy) .with_cache_policy(cache_policy) .with_hydration_policy(hydration_policy) @@ -332,6 +358,11 @@ impl LiquidCacheParquet { self.cache_store.budget().disk_usage_bytes() } + /// Get the max disk bytes of the cache. + pub fn max_disk_bytes(&self) -> usize { + self.cache_store.budget().max_disk_bytes() + } + /// Flush the cache trace to a file. pub fn flush_trace(&self, to_file: impl AsRef) { self.cache_store.observer().flush_cache_trace(to_file); diff --git a/src/datafusion/src/reader/plantime/source.rs b/src/datafusion/src/reader/plantime/source.rs index 6c193c3f..9b3cb5a3 100644 --- a/src/datafusion/src/reader/plantime/source.rs +++ b/src/datafusion/src/reader/plantime/source.rs @@ -110,6 +110,21 @@ pub struct ParquetMetadataCacheReader { path: Path, } +#[cfg(test)] +impl ParquetMetadataCacheReader { + pub(crate) fn from_parts( + file_metrics: ParquetFileMetrics, + inner: ParquetObjectReader, + path: Path, + ) -> Self { + Self { + file_metrics, + inner, + path, + } + } +} + impl AsyncFileReader for ParquetMetadataCacheReader { fn get_byte_ranges( &mut self, diff --git a/src/datafusion/src/reader/runtime/liquid_cache_reader.rs b/src/datafusion/src/reader/runtime/liquid_cache_reader.rs index 87c123ec..ba838223 100644 --- a/src/datafusion/src/reader/runtime/liquid_cache_reader.rs +++ b/src/datafusion/src/reader/runtime/liquid_cache_reader.rs @@ -8,14 +8,45 @@ use arrow::buffer::BooleanBuffer; use arrow::compute::prep_null_mask_filter; use arrow::record_batch::RecordBatchOptions; use arrow_schema::{ArrowError, SchemaRef}; -use futures::{Stream, future::BoxFuture}; -use parquet::arrow::arrow_reader::{RowSelection, RowSelector}; +use futures::{Stream, StreamExt, future::BoxFuture}; +use parquet::arrow::arrow_reader::{ + ArrowReaderMetadata, ArrowReaderOptions, RowSelection, RowSelector, +}; +use parquet::arrow::{ParquetRecordBatchStreamBuilder, ProjectionMask}; +use parquet::file::metadata::ParquetMetaData; use crate::cache::{BatchID, CachedRowGroupRef}; -use crate::reader::plantime::LiquidRowFilter; -use crate::reader::runtime::utils::take_next_batch; +use crate::reader::plantime::{LiquidRowFilter, ParquetMetadataCacheReader}; +use crate::reader::runtime::utils::{get_root_column_ids, take_next_batch}; use crate::utils::{boolean_buffer_and_then, row_selector_to_boolean_buffer}; +#[derive(Clone)] +pub(crate) struct FallbackReader { + input: ParquetMetadataCacheReader, + metadata: Arc, + row_group_idx: usize, + projection: ProjectionMask, + row_count: usize, +} + +impl FallbackReader { + pub(crate) fn new( + input: ParquetMetadataCacheReader, + metadata: Arc, + row_group_idx: usize, + projection: ProjectionMask, + row_count: usize, + ) -> Self { + Self { + input, + metadata, + row_group_idx, + projection, + row_count, + } + } +} + pub(crate) struct LiquidCacheReader { state: ReaderState, row_filter: Option, @@ -48,6 +79,7 @@ struct LiquidCacheReaderInner { schema: SchemaRef, batch_size: usize, projection_columns: Vec, + fallback_reader: Option, } impl LiquidCacheReader { @@ -58,6 +90,7 @@ impl LiquidCacheReader { cached_row_group: CachedRowGroupRef, projection_columns: Vec, schema: SchemaRef, + fallback_reader: Option, ) -> Self { let inner = LiquidCacheReaderInner::new( batch_size, @@ -65,6 +98,7 @@ impl LiquidCacheReader { cached_row_group, projection_columns, Arc::clone(&schema), + fallback_reader, ); Self { state: ReaderState::Ready(inner), @@ -132,6 +166,7 @@ impl LiquidCacheReaderInner { cached_row_group: CachedRowGroupRef, projection_columns: Vec, schema: SchemaRef, + fallback_reader: Option, ) -> Self { Self { cached_row_group, @@ -140,6 +175,7 @@ impl LiquidCacheReaderInner { schema, batch_size, projection_columns, + fallback_reader, } } @@ -232,6 +268,7 @@ impl LiquidCacheReaderInner { } let mut arrays = Vec::with_capacity(self.projection_columns.len()); + let mut need_fallback = false; for &column_idx in &self.projection_columns { let column = self .cached_row_group @@ -244,21 +281,78 @@ impl LiquidCacheReaderInner { let array = column .get_arrow_array_with_filter(self.current_batch_id, selection) - .await - .ok_or_else(|| { - ArrowError::ComputeError(format!( - "column {column_idx} batch {} not cached", - *self.current_batch_id as usize - )) - })?; - + .await; + need_fallback |= array.is_none(); arrays.push(array); } + if need_fallback { + let fb = self.read_batch_fallback(selection).await?; + for (i, slot) in arrays.iter_mut().enumerate() { + if slot.is_none() { + *slot = Some(fb.column(i).clone()); + } + } + } + + let arrays: Vec<_> = arrays.into_iter().flatten().collect(); Ok(Some( RecordBatch::try_new(self.schema.clone(), arrays).unwrap(), )) } + + async fn read_batch_fallback( + &self, + selection: &BooleanBuffer, + ) -> Result { + let fb = self.fallback_reader.as_ref().ok_or_else(|| { + ArrowError::ComputeError("batch evicted from cache with no fallback reader".into()) + })?; + + let idx = usize::from(*self.current_batch_id); + let start = idx * self.batch_size; + let end = ((idx + 1) * self.batch_size).min(fb.row_count); + + let mut sel = Vec::new(); + if start > 0 { + sel.push(RowSelector::skip(start)); + } + sel.push(RowSelector::select(end - start)); + if end < fb.row_count { + sel.push(RowSelector::skip(fb.row_count - end)); + } + + let meta = + ArrowReaderMetadata::try_new(Arc::clone(&fb.metadata), ArrowReaderOptions::new()) + .map_err(|e| ArrowError::ExternalError(Box::new(e)))?; + let mut stream = ParquetRecordBatchStreamBuilder::new_with_metadata(fb.input.clone(), meta) + .with_projection(fb.projection.clone()) + .with_row_groups(vec![fb.row_group_idx]) + .with_row_selection(RowSelection::from(sel)) + .with_batch_size(self.batch_size) + .build() + .map_err(|e| ArrowError::ExternalError(Box::new(e)))?; + + let batch = stream + .next() + .await + .ok_or_else(|| ArrowError::ComputeError("fallback reader produced no data".into()))? + .map_err(|e| ArrowError::ExternalError(Box::new(e)))?; + + let src_cols = + get_root_column_ids(fb.metadata.file_metadata().schema_descr(), &fb.projection); + let mask = arrow::array::BooleanArray::new(selection.clone(), None); + let result: Result, _> = self + .projection_columns + .iter() + .map(|&col| { + let pos = src_cols.iter().position(|&c| c == col).unwrap(); + arrow::compute::filter(batch.column(pos), &mask) + }) + .collect(); + + RecordBatch::try_new(self.schema.clone(), result?) + } } #[cfg(test)] @@ -281,11 +375,14 @@ mod tests { use liquid_cache::cache::{AlwaysHydrate, squeeze_policies::Evict}; use liquid_cache::cache_policies::LiquidPolicy; use parquet::arrow::{ - ArrowWriter, + ArrowWriter, ProjectionMask, arrow_reader::{ArrowReaderMetadata, ArrowReaderOptions, RowSelection, RowSelector}, }; use std::sync::Arc; + use super::FallbackReader; + use crate::reader::plantime::ParquetMetadataCacheReader; + async fn make_row_group( batch_size: usize, batches: &[Vec], @@ -380,8 +477,15 @@ mod tests { let (row_group, schema) = make_row_group(batch_size, &[vec![1, 2], vec![3, 4]]).await; let selection = RowSelection::from(vec![RowSelector::select(4)]); - let reader = - LiquidCacheReader::new(batch_size, selection, None, row_group, vec![0], schema); + let reader = LiquidCacheReader::new( + batch_size, + selection, + None, + row_group, + vec![0], + schema, + None, + ); let batches = collect_batches(reader); assert_eq!(batches.len(), 2); @@ -395,8 +499,15 @@ mod tests { let (row_group, schema) = 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 = - LiquidCacheReader::new(batch_size, selection, None, row_group, vec![0], schema); + let reader = LiquidCacheReader::new( + batch_size, + selection, + None, + row_group, + vec![0], + schema, + None, + ); let batches = collect_batches(reader); assert_eq!(batches.len(), 1); @@ -416,6 +527,7 @@ mod tests { row_group, Vec::new(), Arc::new(Schema::new(Vec::::new())), + None, ); let batches = collect_batches(reader); @@ -439,6 +551,7 @@ mod tests { row_group, vec![0], schema, + None, ); let waker = futures::task::noop_waker(); @@ -467,6 +580,7 @@ mod tests { row_group, vec![0], schema, + None, ); let batches = collect_batches(reader); @@ -510,6 +624,7 @@ mod tests { row_group, vec![0], schema, + None, ); let batches = collect_batches(reader); @@ -538,6 +653,7 @@ mod tests { row_group, vec![0], schema, + None, ); let mut batches = collect_batches(reader); @@ -561,6 +677,7 @@ mod tests { row_group, vec![0], schema, + None, ); let next_batch = futures::executor::block_on(async { @@ -570,4 +687,77 @@ mod tests { assert!(next_batch.is_none()); } + + #[tokio::test] + async fn fallback_reads_evicted_batch_from_parquet() { + use object_store::local::LocalFileSystem; + use parquet::arrow::async_reader::ParquetObjectReader; + + let batch_size = 2; + let schema = Arc::new(Schema::new(vec![Field::new( + "col0", + DataType::Int32, + false, + )])); + + let tmp_dir = tempfile::tempdir().unwrap(); + let parquet_path = tmp_dir.path().join("test.parquet"); + let array: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3, 4])); + let batch = RecordBatch::try_new(Arc::clone(&schema), vec![array]).unwrap(); + let file = std::fs::File::create(&parquet_path).unwrap(); + let mut writer = ArrowWriter::try_new(file, Arc::clone(&schema), None).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + + let store = t4::mount(tmp_dir.path().join("cache.t4")).await.unwrap(); + let cache = LiquidCacheParquet::new( + batch_size, + usize::MAX, + store, + Box::new(LiquidPolicy::new()), + Box::new(Evict), + Box::new(AlwaysHydrate::new()), + ) + .await; + let cached_file = cache.register_or_get_file("test".to_string(), schema.clone()); + let row_group = cached_file.create_row_group(0, vec![]); + let column = row_group.get_column(0).unwrap(); + let b0: ArrayRef = Arc::new(Int32Array::from(vec![1, 2])); + column.insert(BatchID::from_raw(0), b0).await.unwrap(); + + let object_store = Arc::new(LocalFileSystem::new_with_prefix(tmp_dir.path()).unwrap()); + let obj_path = object_store::path::Path::from("test.parquet"); + let parquet_reader = ParquetObjectReader::new(object_store, obj_path.clone()); + + let file_reader = std::fs::File::open(&parquet_path).unwrap(); + let arrow_meta = + ArrowReaderMetadata::load(&file_reader, ArrowReaderOptions::new()).unwrap(); + let metadata = Arc::new(arrow_meta.metadata().clone()); + let projection = ProjectionMask::roots(metadata.file_metadata().schema_descr(), vec![0]); + + use datafusion::datasource::physical_plan::ParquetFileMetrics; + use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet; + let metrics = ExecutionPlanMetricsSet::new(); + let file_metrics = ParquetFileMetrics::new(0, "test.parquet", &metrics); + let pq_reader = + ParquetMetadataCacheReader::from_parts(file_metrics, parquet_reader, obj_path); + + let fallback = FallbackReader::new(pq_reader, Arc::clone(&metadata), 0, projection, 4); + + let selection = RowSelection::from(vec![RowSelector::select(4)]); + let reader = LiquidCacheReader::new( + batch_size, + selection, + None, + row_group, + vec![0], + schema, + Some(fallback), + ); + + 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]); + } } diff --git a/src/datafusion/src/reader/runtime/liquid_stream.rs b/src/datafusion/src/reader/runtime/liquid_stream.rs index b93dfb2f..4f7b715c 100644 --- a/src/datafusion/src/reader/runtime/liquid_stream.rs +++ b/src/datafusion/src/reader/runtime/liquid_stream.rs @@ -22,7 +22,7 @@ use std::{ task::{Context, Poll}, }; -use super::liquid_cache_reader::LiquidCacheReader; +use super::liquid_cache_reader::{FallbackReader, LiquidCacheReader}; use super::utils::{get_root_column_ids, limit_row_selection, offset_row_selection}; type PlanResult = Option; @@ -591,6 +591,27 @@ impl LiquidStream { pub fn schema(&self) -> &SchemaRef { &self.schema } + + fn make_cache_reader(&mut self, context: PlanningContext) -> LiquidCacheReader { + let row_group_idx = context.row_group_idx; + let reader_factory = self.reader.as_mut().unwrap(); + let fallback = FallbackReader::new( + reader_factory.input.clone(), + Arc::clone(&self.metadata), + row_group_idx, + self.projection.clone(), + self.metadata.row_group(row_group_idx).num_rows() as usize, + ); + LiquidCacheReader::new( + context.batch_size, + context.selection, + reader_factory.filter.take(), + context.cached_row_group, + context.projection_column_ids, + Arc::clone(&self.schema), + Some(fallback), + ) + } } impl Stream for LiquidStream { @@ -652,15 +673,7 @@ impl Stream for LiquidStream { self.state = StreamState::FillCache(fut); } else { LocalSpan::add_event(Event::new("LiquidStream::read_from_cache")); - let reader_factory = self.reader.as_mut().unwrap(); - let batch_reader = LiquidCacheReader::new( - context.batch_size, - context.selection, - reader_factory.filter.take(), - context.cached_row_group, - context.projection_column_ids, - Arc::clone(&self.schema), - ); + let batch_reader = self.make_cache_reader(context); self.state = StreamState::ReadFromCache(batch_reader); } } @@ -678,15 +691,7 @@ impl Stream for LiquidStream { Ok((reader_factory, context)) => { self.reader = Some(reader_factory); LocalSpan::add_event(Event::new("LiquidStream::read_from_cache")); - let reader_factory = self.reader.as_mut().unwrap(); - let batch_reader = LiquidCacheReader::new( - context.batch_size, - context.selection, - reader_factory.filter.take(), - context.cached_row_group, - context.projection_column_ids, - Arc::clone(&self.schema), - ); + let batch_reader = self.make_cache_reader(context); self.state = StreamState::ReadFromCache(batch_reader); } Err(e) => {