From 53414cd5e81fd2a558d7a1df65af9b37100a7dcf Mon Sep 17 00:00:00 2001 From: Anoop Date: Mon, 12 Oct 2026 14:22:35 +0530 Subject: [PATCH 1/4] fix(reader): honour bloom_filter_on_read (#524) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The liquid parquet reader ignores `datafusion.execution.parquet.bloom_filter_on_read`. Whenever a scan has a pruning predicate, it loads and applies every remaining row group's bloom filters, even with the option set to `false`. ```rust let config = SessionConfig::new() .set_bool("datafusion.execution.parquet.bloom_filter_on_read", false); let ctx = SessionContext::new_with_config(config); // register a parquet table `t` with bloom filters, then: // SELECT * FROM t WHERE a = 2 ``` With `ParquetSource`, this scan reads no bloom filters. Once the scan goes through `LiquidParquetSource`, it reads them anyway, and `row_groups_pruned_bloom_filter` reports row groups pruned by a filter the session turned off. ## Why it matters Each bloom filter is a separate ranged read per row group and predicate column, and the reads are not cached across queries. Callers turn the option off when bloom checks cannot pay for themselves — for example, many small files where file-level pruning has already selected exactly the files that match. Today they have no way to avoid those reads. The option was honoured in the original opener and dropped in d7f0741 ("remove some flags"). If leaving it out was deliberate, I'm happy to hear the reasoning; otherwise this restores it. ## The fix `LiquidParquetSource` already carries the session's `TableParquetOptions`. Pass `global.bloom_filter_on_read` into the morselizer and skip the bloom-filter load when it is off. As in DataFusion's own opener, the remaining row groups are then counted as matched by bloom filters, so the pruning metrics stay consistent between the two sources. ## Tests The existing `prunes_row_group_with_bloom_filter` test becomes a helper parameterised on the option, and a new test covers the off case: - `prunes_row_group_with_bloom_filter` (on): 1 row, 1 row group pruned, 1 matched. - `skips_bloom_filter_when_bloom_filter_on_read_is_off` (off): 1 row, 0 pruned, 2 matched. The predicate alone still produces the right answer. With the test applied to `main` and the fix left out: ``` $ cargo test -p liquid-cache-datafusion --lib bloom test optimizers::tests::prunes_row_group_with_bloom_filter ... ok test optimizers::tests::skips_bloom_filter_when_bloom_filter_on_read_is_off ... FAILED assertion `left == right` failed left: 1 right: 0 ``` With the fix, `cargo test -p liquid-cache-datafusion` passes (63 tests). --- src/datafusion/src/optimizers/mod.rs | 37 ++++++++++++++++--- .../src/reader/plantime/morselizer.rs | 22 ++++++++++- src/datafusion/src/reader/plantime/source.rs | 5 +++ 3 files changed, 58 insertions(+), 6 deletions(-) diff --git a/src/datafusion/src/optimizers/mod.rs b/src/datafusion/src/optimizers/mod.rs index d8950b11..f6c47068 100644 --- a/src/datafusion/src/optimizers/mod.rs +++ b/src/datafusion/src/optimizers/mod.rs @@ -230,7 +230,7 @@ mod tests { physical_plan::{ PhysicalExpr, collect, display::DisplayableExecutionPlan, filter_pushdown::PushedDown, }, - prelude::SessionContext, + prelude::{SessionConfig, SessionContext}, }; use liquid_cache::{ cache::{AlwaysHydrate, TranscodeEvict}, @@ -391,13 +391,21 @@ mod tests { writer.close().unwrap(); } - #[tokio::test] - async fn prunes_row_group_with_bloom_filter() { + /// Scan `a = 2` over a file whose two row groups both span `1..=4`, so + /// statistics keep both and only the bloom filter can rule one out. Returns + /// the rows produced and the bloom-filter pruning metric. + async fn scan_bloom_file( + bloom_filter_on_read: bool, + ) -> (usize, datafusion::physical_plan::metrics::PruningMetrics) { let tmp_dir = tempfile::tempdir().unwrap(); let parquet_path = tmp_dir.path().join("bloom.parquet"); write_bloom_file(&parquet_path); - let ctx = SessionContext::new(); + let config = SessionConfig::new().set_bool( + "datafusion.execution.parquet.bloom_filter_on_read", + bloom_filter_on_read, + ); + let ctx = SessionContext::new_with_config(config); ctx.register_parquet("t", parquet_path.to_str().unwrap(), Default::default()) .await .unwrap(); @@ -413,7 +421,7 @@ mod tests { let metrics = liquid_source(&rewritten).metrics().clone(); let batches = collect(rewritten, ctx.task_ctx()).await.unwrap(); - assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 1); + let rows = batches.iter().map(RecordBatch::num_rows).sum::(); let metric = metrics .clone_inner() @@ -425,7 +433,26 @@ mod tests { else { panic!("unexpected metric: {metric:?}"); }; + (rows, pruning_metrics) + } + + #[tokio::test] + async fn prunes_row_group_with_bloom_filter() { + let (rows, pruning_metrics) = scan_bloom_file(true).await; + assert_eq!(rows, 1); assert_eq!(pruning_metrics.pruned(), 1); + assert_eq!(pruning_metrics.matched(), 1); + } + + /// `bloom_filter_on_read = false` must keep the reader off the bloom + /// filters entirely: the row group the filter would have pruned is read, + /// and the predicate alone still produces the right answer. + #[tokio::test] + async fn skips_bloom_filter_when_bloom_filter_on_read_is_off() { + let (rows, pruning_metrics) = scan_bloom_file(false).await; + assert_eq!(rows, 1); + assert_eq!(pruning_metrics.pruned(), 0); + assert_eq!(pruning_metrics.matched(), 2); } /// Declining a scan costs it the cache, so the guard must key on what the scan diff --git a/src/datafusion/src/reader/plantime/morselizer.rs b/src/datafusion/src/reader/plantime/morselizer.rs index a8ff2b80..7e75f1b0 100644 --- a/src/datafusion/src/reader/plantime/morselizer.rs +++ b/src/datafusion/src/reader/plantime/morselizer.rs @@ -67,6 +67,7 @@ pub(crate) struct LiquidMorselizer { pub(crate) parquet_file_reader_factory: Arc, pub(crate) object_store_url: ObjectStoreUrl, pub(crate) reorder_filters: bool, + pub(crate) enable_bloom_filter: bool, pub(crate) liquid_cache: LiquidCacheParquetRef, pub(crate) expr_adapter_factory: Arc, pub(crate) span: Option>, @@ -188,6 +189,7 @@ impl Morselizer for LiquidMorselizer { projection, predicate, reorder_filters: self.reorder_filters, + enable_bloom_filter: self.enable_bloom_filter, liquid_cache: self.liquid_cache.clone(), expr_adapter_factory: Arc::clone(&self.expr_adapter_factory), file_identity, @@ -238,6 +240,7 @@ struct PreparedLiquidOpen { projection: ProjectionExprs, predicate: Option>, reorder_filters: bool, + enable_bloom_filter: bool, liquid_cache: LiquidCacheParquetRef, expr_adapter_factory: Arc, file_identity: ParquetFileIdentity, @@ -468,7 +471,10 @@ fn prepare_and_prune_by_stats(mut loaded: MetadataLoadedLiquidOpen) -> Result
  • Result
  • bool { + self.table_parquet_options.global.bloom_filter_on_read + } + /// Set the span for the LiquidParquetSource pub fn with_span(&self, span: fastrace::Span) -> Self { Self { @@ -175,6 +179,7 @@ impl FileSource for LiquidParquetSource { parquet_file_reader_factory, object_store_url: base_config.object_store_url.clone(), reorder_filters: self.reorder_filters(), + enable_bloom_filter: self.bloom_filter_on_read(), expr_adapter_factory, span: execution_span.map(Arc::new), lineages: Arc::clone(&self.lineages), From c946ea460a67e166d0d828940a2cd193b59369f5 Mon Sep 17 00:00:00 2001 From: Anoop Date: Tue, 13 Oct 2026 12:41:54 +0530 Subject: [PATCH 2/4] fix(reader): read the cache in cache-sized batches (#525) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Setting `datafusion.execution.batch_size` to anything other than the cache's batch size makes liquid-cache scans fail or return the wrong rows. ```sql -- cache batch size is the default 8192 SET datafusion.execution.batch_size = 2145; SELECT id FROM t; -- errors SELECT id FROM t LIMIT 4290; -- succeeds, rows from the wrong offsets ``` `LiquidCacheLocalBuilder::build` pins the session batch size to the cache's, so the default configuration does not hit this. A host that changes it afterwards — per query, or by sizing batches to row width — does. ## Why `LiquidCacheReader` walks the row selection in windows of the session batch size, and numbers each window with `current_batch_id`. But that id is a cache chunk index: the cache read and the parquet fallback both turn it back into rows using the *cache* batch size. When the two sizes differ, batch id N names one range of rows to the reader and a different range to everything that resolves it. - Without a `LIMIT`, the scan runs off the end of the row group: `parquet fallback ended before batch N`. - With a `LIMIT` that stops it first, it returns rows from the wrong offsets and succeeds. - On a fully cached row group, a selection mask sized by the session batch size is applied to a chunk sized by the cache's. `arrow::compute::filter` only rejects a mask *longer* than its target, so a shorter one silently selects the wrong rows. The only guard was a `debug_assert_eq!` in `LiquidCacheReader::new`, so debug builds panic and release builds carry on. ## The fix Window the selection by the cached row group's batch size, and drop the session batch size from `LiquidCacheReaderConfig` and `LiquidRowGroupPlanner`, which no longer need it. Callers still get session-sized batches: `DataSourceExec` re-splits every source stream to `datafusion.execution.batch_size`, and `session_batch_size_still_bounds_the_batches_the_caller_receives` checks that. ## Tests `src/datafusion-local/src/tests/batch_size_alignment.rs` covers a narrowed and a widened session batch size, `LIMIT`, the fully cached path, multiple row groups and multiple partitions, all against a cache at the default 8192. On `main` with the fix left out, in release (the debug build stops at the assert): ``` $ cargo test --release -p liquid-cache-datafusion-local --lib batch_size_alignment test result: FAILED. 0 passed; 7 failed ... "parquet fallback read failed: Parquet error: parquet fallback ended before batch 1" ... narrowed_session_batch_size_with_limit_returns_aligned_rows: assertion `left == right` failed ``` With the fix, all 7 pass in debug and release, and `cargo test -p liquid-cache-datafusion -p liquid-cache-datafusion-local` passes. One case is not covered here: page-index pruning that starts the selection partway into a row group. A scan with page-index pruning plus a pushed-down predicate panics in `boolean_buffer_and_then` on `main` even with matching batch sizes, so it cannot be tested until that is fixed; that is a separate issue. --- .../src/tests/batch_size_alignment.rs | 253 ++++++++++++++++++ src/datafusion-local/src/tests/mod.rs | 1 + .../src/reader/plantime/morselizer.rs | 3 - .../src/reader/runtime/liquid_cache_reader.rs | 20 +- src/datafusion/src/reader/runtime/morsel.rs | 2 - 5 files changed, 264 insertions(+), 15 deletions(-) create mode 100644 src/datafusion-local/src/tests/batch_size_alignment.rs diff --git a/src/datafusion-local/src/tests/batch_size_alignment.rs b/src/datafusion-local/src/tests/batch_size_alignment.rs new file mode 100644 index 00000000..5c2e29b3 --- /dev/null +++ b/src/datafusion-local/src/tests/batch_size_alignment.rs @@ -0,0 +1,253 @@ +//! Regression tests for batch-size alignment: the reader indexes the cache by batch +//! id, and the parquet fallback turns that id back into rows using the *cache* batch +//! size. Reading at the session batch size (`datafusion.execution.batch_size`) +//! instead addresses the wrong rows whenever the two differ — the scan either runs +//! off the end of the row group and fails, or silently returns rows from the wrong +//! offsets. +//! +//! `LiquidCacheLocalBuilder::build` pins `execution.batch_size` to the cache batch +//! size, which is why the default configuration never hit this. A host that sets a +//! per-query batch size afterwards — as `SET datafusion.execution.batch_size` does +//! here — used to break the alignment. + +use std::path::Path; +use std::sync::Arc; + +use arrow::array::{Array, Int64Array}; +use arrow::record_batch::RecordBatch; +use arrow_schema::{DataType, Field, Schema}; +use datafusion::error::Result; +use datafusion::prelude::{ParquetReadOptions, SessionConfig, SessionContext}; +use parquet::arrow::ArrowWriter; +use parquet::file::properties::WriterProperties; +use tempfile::TempDir; + +use crate::LiquidCacheLocalBuilder; + +/// One row group of `rows` rows, `id` ascending from 0. +fn write_single_row_group(path: &Path, rows: i64) { + write_parquet(path, rows, None); +} + +/// `rows` rows of ascending `id`, optionally split into row groups of +/// `row_group_rows`. +fn write_parquet(path: &Path, rows: i64, row_group_rows: Option) { + let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int64Array::from((0..rows).collect::>()))], + ) + .unwrap(); + let mut props = WriterProperties::builder(); + if let Some(n) = row_group_rows { + props = props.set_max_row_group_row_count(Some(n)); + } + let file = std::fs::File::create(path).unwrap(); + let mut writer = ArrowWriter::try_new(file, schema, Some(props.build())).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); +} + +/// Cache batch size stays at the 8192 default; the session batch size is then set +/// to `session_batch_size`, as a host sizing batches by row width would do. +async fn ctx_with_session_batch_size( + cache_dir: &Path, + parquet_path: &Path, + session_batch_size: usize, +) -> Result { + ctx_with_session_batch_size_and_partitions(cache_dir, parquet_path, session_batch_size, 1).await +} + +async fn ctx_with_session_batch_size_and_partitions( + cache_dir: &Path, + parquet_path: &Path, + session_batch_size: usize, + target_partitions: usize, +) -> Result { + std::fs::create_dir_all(cache_dir)?; + let mut config = SessionConfig::new(); + config.options_mut().execution.target_partitions = target_partitions; + let (ctx, cache) = LiquidCacheLocalBuilder::new() + .with_cache_dir(cache_dir.to_path_buf()) + .build(config) + .await?; + assert_eq!(cache.batch_size(), 8192, "cache batch size should be 8192"); + + ctx.sql(&format!( + "SET datafusion.execution.batch_size = {session_batch_size}" + )) + .await? + .collect() + .await?; + + ctx.register_parquet( + "t", + parquet_path.to_str().unwrap(), + ParquetReadOptions::default(), + ) + .await?; + Ok(ctx) +} + +async fn collect_batches(ctx: &SessionContext, sql: &str) -> Vec { + ctx.sql(sql).await.unwrap().collect().await.unwrap() +} + +fn ids_of(batches: &[RecordBatch]) -> Vec { + let mut ids = Vec::new(); + for batch in batches { + let col = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + ids.extend((0..col.len()).map(|i| col.value(i))); + } + ids +} + +async fn collect_ids(ctx: &SessionContext, sql: &str) -> Vec { + ids_of(&collect_batches(ctx, sql).await) +} + +/// The reported failure: a 12000-row row group (fewer than `2 * 8192` rows) read +/// with a session batch size of 2145 used to drive the fallback past the end of +/// the row group on query batch 2, panicking with "parquet fallback ended before +/// batch 2". +#[tokio::test] +async fn narrowed_session_batch_size_reads_the_whole_row_group() { + let tmp = TempDir::new().unwrap(); + let parquet_path = tmp.path().join("t.parquet"); + write_single_row_group(&parquet_path, 12000); + let ctx = ctx_with_session_batch_size(&tmp.path().join("cache"), &parquet_path, 2145) + .await + .unwrap(); + + for _ in 0..2 { + let ids = collect_ids(&ctx, "SELECT id FROM t").await; + assert_eq!(ids, (0..12000).collect::>()); + } +} + +/// Silent corruption, not an error: a LIMIT stopped the scan before the fallback +/// ran out of rows, so the query succeeded and returned rows 0-2144 followed by +/// rows 8192-10336 in place of rows 2145-4289. +#[tokio::test] +async fn narrowed_session_batch_size_with_limit_returns_aligned_rows() { + let tmp = TempDir::new().unwrap(); + let parquet_path = tmp.path().join("t.parquet"); + write_single_row_group(&parquet_path, 12000); + let ctx = ctx_with_session_batch_size(&tmp.path().join("cache"), &parquet_path, 2145) + .await + .unwrap(); + + let ids = collect_ids(&ctx, "SELECT id FROM t LIMIT 4290").await; + assert_eq!(ids, (0..4290).collect::>()); +} + +/// The fully-cached path misaligns as well. The cold run +/// populates cache batches 0 and 1; the warm run reads only from the cache and +/// used to emit rows 8192.. for query batch 1, because `read_arrow_array` applied +/// the 2145-bit selection to the 8192-row stored chunk and `arrow::compute::filter` +/// only rejects a mask *longer* than its target. +#[tokio::test] +async fn narrowed_session_batch_size_stays_aligned_on_the_fully_cached_path() { + let tmp = TempDir::new().unwrap(); + let parquet_path = tmp.path().join("t.parquet"); + write_single_row_group(&parquet_path, 12000); + let ctx = ctx_with_session_batch_size(&tmp.path().join("cache"), &parquet_path, 2145) + .await + .unwrap(); + + let sql = "SELECT id FROM t LIMIT 4290"; + let _cold = collect_ids(&ctx, sql).await; + let warm = collect_ids(&ctx, sql).await; + assert_eq!(warm, (0..4290).collect::>()); +} + +/// A session batch size *larger* than the cache batch size misaligned in the other +/// direction: the selection mask outran the stored chunk, and +/// `arrow::compute::filter` rejected it with "Filter predicate of length .. is +/// larger than target array of length ..". +#[tokio::test] +async fn widened_session_batch_size_stays_aligned() { + let tmp = TempDir::new().unwrap(); + let parquet_path = tmp.path().join("t.parquet"); + write_single_row_group(&parquet_path, 12000); + let ctx = ctx_with_session_batch_size(&tmp.path().join("cache"), &parquet_path, 20000) + .await + .unwrap(); + + for _ in 0..2 { + let ids = collect_ids(&ctx, "SELECT id FROM t").await; + assert_eq!(ids, (0..12000).collect::>()); + } +} + +/// The scan now reads in cache-sized batches, but the caller must still receive +/// session-sized ones: `DataSourceExec::execute` wraps every source stream in a +/// `BatchSplitStream` sized by the session config. This is what makes ignoring the +/// session batch size inside the reader safe. +#[tokio::test] +async fn session_batch_size_still_bounds_the_batches_the_caller_receives() { + let tmp = TempDir::new().unwrap(); + let parquet_path = tmp.path().join("t.parquet"); + write_single_row_group(&parquet_path, 12000); + let ctx = ctx_with_session_batch_size(&tmp.path().join("cache"), &parquet_path, 2145) + .await + .unwrap(); + + for _ in 0..2 { + let batches = collect_batches(&ctx, "SELECT id FROM t").await; + let oversized: Vec<_> = batches + .iter() + .map(|b| b.num_rows()) + .filter(|rows| *rows > 2145) + .collect(); + assert!( + oversized.is_empty(), + "batches exceeded the session batch size: {oversized:?}" + ); + assert_eq!(ids_of(&batches), (0..12000).collect::>()); + } +} + +/// `current_batch_id` restarts at 0 for each row group, so a file with several of +/// them exercises the mapping repeatedly rather than once. +#[tokio::test] +async fn multi_row_group_scan_stays_aligned() { + let tmp = TempDir::new().unwrap(); + let parquet_path = tmp.path().join("t.parquet"); + write_parquet(&parquet_path, 30000, Some(5000)); + let ctx = ctx_with_session_batch_size(&tmp.path().join("cache"), &parquet_path, 2145) + .await + .unwrap(); + + for _ in 0..2 { + let ids = collect_ids(&ctx, "SELECT id FROM t").await; + assert_eq!(ids, (0..30000).collect::>()); + } +} + +/// Several partitions read the same file concurrently, each with its own reader +/// and its own `current_batch_id` sequence. +#[tokio::test] +async fn multi_partition_scan_stays_aligned() { + let tmp = TempDir::new().unwrap(); + let parquet_path = tmp.path().join("t.parquet"); + write_parquet(&parquet_path, 30000, Some(5000)); + let ctx = ctx_with_session_batch_size_and_partitions( + &tmp.path().join("cache"), + &parquet_path, + 2145, + 4, + ) + .await + .unwrap(); + + for _ in 0..2 { + let mut ids = collect_ids(&ctx, "SELECT id FROM t").await; + ids.sort_unstable(); + assert_eq!(ids, (0..30000).collect::>()); + } +} diff --git a/src/datafusion-local/src/tests/mod.rs b/src/datafusion-local/src/tests/mod.rs index a08b76c5..3e6a90cd 100644 --- a/src/datafusion-local/src/tests/mod.rs +++ b/src/datafusion-local/src/tests/mod.rs @@ -19,6 +19,7 @@ use datafusion::{ }; use crate::LiquidCacheLocalBuilder; +mod batch_size_alignment; mod constant_conjunct; mod date_optimizer; mod filter_limit; diff --git a/src/datafusion/src/reader/plantime/morselizer.rs b/src/datafusion/src/reader/plantime/morselizer.rs index 7e75f1b0..6e9ee546 100644 --- a/src/datafusion/src/reader/plantime/morselizer.rs +++ b/src/datafusion/src/reader/plantime/morselizer.rs @@ -183,7 +183,6 @@ impl Morselizer for LiquidMorselizer { metrics, file_pruner, reader_factory, - batch_size: self.batch_size, logical_file_schema, output_schema, projection, @@ -234,7 +233,6 @@ struct PreparedLiquidOpen { metrics: LiquidFileMetrics, file_pruner: Option, reader_factory: Arc, - batch_size: usize, logical_file_schema: SchemaRef, output_schema: SchemaRef, projection: ProjectionExprs, @@ -616,7 +614,6 @@ fn plan_row_group_morsels(planned: PlannedRowGroups) -> Result, pub(crate) cached_row_group: CachedRowGroupRef, @@ -92,13 +91,17 @@ pub(crate) struct ParquetFallback { impl LiquidCacheReader { pub(crate) fn new(config: LiquidCacheReaderConfig) -> Self { - debug_assert_eq!( - config.batch_size, - config.cached_row_group.batch_size(), - "DataFusion and LiquidCache batch sizes must agree" - ); + // The selection is walked in cache-sized windows, not session-sized + // ones. `current_batch_id` indexes stored chunks, and both the cache + // read and the parquet fallback turn that id back into rows with the + // cache batch size; windowing at `datafusion.execution.batch_size` + // would address different rows than the id names whenever a caller + // sets it to anything other than the cache's own size. The caller is + // unaffected either way: `DataSourceExec` re-splits every source + // stream to the session batch size. + let batch_size = config.cached_row_group.batch_size(); let inner = LiquidCacheReaderInner::new( - config.batch_size, + batch_size, config.selection, config.cached_row_group, config.projection_columns, @@ -526,7 +529,6 @@ mod tests { use std::sync::Arc; struct TestRowGroup { - batch_size: usize, row_group: CachedRowGroupRef, schema: SchemaRef, fallback: ParquetFallbackConfig, @@ -543,7 +545,6 @@ mod tests { 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), @@ -622,7 +623,6 @@ mod tests { } TestRowGroup { - batch_size, row_group, schema, fallback: ParquetFallbackConfig { diff --git a/src/datafusion/src/reader/runtime/morsel.rs b/src/datafusion/src/reader/runtime/morsel.rs index 271f53ff..035a0a4e 100644 --- a/src/datafusion/src/reader/runtime/morsel.rs +++ b/src/datafusion/src/reader/runtime/morsel.rs @@ -31,7 +31,6 @@ pub(crate) struct LiquidRowGroupPlanner { pub(crate) row_filter: Option, pub(crate) cached_file: CachedFileRef, pub(crate) projection: ProjectionMask, - pub(crate) batch_size: usize, pub(crate) stream_schema: SchemaRef, pub(crate) output_schema: SchemaRef, pub(crate) projector: Arc, @@ -151,7 +150,6 @@ impl LiquidRowGroupPlanner { Some(LiquidRowGroupMorsel { config: LiquidCacheReaderConfig { - batch_size: self.batch_size, selection, row_filter: self.row_filter.clone(), cached_row_group, From 25358fff355104e5b3c50ce578cb68290510b0d0 Mon Sep 17 00:00:00 2001 From: Anoop Date: Fri, 16 Oct 2026 10:45:18 +0530 Subject: [PATCH 3/4] fix(cache): serve cached entries only to the file they came from (#526) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A process that reads more than 65,536 distinct parquet files gets another file's cached data back. A cache key packs the file id into 16 bits (`ParquetArrayID`), but `LiquidCacheParquet::register_or_get_file` hands out file ids from a counter that only climbs. The 65,537th distinct `(object store, path)` gets id 65536, which narrows to 0, so every key it computes is a key of the first file. Reads then return the first file's arrays: - silently wrong rows when the column types match; - a panic in `CachedColumn::array_to_record_batch` (`RecordBatch::try_new(..).unwrap()`) when they don't; - and `contains` reports the key as cached, so the newer file never caches its own data. The narrowing is guarded only by `debug_assert!`, so release builds hit this with no signal. A long-running server reading a stream of new files reaches it in normal use. ## Repro `files_whose_keys_collide_do_not_share_entries` (in `src/datafusion/src/cache/mod.rs`) caches a batch for one file, registers 65,535 more, then reads the same batch through the 65,537th. On `main`, with the test alone: - debug: fails at `debug_assert!(file_id <= u16::MAX as u64)` (`id.rs:150`); - release: fails with "a file was served another file's cached data" — the read returns the first file's array. ## Fix Store the full file id (an "identity") with every cached entry and check it on every lookup. A mismatch reads as a miss, so the caller falls back to parquet. An insert replaces whatever held the key. - **Index:** each slot holds `(identity, entry)`. - **Disk:** every object gets an 8-byte identity trailer. Objects stay keyed by entry id. An eviction/flush/hydration that lands after the key changed hands leaves bytes naming the old identity, and the new owner's read of them misses, so these races end in a miss rather than in wrong data and need no extra coordination. The trailer follows the payload so the payload keeps its alignment; the disk budget is charged for it. Nothing reads objects across a restart (the index is in-memory), so there is no on-disk compatibility concern. - **DataFusion layer:** `CachedFile` already holds the full `u64` file id; `CachedColumn` now carries it and passes it with every call. The key width is unchanged: the index (`congee`) supports 8-byte keys only, so widening the key would mean replacing the index. ## API change `LiquidCache::insert`, `get`, `eval_predicate`, `prefetch`, `try_read_liquid`, `contains`, `read_entry` and `eval_predicate_on_entry` take an `identity: u64`. Callers whose keys are already unique can pass any constant; the README, example and study bench pass `0`. `read_entry` / `eval_predicate_on_entry` act on an entry the caller already holds, so they trust that entry and check the identity only for data read from disk. ## Cost 8 bytes per index slot and per disk object, and one integer comparison per lookup. ## Tests - New: the repro above, `an_entry_is_served_only_to_its_identity`, `disk_bytes_written_for_another_identity_are_a_miss`, `a_slot_is_visible_only_to_its_identity`. - Replaced: `test_cache_entry_id_new_panic_file_id` pinned the debug-only guard; `test_cache_entry_id_new_keeps_low_16_bits_of_file_id` pins the narrowing the identity check relies on. - Two disk-budget tests now size the budget to include the trailer. - Workspace tests, doctests and the shuttle tests pass; test inventory is otherwise unchanged. ## Not in this PR - After 65,536 files, two files that share a narrowed id and are both being read keep replacing each other's entries. Results are correct; caching just stops helping for that pair. - `register_or_get_file`'s map keeps one entry per distinct file forever. - Replacing a disk entry (by hydration, or now by another identity) does not release its disk reservation. That predates this change and is separate. --- README.md | 11 +- examples/core.rs | 10 +- src/core/README.md | 16 +- src/core/src/cache/builders.rs | 27 +- src/core/src/cache/core.rs | 422 +++++++++++++++++++-------- src/core/src/cache/index.rs | 89 ++++-- src/core/src/cache/tests/policies.rs | 24 +- src/core/study/cache_storage.rs | 4 +- src/datafusion/src/cache/column.rs | 30 +- src/datafusion/src/cache/id.rs | 14 +- src/datafusion/src/cache/mod.rs | 73 ++++- 11 files changed, 528 insertions(+), 192 deletions(-) diff --git a/README.md b/README.md index 1cf6210b..fcd42aca 100644 --- a/README.md +++ b/README.md @@ -40,18 +40,21 @@ use std::sync::Arc; tokio_test::block_on(async { let cache = LiquidCacheBuilder::new().build().await; let entry_id = EntryID::from(1); + // The identity names the source the data belongs to; reads under any + // other identity miss. One source here, so any constant will do. + let identity = 0; let values = Arc::new(UInt64Array::from(vec![10, 11, 12, 13, 14, 15])); // 1) insert - cache.insert(entry_id, values.clone()).await; + cache.insert(entry_id, identity, values.clone()).await; // 2) get - let all_rows = cache.get(&entry_id).await.expect("entry should exist"); + let all_rows = cache.get(&entry_id, identity).await.expect("entry should exist"); // 3) get filtered (selection pushdown): keep rows 0, 2, 4 let selection = BooleanBuffer::from(vec![true, false, true, false, true, false]); let selected_rows = cache - .get(&entry_id) + .get(&entry_id, identity) .with_selection(&selection) .await .expect("entry should exist"); @@ -63,7 +66,7 @@ tokio_test::block_on(async { Arc::new(Literal::new(ScalarValue::UInt64(Some(12)))), )); let predicate_mask = cache - .eval_predicate(&entry_id, &predicate) + .eval_predicate(&entry_id, identity, &predicate) .await .expect("entry should exist") .expect("predicate should be evaluated in cache"); diff --git a/examples/core.rs b/examples/core.rs index 1df164c7..f5823291 100644 --- a/examples/core.rs +++ b/examples/core.rs @@ -18,13 +18,19 @@ async fn main() -> Result<(), Box> { .await; let entry_id = EntryID::from(7); + // The identity names the source this data belongs to. One source here, + // so any constant will do. + let identity = 0; let arrow_array = Arc::new(UInt64Array::from_iter_values(0..16)); - storage.insert(entry_id, arrow_array.clone()).await.unwrap(); + storage + .insert(entry_id, identity, arrow_array.clone()) + .await + .unwrap(); // Move data to disk so the read demonstrates async I/O storage.flush_all_to_disk().await.unwrap(); - let retrieved = storage.get(&entry_id).await.unwrap(); + let retrieved = storage.get(&entry_id, identity).await.unwrap(); assert_eq!(retrieved.as_ref(), arrow_array.as_ref()); Ok(()) diff --git a/src/core/README.md b/src/core/README.md index e0e75212..26acdce7 100644 --- a/src/core/README.md +++ b/src/core/README.md @@ -24,10 +24,12 @@ let storage = LiquidCacheBuilder::new().build().await; let entry_id = EntryID::from(42); let arrow_array = Arc::new(UInt64Array::from_iter_values(0..1000)); -// Insert once; replacement/placement is handled by the cache policy -storage.insert(entry_id, arrow_array.clone()).await; +// Insert once; replacement/placement is handled by the cache policy. +// `0` is the identity: the source this data belongs to. Reads under any +// other identity miss. One source here, so any constant will do. +storage.insert(entry_id, 0, arrow_array.clone()).await; -assert!(storage.contains(&entry_id)); +assert!(storage.contains(&entry_id, 0)); }); ``` @@ -43,13 +45,13 @@ let storage = LiquidCacheBuilder::new().build().await; let entry_id = EntryID::from(7); let arrow_array = Arc::new(UInt64Array::from_iter_values(0..16)); -storage.insert(entry_id, arrow_array.clone()).await; +storage.insert(entry_id, 0, arrow_array.clone()).await; // Move data to disk so the read will demonstrate async I/O storage.flush_all_to_disk().await; // Read asynchronously -let retrieved = storage.get(&entry_id).await.unwrap(); +let retrieved = storage.get(&entry_id, 0).await.unwrap(); assert_eq!(retrieved.as_ref(), arrow_array.as_ref()); }); ``` @@ -74,7 +76,7 @@ let entry_id = EntryID::from(8); let data = Arc::new(StringArray::from(vec![ Some("apple"), Some("banana"), None, Some("apple"), Some("cherry"), ])); -storage.insert(entry_id, data.clone()).await; +storage.insert(entry_id, 0, data.clone()).await; // Move data to disk so the read will demonstrate async I/O storage.flush_all_to_disk().await; @@ -93,7 +95,7 @@ let liquid_expr = liquid_cache::cache::LiquidExpr::try_new( // Read with predicate pushdown let mask = storage - .eval_predicate(&entry_id, &liquid_expr) + .eval_predicate(&entry_id, 0, &liquid_expr) .with_selection(&selection) .await .unwrap(); diff --git a/src/core/src/cache/builders.rs b/src/core/src/cache/builders.rs index 9bb4a843..4a08b9be 100644 --- a/src/core/src/cache/builders.rs +++ b/src/core/src/cache/builders.rs @@ -175,16 +175,23 @@ pub fn default_max_memory_bytes() -> usize { pub struct Insert<'a> { pub(super) storage: &'a Arc, pub(super) entry_id: EntryID, + pub(super) identity: u64, pub(super) batch: ArrayRef, pub(super) skip_gc: bool, pub(super) lineage: Option>, } impl<'a> Insert<'a> { - pub(super) fn new(storage: &'a Arc, entry_id: EntryID, batch: ArrayRef) -> Self { + pub(super) fn new( + storage: &'a Arc, + entry_id: EntryID, + identity: u64, + batch: ArrayRef, + ) -> Self { Self { storage, entry_id, + identity, batch, skip_gc: false, lineage: None, @@ -213,7 +220,9 @@ impl<'a> Insert<'a> { self.storage.add_lineage(&self.entry_id, lineage); } let batch = CacheEntry::memory_arrow(batch); - self.storage.insert_inner(self.entry_id, batch).await + self.storage + .insert_inner(self.entry_id, self.identity, batch) + .await } } @@ -231,15 +240,17 @@ impl<'a> IntoFuture for Insert<'a> { pub struct Get<'a> { pub(super) storage: &'a LiquidCache, pub(super) entry_id: &'a EntryID, + pub(super) identity: u64, pub(super) selection: Option<&'a BooleanBuffer>, pub(super) expression_hint: Option>, } impl<'a> Get<'a> { - pub(super) fn new(storage: &'a LiquidCache, entry_id: &'a EntryID) -> Self { + pub(super) fn new(storage: &'a LiquidCache, entry_id: &'a EntryID, identity: u64) -> Self { Self { storage, entry_id, + identity, selection: None, expression_hint: None, } @@ -271,6 +282,7 @@ impl<'a> Get<'a> { self.storage .read_arrow_array( self.entry_id, + self.identity, self.selection, self.expression_hint.as_deref(), ) @@ -326,6 +338,7 @@ fn maybe_gc_view_arrays(array: &ArrayRef) -> Option { pub struct EvaluatePredicate<'a> { pub(super) storage: &'a LiquidCache, pub(super) entry_id: &'a EntryID, + pub(super) identity: u64, pub(super) predicate: &'a LiquidExpr, pub(super) selection: Option<&'a BooleanBuffer>, } @@ -334,11 +347,13 @@ impl<'a> EvaluatePredicate<'a> { pub(super) fn new( storage: &'a LiquidCache, entry_id: &'a EntryID, + identity: u64, predicate: &'a LiquidExpr, ) -> Self { Self { storage, entry_id, + identity, predicate, selection: None, } @@ -353,7 +368,7 @@ impl<'a> EvaluatePredicate<'a> { /// Evaluate the predicate against the cached data. pub async fn read(self) -> Option { self.storage - .eval_predicate_internal(self.entry_id, self.selection, self.predicate) + .eval_predicate_internal(self.entry_id, self.identity, self.selection, self.predicate) .await } } @@ -451,9 +466,9 @@ mod tests { let cache = LiquidCacheBuilder::new().build().await; let entry_id = EntryID::from(123usize); - cache.insert(entry_id, root.clone()).await.unwrap(); + cache.insert(entry_id, 0, root.clone()).await.unwrap(); - let stored = cache.get(&entry_id).await.expect("array present"); + let stored = cache.get(&entry_id, 0).await.expect("array present"); let post_size = stored.get_array_memory_size(); // GC should have compacted the view arrays, reducing memory footprint. diff --git a/src/core/src/cache/core.rs b/src/core/src/cache/core.rs index d007d316..a1b77d8e 100644 --- a/src/core/src/cache/core.rs +++ b/src/core/src/cache/core.rs @@ -25,6 +25,15 @@ use crate::sync::Arc; /// Cache storage for liquid cache. /// +/// # Identity +/// +/// Every entry is cached for an *identity*: a `u64` naming the source the data +/// belongs to. Callers that pack several fields into a fixed-width +/// [`EntryID`] can have two sources compute the same key; the identity is how +/// the cache tells them apart. A lookup under any identity other than the one +/// an entry was cached for is a miss, and an insert replaces whatever held the +/// key. Callers whose keys are already unique can pass any constant. +/// /// Example (async read): /// ```rust /// use liquid_cache::cache::{LiquidCacheBuilder, EntryID}; @@ -36,10 +45,12 @@ use crate::sync::Arc; /// /// let entry_id = EntryID::from(0); /// let arrow_array = Arc::new(UInt64Array::from_iter_values(0..32)); -/// storage.insert(entry_id, arrow_array.clone()).await; +/// // The identity names the source this data belongs to. One source here, +/// // so any constant will do. +/// storage.insert(entry_id, 0, arrow_array.clone()).await; /// /// // Get the arrow array back asynchronously -/// let retrieved = storage.get(&entry_id).await.unwrap(); +/// let retrieved = storage.get(&entry_id, 0).await.unwrap(); /// assert_eq!(retrieved.as_ref(), arrow_array.as_ref()); /// }); /// ``` @@ -65,6 +76,9 @@ pub enum PrefetchResult { Absent, } +/// Length of the identity trailer appended to every object written to disk. +const IDENTITY_TRAILER_LEN: usize = std::mem::size_of::(); + /// Builder returned by [`LiquidCache::insert`] for configuring cache writes. impl LiquidCache { /// Return current cache statistics: counts and resource usage. @@ -80,7 +94,7 @@ impl LiquidCache { let mut memory_arrow_bytes = 0usize; let mut memory_liquid_bytes = 0usize; - self.index.for_each(|_, batch| match batch { + self.index.for_each(|_, slot| match &slot.entry { CacheEntry::MemoryArrow(array) => { memory_arrow_entries += 1; memory_arrow_bytes += array.get_array_memory_size(); @@ -113,52 +127,66 @@ impl LiquidCache { } } - /// Insert a batch into the cache. + /// Insert a batch into the cache for `identity`. pub fn insert<'a>( self: &'a Arc, entry_id: EntryID, + identity: u64, batch_to_cache: ArrayRef, ) -> Insert<'a> { - Insert::new(self, entry_id, batch_to_cache) + Insert::new(self, entry_id, identity, 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 a [`Get`] builder for the provided entry, as cached for `identity`. + pub fn get<'a>(&'a self, entry_id: &'a EntryID, identity: u64) -> Get<'a> { + Get::new(self, entry_id, identity) } /// Create an [`EvaluatePredicate`] builder for evaluating predicates on cached data. pub fn eval_predicate<'a>( &'a self, entry_id: &'a EntryID, + identity: u64, predicate: &'a LiquidExpr, ) -> EvaluatePredicate<'a> { - EvaluatePredicate::new(self, entry_id, predicate) + EvaluatePredicate::new(self, entry_id, identity, predicate) } /// Prefetch an entry into a memory-form snapshot without recording an access. - pub async fn prefetch(&self, entry_id: &EntryID) -> PrefetchResult { - let Some(entry) = self.index.get(entry_id) else { + pub async fn prefetch(&self, entry_id: &EntryID, identity: u64) -> PrefetchResult { + let Some(slot) = self.index.get_for(entry_id, identity) else { return PrefetchResult::Absent; }; - match entry.as_ref() { - CacheEntry::MemoryArrow(_) | CacheEntry::MemoryLiquid(_) => { - PrefetchResult::Snapshot(entry) + match &slot.entry { + entry @ (CacheEntry::MemoryArrow(_) | CacheEntry::MemoryLiquid(_)) => { + PrefetchResult::Snapshot(Arc::new(entry.clone())) } disk @ CacheEntry::DiskArrow { .. } => { - let Some(array) = self.read_disk_arrow_array(entry_id).await else { + let Some(array) = self.read_disk_arrow_array(entry_id, identity).await else { return PrefetchResult::Absent; }; - self.maybe_hydrate(entry_id, disk, MaterializedEntry::Arrow(&array), None) - .await; + self.maybe_hydrate( + entry_id, + identity, + disk, + MaterializedEntry::Arrow(&array), + None, + ) + .await; PrefetchResult::Snapshot(Arc::new(CacheEntry::memory_arrow(array))) } disk @ CacheEntry::DiskLiquid { .. } => { - let Some(array) = self.read_disk_liquid_array(entry_id).await else { + let Some(array) = self.read_disk_liquid_array(entry_id, identity).await else { return PrefetchResult::Absent; }; - self.maybe_hydrate(entry_id, disk, MaterializedEntry::Liquid(&array), None) - .await; + self.maybe_hydrate( + entry_id, + identity, + disk, + MaterializedEntry::Liquid(&array), + None, + ) + .await; PrefetchResult::Snapshot(Arc::new(CacheEntry::memory_liquid(array))) } } @@ -169,30 +197,38 @@ impl LiquidCache { pub async fn try_read_liquid( &self, entry_id: &EntryID, + identity: u64, ) -> Option { self.observer.on_try_read_liquid(); self.trace(InternalEvent::TryReadLiquid { entry: *entry_id }); - let batch = self.index.get(entry_id)?; + let slot = self.index.get_for(entry_id, identity)?; self.cache_policy - .notify_access(entry_id, CachedBatchType::from(batch.as_ref())); + .notify_access(entry_id, CachedBatchType::from(&slot.entry)); - match batch.as_ref() { + match &slot.entry { CacheEntry::MemoryLiquid(array) => Some(array.clone()), entry @ CacheEntry::DiskLiquid { .. } => { - let liquid = self.read_disk_liquid_array(entry_id).await?; - self.maybe_hydrate(entry_id, entry, MaterializedEntry::Liquid(&liquid), None) - .await; + let liquid = self.read_disk_liquid_array(entry_id, identity).await?; + self.maybe_hydrate( + entry_id, + identity, + entry, + MaterializedEntry::Liquid(&liquid), + None, + ) + .await; Some(liquid) } CacheEntry::DiskArrow { .. } | CacheEntry::MemoryArrow(_) => None, } } - /// Iterate over all entries in the cache. + /// Iterate over all entries in the cache, whichever identity holds them. /// 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); + self.index + .for_each(|entry_id, slot| f(entry_id, &slot.entry)); } /// Reset the cache. @@ -201,9 +237,9 @@ impl LiquidCache { self.budget.reset_usage(); } - /// Check whether the cache contains a batch. - pub fn contains(&self, entry_id: &EntryID) -> bool { - self.index.contains(entry_id) + /// Check whether the cache holds a batch under `entry_id` for `identity`. + pub fn contains(&self, entry_id: &EntryID, identity: u64) -> bool { + self.index.contains(entry_id, identity) } /// Get the config of the cache. @@ -234,18 +270,22 @@ impl LiquidCache { /// Flush all entries to disk. pub async fn flush_all_to_disk(&self) -> Result<(), CacheFull> { let mut entires = Vec::new(); - self.for_each_entry(|entry_id, batch| { - entires.push((*entry_id, batch.clone())); + self.index.for_each(|entry_id, slot| { + entires.push((*entry_id, slot.identity, slot.entry.clone())); }); - for (entry_id, batch) in entires { + for (entry_id, identity, batch) in entires { match &batch { CacheEntry::MemoryArrow(array) => { let bytes = arrow_to_bytes(array).expect("failed to convert arrow to bytes"); let disk_bytes = bytes.len(); - match self.write_batch_to_disk(entry_id, &batch, bytes).await { + match self + .write_batch_to_disk(entry_id, identity, &batch, bytes) + .await + { Ok(()) => { self.try_insert( entry_id, + identity, CacheEntry::disk_arrow(array.data_type().clone(), disk_bytes), ) .expect("failed to insert disk arrow entry"); @@ -257,12 +297,13 @@ impl LiquidCache { let liquid_bytes = liquid_array.to_bytes(); let disk_bytes = liquid_bytes.len(); match self - .write_batch_to_disk(entry_id, &batch, Bytes::from(liquid_bytes)) + .write_batch_to_disk(entry_id, identity, &batch, Bytes::from(liquid_bytes)) .await { Ok(()) => { self.try_insert( entry_id, + identity, CacheEntry::disk_liquid( liquid_array.original_arrow_data_type(), disk_bytes, @@ -287,6 +328,7 @@ impl LiquidCache { async fn write_in_memory_batch_to_disk( &self, entry_id: EntryID, + identity: u64, batch: CacheEntry, ) -> Result { match &batch { @@ -302,7 +344,7 @@ impl LiquidCache { unreachable!("memory Arrow eviction cannot remove entry"); }; if let Some(bytes_to_write) = bytes_to_write { - self.write_batch_to_disk(entry_id, &new_batch, bytes_to_write) + self.write_batch_to_disk(entry_id, identity, &new_batch, bytes_to_write) .await?; } Ok(new_batch) @@ -310,7 +352,7 @@ impl LiquidCache { CacheEntry::MemoryLiquid(liquid_array) => { let liquid_bytes = Bytes::from(liquid_array.to_bytes()); let disk_bytes = liquid_bytes.len(); - self.write_batch_to_disk(entry_id, &batch, liquid_bytes) + self.write_batch_to_disk(entry_id, identity, &batch, liquid_bytes) .await?; Ok(CacheEntry::disk_liquid( liquid_array.original_arrow_data_type(), @@ -327,10 +369,11 @@ impl LiquidCache { pub(crate) async fn insert_inner( &self, entry_id: EntryID, + identity: u64, mut batch_to_cache: CacheEntry, ) -> Result<(), CacheFull> { loop { - let Err(not_inserted) = self.try_insert(entry_id, batch_to_cache) else { + let Err(not_inserted) = self.try_insert(entry_id, identity, batch_to_cache) else { return Ok(()); }; self.trace(InternalEvent::InsertFailed { @@ -344,7 +387,7 @@ impl LiquidCache { // this can happen if the entry to be inserted is too large, in that case, // we write it to disk let on_disk_batch = self - .write_in_memory_batch_to_disk(entry_id, not_inserted) + .write_in_memory_batch_to_disk(entry_id, identity, not_inserted) .await?; batch_to_cache = on_disk_batch; continue; @@ -389,10 +432,15 @@ impl LiquidCache { } } - fn try_insert(&self, entry_id: EntryID, to_insert: CacheEntry) -> Result<(), CacheEntry> { + fn try_insert( + &self, + entry_id: EntryID, + identity: u64, + 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(); + let cached_batch_type = if let Some(slot) = self.index.get(&entry_id) { + let old_memory_size = slot.entry.memory_usage_bytes(); if self .budget .try_update_memory_usage(old_memory_size, new_memory_size) @@ -401,14 +449,14 @@ impl LiquidCache { return Err(to_insert); } let batch_type = CachedBatchType::from(&to_insert); - self.index.insert(&entry_id, to_insert); + self.index.insert(&entry_id, identity, 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); + self.index.insert(&entry_id, identity, to_insert); batch_type }; @@ -428,13 +476,13 @@ impl LiquidCache { }; assert!( matches!( - removed.as_ref(), + &removed.entry, CacheEntry::MemoryArrow(_) | CacheEntry::MemoryLiquid(_) ), "flush should only drop memory entries" ); self.budget - .try_update_memory_usage(removed.memory_usage_bytes(), 0) + .try_update_memory_usage(removed.entry.memory_usage_bytes(), 0) .expect("memory release cannot fail"); self.cache_policy.notify_remove(&entry_id); } @@ -443,7 +491,7 @@ impl LiquidCache { let Some(removed) = self.index.remove(&entry_id) else { return; }; - let disk_bytes = match removed.as_ref() { + let disk_bytes = match &removed.entry { CacheEntry::DiskLiquid { disk_bytes, .. } | CacheEntry::DiskArrow { disk_bytes, .. } => *disk_bytes, _ => panic!("remove_disk_entry called for non-disk entry"), @@ -452,7 +500,7 @@ impl LiquidCache { .remove(&entry_id_to_key(&entry_id)) .await .expect("disk remove failed"); - self.budget.release_disk(disk_bytes); + self.budget.release_disk(disk_bytes + IDENTITY_TRAILER_LEN); self.cache_policy.notify_remove(&entry_id); self.trace(InternalEvent::DiskEvict { entry: entry_id, @@ -496,15 +544,19 @@ impl LiquidCache { } async fn evict_victim_inner(&self, victim: EntryID) -> Result<(), CacheFull> { - let Some(mut victim_entry) = self.index.get(&victim) else { + let Some(slot) = self.index.get(&victim) else { return Ok(()); }; + // The rewrite is made for the identity the victim was read under. If + // another identity takes the key over meanwhile, the rewrite still + // carries the old identity, so the new owner reads it as a miss. + let identity = slot.identity; + let mut victim_entry = slot.entry.clone(); self.trace(InternalEvent::EvictionVictim { entry: victim }); loop { - let outcome = self.eviction_policy.evict( - victim_entry.as_ref(), - self.metadata.lineage(&victim).as_deref(), - ); + let outcome = self + .eviction_policy + .evict(&victim_entry, self.metadata.lineage(&victim).as_deref()); match outcome { EvictionOutcome::Replace { @@ -512,15 +564,15 @@ impl LiquidCache { bytes_to_write, } => { if let Some(bytes_to_write) = bytes_to_write { - self.write_batch_to_disk(victim, &new_batch, bytes_to_write) + self.write_batch_to_disk(victim, identity, &new_batch, bytes_to_write) .await?; } - match self.try_insert(victim, new_batch) { + match self.try_insert(victim, identity, new_batch) { Ok(()) => { break; } Err(batch) => { - victim_entry = Arc::new(batch); + victim_entry = batch; } } } @@ -536,6 +588,7 @@ impl LiquidCache { async fn maybe_hydrate( &self, entry_id: &EntryID, + identity: u64, cached: &CacheEntry, materialized: MaterializedEntry<'_>, expression: Option<&CacheExpression>, @@ -553,40 +606,48 @@ impl LiquidCache { cached: cached_type, new: new_type, }); - let _ = self.insert_inner(*entry_id, new_entry).await; + let _ = self.insert_inner(*entry_id, identity, new_entry).await; } } pub(crate) async fn read_arrow_array( &self, entry_id: &EntryID, + identity: u64, selection: Option<&BooleanBuffer>, expression: Option<&CacheExpression>, ) -> Option { self.observer.on_get(selection.is_some()); - let batch = self.index.get(entry_id)?; + let slot = self.index.get_for(entry_id, identity)?; self.cache_policy - .notify_access(entry_id, CachedBatchType::from(batch.as_ref())); - self.read_entry_inner(entry_id, batch.as_ref(), selection, expression) + .notify_access(entry_id, CachedBatchType::from(&slot.entry)); + self.read_entry_inner(entry_id, identity, &slot.entry, selection, expression) .await } - /// Read an already-looked-up cache entry. + /// Read an already-looked-up cache entry, cached for `identity`. + /// + /// `entry` is a value the caller already holds, so it is not checked + /// against `identity`: the caller must have obtained it for that identity, + /// e.g. from [`Self::prefetch`]. `identity` is checked when the entry's + /// data has to be read from disk. pub async fn read_entry( &self, entry_id: &EntryID, + identity: u64, entry: &CacheEntry, selection: Option<&BooleanBuffer>, expression: Option<&CacheExpression>, ) -> Option { self.observer.on_get(selection.is_some()); - self.read_entry_inner(entry_id, entry, selection, expression) + self.read_entry_inner(entry_id, identity, entry, selection, expression) .await } async fn read_entry_inner( &self, entry_id: &EntryID, + identity: u64, entry: &CacheEntry, selection: Option<&BooleanBuffer>, expression: Option<&CacheExpression>, @@ -613,7 +674,7 @@ impl LiquidCache { None => Some(array.to_arrow_array()), }, CacheEntry::DiskArrow { .. } | CacheEntry::DiskLiquid { .. } => { - self.read_disk_array(entry, entry_id, expression, selection) + self.read_disk_array(entry, entry_id, identity, expression, selection) .await } } @@ -623,6 +684,7 @@ impl LiquidCache { &self, entry: &CacheEntry, entry_id: &EntryID, + identity: u64, expression: Option<&CacheExpression>, selection: Option<&BooleanBuffer>, ) -> Option { @@ -633,9 +695,10 @@ impl LiquidCache { { return Some(arrow::array::new_empty_array(data_type)); } - let full_array = self.read_disk_arrow_array(entry_id).await?; + let full_array = self.read_disk_arrow_array(entry_id, identity).await?; self.maybe_hydrate( entry_id, + identity, entry, MaterializedEntry::Arrow(&full_array), expression, @@ -655,9 +718,10 @@ impl LiquidCache { { return Some(arrow::array::new_empty_array(data_type)); } - let liquid = self.read_disk_liquid_array(entry_id).await?; + let liquid = self.read_disk_liquid_array(entry_id, identity).await?; self.maybe_hydrate( entry_id, + identity, entry, MaterializedEntry::Liquid(&liquid), expression, @@ -672,15 +736,27 @@ impl LiquidCache { } } + /// Write `bytes` to disk for `identity`, followed by the identity itself. + /// + /// Objects are keyed by entry id alone, so a write landing after the key + /// changed hands can leave one identity's bytes where another's entry + /// points. The trailer is what lets a read notice that and miss. It sits + /// after the payload so the payload keeps the alignment it is read with. + /// The disk budget is charged for the trailer as well as the payload. async fn write_batch_to_disk( &self, entry_id: EntryID, + identity: u64, batch: &CacheEntry, bytes: Bytes, ) -> Result<(), CacheFull> { let len = bytes.len(); loop { - if self.budget.try_reserve_disk(len).is_ok() { + if self + .budget + .try_reserve_disk(len + IDENTITY_TRAILER_LEN) + .is_ok() + { break; } let victims = self.cache_policy.find_disk_victim(8); @@ -696,19 +772,34 @@ impl LiquidCache { kind: CachedBatchType::from(batch), bytes: len, }); + let mut object = Vec::with_capacity(len + IDENTITY_TRAILER_LEN); + object.extend_from_slice(&bytes); + object.extend_from_slice(&identity.to_le_bytes()); self.store - .put(entry_id_to_key(&entry_id), bytes.to_vec()) + .put(entry_id_to_key(&entry_id), object) .await .expect("write failed"); Ok(()) } - async fn read_disk_arrow_array(&self, entry_id: &EntryID) -> Option { - let bytes = match self.store.get(&entry_id_to_key(entry_id)).await { - Ok(bytes) => bytes, + /// The payload stored under `entry_id`, if it was written for `identity`. + async fn read_disk_bytes(&self, entry_id: &EntryID, identity: u64) -> Option { + let mut object = match self.store.get(&entry_id_to_key(entry_id)).await { + Ok(object) => object, Err(t4::Error::NotFound) => return None, Err(error) => panic!("read failed: {error}"), }; + let payload_len = object.len().checked_sub(IDENTITY_TRAILER_LEN)?; + let trailer: [u8; IDENTITY_TRAILER_LEN] = object[payload_len..].try_into().unwrap(); + if u64::from_le_bytes(trailer) != identity { + return None; + } + object.truncate(payload_len); + Some(Bytes::from(object)) + } + + async fn read_disk_arrow_array(&self, entry_id: &EntryID, identity: u64) -> Option { + let bytes = self.read_disk_bytes(entry_id, identity).await?; let bytes_len = bytes.len(); let cursor = std::io::Cursor::new(bytes); let mut reader = @@ -725,51 +816,60 @@ impl LiquidCache { async fn read_disk_liquid_array( &self, entry_id: &EntryID, + identity: u64, ) -> Option { - let bytes = match self.store.get(&entry_id_to_key(entry_id)).await { - Ok(bytes) => bytes, - Err(t4::Error::NotFound) => return None, - Err(error) => panic!("read failed: {error}"), - }; + let bytes = self.read_disk_bytes(entry_id, identity).await?; self.trace(InternalEvent::IoReadLiquid { entry: *entry_id, bytes: bytes.len(), }); Some(Arc::new(crate::liquid_array::LiquidArray::from_bytes( - Bytes::from(bytes), + bytes, ))) } pub(crate) async fn eval_predicate_internal( &self, entry_id: &EntryID, + identity: u64, selection_opt: Option<&BooleanBuffer>, predicate: &LiquidExpr, ) -> Option { self.observer.on_eval_predicate(); - let batch = self.index.get(entry_id)?; + let slot = self.index.get_for(entry_id, identity)?; self.cache_policy - .notify_access(entry_id, CachedBatchType::from(batch.as_ref())); - self.eval_predicate_on_entry_inner(entry_id, batch.as_ref(), selection_opt, predicate) - .await + .notify_access(entry_id, CachedBatchType::from(&slot.entry)); + self.eval_predicate_on_entry_inner( + entry_id, + identity, + &slot.entry, + selection_opt, + predicate, + ) + .await } - /// Evaluate a predicate on an already-looked-up cache entry. + /// Evaluate a predicate on an already-looked-up cache entry, cached for `identity`. + /// + /// As with [`Self::read_entry`], `entry` itself is not checked against + /// `identity`; only data read from disk is. pub async fn eval_predicate_on_entry( &self, entry_id: &EntryID, + identity: u64, entry: &CacheEntry, selection_opt: Option<&BooleanBuffer>, predicate: &LiquidExpr, ) -> Option { self.observer.on_eval_predicate(); - self.eval_predicate_on_entry_inner(entry_id, entry, selection_opt, predicate) + self.eval_predicate_on_entry_inner(entry_id, identity, entry, selection_opt, predicate) .await } async fn eval_predicate_on_entry_inner( &self, entry_id: &EntryID, + identity: u64, entry: &CacheEntry, selection_opt: Option<&BooleanBuffer>, predicate: &LiquidExpr, @@ -793,9 +893,15 @@ impl LiquidCache { Some(self.eval_predicate_on_array(filtered, predicate)) } entry @ CacheEntry::DiskArrow { .. } => { - let array = self.read_disk_arrow_array(entry_id).await?; - self.maybe_hydrate(entry_id, entry, MaterializedEntry::Arrow(&array), None) - .await; + let array = self.read_disk_arrow_array(entry_id, identity).await?; + self.maybe_hydrate( + entry_id, + identity, + entry, + MaterializedEntry::Arrow(&array), + None, + ) + .await; let mut owned = None; let selection = selection_opt.unwrap_or_else(|| { owned = Some(BooleanBuffer::new_set(array.len())); @@ -815,9 +921,15 @@ impl LiquidCache { Some(array.try_eval_predicate(predicate, selection)) } entry @ CacheEntry::DiskLiquid { .. } => { - let liquid = self.read_disk_liquid_array(entry_id).await?; - self.maybe_hydrate(entry_id, entry, MaterializedEntry::Liquid(&liquid), None) - .await; + let liquid = self.read_disk_liquid_array(entry_id, identity).await?; + self.maybe_hydrate( + entry_id, + identity, + entry, + MaterializedEntry::Liquid(&liquid), + None, + ) + .await; let mut owned = None; let selection = selection_opt.unwrap_or_else(|| { owned = Some(BooleanBuffer::new_set(liquid.len())); @@ -897,12 +1009,12 @@ mod tests { 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).await.unwrap(); + store.insert_inner(entry_id1, 0, array1).await.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() { + match &retrieved1.entry { CacheEntry::MemoryArrow(arr) => assert_eq!(arr.len(), 100), _ => panic!("Expected ArrowMemory"), } @@ -910,13 +1022,13 @@ mod tests { 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).await.unwrap(); + store.insert_inner(entry_id2, 0, array2).await.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).await.unwrap(); + store.insert_inner(entry_id1, 0, array3).await.unwrap(); assert_eq!(store.budget.memory_usage_bytes(), size3 + size2); assert!(store.index().get(&EntryID::from(999)).is_none()); @@ -936,19 +1048,19 @@ mod tests { let store = create_cache_store(8000, Box::new(advisor)).await; // Small budget to force advice store - .insert_inner(entry_id1, create_test_array(800)) + .insert_inner(entry_id1, 0, create_test_array(800)) .await .unwrap(); - match store.index().get(&entry_id1).unwrap().as_ref() { + match &store.index().get(&entry_id1).unwrap().entry { CacheEntry::MemoryArrow(_) => {} other => panic!("Expected ArrowMemory, got {other:?}"), } store - .insert_inner(entry_id2, create_test_array(800)) + .insert_inner(entry_id2, 0, create_test_array(800)) .await .unwrap(); - match store.index().get(&entry_id1).unwrap().as_ref() { + match &store.index().get(&entry_id1).unwrap().entry { CacheEntry::MemoryLiquid(_) => {} other => panic!("Expected LiquidMemory after eviction, got {other:?}"), } @@ -995,7 +1107,7 @@ mod tests { 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).await.unwrap(); + store.insert(entry_id, 0, array).await.unwrap(); } }); })); @@ -1029,8 +1141,14 @@ mod tests { // 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).await.unwrap(); - storage.insert(EntryID::from(2usize), arr2).await.unwrap(); + storage + .insert(EntryID::from(1usize), 0, arr1) + .await + .unwrap(); + storage + .insert(EntryID::from(2usize), 0, arr2) + .await + .unwrap(); // Stats after insert: 2 entries, memory usage > 0, disk usage == 0 let s = storage.stats(); @@ -1054,21 +1172,68 @@ mod tests { let entry_id = EntryID::from(321usize); let array = create_test_arrow_array(8); - store.insert(entry_id, array.clone()).await.unwrap(); + store.insert(entry_id, 0, array.clone()).await.unwrap(); store.flush_all_to_disk().await.unwrap(); { let entry = store.index().get(&entry_id).unwrap(); - assert!(matches!(entry.as_ref(), CacheEntry::DiskArrow { .. })); + assert!(matches!(&entry.entry, CacheEntry::DiskArrow { .. })); } - let result = store.get(&entry_id).await.expect("present"); + let result = store.get(&entry_id, 0).await.expect("present"); assert_eq!(result.as_ref(), array.as_ref()); { let entry = store.index().get(&entry_id).unwrap(); - assert!(matches!(entry.as_ref(), CacheEntry::MemoryArrow(_))); + assert!(matches!(&entry.entry, CacheEntry::MemoryArrow(_))); } } + #[tokio::test] + async fn an_entry_is_served_only_to_its_identity() { + let store = create_cache_store(1 << 20, Box::new(LiquidPolicy::new())).await; + let entry_id = EntryID::from(321usize); + let array = create_test_arrow_array(8); + + store.insert(entry_id, 1, array.clone()).await.unwrap(); + assert!(store.contains(&entry_id, 1)); + assert!(!store.contains(&entry_id, 2)); + assert!(store.get(&entry_id, 2).await.is_none()); + assert!(matches!( + store.prefetch(&entry_id, 2).await, + PrefetchResult::Absent + )); + + store.flush_all_to_disk().await.unwrap(); + assert!(store.get(&entry_id, 2).await.is_none()); + let result = store.get(&entry_id, 1).await.expect("present"); + assert_eq!(result.as_ref(), array.as_ref()); + + let replacement = create_test_arrow_array(16); + store + .insert(entry_id, 2, replacement.clone()) + .await + .unwrap(); + assert!(!store.contains(&entry_id, 1)); + let result = store.get(&entry_id, 2).await.expect("present"); + assert_eq!(result.as_ref(), replacement.as_ref()); + } + + #[tokio::test] + async fn disk_bytes_written_for_another_identity_are_a_miss() { + let store = create_cache_store(1 << 20, Box::new(LiquidPolicy::new())).await; + let entry_id = EntryID::from(321usize); + let array = create_test_arrow_array(8); + + store.insert(entry_id, 1, array.clone()).await.unwrap(); + store.flush_all_to_disk().await.unwrap(); + let disk_entry = store.index().get(&entry_id).unwrap().entry.clone(); + assert!(matches!(disk_entry, CacheEntry::DiskArrow { .. })); + + // A record naming identity 2 over bytes identity 1 wrote: what a + // rewrite landing after the key changed hands leaves behind. + store.index().insert(&entry_id, 2, disk_entry); + assert!(store.get(&entry_id, 2).await.is_none()); + } + #[tokio::test] async fn missing_disk_blob_is_a_cache_miss() { let directory = tempfile::tempdir().unwrap(); @@ -1079,11 +1244,14 @@ mod tests { .await; let id = EntryID::from(320usize); - cache.insert(id, create_test_arrow_array(8)).await.unwrap(); + cache + .insert(id, 0, create_test_arrow_array(8)) + .await + .unwrap(); cache.flush_all_to_disk().await.unwrap(); store.remove(&entry_id_to_key(&id)).await.unwrap(); - assert!(cache.get(&id).await.is_none()); + assert!(cache.get(&id, 0).await.is_none()); } #[tokio::test] @@ -1095,20 +1263,20 @@ mod tests { Arc::new(crate::liquid_array::LiquidArray::from_arrow_array(&arrow_array).unwrap()); store - .insert_inner(entry_id, CacheEntry::memory_liquid(liquid.clone())) + .insert_inner(entry_id, 0, CacheEntry::memory_liquid(liquid.clone())) .await .unwrap(); store.flush_all_to_disk().await.unwrap(); { let entry = store.index().get(&entry_id).unwrap(); - assert!(matches!(entry.as_ref(), CacheEntry::DiskLiquid { .. })); + assert!(matches!(&entry.entry, CacheEntry::DiskLiquid { .. })); } - let result = store.get(&entry_id).await.expect("present"); + let result = store.get(&entry_id, 0).await.expect("present"); assert_eq!(result.as_ref(), arrow_array.as_ref()); { let entry = store.index().get(&entry_id).unwrap(); - assert!(matches!(entry.as_ref(), CacheEntry::MemoryLiquid(_))); + assert!(matches!(&entry.entry, CacheEntry::MemoryLiquid(_))); } } @@ -1122,18 +1290,18 @@ mod tests { .await; let array: ArrayRef = Arc::new(Int32Array::from_iter_values(0..16)); - let err = cache.insert(EntryID::from(900usize), array).await; + let err = cache.insert(EntryID::from(900usize), 0, array).await; assert_eq!(err, Err(CacheFull)); - assert!(!cache.contains(&EntryID::from(900usize))); + assert!(!cache.contains(&EntryID::from(900usize), 0)); } #[tokio::test] async fn insert_until_disk_full_then_evicts_oldest_disk_entry() { let first_array: ArrayRef = Arc::new(Int32Array::from_iter_values(0..16)); let second_array: ArrayRef = Arc::new(Int32Array::from_iter_values(16..32)); - let first_bytes = arrow_to_bytes(&first_array).unwrap().len(); - let second_bytes = arrow_to_bytes(&second_array).unwrap().len(); + let first_bytes = arrow_to_bytes(&first_array).unwrap().len() + IDENTITY_TRAILER_LEN; + let second_bytes = arrow_to_bytes(&second_array).unwrap().len() + IDENTITY_TRAILER_LEN; let cache = LiquidCacheBuilder::new() .with_max_memory_bytes(1 << 20) .with_max_disk_bytes(first_bytes.max(second_bytes)) @@ -1144,16 +1312,16 @@ mod tests { let first = EntryID::from(910usize); let second = EntryID::from(911usize); - cache.insert(first, first_array).await.unwrap(); + cache.insert(first, 0, first_array).await.unwrap(); cache.flush_all_to_disk().await.unwrap(); - assert!(cache.contains(&first)); + assert!(cache.contains(&first, 0)); - cache.insert(second, second_array).await.unwrap(); + cache.insert(second, 0, second_array).await.unwrap(); cache.flush_all_to_disk().await.unwrap(); - assert!(!cache.contains(&first)); + assert!(!cache.contains(&first, 0)); assert!(matches!( - cache.index().get(&second).unwrap().as_ref(), + &cache.index().get(&second).unwrap().entry, CacheEntry::DiskArrow { .. } )); } @@ -1172,19 +1340,19 @@ mod tests { .await; let first = EntryID::from(912usize); let second = EntryID::from(913usize); - cache.insert(first, first_array).await.unwrap(); + cache.insert(first, 0, first_array).await.unwrap(); cache.flush_all_to_disk().await.unwrap(); - cache.insert(second, second_array).await.unwrap(); + cache.insert(second, 0, second_array).await.unwrap(); cache.flush_all_to_disk().await.unwrap(); - assert!(!cache.contains(&first) || !cache.contains(&second)); + assert!(!cache.contains(&first, 0) || !cache.contains(&second, 0)); } #[tokio::test] async fn disk_eviction_releases_budget() { let array: ArrayRef = Arc::new(Int32Array::from_iter_values(0..16)); - let disk_bytes = arrow_to_bytes(&array).unwrap().len(); + let disk_bytes = arrow_to_bytes(&array).unwrap().len() + IDENTITY_TRAILER_LEN; let cache = LiquidCacheBuilder::new() .with_max_memory_bytes(1 << 20) .with_max_disk_bytes(disk_bytes) @@ -1193,14 +1361,14 @@ mod tests { .build() .await; let entry = EntryID::from(914usize); - cache.insert(entry, array).await.unwrap(); + cache.insert(entry, 0, array).await.unwrap(); cache.flush_all_to_disk().await.unwrap(); let before = cache.stats().disk_usage_bytes; cache.remove_disk_entry(entry).await; assert_eq!(cache.stats().disk_usage_bytes, before - disk_bytes); - assert!(!cache.contains(&entry)); + assert!(!cache.contains(&entry, 0)); } #[tokio::test] @@ -1213,11 +1381,11 @@ mod tests { .await; let entry_id = EntryID::from(901usize); let array: ArrayRef = Arc::new(Int32Array::from_iter_values(0..16)); - cache.insert(entry_id, array).await.unwrap(); + cache.insert(entry_id, 0, array).await.unwrap(); let result = cache.flush_all_to_disk().await; assert_eq!(result, Ok(())); - assert!(!cache.contains(&entry_id)); + assert!(!cache.contains(&entry_id, 0)); } } diff --git a/src/core/src/cache/index.rs b/src/core/src/cache/index.rs index 62acca36..8607acec 100644 --- a/src/core/src/cache/index.rs +++ b/src/core/src/cache/index.rs @@ -9,8 +9,27 @@ use std::{ use crate::cache::{cached_batch::CacheEntry, utils::EntryID}; +/// What the index holds for a key: the entry, and the identity of the source +/// it was cached for. +/// +/// A key does not name its source uniquely: callers pack fields into fixed +/// widths, so two sources can compute the same key. The identity is what tells +/// them apart, so a read under any other identity must treat the slot as +/// absent. +#[derive(Debug)] +pub(crate) struct Slot { + pub(crate) identity: u64, + pub(crate) entry: CacheEntry, +} + +impl Slot { + pub(crate) fn belongs_to(&self, identity: u64) -> bool { + self.identity == identity + } +} + pub(crate) struct ArtIndex { - art: CongeeArc, + art: CongeeArc, entry_count: AtomicUsize, } @@ -22,36 +41,46 @@ impl Debug for ArtIndex { impl ArtIndex { pub(crate) fn new() -> Self { - let art: CongeeArc = CongeeArc::new(); + let art: CongeeArc = CongeeArc::new(); Self { art, entry_count: AtomicUsize::new(0), } } - pub(crate) fn get(&self, entry_id: &EntryID) -> Option> { + /// The slot under `entry_id`, whichever identity holds it. + pub(crate) fn get(&self, entry_id: &EntryID) -> Option> { let guard = self.art.pin(); - let batch = self.art.get(*entry_id, &guard)?; - Some(batch) + self.art.get(*entry_id, &guard) } - pub(crate) fn contains(&self, entry_id: &EntryID) -> bool { - let guard = self.art.pin(); - self.art.get(*entry_id, &guard).is_some() + /// The slot under `entry_id`, only if `identity` holds it. + pub(crate) fn get_for(&self, entry_id: &EntryID, identity: u64) -> Option> { + self.get(entry_id).filter(|slot| slot.belongs_to(identity)) + } + + pub(crate) fn contains(&self, entry_id: &EntryID, identity: u64) -> bool { + self.get_for(entry_id, identity).is_some() } - pub(crate) fn insert(&self, entry_id: &EntryID, batch: CacheEntry) { + /// Store `batch` under `entry_id` for `identity`, replacing whatever held + /// the key, including a slot held by another identity. + pub(crate) fn insert(&self, entry_id: &EntryID, identity: u64, batch: CacheEntry) { let guard = self.art.pin(); + let slot = Slot { + identity, + entry: batch, + }; let existing = self .art - .insert(*entry_id, Arc::new(batch), &guard) + .insert(*entry_id, Arc::new(slot), &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> { + 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() { @@ -68,7 +97,7 @@ impl ArtIndex { self.entry_count.store(0, Ordering::Relaxed); } - pub(crate) fn for_each(&self, mut f: impl FnMut(&EntryID, &CacheEntry)) { + pub(crate) fn for_each(&self, mut f: impl FnMut(&EntryID, &Slot)) { let guard = self.art.pin(); for id in self.art.keys().into_iter() { f( @@ -106,21 +135,21 @@ mod tests { let array1 = create_test_array(100); // Initially, entries should not be cached - assert!(!store.contains(&entry_id1)); - assert!(!store.contains(&entry_id2)); + assert!(!store.contains(&entry_id1, 0)); + assert!(!store.contains(&entry_id2, 0)); assert!(store.get(&entry_id1).is_none()); // Insert an entry and verify it's cached { - store.insert(&entry_id1, array1.clone()); + store.insert(&entry_id1, 0, array1.clone()); } - assert!(store.contains(&entry_id1)); - assert!(!store.contains(&entry_id2)); + assert!(store.contains(&entry_id1, 0)); + assert!(!store.contains(&entry_id2, 0)); // Get should return the cached value - match store.get(&entry_id1) { - Some(batch) => match batch.as_ref() { + match store.get_for(&entry_id1, 0) { + Some(slot) => match &slot.entry { CacheEntry::MemoryArrow(arr) => assert_eq!(arr.len(), 100), _ => panic!("Expected ArrowMemory batch"), }, @@ -134,13 +163,29 @@ mod tests { let entry_id: EntryID = EntryID::from(1); let array = create_test_array(100); - store.insert(&entry_id, array.clone()); + store.insert(&entry_id, 0, array.clone()); let entry_id: EntryID = EntryID::from(1); - assert!(store.contains(&entry_id)); + assert!(store.contains(&entry_id, 0)); store.reset(); let entry_id: EntryID = EntryID::from(1); - assert!(!store.contains(&entry_id)); + assert!(!store.contains(&entry_id, 0)); + } + + #[test] + fn a_slot_is_visible_only_to_its_identity() { + let store = ArtIndex::new(); + let entry_id = EntryID::from(1); + store.insert(&entry_id, 7, create_test_array(10)); + + assert!(store.contains(&entry_id, 7)); + assert!(!store.contains(&entry_id, 8)); + assert!(store.get_for(&entry_id, 8).is_none()); + + store.insert(&entry_id, 8, create_test_array(20)); + assert!(!store.contains(&entry_id, 7)); + assert!(store.contains(&entry_id, 8)); + assert_eq!(store.entry_count(), 1); } } diff --git a/src/core/src/cache/tests/policies.rs b/src/core/src/cache/tests/policies.rs index ee3681a4..e074adbd 100644 --- a/src/core/src/cache/tests/policies.rs +++ b/src/core/src/cache/tests/policies.rs @@ -26,12 +26,12 @@ async fn default_policies() { for i in 0..5 { let entry_id = EntryID::from(i); - cache.insert(entry_id, test_array.clone()).await.unwrap(); + cache.insert(entry_id, 0, test_array.clone()).await.unwrap(); } for i in 0..5 { let entry_id = EntryID::from(i); - let array = cache.get(&entry_id).read().await.unwrap(); + let array = cache.get(&entry_id, 0).read().await.unwrap(); assert_eq!(array.len(), test_array.len()); } @@ -54,17 +54,17 @@ async fn insert_wont_fit_cache() { .build() .await; cache - .insert(EntryID::from(0), test_array.clone()) + .insert(EntryID::from(0), 0, test_array.clone()) .await .unwrap(); let array_3x = arrow::compute::concat(&[&test_array, &test_array, &test_array]).unwrap(); let array_9x = arrow::compute::concat(&[&array_3x, &array_3x, &array_3x]).unwrap(); let array_27x = arrow::compute::concat(&[&array_9x, &array_9x, &array_9x]).unwrap(); cache - .insert(EntryID::from(1), array_27x.clone()) + .insert(EntryID::from(1), 0, array_27x.clone()) .await .unwrap(); - cache.get(&EntryID::from(1)).read().await.unwrap(); + cache.get(&EntryID::from(1), 0).read().await.unwrap(); let trace = cache.consume_event_trace(); let json_trace = serde_json::to_string(&trace).unwrap(); @@ -100,7 +100,7 @@ async fn liquid_eviction_reads_memory_and_disk() { .into_iter() .enumerate() { - cache.insert(EntryID::from(id), array).await.unwrap(); + cache.insert(EntryID::from(id), 0, array).await.unwrap(); } let mut states = Vec::new(); @@ -113,13 +113,13 @@ async fn liquid_eviction_reads_memory_and_disk() { let id = EntryID::from(0); assert_eq!( - cache.get(&id).read().await.unwrap().as_ref(), + cache.get(&id, 0).read().await.unwrap().as_ref(), integers.as_ref() ); let selection = BooleanBuffer::from_iter((0..integers.len()).map(|index| index % 3 == 0)); let selected = cache - .get(&id) + .get(&id, 0) .with_selection(&selection) .read() .await @@ -137,7 +137,11 @@ async fn liquid_eviction_reads_memory_and_disk() { Arc::new(Literal::new(ScalarValue::Int64(Some(1_000)))), )); let predicate = LiquidExpr::try_new(physical, &DataType::Int64).unwrap(); - let actual = cache.eval_predicate(&id, &predicate).read().await.unwrap(); + let actual = cache + .eval_predicate(&id, 0, &predicate) + .read() + .await + .unwrap(); let expected = arrow::array::BooleanArray::from_iter( integers .as_any() @@ -157,7 +161,7 @@ async fn liquid_eviction_reads_memory_and_disk() { ))); assert!(disk_states.contains(&CachedBatchType::DiskLiquid)); assert_eq!( - cache.get(&id).read().await.unwrap().as_ref(), + cache.get(&id, 0).read().await.unwrap().as_ref(), integers.as_ref() ); } diff --git a/src/core/study/cache_storage.rs b/src/core/study/cache_storage.rs index 18e01d5c..b2571753 100644 --- a/src/core/study/cache_storage.rs +++ b/src/core/study/cache_storage.rs @@ -108,7 +108,7 @@ fn main() { continue; }; if storage - .eval_predicate(id, &liquid_expr) + .eval_predicate(id, 0, &liquid_expr) .with_selection(&selection) .await .is_some() @@ -161,7 +161,7 @@ fn load_and_insert_referer( let id = EntryID::from(idx); ids.push(id); total_size += array.get_array_memory_size(); - storage.insert(id, array).await.unwrap(); + storage.insert(id, 0, array).await.unwrap(); idx += 1; } diff --git a/src/datafusion/src/cache/column.rs b/src/datafusion/src/cache/column.rs index c08cf9af..a0cdaba6 100644 --- a/src/datafusion/src/cache/column.rs +++ b/src/datafusion/src/cache/column.rs @@ -22,6 +22,11 @@ pub struct CachedColumn { cache_store: Arc, field: Arc, column_path: ColumnAccessPath, + /// The unnarrowed id of the file this column belongs to. `column_path` + /// keeps only its low 16 bits, so another file can compute the same + /// entry ids; the cache serves an entry only to the identity it was + /// cached for. + file_identity: u64, expression: Option>, snapshots: Arc, } @@ -55,6 +60,7 @@ impl CachedColumn { field: Arc, cache_store: Arc, column_access_path: ColumnAccessPath, + file_identity: u64, expression: Option>, is_predicate_column: bool, snapshots: Arc, @@ -79,6 +85,7 @@ impl CachedColumn { field, cache_store, column_path: column_access_path, + file_identity, expression, snapshots, } @@ -89,8 +96,13 @@ impl CachedColumn { self.column_path.entry_id(batch_id) } + pub(crate) fn file_identity(&self) -> u64 { + self.file_identity + } + pub(crate) fn contains(&self, batch_id: BatchID) -> bool { - self.cache_store.contains(&self.entry_id(batch_id).into()) + self.cache_store + .contains(&self.entry_id(batch_id).into(), self.file_identity) } pub(crate) fn snapshot_entry(&self, batch_id: BatchID) -> Option> { @@ -131,6 +143,7 @@ impl CachedColumn { self.cache_store .eval_predicate_on_entry( &entry_id, + self.file_identity, entry.as_ref(), Some(filter), &liquid_expr, @@ -139,7 +152,7 @@ impl CachedColumn { } None => { self.cache_store - .eval_predicate(&entry_id, &liquid_expr) + .eval_predicate(&entry_id, self.file_identity, &liquid_expr) .with_selection(filter) .await } @@ -191,6 +204,7 @@ impl CachedColumn { .cache_store .read_entry( &entry_id, + self.file_identity, entry.as_ref(), Some(filter), self.expression.as_deref(), @@ -198,7 +212,7 @@ impl CachedColumn { .await; } self.cache_store - .get(&entry_id) + .get(&entry_id, self.file_identity) .with_selection(filter) .with_optional_expression_hint(self.expression()) .read() @@ -208,7 +222,7 @@ impl CachedColumn { #[cfg(test)] pub(crate) async 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).await + self.cache_store.get(&entry_id, self.file_identity).await } /// Insert an array into the cache. @@ -222,7 +236,7 @@ impl CachedColumn { } self.cache_store - .insert(self.entry_id(batch_id).into(), array) + .insert(self.entry_id(batch_id).into(), self.file_identity, array) .await?; Ok(()) } @@ -239,7 +253,11 @@ impl CachedColumn { if self.snapshots.get(&entry_id).is_some() { return PrefetchOutcome::AlreadySnapshotted; } - match self.cache_store.prefetch(&entry_id).await { + match self + .cache_store + .prefetch(&entry_id, self.file_identity) + .await + { PrefetchResult::Snapshot(entry) => { self.snapshots.insert(entry_id, entry); PrefetchOutcome::Snapshotted diff --git a/src/datafusion/src/cache/id.rs b/src/datafusion/src/cache/id.rs index b52379de..84674589 100644 --- a/src/datafusion/src/cache/id.rs +++ b/src/datafusion/src/cache/id.rs @@ -51,8 +51,11 @@ const _: () = assert!(std::mem::align_of::() == 8); impl ParquetArrayID { /// Creates a new CacheEntryID. + /// + /// Only the low 16 bits of `file_id` are kept, so distinct files can + /// produce the same id. The cache tells them apart by the full file id it + /// is given alongside the entry id. 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 { @@ -146,8 +149,9 @@ pub struct ColumnAccessPath { impl ColumnAccessPath { /// Create a new instance of ColumnAccessPath. + /// + /// Only the low 16 bits of `file_id` are kept; see [`ParquetArrayID::new`]. 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 { @@ -226,9 +230,9 @@ mod tests { } #[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)); + fn test_cache_entry_id_new_keeps_low_16_bits_of_file_id() { + let wrapped = ParquetArrayID::new((u16::MAX as u64) + 6, 2, 3, BatchID::from_raw(4)); + assert_eq!(wrapped, ParquetArrayID::new(5, 2, 3, BatchID::from_raw(4))); } #[test] diff --git a/src/datafusion/src/cache/mod.rs b/src/datafusion/src/cache/mod.rs index d51bc611..e4f751ff 100644 --- a/src/datafusion/src/cache/mod.rs +++ b/src/datafusion/src/cache/mod.rs @@ -119,6 +119,7 @@ impl CachedRowGroup { Arc::clone(field), Arc::clone(&cache_store), column_access_path, + file_idx, expression.clone(), *is_predicate_column, Arc::clone(&snapshots), @@ -218,7 +219,11 @@ impl CachedRowGroup { let entry_id = column.entry_id(batch_id).into(); let liquid_array = match snapshot_liquid { Some(array) => Some(array), - None => self.cache_store.try_read_liquid(&entry_id).await, + None => { + self.cache_store + .try_read_liquid(&entry_id, column.file_identity()) + .await + } }; let liquid_array = match liquid_array { None => { @@ -928,4 +933,70 @@ mod tests { let expected = BooleanBuffer::collect_bool(batch_size, |i| i == 1 || i == 3).into(); assert_eq!(result, expected); } + + /// A cache key holds 16 bits of file id, so the 65,537th file registered + /// computes the same keys as the first. Its reads must not be served the + /// first file's data, and its own data must still be cacheable. + #[tokio::test] + async fn files_whose_keys_collide_do_not_share_entries() { + let batch_size = 64; + let tmp_dir = tempfile::tempdir().unwrap(); + let store = t4::mount(tmp_dir.path().join("liquid_cache.t4")) + .await + .unwrap(); + let cache = LiquidCacheParquet::new( + batch_size, + usize::MAX, + usize::MAX, + store, + Box::new(LiquidPolicy::new()), + Box::new(TranscodeEvict), + Box::new(AlwaysHydrate::new()), + ) + .await; + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let identity = + |path: String| ParquetFileIdentity::new(ObjectStoreUrl::local_filesystem(), path); + let batch_id = BatchID::from_row_id(0, batch_size); + + let first = cache.register_or_get_file(identity("first".to_string()), schema.clone()); + let first_column = first.create_row_group(0, vec![]).get_column(0).unwrap(); + let first_data: ArrayRef = Arc::new(Int32Array::from_iter_values(0..batch_size as i32)); + first_column + .insert(batch_id, first_data.clone()) + .await + .unwrap(); + + for i in 1..=u16::MAX as u32 { + cache.register_or_get_file(identity(format!("filler-{i}")), schema.clone()); + } + + let colliding = cache.register_or_get_file(identity("colliding".to_string()), schema); + let colliding_column = colliding.create_row_group(0, vec![]).get_column(0).unwrap(); + assert_eq!( + colliding_column.entry_id(batch_id), + first_column.entry_id(batch_id), + "the 65,537th file should compute the first file's key" + ); + + assert!( + colliding_column + .get_arrow_array_test_only(batch_id) + .await + .is_none(), + "a file was served another file's cached data" + ); + + let colliding_data: ArrayRef = + Arc::new(Int32Array::from_iter_values((0..batch_size as i32).rev())); + colliding_column + .insert(batch_id, colliding_data.clone()) + .await + .expect("a file could not cache its own data"); + let read = colliding_column + .get_arrow_array_test_only(batch_id) + .await + .unwrap(); + assert_eq!(read.as_ref(), colliding_data.as_ref()); + } } From ead34d22b07696d9be261010f9f9c047284d8036 Mon Sep 17 00:00:00 2001 From: Anoop Date: Fri, 16 Oct 2026 13:34:56 +0530 Subject: [PATCH 4/4] build(deps): keep in-process mode off the RPC and HTTP stacks (#527) Closes #522. In-process mode (`liquid-cache-datafusion-local`) no longer pulls in the client/server network stack. Its normal dependency graph drops from 387 to 343 crates; `arrow-flight`, `tonic`, `axum`, `hyper`, `h2`, `reqwest` and `rustls` are all gone. ## Changes - **`liquid-cache-common`: new `rpc` feature, off by default.** It gates `pub mod rpc` and makes `arrow-flight` and `prost` optional. `liquid-cache-datafusion-client` and `liquid-cache-datafusion-server` enable it. - **`ExecutionMetricsResponse` moves to `liquid_cache_common::metrics`.** It is a plain serde struct and the only part of `rpc` the in-process crates used. `extract_execution_metrics` stays where it is, with the same signature. - **`liquid-cache-datafusion` no longer enables `object_store`'s `http` feature.** Nothing in that crate uses it, and it was what pulled in `reqwest`/`hyper`/`h2`/`rustls`, so gating `rpc` alone would leave those in place. `liquid-cache-datafusion-client` calls `object_store::parse_url_opts` and was relying on this transitively, so it now enables `http` itself. The server, benchmark and examples already did. - **`liquid-cache-common` gets `tokio` with `macros` as a dev-dependency.** Its `#[tokio::test]`s only compiled before because `tonic` enabled that feature. - **docs.rs builds `liquid-cache-common` with `rpc`**, and its README notes the feature. `prost` remains in the graph through `vortex`; that is outside this change. ## Breaking changes - **The metrics struct has a new path.** `liquid_cache_common::rpc::ExecutionMetricsResponse` is now `liquid_cache_common::metrics::ExecutionMetricsResponse`. The serialized JSON is unchanged. - **`rpc` is now opt-in.** Code that uses `liquid_cache_common::rpc` directly, or through the `common` re-exports of `liquid-cache` / `liquid-cache-datafusion` / `liquid-cache-datafusion-local`, needs `liquid-cache-common = { ..., features = ["rpc"] }`. - **No more transitive `http`.** Users of `liquid-cache-datafusion` or `-local` who registered `http(s)://` object stores and got `object_store`'s `http` feature through this crate need to enable it themselves. ## Testing - **Workspace checks.** `cargo test --workspace --exclude dev-tools` gives 236 passed, 0 failed. The test list is identical to `main` (239 entries). `cargo clippy --workspace --all-targets --all-features -D warnings`, `cargo fmt --check`, `cargo doc` and `cargo shear` are clean. - **Each crate built on its own.** `cargo check -p --all-targets` passes for every crate. Workspace feature unification would otherwise hide a missing feature declaration. - **Standalone in-process crate.** A crate outside the workspace depending only on `liquid-cache-datafusion-local`: - It builds and runs the `examples/datafusion-local.rs` flow on `nano_hits.parquet`; the second run is served from cache and `extract_execution_metrics` reports correctly. - `cargo tree` shows none of the crates listed above. - **Standalone client crate.** A crate depending only on `liquid-cache-datafusion-client` registers `http://` and `https://` object stores. - With the client's new `http` line removed, the same build fails at runtime with `feature for Http not enabled`, so that line is required. - **End-to-end client/server.** - `bench_server` with `clickbench_client` returns correct query results. - The example client/server reads an `https://` parquet file. - The admin `/execution_metrics` endpoint returns the same JSON keys as before. --- benchmark/src/lib.rs | 2 +- src/common/Cargo.toml | 12 ++++++++++-- src/common/README.md | 2 ++ src/common/src/lib.rs | 2 ++ src/common/src/metrics.rs | 18 ++++++++++++++++++ src/common/src/rpc.rs | 17 ----------------- src/datafusion-client/Cargo.toml | 4 ++-- src/datafusion-server/Cargo.toml | 2 +- .../src/admin_server/handlers.rs | 2 +- src/datafusion-server/src/service.rs | 2 +- src/datafusion/Cargo.toml | 2 +- src/datafusion/src/utils.rs | 2 +- 12 files changed, 40 insertions(+), 27 deletions(-) create mode 100644 src/common/src/metrics.rs diff --git a/benchmark/src/lib.rs b/benchmark/src/lib.rs index a72185b0..dec471f6 100644 --- a/benchmark/src/lib.rs +++ b/benchmark/src/lib.rs @@ -8,7 +8,7 @@ use fastrace::Span; use fastrace::future::FutureExt as _; use liquid_cache::cache::CacheStats; use liquid_cache::cache::{Evict, EvictionPolicy, TranscodeEvict}; -use liquid_cache_common::rpc::ExecutionMetricsResponse; +use liquid_cache_common::metrics::ExecutionMetricsResponse; use liquid_cache_datafusion_server::{ApiResponse, ExecutionStats}; use log::info; use serde::Serialize; diff --git a/src/common/Cargo.toml b/src/common/Cargo.toml index 582f4bc9..ba612de0 100644 --- a/src/common/Cargo.toml +++ b/src/common/Cargo.toml @@ -7,14 +7,18 @@ readme = "README.md" description = { workspace = true } repository = { workspace = true } +[features] +# Arrow Flight RPC message types used between the client and the server. +rpc = ["dep:arrow-flight", "dep:prost"] + [dependencies] -arrow-flight = { workspace = true } +arrow-flight = { workspace = true, optional = true } async-trait = { workspace = true } bytes = { workspace = true } chrono = "0.4.45" futures = { workspace = true } object_store = { workspace = true } -prost = { workspace = true } +prost = { workspace = true, optional = true } serde = { workspace = true } thiserror = "2.0.18" tokio = { workspace = true } @@ -22,3 +26,7 @@ url = { workspace = true } [dev-dependencies] tempfile = { workspace = true } +tokio = { workspace = true, features = ["macros"] } + +[package.metadata.docs.rs] +features = ["rpc"] diff --git a/src/common/README.md b/src/common/README.md index 502f4227..9a2875f5 100644 --- a/src/common/README.md +++ b/src/common/README.md @@ -3,3 +3,5 @@ Common types and utilities shared across LiquidCache components. This crate contains shared data structures, enums, and utility functions used throughout the LiquidCache ecosystem. It defines the core abstractions for cache modes, schema transformations, and RPC message types that enable communication between client and server components. + +The RPC message types (`rpc` module) are behind the `rpc` feature, which is off by default. The client and server crates enable it; in-process use through `liquid-cache-datafusion-local` does not need it, and so does not depend on Arrow Flight or gRPC. diff --git a/src/common/src/lib.rs b/src/common/src/lib.rs index 5b4b43b8..57593084 100644 --- a/src/common/src/lib.rs +++ b/src/common/src/lib.rs @@ -1,5 +1,7 @@ #![doc = include_str!("../README.md")] +pub mod metrics; pub mod mock_store; +#[cfg(feature = "rpc")] pub mod rpc; pub mod utils; diff --git a/src/common/src/metrics.rs b/src/common/src/metrics.rs new file mode 100644 index 00000000..00960a77 --- /dev/null +++ b/src/common/src/metrics.rs @@ -0,0 +1,18 @@ +//! Execution metrics reported for a query plan. + +#[derive(Clone, PartialEq, serde::Serialize, serde::Deserialize, Debug)] +pub struct ExecutionMetricsResponse { + pub pushdown_eval_time: u64, + pub cache_memory_usage: u64, + pub liquid_cache_usage: u64, +} + +impl ExecutionMetricsResponse { + pub fn zero() -> Self { + Self { + pushdown_eval_time: 0, + cache_memory_usage: 0, + liquid_cache_usage: 0, + } + } +} diff --git a/src/common/src/rpc.rs b/src/common/src/rpc.rs index 452d4ef5..0907085a 100644 --- a/src/common/src/rpc.rs +++ b/src/common/src/rpc.rs @@ -246,20 +246,3 @@ impl ProstMessageExt for FetchResults { } } } - -#[derive(Clone, PartialEq, serde::Serialize, serde::Deserialize, Debug)] -pub struct ExecutionMetricsResponse { - pub pushdown_eval_time: u64, - pub cache_memory_usage: u64, - pub liquid_cache_usage: u64, -} - -impl ExecutionMetricsResponse { - pub fn zero() -> Self { - Self { - pushdown_eval_time: 0, - cache_memory_usage: 0, - liquid_cache_usage: 0, - } - } -} diff --git a/src/datafusion-client/Cargo.toml b/src/datafusion-client/Cargo.toml index ddd207cd..42e04186 100644 --- a/src/datafusion-client/Cargo.toml +++ b/src/datafusion-client/Cargo.toml @@ -17,13 +17,13 @@ arrow-schema = { workspace = true } futures = { workspace = true } tonic = { workspace = true } uuid = { workspace = true } -liquid-cache-common = { workspace = true } +liquid-cache-common = { workspace = true, features = ["rpc"] } liquid-cache-datafusion = { workspace = true } fastrace = { workspace = true } fastrace-tonic = { workspace = true } tower = "0.5.3" tokio = { workspace = true } -object_store = { workspace = true } +object_store = { workspace = true, features = ["http"] } [dev-dependencies] insta = { workspace = true } diff --git a/src/datafusion-server/Cargo.toml b/src/datafusion-server/Cargo.toml index 804b48dc..a391c48f 100644 --- a/src/datafusion-server/Cargo.toml +++ b/src/datafusion-server/Cargo.toml @@ -21,7 +21,7 @@ tokio = { workspace = true } url = { workspace = true } liquid-cache-datafusion = { workspace = true } object_store = { workspace = true, features = ["aws", "http"] } -liquid-cache-common = { workspace = true } +liquid-cache-common = { workspace = true, features = ["rpc"] } tempfile = { workspace = true } axum = "0.8.9" serde = { workspace = true } diff --git a/src/datafusion-server/src/admin_server/handlers.rs b/src/datafusion-server/src/admin_server/handlers.rs index 029f928a..4614fe96 100644 --- a/src/datafusion-server/src/admin_server/handlers.rs +++ b/src/datafusion-server/src/admin_server/handlers.rs @@ -17,7 +17,7 @@ use datafusion::{ datasource::physical_plan::FileScanConfig, physical_plan::{ExecutionPlan, StatisticsArgs, StatisticsContext}, }; -use liquid_cache_common::rpc::ExecutionMetricsResponse; +use liquid_cache_common::metrics::ExecutionMetricsResponse; use liquid_cache_datafusion::LiquidParquetSource; use log::info; use serde::Serialize; diff --git a/src/datafusion-server/src/service.rs b/src/datafusion-server/src/service.rs index 04c67124..479daca0 100644 --- a/src/datafusion-server/src/service.rs +++ b/src/datafusion-server/src/service.rs @@ -7,7 +7,7 @@ use datafusion::{ }; use liquid_cache::{ByteCache, cache::EvictionPolicy}; use liquid_cache::{cache::HydrationPolicy, cache_policies::CachePolicy}; -use liquid_cache_common::rpc::ExecutionMetricsResponse; +use liquid_cache_common::metrics::ExecutionMetricsResponse; use liquid_cache_datafusion::{ cache::{ColumnLineages, LiquidCacheParquet, LiquidCacheParquetRef}, extract_execution_metrics, diff --git a/src/datafusion/Cargo.toml b/src/datafusion/Cargo.toml index 53e89258..76f1cbdf 100644 --- a/src/datafusion/Cargo.toml +++ b/src/datafusion/Cargo.toml @@ -17,7 +17,7 @@ futures = { workspace = true } tokio = { workspace = true } ahash = { workspace = true } log = { workspace = true } -object_store = { workspace = true, features = ["http"] } +object_store = { workspace = true } liquid-cache-common = { workspace = true } liquid-cache = { workspace = true } fastrace = { workspace = true } diff --git a/src/datafusion/src/utils.rs b/src/datafusion/src/utils.rs index e6d903cc..5d5b3aad 100644 --- a/src/datafusion/src/utils.rs +++ b/src/datafusion/src/utils.rs @@ -9,7 +9,7 @@ use datafusion::{ datasource::source::DataSourceExec, physical_plan::{ExecutionPlan, metrics::MetricValue}, }; -use liquid_cache_common::rpc::ExecutionMetricsResponse; +use liquid_cache_common::metrics::ExecutionMetricsResponse; use parquet::arrow::arrow_reader::RowSelector; use crate::cache::LiquidCacheParquetRef;