From 8a579025b6eba63e748a0dd4b5a908818e43c103 Mon Sep 17 00:00:00 2001 From: G Date: Tue, 16 Jun 2026 19:18:17 +0530 Subject: [PATCH 01/18] Step 1a: unified scoped parquet page-index cache (pure module) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add indexed_table::page_index_loader — a process-global, byte-bounded LRU that loads the parquet page index scoped to a query's predicate columns and is shared by both scan paths (DataFusion ListingTable + custom indexed executor). No scan-path wiring yet; this lands the cache + its unit tests in isolation. Design: - ColumnIndex scoped to predicate columns (NONE placeholders elsewhere); OffsetIndex kept for all columns / all row groups (read at scan time for every projected column — placeholders there break reads). - Cache stores the decoded (ColumnIndex, OffsetIndex) PAIR + a structural size estimate, NOT a full ParquetMetaData, and grafts the pair onto the caller's resident footer at lookup. No footer duplication / over-allocation. - Key is (path, predicate parquet-cols) ONLY — both paths build the identical all-RG artifact, so entries are shared across paths. Row-group scoping is a future axis (gated Step 2/3) and is deliberately not in the key yet. - Cache hit returns a fresh Arc (graft), so hits are asserted via the stats counter, not Arc::ptr_eq. Uses the deprecated read_columns_indexes/read_offset_indexes — the only public subset-decode API today; arrow-rs#8643/#8797/#8714 track the first-class replacement. 9 unit tests pass (cargo test --lib page_index_loader --test-threads=1). --- .../rust/src/indexed_table/mod.rs | 1 + .../src/indexed_table/page_index_loader.rs | 1028 +++++++++++++++++ 2 files changed, 1029 insertions(+) create mode 100644 sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/page_index_loader.rs diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/mod.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/mod.rs index a2c5c5eab6ade..05ba3acc17773 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/mod.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/mod.rs @@ -62,6 +62,7 @@ pub mod row_id_injection; pub mod ffm_callbacks; pub mod index; pub mod metrics; +pub mod page_index_loader; pub mod page_pruner; pub mod parquet_bridge; pub mod partitioning; diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/page_index_loader.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/page_index_loader.rs new file mode 100644 index 0000000000000..4ffa50677f84d --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/page_index_loader.rs @@ -0,0 +1,1028 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +//! Unified, scoped parquet page-index cache (Step 1a). +//! +//! # Why this exists +//! +//! Parquet metadata loading pulls the **entire page index** — `ColumnIndex` +//! (per-page min/max; the per-page *string* min/max is the heap hog) plus +//! `OffsetIndex` (per-page byte offsets), for every column of every row group. +//! On wide schemas (the production `textbench` index has 402 columns) this is +//! ~82% of native heap and is re-decoded per query, even when the query filters +//! a single column. +//! +//! This module loads the page index **scoped to the query's predicate columns** +//! and caches the result in **one** process-global cache that is shared by both +//! scan paths (the DataFusion `ListingTable` path and the custom indexed-table +//! executor). For the same `(file, predicate columns)` an entry built by one +//! path is reusable by the other. +//! +//! ## What is scoped, and why each scope is safe +//! +//! - **`ColumnIndex` → predicate columns only.** It is read only at *prune* +//! time, and the pruner (`StatisticsConverter`, and our [`super::page_pruner`]) +//! only ever dereferences predicate-column positions. Non-predicate positions +//! carry [`ColumnIndexMetaData::NONE`] placeholders, which keep absolute +//! `index[rg][col]` indexing valid and cost nothing. This scoping is where the +//! heap win comes from. +//! - **`OffsetIndex` → all columns (this step also: all row groups).** Unlike +//! the `ColumnIndex`, the `OffsetIndex` is read at *scan* time: when a +//! `RowSelection` is active, arrow-rs's `InMemoryRowGroup::fetch_ranges` +//! dereferences `offset_index[col].page_locations()` for every **projected** +//! column (by absolute index) to compute which page byte ranges to fetch. A +//! placeholder there for a projected column fetches zero ranges and the read +//! fails ("failed to skip rows, expected N, got 0"). Because this loader runs +//! before the projection is known, it keeps a real `OffsetIndex` for every +//! column. That is cheap relative to the `ColumnIndex`: fixed-width page +//! offsets/sizes, no per-page string stats. +//! +//! ## Why the cache stores the (ColumnIndex, OffsetIndex) PAIR, not metadata +//! +//! `ParquetMetaData` owns its `row_groups: Vec` (~5–6 MB on +//! the 402-column schema). Caching a full augmented `ParquetMetaData` per entry +//! would **duplicate that footer** in every entry (the prior iteration's +//! over-allocation bug: a row-group-pruned entry that read 0 row groups still +//! cost a 60 MB footer clone). Instead this cache stores only the decoded +//! `(ParquetColumnIndex, ParquetOffsetIndex)` pair plus a size estimate, and +//! **grafts** the pair onto the caller's already-resident footer at lookup via +//! [`ParquetMetaData::into_builder`] → `set_column_index`/`set_offset_index`. +//! The footer is never duplicated into the cache; the only footer copy is the +//! one transient graft handed to the scan (one deep clone per file per query, +//! the same order as a cold metadata fetch). +//! +//! **Consequence for tests:** a cache hit returns a *fresh* `Arc`, so +//! `Arc::ptr_eq` is the wrong signal for "served from cache" — assert via the +//! [`ScopedCacheStats::hits`] counter instead. +//! +//! ## Key +//! +//! The cache key is `(file path, predicate parquet-column indices)` **only**. +//! Both scan paths resolve predicate columns the same way +//! ([`resolve_predicate_parquet_columns`]) and build the identical artifact +//! (all-RG, all-column `OffsetIndex`; predicate-scoped `ColumnIndex`), so an +//! entry is shared across paths. Row-group scoping is a *future* axis that would +//! extend the key — deliberately omitted here so Step 1 keeps one unified key. +//! +//! ## Correctness / fallback +//! +//! Any failure (file has no page index, a predicate column lacks an index range, +//! a decode/IO error) makes the load return `None`. The caller keeps its +//! footer-only metadata and the pruner conservatively no-ops (scans the whole +//! row group) — never a wrong result. +//! +//! ## Upstream note +//! +//! arrow-rs is moving toward first-class selective metadata decoding: a unified +//! options mechanism that would include "decode column indexes only for +//! predicate columns" and "row-group selection" (apache/arrow-rs#8643, open), +//! and the `ParquetMetaDataOptions` / `ParquetStatisticsPolicy::skip_except` +//! pattern just merged for page *encoding* statistics (apache/arrow-rs#8797, +//! built on the selective-decode metadata "index" of #8714). None of these yet +//! expose a page-index column/row-group projection. Until one does, the only +//! **public** API that decodes a *subset* of columns is the deprecated +//! [`read_columns_indexes`]/[`read_offset_indexes`] (the per-column primitives +//! are `pub(crate)`; `ParquetMetaDataReader`/`PageIndexPolicy` is +//! all-columns-or-nothing). This module is the interim hand-rolled equivalent; +//! when a `ParquetMetaDataOptions` page-index projection lands, migrate +//! [`build_scoped_page_index`] to it and drop the `#[allow(deprecated)]`. + +use std::collections::HashMap; +use std::ops::Range; +use std::sync::{Arc, Mutex}; + +use once_cell::sync::Lazy; + +use arrow::datatypes::SchemaRef; +use datafusion::parquet::arrow::arrow_reader::statistics::StatisticsConverter; +use datafusion::parquet::errors::{ParquetError, Result as ParquetResult}; +use datafusion::parquet::file::metadata::{ + ColumnChunkMetaData, OffsetIndexBuilder, ParquetColumnIndex, ParquetMetaData, + ParquetOffsetIndex, +}; +use datafusion::parquet::file::page_index::column_index::ColumnIndexMetaData; +use datafusion::parquet::file::page_index::index_reader::{ + read_columns_indexes, read_offset_indexes, +}; +use datafusion::parquet::file::reader::{ChunkReader, Length}; +use object_store::ObjectStore; +use prost::bytes::{Buf, Bytes}; + +/// Default byte budget for the scoped page-index cache, used until the caller +/// sets one from the runtime's configured limit (see [`set_scoped_cache_limit`]). +/// 64 MiB is generous for a predicate-column-only page index yet a tiny fraction +/// of the footer-only baseline the page-index strip buys back. +const DEFAULT_SCOPED_CACHE_LIMIT: usize = 64 * 1024 * 1024; + +// ── Cache types ────────────────────────────────────────────────────────── + +/// Cache key: object-store path + the sorted/deduped set of parquet column +/// indices the page index was built for. Row-group scoping is intentionally +/// NOT part of the key in Step 1 — both paths build an all-row-group artifact so +/// the key (and the artifact) is identical across paths. +#[derive(Clone, PartialEq, Eq, Hash, Debug)] +struct ScopedKey { + path: String, + parquet_cols: Vec, +} + +/// One cached entry: the decoded page-index pair (full-width over all row groups +/// and all columns; `ColumnIndex` real only at predicate positions), a size +/// estimate, and a last-used tick for LRU ordering. Deliberately does NOT hold a +/// `ParquetMetaData` — see the module docs (no footer duplication). +struct ScopedEntry { + column_index: ParquetColumnIndex, + offset_index: ParquetOffsetIndex, + size: usize, + last_used: u64, +} + +/// Snapshot of scoped page-index cache counters. Surfaced on node-stats and used +/// by tests to assert hits/misses without relying on `Arc::ptr_eq`. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub struct ScopedCacheStats { + pub hits: u64, + pub misses: u64, + pub evictions: u64, + pub entries: usize, + pub used_bytes: usize, + pub limit_bytes: usize, +} + +/// Byte-bounded LRU over scoped page-index pairs, keyed by +/// `(file, predicate-column-set)`. +/// +/// Evicts the least-recently-used entries once `used > limit`, so it always +/// serves cached data within a memory budget rather than silently degrading to +/// "decode every query" when full. Entry size is a deterministic structural +/// estimate (see [`scoped_page_index_size`]). +struct ScopedLru { + map: HashMap, + used: usize, + limit: usize, + /// Monotonic clock for LRU ordering (avoids `Instant`; fine single-process). + tick: u64, + hits: u64, + misses: u64, + evictions: u64, +} + +impl ScopedLru { + fn new(limit: usize) -> Self { + Self { + map: HashMap::new(), + used: 0, + limit, + tick: 0, + hits: 0, + misses: 0, + evictions: 0, + } + } + + fn next_tick(&mut self) -> u64 { + self.tick += 1; + self.tick + } + + /// On hit, bump recency and return a clone of the pair (cheap relative to a + /// footer: predicate-scoped `ColumnIndex` + fixed-width `OffsetIndex`). + fn get(&mut self, key: &ScopedKey) -> Option<(ParquetColumnIndex, ParquetOffsetIndex)> { + let t = self.next_tick(); + match self.map.get_mut(key) { + Some(entry) => { + entry.last_used = t; + self.hits += 1; + Some((entry.column_index.clone(), entry.offset_index.clone())) + } + None => { + self.misses += 1; + None + } + } + } + + fn insert( + &mut self, + key: ScopedKey, + column_index: ParquetColumnIndex, + offset_index: ParquetOffsetIndex, + size: usize, + ) { + // An entry larger than the whole budget can never be retained; skip it + // rather than evicting everything else for something we'd drop anyway. + if size > self.limit { + return; + } + let t = self.next_tick(); + if let Some(old) = self.map.insert( + key, + ScopedEntry { + column_index, + offset_index, + size, + last_used: t, + }, + ) { + self.used -= old.size; + } + self.used += size; + self.evict(); + } + + fn evict(&mut self) { + while self.used > self.limit { + let Some(victim) = self + .map + .iter() + .min_by_key(|(_, e)| e.last_used) + .map(|(k, _)| k.clone()) + else { + break; + }; + if let Some(removed) = self.map.remove(&victim) { + self.used -= removed.size; + self.evictions += 1; + } + } + } + + fn stats(&self) -> ScopedCacheStats { + ScopedCacheStats { + hits: self.hits, + misses: self.misses, + evictions: self.evictions, + entries: self.map.len(), + used_bytes: self.used, + limit_bytes: self.limit, + } + } + + fn set_limit(&mut self, limit: usize) { + self.limit = limit; + self.evict(); + } +} + +/// Process-wide scoped page-index cache. +static SCOPED_CACHE: Lazy> = + Lazy::new(|| Mutex::new(ScopedLru::new(DEFAULT_SCOPED_CACHE_LIMIT))); + +/// Set the scoped cache's byte budget. Called from startup wiring with the +/// configured limit. Idempotent; shrinking evicts immediately. Zero is ignored +/// (keeps the existing budget). +pub fn set_scoped_cache_limit(limit: usize) { + if limit == 0 { + return; + } + if let Ok(mut c) = SCOPED_CACHE.lock() { + c.set_limit(limit); + } +} + +/// Snapshot of the scoped page-index cache counters plus occupancy. For +/// node-stats / observability and tests. +pub fn scoped_cache_stats() -> ScopedCacheStats { + SCOPED_CACHE.lock().map(|c| c.stats()).unwrap_or_default() +} + +// ── Public API ───────────────────────────────────────────────────────────── + +/// Map the query's arrow predicate-column names to this file's parquet column +/// indices, using the same resolution the pruner uses +/// (`StatisticsConverter::parquet_column_index`). Columns absent from the +/// parquet file (schema evolution) are skipped. Returns a sorted, deduped set so +/// both scan paths produce an identical key for the same logical predicate. +pub fn resolve_predicate_parquet_columns( + arrow_schema: &SchemaRef, + metadata: &ParquetMetaData, + predicate_column_names: &[String], +) -> Vec { + let parquet_schema = metadata.file_metadata().schema_descr(); + let mut set = std::collections::BTreeSet::new(); + for name in predicate_column_names { + if let Ok(conv) = StatisticsConverter::try_new(name, arrow_schema, parquet_schema) { + if let Some(idx) = conv.parquet_column_index() { + set.insert(idx); + } + } + } + set.into_iter().collect() +} + +/// Load (or build + cache) the scoped page index for `parquet_cols` and graft it +/// onto `footer_meta`, returning fresh metadata that carries: +/// - a real `ColumnIndex` at the predicate-column positions (NONE elsewhere), +/// - a real `OffsetIndex` for every column, across all row groups. +/// +/// Returns `None` (caller keeps footer-only metadata) on any condition that +/// would make page pruning unsafe or impossible: empty column set, a file +/// without a page index, an out-of-range predicate column, or a decode/IO error. +/// +/// Consults and populates the shared `(file, predicate-cols)` cache. +pub async fn load_scoped_page_index( + store: &Arc, + location: &object_store::path::Path, + footer_meta: &Arc, + parquet_cols: &[usize], +) -> Option> { + if parquet_cols.is_empty() { + return None; + } + let key = ScopedKey { + path: location.as_ref().to_string(), + parquet_cols: parquet_cols.to_vec(), + }; + + // Cache hit: graft the cached pair onto the caller's footer. No I/O, no + // decode. (Returns a fresh Arc — assert hits via the counter, not ptr_eq.) + if let Ok(mut cache) = SCOPED_CACHE.lock() { + if let Some((ci, oi)) = cache.get(&key) { + return Some(graft(footer_meta, ci, oi)); + } + } + + // Miss: range-read + decode the scoped page index. + let (column_index, offset_index, size) = + build_scoped_page_index(store, location, footer_meta, parquet_cols).await?; + + // Graft a copy for the caller; store the originals in the cache. + let grafted = graft(footer_meta, column_index.clone(), offset_index.clone()); + if let Ok(mut cache) = SCOPED_CACHE.lock() { + cache.insert(key, column_index, offset_index, size); + } + Some(grafted) +} + +/// Build a fresh `ParquetMetaData` = `footer` with the scoped page-index pair +/// grafted on. This deep-clones the footer (the builder consumes an owned +/// `ParquetMetaData`); that transient clone is the only footer copy and is never +/// stored in the cache. +fn graft( + footer_meta: &Arc, + column_index: ParquetColumnIndex, + offset_index: ParquetOffsetIndex, +) -> Arc { + let base = ParquetMetaData::clone(footer_meta); + let rebuilt = base + .into_builder() + .set_column_index(Some(column_index)) + .set_offset_index(Some(offset_index)) + .build(); + Arc::new(rebuilt) +} + +// ── Build ──────────────────────────────────────────────────────────────── + +/// Range-read and decode the page index scoped to `parquet_cols`, returning the +/// full-width `(ColumnIndex, OffsetIndex)` pair plus a size estimate. All row +/// groups, all columns for the `OffsetIndex`; predicate columns only (NONE +/// elsewhere) for the `ColumnIndex`. `None` on any unsafe/impossible condition. +async fn build_scoped_page_index( + store: &Arc, + location: &object_store::path::Path, + footer_meta: &Arc, + parquet_cols: &[usize], +) -> Option<(ParquetColumnIndex, ParquetOffsetIndex, usize)> { + let num_rgs = footer_meta.num_row_groups(); + if num_rgs == 0 { + return None; + } + let num_cols = footer_meta.file_metadata().schema_descr().num_columns(); + // Predicate indices index into the file schema (shared across RGs). + if parquet_cols.iter().any(|&i| i >= num_cols) { + return None; + } + + // Phase 1: per RG, gather the predicate columns' ColumnIndex chunks and ALL + // columns' OffsetIndex chunks, and compute their union byte ranges for a + // single vectored fetch. Bail to footer-only if any required index range is + // missing (the file has no page index for it). + struct RgPlan { + rg_idx: usize, + pred_chunks: Vec, + all_chunks: Vec, + col_range: Range, + off_range: Range, + } + let mut plans: Vec = Vec::with_capacity(num_rgs); + let mut fetch_ranges: Vec> = Vec::with_capacity(num_rgs * 2); + for rg_idx in 0..num_rgs { + let rg = footer_meta.row_group(rg_idx); + let pred_chunks: Vec = + parquet_cols.iter().map(|&i| rg.column(i).clone()).collect(); + let all_chunks: Vec = + (0..num_cols).map(|i| rg.column(i).clone()).collect(); + let col_range = column_index_union(&pred_chunks)?; + let off_range = offset_index_union(&all_chunks)?; + fetch_ranges.push(col_range.clone()); + fetch_ranges.push(off_range.clone()); + plans.push(RgPlan { + rg_idx, + pred_chunks, + all_chunks, + col_range, + off_range, + }); + } + if plans.is_empty() { + return None; + } + + // Phase 2: one vectored fetch of all RGs' index byte ranges. + let buffers = store.get_ranges(location, &fetch_ranges).await.ok()?; + if buffers.len() != fetch_ranges.len() { + return None; + } + + // Phase 3: decode and scatter into full-width per-RG vectors. Pre-fill with + // placeholders so absolute `index[rg][col]` indexing is always valid. + let mut column_index: ParquetColumnIndex = (0..num_rgs) + .map(|_| (0..num_cols).map(|_| ColumnIndexMetaData::NONE).collect()) + .collect(); + let mut offset_index: ParquetOffsetIndex = (0..num_rgs) + .map(|_| { + (0..num_cols) + .map(|_| OffsetIndexBuilder::new().build()) + .collect() + }) + .collect(); + + for (i, plan) in plans.iter().enumerate() { + let col_buf = buffers[i * 2].clone(); + let off_buf = buffers[i * 2 + 1].clone(); + + let col_reader = BufferChunkReader { + base: plan.col_range.start, + bytes: col_buf, + }; + let off_reader = BufferChunkReader { + base: plan.off_range.start, + bytes: off_buf, + }; + + // `read_columns_indexes` / `read_offset_indexes` are deprecated in + // arrow-rs but are the only PUBLIC API that decodes a *column subset*. + // See module docs + apache/arrow-rs#8643. + #[allow(deprecated)] + let decoded_cols = read_columns_indexes(&col_reader, &plan.pred_chunks).ok()??; + #[allow(deprecated)] + let decoded_offs = read_offset_indexes(&off_reader, &plan.all_chunks).ok()??; + if decoded_cols.len() != parquet_cols.len() || decoded_offs.len() != num_cols { + return None; + } + + let col_row = &mut column_index[plan.rg_idx]; + for (k, &parquet_col) in parquet_cols.iter().enumerate() { + col_row[parquet_col] = decoded_cols[k].clone(); + } + offset_index[plan.rg_idx] = decoded_offs; + } + + let size = scoped_page_index_size(footer_meta, parquet_cols, num_cols); + Some((column_index, offset_index, size)) +} + +/// Deterministic size estimate for one cached pair, in bytes. +/// +/// Uses the **on-disk serialized lengths** from the footer +/// (`column_index_length` for predicate columns + `offset_index_length` for all +/// columns, summed over all row groups). This is a robust, monotonic proxy for +/// the decoded heap: it is never zero for a real page index, grows with the +/// number of predicate columns / row groups / pages, and — unlike +/// `ParquetMetaData::memory_size()` (which includes the footer we deliberately +/// don't store and rounds to ~0 on small files) — does not depend on private +/// `HeapSize` internals or on arrow-rs decoded-struct layout. The decoded form +/// is a small constant multiple of this; the configured limit is interpreted in +/// these units. +fn scoped_page_index_size( + footer_meta: &ParquetMetaData, + parquet_cols: &[usize], + num_cols: usize, +) -> usize { + let mut total = 0usize; + for rg in footer_meta.row_groups() { + for &pc in parquet_cols { + total += rg.column(pc).column_index_length().unwrap_or(0).max(0) as usize; + } + for c in 0..num_cols { + total += rg.column(c).offset_index_length().unwrap_or(0).max(0) as usize; + } + } + total +} + +/// Union of `column_index` byte ranges across the given column chunks. `None` if +/// any chunk lacks a column index (we require all predicate columns to have one, +/// else we fall back to footer-only). +fn column_index_union(chunks: &[ColumnChunkMetaData]) -> Option> { + range_union(chunks, |c| { + let off = u64::try_from(c.column_index_offset()?).ok()?; + let len = u64::try_from(c.column_index_length()?).ok()?; + Some(off..off + len) + }) +} + +/// Union of `offset_index` byte ranges across the given column chunks. +fn offset_index_union(chunks: &[ColumnChunkMetaData]) -> Option> { + range_union(chunks, |c| { + let off = u64::try_from(c.offset_index_offset()?).ok()?; + let len = u64::try_from(c.offset_index_length()?).ok()?; + Some(off..off + len) + }) +} + +fn range_union( + chunks: &[ColumnChunkMetaData], + f: impl Fn(&ColumnChunkMetaData) -> Option>, +) -> Option> { + let mut acc: Option> = None; + for c in chunks { + let r = f(c)?; // any missing range → bail (caller falls back) + acc = Some(match acc { + None => r, + Some(a) => a.start.min(r.start)..a.end.max(r.end), + }); + } + acc +} + +/// A [`ChunkReader`] over an in-memory byte buffer representing the file region +/// `[base, base + bytes.len())`. The arrow-rs page-index readers call +/// `get_bytes(absolute_offset, len)`; we translate the absolute file offset into +/// the buffer. +struct BufferChunkReader { + /// Absolute file offset of `bytes[0]`. + base: u64, + bytes: Bytes, +} + +impl Length for BufferChunkReader { + fn len(&self) -> u64 { + self.base + self.bytes.len() as u64 + } +} + +impl ChunkReader for BufferChunkReader { + type T = prost::bytes::buf::Reader; + + fn get_read(&self, start: u64) -> ParquetResult { + let rel = self.rel(start, 0)?; + Ok(self.bytes.slice(rel..).reader()) + } + + fn get_bytes(&self, start: u64, length: usize) -> ParquetResult { + let rel = self.rel(start, length)?; + Ok(self.bytes.slice(rel..rel + length)) + } +} + +impl BufferChunkReader { + /// Translate an absolute file offset (and a length to validate) into a + /// buffer-relative start, erroring if it falls outside the prefetched span. + fn rel(&self, start: u64, length: usize) -> ParquetResult { + let rel = start.checked_sub(self.base).ok_or_else(|| { + ParquetError::General(format!( + "page-index read offset {} precedes buffer base {}", + start, self.base + )) + })?; + let rel = usize::try_from(rel) + .map_err(|e| ParquetError::General(format!("offset overflow: {}", e)))?; + if rel + length > self.bytes.len() { + return Err(ParquetError::General(format!( + "page-index read [{}..{}) exceeds buffer of len {}", + rel, + rel + length, + self.bytes.len() + ))); + } + Ok(rel) + } +} + +// ── Test-only helpers (the cache is process-global; see SCOPED_CACHE_TEST_GUARD) ── + +/// Crate-wide guard so every test that touches the global scoped cache mutually +/// excludes — distinct fixtures alone aren't enough because the `InMemory` path +/// is always "data.parquet". Shared (not per-module) so future cache users +/// (`shard_scoped_reader`, the optimizer) serialize against the same lock. +#[cfg(test)] +pub(crate) static SCOPED_CACHE_TEST_GUARD: std::sync::Mutex<()> = std::sync::Mutex::new(()); + +/// Reset entries + counters AND restore the default limit. +#[cfg(test)] +pub(crate) fn clear_scoped_cache_for_test() { + if let Ok(mut c) = SCOPED_CACHE.lock() { + c.map.clear(); + c.used = 0; + c.tick = 0; + c.limit = DEFAULT_SCOPED_CACHE_LIMIT; + c.hits = 0; + c.misses = 0; + c.evictions = 0; + } +} + +#[cfg(test)] +pub(crate) fn scoped_cache_len_for_test() -> usize { + SCOPED_CACHE.lock().map(|c| c.map.len()).unwrap_or(0) +} + +#[cfg(test)] +pub(crate) fn scoped_cache_bytes_for_test() -> usize { + SCOPED_CACHE.lock().map(|c| c.used).unwrap_or(0) +} + +#[cfg(test)] +pub(crate) fn set_scoped_cache_limit_for_test(limit: usize) { + if let Ok(mut c) = SCOPED_CACHE.lock() { + c.set_limit(limit); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::indexed_table::page_pruner::{build_pruning_predicate, PagePruner}; + use arrow::array::{Int32Array, RecordBatch}; + use arrow::datatypes::{DataType, Field, Schema}; + use datafusion::common::ScalarValue; + use datafusion::logical_expr::Operator; + use datafusion::parquet::arrow::arrow_reader::{ + ArrowReaderMetadata, ArrowReaderOptions, RowSelection, RowSelector, + }; + use datafusion::parquet::arrow::ArrowWriter; + use datafusion::parquet::file::properties::{EnabledStatistics, WriterProperties}; + use datafusion::physical_expr::expressions::{BinaryExpr, Column as PhysColumn, Literal}; + use datafusion::physical_expr::PhysicalExpr; + use object_store::memory::InMemory; + use object_store::path::Path as ObjPath; + use object_store::{ObjectStoreExt, PutPayload}; + + // Aliased so every cache-touching test serializes on one lock. + use super::SCOPED_CACHE_TEST_GUARD as CACHE_TEST_GUARD; + + // ── fixtures + expr helpers ────────────────────────────────────────── + + /// 2 columns (`price`, `qty`), 32 rows, 1 row group, 4 pages of 8 rows. + fn two_col_parquet() -> (Bytes, SchemaRef) { + let schema = Arc::new(Schema::new(vec![ + Field::new("price", DataType::Int32, false), + Field::new("qty", DataType::Int32, false), + ])); + let prices: Vec = (0..32).collect(); + let qtys: Vec = (100..132).collect(); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(prices)), + Arc::new(Int32Array::from(qtys)), + ], + ) + .unwrap(); + let props = WriterProperties::builder() + .set_max_row_group_size(32) + .set_data_page_row_count_limit(8) + .set_write_batch_size(8) + .set_statistics_enabled(EnabledStatistics::Page) + .build(); + let mut buf: Vec = Vec::new(); + let mut w = ArrowWriter::try_new(&mut buf, schema.clone(), Some(props)).unwrap(); + w.write(&batch).unwrap(); + w.close().unwrap(); + (Bytes::from(buf), schema) + } + + async fn stage(bytes: Bytes) -> (Arc, ObjPath) { + let store: Arc = Arc::new(InMemory::new()); + let loc = ObjPath::from("data.parquet"); + store.put(&loc, PutPayload::from_bytes(bytes)).await.unwrap(); + (store, loc) + } + + fn footer_only(bytes: &Bytes) -> Arc { + ArrowReaderMetadata::load( + &bytes.clone(), + ArrowReaderOptions::new().with_page_index(false), + ) + .unwrap() + .metadata() + .clone() + } + + fn full_index(bytes: &Bytes) -> Arc { + ArrowReaderMetadata::load( + &bytes.clone(), + ArrowReaderOptions::new().with_page_index(true), + ) + .unwrap() + .metadata() + .clone() + } + + fn col(name: &str, idx: usize) -> Arc { + Arc::new(PhysColumn::new(name, idx)) + } + fn lit_int(v: i32) -> Arc { + Arc::new(Literal::new(ScalarValue::Int32(Some(v)))) + } + fn pred(name: &str, idx: usize, op: Operator, v: i32) -> Arc { + Arc::new(BinaryExpr::new(col(name, idx), op, lit_int(v))) + } + fn kept(sel: &RowSelection) -> usize { + sel.iter().filter(|s| !s.skip).map(|s| s.row_count).sum() + } + + /// Read the rows kept by `selection` for one projected leaf column — drives + /// the offset-index read path so we can prove non-predicate projected reads + /// succeed with scoped metadata. + fn read_selected_column( + bytes: &Bytes, + meta: &Arc, + leaf_col: usize, + selection: RowSelection, + ) -> std::result::Result, String> { + use datafusion::parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; + use datafusion::parquet::arrow::ProjectionMask; + + let arm = ArrowReaderMetadata::try_new(Arc::clone(meta), ArrowReaderOptions::new()) + .map_err(|e| format!("try_new metadata: {e}"))?; + let builder = ParquetRecordBatchReaderBuilder::new_with_metadata(bytes.clone(), arm); + let proj = ProjectionMask::leaves(builder.parquet_schema(), [leaf_col]); + let mut reader = builder + .with_row_groups(vec![0]) + .with_projection(proj) + .with_row_selection(selection) + .build() + .map_err(|e| format!("build reader: {e}"))?; + + let mut out = Vec::new(); + while let Some(next) = reader.next() { + let batch = next.map_err(|e| format!("read batch: {e}"))?; + let a = batch + .column(0) + .as_any() + .downcast_ref::() + .ok_or("projected column was not Int32")?; + for i in 0..a.len() { + out.push(a.value(i)); + } + } + Ok(out) + } + + // ── tests ──────────────────────────────────────────────────────────── + + /// Footer-only load carries no page index — the baseline the cache restores. + #[tokio::test] + async fn footer_only_has_no_page_index() { + let (bytes, _schema) = two_col_parquet(); + let fo = footer_only(&bytes); + assert!(fo.column_index().is_none()); + assert!(fo.offset_index().is_none()); + } + + /// Empty predicate-column set never builds an entry. + #[tokio::test] + async fn empty_column_set_returns_none() { + let _g = CACHE_TEST_GUARD.lock().unwrap(); + clear_scoped_cache_for_test(); + let (bytes, _schema) = two_col_parquet(); + let (store, loc) = stage(bytes.clone()).await; + let fo = footer_only(&bytes); + assert!(load_scoped_page_index(&store, &loc, &fo, &[]).await.is_none()); + assert_eq!(scoped_cache_len_for_test(), 0); + } + + /// The grafted metadata has a real ColumnIndex only at predicate positions + /// (NONE elsewhere) and a real OffsetIndex for every column. + #[tokio::test] + async fn scoped_index_is_predicate_scoped_for_column_index() { + let _g = CACHE_TEST_GUARD.lock().unwrap(); + clear_scoped_cache_for_test(); + let (bytes, schema) = two_col_parquet(); + let (store, loc) = stage(bytes.clone()).await; + let fo = footer_only(&bytes); + + let cols = resolve_predicate_parquet_columns(&schema, &fo, &["price".to_string()]); + assert_eq!(cols, vec![0]); + + let aug = load_scoped_page_index(&store, &loc, &fo, &cols) + .await + .expect("augmentation must succeed"); + let ci = aug.column_index().expect("has column index"); + let oi = aug.offset_index().expect("has offset index"); + + assert!( + !matches!(ci[0][0], ColumnIndexMetaData::NONE), + "predicate col (price) must have a real ColumnIndex" + ); + assert!( + matches!(ci[0][1], ColumnIndexMetaData::NONE), + "non-predicate col (qty) ColumnIndex must be a NONE placeholder" + ); + assert!( + !oi[0][0].page_locations().is_empty() && !oi[0][1].page_locations().is_empty(), + "OffsetIndex must be real for every column" + ); + } + + /// Scoped pruning must match full-index pruning on the predicate column. + #[tokio::test] + async fn scoped_pruning_matches_full_index() { + let _g = CACHE_TEST_GUARD.lock().unwrap(); + clear_scoped_cache_for_test(); + let (bytes, schema) = two_col_parquet(); + let (store, loc) = stage(bytes.clone()).await; + let fo = footer_only(&bytes); + + let cols = resolve_predicate_parquet_columns(&schema, &fo, &["price".to_string()]); + let aug = load_scoped_page_index(&store, &loc, &fo, &cols).await.unwrap(); + let full = full_index(&bytes); + + // price pages: 0..8,8..16,16..24,24..32; `price >= 20` keeps the last two. + let pp = build_pruning_predicate(&pred("price", 0, Operator::GtEq, 20), schema.clone()) + .unwrap(); + let s = PagePruner::new(&schema, Arc::clone(&aug)).prune_rg(&pp, 0, None); + let f = PagePruner::new(&schema, full).prune_rg(&pp, 0, None); + assert_eq!(s.as_ref().map(kept), f.as_ref().map(kept)); + assert_eq!(s.as_ref().map(kept), Some(16)); + } + + /// A non-predicate but PROJECTED column must read correctly through the + /// scoped metadata — its OffsetIndex is real, so the offset-index read path + /// fetches the right page ranges (regression for "failed to skip rows"). + #[tokio::test] + async fn scoped_index_reads_non_predicate_projected_column() { + let _g = CACHE_TEST_GUARD.lock().unwrap(); + clear_scoped_cache_for_test(); + let (bytes, schema) = two_col_parquet(); + let (store, loc) = stage(bytes.clone()).await; + let fo = footer_only(&bytes); + + let cols = resolve_predicate_parquet_columns(&schema, &fo, &["price".to_string()]); + let aug = load_scoped_page_index(&store, &loc, &fo, &cols).await.unwrap(); + + let selection = RowSelection::from(vec![RowSelector::skip(16), RowSelector::select(16)]); + // Project the NON-predicate column (qty = leaf 1). + let scoped_vals = read_selected_column(&bytes, &aug, 1, selection.clone()) + .expect("non-predicate projected read must succeed with scoped metadata"); + let full = full_index(&bytes); + let full_vals = read_selected_column(&bytes, &full, 1, selection).unwrap(); + + let expected: Vec = (116..132).collect(); + assert_eq!(scoped_vals, expected); + assert_eq!(scoped_vals, full_vals); + } + + /// First load of a (file, column-set) is a miss; the next identical load is a + /// HIT served from cache — asserted via the counter (graft returns a fresh + /// Arc, so ptr_eq would be wrong), with no new entry and no byte growth. + #[tokio::test] + async fn second_load_is_cache_hit() { + let _g = CACHE_TEST_GUARD.lock().unwrap(); + clear_scoped_cache_for_test(); + let (bytes, schema) = two_col_parquet(); + let (store, loc) = stage(bytes.clone()).await; + let fo = footer_only(&bytes); + let cols = resolve_predicate_parquet_columns(&schema, &fo, &["price".to_string()]); + + let _first = load_scoped_page_index(&store, &loc, &fo, &cols).await.unwrap(); + let s1 = scoped_cache_stats(); + assert_eq!((s1.hits, s1.misses, s1.entries), (0, 1, 1)); + let bytes_after_first = s1.used_bytes; + assert!(bytes_after_first > 0, "a real page index must cost > 0 bytes"); + + let _second = load_scoped_page_index(&store, &loc, &fo, &cols).await.unwrap(); + let s2 = scoped_cache_stats(); + assert_eq!( + (s2.hits, s2.misses, s2.entries, s2.used_bytes), + (1, 1, 1, bytes_after_first), + "second identical load: pure hit, no new entry, no byte growth" + ); + } + + /// Distinct predicate-column sets are cached independently. + #[tokio::test] + async fn distinct_column_sets_cached_independently() { + let _g = CACHE_TEST_GUARD.lock().unwrap(); + clear_scoped_cache_for_test(); + let (bytes, schema) = two_col_parquet(); + let (store, loc) = stage(bytes.clone()).await; + let fo = footer_only(&bytes); + + let c_price = resolve_predicate_parquet_columns(&schema, &fo, &["price".to_string()]); + let c_qty = resolve_predicate_parquet_columns(&schema, &fo, &["qty".to_string()]); + assert_eq!(c_price, vec![0]); + assert_eq!(c_qty, vec![1]); + + let a_price = load_scoped_page_index(&store, &loc, &fo, &c_price).await.unwrap(); + let _a_qty = load_scoped_page_index(&store, &loc, &fo, &c_qty).await.unwrap(); + assert_eq!(scoped_cache_len_for_test(), 2); + + // qty-scoped index prunes a qty predicate; price-scoped prunes a price one. + let pp_qty = + build_pruning_predicate(&pred("qty", 1, Operator::GtEq, 120), schema.clone()).unwrap(); + // The qty entry must be usable; reload (hit) then prune. + let a_qty2 = load_scoped_page_index(&store, &loc, &fo, &c_qty).await.unwrap(); + let sel_qty = PagePruner::new(&schema, a_qty2).prune_rg(&pp_qty, 0, None).unwrap(); + assert_eq!(kept(&sel_qty), 16); + + let pp_price = + build_pruning_predicate(&pred("price", 0, Operator::GtEq, 20), schema.clone()).unwrap(); + let sel_price = PagePruner::new(&schema, a_price).prune_rg(&pp_price, 0, None).unwrap(); + assert_eq!(kept(&sel_price), 16); + } + + /// Hit/miss accounting across two distinct column-sets. + #[tokio::test] + async fn stats_count_hits_and_misses() { + let _g = CACHE_TEST_GUARD.lock().unwrap(); + clear_scoped_cache_for_test(); + let (bytes, schema) = two_col_parquet(); + let (store, loc) = stage(bytes.clone()).await; + let fo = footer_only(&bytes); + let c_price = resolve_predicate_parquet_columns(&schema, &fo, &["price".to_string()]); + let c_qty = resolve_predicate_parquet_columns(&schema, &fo, &["qty".to_string()]); + + let _ = load_scoped_page_index(&store, &loc, &fo, &c_price).await.unwrap(); + assert_eq!( + (scoped_cache_stats().hits, scoped_cache_stats().misses), + (0, 1) + ); + let _ = load_scoped_page_index(&store, &loc, &fo, &c_price).await.unwrap(); + assert_eq!( + (scoped_cache_stats().hits, scoped_cache_stats().misses), + (1, 1) + ); + let _ = load_scoped_page_index(&store, &loc, &fo, &c_qty).await.unwrap(); + assert_eq!( + (scoped_cache_stats().hits, scoped_cache_stats().misses), + (1, 2) + ); + let _ = load_scoped_page_index(&store, &loc, &fo, &c_price).await.unwrap(); + let _ = load_scoped_page_index(&store, &loc, &fo, &c_qty).await.unwrap(); + let s = scoped_cache_stats(); + assert_eq!((s.hits, s.misses, s.entries, s.evictions), (3, 2, 2, 0)); + } + + /// Byte-bounded LRU evicts the least-recently-used over budget, never + /// exceeding the limit, and never degrades to "cache nothing". + #[tokio::test] + async fn lru_evicts_over_byte_budget() { + let _g = CACHE_TEST_GUARD.lock().unwrap(); + clear_scoped_cache_for_test(); + let (bytes, schema) = two_col_parquet(); + let (store, loc) = stage(bytes.clone()).await; + let fo = footer_only(&bytes); + let c_price = resolve_predicate_parquet_columns(&schema, &fo, &["price".to_string()]); + let c_qty = resolve_predicate_parquet_columns(&schema, &fo, &["qty".to_string()]); + let c_both = resolve_predicate_parquet_columns( + &schema, + &fo, + &["price".to_string(), "qty".to_string()], + ); + + // Size one entry, then set the budget to ~1.5 entries so only one fits. + let _ = load_scoped_page_index(&store, &loc, &fo, &c_price).await.unwrap(); + let one_size = scoped_cache_bytes_for_test(); + assert!(one_size > 0); + let budget = one_size + one_size / 2; + // Clear, set the budget, and fill from empty. + clear_scoped_cache_for_test(); + set_scoped_cache_limit_for_test(budget); + + let _ = load_scoped_page_index(&store, &loc, &fo, &c_price).await.unwrap(); + let _ = load_scoped_page_index(&store, &loc, &fo, &c_qty).await.unwrap(); + let _ = load_scoped_page_index(&store, &loc, &fo, &c_both).await.unwrap(); + + assert!( + scoped_cache_bytes_for_test() <= budget, + "cache bytes {} must stay within budget {}", + scoped_cache_bytes_for_test(), + budget + ); + assert!( + scoped_cache_len_for_test() >= 1, + "LRU must retain at least the most-recent entry" + ); + assert!(scoped_cache_stats().evictions >= 1, "something must have evicted"); + + // The most-recently-used (c_both) must still be a hit. + let hits_before = scoped_cache_stats().hits; + let _ = load_scoped_page_index(&store, &loc, &fo, &c_both).await.unwrap(); + assert_eq!( + scoped_cache_stats().hits, + hits_before + 1, + "most-recently-used entry must remain cached" + ); + + clear_scoped_cache_for_test(); + } +} From dde881f3d15c44aea49eff6639e6b40b505afe07 Mon Sep 17 00:00:00 2001 From: G Date: Tue, 16 Jun 2026 19:31:47 +0530 Subject: [PATCH 02/18] Step 1b: surface scoped page-index cache on node-stats + add limit setting MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Make the scoped cache (Step 1a) observable for A/B testing, and let its byte budget be configured. Rust: - stats.rs: CacheStatsRepr grows a third CacheGroupRepr (scoped_page_index_cache), packed from page_index_loader::scoped_cache_stats(). DfStatsBuffer 75->80 longs (600->640 bytes); asserts + tests updated. New test asserts the packed group reflects the live global counters/limit. - ffm.rs: df_set_scoped_page_index_cache_limit(i64) — process-global, no manager pointer; negative rejected, zero ignored. Java: - StatsLayout: cache_stats group now nests 3 sub-caches (metadata, statistics, scoped_page_index); layout 80 longs/640 bytes; 5 new VarHandles; readCacheStats builds the 3rd group. - CacheStats: third CacheGroupStats threaded through ctor/StreamInput/writeTo/ XContent ("scoped_page_index_cache")/equals/hashCode + getter. - NativeBridge.setScopedPageIndexCacheLimit -> df_set_scoped_page_index_cache_limit. - CacheSettings datafusion.scoped_page_index.cache.size.limit (default 64mb); CacheUtils pushes it to native at startup (NodeScope, not yet dynamic). - Updated StatsLayout(Property)Tests offsets/FIELD_COUNT (75->80) + CacheStats(*) tests for the 3-arg shape. Tests: 19 Rust stats + 9 page_index_loader pass; 7 Java stats suites (57 tests) pass under -PrustDebug -Dsandbox.enabled=true. No scan-path behavior change yet. --- .../rust/src/ffm.rs | 20 ++++++ .../rust/src/stats.rs | 61 ++++++++++++++++--- .../be/datafusion/cache/CacheSettings.java | 19 +++++- .../be/datafusion/cache/CacheUtils.java | 10 +++ .../be/datafusion/nativelib/NativeBridge.java | 21 +++++++ .../be/datafusion/nativelib/StatsLayout.java | 36 ++++++++--- .../be/datafusion/stats/CacheStats.java | 25 ++++++-- .../nativelib/StatsLayoutPropertyTests.java | 51 ++++++++++------ .../nativelib/StatsLayoutTests.java | 6 +- .../be/datafusion/stats/CacheStatsTests.java | 55 +++++++++++++---- .../stats/DataFusionStatsPropertyTests.java | 2 +- .../stats/DataFusionStatsTests.java | 3 +- 12 files changed, 249 insertions(+), 60 deletions(-) diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs index 55e8719813e64..7e481846707e5 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs @@ -1040,6 +1040,26 @@ pub unsafe extern "C" fn df_cache_manager_get_total_memory(runtime_ptr: i64) -> Ok(manager.get_total_memory_consumed() as i64) } +/// Set the byte budget of the process-global scoped page-index cache. +/// +/// Unlike the metadata/statistics caches, the scoped page-index cache is a +/// process-wide singleton (see [`crate::indexed_table::page_index_loader`]), so +/// there is no manager pointer — the limit applies globally. Negative values are +/// rejected; zero is ignored (keeps the existing budget). Shrinking the limit +/// evicts least-recently-used entries immediately. +#[ffm_safe] +#[no_mangle] +pub extern "C" fn df_set_scoped_page_index_cache_limit(size_limit: i64) -> i64 { + if size_limit < 0 { + return Err(format!( + "df_set_scoped_page_index_cache_limit: negative limit {}", + size_limit + )); + } + crate::indexed_table::page_index_loader::set_scoped_cache_limit(size_limit as usize); + Ok(0) +} + #[ffm_safe] #[no_mangle] pub unsafe extern "C" fn df_cache_manager_contains_by_type( diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/stats.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/stats.rs index 2fcec9ce4d86e..58f83cafaa09d 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/stats.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/stats.rs @@ -20,7 +20,7 @@ //! | `plan_setup` | `TaskMonitorRepr` | 5 × i64 | //! | `fragment_executor_gate` | `PartitionGateRepr` | 8 × i64 | //! | `adaptive_budget` | `AdaptiveBudgetRepr` | 2 × i64 | -//! | `cache_stats` | `CacheStatsRepr` | 10 × i64 (2 × 5: metadata + statistics caches) | +//! | `cache_stats` | `CacheStatsRepr` | 15 × i64 (3 × 5: metadata + statistics + scoped-page-index caches) | //! | `search_stats` | `SearchStatsRepr` | 17 × i64 | use tokio::runtime::Handle; @@ -92,6 +92,12 @@ pub struct CacheGroupRepr { pub struct CacheStatsRepr { pub metadata_cache: CacheGroupRepr, pub statistics_cache: CacheGroupRepr, + /// Process-wide scoped page-index cache (see + /// [`crate::indexed_table::page_index_loader`]). `hit_count`/`miss_count` + /// are cumulative since process start; `entry_count`/`memory_bytes`/ + /// `size_limit_bytes` are point-in-time. (Evictions are tracked internally + /// but not surfaced in this 5-field group.) + pub scoped_page_index_cache: CacheGroupRepr, } impl Default for CacheGroupRepr { @@ -111,6 +117,7 @@ impl Default for CacheStatsRepr { Self { metadata_cache: CacheGroupRepr::default(), statistics_cache: CacheGroupRepr::default(), + scoped_page_index_cache: CacheGroupRepr::default(), } } } @@ -161,14 +168,14 @@ const _: () = assert!(std::mem::size_of::() == 5 * 8); const _: () = assert!(std::mem::size_of::() == 8 * 8); const _: () = assert!(std::mem::size_of::() == 2 * 8); const _: () = assert!(std::mem::size_of::() == 5 * 8); -const _: () = assert!(std::mem::size_of::() == 10 * 8); +const _: () = assert!(std::mem::size_of::() == 15 * 8); const _: () = assert!(std::mem::size_of::() == 17 * 8); -const _: () = assert!(std::mem::size_of::() == 75 * 8); +const _: () = assert!(std::mem::size_of::() == 80 * 8); pub mod layout { use super::*; pub const BUFFER_BYTE_SIZE: usize = std::mem::size_of::(); - const _: () = assert!(BUFFER_BYTE_SIZE == 600); + const _: () = assert!(BUFFER_BYTE_SIZE == 640); } /// Snapshot a `RuntimeMonitor` and return a populated `RuntimeMetricsRepr`. @@ -304,6 +311,17 @@ pub fn pack_cache_stats(mgr: &CustomCacheManager) -> CacheStatsRepr { memory_bytes: statistics_memory, size_limit_bytes: mgr.statistics_cache_size_limit() as i64, }, + scoped_page_index_cache: { + // Process-global cache, not owned by the manager — read it directly. + let s = crate::indexed_table::page_index_loader::scoped_cache_stats(); + CacheGroupRepr { + hit_count: s.hits as i64, + miss_count: s.misses as i64, + entry_count: s.entries as i64, + memory_bytes: s.used_bytes as i64, + size_limit_bytes: s.limit_bytes as i64, + } + }, } } @@ -396,7 +414,7 @@ mod tests { search_stats: crate::search_stats::snapshot(), }; - assert_eq!(layout::BUFFER_BYTE_SIZE, 600); + assert_eq!(layout::BUFFER_BYTE_SIZE, 640); assert!(buf.io_runtime.workers_count > 0, "IO runtime workers_count should be > 0, got {}", buf.io_runtime.workers_count); assert!(buf.fragment_executor_gate.max_permits > 0, "fragment_executor_gate max_permits should be > 0, got {}", buf.fragment_executor_gate.max_permits); @@ -411,9 +429,9 @@ mod tests { #[test] fn test_df_stats_buffer_too_small() { // Verify that the buffer size assertion holds - assert_eq!(std::mem::size_of::(), 600); - assert_eq!(layout::BUFFER_BYTE_SIZE, 600); - // A buffer smaller than 600 bytes should be rejected by df_stats. + assert_eq!(std::mem::size_of::(), 640); + assert_eq!(layout::BUFFER_BYTE_SIZE, 640); + // A buffer smaller than 640 bytes should be rejected by df_stats. // We can't call df_stats directly without a runtime manager, // but we verify the constant is correct. assert!(layout::BUFFER_BYTE_SIZE > 0); @@ -431,6 +449,33 @@ mod tests { assert_eq!(g.memory_bytes, 0); assert_eq!(g.size_limit_bytes, 0); } + // The scoped page-index cache is process-global, not manager-owned, so it + // is NOT necessarily zero here (other tests may have populated it). It + // must always advertise a positive byte budget, though. + assert!(repr.scoped_page_index_cache.size_limit_bytes > 0); + } + + #[test] + fn test_pack_cache_stats_reflects_scoped_page_index_counters() { + use crate::custom_cache_manager::CustomCacheManager; + use crate::indexed_table::page_index_loader::{ + scoped_cache_stats, set_scoped_cache_limit, SCOPED_CACHE_TEST_GUARD, + }; + + // Serialize against the process-global scoped cache. + let _g = SCOPED_CACHE_TEST_GUARD.lock().unwrap(); + // Set a known budget and read it back through the packed repr. + set_scoped_cache_limit(123 * 1024 * 1024); + + let mgr = CustomCacheManager::new(); + let repr = pack_cache_stats(&mgr); + let live = scoped_cache_stats(); + + assert_eq!(repr.scoped_page_index_cache.size_limit_bytes, 123 * 1024 * 1024); + assert_eq!(repr.scoped_page_index_cache.hit_count, live.hits as i64); + assert_eq!(repr.scoped_page_index_cache.miss_count, live.misses as i64); + assert_eq!(repr.scoped_page_index_cache.entry_count, live.entries as i64); + assert_eq!(repr.scoped_page_index_cache.memory_bytes, live.used_bytes as i64); } #[test] diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/cache/CacheSettings.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/cache/CacheSettings.java index 3f1f9e969b880..c849df35480a6 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/cache/CacheSettings.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/cache/CacheSettings.java @@ -20,6 +20,7 @@ public class CacheSettings { public static final String METADATA_CACHE_SIZE_LIMIT_KEY = "datafusion.metadata.cache.size.limit"; public static final String STATISTICS_CACHE_SIZE_LIMIT_KEY = "datafusion.statistics.cache.size.limit"; + public static final String SCOPED_PAGE_INDEX_CACHE_SIZE_LIMIT_KEY = "datafusion.scoped_page_index.cache.size.limit"; public static final Setting METADATA_CACHE_SIZE_LIMIT = new Setting<>( METADATA_CACHE_SIZE_LIMIT_KEY, "250mb", @@ -36,6 +37,21 @@ public class CacheSettings { Setting.Property.Dynamic ); + /** + * Byte budget for the process-global scoped page-index cache (column index + * scoped to predicate columns, offset index for all columns). Pushed to + * native at startup via {@link CacheUtils#createCacheConfig}. Unlike the + * metadata/statistics caches this is a process-wide singleton, not owned by + * the cache manager. + */ + public static final Setting SCOPED_PAGE_INDEX_CACHE_SIZE_LIMIT = new Setting<>( + SCOPED_PAGE_INDEX_CACHE_SIZE_LIMIT_KEY, + "64mb", + (s) -> ByteSizeValue.parseBytesSizeValue(s, new ByteSizeValue(0, ByteSizeUnit.KB), SCOPED_PAGE_INDEX_CACHE_SIZE_LIMIT_KEY), + Setting.Property.NodeScope, + Setting.Property.Dynamic + ); + public static final Setting METADATA_CACHE_EVICTION_TYPE = new Setting( "datafusion.metadata.cache.eviction.type", "LRU", @@ -74,7 +90,8 @@ public class CacheSettings { METADATA_CACHE_EVICTION_TYPE, STATISTICS_CACHE_ENABLED, STATISTICS_CACHE_SIZE_LIMIT, - STATISTICS_CACHE_EVICTION_TYPE + STATISTICS_CACHE_EVICTION_TYPE, + SCOPED_PAGE_INDEX_CACHE_SIZE_LIMIT ); private static String validateEvictionType(String value) { diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/cache/CacheUtils.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/cache/CacheUtils.java index 1b95478fd83aa..2b55bb514728a 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/cache/CacheUtils.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/cache/CacheUtils.java @@ -117,6 +117,16 @@ public static NativeCacheManagerHandle createCacheConfig(ClusterSettings cluster logger.debug("Cache type {} is disabled", type.getCacheTypeName()); } } + + // The scoped page-index cache is a process-global singleton (not owned by + // the cache manager), so its limit is pushed straight to native rather + // than created on the manager. Startup-only for now; a settings-update + // consumer can call NativeBridge.setScopedPageIndexCacheLimit to make it + // dynamic. + long scopedLimit = clusterSettings.get(CacheSettings.SCOPED_PAGE_INDEX_CACHE_SIZE_LIMIT).getBytes(); + logger.info("Configuring scoped page-index cache: size={} bytes", scopedLimit); + NativeBridge.setScopedPageIndexCacheLimit(scopedLimit); + logger.info("Cache configuration completed"); return handle; } diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/nativelib/NativeBridge.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/nativelib/NativeBridge.java index 475f8d228191b..b0d1b98edf01a 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/nativelib/NativeBridge.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/nativelib/NativeBridge.java @@ -130,6 +130,7 @@ private static RuntimeException rethrowConverted(RuntimeException e) { private static final MethodHandle CACHE_MANAGER_GET_MEMORY_BY_TYPE; private static final MethodHandle CACHE_MANAGER_GET_TOTAL_MEMORY; private static final MethodHandle CACHE_MANAGER_CONTAINS_BY_TYPE; + private static final MethodHandle SET_SCOPED_PAGE_INDEX_CACHE_LIMIT; private static final MethodHandle CREATE_SESSION_CONTEXT; private static final MethodHandle CREATE_SESSION_CONTEXT_INDEXED; private static final MethodHandle CLOSE_SESSION_CONTEXT; @@ -501,6 +502,11 @@ private static RuntimeException rethrowConverted(RuntimeException e) { ) ); + SET_SCOPED_PAGE_INDEX_CACHE_LIMIT = linker.downcallHandle( + lib.find("df_set_scoped_page_index_cache_limit").orElseThrow(), + FunctionDescriptor.of(ValueLayout.JAVA_LONG, ValueLayout.JAVA_LONG) + ); + CANCEL_QUERY = linker.downcallHandle(lib.find("df_cancel_query").orElseThrow(), FunctionDescriptor.ofVoid(ValueLayout.JAVA_LONG)); SET_CANCEL_STATS_THRESHOLD_MS = linker.downcallHandle( @@ -762,6 +768,21 @@ public static void setMemoryPoolLimit(long runtimePtr, long newLimitBytes) { } } + /** + * Sets the byte budget of the process-global scoped page-index cache. + * + *

This cache is a process-wide singleton (not owned by any cache manager), + * so there is no runtime/manager pointer — the limit applies globally. + * Shrinking the limit evicts least-recently-used entries immediately. + * + * @param sizeLimitBytes the new byte budget (zero is ignored; keeps the current budget) + */ + public static void setScopedPageIndexCacheLimit(long sizeLimitBytes) { + try (var call = new NativeCall()) { + call.invoke(SET_SCOPED_PAGE_INDEX_CACHE_LIMIT, sizeLimitBytes); + } + } + /** * Returns true if the loaded native library exports {@code df_set_spill_limit}. * When false, spill-cap updates require a node restart — the public diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/nativelib/StatsLayout.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/nativelib/StatsLayout.java index 8e90357fcac53..cd1821fc224f4 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/nativelib/StatsLayout.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/nativelib/StatsLayout.java @@ -28,8 +28,9 @@ * and provides {@link VarHandle} accessors for each field via layout path navigation. * *

The layout contains 10 named groups (2 runtime × 9 fields + 4 task monitor × 5 fields - * + 1 partition gate × 8 fields + 1 adaptive budget × 2 fields + 1 cache stats × 10 fields - * + 1 search stats × 17 fields = 75 longs = 600 bytes). + * + 1 partition gate × 8 fields + 1 adaptive budget × 2 fields + 1 cache stats × 15 fields + * + 1 search stats × 17 fields = 80 longs = 640 bytes). The cache stats group holds three + * sub-caches × 5 fields each: metadata, statistics, and scoped page-index. */ public final class StatsLayout { @@ -99,8 +100,8 @@ public final class StatsLayout { ); static { - if (LAYOUT.byteSize() != 75 * Long.BYTES) { - throw new AssertionError("StatsLayout size mismatch: expected " + (75 * Long.BYTES) + " but got " + LAYOUT.byteSize()); + if (LAYOUT.byteSize() != 80 * Long.BYTES) { + throw new AssertionError("StatsLayout size mismatch: expected " + (80 * Long.BYTES) + " but got " + LAYOUT.byteSize()); } } @@ -182,6 +183,13 @@ public final class StatsLayout { private static final VarHandle CACHE_STATS_MEMORY_BYTES = cacheHandle("statistics_cache", "memory_bytes"); private static final VarHandle CACHE_STATS_SIZE_LIMIT_BYTES = cacheHandle("statistics_cache", "size_limit_bytes"); + // ---- VarHandles for cache_stats.scoped_page_index_cache fields ---- + private static final VarHandle CACHE_SPI_HIT_COUNT = cacheHandle("scoped_page_index_cache", "hit_count"); + private static final VarHandle CACHE_SPI_MISS_COUNT = cacheHandle("scoped_page_index_cache", "miss_count"); + private static final VarHandle CACHE_SPI_ENTRY_COUNT = cacheHandle("scoped_page_index_cache", "entry_count"); + private static final VarHandle CACHE_SPI_MEMORY_BYTES = cacheHandle("scoped_page_index_cache", "memory_bytes"); + private static final VarHandle CACHE_SPI_SIZE_LIMIT_BYTES = cacheHandle("scoped_page_index_cache", "size_limit_bytes"); + // ---- VarHandles for search_stats fields ---- private static final VarHandle SS_LISTING_TABLE_SCAN = handle("search_stats", "listing_table_scan"); private static final VarHandle SS_SINGLE_COLLECTOR_SCAN = handle("search_stats", "single_collector_scan"); @@ -290,11 +298,10 @@ public static PartitionGateStats readPartitionGate(MemorySegment seg, String gro } /** - /** - * Read the cache_stats group (10 fields, 2 sub-caches × 5 fields each). + * Read the cache_stats group (15 fields, 3 sub-caches × 5 fields each). * * @param seg the memory segment containing the DfStatsBuffer - * @return a populated CacheStats instance with metadata + statistics sub-groups + * @return a populated CacheStats instance with metadata + statistics + scoped-page-index sub-groups */ public static CacheStats readCacheStats(MemorySegment seg) { CacheGroupStats metadata = new CacheGroupStats( @@ -311,7 +318,14 @@ public static CacheStats readCacheStats(MemorySegment seg) { (long) CACHE_STATS_MEMORY_BYTES.get(seg, 0L), (long) CACHE_STATS_SIZE_LIMIT_BYTES.get(seg, 0L) ); - return new CacheStats(metadata, statistics); + CacheGroupStats scopedPageIndex = new CacheGroupStats( + (long) CACHE_SPI_HIT_COUNT.get(seg, 0L), + (long) CACHE_SPI_MISS_COUNT.get(seg, 0L), + (long) CACHE_SPI_ENTRY_COUNT.get(seg, 0L), + (long) CACHE_SPI_MEMORY_BYTES.get(seg, 0L), + (long) CACHE_SPI_SIZE_LIMIT_BYTES.get(seg, 0L) + ); + return new CacheStats(metadata, statistics, scopedPageIndex); } /** @@ -402,7 +416,11 @@ private static StructLayout cacheGroup(String name) { } private static StructLayout cacheStatsGroup(String name) { - return MemoryLayout.structLayout(cacheGroup("metadata_cache"), cacheGroup("statistics_cache")).withName(name); + return MemoryLayout.structLayout( + cacheGroup("metadata_cache"), + cacheGroup("statistics_cache"), + cacheGroup("scoped_page_index_cache") + ).withName(name); } private static StructLayout searchStatsGroup(String name) { diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/stats/CacheStats.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/stats/CacheStats.java index 3f01c221f3c68..f62ce980b21bd 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/stats/CacheStats.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/stats/CacheStats.java @@ -29,16 +29,19 @@ public class CacheStats implements Writeable, ToXContentFragment { private final CacheGroupStats metadataCache; private final CacheGroupStats statisticsCache; + private final CacheGroupStats scopedPageIndexCache; /** * Construct from individual sub-group stats. * - * @param metadataCache metadata cache counters (must not be null) - * @param statisticsCache statistics cache counters (must not be null) + * @param metadataCache metadata cache counters (must not be null) + * @param statisticsCache statistics cache counters (must not be null) + * @param scopedPageIndexCache scoped page-index cache counters (must not be null) */ - public CacheStats(CacheGroupStats metadataCache, CacheGroupStats statisticsCache) { + public CacheStats(CacheGroupStats metadataCache, CacheGroupStats statisticsCache, CacheGroupStats scopedPageIndexCache) { this.metadataCache = Objects.requireNonNull(metadataCache); this.statisticsCache = Objects.requireNonNull(statisticsCache); + this.scopedPageIndexCache = Objects.requireNonNull(scopedPageIndexCache); } /** @@ -50,12 +53,14 @@ public CacheStats(CacheGroupStats metadataCache, CacheGroupStats statisticsCache public CacheStats(StreamInput in) throws IOException { this.metadataCache = new CacheGroupStats(in); this.statisticsCache = new CacheGroupStats(in); + this.scopedPageIndexCache = new CacheGroupStats(in); } @Override public void writeTo(StreamOutput out) throws IOException { metadataCache.writeTo(out); statisticsCache.writeTo(out); + scopedPageIndexCache.writeTo(out); } @Override @@ -67,6 +72,9 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws builder.startObject("statistics_cache"); statisticsCache.toXContent(builder); builder.endObject(); + builder.startObject("scoped_page_index_cache"); + scopedPageIndexCache.toXContent(builder); + builder.endObject(); builder.endObject(); return builder; } @@ -81,16 +89,23 @@ public CacheGroupStats getStatisticsCache() { return statisticsCache; } + /** Returns the scoped page-index cache counters. */ + public CacheGroupStats getScopedPageIndexCache() { + return scopedPageIndexCache; + } + @Override public boolean equals(Object o) { if (this == o) return true; if (o == null || getClass() != o.getClass()) return false; CacheStats that = (CacheStats) o; - return Objects.equals(metadataCache, that.metadataCache) && Objects.equals(statisticsCache, that.statisticsCache); + return Objects.equals(metadataCache, that.metadataCache) + && Objects.equals(statisticsCache, that.statisticsCache) + && Objects.equals(scopedPageIndexCache, that.scopedPageIndexCache); } @Override public int hashCode() { - return Objects.hash(metadataCache, statisticsCache); + return Objects.hash(metadataCache, statisticsCache, scopedPageIndexCache); } } diff --git a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/nativelib/StatsLayoutPropertyTests.java b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/nativelib/StatsLayoutPropertyTests.java index 4fd7679ce5b9b..bb1c0afea1229 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/nativelib/StatsLayoutPropertyTests.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/nativelib/StatsLayoutPropertyTests.java @@ -35,7 +35,7 @@ public class StatsLayoutPropertyTests extends OpenSearchTestCase { private static final int TRIES = 100; - private static final int FIELD_COUNT = 75; + private static final int FIELD_COUNT = 80; // ---- Generators ---- @@ -175,7 +175,7 @@ public void testPackThenDecodeRoundTripWithCpu() { assertEquals(values[46], bs.fallbacks); assertEquals(values[47], bs.rejections); - // Cache stats (offsets 48-57) + // Cache stats (offsets 48-62; 3 sub-caches × 5) var cs = StatsLayout.readCacheStats(seg); assertEquals(values[48], cs.getMetadataCache().hitCount); assertEquals(values[49], cs.getMetadataCache().missCount); @@ -187,26 +187,31 @@ public void testPackThenDecodeRoundTripWithCpu() { assertEquals(values[55], cs.getStatisticsCache().entryCount); assertEquals(values[56], cs.getStatisticsCache().memoryBytes); assertEquals(values[57], cs.getStatisticsCache().sizeLimitBytes); + assertEquals(values[58], cs.getScopedPageIndexCache().hitCount); + assertEquals(values[59], cs.getScopedPageIndexCache().missCount); + assertEquals(values[60], cs.getScopedPageIndexCache().entryCount); + assertEquals(values[61], cs.getScopedPageIndexCache().memoryBytes); + assertEquals(values[62], cs.getScopedPageIndexCache().sizeLimitBytes); - // Search stats (offsets 58-74) + // Search stats (offsets 63-79) var ss = StatsLayout.readSearchStats(seg); - assertEquals(values[58], ss.listingTableScan); - assertEquals(values[59], ss.singleCollectorScan); - assertEquals(values[60], ss.bitmapTreeScan); - assertEquals(values[61], ss.delegationCalls); - assertEquals(values[62], ss.rgProcessed); - assertEquals(values[63], ss.rgSkipped); - assertEquals(values[64], ss.parquetScanTotalTimeMs); - assertEquals(values[65], ss.parquetScanUntilDataTimeMs); - assertEquals(values[66], ss.parquetProcessingTimeMs); - assertEquals(values[67], ss.parquetBytesScanned); - assertEquals(values[68], ss.prefetchWaitTimeMs); - assertEquals(values[69], ss.prefetchWaitCount); - assertEquals(values[70], ss.elapsedComputeMs); - assertEquals(values[71], ss.buildMaskTimeMs); - assertEquals(values[72], ss.onBatchMaskTimeMs); - assertEquals(values[73], ss.filterRecordBatchTimeMs); - assertEquals(values[74], ss.objectStoreReadTimeMs); + assertEquals(values[63], ss.listingTableScan); + assertEquals(values[64], ss.singleCollectorScan); + assertEquals(values[65], ss.bitmapTreeScan); + assertEquals(values[66], ss.delegationCalls); + assertEquals(values[67], ss.rgProcessed); + assertEquals(values[68], ss.rgSkipped); + assertEquals(values[69], ss.parquetScanTotalTimeMs); + assertEquals(values[70], ss.parquetScanUntilDataTimeMs); + assertEquals(values[71], ss.parquetProcessingTimeMs); + assertEquals(values[72], ss.parquetBytesScanned); + assertEquals(values[73], ss.prefetchWaitTimeMs); + assertEquals(values[74], ss.prefetchWaitCount); + assertEquals(values[75], ss.elapsedComputeMs); + assertEquals(values[76], ss.buildMaskTimeMs); + assertEquals(values[77], ss.onBatchMaskTimeMs); + assertEquals(values[78], ss.filterRecordBatchTimeMs); + assertEquals(values[79], ss.objectStoreReadTimeMs); } } } @@ -328,6 +333,12 @@ public void testDecodeThenReencodeIdentity() { cs.getStatisticsCache().entryCount, cs.getStatisticsCache().memoryBytes, cs.getStatisticsCache().sizeLimitBytes, + // cache_stats.scoped_page_index_cache (5) + cs.getScopedPageIndexCache().hitCount, + cs.getScopedPageIndexCache().missCount, + cs.getScopedPageIndexCache().entryCount, + cs.getScopedPageIndexCache().memoryBytes, + cs.getScopedPageIndexCache().sizeLimitBytes, // search_stats (17) ss.listingTableScan, ss.singleCollectorScan, diff --git a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/nativelib/StatsLayoutTests.java b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/nativelib/StatsLayoutTests.java index 9a72c2af21a26..b938e1fd17f43 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/nativelib/StatsLayoutTests.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/nativelib/StatsLayoutTests.java @@ -22,10 +22,10 @@ */ public class StatsLayoutTests extends OpenSearchTestCase { - /** 7.1: Layout byte size must be 600 (75 × 8). */ + /** 7.1: Layout byte size must be 640 (80 × 8). */ public void testLayoutByteSize() { - assertEquals(600L, StatsLayout.LAYOUT.byteSize()); - assertEquals(75 * Long.BYTES, (int) StatsLayout.LAYOUT.byteSize()); + assertEquals(640L, StatsLayout.LAYOUT.byteSize()); + assertEquals(80 * Long.BYTES, (int) StatsLayout.LAYOUT.byteSize()); } /** 7.2: readRuntimeMetrics decodes 9 known values from io_runtime group. */ diff --git a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/stats/CacheStatsTests.java b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/stats/CacheStatsTests.java index 77a839c377b9c..3b51f80648697 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/stats/CacheStatsTests.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/stats/CacheStatsTests.java @@ -90,21 +90,29 @@ public void testCacheGroupStatsEqualsAndHashCode() { // ---- CacheStats ---- public void testCacheStatsConstructorRejectsNullSubGroups() { - expectThrows(NullPointerException.class, () -> new CacheStats(null, new CacheGroupStats(0, 0, 0, 0, 0))); - expectThrows(NullPointerException.class, () -> new CacheStats(new CacheGroupStats(0, 0, 0, 0, 0), null)); + CacheGroupStats g = new CacheGroupStats(0, 0, 0, 0, 0); + expectThrows(NullPointerException.class, () -> new CacheStats(null, g, g)); + expectThrows(NullPointerException.class, () -> new CacheStats(g, null, g)); + expectThrows(NullPointerException.class, () -> new CacheStats(g, g, null)); } public void testCacheStatsAccessors() { CacheGroupStats meta = new CacheGroupStats(1, 2, 3, 4, 5); CacheGroupStats stats = new CacheGroupStats(6, 7, 8, 9, 10); - CacheStats c = new CacheStats(meta, stats); + CacheGroupStats scoped = new CacheGroupStats(11, 12, 13, 14, 15); + CacheStats c = new CacheStats(meta, stats, scoped); assertSame(meta, c.getMetadataCache()); assertSame(stats, c.getStatisticsCache()); + assertSame(scoped, c.getScopedPageIndexCache()); } public void testCacheStatsWriteableRoundTrip() throws IOException { - CacheStats original = new CacheStats(new CacheGroupStats(11, 12, 13, 14, 15), new CacheGroupStats(21, 22, 23, 24, 25)); + CacheStats original = new CacheStats( + new CacheGroupStats(11, 12, 13, 14, 15), + new CacheGroupStats(21, 22, 23, 24, 25), + new CacheGroupStats(31, 32, 33, 34, 35) + ); BytesStreamOutput out = new BytesStreamOutput(); original.writeTo(out); StreamInput in = out.bytes().streamInput(); @@ -113,7 +121,11 @@ public void testCacheStatsWriteableRoundTrip() throws IOException { } public void testCacheStatsToXContentShape() throws IOException { - CacheStats c = new CacheStats(new CacheGroupStats(11, 0, 3, 1024, 250_000_000), new CacheGroupStats(0, 7, 0, 0, 100_000_000)); + CacheStats c = new CacheStats( + new CacheGroupStats(11, 0, 3, 1024, 250_000_000), + new CacheGroupStats(0, 7, 0, 0, 100_000_000), + new CacheGroupStats(5, 1, 2, 512, 64_000_000) + ); XContentBuilder builder = XContentFactory.jsonBuilder(); builder.startObject(); c.toXContent(builder, ToXContent.EMPTY_PARAMS); @@ -122,9 +134,10 @@ public void testCacheStatsToXContentShape() throws IOException { // Top-level wrapper assertTrue("expected cache_stats wrapper, got: " + json, json.contains("\"cache_stats\"")); - // Both sub-groups present + // All three sub-groups present assertTrue(json.contains("\"metadata_cache\"")); assertTrue(json.contains("\"statistics_cache\"")); + assertTrue("expected scoped_page_index_cache group, got: " + json, json.contains("\"scoped_page_index_cache\"")); // Per-group fields assertTrue(json.contains("\"hit_count\":11")); assertTrue(json.contains("\"miss_count\":7")); @@ -132,27 +145,45 @@ public void testCacheStatsToXContentShape() throws IOException { assertTrue(json.contains("\"memory_bytes\":1024")); assertTrue(json.contains("\"size_limit_bytes\":250000000")); assertTrue(json.contains("\"size_limit_bytes\":100000000")); + assertTrue(json.contains("\"size_limit_bytes\":64000000")); } public void testCacheStatsZeroedRendersAllZeros() throws IOException { - CacheStats zero = new CacheStats(new CacheGroupStats(0, 0, 0, 0, 0), new CacheGroupStats(0, 0, 0, 0, 0)); + CacheStats zero = new CacheStats( + new CacheGroupStats(0, 0, 0, 0, 0), + new CacheGroupStats(0, 0, 0, 0, 0), + new CacheGroupStats(0, 0, 0, 0, 0) + ); XContentBuilder builder = XContentFactory.jsonBuilder(); builder.startObject(); zero.toXContent(builder, ToXContent.EMPTY_PARAMS); builder.endObject(); String json = builder.toString(); - // Disabled-cache sentinel: both size_limit_bytes are 0 + // Disabled-cache sentinel: all three size_limit_bytes are 0 long sizeLimitOccurrences = json.split("\"size_limit_bytes\":0", -1).length - 1; - assertEquals("expected size_limit_bytes:0 to appear twice (one per sub-cache)", 2, sizeLimitOccurrences); + assertEquals("expected size_limit_bytes:0 to appear three times (one per sub-cache)", 3, sizeLimitOccurrences); // hit_rate must not be NaN assertFalse("hit_rate must not be NaN: " + json, json.contains("NaN")); } public void testCacheStatsEqualsAndHashCode() { - CacheStats a = new CacheStats(new CacheGroupStats(1, 2, 3, 4, 5), new CacheGroupStats(6, 7, 8, 9, 10)); - CacheStats b = new CacheStats(new CacheGroupStats(1, 2, 3, 4, 5), new CacheGroupStats(6, 7, 8, 9, 10)); - CacheStats c = new CacheStats(new CacheGroupStats(1, 2, 3, 4, 5), new CacheGroupStats(6, 7, 8, 9, 99)); + CacheStats a = new CacheStats( + new CacheGroupStats(1, 2, 3, 4, 5), + new CacheGroupStats(6, 7, 8, 9, 10), + new CacheGroupStats(11, 12, 13, 14, 15) + ); + CacheStats b = new CacheStats( + new CacheGroupStats(1, 2, 3, 4, 5), + new CacheGroupStats(6, 7, 8, 9, 10), + new CacheGroupStats(11, 12, 13, 14, 15) + ); + // Differs only in the scoped group → must be unequal. + CacheStats c = new CacheStats( + new CacheGroupStats(1, 2, 3, 4, 5), + new CacheGroupStats(6, 7, 8, 9, 10), + new CacheGroupStats(11, 12, 13, 14, 99) + ); assertEquals(a, b); assertEquals(a.hashCode(), b.hashCode()); diff --git a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/stats/DataFusionStatsPropertyTests.java b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/stats/DataFusionStatsPropertyTests.java index 972deafcbdce7..9943a0d993780 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/stats/DataFusionStatsPropertyTests.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/stats/DataFusionStatsPropertyTests.java @@ -108,7 +108,7 @@ private CacheGroupStats randomCacheGroupStats() { } private CacheStats randomCacheStats() { - return new CacheStats(randomCacheGroupStats(), randomCacheGroupStats()); + return new CacheStats(randomCacheGroupStats(), randomCacheGroupStats(), randomCacheGroupStats()); } private SearchStats randomSearchStats() { diff --git a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/stats/DataFusionStatsTests.java b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/stats/DataFusionStatsTests.java index 075a48c5ebcc5..a7c0d93104dc6 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/stats/DataFusionStatsTests.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/stats/DataFusionStatsTests.java @@ -278,7 +278,8 @@ public void testCacheStatsPresentRendersIntoJson() throws IOException { CacheStats cache = new CacheStats( new CacheGroupStats(100, 5, 25, 4096, 250_000_000), - new CacheGroupStats(50, 0, 25, 2048, 100_000_000) + new CacheGroupStats(50, 0, 25, 2048, 100_000_000), + new CacheGroupStats(20, 3, 8, 1024, 64_000_000) ); DataFusionStats stats = new DataFusionStats( new NativeExecutorsStats(io, null, taskMonitors), From 2b3456221ca99c99d85f7864e7274b1aebb15203 Mon Sep 17 00:00:00 2001 From: G Date: Tue, 16 Jun 2026 19:47:33 +0530 Subject: [PATCH 03/18] Step 1c: wire the listing-table path to the scoped page-index cache MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Populate the unified scoped cache (1a) on the DataFusion ListingTable / vanilla parquet scan path so it is exercised + observable via node-stats (1b). - shard_scoped_reader.rs: ScopedPageIndexReaderFactory — a ParquetFileReaderFactory whose get_metadata loads footer-only metadata, resolves predicate column names to per-file parquet leaf indices, and grafts an all-RG, column-scoped page index via load_scoped_page_index (falls back to footer-only on no-predicate / failure). All row groups (not RG-scoped): DataFusion owns RG selection on this path, so a placeholder on an RG it scans would panic. - scoped_page_index_optimizer.rs: ScopedPageIndexOptimizer — a PhysicalOptimizerRule that finds each parquet DataSourceExec, reads the pushed predicate, derives predicate column names, and REPLACES the reader factory (DataFusion always pre-installs its own full-index CachedParquetFileReaderFactory — that is what we swap). Provider-agnostic; no-op without a predicate. - session_context.rs: register the rule after ProjectRowIdOptimizer so it sees the final DataSourceExec, using this shard's store + the shared metadata cache. Tests (15 pass, --test-threads=1): the optimizer end_to_end_listing_scan_fills_ scoped_cache drives a real SessionContext + stock ListingTable and asserts the scoped cache fills with a miss; reader tests assert predicate-scoped ColumnIndex + all-column OffsetIndex and the no-predicate no-scope path. No metadata-cache strip yet (1e) — scoped-cache population does not depend on it. --- .../rust/src/lib.rs | 2 + .../rust/src/scoped_page_index_optimizer.rs | 356 ++++++++++++++++ .../rust/src/session_context.rs | 13 + .../rust/src/shard_scoped_reader.rs | 379 ++++++++++++++++++ .../be/datafusion/DatafusionSettings.java | 1 + 5 files changed, 751 insertions(+) create mode 100644 sandbox/plugins/analytics-backend-datafusion/rust/src/scoped_page_index_optimizer.rs create mode 100644 sandbox/plugins/analytics-backend-datafusion/rust/src/shard_scoped_reader.rs diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/lib.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/lib.rs index c1c853db2c04a..806e6bec31450 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/lib.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/lib.rs @@ -46,6 +46,8 @@ pub mod query_budget; pub mod query_executor; pub mod query_tracker; pub mod relabel_exec; +pub mod scoped_page_index_optimizer; +pub mod shard_scoped_reader; pub mod shard_table_provider; pub mod runtime_manager; pub mod schema_coerce; diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/scoped_page_index_optimizer.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/scoped_page_index_optimizer.rs new file mode 100644 index 0000000000000..643a2070052fd --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/scoped_page_index_optimizer.rs @@ -0,0 +1,356 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +//! Physical optimizer rule that installs the scoped page-index reader factory on +//! **every** parquet scan in the plan — provider-agnostic. +//! +//! # Why a rule (not a TableProvider) +//! +//! The scoped page-index loader is a property of *how we read parquet*, not of +//! *which TableProvider* produced the scan. Wiring it into a specific provider +//! leaves other scan paths on DataFusion's default reader, which loads the full +//! all-column page index every query and caches none of it. +//! +//! This rule walks the physical plan, finds each parquet `DataSourceExec`, reads +//! the predicate already pushed onto its `ParquetSource`, derives the predicate +//! columns, and swaps in a [`ScopedPageIndexReaderFactory`] scoped to those +//! columns. It runs after DataFusion's own optimizers (which is when filter +//! pushdown has populated `ParquetSource::predicate`), so it works uniformly for +//! `ListingTable`, `ShardTableProvider`, and any future parquet provider. +//! +//! # Replace, do NOT skip-if-present +//! +//! DataFusion's `ParquetFormat::create_physical_plan` ALWAYS pre-installs its own +//! `CachedParquetFileReaderFactory` (the full all-column page-index loader). That +//! is exactly the factory we want to replace, so this rule does not skip a scan +//! just because a factory is already set. (A skip-if-present guard was the +//! original bug that made the end-to-end listing scan never use the scoped +//! reader.) The indexed path does not run this rule — it uses its own executor. +//! +//! # No-op cases (left exactly as DataFusion would run them) +//! +//! - A `DataSourceExec` that isn't parquet — skipped. +//! - A parquet scan with no predicate, or whose predicate references no file +//! columns — skipped (nothing to scope; the opener loads on demand as today). + +use std::sync::Arc; + +use datafusion::common::config::ConfigOptions; +use datafusion::common::tree_node::{Transformed, TreeNode}; +use datafusion::common::Result; +use datafusion::datasource::physical_plan::ParquetSource; +use datafusion::datasource::source::DataSourceExec; +use datafusion::execution::cache::cache_manager::FileMetadataCache; +use datafusion::physical_expr::utils::collect_columns; +use datafusion::physical_optimizer::PhysicalOptimizerRule; +use datafusion::physical_plan::ExecutionPlan; +use datafusion_datasource::file_scan_config::{FileScanConfig, FileScanConfigBuilder}; +use object_store::ObjectStore; + +use crate::shard_scoped_reader::ScopedPageIndexReaderFactory; + +/// Installs the scoped page-index reader factory on parquet scans. +/// +/// Carries the object store and shared metadata cache because a +/// `PhysicalOptimizerRule` has no access to the session; the caller constructs +/// it from the query's `RuntimeEnv`. +#[derive(Debug)] +pub struct ScopedPageIndexOptimizer { + store: Arc, + metadata_cache: Arc, +} + +impl ScopedPageIndexOptimizer { + pub fn new(store: Arc, metadata_cache: Arc) -> Self { + Self { store, metadata_cache } + } +} + +impl PhysicalOptimizerRule for ScopedPageIndexOptimizer { + fn optimize( + &self, + plan: Arc, + _config: &ConfigOptions, + ) -> Result> { + let rewritten = plan.transform_up(|node| { + let Some(dse) = node.downcast_ref::() else { + return Ok(Transformed::no(node)); + }; + let Some(config) = dse.data_source().as_ref().downcast_ref::() else { + return Ok(Transformed::no(node)); + }; + let Some(parquet) = (config.file_source().as_ref() as &dyn std::any::Any) + .downcast_ref::() + else { + return Ok(Transformed::no(node)); + }; + + // Need a predicate to scope to. No predicate → nothing to do. + let Some(predicate) = parquet.predicate() else { + return Ok(Transformed::no(node)); + }; + + // Derive predicate column NAMES that exist in the file schema. The + // reader resolves them to per-file parquet leaf indices. + let file_schema = config.file_schema(); + let mut names: Vec = collect_columns(predicate) + .into_iter() + .map(|c| c.name().to_string()) + .filter(|n| file_schema.index_of(n).is_ok()) + .collect(); + names.sort(); + names.dedup(); + if names.is_empty() { + return Ok(Transformed::no(node)); + } + + // Build the scoped factory and reinstall the source. The predicate is + // retained for parity but not used for RG scoping (Step 1 builds an + // all-row-group, column-scoped page index — see the reader's docs). + let factory = Arc::new(ScopedPageIndexReaderFactory::new( + Arc::clone(&self.store), + Arc::clone(&self.metadata_cache), + names, + Some(Arc::clone(predicate)), + Arc::clone(file_schema), + )); + let new_source = parquet.clone().with_parquet_file_reader_factory(factory); + let new_config = FileScanConfigBuilder::from(config.clone()) + .with_source(Arc::new(new_source)) + .build(); + let new_dse: Arc = DataSourceExec::from_data_source(new_config); + Ok(Transformed::yes(new_dse)) + })?; + Ok(rewritten.data) + } + + fn name(&self) -> &str { + "ScopedPageIndexOptimizer" + } + + /// We swap a reader factory only; the scan's output schema is unchanged. + fn schema_check(&self) -> bool { + true + } +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::datatypes::{DataType, Field, Schema}; + use datafusion::execution::cache::DefaultFilesMetadataCache; + use datafusion::execution::object_store::ObjectStoreUrl; + use datafusion::logical_expr::Operator; + use datafusion::physical_expr::expressions::{lit, BinaryExpr, Column}; + use datafusion::physical_expr::PhysicalExpr; + use datafusion_datasource::table_schema::TableSchema; + use object_store::memory::InMemory; + + fn schema() -> Arc { + Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + ])) + } + + fn deps() -> (Arc, Arc) { + ( + Arc::new(InMemory::new()), + Arc::new(DefaultFilesMetadataCache::new(64 * 1024 * 1024)), + ) + } + + fn datasource_exec(parquet: ParquetSource) -> Arc { + let config = + FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), Arc::new(parquet)).build(); + DataSourceExec::from_data_source(config) + } + + fn predicate_on_a() -> Arc { + Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::Gt, + lit(5i32), + )) + } + + fn parquet_for(sch: &Arc) -> ParquetSource { + ParquetSource::new(TableSchema::new(sch.clone(), vec![])) + } + + fn get_factory(plan: &Arc) -> Option { + let dse = plan.downcast_ref::()?; + let cfg = (dse.data_source().as_ref() as &dyn std::any::Any) + .downcast_ref::()?; + let pq = (cfg.file_source().as_ref() as &dyn std::any::Any) + .downcast_ref::()?; + Some(pq.parquet_file_reader_factory().is_some()) + } + + #[test] + fn installs_factory_when_predicate_present() { + let sch = schema(); + let (store, cache) = deps(); + let parquet = parquet_for(&sch).with_predicate(predicate_on_a()); + let plan = datasource_exec(parquet); + assert_eq!(get_factory(&plan), Some(false), "precondition: no factory yet"); + + let rule = ScopedPageIndexOptimizer::new(store, cache); + let out = rule.optimize(plan, &ConfigOptions::default()).unwrap(); + assert_eq!( + get_factory(&out), + Some(true), + "optimizer must install a scoped reader factory when a predicate is present" + ); + } + + #[test] + fn noop_without_predicate() { + let sch = schema(); + let (store, cache) = deps(); + let plan = datasource_exec(parquet_for(&sch)); // no predicate + let rule = ScopedPageIndexOptimizer::new(store, cache); + let out = rule.optimize(plan, &ConfigOptions::default()).unwrap(); + assert_eq!( + get_factory(&out), + Some(false), + "no predicate → nothing to scope → no factory installed" + ); + } + + /// The rule REPLACES an already-installed factory when a predicate is present + /// (DataFusion's `ParquetFormat` always pre-installs its own). + #[test] + fn replaces_existing_default_factory() { + let sch = schema(); + let (store, cache) = deps(); + let pre = Arc::new(ScopedPageIndexReaderFactory::new( + Arc::clone(&store), + Arc::clone(&cache), + vec!["a".to_string()], + None, + sch.clone(), + )); + let parquet = parquet_for(&sch) + .with_predicate(predicate_on_a()) + .with_parquet_file_reader_factory(pre); + let plan = datasource_exec(parquet); + assert_eq!(get_factory(&plan), Some(true), "precondition: a factory is present"); + + let rule = ScopedPageIndexOptimizer::new(store, cache); + let out = rule.optimize(Arc::clone(&plan), &ConfigOptions::default()).unwrap(); + assert_eq!(get_factory(&out), Some(true), "scoped factory present after rule"); + assert!( + !Arc::ptr_eq(&plan, &out), + "rule must rewrite the scan to install the scoped factory, replacing the default" + ); + } + + /// End-to-end through a real `SessionContext` + stock `ListingTable`: write a + /// parquet file, register it, plan `SELECT s0 WHERE n1 >= k`, apply the rule, + /// execute, and assert (a) results are correct and (b) the shared scoped + /// page-index cache filled — proving the rule installs a working scoped reader + /// on the vanilla listing path. + #[tokio::test] + async fn end_to_end_listing_scan_fills_scoped_cache() { + use arrow::array::{Int32Array, StringArray}; + use arrow::record_batch::RecordBatch; + use datafusion::datasource::file_format::parquet::ParquetFormat; + use datafusion::datasource::listing::{ + ListingOptions, ListingTable, ListingTableConfig, ListingTableUrl, + }; + use datafusion::parquet::arrow::ArrowWriter; + use datafusion::parquet::file::properties::{EnabledStatistics, WriterProperties}; + use datafusion::prelude::SessionContext; + use futures::StreamExt; + + // Serialize on the shared guard — this asserts on the global cache. + let _g = crate::indexed_table::page_index_loader::SCOPED_CACHE_TEST_GUARD + .lock() + .unwrap(); + crate::indexed_table::page_index_loader::clear_scoped_cache_for_test(); + + let sch = Arc::new(Schema::new(vec![ + Field::new("n0", DataType::Int32, false), + Field::new("n1", DataType::Int32, false), + Field::new("s0", DataType::Utf8, false), + Field::new("s1", DataType::Utf8, false), + ])); + const ROWS: i32 = 4096; + let n0: Vec = (0..ROWS).collect(); + let n1: Vec = (0..ROWS).collect(); + let s0: Vec = (0..ROWS).map(|r| format!("s0_{r:06}_padding_padding")).collect(); + let s1: Vec = (0..ROWS).map(|r| format!("s1_{r:06}_padding_padding")).collect(); + let batch = RecordBatch::try_new( + sch.clone(), + vec![ + Arc::new(Int32Array::from(n0)), + Arc::new(Int32Array::from(n1)), + Arc::new(StringArray::from(s0)), + Arc::new(StringArray::from(s1)), + ], + ) + .unwrap(); + + let dir = std::env::temp_dir().join(format!("scoped_e2e_{}", std::process::id())); + let _ = std::fs::create_dir_all(&dir); + let file_path = dir.join("data.parquet"); + { + let props = WriterProperties::builder() + .set_data_page_row_count_limit(256) + .set_write_batch_size(256) + .set_statistics_enabled(EnabledStatistics::Page) + .build(); + let f = std::fs::File::create(&file_path).unwrap(); + let mut w = ArrowWriter::try_new(f, sch.clone(), Some(props)).unwrap(); + w.write(&batch).unwrap(); + w.close().unwrap(); + } + + let ctx = SessionContext::new(); + let store: Arc = Arc::new(object_store::local::LocalFileSystem::new()); + let table_url = ListingTableUrl::parse(format!("file://{}", dir.to_str().unwrap())).unwrap(); + ctx.register_object_store(table_url.as_ref(), Arc::clone(&store)); + let listing_options = ListingOptions::new(Arc::new(ParquetFormat::new())) + .with_file_extension(".parquet") + .with_collect_stat(true); + let resolved = listing_options.infer_schema(&ctx.state(), &table_url).await.unwrap(); + let config = ListingTableConfig::new(table_url.clone()) + .with_listing_options(listing_options) + .with_schema(resolved); + let provider = Arc::new(ListingTable::try_new(config).unwrap()); + ctx.register_table("t", provider).unwrap(); + + // n1 >= 4080 keeps rows 4080..4096 (16 rows); project the non-predicate s0. + let df = ctx.sql("SELECT s0 FROM t WHERE n1 >= 4080").await.unwrap(); + let physical = df.create_physical_plan().await.unwrap(); + + let metadata_cache = ctx.runtime_env().cache_manager.get_file_metadata_cache(); + let rule = ScopedPageIndexOptimizer::new(Arc::clone(&store), metadata_cache); + let physical = rule.optimize(physical, &ConfigOptions::default()).unwrap(); + + let mut stream = + datafusion::physical_plan::execute_stream(physical, ctx.task_ctx()).unwrap(); + let mut rows = 0usize; + while let Some(b) = stream.next().await { + rows += b.unwrap().num_rows(); + } + + assert_eq!(rows, 16, "predicate n1>=4080 must keep 16 rows"); + + let stats = crate::indexed_table::page_index_loader::scoped_cache_stats(); + assert!( + stats.entries >= 1 && stats.used_bytes > 0, + "scoped cache must have filled on the listing path: {stats:?}" + ); + assert!(stats.misses >= 1, "first scan must register a scoped-cache miss: {stats:?}"); + + crate::indexed_table::page_index_loader::clear_scoped_cache_for_test(); + let _ = std::fs::remove_dir_all(&dir); + } +} diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/session_context.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/session_context.rs index bfae9428738bf..749b532ab4044 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/session_context.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/session_context.rs @@ -229,6 +229,19 @@ pub async unsafe fn create_session_context( ); } + // Install the scoped page-index reader factory on every parquet scan + // (provider-agnostic — covers stock ListingTable for `None` and + // ShardTableProvider for `ListingTable`). Registered AFTER ProjectRowIdOptimizer + // so it sees the final DataSourceExec (ProjectRowIdOptimizer rebuilds the scan, + // which would otherwise drop a factory installed before it). The metadata cache + // is the same shared footer cache; the store is this shard's object store. + state_builder = state_builder.with_physical_optimizer_rule(Arc::new( + crate::scoped_page_index_optimizer::ScopedPageIndexOptimizer::new( + Arc::clone(&shard_view.store), + runtime.runtime_env.cache_manager.get_file_metadata_cache(), + ), + )); + let state = state_builder.build(); let ctx = SessionContext::new_with_state(state); diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/shard_scoped_reader.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/shard_scoped_reader.rs new file mode 100644 index 0000000000000..94222622d766b --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/shard_scoped_reader.rs @@ -0,0 +1,379 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +//! Scoped page-index reader factory for the **listing-table** scan path. +//! +//! # Why this exists +//! +//! The listing-table path (`ShardTableProvider` / vanilla `ListingTable`) uses +//! DataFusion's default reader factory, so when page pruning is enabled the +//! `ParquetOpener` loads the **entire** page index (`ColumnIndex` + `OffsetIndex` +//! for *every* column) of each surviving file, every query, and caches none of +//! it. On wide schemas the `ColumnIndex` (per-page string min/max) dominates the +//! native heap. +//! +//! This factory closes that gap using the unified scoped cache +//! ([`crate::indexed_table::page_index_loader`]). The seam is DataFusion's +//! [`ParquetFileReaderFactory`]: the `ParquetOpener` asks the reader for metadata +//! via `get_metadata`, and — per `opener::load_page_index` — if the returned +//! `ParquetMetaData` *already* carries a page index, the opener uses it and skips +//! the full, all-column load. So our reader's `get_metadata`: +//! +//! 1. loads footer-only metadata (shared metadata-cache hit — see +//! [`crate::indexed_table::parquet_bridge::load_parquet_metadata`]), then +//! 2. augments it with a page index scoped to the predicate columns via +//! [`crate::indexed_table::page_index_loader::load_scoped_page_index`] +//! (real `ColumnIndex` for predicate columns, real `OffsetIndex` for all +//! columns), and +//! 3. returns that augmented metadata. +//! +//! The scoped `(file, predicate-columns)` cache is shared with the indexed path, +//! so repeated queries reuse the decoded index across both scan paths. +//! +//! # Why all row groups (no RG scoping here) +//! +//! `ScopedPageIndexOptimizer` only swaps the reader factory; DataFusion still +//! selects which row groups to scan via its OWN RG-statistics pruning, and its +//! page-pruner + reader then dereference `column_index[rg][col]` / +//! `offset_index[rg][col]` for *its* chosen RGs — a set independent of anything +//! we could compute here. Leaving a placeholder entry on an RG DataFusion still +//! touches panics (`page_row_counts.first().unwrap()` on an empty `OffsetIndex`), +//! and its page-index gate is per-FILE (both indexes must be `Some`), so a +//! partial page index would lie to it. So the page index is built for ALL row +//! groups, column-scoped only — heap stays bounded because the heavy +//! `ColumnIndex` is scoped to predicate columns and only the cheap all-column +//! `OffsetIndex` spans every RG (and that is required for correctness at read +//! time anyway). +//! +//! # Fallback +//! +//! If there are no predicate columns, or scoped augmentation fails for a file +//! (no page index, decode/IO error), `get_metadata` returns the footer-only +//! metadata and the opener loads the page index on demand exactly as today — +//! correct, just without the scoping benefit for that file. Never a wrong result. + +use std::sync::Arc; + +use arrow::datatypes::SchemaRef; +use datafusion::datasource::physical_plan::parquet::{ParquetFileMetrics, ParquetFileReaderFactory}; +use datafusion::execution::cache::cache_manager::FileMetadataCache; +use datafusion::parquet::arrow::arrow_reader::ArrowReaderOptions; +use datafusion::parquet::arrow::async_reader::AsyncFileReader; +use datafusion::parquet::file::metadata::ParquetMetaData; +use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet; +use datafusion_datasource::PartitionedFile; +use futures::future::BoxFuture; +use futures::FutureExt; +use object_store::{ObjectStore, ObjectStoreExt}; +use prost::bytes::Bytes; + +use crate::indexed_table::page_index_loader::{ + load_scoped_page_index, resolve_predicate_parquet_columns, +}; +use crate::indexed_table::parquet_bridge::load_parquet_metadata; + +/// A [`ParquetFileReaderFactory`] that, on `get_metadata`, returns metadata whose +/// page index is scoped to the query's predicate columns. Data reads go straight +/// to the object store. +/// +/// Carries predicate column *names* + the file schema rather than pre-resolved +/// parquet indices: the reader resolves names → parquet leaf indices per file via +/// the same `resolve_predicate_parquet_columns` the indexed path uses, robust to +/// schema evolution across files (a column absent from one file is just skipped). +#[derive(Debug)] +pub struct ScopedPageIndexReaderFactory { + store: Arc, + metadata_cache: Arc, + /// File-column names referenced by the query predicate. Empty means "no + /// scoping" — `get_metadata` returns footer-only and the opener loads the + /// page index on demand as usual. + predicate_column_names: Arc>, + /// The physical predicate (if any). Retained in the constructor signature for + /// parity with the indexed path, but intentionally NOT used for RG scoping + /// here (see module docs). + #[allow(dead_code)] + predicate: Option>, + /// File schema (no partition columns), for per-file column resolution. + file_schema: SchemaRef, +} + +impl ScopedPageIndexReaderFactory { + pub fn new( + store: Arc, + metadata_cache: Arc, + predicate_column_names: Vec, + predicate: Option>, + file_schema: SchemaRef, + ) -> Self { + Self { + store, + metadata_cache, + predicate_column_names: Arc::new(predicate_column_names), + predicate, + file_schema, + } + } +} + +impl ParquetFileReaderFactory for ScopedPageIndexReaderFactory { + fn create_reader( + &self, + partition_index: usize, + file: PartitionedFile, + _metadata_size_hint: Option, + metrics: &ExecutionPlanMetricsSet, + ) -> datafusion::common::Result> { + let file_metrics = + ParquetFileMetrics::new(partition_index, file.object_meta.location.as_ref(), metrics); + Ok(Box::new(ScopedPageIndexReader { + store: Arc::clone(&self.store), + metadata_cache: Arc::clone(&self.metadata_cache), + predicate_column_names: Arc::clone(&self.predicate_column_names), + file_schema: Arc::clone(&self.file_schema), + location: file.object_meta.location.clone(), + metrics: file_metrics, + })) + } +} + +struct ScopedPageIndexReader { + store: Arc, + metadata_cache: Arc, + predicate_column_names: Arc>, + file_schema: SchemaRef, + location: object_store::path::Path, + metrics: ParquetFileMetrics, +} + +impl AsyncFileReader for ScopedPageIndexReader { + fn get_bytes( + &mut self, + range: std::ops::Range, + ) -> BoxFuture<'_, datafusion::parquet::errors::Result> { + self.metrics.bytes_scanned.add((range.end - range.start) as usize); + let store = Arc::clone(&self.store); + let location = self.location.clone(); + // IO-runtime dispatch is handled by the store wrapper around the + // registered store, so a plain `.await` already runs on the IO runtime. + async move { + store + .get_range(&location, range) + .await + .map_err(|e| datafusion::parquet::errors::ParquetError::External(Box::new(e))) + } + .boxed() + } + + fn get_byte_ranges( + &mut self, + ranges: Vec>, + ) -> BoxFuture<'_, datafusion::parquet::errors::Result>> { + let total: u64 = ranges.iter().map(|r| r.end - r.start).sum(); + self.metrics.bytes_scanned.add(total as usize); + let store = Arc::clone(&self.store); + let location = self.location.clone(); + async move { + store + .get_ranges(&location, &ranges) + .await + .map_err(|e| datafusion::parquet::errors::ParquetError::External(Box::new(e))) + } + .boxed() + } + + fn get_metadata( + &mut self, + _options: Option<&ArrowReaderOptions>, + ) -> BoxFuture<'_, datafusion::parquet::errors::Result>> { + let store = Arc::clone(&self.store); + let metadata_cache = Arc::clone(&self.metadata_cache); + let predicate_names = Arc::clone(&self.predicate_column_names); + let file_schema = Arc::clone(&self.file_schema); + let location = self.location.clone(); + async move { + // 1. Footer-only metadata (shared metadata-cache hit if pre-seeded). + let (_schema, _size, footer) = + load_parquet_metadata(Arc::clone(&store), &location, Arc::clone(&metadata_cache)) + .await + .map_err(|e| { + datafusion::parquet::errors::ParquetError::General(format!( + "footer metadata {}: {e}", + location + )) + })?; + + // 2. Resolve predicate names → this file's parquet leaf indices + // (per-file, so schema evolution across files is handled), then + // augment with an all-RG, column-scoped page index. + if !predicate_names.is_empty() { + let parquet_cols = + resolve_predicate_parquet_columns(&file_schema, &footer, &predicate_names); + if let Some(augmented) = + load_scoped_page_index(&store, &location, &footer, &parquet_cols).await + { + log::debug!( + "scoped-pageidx[listing]: {} rgs={} cols={}", + location, + footer.num_row_groups(), + parquet_cols.len() + ); + return Ok(augmented); + } + } + + Ok(footer) + } + .boxed() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{Int32Array, RecordBatch}; + use arrow::datatypes::{DataType, Field, Schema}; + use datafusion::parquet::arrow::ArrowWriter; + use datafusion::parquet::file::page_index::column_index::ColumnIndexMetaData; + use datafusion::parquet::file::properties::{EnabledStatistics, WriterProperties}; + use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet; + use object_store::memory::InMemory; + use object_store::path::Path as ObjPath; + use object_store::{ObjectStore, ObjectStoreExt, PutPayload}; + + // Shared crate-wide guard so all users of the one process-global scoped cache + // mutually exclude. + use crate::indexed_table::page_index_loader::SCOPED_CACHE_TEST_GUARD as SCOPED_TEST_GUARD; + + /// Two int columns (`price`, `qty`), one row group, four 8-row data pages. + fn two_col_parquet() -> (Bytes, SchemaRef) { + let schema = Arc::new(Schema::new(vec![ + Field::new("price", DataType::Int32, false), + Field::new("qty", DataType::Int32, false), + ])); + let prices: Vec = (0..32).collect(); + let qtys: Vec = (100..132).collect(); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(prices)), + Arc::new(Int32Array::from(qtys)), + ], + ) + .unwrap(); + let props = WriterProperties::builder() + .set_max_row_group_size(32) + .set_data_page_row_count_limit(8) + .set_write_batch_size(8) + .set_statistics_enabled(EnabledStatistics::Page) + .build(); + let mut buf: Vec = Vec::new(); + let mut w = ArrowWriter::try_new(&mut buf, schema.clone(), Some(props)).unwrap(); + w.write(&batch).unwrap(); + w.close().unwrap(); + (Bytes::from(buf), schema) + } + + async fn stage(bytes: Bytes) -> (Arc, ObjPath) { + let store: Arc = Arc::new(InMemory::new()); + let loc = ObjPath::from("data.parquet"); + store.put(&loc, PutPayload::from_bytes(bytes)).await.unwrap(); + (store, loc) + } + + fn fresh_cache() -> Arc { + Arc::new(crate::cache::MutexFileMetadataCache::new( + datafusion::execution::cache::DefaultFilesMetadataCache::new(64 * 1024 * 1024), + )) + } + + fn metrics() -> ExecutionPlanMetricsSet { + ExecutionPlanMetricsSet::new() + } + + /// The factory's reader must, on `get_metadata`, return metadata whose page + /// index is scoped to the predicate column (`price`) — real ColumnIndex for + /// `price`, NONE placeholder for `qty` — while keeping a REAL OffsetIndex for + /// BOTH columns. Also fills the shared scoped cache. + #[tokio::test] + async fn get_metadata_returns_scoped_page_index() { + let _g = SCOPED_TEST_GUARD.lock().unwrap(); + crate::indexed_table::page_index_loader::clear_scoped_cache_for_test(); + + let (bytes, schema) = two_col_parquet(); + let (store, loc) = stage(bytes).await; + let factory = ScopedPageIndexReaderFactory::new( + Arc::clone(&store), + fresh_cache(), + vec!["price".to_string()], + None, + schema, + ); + let pf = PartitionedFile::new(loc.as_ref().to_string(), 0); + let m = metrics(); + let mut reader = factory.create_reader(0, pf, None, &m).unwrap(); + + let meta = reader.get_metadata(None).await.unwrap(); + let ci = meta.column_index().expect("augmented metadata has column index"); + let oi = meta.offset_index().expect("augmented metadata has offset index"); + assert!( + !matches!(ci[0][0], ColumnIndexMetaData::NONE), + "predicate col (price) must have a real ColumnIndex" + ); + assert!( + matches!(ci[0][1], ColumnIndexMetaData::NONE), + "non-predicate col (qty) ColumnIndex must be a NONE placeholder" + ); + assert!( + !oi[0][0].page_locations().is_empty() && !oi[0][1].page_locations().is_empty(), + "OffsetIndex must be real for every column" + ); + + let stats = crate::indexed_table::page_index_loader::scoped_cache_stats(); + assert!(stats.entries >= 1 && stats.misses >= 1 && stats.used_bytes > 0); + + crate::indexed_table::page_index_loader::clear_scoped_cache_for_test(); + } + + /// No predicate columns → no scoping happens: `get_metadata` returns the + /// footer load as-is and the scoped cache is never touched. + /// + /// Note: we deliberately do NOT assert the returned metadata has no page + /// index. Until the base metadata-cache strip lands (Step 1e), the shared + /// `load_parquet_metadata` still loads the full page index when a metadata + /// cache is present (DataFusion's `PageIndexPolicy::Optional`). The invariant + /// this reader guarantees with no predicate is "no scoping", i.e. the scoped + /// cache stays empty — which holds before and after 1e. + #[tokio::test] + async fn get_metadata_no_predicate_does_not_scope() { + let _g = SCOPED_TEST_GUARD.lock().unwrap(); + crate::indexed_table::page_index_loader::clear_scoped_cache_for_test(); + + let (bytes, schema) = two_col_parquet(); + let (store, loc) = stage(bytes).await; + let factory = ScopedPageIndexReaderFactory::new( + Arc::clone(&store), + fresh_cache(), + vec![], + None, + schema, + ); + let pf = PartitionedFile::new(loc.as_ref().to_string(), 0); + let m = metrics(); + let mut reader = factory.create_reader(0, pf, None, &m).unwrap(); + + let _meta = reader.get_metadata(None).await.unwrap(); + let stats = crate::indexed_table::page_index_loader::scoped_cache_stats(); + assert_eq!( + (stats.entries, stats.misses, stats.hits), + (0, 0, 0), + "no predicate → scoped cache must be untouched" + ); + + crate::indexed_table::page_index_loader::clear_scoped_cache_for_test(); + } +} diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionSettings.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionSettings.java index 06de123d0a675..417bd55d6469e 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionSettings.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionSettings.java @@ -298,6 +298,7 @@ static String derivePoolMinDefault(Settings settings, int percent) { CacheSettings.STATISTICS_CACHE_EVICTION_TYPE, CacheSettings.METADATA_CACHE_ENABLED, CacheSettings.STATISTICS_CACHE_ENABLED, + CacheSettings.SCOPED_PAGE_INDEX_CACHE_SIZE_LIMIT, // Concurrency gate settings CONCURRENCY_DATANODE_MULTIPLIER, From 591aa4191f69cff9ea9c308fd6eee71dc6af8aad Mon Sep 17 00:00:00 2001 From: G Date: Tue, 16 Jun 2026 20:02:07 +0530 Subject: [PATCH 04/18] Step 1f (listing): stats-backed e2e IT for the scoped page-index cache MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ScopedPageIndexCacheIT provisions app_logs (parquet/composite) and exercises the listing-table scan path (numeric predicate — native, not Lucene- delegated), asserting via GET /_plugins/_analytics_backend_datafusion/stats (cache_stats.scoped_page_index_cache): - the scoped group is always exposed with a positive size_limit_bytes (64mb); - a listing query populates the cache (entry + memory_bytes > 0) at QUERY time; - re-running the IDENTICAL query is a pure cache hit: hits increase while misses, entries, and memory_bytes stay flat (no duplication / no over-allocation); - 10 repeated runs do not grow entries or bytes (bounded, predictable). Assertions are same-method twice-run deltas, not cold-cache baselines: the scoped cache is a process-global singleton with no reset endpoint and IT methods run in randomized order on a preserved cluster, so an entry may already exist from a prior method. Verified the scoped cache is populated ONLY at query time — its sole writer (load_scoped_page_index) is reachable only from the query-time scoped reader, never the refresh/warm path (custom_cache_manager only warms the level-1 metadata cache). 3/3 pass under -Dsandbox.enabled=true -PrustDebug. --- .../analytics/qa/ScopedPageIndexCacheIT.java | 221 ++++++++++++++++++ 1 file changed, 221 insertions(+) create mode 100644 sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/ScopedPageIndexCacheIT.java diff --git a/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/ScopedPageIndexCacheIT.java b/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/ScopedPageIndexCacheIT.java new file mode 100644 index 0000000000000..34031acd26134 --- /dev/null +++ b/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/ScopedPageIndexCacheIT.java @@ -0,0 +1,221 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.analytics.qa; + +import org.opensearch.client.Request; +import org.opensearch.client.Response; + +import java.io.IOException; +import java.util.Locale; +import java.util.Map; + +/** + * End-to-end integration test for the unified scoped parquet page-index cache, + * verified through the node-stats API + * ({@code GET /_plugins/_analytics_backend_datafusion/stats}) — specifically the + * {@code cache_stats.scoped_page_index_cache} group. + * + *

This IT exercises the listing-table scan path (a numeric predicate + * DataFusion evaluates natively over parquet, NOT delegated to Lucene) and + * asserts the empirically-observable cache behaviour. + * + *

Why the assertions are written as same-method, twice-run deltas

+ * + * The scoped cache is a process-global singleton that persists for the + * life of the node, and the test cluster is preserved across methods, which run + * in randomized order. There is no REST endpoint to reset it. On top of that, + * provisioning a dataset triggers a refresh that warms the metadata cache. So a + * test must NOT assume the cache is cold when it starts — an entry for a given + * predicate column may already exist from an earlier method. + * + *

The robust, order-independent signal is therefore measured within + * one method: run a query, snapshot, run the same query again, snapshot. + * After the first run the file's scoped entry is guaranteed present, so the + * second run must be a pure cache hit — hits increase while misses, entries, and + * memory_bytes stay flat. That measures all three (hits / misses / size) without + * depending on global history. + * + *

Run (fast): {@code ./gradlew :sandbox:qa:analytics-engine-rest:integTest + * --tests "*ScopedPageIndexCacheIT" -Dsandbox.enabled=true -PrustDebug}. + */ +public class ScopedPageIndexCacheIT extends AnalyticsRestTestCase { + + private static final Dataset DATASET = new Dataset("app_logs", "app_logs"); + private static boolean dataProvisioned = false; + + @Override + protected void onBeforeQuery() throws IOException { + if (dataProvisioned == false) { + DatasetProvisioner.provision(client(), DATASET); + dataProvisioned = true; + } + } + + // ---- stats helpers --------------------------------------------------- + + /** A point-in-time snapshot of the scoped page-index cache, summed across all nodes. */ + private record ScopedCacheSnapshot(long hits, long misses, long entries, long memoryBytes, long sizeLimitBytes) { + long lookups() { + return hits + misses; + } + } + + @SuppressWarnings("unchecked") + private ScopedCacheSnapshot fetchScopedCache() throws IOException { + Request request = new Request("GET", "/_plugins/_analytics_backend_datafusion/stats"); + Response response = client().performRequest(request); + Map body = assertOkAndParse(response, "datafusion-backend stats"); + + Map nodes = (Map) body.get("nodes"); + assertNotNull("stats response must contain 'nodes'", nodes); + assertFalse("at least one node in stats response", nodes.isEmpty()); + + long hits = 0, misses = 0, entries = 0, memory = 0, limit = 0; + for (Object nodeObj : nodes.values()) { + Map node = (Map) nodeObj; + Map cacheStats = (Map) node.get("cache_stats"); + assertNotNull("node must report cache_stats", cacheStats); + Map scoped = (Map) cacheStats.get("scoped_page_index_cache"); + assertNotNull("cache_stats must contain scoped_page_index_cache", scoped); + + hits += num(scoped, "hit_count"); + misses += num(scoped, "miss_count"); + entries += num(scoped, "entry_count"); + memory += num(scoped, "memory_bytes"); + limit += num(scoped, "size_limit_bytes"); + } + return new ScopedCacheSnapshot(hits, misses, entries, memory, limit); + } + + private static long num(Map obj, String key) { + Object v = obj.get(key); + assertNotNull("scoped_page_index_cache missing field '" + key + "'", v); + return ((Number) v).longValue(); + } + + // ---- tests ----------------------------------------------------------- + + /** + * The scoped page-index cache group must always be present in the stats + * response and advertise a positive byte budget (its configured limit), even + * before any query has populated it. + */ + public void testScopedCacheGroupIsExposedWithBudget() throws IOException { + ScopedCacheSnapshot snap = fetchScopedCache(); + assertTrue( + "scoped_page_index_cache must advertise a positive size_limit_bytes, got " + snap.sizeLimitBytes(), + snap.sizeLimitBytes() > 0 + ); + assertTrue("hit_count must be >= 0", snap.hits() >= 0); + assertTrue("miss_count must be >= 0", snap.misses() >= 0); + assertTrue("entry_count must be >= 0", snap.entries() >= 0); + } + + /** + * A filtered listing-path query populates the scoped cache (after at least one + * run there is a non-empty entry consuming bytes), and re-running the IDENTICAL + * query is served entirely from cache: hits increase while misses, entries, and + * memory_bytes stay flat — the "predictable, no duplication / no + * over-allocation" guarantee. Measures hits, misses, and size together. + */ + public void testSameListingQueryReRunIsPureCacheHit() throws IOException { + // A numeric predicate on `status` stays on the native listing path (it is + // not delegated to Lucene), so it flows through ScopedPageIndexOptimizer + // and the scoped reader. The scoped-cache key is (file, predicate-columns) + // — value-independent — so any `status` predicate maps to the same entry. + String query = "source=" + DATASET.indexName + " | where status >= 400 | stats count() by service_name"; + + // Run #1 — guarantees the file's scoped entry is present afterwards + // (whether this run was a cold miss or already warm from a prior method). + executePpl(query); + ScopedCacheSnapshot afterFirst = fetchScopedCache(); + + assertTrue( + "after a listing query the scoped cache must hold at least one entry, got " + afterFirst.entries(), + afterFirst.entries() >= 1 + ); + assertTrue( + "a populated scoped cache must consume memory_bytes, got " + afterFirst.memoryBytes(), + afterFirst.memoryBytes() > 0 + ); + + // Run #2 — identical query. The entry is already cached, so this run must + // be a pure hit: hits up, everything else flat. + executePpl(query); + ScopedCacheSnapshot afterSecond = fetchScopedCache(); + + assertTrue( + String.format( + Locale.ROOT, + "re-running the same listing query must register cache hits (run1 lookups: h=%d m=%d; run2: h=%d m=%d)", + afterFirst.hits(), + afterFirst.misses(), + afterSecond.hits(), + afterSecond.misses() + ), + afterSecond.hits() > afterFirst.hits() + ); + assertEquals( + "re-running the same query must NOT add scoped-cache misses", + afterFirst.misses(), + afterSecond.misses() + ); + assertEquals( + "re-running the same query must NOT add entries (no duplication)", + afterFirst.entries(), + afterSecond.entries() + ); + assertEquals( + "re-running the same query must NOT grow memory_bytes (no over-allocation)", + afterFirst.memoryBytes(), + afterSecond.memoryBytes() + ); + } + + /** + * The cache stays bounded under repeated identical queries: ten runs of the + * same listing query produce ten lookups but the entry set and the resident + * bytes do not grow after the first populating run. This is the strongest + * "predictable, no over-allocation" assertion the stats API can make from REST. + */ + public void testRepeatedListingQueriesDoNotGrowTheCache() throws IOException { + String query = "source=" + DATASET.indexName + " | where status >= 200 | stats count()"; + + // Populate once, then take the steady-state baseline. + executePpl(query); + ScopedCacheSnapshot baseline = fetchScopedCache(); + assertTrue("entry must exist after first run", baseline.entries() >= 1); + + for (int i = 0; i < 9; i++) { + executePpl(query); + } + ScopedCacheSnapshot after = fetchScopedCache(); + + assertEquals( + "repeated identical queries must not add entries", + baseline.entries(), + after.entries() + ); + assertEquals( + "repeated identical queries must not grow memory_bytes", + baseline.memoryBytes(), + after.memoryBytes() + ); + assertEquals( + "repeated identical queries must not add misses", + baseline.misses(), + after.misses() + ); + assertTrue( + String.format(Locale.ROOT, "the 9 re-runs must register hits (baseline=%d after=%d)", + baseline.hits(), after.hits()), + after.hits() >= baseline.hits() + 9 + ); + } +} From 1afa35a708c841909194e8e436b161b698cb8b1a Mon Sep 17 00:00:00 2001 From: G Date: Tue, 16 Jun 2026 20:08:54 +0530 Subject: [PATCH 05/18] Step 1e: strip the page index at the metadata-cache chokepoint (footer-only) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Make the level-1 metadata cache footer-only (row-group + file stats), so the heavy decoded page index is never retained there. This is the flip that makes the scoped page-index cache (1a/1c) load-bearing and kills the wide-schema heap blowup. - cache.rs: strip_page_index(entry) downcasts CachedParquetMetaData and rebuilds it with set_column_index(None)/set_offset_index(None); no-op (same Arc) if already footer-only. Called in MutexFileMetadataCache::put — the single chokepoint every put flows through (DataFusion scan opener, infer_schema, fetch_statistics, and the warm path all force-decode the full index before put; we refuse to retain it). - custom_cache_manager.rs: corrected the warm-path comment — warming now lands a footer-only entry through the same stripping put. Correctness: page pruning is an optimization; PagePruner::prune_rg early-returns (scans the whole RG) when no page index is present, never a wrong result. Both scan paths rebuild a predicate-scoped page index per query via the shared scoped cache, at query time only — ColumnIndex is never warmed at refresh. Tests: 2 new strip tests (put_strips_page_index_and_get_returns_footer_only, strip_is_noop_for_footer_only_entry); 295 indexed_table::tests_e2e pass with the strip in place (correctness preserved); cache/page_index_loader/shard_scoped_ reader/optimizer/stats all green. --- .../rust/src/cache.rs | 181 +++++++++++++++++- .../rust/src/custom_cache_manager.rs | 13 +- 2 files changed, 187 insertions(+), 7 deletions(-) diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/cache.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/cache.rs index d5fb186acbc51..9e0fc21cbe720 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/cache.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/cache.rs @@ -9,11 +9,13 @@ use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; +use datafusion::datasource::physical_plan::parquet::metadata::CachedParquetMetaData; use datafusion::execution::cache::cache_manager::{ CachedFileMetadataEntry, FileMetadataCache, FileMetadataCacheEntry, }; -use datafusion::execution::cache::DefaultFilesMetadataCache; use datafusion::execution::cache::CacheAccessor; +use datafusion::execution::cache::DefaultFilesMetadataCache; +use datafusion::parquet::file::metadata::ParquetMetaData; use log::error; use object_store::path::Path; @@ -26,6 +28,41 @@ fn log_cache_error(operation: &str, error: &str) { error!("[CACHE ERROR] {} operation failed: {}", operation, error); } +/// Return a cache entry whose `ParquetMetaData` carries footer-only metadata (no +/// `ColumnIndex` / `OffsetIndex`). If the entry already lacks a page index — or +/// isn't a `CachedParquetMetaData` at all — it's returned unchanged (no clone, no +/// rebuild). +/// +/// This is the single chokepoint that enforces the footer-only invariant: every +/// `put` runs the entry through here before it lands in the shared LRU. +fn strip_page_index(entry: CachedFileMetadataEntry) -> CachedFileMetadataEntry { + let Some(cached) = entry + .file_metadata + .as_any() + .downcast_ref::() + else { + return entry; + }; + let meta = cached.parquet_metadata(); + if meta.column_index().is_none() && meta.offset_index().is_none() { + // Already footer-only — keep the existing Arc, avoid a rebuild. + return entry; + } + + // Rebuild without the page index. The heavy decoded `ColumnIndex` / + // `OffsetIndex` are released when the original Arc drops; the footer + // (row-group + column chunk stats) is preserved. + let stripped = ParquetMetaData::clone(meta) + .into_builder() + .set_column_index(None) + .set_offset_index(None) + .build(); + CachedFileMetadataEntry::new( + entry.meta, + Arc::new(CachedParquetMetaData::new(Arc::new(stripped))), + ) +} + // Wrapper to make Mutex implement FileMetadataCache pub struct MutexFileMetadataCache { pub inner: Mutex, @@ -96,6 +133,24 @@ impl CacheAccessor for MutexFileMetadataCache { } fn put(&self, k: &Path, v: CachedFileMetadataEntry) -> Option { + // Enforce the footer-only invariant at the single cache chokepoint. + // + // DataFusion's parquet paths (`infer_schema`, the scan opener, + // `fetch_statistics`) hand this cache to `DFParquetMetadata::fetch_metadata`, + // which force-decodes the FULL page index (`ColumnIndex` + `OffsetIndex` + // for every column of every row group) before calling `put` + // (`datafusion-datasource-parquet/src/metadata.rs`). On wide schemas that + // decoded index dominates the native heap and, since this is a shared LRU + // keyed by path, also evicts the small footer-only entries the scan paths + // depend on. + // + // We can't stop DataFusion from decoding it, but we can refuse to *retain* + // it: strip the page index here so the level-1 cache only ever holds + // footer-only metadata (row-group + file stats). Page-level pruning is + // unaffected — both scan paths rebuild a predicate-scoped page index per + // query through the shared scoped cache + // (`indexed_table::page_index_loader`), at query time only. + let v = strip_page_index(v); match self.inner.lock() { Ok(cache) => cache.put(k, v), Err(e) => { @@ -181,3 +236,127 @@ impl FileMetadataCache for MutexFileMetadataCache { } } } + +#[cfg(test)] +mod strip_page_index_tests { + use super::*; + use datafusion::arrow::array::{Int64Array, RecordBatch}; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::parquet::arrow::arrow_reader::{ArrowReaderMetadata, ArrowReaderOptions}; + use datafusion::parquet::arrow::ArrowWriter; + use datafusion::parquet::file::properties::{EnabledStatistics, WriterProperties}; + use object_store::ObjectMeta; + use prost::bytes::Bytes; + + /// Multi-page parquet bytes (page-level stats → a real page index). + fn parquet_with_page_index() -> Bytes { + let schema = Arc::new(Schema::new(vec![Field::new("v", DataType::Int64, false)])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int64Array::from((0..4096i64).collect::>()))], + ) + .unwrap(); + let props = WriterProperties::builder() + .set_statistics_enabled(EnabledStatistics::Page) + .set_data_page_row_count_limit(128) + .build(); + let mut buf: Vec = Vec::new(); + let mut w = ArrowWriter::try_new(&mut buf, schema, Some(props)).unwrap(); + w.write(&batch).unwrap(); + w.close().unwrap(); + Bytes::from(buf) + } + + fn object_meta(bytes: &Bytes) -> ObjectMeta { + ObjectMeta { + location: Path::from("data.parquet"), + last_modified: chrono::Utc::now(), + size: bytes.len() as u64, + e_tag: None, + version: None, + } + } + + fn full_index_entry(bytes: &Bytes) -> CachedFileMetadataEntry { + let meta = ArrowReaderMetadata::load( + &bytes.clone(), + ArrowReaderOptions::new().with_page_index(true), + ) + .unwrap(); + let pq = meta.metadata().clone(); + assert!( + pq.column_index().is_some() && pq.offset_index().is_some(), + "fixture must carry a page index" + ); + CachedFileMetadataEntry::new(object_meta(bytes), Arc::new(CachedParquetMetaData::new(pq))) + } + + fn page_index_present(entry: &CachedFileMetadataEntry) -> bool { + let cached = entry + .file_metadata + .as_any() + .downcast_ref::() + .unwrap(); + let m = cached.parquet_metadata(); + m.column_index().is_some() || m.offset_index().is_some() + } + + /// `put` of a full-index entry must store it footer-only; a subsequent `get` + /// returns metadata without a page index, but with footer row-group stats. + #[test] + fn put_strips_page_index_and_get_returns_footer_only() { + let bytes = parquet_with_page_index(); + let entry = full_index_entry(&bytes); + assert!(page_index_present(&entry), "precondition: entry has page index"); + + let cache = MutexFileMetadataCache::new(DefaultFilesMetadataCache::new(64 * 1024 * 1024)); + let key = Path::from("data.parquet"); + cache.put(&key, entry); + + let got = cache.get(&key).expect("entry must be retrievable"); + assert!( + !page_index_present(&got), + "cached entry must be footer-only after put" + ); + let cached = got + .file_metadata + .as_any() + .downcast_ref::() + .unwrap(); + let m = cached.parquet_metadata(); + assert!(m.num_row_groups() > 0); + assert!( + m.row_group(0).column(0).statistics().is_some(), + "footer row-group stats must survive the strip" + ); + } + + /// A strip of an already-footer-only entry keeps the SAME inner Arc (no + /// pointless rebuild). + #[test] + fn strip_is_noop_for_footer_only_entry() { + let bytes = parquet_with_page_index(); + let meta = ArrowReaderMetadata::load( + &bytes.clone(), + ArrowReaderOptions::new().with_page_index(false), + ) + .unwrap(); + let pq = meta.metadata().clone(); + assert!(pq.column_index().is_none() && pq.offset_index().is_none()); + let entry = CachedFileMetadataEntry::new( + object_meta(&bytes), + Arc::new(CachedParquetMetaData::new(Arc::clone(&pq))), + ); + + let stripped = strip_page_index(entry); + let cached = stripped + .file_metadata + .as_any() + .downcast_ref::() + .unwrap(); + assert!( + Arc::ptr_eq(cached.parquet_metadata(), &pq), + "footer-only entry must be returned unchanged (same Arc, no rebuild)" + ); + } +} diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/custom_cache_manager.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/custom_cache_manager.rs index 097d3657b9e8e..4673e3e2539b2 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/custom_cache_manager.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/custom_cache_manager.rs @@ -359,12 +359,13 @@ impl CustomCacheManager { let metadata_cache = cache_ref.clone() as Arc; - // Use DataFusion's metadata loading by passing reference to file_metadata_cache to get complete metadata - // IMPORTANT: When a cache is provided to DFParquetMetadata, fetch_metadata() will: - // 1. Enable page index loading (with_page_indexes(true)) - // 2. Load the complete metadata including column and offset indexes - // 3. Automatically put the metadata into the cache (lines 155-160 in datafusion's metadata.rs) - // This ensures we cache exactly what DataFusion would cache during query execution + // Use DataFusion's metadata loading by passing the file_metadata_cache. + // NOTE: When a cache is provided, fetch_metadata() force-decodes the full + // page index (ColumnIndex + OffsetIndex) and calls cache.put() internally. + // Our MutexFileMetadataCache::put strips the page index at that chokepoint, + // so warming this cache lands a FOOTER-ONLY entry (row-group + file stats). + // The heavy decoded page index is released immediately; it is never warmed + // here. The scoped page-index cache is built per query, at query time only. let _parquet_metadata = rt_handle.block_on(async { let df_metadata = DFParquetMetadata::new(store.as_ref(), object_meta) .with_file_metadata_cache(Some(metadata_cache)); From 91e74bdf62db0fd975a67b9d005683d6d4073d6c Mon Sep 17 00:00:00 2001 From: G Date: Tue, 16 Jun 2026 20:16:23 +0530 Subject: [PATCH 06/18] Step 1e: level-1 metadata cache is footer-only WITHOUT decoding the page index MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The decisive fix: do not merely strip the page index after the fact — never decode it for the level-1 cache in the first place. DataFusion 54 hardcodes the page-index policy in fetch_metadata: cache present => PageIndexPolicy::Optional (force-decodes the FULL ColumnIndex+OffsetIndex before caching), no cache => Skip. So passing our cache made it decode the whole page index every footer load, only for put() to throw it away. On wide schemas that decode is the dominant cost the whole effort targets. - parquet_bridge::load_parquet_metadata now manages the cache itself: cache get + is_valid_for => return cached footer (zero I/O, zero decode); miss => fetch WITHOUT a cache (PageIndexPolicy::Skip => page index never decoded), then put a footer-only entry. This is the loader both scoped-reader paths and the warm path use. - custom_cache_manager: warm path now delegates to load_parquet_metadata (footer-only, no decode) instead of DFParquetMetadata::fetch_metadata-with-cache. - cache.rs MutexFileMetadataCache::put keeps strip_page_index as a defensive backstop for the scan paths DataFusion drives directly (its own opener / infer_schema), which still construct fetch_metadata-with-cache out of our reach. Correctness preserved: PagePruner no-ops without a page index; scoped page index is rebuilt per query for predicate columns only. 49 unit + 295 indexed e2e pass. --- .../rust/src/custom_cache_manager.rs | 28 ++++----- .../rust/src/indexed_table/parquet_bridge.rs | 62 +++++++++++++++++-- 2 files changed, 70 insertions(+), 20 deletions(-) diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/custom_cache_manager.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/custom_cache_manager.rs index 4673e3e2539b2..40f2bc1264e95 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/custom_cache_manager.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/custom_cache_manager.rs @@ -15,7 +15,7 @@ use crate::cache::MutexFileMetadataCache; use crate::statistics_cache::CustomStatisticsCache; use object_store::path::Path; use object_store::ObjectMeta; -use datafusion::datasource::physical_plan::parquet::metadata::DFParquetMetadata; +use object_store::ObjectStore; use log::{debug, error}; /// Create ObjectMeta from a local file path. @@ -359,20 +359,18 @@ impl CustomCacheManager { let metadata_cache = cache_ref.clone() as Arc; - // Use DataFusion's metadata loading by passing the file_metadata_cache. - // NOTE: When a cache is provided, fetch_metadata() force-decodes the full - // page index (ColumnIndex + OffsetIndex) and calls cache.put() internally. - // Our MutexFileMetadataCache::put strips the page index at that chokepoint, - // so warming this cache lands a FOOTER-ONLY entry (row-group + file stats). - // The heavy decoded page index is released immediately; it is never warmed - // here. The scoped page-index cache is built per query, at query time only. - let _parquet_metadata = rt_handle.block_on(async { - let df_metadata = DFParquetMetadata::new(store.as_ref(), object_meta) - .with_file_metadata_cache(Some(metadata_cache)); - - // fetch_metadata() performs the cache put operation internally - df_metadata.fetch_metadata().await - .map_err(|e| format!("Failed to fetch metadata: {}", e)) + // Warm the level-1 metadata cache FOOTER-ONLY. `load_parquet_metadata` + // manages the cache itself and fetches with PageIndexPolicy::Skip, so the + // heavy page index (ColumnIndex + OffsetIndex) is NEVER decoded here — only + // row-group + file stats are read and cached. The scoped page index is + // built separately, per query, only for the predicate columns. + let store_dyn: Arc = store.clone(); + let location = object_meta.location.clone(); + let _footer = rt_handle.block_on(async { + crate::indexed_table::parquet_bridge::load_parquet_metadata( + store_dyn, &location, metadata_cache, + ) + .await })?; // Verify the metadata was cached properly diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/parquet_bridge.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/parquet_bridge.rs index 5231187a694e2..72fcb80546dfb 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/parquet_bridge.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/parquet_bridge.rs @@ -24,12 +24,15 @@ use std::time::{Duration, Instant}; use datafusion::arrow::datatypes::SchemaRef; use datafusion::common::Result; -use datafusion::datasource::physical_plan::parquet::metadata::DFParquetMetadata; +use datafusion::datasource::physical_plan::parquet::metadata::{ + CachedParquetMetaData, DFParquetMetadata, +}; use datafusion::datasource::physical_plan::parquet::{ ParquetAccessPlan, ParquetFileMetrics, ParquetFileReaderFactory, RowGroupAccess, }; use datafusion::datasource::physical_plan::ParquetSource; -use datafusion::execution::cache::cache_manager::FileMetadataCache; +use datafusion::execution::cache::cache_manager::{CachedFileMetadataEntry, FileMetadataCache}; +use datafusion::execution::cache::CacheAccessor; use datafusion::execution::object_store::ObjectStoreUrl; use datafusion::execution::SendableRecordBatchStream; use datafusion::parquet::arrow::arrow_reader::{ArrowReaderOptions, RowSelection}; @@ -48,8 +51,29 @@ use prost::bytes::Bytes; // ── Parquet Metadata Loading ───────────────────────────────────────── -/// Load parquet metadata via DataFusion's `DFParquetMetadata`, consulting the -/// caller-supplied `FileMetadataCache`. +/// Load **footer-only** parquet metadata (row-group + file stats, no page index), +/// managing the `FileMetadataCache` ourselves so the page index is **never +/// decoded**. +/// +/// Why we don't just hand the cache to `DFParquetMetadata::fetch_metadata`: in +/// DataFusion 54 that method hardcodes the page-index policy — when a cache is +/// present it uses `PageIndexPolicy::Optional` and force-decodes the FULL page +/// index (ColumnIndex + OffsetIndex, every column, every row group) before +/// caching, and only uses `Skip` when NO cache is passed +/// (`datafusion-datasource-parquet/src/metadata.rs`). Stripping on `put` would +/// discard that index but the expensive decode already happened. On wide schemas +/// that decode is the dominant cost we're trying to avoid. +/// +/// So this fn: consults the cache directly; on a valid hit returns the cached +/// (footer-only) metadata with zero I/O; on a miss it fetches **without** a cache +/// (→ `Skip` policy → the page index is never decoded) and puts the footer-only +/// entry itself. The scoped page index is built separately, per query, only for +/// the predicate columns (see [`super::page_index_loader`]). +/// +/// `MutexFileMetadataCache::put` still strips defensively as a backstop for the +/// scan paths DataFusion drives directly (the opener / `infer_schema`), but this +/// loader makes the warm path and both scoped-reader paths skip the decode +/// entirely. pub async fn load_parquet_metadata( store: Arc, location: &object_store::path::Path, @@ -61,12 +85,40 @@ pub async fn load_parquet_metadata( .map_err(|e| format!("object-store head {}: {}", location, e))?; let size = meta.size; + // Cache hit (validated against current size/last_modified) → no I/O, no decode. + if let Some(cached) = metadata_cache.get(location) { + if cached.is_valid_for(&meta) { + if let Some(cp) = cached + .file_metadata + .as_any() + .downcast_ref::() + { + let pq_meta = Arc::clone(cp.parquet_metadata()); + let file_meta = pq_meta.file_metadata(); + let schema = + parquet_to_arrow_schema(file_meta.schema_descr(), file_meta.key_value_metadata()) + .map_err(|e| format!("parquet_to_arrow_schema {}: {}", location, e))?; + return Ok((Arc::new(schema), size, pq_meta)); + } + } + } + + // Miss → fetch WITHOUT a cache so DataFusion uses PageIndexPolicy::Skip and + // never decodes the page index. let pq_meta = DFParquetMetadata::new(&*store, &meta) - .with_file_metadata_cache(Some(metadata_cache)) .fetch_metadata() .await .map_err(|e| format!("load parquet metadata {}: {}", location, e))?; + // Publish the footer-only entry to the shared cache ourselves. + metadata_cache.put( + location, + CachedFileMetadataEntry::new( + meta.clone(), + Arc::new(CachedParquetMetaData::new(Arc::clone(&pq_meta))), + ), + ); + let file_meta = pq_meta.file_metadata(); let schema = parquet_to_arrow_schema(file_meta.schema_descr(), file_meta.key_value_metadata()) .map_err(|e| format!("parquet_to_arrow_schema {}: {}", location, e))?; From e66e7cfbfa1dafba564e508d3c7795ce629f2768 Mon Sep 17 00:00:00 2001 From: G Date: Tue, 16 Jun 2026 20:21:25 +0530 Subject: [PATCH 07/18] Step 1f: expand stats-backed IT into a thorough no-breakage verification MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Broaden ScopedPageIndexCacheIT from 3 to 8 tests so the page-index changes (scoped cache + footer-only level-1 cache) are verified to break nothing, all via GET /_plugins/_analytics_backend_datafusion/stats: Correctness (exact counts vs known app_logs cardinalities — 200 docs): - total=200; status>=400=103 (numeric listing predicate); log_level=ERROR=115 (keyword); answers stable across cold vs cached re-runs. Level-1 metadata cache still works after the strip: - holds footer entries and registers hits across repeated queries. Scoped cache (hits / misses / size): - same query re-run = pure hit (misses/entries/bytes flat); 10 repeats stay bounded; occupancy never exceeds the 64mb budget. No-breakage query sweep: - plain projection, grouped aggregation, multi-column native filter, full-text match() (Lucene-delegated), and mixed native+match predicate all execute. 8/8 pass under -Dsandbox.enabled=true -PrustDebug; zero panics/errors in node logs; caches configured metadata=250mb stats=100mb scoped=64mb. --- .../analytics/qa/ScopedPageIndexCacheIT.java | 328 ++++++++++++------ 1 file changed, 215 insertions(+), 113 deletions(-) diff --git a/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/ScopedPageIndexCacheIT.java b/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/ScopedPageIndexCacheIT.java index 34031acd26134..6e254abf5793e 100644 --- a/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/ScopedPageIndexCacheIT.java +++ b/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/ScopedPageIndexCacheIT.java @@ -12,34 +12,32 @@ import org.opensearch.client.Response; import java.io.IOException; +import java.util.List; import java.util.Locale; import java.util.Map; /** - * End-to-end integration test for the unified scoped parquet page-index cache, - * verified through the node-stats API - * ({@code GET /_plugins/_analytics_backend_datafusion/stats}) — specifically the - * {@code cache_stats.scoped_page_index_cache} group. + * End-to-end integration test for the unified scoped parquet page-index cache and + * the footer-only level-1 metadata cache, verified through the node-stats API + * ({@code GET /_plugins/_analytics_backend_datafusion/stats}) — the + * {@code cache_stats.metadata_cache} and {@code cache_stats.scoped_page_index_cache} + * groups. * - *

This IT exercises the listing-table scan path (a numeric predicate - * DataFusion evaluates natively over parquet, NOT delegated to Lucene) and - * asserts the empirically-observable cache behaviour. + *

This suite is deliberately broad: it checks query correctness (not + * just cache counters), that the level-1 metadata cache still works after + * the page-index strip, the full scoped-cache hit/miss/size story, and that + * a spread of query shapes (aggregations, multi-column filters, full-text + * {@code match}) all still execute. The goal is to prove the page-index changes + * broke nothing. * - *

Why the assertions are written as same-method, twice-run deltas

+ *

Why assertions are same-method, twice-run deltas

* - * The scoped cache is a process-global singleton that persists for the - * life of the node, and the test cluster is preserved across methods, which run - * in randomized order. There is no REST endpoint to reset it. On top of that, - * provisioning a dataset triggers a refresh that warms the metadata cache. So a - * test must NOT assume the cache is cold when it starts — an entry for a given - * predicate column may already exist from an earlier method. - * - *

The robust, order-independent signal is therefore measured within - * one method: run a query, snapshot, run the same query again, snapshot. - * After the first run the file's scoped entry is guaranteed present, so the - * second run must be a pure cache hit — hits increase while misses, entries, and - * memory_bytes stay flat. That measures all three (hits / misses / size) without - * depending on global history. + * The scoped cache is a process-global singleton that persists for the life + * of the node; the cluster is preserved across methods which run in randomized + * order; there is no reset endpoint; and provisioning triggers a refresh that + * warms the (footer-only) metadata cache. So a method must NOT assume a cold + * cache. The robust signal is measured within one method: run a query, snapshot, + * run the SAME query again, snapshot — the second run must be a pure hit. * *

Run (fast): {@code ./gradlew :sandbox:qa:analytics-engine-rest:integTest * --tests "*ScopedPageIndexCacheIT" -Dsandbox.enabled=true -PrustDebug}. @@ -49,6 +47,11 @@ public class ScopedPageIndexCacheIT extends AnalyticsRestTestCase { private static final Dataset DATASET = new Dataset("app_logs", "app_logs"); private static boolean dataProvisioned = false; + // Ground truth from datasets/app_logs/bulk.json (200 docs). + private static final long TOTAL_DOCS = 200; + private static final long STATUS_GE_400 = 103; + private static final long LEVEL_ERROR = 115; + @Override protected void onBeforeQuery() throws IOException { if (dataProvisioned == false) { @@ -59,15 +62,11 @@ protected void onBeforeQuery() throws IOException { // ---- stats helpers --------------------------------------------------- - /** A point-in-time snapshot of the scoped page-index cache, summed across all nodes. */ - private record ScopedCacheSnapshot(long hits, long misses, long entries, long memoryBytes, long sizeLimitBytes) { - long lookups() { - return hits + misses; - } - } + /** A point-in-time snapshot of one cache group, summed across all nodes. */ + private record CacheGroup(long hits, long misses, long entries, long memoryBytes, long sizeLimitBytes) {} @SuppressWarnings("unchecked") - private ScopedCacheSnapshot fetchScopedCache() throws IOException { + private CacheGroup fetchGroup(String groupName) throws IOException { Request request = new Request("GET", "/_plugins/_analytics_backend_datafusion/stats"); Response response = client().performRequest(request); Map body = assertOkAndParse(response, "datafusion-backend stats"); @@ -81,141 +80,244 @@ private ScopedCacheSnapshot fetchScopedCache() throws IOException { Map node = (Map) nodeObj; Map cacheStats = (Map) node.get("cache_stats"); assertNotNull("node must report cache_stats", cacheStats); - Map scoped = (Map) cacheStats.get("scoped_page_index_cache"); - assertNotNull("cache_stats must contain scoped_page_index_cache", scoped); - - hits += num(scoped, "hit_count"); - misses += num(scoped, "miss_count"); - entries += num(scoped, "entry_count"); - memory += num(scoped, "memory_bytes"); - limit += num(scoped, "size_limit_bytes"); + Map group = (Map) cacheStats.get(groupName); + assertNotNull("cache_stats must contain " + groupName, group); + + hits += num(group, "hit_count"); + misses += num(group, "miss_count"); + entries += num(group, "entry_count"); + memory += num(group, "memory_bytes"); + limit += num(group, "size_limit_bytes"); } - return new ScopedCacheSnapshot(hits, misses, entries, memory, limit); + return new CacheGroup(hits, misses, entries, memory, limit); + } + + private CacheGroup scoped() throws IOException { + return fetchGroup("scoped_page_index_cache"); + } + + private CacheGroup metadata() throws IOException { + return fetchGroup("metadata_cache"); } private static long num(Map obj, String key) { Object v = obj.get(key); - assertNotNull("scoped_page_index_cache missing field '" + key + "'", v); + assertNotNull(key + " missing", v); return ((Number) v).longValue(); } - // ---- tests ----------------------------------------------------------- + /** Run a PPL query and return the number of result rows (datarows). */ + @SuppressWarnings("unchecked") + private long rowCount(String ppl) throws IOException { + Map body = executePpl(ppl); + Object dr = body.get("datarows"); + assertNotNull("PPL response must carry datarows: " + ppl, dr); + return ((List) dr).size(); + } + + /** Run a single-row `stats count()`-style aggregation and return its scalar long. */ + @SuppressWarnings("unchecked") + private long scalarAgg(String ppl) throws IOException { + Map body = executePpl(ppl); + List rows = (List) body.get("datarows"); + assertNotNull("agg must carry datarows: " + ppl, rows); + assertEquals("agg must return exactly one row: " + ppl, 1, rows.size()); + List first = (List) rows.get(0); + assertFalse("agg row must have a column", first.isEmpty()); + return ((Number) first.get(0)).longValue(); + } + + // ---- exposure + correctness ----------------------------------------- /** - * The scoped page-index cache group must always be present in the stats - * response and advertise a positive byte budget (its configured limit), even - * before any query has populated it. + * The scoped page-index cache group must always be present with a positive + * byte budget (its configured 64mb limit), even before any query runs. */ public void testScopedCacheGroupIsExposedWithBudget() throws IOException { - ScopedCacheSnapshot snap = fetchScopedCache(); + CacheGroup snap = scoped(); assertTrue( "scoped_page_index_cache must advertise a positive size_limit_bytes, got " + snap.sizeLimitBytes(), snap.sizeLimitBytes() > 0 ); - assertTrue("hit_count must be >= 0", snap.hits() >= 0); - assertTrue("miss_count must be >= 0", snap.misses() >= 0); - assertTrue("entry_count must be >= 0", snap.entries() >= 0); + assertTrue("hit_count >= 0", snap.hits() >= 0); + assertTrue("miss_count >= 0", snap.misses() >= 0); + assertTrue("entry_count >= 0", snap.entries() >= 0); } /** - * A filtered listing-path query populates the scoped cache (after at least one - * run there is a non-empty entry consuming bytes), and re-running the IDENTICAL - * query is served entirely from cache: hits increase while misses, entries, and - * memory_bytes stay flat — the "predictable, no duplication / no - * over-allocation" guarantee. Measures hits, misses, and size together. + * The page-index changes must not change query answers. Exact-count + * assertions over the listing path (numeric + keyword predicates) against + * known dataset cardinalities. */ - public void testSameListingQueryReRunIsPureCacheHit() throws IOException { - // A numeric predicate on `status` stays on the native listing path (it is - // not delegated to Lucene), so it flows through ScopedPageIndexOptimizer - // and the scoped reader. The scoped-cache key is (file, predicate-columns) - // — value-independent — so any `status` predicate maps to the same entry. - String query = "source=" + DATASET.indexName + " | where status >= 400 | stats count() by service_name"; + public void testListingQueryCorrectnessUnchanged() throws IOException { + assertEquals( + "total doc count must be exact", + TOTAL_DOCS, + scalarAgg("source=" + DATASET.indexName + " | stats count()") + ); + assertEquals( + "status >= 400 count must be exact (numeric listing predicate)", + STATUS_GE_400, + scalarAgg("source=" + DATASET.indexName + " | where status >= 400 | stats count()") + ); + assertEquals( + "log_level = 'ERROR' count must be exact (keyword predicate)", + LEVEL_ERROR, + scalarAgg("source=" + DATASET.indexName + " | where log_level = 'ERROR' | stats count()") + ); + } - // Run #1 — guarantees the file's scoped entry is present afterwards - // (whether this run was a cold miss or already warm from a prior method). - executePpl(query); - ScopedCacheSnapshot afterFirst = fetchScopedCache(); + /** + * The same correctness must hold when the SAME query is re-run (served partly + * from the scoped cache) — a cached page index must never change the answer. + */ + public void testCorrectnessIsStableAcrossCachedReRuns() throws IOException { + String q = "source=" + DATASET.indexName + " | where status >= 400 | stats count()"; + long first = scalarAgg(q); + long second = scalarAgg(q); + assertEquals("cold and warm runs must agree", first, second); + assertEquals("warm run must still be exact", STATUS_GE_400, second); + } + + // ---- level-1 metadata cache still works ----------------------------- + + /** + * The footer-only level-1 metadata cache must still function: after repeated + * queries it holds entries and registers hits (footers are reused, not + * re-read). This guards against the page-index strip accidentally breaking + * normal footer caching. + */ + public void testMetadataCacheStillServesFooters() throws IOException { + String q = "source=" + DATASET.indexName + " | where status >= 200 | stats count() by service_name"; + // Warm. + executePpl(q); + CacheGroup before = metadata(); + // Several more runs — footers must come from cache, driving hits up. + for (int i = 0; i < 5; i++) { + executePpl(q); + } + CacheGroup after = metadata(); assertTrue( - "after a listing query the scoped cache must hold at least one entry, got " + afterFirst.entries(), - afterFirst.entries() >= 1 + "metadata cache must hold at least one footer entry, got " + after.entries(), + after.entries() >= 1 ); assertTrue( - "a populated scoped cache must consume memory_bytes, got " + afterFirst.memoryBytes(), - afterFirst.memoryBytes() > 0 + String.format(Locale.ROOT, "metadata cache must register hits across repeated queries (before=%d after=%d)", + before.hits(), after.hits()), + after.hits() > before.hits() ); + assertTrue( + "metadata cache must advertise its configured byte budget", + after.sizeLimitBytes() > 0 + ); + } + + // ---- scoped cache: populate, hit, bounded --------------------------- + + /** + * A filtered listing query populates the scoped cache (entry + bytes > 0) + * at query time, and re-running the IDENTICAL query is a pure hit: hits up; + * misses, entries, and memory_bytes flat. Measures hits, misses, AND size. + */ + public void testSameListingQueryReRunIsPureCacheHit() throws IOException { + String query = "source=" + DATASET.indexName + " | where status >= 400 | stats count() by service_name"; + + executePpl(query); + CacheGroup afterFirst = scoped(); + assertTrue("scoped cache must hold >= 1 entry after a listing query", afterFirst.entries() >= 1); + assertTrue("populated scoped cache must consume memory_bytes", afterFirst.memoryBytes() > 0); - // Run #2 — identical query. The entry is already cached, so this run must - // be a pure hit: hits up, everything else flat. executePpl(query); - ScopedCacheSnapshot afterSecond = fetchScopedCache(); + CacheGroup afterSecond = scoped(); assertTrue( - String.format( - Locale.ROOT, - "re-running the same listing query must register cache hits (run1 lookups: h=%d m=%d; run2: h=%d m=%d)", - afterFirst.hits(), - afterFirst.misses(), - afterSecond.hits(), - afterSecond.misses() - ), + String.format(Locale.ROOT, "re-run must register hits (h1=%d m1=%d h2=%d m2=%d)", + afterFirst.hits(), afterFirst.misses(), afterSecond.hits(), afterSecond.misses()), afterSecond.hits() > afterFirst.hits() ); - assertEquals( - "re-running the same query must NOT add scoped-cache misses", - afterFirst.misses(), - afterSecond.misses() - ); - assertEquals( - "re-running the same query must NOT add entries (no duplication)", - afterFirst.entries(), - afterSecond.entries() - ); - assertEquals( - "re-running the same query must NOT grow memory_bytes (no over-allocation)", - afterFirst.memoryBytes(), - afterSecond.memoryBytes() - ); + assertEquals("re-run must NOT add misses", afterFirst.misses(), afterSecond.misses()); + assertEquals("re-run must NOT add entries (no duplication)", afterFirst.entries(), afterSecond.entries()); + assertEquals("re-run must NOT grow memory_bytes (no over-alloc)", afterFirst.memoryBytes(), afterSecond.memoryBytes()); } /** - * The cache stays bounded under repeated identical queries: ten runs of the - * same listing query produce ten lookups but the entry set and the resident - * bytes do not grow after the first populating run. This is the strongest - * "predictable, no over-allocation" assertion the stats API can make from REST. + * Repeated identical queries keep the cache bounded: entries and bytes do not + * grow after the first populating run; the re-runs all register hits. */ public void testRepeatedListingQueriesDoNotGrowTheCache() throws IOException { String query = "source=" + DATASET.indexName + " | where status >= 200 | stats count()"; - // Populate once, then take the steady-state baseline. executePpl(query); - ScopedCacheSnapshot baseline = fetchScopedCache(); + CacheGroup baseline = scoped(); assertTrue("entry must exist after first run", baseline.entries() >= 1); for (int i = 0; i < 9; i++) { executePpl(query); } - ScopedCacheSnapshot after = fetchScopedCache(); + CacheGroup after = scoped(); - assertEquals( - "repeated identical queries must not add entries", - baseline.entries(), - after.entries() - ); - assertEquals( - "repeated identical queries must not grow memory_bytes", - baseline.memoryBytes(), - after.memoryBytes() - ); - assertEquals( - "repeated identical queries must not add misses", - baseline.misses(), - after.misses() - ); + assertEquals("repeated queries must not add entries", baseline.entries(), after.entries()); + assertEquals("repeated queries must not grow memory_bytes", baseline.memoryBytes(), after.memoryBytes()); + assertEquals("repeated queries must not add misses", baseline.misses(), after.misses()); assertTrue( String.format(Locale.ROOT, "the 9 re-runs must register hits (baseline=%d after=%d)", baseline.hits(), after.hits()), after.hits() >= baseline.hits() + 9 ); } + + /** + * The scoped cache must never exceed its configured byte budget — a basic + * "no over-allocation" invariant readable from the stats API. + */ + public void testScopedCacheStaysWithinBudget() throws IOException { + // Exercise a few distinct predicate shapes to build whatever entries the + // listing path produces, then assert occupancy <= budget. + executePpl("source=" + DATASET.indexName + " | where status >= 400 | stats count()"); + executePpl("source=" + DATASET.indexName + " | where status < 300 | stats count()"); + CacheGroup snap = scoped(); + assertTrue( + String.format(Locale.ROOT, "scoped cache memory_bytes (%d) must stay within size_limit_bytes (%d)", + snap.memoryBytes(), snap.sizeLimitBytes()), + snap.memoryBytes() <= snap.sizeLimitBytes() + ); + } + + // ---- no-breakage query sweep ---------------------------------------- + + /** + * A spread of query shapes must all execute successfully (HTTP 200, parseable + * datarows) with the page-index changes in place: plain projection, + * aggregation, multi-column filter, full-text match (Lucene-delegated path), + * and a mixed predicate. This is the broad "nothing is broken" guard. + */ + public void testVariedQueryShapesAllExecute() throws IOException { + String idx = DATASET.indexName; + + // Plain projection (no predicate). + assertEquals("plain projection returns all docs", TOTAL_DOCS, rowCount("source=" + idx + " | fields service_name, status")); + + // Aggregation with grouping. + assertTrue( + "grouped aggregation must return at least one bucket", + rowCount("source=" + idx + " | stats count() by service_name") >= 1 + ); + + // Multi-column native predicate (listing path). + assertEquals( + "multi-column filter must be exact", + scalarAgg("source=" + idx + " | where status >= 400 and log_level = 'ERROR' | stats count()"), + scalarAgg("source=" + idx + " | where log_level = 'ERROR' and status >= 400 | stats count()") + ); + + // Full-text match — exercises the Lucene-delegated path; must not error. + // (Count is data-dependent; we only assert it executes and is non-negative.) + long matchCount = scalarAgg("source=" + idx + " | where match(message, 'timeout') | stats count()"); + assertTrue("match() query must execute and return a non-negative count", matchCount >= 0); + + // Mixed: native predicate + full-text, then aggregate. + long mixed = scalarAgg("source=" + idx + " | where status >= 400 or match(message, 'error') | stats count()"); + assertTrue("mixed predicate query must execute", mixed >= 0); + } } From 6c403ef0265711a1a2454ec27ac1fed4d7c233ba Mon Sep 17 00:00:00 2001 From: G Date: Tue, 16 Jun 2026 20:30:59 +0530 Subject: [PATCH 08/18] Step 1d: wire the indexed path to the scoped page-index cache MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit After Step 1e the level-1 metadata cache is footer-only, so the indexed PagePruner (which reads column_index()/offset_index() straight off segment.metadata) would see no page index and conservatively no-op. Restore page pruning by re-augmenting each segment metadata with a predicate-scoped page index, built through the SAME shared scoped cache as the listing path. - indexed_executor.rs: after predicate columns are known, collect predicate column NAMES and, per segment, resolve_predicate_parquet_columns + load_scoped_page_index, replacing segment.metadata with the augmented copy. - Resolving NAMES -> per-file parquet leaf indices via the same resolve_predicate_parquet_columns the listing reader uses means both paths compute the identical scoped-cache key for the same (file, predicate columns), so entries are SHARED across paths (constraint #1). All row groups (no RG scoping in Step 1); on any failure the segment keeps footer-only metadata and pruning no-ops — never a wrong result. Tests: 295 indexed_table::tests_e2e pass with re-augmentation; page_index_loader / shard_scoped_reader / cache modules green. --- .../rust/src/indexed_executor.rs | 65 +++++++++++++++++++ 1 file changed, 65 insertions(+) diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_executor.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_executor.rs index faf1615d0b3b0..65e9804b1a7b4 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_executor.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_executor.rs @@ -315,6 +315,34 @@ fn collect_predicate_column_indices(extraction: Option<&ExtractionResult>) -> Ve } indices.into_iter().collect() } + +/// Collect the **names** of the columns referenced by the query predicate, +/// resolved against `schema` (the plan/arrow schema the `Column` indices index +/// into). Used to build the scoped page index per segment via +/// `resolve_predicate_parquet_columns`, which maps names → per-file parquet leaf +/// indices the SAME way the listing path does — so both paths compute the +/// identical scoped-cache key for the same `(file, predicate columns)` and share +/// entries. Returns a sorted, deduped set. +fn collect_predicate_column_names( + extraction: Option<&ExtractionResult>, + schema: &arrow::datatypes::SchemaRef, +) -> Vec { + let Some(e) = extraction else { return vec![] }; + let mut exprs = Vec::new(); + collect_predicate_exprs(&e.tree, &mut exprs); + let mut names = BTreeSet::new(); + for expr in &exprs { + let _ = expr.apply(|node| { + if let Some(col) = node.downcast_ref::() { + if let Some(field) = schema.fields().get(col.index()) { + names.insert(field.name().to_string()); + } + } + Ok(TreeNodeRecursion::Continue) + }); + } + names.into_iter().collect() +} /// For a tree classified as `SingleCollector`, walk it to find the single /// Collector leaf and return its query bytes. fn single_collector_id(tree: &BoolNode) -> Option { @@ -921,6 +949,43 @@ async unsafe fn execute_indexed_with_context_inner( let predicate_columns = collect_predicate_column_indices(extraction.as_ref()); + // Augment each segment's footer-only metadata with a page index scoped to the + // query's predicate columns, so the indexed PagePruner (which reads + // `column_index()` / `offset_index()` straight off `segment.metadata`) can + // page-prune. Since Step 1e the level-1 cache is footer-only, so without this + // the page index would be absent and PagePruner would conservatively no-op + // (scan whole RG) — correct but slow. + // + // The scoped page index is resolved per file via `resolve_predicate_parquet_columns` + // (predicate column NAMES → this file's parquet leaf indices) — identical to + // the listing path's reader — so both paths compute the SAME scoped-cache key + // for the same `(file, predicate columns)` and share entries. All row groups + // (no RG scoping in Step 1). On any failure the segment keeps footer-only + // metadata and pruning no-ops. + let predicate_column_names = collect_predicate_column_names(extraction.as_ref(), &schema); + if !predicate_column_names.is_empty() { + for segment in segments.iter_mut() { + let parquet_cols = crate::indexed_table::page_index_loader::resolve_predicate_parquet_columns( + &schema, + &segment.metadata, + &predicate_column_names, + ); + if parquet_cols.is_empty() { + continue; + } + if let Some(augmented) = crate::indexed_table::page_index_loader::load_scoped_page_index( + &store, + &segment.object_path, + &segment.metadata, + &parquet_cols, + ) + .await + { + segment.metadata = augmented; + } + } + } + let factory: EvaluatorFactory = match classification { FilterClass::None => { // Predicate-only scan: page-pruned universe, residual applied in From bc102bf43688a44ac07be2556d2f394158929ccd Mon Sep 17 00:00:00 2001 From: G Date: Tue, 16 Jun 2026 20:33:01 +0530 Subject: [PATCH 09/18] Step 1d/1f: cross-path sharing IT (listing then indexed = cache HIT) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add testCrossPathSharingListingThenIndexed: a listing query filters `status`, then an indexed query (forced onto the indexed path by match(message,...)) also filters `status`. Asserts via the stats API that the indexed query HITS the scoped entry the listing query built — hits increase, entries and memory_bytes stay flat (no second, path-specific entry). This is constraint #1 (one unified cache shared by both paths), the central failure of the prior iteration, now verified on a real cluster. 9/9 ScopedPageIndexCacheIT pass under -Dsandbox.enabled=true -PrustDebug. --- .../analytics/qa/ScopedPageIndexCacheIT.java | 46 +++++++++++++++++++ 1 file changed, 46 insertions(+) diff --git a/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/ScopedPageIndexCacheIT.java b/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/ScopedPageIndexCacheIT.java index 6e254abf5793e..416547da41566 100644 --- a/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/ScopedPageIndexCacheIT.java +++ b/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/ScopedPageIndexCacheIT.java @@ -320,4 +320,50 @@ public void testVariedQueryShapesAllExecute() throws IOException { long mixed = scalarAgg("source=" + idx + " | where status >= 400 or match(message, 'error') | stats count()"); assertTrue("mixed predicate query must execute", mixed >= 0); } + + // ---- cross-path sharing (constraint #1) ------------------------------ + + /** + * The unified cache must be shared across scan paths: an entry built for a + * predicate column on one path is reused by the other for the same column. + * + *

Here a listing query filters {@code status}, then an indexed query + * (forced onto the indexed path by a {@code match(message, ...)} full-text + * filter) ALSO filters {@code status}. Because both paths resolve the same + * file + predicate column to the same scoped-cache key, the indexed query + * must HIT the entry the listing query built: hits increase while entries do + * not grow (no second, path-specific entry). + */ + public void testCrossPathSharingListingThenIndexed() throws IOException { + String idx = DATASET.indexName; + // Listing path: numeric predicate on `status` populates the scoped entry. + String listingQuery = "source=" + idx + " | where status >= 400 | stats count()"; + executePpl(listingQuery); + CacheGroup afterListing = scoped(); + assertTrue("listing query must populate a scoped entry", afterListing.entries() >= 1); + + // Indexed path: a match() filter forces indexed routing; it also filters + // `status`, so it resolves to the SAME (file, status) scoped-cache key. + String indexedQuery = + "source=" + idx + " | where match(message, 'timeout') and status >= 400 | stats count()"; + executePpl(indexedQuery); + CacheGroup afterIndexed = scoped(); + + assertTrue( + String.format(Locale.ROOT, + "indexed query on the same predicate column must HIT the listing entry (listing hits=%d indexed hits=%d)", + afterListing.hits(), afterIndexed.hits()), + afterIndexed.hits() > afterListing.hits() + ); + assertEquals( + "cross-path reuse must NOT create a second entry for the same (file, predicate column)", + afterListing.entries(), + afterIndexed.entries() + ); + assertEquals( + "cross-path reuse must NOT grow memory_bytes", + afterListing.memoryBytes(), + afterIndexed.memoryBytes() + ); + } } From ccaa1c13a2ea2ec0faf6b2f137e2908002d40800 Mon Sep 17 00:00:00 2001 From: G Date: Tue, 16 Jun 2026 20:55:27 +0530 Subject: [PATCH 10/18] Add scoped page-index cache clear endpoint (testing convenience) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Reset the scoped cache and re-measure via stats without a cluster restart. - page_index_loader::clear_scoped_cache() — drop entries + reset counters, keep the byte budget (the next query repopulates lazily). - ffm.rs df_clear_scoped_page_index_cache() FFI. - NativeBridge.clearScopedPageIndexCache(). - RestClearScopedPageIndexCacheAction: POST /_plugins/_analytics_backend_datafusion/cache/scoped_page_index/_clear (local node; for the single-node benchmark cluster). Registered in DataFusionPlugin.getRestHandlers. Verified live on the 100M-doc clickbench index: populate -> 20 entries/8.1MB; clear -> {"acknowledged":true} -> 0/0/0; repopulate -> identical 20/8.1MB. --- .../rust/src/ffm.rs | 10 +++ .../src/indexed_table/page_index_loader.rs | 14 ++++ .../be/datafusion/DataFusionPlugin.java | 5 +- .../RestClearScopedPageIndexCacheAction.java | 66 +++++++++++++++++++ .../be/datafusion/nativelib/NativeBridge.java | 17 +++++ 5 files changed, 111 insertions(+), 1 deletion(-) create mode 100644 sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/action/stats/RestClearScopedPageIndexCacheAction.java diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs index 7e481846707e5..9d9ad397c544d 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs @@ -1060,6 +1060,16 @@ pub extern "C" fn df_set_scoped_page_index_cache_limit(size_limit: i64) -> i64 { Ok(0) } +/// Clear the process-global scoped page-index cache (drop entries + reset +/// counters, keep the budget). For operational testing — reset and re-measure +/// without a cluster restart. +#[ffm_safe] +#[no_mangle] +pub extern "C" fn df_clear_scoped_page_index_cache() -> i64 { + crate::indexed_table::page_index_loader::clear_scoped_cache(); + Ok(0) +} + #[ffm_safe] #[no_mangle] pub unsafe extern "C" fn df_cache_manager_contains_by_type( diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/page_index_loader.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/page_index_loader.rs index 4ffa50677f84d..1d8272da422b7 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/page_index_loader.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/page_index_loader.rs @@ -291,6 +291,20 @@ pub fn scoped_cache_stats() -> ScopedCacheStats { SCOPED_CACHE.lock().map(|c| c.stats()).unwrap_or_default() } +/// Drop all cached entries and reset the hit/miss/eviction counters, keeping the +/// configured byte budget. Exposed for operational testing (clear the cache and +/// re-measure without a cluster restart). The next query repopulates lazily. +pub fn clear_scoped_cache() { + if let Ok(mut c) = SCOPED_CACHE.lock() { + c.map.clear(); + c.used = 0; + c.tick = 0; + c.hits = 0; + c.misses = 0; + c.evictions = 0; + } +} + // ── Public API ───────────────────────────────────────────────────────────── /// Map the query's arrow predicate-column names to this file's parquet column diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionPlugin.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionPlugin.java index ab8892c23cbfd..7422149267f31 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionPlugin.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionPlugin.java @@ -764,7 +764,10 @@ public List getRestHandlers( if (dataFusionService == null) { return Collections.emptyList(); } - return List.of(new RestDataFusionStatsAction()); + return List.of( + new RestDataFusionStatsAction(), + new org.opensearch.be.datafusion.action.stats.RestClearScopedPageIndexCacheAction() + ); } @Override diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/action/stats/RestClearScopedPageIndexCacheAction.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/action/stats/RestClearScopedPageIndexCacheAction.java new file mode 100644 index 0000000000000..81fbb9de09402 --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/action/stats/RestClearScopedPageIndexCacheAction.java @@ -0,0 +1,66 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.be.datafusion.action.stats; + +import org.opensearch.be.datafusion.nativelib.NativeBridge; +import org.opensearch.transport.client.node.NodeClient; +import org.opensearch.core.rest.RestStatus; +import org.opensearch.core.xcontent.XContentBuilder; +import org.opensearch.rest.BaseRestHandler; +import org.opensearch.rest.BytesRestResponse; +import org.opensearch.rest.RestRequest; + +import java.io.IOException; +import java.util.List; + +import static java.util.Collections.singletonList; +import static org.opensearch.rest.RestRequest.Method.POST; + +/** + * Clears the process-global scoped page-index cache on the local node (drops + * entries + resets counters, keeps the configured budget). + * + *

Operational/testing convenience: reset the cache and re-measure via the + * stats endpoint without a cluster restart. + * + * 

POST /_plugins/_analytics_backend_datafusion/cache/scoped_page_index/_clear
+ * + *

Note: this acts on the node that receives the request only. For a single-node + * benchmark cluster that is exactly what's wanted; for multi-node clusters, send + * it to each node (or restart). + * + * @opensearch.internal + */ +public class RestClearScopedPageIndexCacheAction extends BaseRestHandler { + + private static final String ROUTE = "/_plugins/_analytics_backend_datafusion/cache/scoped_page_index/_clear"; + + @Override + public String getName() { + return "datafusion_clear_scoped_page_index_cache_action"; + } + + @Override + public List routes() { + return singletonList(new Route(POST, ROUTE)); + } + + @Override + protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) throws IOException { + return channel -> { + NativeBridge.clearScopedPageIndexCache(); + XContentBuilder builder = channel.newBuilder(); + builder.startObject(); + builder.field("acknowledged", true); + builder.field("cleared", "scoped_page_index_cache"); + builder.endObject(); + channel.sendResponse(new BytesRestResponse(RestStatus.OK, builder)); + }; + } +} diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/nativelib/NativeBridge.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/nativelib/NativeBridge.java index b0d1b98edf01a..195fa919f7bb2 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/nativelib/NativeBridge.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/nativelib/NativeBridge.java @@ -131,6 +131,7 @@ private static RuntimeException rethrowConverted(RuntimeException e) { private static final MethodHandle CACHE_MANAGER_GET_TOTAL_MEMORY; private static final MethodHandle CACHE_MANAGER_CONTAINS_BY_TYPE; private static final MethodHandle SET_SCOPED_PAGE_INDEX_CACHE_LIMIT; + private static final MethodHandle CLEAR_SCOPED_PAGE_INDEX_CACHE; private static final MethodHandle CREATE_SESSION_CONTEXT; private static final MethodHandle CREATE_SESSION_CONTEXT_INDEXED; private static final MethodHandle CLOSE_SESSION_CONTEXT; @@ -507,6 +508,11 @@ private static RuntimeException rethrowConverted(RuntimeException e) { FunctionDescriptor.of(ValueLayout.JAVA_LONG, ValueLayout.JAVA_LONG) ); + CLEAR_SCOPED_PAGE_INDEX_CACHE = linker.downcallHandle( + lib.find("df_clear_scoped_page_index_cache").orElseThrow(), + FunctionDescriptor.of(ValueLayout.JAVA_LONG) + ); + CANCEL_QUERY = linker.downcallHandle(lib.find("df_cancel_query").orElseThrow(), FunctionDescriptor.ofVoid(ValueLayout.JAVA_LONG)); SET_CANCEL_STATS_THRESHOLD_MS = linker.downcallHandle( @@ -783,6 +789,17 @@ public static void setScopedPageIndexCacheLimit(long sizeLimitBytes) { } } + /** + * Clears the process-global scoped page-index cache (drops entries + resets + * counters, keeps the budget). For operational testing — reset and re-measure + * without a cluster restart. + */ + public static void clearScopedPageIndexCache() { + try (var call = new NativeCall()) { + call.invoke(CLEAR_SCOPED_PAGE_INDEX_CACHE); + } + } + /** * Returns true if the loaded native library exports {@code df_set_spill_limit}. * When false, spill-cap updates require a node restart — the public From 9057b468683322b2fc8d09f4182477c43cc3aac2 Mon Sep 17 00:00:00 2001 From: G Date: Tue, 16 Jun 2026 21:10:49 +0530 Subject: [PATCH 11/18] Step 2 DRAFT (open for review): RG-scope the ColumnIndex (not wired) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Initial, test-backed draft of row-group scoping for the scoped page index, grounded in a close read of DF54 + parquet58 source (HANDOFF_step2_rg_scoping.md has the full analysis + citations). NOT wired into the scan paths yet — the live paths still call the unchanged all-RG load_scoped_page_index. Design (safe per the source analysis): - ColumnIndex (the heap hog: per-page string min/max) is RG-scopable — decode it only for row groups that pass footer RG-stats pruning; NONE elsewhere. NONE is handled gracefully by DF (statistics.rs get_data_page_statistics! _ arm) and the RGs we placeholder are exactly those DF also prunes. - OffsetIndex MUST stay all-RG: an empty OffsetIndex on any RG DataFusion scans panics (statistics.rs:1835 page_locations.last().unwrap()) or breaks reads (failed to skip rows). DF picks the scanned set itself, after + independent of our index load, so we cannot safely omit any RG OffsetIndex. - Survivor set = footer-stats prune (superset of DF scanned set, since DF further applies bloom/range/limit) — deterministic in (file, predicate), so both scan paths compute the identical set → identical key → cross-path sharing preserved. Code: - ScopedKey gains surviving_rgs (empty = Step-1 all-RG; non-empty = Step-2). - surviving_row_groups(footer, schema, predicate): PruningPredicate over a footer-stats PruningStatistics (mirrors DF RowGroupAccessPlanFilter; no page index); conservative all-RG fallback. - load_scoped_page_index_rgs + build_scoped_page_index refactor: ColumnIndex gated per-RG, OffsetIndex always all-RG; size estimate reflects the RG scope. - 4 unit tests (survivor calc, RG-scoped CI / all-RG OI, byte reduction, key identity). load_scoped_page_index unchanged → Step 1 fully backward compatible. Tests: 13 page_index_loader + 295 indexed e2e + cache/reader/optimizer all pass. Step 3 (RG-scope OffsetIndex) deferred — needs an adversarial repro; see handoff §4. --- HANDOFF_step2_rg_scoping.md | 223 +++++++++ .../src/indexed_table/page_index_loader.rs | 464 ++++++++++++++++-- 2 files changed, 646 insertions(+), 41 deletions(-) create mode 100644 HANDOFF_step2_rg_scoping.md diff --git a/HANDOFF_step2_rg_scoping.md b/HANDOFF_step2_rg_scoping.md new file mode 100644 index 0000000000000..3fb73167d5ee8 --- /dev/null +++ b/HANDOFF_step2_rg_scoping.md @@ -0,0 +1,223 @@ +# Step 2/3: Row-Group Scoping the Scoped Page Index — Design Draft (OPEN for review) + +**Status:** DRAFT — research + design complete, initial code sketch only. NOT wired +into the live paths. Built on the verified Step 1 (commits up to `f65d3fbea68`). + +This is the "tricky" step the user flagged. The design below is grounded in a +close read of DataFusion 54.0.0 + parquet 58.1.0 source (citations inline), not +guesswork. **No hacks**: it relies only on documented arrow-rs/DF behavior. + +--- + +## 0. What Step 1 already does (baseline) + +The scoped cache (`indexed_table::page_index_loader`) builds, per `(file, predicate +parquet-cols)`: +- `ColumnIndex`: real for predicate columns, `NONE` placeholder elsewhere — **all + row groups**. +- `OffsetIndex`: real for **all columns, all row groups**. + +Key = `(path, parquet_cols)`. Both scan paths build the identical artifact, so +entries are shared (verified on 100M-doc clickbench: listing then indexed `match` +hits the same 20 entries). + +The heap hog is the **ColumnIndex** (per-page *string* min/max). Step 1 already +scopes it by COLUMN. Step 2 scopes it by ROW GROUP too — decode predicate-column +ColumnIndex only for row groups that can match, not all of them. + +--- + +## 1. The hard constraints (from DF54 + parquet58 source) + +### 1a. Page index is loaded BEFORE DataFusion's RG pruning +`opener/mod.rs` state machine order: +`PrepareFilters → LoadPageIndex (mod.rs:387-389) → PruneWithStatistics +(RowGroupAccessPlanFilter, mod.rs:896-916) → LoadBloomFilters → PruneWithBloomFilters +→ row_groups.build() (mod.rs:1097) → PAGE PRUNE (mod.rs:1102-1114)`. + +⇒ **We do not know DataFusion's surviving-RG set at index-build time.** The page +index must already cover whatever RGs survive DF's stats+bloom pruning. + +### 1b. RowGroupAccessPlanFilter uses NO page index +`row_group_filter.rs` `RowGroupPruningStatistics` uses only footer RG stats + +bloom filters (no `column_index`/`offset_index` refs anywhere). So DF's surviving +set is a function of **footer RG stats + bloom filters + the predicate** — inputs +we also have (minus bloom, which only prunes MORE). + +### 1c. Page pruner iterates DF's surviving set, not all RGs +`page_filter.rs:240-242`: `for row_group_index in access_plan.row_group_indexes()`. +Per-file gate `page_filter.rs:217-226`: if `offset_index().is_none() || +column_index().is_none()` → skip page pruning for the whole file. A `Some(vec)` +with placeholder contents PASSES this gate. + +### 1d. NONE ColumnIndex is SAFE (graceful) +`statistics.rs:603-621` `get_data_page_statistics!`: a `ColumnIndexMetaData::NONE` +hits the `_` arm → `append_nulls(len)` → "no usable stats, keep all pages". No +panic. **BUT** `len = column_offset_index[rg][col].page_locations().len()` +(statistics.rs:1719-1721) — so the **OffsetIndex for that rg/col must still be +real and correct** for the NONE arm to emit a correctly-sized null array. + +### 1e. Empty OffsetIndex is NOT SAFE for any scanned RG +- Prune time: `data_page_row_counts` panics at `statistics.rs:1835` + (`page_locations.last().unwrap()`) if a scanned column's `page_locations` is empty. +- Read time: `in_memory_row_group.rs:88-103` — with a RowSelection active, + `selection.scan_ranges(&offset_index[idx].page_locations)` returns zero ranges → + `arrow_reader/mod.rs:1383/1449` "failed to skip rows, expected N got 0". +- A *missing* inner column (`.get(col) == None`) degrades gracefully + (`PagesPruningStatistics::try_new` returns None, page_filter.rs:508-520), but a + *present-but-empty* `page_locations` panics. + +### Conclusion +| Index | RG-scopable? | Why | +|---|---|---| +| **ColumnIndex** (predicate cols) | **YES** — NONE on non-surviving RGs | NONE is graceful; only loses pruning on RGs we placeholder (and we only placeholder RGs that can't match anyway) | +| **OffsetIndex** (all cols) | **NO** (must stay all-RG) | empty placeholder panics / breaks reads on any RG DF scans, and we can't know DF's set first | + +So **Step 2 = RG-scope the ColumnIndex; keep OffsetIndex all-RG.** This is the +real, safe win: the ColumnIndex is the heap hog. Step 3 (RG-scope the OffsetIndex) +is **probably not safely achievable** on the listing path (see §4) — likely we +stop at Step 2 and explicitly document why OffsetIndex stays all-RG. + +--- + +## 2. Which RG set do we scope the ColumnIndex to? + +We can't read DF's survivor set, but we can **recompute a superset** of it: + +> `surviving_rgs(file, predicate)` = the row groups that pass **footer RG-stats +> pruning** for the predicate. + +- We build the predicate-column ColumnIndex only for `surviving_rgs`; `NONE` + elsewhere. +- DF's actual scanned set = footer-stats pruning **∩ bloom pruning ∩ range/limit** + ⊆ our footer-stats set. So **every RG DF scans has a real ColumnIndex** → full + page-pruning power, zero degradation. +- RGs we placeholdered (`NONE`) are exactly the ones DF also prunes by footer + stats → DF never looks at them at page-prune time anyway → the NONE is never + even dereferenced. And if DF *did* look (it won't, since stats agree), NONE is + still graceful (§1d) because we keep the real OffsetIndex for all RGs. + +**Cross-path sharing preserved:** both paths compute `surviving_rgs` from the +SAME footer stats + the SAME predicate ⇒ identical set ⇒ identical key+artifact. +The key gains a `surviving_rgs` component, but it is a deterministic function of +`(file, predicate)`, so it does not diverge by path (this is precisely what the +prior rejected iteration got wrong — it keyed on path-specific RG sets). + +### Computing footer-stats survivors +DataFusion exposes `PruningPredicate` + `RowGroupAccessPlanFilter`, but the +cleanest reusable primitive is the existing `build_pruning_predicate` + +`PruningPredicate::prune` over a `RowGroupPruningStatistics`-equivalent. The +indexed path ALREADY computes this: `StatsPruneTree.rg_can_match: Vec` +(`page_pruner.rs:642`, built `table_provider.rs:649`). The listing path would need +the same footer-stats prune. Options: +- (a) Reuse `PruningPredicate` against footer RG stats (a small wrapper around + `StatisticsConverter::row_group_mins/maxes`), computing `Vec` per RG. +- (b) Only RG-scope on the indexed path (where `rg_can_match` already exists) and + leave the listing path all-RG. **Rejected** — breaks the "same artifact both + paths" rule and re-introduces key divergence. + +⇒ Implement (a): a shared `surviving_row_groups(footer_meta, arrow_schema, +predicate) -> Vec` used by BOTH paths before `load_scoped_page_index`. + +--- + +## 3. Proposed API changes (incremental, key stays unified) + +```rust +// page_index_loader.rs + +/// RGs that pass footer RG-stats pruning for `predicate`. Superset of DataFusion's +/// scanned set (which further applies bloom/range/limit). Deterministic in +/// (file, predicate) so both scan paths agree. +pub fn surviving_row_groups( + footer_meta: &ParquetMetaData, + arrow_schema: &SchemaRef, + predicate: &Arc, +) -> Vec; + +/// Like load_scoped_page_index, but the predicate-column ColumnIndex is built +/// ONLY for `surviving_rgs` (NONE elsewhere). OffsetIndex stays real for ALL RGs +/// (correctness — see §1e). `surviving_rgs` becomes part of the cache key. +pub async fn load_scoped_page_index_rgs( + store, location, footer_meta, parquet_cols, + surviving_rgs: &[usize], // sorted/deduped; from surviving_row_groups() +) -> Option>; +``` + +Cache key grows to `(path, parquet_cols, surviving_rgs)`. Because `surviving_rgs` +is derived identically on both paths, sharing holds. `load_scoped_page_index` +(Step 1, all-RG) stays as the fallback when no predicate stats are prunable. + +`build_scoped_page_index` change: the `read_rg: Vec` gate already exists in +the Step-1-era reference (`build_augmented_metadata` had a `surviving_rgs` param); +re-introduce it but ONLY for the ColumnIndex fetch/scatter. The OffsetIndex fetch +stays over all RGs. (Step 1's current code fetches both for all RGs — the diff is: +skip the ColumnIndex range-read for non-surviving RGs, leave NONE.) + +Size estimate: `scoped_page_index_size` already sums per-RG; restrict the +ColumnIndex term to `surviving_rgs`, keep the OffsetIndex term over all RGs. + +--- + +## 4. Step 3 (RG-scope the OffsetIndex) — likely NOT safe; decision pending + +To RG-scope the OffsetIndex we'd need every RG with an empty OffsetIndex to be +guaranteed-unscanned by DF. We can't guarantee that on the **listing** path: DF +owns the access plan and decides scanned RGs after the index is loaded; even if +our footer-stats survivor set ⊇ DF's, the RGs we *omit* are ones DF also prunes — +so in theory an omitted RG is never scanned. **But** the panic in §1e fires at +prune time (`data_page_row_counts`) for any RG in `access_plan.row_group_indexes()`, +and that set = our-superset minus bloom/range. The omitted RGs are not in it, so +they're never iterated... which suggests OffsetIndex RG-scoping to the SAME +`surviving_rgs` set might actually be safe (the omitted RGs are exactly the +footer-pruned ones DF also drops before page-pruning AND before read). + +**This needs a dedicated test to confirm**, because it hinges on "DF's +access-plan RG set ⊆ our footer-stats survivor set" being airtight including the +`prune_by_range`/`prune_by_limit`/initial-plan edge cases (mod.rs:896-916, +1090-1097). If confirmed, Step 3 = scope BOTH indexes to `surviving_rgs` (drop the +"OffsetIndex all-RG" special-case). If not, Step 2 is the final state. + +**Recommendation:** ship Step 2 (ColumnIndex RG-scoped, OffsetIndex all-RG) which +is unambiguously safe and captures the heap win. Treat Step 3 as a separate, +test-gated follow-up with an adversarial "DF scans an RG we omitted" repro. + +--- + +## 5. Test plan (Rust unit + stats-backed IT) + +Rust (`page_index_loader` tests): +1. `surviving_row_groups` matches a known footer-stats prune on a 4-RG fixture + (`four_rg_parquet`): `id >= 25` keeps only the RGs whose min/max overlap. +2. `load_scoped_page_index_rgs`: ColumnIndex real on surviving RGs, `NONE` on + pruned RGs; OffsetIndex real on ALL RGs (the §1e safety invariant). +3. Pruning parity: `PagePruner` on the RG-scoped metadata gives the SAME + RowSelection as the full index, for a surviving RG. +4. **Adversarial (the §4 question):** force a NONE-ColumnIndex RG to be page-pruned + by DF and assert no panic + correct rows (drive a real `DataSourceExec` like + `factory_at_construction_executes_through_scoped_reader`). +5. Cross-path key identity: `surviving_row_groups` + `parquet_cols` identical for + the same predicate regardless of which path computes it. + +IT (`ScopedPageIndexCacheIT`, stats-backed, restart for clean baseline): +- A selective predicate (`AdvEngineID = `) should yield a SMALLER scoped + `memory_bytes` than a non-selective one (`AdvEngineID > 0`) over the same file + set — proving RG-scoping shrank the ColumnIndex. (Both correct vs known counts.) +- Re-run = pure hit (entries/bytes flat), as Step 1. +- Cross-path: listing `AdvEngineID=` then indexed `match(URL,..) AND + AdvEngineID=` HITS (same `surviving_rgs` key). + +--- + +## 6. Files to touch (Step 2) +- `rust/src/indexed_table/page_index_loader.rs`: add `surviving_row_groups`, + `load_scoped_page_index_rgs`; RG-gate the ColumnIndex in `build_scoped_page_index`; + extend `ScopedKey` with `surviving_rgs`; size estimate. +- `rust/src/shard_scoped_reader.rs`: compute `surviving_row_groups` (needs the + predicate — already held as `self.predicate`, currently `#[allow(dead_code)]`) + and call `load_scoped_page_index_rgs`. +- `rust/src/indexed_executor.rs`: same — compute survivors from the bool-tree + predicate (or reuse `StatsPruneTree.rg_can_match`) and call the `_rgs` variant. +- Tests as in §5. + +Keep `load_scoped_page_index` (all-RG) as the no-predicate-stats fallback. diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/page_index_loader.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/page_index_loader.rs index 1d8272da422b7..e9be73c6259d8 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/page_index_loader.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/page_index_loader.rs @@ -122,13 +122,21 @@ const DEFAULT_SCOPED_CACHE_LIMIT: usize = 64 * 1024 * 1024; // ── Cache types ────────────────────────────────────────────────────────── /// Cache key: object-store path + the sorted/deduped set of parquet column -/// indices the page index was built for. Row-group scoping is intentionally -/// NOT part of the key in Step 1 — both paths build an all-row-group artifact so -/// the key (and the artifact) is identical across paths. +/// indices the page index was built for + the sorted/deduped set of surviving +/// row groups the (predicate-column) ColumnIndex was built for. +/// +/// `surviving_rgs` is empty for the Step-1 all-row-group artifact +/// ([`load_scoped_page_index`]); non-empty for the Step-2 RG-scoped artifact +/// ([`load_scoped_page_index_rgs`]). It is part of the key because a ColumnIndex +/// built for a different RG set is a different artifact — but because both scan +/// paths derive `surviving_rgs` deterministically from the SAME `(file, +/// predicate)` via [`surviving_row_groups`], they compute the identical key, so +/// cross-path sharing is preserved. #[derive(Clone, PartialEq, Eq, Hash, Debug)] struct ScopedKey { path: String, parquet_cols: Vec, + surviving_rgs: Vec, } /// One cached entry: the decoded page-index pair (full-width over all row groups @@ -344,13 +352,64 @@ pub async fn load_scoped_page_index( location: &object_store::path::Path, footer_meta: &Arc, parquet_cols: &[usize], +) -> Option> { + load_scoped_page_index_inner(store, location, footer_meta, parquet_cols, None).await +} + +/// Like [`load_scoped_page_index`], but the predicate-column `ColumnIndex` is +/// decoded ONLY for the row groups in `surviving_rgs` (those that pass footer +/// RG-stats pruning for the predicate — see [`surviving_row_groups`]); all other +/// row groups get a `NONE` ColumnIndex placeholder. The `OffsetIndex` is STILL +/// built for every column of every row group — it is read at scan time for any +/// row group DataFusion chooses to scan, and DataFusion picks that set itself +/// (after, and independently of, our page-index load), so an empty OffsetIndex on +/// a scanned RG would panic / break reads (see HANDOFF_step2_rg_scoping.md §1e). +/// +/// Safety of the ColumnIndex RG-scoping: `surviving_rgs` is a **superset** of the +/// RGs DataFusion actually scans (DF further applies bloom/range/limit pruning), +/// so every RG DF page-prunes has a real ColumnIndex. RGs we placeholdered are +/// exactly those DF also prunes by footer stats, so the `NONE` is never +/// dereferenced — and even if it were, `NONE` is handled gracefully ("no stats, +/// keep pages") as long as the OffsetIndex stays valid, which it does. +/// +/// `surviving_rgs` becomes part of the cache key, but since both scan paths derive +/// it deterministically from the same `(file, predicate)`, they share entries. +pub async fn load_scoped_page_index_rgs( + store: &Arc, + location: &object_store::path::Path, + footer_meta: &Arc, + parquet_cols: &[usize], + surviving_rgs: &[usize], +) -> Option> { + load_scoped_page_index_inner(store, location, footer_meta, parquet_cols, Some(surviving_rgs)) + .await +} + +async fn load_scoped_page_index_inner( + store: &Arc, + location: &object_store::path::Path, + footer_meta: &Arc, + parquet_cols: &[usize], + surviving_rgs: Option<&[usize]>, ) -> Option> { if parquet_cols.is_empty() { return None; } + // Normalize the surviving-RG set into the key (sorted/deduped). `None` (all + // RGs, Step 1) is encoded as an empty vec, distinct from a real subset. + let key_rgs: Vec = match surviving_rgs { + Some(rgs) => { + let mut v = rgs.to_vec(); + v.sort_unstable(); + v.dedup(); + v + } + None => Vec::new(), + }; let key = ScopedKey { path: location.as_ref().to_string(), parquet_cols: parquet_cols.to_vec(), + surviving_rgs: key_rgs, }; // Cache hit: graft the cached pair onto the caller's footer. No I/O, no @@ -363,7 +422,7 @@ pub async fn load_scoped_page_index( // Miss: range-read + decode the scoped page index. let (column_index, offset_index, size) = - build_scoped_page_index(store, location, footer_meta, parquet_cols).await?; + build_scoped_page_index(store, location, footer_meta, parquet_cols, surviving_rgs).await?; // Graft a copy for the caller; store the originals in the cache. let grafted = graft(footer_meta, column_index.clone(), offset_index.clone()); @@ -394,14 +453,22 @@ fn graft( // ── Build ──────────────────────────────────────────────────────────────── /// Range-read and decode the page index scoped to `parquet_cols`, returning the -/// full-width `(ColumnIndex, OffsetIndex)` pair plus a size estimate. All row -/// groups, all columns for the `OffsetIndex`; predicate columns only (NONE -/// elsewhere) for the `ColumnIndex`. `None` on any unsafe/impossible condition. +/// full-width `(ColumnIndex, OffsetIndex)` pair plus a size estimate. +/// +/// - `OffsetIndex`: real for every column of EVERY row group (always — required +/// for correctness at scan time, see [`load_scoped_page_index_rgs`]). +/// - `ColumnIndex`: real at the predicate-column positions, `NONE` elsewhere. +/// When `surviving_rgs` is `Some(set)`, the (heavy) predicate-column ColumnIndex +/// is decoded ONLY for row groups in `set`; other RGs get `NONE` (Step-2 +/// RG-scoping). When `None`, it is decoded for all RGs (Step-1). +/// +/// `None` on any unsafe/impossible condition. async fn build_scoped_page_index( store: &Arc, location: &object_store::path::Path, footer_meta: &Arc, parquet_cols: &[usize], + surviving_rgs: Option<&[usize]>, ) -> Option<(ParquetColumnIndex, ParquetOffsetIndex, usize)> { let num_rgs = footer_meta.num_row_groups(); if num_rgs == 0 { @@ -413,31 +480,66 @@ async fn build_scoped_page_index( return None; } - // Phase 1: per RG, gather the predicate columns' ColumnIndex chunks and ALL - // columns' OffsetIndex chunks, and compute their union byte ranges for a - // single vectored fetch. Bail to footer-only if any required index range is - // missing (the file has no page index for it). + // Per-RG mask: build the predicate-column ColumnIndex for this RG? + // `None` → all RGs (Step 1). `Some(set)` → only RGs in the set (Step 2); + // out-of-range entries are ignored. The OffsetIndex is built for all RGs + // regardless. + let build_ci: Vec = match surviving_rgs { + None => vec![true; num_rgs], + Some(set) => { + let mut v = vec![false; num_rgs]; + for &r in set { + if r < num_rgs { + v[r] = true; + } + } + v + } + }; + + // Phase 1: per RG, gather ALL columns' OffsetIndex chunks (always) and — for + // RGs we're building the ColumnIndex for — the predicate columns' ColumnIndex + // chunks. Compute union byte ranges for a single vectored fetch. Bail to + // footer-only if any required index range is missing. struct RgPlan { rg_idx: usize, + build_ci: bool, pred_chunks: Vec, all_chunks: Vec, - col_range: Range, + col_range: Option>, off_range: Range, } let mut plans: Vec = Vec::with_capacity(num_rgs); + // Flat fetch list; per RG we always push the offset range, and push the + // column range only when building the CI for that RG. We record each plan's + // buffer offsets explicitly so the decode phase can find them. let mut fetch_ranges: Vec> = Vec::with_capacity(num_rgs * 2); + // (off_buf_idx, Option) per plan. + let mut buf_idx: Vec<(usize, Option)> = Vec::with_capacity(num_rgs); for rg_idx in 0..num_rgs { let rg = footer_meta.row_group(rg_idx); - let pred_chunks: Vec = - parquet_cols.iter().map(|&i| rg.column(i).clone()).collect(); let all_chunks: Vec = (0..num_cols).map(|i| rg.column(i).clone()).collect(); - let col_range = column_index_union(&pred_chunks)?; let off_range = offset_index_union(&all_chunks)?; - fetch_ranges.push(col_range.clone()); + + let off_i = fetch_ranges.len(); fetch_ranges.push(off_range.clone()); + + let (pred_chunks, col_range, col_i) = if build_ci[rg_idx] { + let pred_chunks: Vec = + parquet_cols.iter().map(|&i| rg.column(i).clone()).collect(); + let col_range = column_index_union(&pred_chunks)?; + let col_i = fetch_ranges.len(); + fetch_ranges.push(col_range.clone()); + (pred_chunks, Some(col_range), Some(col_i)) + } else { + (Vec::new(), None, None) + }; + + buf_idx.push((off_i, col_i)); plans.push(RgPlan { rg_idx, + build_ci: build_ci[rg_idx], pred_chunks, all_chunks, col_range, @@ -448,7 +550,7 @@ async fn build_scoped_page_index( return None; } - // Phase 2: one vectored fetch of all RGs' index byte ranges. + // Phase 2: one vectored fetch of all required index byte ranges. let buffers = store.get_ranges(location, &fetch_ranges).await.ok()?; if buffers.len() != fetch_ranges.len() { return None; @@ -467,41 +569,156 @@ async fn build_scoped_page_index( }) .collect(); - for (i, plan) in plans.iter().enumerate() { - let col_buf = buffers[i * 2].clone(); - let off_buf = buffers[i * 2 + 1].clone(); - - let col_reader = BufferChunkReader { - base: plan.col_range.start, - bytes: col_buf, - }; + for (plan, &(off_i, col_i)) in plans.iter().zip(buf_idx.iter()) { + // OffsetIndex — always, for every column of this RG. let off_reader = BufferChunkReader { base: plan.off_range.start, - bytes: off_buf, + bytes: buffers[off_i].clone(), }; - - // `read_columns_indexes` / `read_offset_indexes` are deprecated in - // arrow-rs but are the only PUBLIC API that decodes a *column subset*. - // See module docs + apache/arrow-rs#8643. - #[allow(deprecated)] - let decoded_cols = read_columns_indexes(&col_reader, &plan.pred_chunks).ok()??; #[allow(deprecated)] let decoded_offs = read_offset_indexes(&off_reader, &plan.all_chunks).ok()??; - if decoded_cols.len() != parquet_cols.len() || decoded_offs.len() != num_cols { + if decoded_offs.len() != num_cols { return None; } + offset_index[plan.rg_idx] = decoded_offs; - let col_row = &mut column_index[plan.rg_idx]; - for (k, &parquet_col) in parquet_cols.iter().enumerate() { - col_row[parquet_col] = decoded_cols[k].clone(); + // ColumnIndex — only for RGs we're building it for. + if plan.build_ci { + let col_range = plan.col_range.clone()?; + let col_i = col_i?; + let col_reader = BufferChunkReader { + base: col_range.start, + bytes: buffers[col_i].clone(), + }; + // `read_columns_indexes` / `read_offset_indexes` are deprecated in + // arrow-rs but are the only PUBLIC API that decodes a *column subset*. + // See module docs + apache/arrow-rs#8643. + #[allow(deprecated)] + let decoded_cols = read_columns_indexes(&col_reader, &plan.pred_chunks).ok()??; + if decoded_cols.len() != parquet_cols.len() { + return None; + } + let col_row = &mut column_index[plan.rg_idx]; + for (k, &parquet_col) in parquet_cols.iter().enumerate() { + col_row[parquet_col] = decoded_cols[k].clone(); + } } - offset_index[plan.rg_idx] = decoded_offs; } - let size = scoped_page_index_size(footer_meta, parquet_cols, num_cols); + let size = scoped_page_index_size(footer_meta, parquet_cols, num_cols, surviving_rgs); Some((column_index, offset_index, size)) } +/// Compute the row groups that pass footer RG-statistics pruning for `predicate`. +/// +/// This is a **superset** of the row groups DataFusion will actually scan: +/// DataFusion applies the same footer-stats pruning (`RowGroupAccessPlanFilter`) +/// plus bloom-filter / range / limit pruning, all of which only remove MORE row +/// groups. So scoping the predicate-column ColumnIndex to this set is safe — every +/// RG DataFusion page-prunes or scans has a real ColumnIndex. +/// +/// Returns all row groups (no scoping benefit) if the predicate can't be lowered +/// to a `PruningPredicate` or any RG lacks the stats to evaluate it — never fewer +/// than DataFusion would scan. Deterministic in `(footer_meta, schema, +/// predicate)`, so both scan paths compute the identical set (→ identical cache +/// key → cross-path sharing). +pub fn surviving_row_groups( + footer_meta: &ParquetMetaData, + arrow_schema: &SchemaRef, + predicate: &Arc, +) -> Vec { + use datafusion::parquet::arrow::arrow_reader::statistics::StatisticsConverter; + use datafusion::physical_optimizer::pruning::{PruningPredicate, PruningStatistics}; + use datafusion::scalar::ScalarValue; + use arrow::array::{ArrayRef, BooleanArray, UInt64Array}; + use arrow::datatypes::FieldRef; + + let num_rgs = footer_meta.num_row_groups(); + let all: Vec = (0..num_rgs).collect(); + if num_rgs == 0 { + return all; + } + + // Lower the predicate to a PruningPredicate; on failure, keep all RGs. + let Ok(pp) = PruningPredicate::try_new(Arc::clone(predicate), Arc::clone(arrow_schema)) else { + return all; + }; + + // Footer RG-stats source: min/max/null_count/row_count per RG per column, + // mirroring DataFusion's `RowGroupPruningStatistics` (footer stats only — no + // page index), via the public `StatisticsConverter`. + struct RgStats<'a> { + meta: &'a ParquetMetaData, + schema: &'a SchemaRef, + num_rgs: usize, + } + impl<'a> RgStats<'a> { + fn conv(&self, col: &str) -> Option> { + StatisticsConverter::try_new( + col, + self.schema, + self.meta.file_metadata().schema_descr(), + ) + .ok() + } + } + impl<'a> PruningStatistics for RgStats<'a> { + fn min_values(&self, column: &datafusion::common::Column) -> Option { + let c = self.conv(&column.name)?; + c.row_group_mins(self.meta.row_groups().iter()).ok() + } + fn max_values(&self, column: &datafusion::common::Column) -> Option { + let c = self.conv(&column.name)?; + c.row_group_maxes(self.meta.row_groups().iter()).ok() + } + fn num_containers(&self) -> usize { + self.num_rgs + } + fn null_counts(&self, column: &datafusion::common::Column) -> Option { + let c = self.conv(&column.name)?; + c.row_group_null_counts(self.meta.row_groups().iter()) + .ok() + .map(|a| Arc::new(a) as ArrayRef) + } + fn row_counts(&self) -> Option { + let counts: Vec = self + .meta + .row_groups() + .iter() + .map(|rg| rg.num_rows() as u64) + .collect(); + Some(Arc::new(UInt64Array::from(counts)) as ArrayRef) + } + fn contained( + &self, + _column: &datafusion::common::Column, + _values: &std::collections::HashSet, + ) -> Option { + None + } + } + // Suppress unused-import lint for FieldRef on toolchains where it's not needed. + let _ = std::marker::PhantomData::; + + let stats = RgStats { + meta: footer_meta, + schema: arrow_schema, + num_rgs, + }; + match pp.prune(&stats) { + Ok(mask) => { + let survivors: Vec = mask + .iter() + .enumerate() + .filter_map(|(i, keep)| if *keep { Some(i) } else { None }) + .collect(); + survivors + } + // Pruning evaluation failed → conservatively keep all RGs. + Err(_) => all, + } +} + /// Deterministic size estimate for one cached pair, in bytes. /// /// Uses the **on-disk serialized lengths** from the footer @@ -518,11 +735,21 @@ fn scoped_page_index_size( footer_meta: &ParquetMetaData, parquet_cols: &[usize], num_cols: usize, + surviving_rgs: Option<&[usize]>, ) -> usize { + // ColumnIndex is only built for surviving RGs (Step 2); OffsetIndex for all. + let build_ci = |rg_idx: usize| -> bool { + match surviving_rgs { + None => true, + Some(set) => set.contains(&rg_idx), + } + }; let mut total = 0usize; - for rg in footer_meta.row_groups() { - for &pc in parquet_cols { - total += rg.column(pc).column_index_length().unwrap_or(0).max(0) as usize; + for (rg_idx, rg) in footer_meta.row_groups().iter().enumerate() { + if build_ci(rg_idx) { + for &pc in parquet_cols { + total += rg.column(pc).column_index_length().unwrap_or(0).max(0) as usize; + } } for c in 0..num_cols { total += rg.column(c).offset_index_length().unwrap_or(0).max(0) as usize; @@ -1039,4 +1266,159 @@ mod tests { clear_scoped_cache_for_test(); } + + // ── Step 2: row-group scoping the ColumnIndex ───────────────────────── + + /// 4 row groups of 10 rows (`id` 0..40, `v` = id*2), page size 5 → real page + /// index, RG min/max well separated for footer-stats pruning. + fn four_rg_parquet() -> (Bytes, SchemaRef) { + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("v", DataType::Int32, false), + ])); + let ids: Vec = (0..40).collect(); + let vs: Vec = (0..40).map(|x| x * 2).collect(); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(ids)), Arc::new(Int32Array::from(vs))], + ) + .unwrap(); + let props = WriterProperties::builder() + .set_max_row_group_size(10) + .set_data_page_row_count_limit(5) + .set_write_batch_size(5) + .set_statistics_enabled(EnabledStatistics::Page) + .build(); + let mut buf: Vec = Vec::new(); + let mut w = ArrowWriter::try_new(&mut buf, schema.clone(), Some(props)).unwrap(); + w.write(&batch).unwrap(); + w.close().unwrap(); + (Bytes::from(buf), schema) + } + + /// `surviving_row_groups` returns exactly the RGs whose footer min/max overlap + /// the predicate — a superset of what DataFusion would scan, never fewer. + #[tokio::test] + async fn surviving_row_groups_matches_footer_stats_prune() { + let (bytes, schema) = four_rg_parquet(); + let fo = footer_only(&bytes); + assert_eq!(fo.num_row_groups(), 4); + + // RG0 ids 0..10, RG1 10..20, RG2 20..30, RG3 30..40. + // `id >= 25` → RG2 (20..30, overlaps) + RG3 (30..40) survive; RG0,RG1 pruned. + let p = pred("id", 0, Operator::GtEq, 25); + let survivors = surviving_row_groups(&fo, &schema, &p); + assert_eq!(survivors, vec![2, 3], "id>=25 must keep only RG2,RG3"); + + // `id < 12` → RG0 (0..10) + RG1 (10..20, overlaps) survive. + let p2 = pred("id", 0, Operator::Lt, 12); + assert_eq!(surviving_row_groups(&fo, &schema, &p2), vec![0, 1]); + } + + /// RG-scoped load: ColumnIndex real ONLY on surviving RGs (NONE on pruned), + /// OffsetIndex real on ALL RGs (the §1e safety invariant), and pruning on a + /// surviving RG matches the full index. + #[tokio::test] + async fn rg_scoped_load_builds_column_index_only_for_survivors() { + let _g = CACHE_TEST_GUARD.lock().unwrap(); + clear_scoped_cache_for_test(); + let (bytes, schema) = four_rg_parquet(); + let (store, loc) = stage(bytes.clone()).await; + let fo = footer_only(&bytes); + + let cols = resolve_predicate_parquet_columns(&schema, &fo, &["id".to_string()]); + assert_eq!(cols, vec![0]); + let surviving = vec![2usize, 3usize]; + + let aug = load_scoped_page_index_rgs(&store, &loc, &fo, &cols, &surviving) + .await + .expect("rg-scoped load must succeed"); + let ci = aug.column_index().expect("has column index"); + let oi = aug.offset_index().expect("has offset index"); + assert_eq!(ci.len(), 4); + + for &rg in &surviving { + assert!( + !matches!(ci[rg][0], ColumnIndexMetaData::NONE), + "surviving RG {rg} must have a real ColumnIndex for `id`" + ); + } + for &rg in &[0usize, 1usize] { + assert!( + matches!(ci[rg][0], ColumnIndexMetaData::NONE), + "pruned RG {rg} ColumnIndex must be NONE" + ); + } + // OffsetIndex must be REAL for EVERY RG and EVERY column (safety). + for rg in 0..4 { + for c in 0..2 { + assert!( + !oi[rg][c].page_locations().is_empty(), + "OffsetIndex must be real for all RGs/cols (rg={rg} col={c})" + ); + } + } + + // Pruning on a surviving RG matches the full index. + let full = full_index(&bytes); + let pp = build_pruning_predicate(&pred("id", 0, Operator::GtEq, 25), schema.clone()).unwrap(); + let s = PagePruner::new(&schema, Arc::clone(&aug)).prune_rg(&pp, 2, None); + let f = PagePruner::new(&schema, full).prune_rg(&pp, 2, None); + assert_eq!(s.as_ref().map(kept), f.as_ref().map(kept)); + + clear_scoped_cache_for_test(); + } + + /// RG-scoping shrinks the cached ColumnIndex: a 2-of-4-RG scoped entry uses + /// fewer bytes than the all-RG (Step 1) entry for the same predicate column. + #[tokio::test] + async fn rg_scoping_reduces_cached_bytes() { + let _g = CACHE_TEST_GUARD.lock().unwrap(); + clear_scoped_cache_for_test(); + let (bytes, schema) = four_rg_parquet(); + let (store, loc) = stage(bytes.clone()).await; + let fo = footer_only(&bytes); + let cols = resolve_predicate_parquet_columns(&schema, &fo, &["id".to_string()]); + + // All-RG (Step 1). + let _ = load_scoped_page_index(&store, &loc, &fo, &cols).await.unwrap(); + let all_rg_bytes = scoped_cache_bytes_for_test(); + clear_scoped_cache_for_test(); + + // RG-scoped to 2 of 4 RGs (Step 2). + let _ = load_scoped_page_index_rgs(&store, &loc, &fo, &cols, &[2, 3]).await.unwrap(); + let scoped_bytes = scoped_cache_bytes_for_test(); + + assert!( + scoped_bytes < all_rg_bytes, + "RG-scoped bytes ({scoped_bytes}) must be < all-RG bytes ({all_rg_bytes})" + ); + clear_scoped_cache_for_test(); + } + + /// The RG-scoped entry is keyed by its surviving-RG set: the same `(file, + /// cols, surviving_rgs)` hits; a different surviving-RG set is a distinct + /// entry (so two paths agree only when they compute the same survivor set). + #[tokio::test] + async fn rg_scoped_key_includes_surviving_rgs() { + let _g = CACHE_TEST_GUARD.lock().unwrap(); + clear_scoped_cache_for_test(); + let (bytes, schema) = four_rg_parquet(); + let (store, loc) = stage(bytes.clone()).await; + let fo = footer_only(&bytes); + let cols = resolve_predicate_parquet_columns(&schema, &fo, &["id".to_string()]); + + let _ = load_scoped_page_index_rgs(&store, &loc, &fo, &cols, &[2, 3]).await.unwrap(); + assert_eq!((scoped_cache_stats().misses, scoped_cache_stats().entries), (1, 1)); + + // Same survivor set → hit, no new entry. + let _ = load_scoped_page_index_rgs(&store, &loc, &fo, &cols, &[2, 3]).await.unwrap(); + assert_eq!((scoped_cache_stats().hits, scoped_cache_stats().entries), (1, 1)); + + // Different survivor set → distinct entry (miss). + let _ = load_scoped_page_index_rgs(&store, &loc, &fo, &cols, &[0, 1]).await.unwrap(); + assert_eq!((scoped_cache_stats().misses, scoped_cache_stats().entries), (2, 2)); + + clear_scoped_cache_for_test(); + } } From 12f78155577e2bdcf35d321ca74af1c7a918a600 Mon Sep 17 00:00:00 2001 From: G Date: Tue, 16 Jun 2026 21:55:31 +0530 Subject: [PATCH 12/18] Step 2 DRAFT: column-scope the OffsetIndex to predicate + projection + col 0 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Answer to "do we need the OffsetIndex for ALL columns?": no — only predicate ∪ projection ∪ {0}. Verified against DF54/parquet58 source: the OffsetIndex is dereferenced only by (a) reads, for PROJECTED columns (in_memory_row_group.rs:80), (b) page pruning, for the PREDICATE column (page_filter.rs:512), and (c) the page-skip METRIC, for column 0 (page_filter.rs:255 — affects only the counter, not the RowSelection). Including col 0 keeps the metric identical to stock DataFusion. DRAFT (not wired into scan paths): - ScopedKey gains offset_cols (empty = all columns / Step 1). - load_scoped_page_index_cols(.., offset_cols): defensively unions predicate cols + col 0, clamps/sorts/dedups; collapses to the all-columns key when the union covers every column (shares the Step-1 entry). - build_scoped_page_index decodes the OffsetIndex only for off_cols, scatters to absolute column positions, empty placeholders elsewhere; size estimate follows. - 4 unit tests: offset real only for the scoped set; projected non-predicate read succeeds; cached bytes reduced; full-coverage collapse to all-columns entry. This is orthogonal to RG-scoping (CI by row group) and is the bigger win on wide schemas (402 cols, project a few). Wiring TODO: listing path has the projection in the optimizer; indexed path must thread read_projection to the augmentation site (both paths must pass the same offset_cols to keep cross-path sharing). See HANDOFF_step2_rg_scoping.md §3b. load_scoped_page_index unchanged → Step 1 + the RG-scoping draft remain backward compatible. Tests: 17 page_index_loader + 295 indexed e2e + cache/reader/optimizer all pass. --- HANDOFF_step2_rg_scoping.md | 42 +++ .../src/indexed_table/page_index_loader.rs | 322 ++++++++++++++++-- 2 files changed, 337 insertions(+), 27 deletions(-) diff --git a/HANDOFF_step2_rg_scoping.md b/HANDOFF_step2_rg_scoping.md index 3fb73167d5ee8..cadf57cd3eb2c 100644 --- a/HANDOFF_step2_rg_scoping.md +++ b/HANDOFF_step2_rg_scoping.md @@ -159,6 +159,48 @@ ColumnIndex term to `surviving_rgs`, keep the OffsetIndex term over all RGs. --- +## 3b. OffsetIndex COLUMN-scoping (better than RG-scoping it) — DRAFTED & safe + +User question: "do we need the OffsetIndex for ALL columns, or just projection + +predicate?" Answer from the source: **just `predicate ∪ projection ∪ {0}`.** The +OffsetIndex is dereferenced in exactly three places (DF54 / parquet58): + +| Consumer | Columns | Source | +|---|---|---| +| Read (`InMemoryRowGroup::fetch_ranges`) | **projected** only (`projection.leaf_included(idx)`) | `in_memory_row_group.rs:80-81` | +| Prune (`PagesPruningStatistics::try_new`) | the **predicate** column only (`offset_index[rg][parquet_column_index]`) | `page_filter.rs:512-519` | +| Metric (`total_pages_in_group`) | **column 0** (`offset_index[rg].first()`) — feeds only the page-skip COUNTER, not the produced RowSelection | `page_filter.rs:255-260, 359-361` | + +So: real OffsetIndex needed for `predicate ∪ projection`; add **column 0** to keep +the page-skip metric identical to stock DF (stock DF also keys it off col 0). Any +other column gets an empty placeholder — never projected, never predicate, never +0, so never dereferenced. This is column-scoping, orthogonal to RG-scoping, and is +the bigger win on the 402-col textbench schema (project a handful → tiny offset set). + +**Implemented (DRAFT, not wired):** +- `ScopedKey` gains `offset_cols` (empty = all columns / Step 1). +- `load_scoped_page_index_cols(.., offset_cols)`: the fn defensively unions in + predicate cols + col 0, clamps/sorts/dedups; if the union covers all columns it + COLLAPSES to the empty "all columns" key (shares the Step-1 entry). +- `build_scoped_page_index` decodes the OffsetIndex only for `off_cols`, scatters + to absolute positions, empty placeholders elsewhere; size estimate follows. +- 4 unit tests: offset real only for the scoped set; projected non-predicate read + works; bytes reduced; full-coverage collapse. + +**Wiring (TODO, same projection-availability split as RG-scoping):** +- Listing path: projection is known in the optimizer (`config.projected_schema()` / + FileScanConfig) → resolve to parquet leaf indices, pass to + `load_scoped_page_index_cols`. Easy. +- Indexed path: projection (`read_projection = output ∪ predicate_columns`) is + known in `scan()`, AFTER the `indexed_executor` augmentation site. Thread it to + the augmentation site, OR keep the indexed path on all-columns OffsetIndex (call + `load_scoped_page_index`) until threaded. Cross-path sharing requires BOTH paths + to pass the SAME `offset_cols` for the same query — so wire both together, or + accept divergent keys (separate entries) until both are wired. + +This combines with RG-scoping: a future `load_scoped_page_index_scoped(parquet_cols, +surviving_rgs, offset_cols)` would do both at once (CI RG-scoped, OI col-scoped). + ## 4. Step 3 (RG-scope the OffsetIndex) — likely NOT safe; decision pending To RG-scope the OffsetIndex we'd need every RG with an empty OffsetIndex to be diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/page_index_loader.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/page_index_loader.rs index e9be73c6259d8..579ac5757aa53 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/page_index_loader.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/page_index_loader.rs @@ -137,6 +137,13 @@ struct ScopedKey { path: String, parquet_cols: Vec, surviving_rgs: Vec, + /// The (normalized) set of columns the `OffsetIndex` was built for. Empty = + /// all columns (Step 1). Non-empty = `predicate ∪ projection ∪ {0}` (Step 2 + /// column-scoping). Part of the key because an OffsetIndex built for a + /// different column set is a different artifact; both scan paths derive it + /// deterministically (predicate + projection are known per query), so the key + /// stays identical across paths for the same logical request. + offset_cols: Vec, } /// One cached entry: the decoded page-index pair (full-width over all row groups @@ -353,7 +360,44 @@ pub async fn load_scoped_page_index( footer_meta: &Arc, parquet_cols: &[usize], ) -> Option> { - load_scoped_page_index_inner(store, location, footer_meta, parquet_cols, None).await + load_scoped_page_index_inner(store, location, footer_meta, parquet_cols, None, None).await +} + +/// Like [`load_scoped_page_index`], but the `OffsetIndex` is decoded only for the +/// columns in `offset_cols` (others get an empty placeholder), instead of all +/// columns. +/// +/// # Why this is safe (and which columns must be included) +/// +/// The `OffsetIndex` is dereferenced in exactly three places in DataFusion 54 / +/// parquet 58: +/// - **Read** (`InMemoryRowGroup::fetch_ranges`): only for **projected** columns +/// (`projection.leaf_included(idx)`), and only when a `RowSelection` is active. +/// - **Prune** (`page_filter::PagesPruningStatistics`): only for the **predicate** +/// column being pruned (`offset_index[rg][parquet_column_index]`). +/// - **Metric** (`page_filter` `total_pages_in_group`): reads **column 0**'s page +/// count to seed the `page_index_rows_pruned` counter. This does NOT affect the +/// produced `RowSelection` (correctness), only the skipped-pages metric. +/// +/// So a real `OffsetIndex` is required for `predicate ∪ projection`; including +/// **column 0** additionally keeps the page-skip metric identical to stock +/// DataFusion. This fn always unions in the predicate columns and column 0 +/// defensively, so callers only need to pass the projection (∪ predicate). Any +/// column NOT in the final set gets an empty placeholder — safe because it is +/// never projected, never the predicate, and never column 0, so it is never +/// dereferenced. +/// +/// `offset_cols` joins the cache key. Callers MUST derive it deterministically +/// from `(predicate, projection)` so both scan paths agree and share entries. +pub async fn load_scoped_page_index_cols( + store: &Arc, + location: &object_store::path::Path, + footer_meta: &Arc, + parquet_cols: &[usize], + offset_cols: &[usize], +) -> Option> { + load_scoped_page_index_inner(store, location, footer_meta, parquet_cols, None, Some(offset_cols)) + .await } /// Like [`load_scoped_page_index`], but the predicate-column `ColumnIndex` is @@ -381,8 +425,15 @@ pub async fn load_scoped_page_index_rgs( parquet_cols: &[usize], surviving_rgs: &[usize], ) -> Option> { - load_scoped_page_index_inner(store, location, footer_meta, parquet_cols, Some(surviving_rgs)) - .await + load_scoped_page_index_inner( + store, + location, + footer_meta, + parquet_cols, + Some(surviving_rgs), + None, + ) + .await } async fn load_scoped_page_index_inner( @@ -391,10 +442,13 @@ async fn load_scoped_page_index_inner( footer_meta: &Arc, parquet_cols: &[usize], surviving_rgs: Option<&[usize]>, + offset_cols: Option<&[usize]>, ) -> Option> { if parquet_cols.is_empty() { return None; } + let num_cols = footer_meta.file_metadata().schema_descr().num_columns(); + // Normalize the surviving-RG set into the key (sorted/deduped). `None` (all // RGs, Step 1) is encoded as an empty vec, distinct from a real subset. let key_rgs: Vec = match surviving_rgs { @@ -406,10 +460,39 @@ async fn load_scoped_page_index_inner( } None => Vec::new(), }; + + // Normalize the OffsetIndex column set. `None` (all columns, Step 1) → empty + // sentinel. `Some(cols)` → cols ∪ predicate-cols ∪ {0}, clamped to in-range + // and sorted/deduped (the defensive union guarantees prune + metric safety + // regardless of what the caller passed). If the union covers every column we + // collapse to the empty "all columns" sentinel so it shares the Step-1 entry. + let key_offset_cols: Vec = match offset_cols { + None => Vec::new(), + Some(cols) => { + let mut set: std::collections::BTreeSet = std::collections::BTreeSet::new(); + set.insert(0); // metric reads column 0 + for &c in parquet_cols { + set.insert(c); + } // prune reads predicate cols + for &c in cols { + if c < num_cols { + set.insert(c); + } + } // read needs projected cols + let v: Vec = set.into_iter().filter(|&c| c < num_cols).collect(); + if v.len() == num_cols { + Vec::new() // full coverage → reuse the all-columns artifact + } else { + v + } + } + }; + let key = ScopedKey { path: location.as_ref().to_string(), parquet_cols: parquet_cols.to_vec(), surviving_rgs: key_rgs, + offset_cols: key_offset_cols.clone(), }; // Cache hit: graft the cached pair onto the caller's footer. No I/O, no @@ -420,9 +503,20 @@ async fn load_scoped_page_index_inner( } } + // Pass the normalized offset-col set (empty = all columns) to the builder. + let offset_cols_for_build: Option<&[usize]> = + if key_offset_cols.is_empty() { None } else { Some(&key_offset_cols) }; + // Miss: range-read + decode the scoped page index. - let (column_index, offset_index, size) = - build_scoped_page_index(store, location, footer_meta, parquet_cols, surviving_rgs).await?; + let (column_index, offset_index, size) = build_scoped_page_index( + store, + location, + footer_meta, + parquet_cols, + surviving_rgs, + offset_cols_for_build, + ) + .await?; // Graft a copy for the caller; store the originals in the cache. let grafted = graft(footer_meta, column_index.clone(), offset_index.clone()); @@ -455,8 +549,11 @@ fn graft( /// Range-read and decode the page index scoped to `parquet_cols`, returning the /// full-width `(ColumnIndex, OffsetIndex)` pair plus a size estimate. /// -/// - `OffsetIndex`: real for every column of EVERY row group (always — required -/// for correctness at scan time, see [`load_scoped_page_index_rgs`]). +/// - `OffsetIndex`: real for the columns in `offset_cols` (the caller-normalized +/// `predicate ∪ projection ∪ {0}`), empty placeholder elsewhere. `None` → all +/// columns (Step 1 / no projection known). Built for EVERY row group (an +/// OffsetIndex omitted on a scanned RG would break reads — DataFusion chooses +/// the scanned set itself, after our load). /// - `ColumnIndex`: real at the predicate-column positions, `NONE` elsewhere. /// When `surviving_rgs` is `Some(set)`, the (heavy) predicate-column ColumnIndex /// is decoded ONLY for row groups in `set`; other RGs get `NONE` (Step-2 @@ -469,6 +566,7 @@ async fn build_scoped_page_index( footer_meta: &Arc, parquet_cols: &[usize], surviving_rgs: Option<&[usize]>, + offset_cols: Option<&[usize]>, ) -> Option<(ParquetColumnIndex, ParquetOffsetIndex, usize)> { let num_rgs = footer_meta.num_row_groups(); if num_rgs == 0 { @@ -483,7 +581,7 @@ async fn build_scoped_page_index( // Per-RG mask: build the predicate-column ColumnIndex for this RG? // `None` → all RGs (Step 1). `Some(set)` → only RGs in the set (Step 2); // out-of-range entries are ignored. The OffsetIndex is built for all RGs - // regardless. + // regardless (see fn docs). let build_ci: Vec = match surviving_rgs { None => vec![true; num_rgs], Some(set) => { @@ -497,30 +595,44 @@ async fn build_scoped_page_index( } }; - // Phase 1: per RG, gather ALL columns' OffsetIndex chunks (always) and — for - // RGs we're building the ColumnIndex for — the predicate columns' ColumnIndex + // Which columns to decode the OffsetIndex for. `None` → all columns. The + // caller already normalized `offset_cols` to include predicate cols + col 0, + // so this is just clamp/sort/dedup. Same indices for every RG (the column set + // is per-file). + let off_cols: Vec = match offset_cols { + None => (0..num_cols).collect(), + Some(cols) => { + let mut v: Vec = cols.iter().copied().filter(|&c| c < num_cols).collect(); + v.sort_unstable(); + v.dedup(); + v + } + }; + if off_cols.is_empty() { + return None; + } + + // Phase 1: per RG, gather the OffsetIndex chunks for `off_cols` and — for RGs + // we're building the ColumnIndex for — the predicate columns' ColumnIndex // chunks. Compute union byte ranges for a single vectored fetch. Bail to // footer-only if any required index range is missing. struct RgPlan { rg_idx: usize, build_ci: bool, pred_chunks: Vec, - all_chunks: Vec, + off_chunks: Vec, col_range: Option>, off_range: Range, } let mut plans: Vec = Vec::with_capacity(num_rgs); - // Flat fetch list; per RG we always push the offset range, and push the - // column range only when building the CI for that RG. We record each plan's - // buffer offsets explicitly so the decode phase can find them. let mut fetch_ranges: Vec> = Vec::with_capacity(num_rgs * 2); // (off_buf_idx, Option) per plan. let mut buf_idx: Vec<(usize, Option)> = Vec::with_capacity(num_rgs); for rg_idx in 0..num_rgs { let rg = footer_meta.row_group(rg_idx); - let all_chunks: Vec = - (0..num_cols).map(|i| rg.column(i).clone()).collect(); - let off_range = offset_index_union(&all_chunks)?; + let off_chunks: Vec = + off_cols.iter().map(|&i| rg.column(i).clone()).collect(); + let off_range = offset_index_union(&off_chunks)?; let off_i = fetch_ranges.len(); fetch_ranges.push(off_range.clone()); @@ -541,7 +653,7 @@ async fn build_scoped_page_index( rg_idx, build_ci: build_ci[rg_idx], pred_chunks, - all_chunks, + off_chunks, col_range, off_range, }); @@ -557,7 +669,9 @@ async fn build_scoped_page_index( } // Phase 3: decode and scatter into full-width per-RG vectors. Pre-fill with - // placeholders so absolute `index[rg][col]` indexing is always valid. + // placeholders so absolute `index[rg][col]` indexing is always valid; only the + // columns in `off_cols` (OffsetIndex) / predicate cols (ColumnIndex) are + // overwritten with real entries. let mut column_index: ParquetColumnIndex = (0..num_rgs) .map(|_| (0..num_cols).map(|_| ColumnIndexMetaData::NONE).collect()) .collect(); @@ -570,17 +684,21 @@ async fn build_scoped_page_index( .collect(); for (plan, &(off_i, col_i)) in plans.iter().zip(buf_idx.iter()) { - // OffsetIndex — always, for every column of this RG. + // OffsetIndex — for the scoped columns of this RG, scattered to absolute + // column positions. let off_reader = BufferChunkReader { base: plan.off_range.start, bytes: buffers[off_i].clone(), }; #[allow(deprecated)] - let decoded_offs = read_offset_indexes(&off_reader, &plan.all_chunks).ok()??; - if decoded_offs.len() != num_cols { + let decoded_offs = read_offset_indexes(&off_reader, &plan.off_chunks).ok()??; + if decoded_offs.len() != off_cols.len() { return None; } - offset_index[plan.rg_idx] = decoded_offs; + let off_row = &mut offset_index[plan.rg_idx]; + for (k, &col) in off_cols.iter().enumerate() { + off_row[col] = decoded_offs[k].clone(); + } // ColumnIndex — only for RGs we're building it for. if plan.build_ci { @@ -605,7 +723,7 @@ async fn build_scoped_page_index( } } - let size = scoped_page_index_size(footer_meta, parquet_cols, num_cols, surviving_rgs); + let size = scoped_page_index_size(footer_meta, parquet_cols, surviving_rgs, &off_cols); Some((column_index, offset_index, size)) } @@ -734,10 +852,11 @@ pub fn surviving_row_groups( fn scoped_page_index_size( footer_meta: &ParquetMetaData, parquet_cols: &[usize], - num_cols: usize, surviving_rgs: Option<&[usize]>, + off_cols: &[usize], ) -> usize { - // ColumnIndex is only built for surviving RGs (Step 2); OffsetIndex for all. + // ColumnIndex is only built for surviving RGs (Step 2); OffsetIndex only for + // off_cols (the predicate ∪ projection ∪ {0} set). let build_ci = |rg_idx: usize| -> bool { match surviving_rgs { None => true, @@ -751,7 +870,7 @@ fn scoped_page_index_size( total += rg.column(pc).column_index_length().unwrap_or(0).max(0) as usize; } } - for c in 0..num_cols { + for &c in off_cols { total += rg.column(c).offset_index_length().unwrap_or(0).max(0) as usize; } } @@ -1421,4 +1540,153 @@ mod tests { clear_scoped_cache_for_test(); } + + // ── Step 2: column-scoping the OffsetIndex ──────────────────────────── + + /// Column-scoped load: OffsetIndex real ONLY for the scoped columns (the + /// caller's set ∪ predicate ∪ {0}); empty placeholder for the rest. ColumnIndex + /// still predicate-scoped. + #[tokio::test] + async fn col_scoped_offset_index_only_for_requested_columns() { + let _g = CACHE_TEST_GUARD.lock().unwrap(); + clear_scoped_cache_for_test(); + // 4 columns: predicate on n1, project s0 (leaf 2). Offset set should be + // {0 (metric), 1 (predicate), 2 (projection)} — NOT column 3 (s1). + let (bytes, schema) = wide4_parquet(); + let (store, loc) = stage(bytes.clone()).await; + let fo = footer_only(&bytes); + let pred_cols = resolve_predicate_parquet_columns(&schema, &fo, &["n1".to_string()]); + assert_eq!(pred_cols, vec![1]); + + // Caller passes projection = {2}; fn unions in predicate {1} + {0}. + let aug = load_scoped_page_index_cols(&store, &loc, &fo, &pred_cols, &[2]) + .await + .expect("col-scoped load must succeed"); + let oi = aug.offset_index().expect("has offset index"); + + for &c in &[0usize, 1, 2] { + assert!( + !oi[0][c].page_locations().is_empty(), + "col {c} (predicate/projection/metric) must have a real OffsetIndex" + ); + } + assert!( + oi[0][3].page_locations().is_empty(), + "col 3 (not projected/predicate/0) OffsetIndex must be an empty placeholder" + ); + } + + /// A projected non-predicate column reads correctly through the col-scoped + /// metadata (its OffsetIndex is real because it's in the projection set) — the + /// read-path safety invariant. + #[tokio::test] + async fn col_scoped_reads_projected_non_predicate_column() { + let _g = CACHE_TEST_GUARD.lock().unwrap(); + clear_scoped_cache_for_test(); + let (bytes, schema) = two_col_parquet(); // price (0), qty (1) + let (store, loc) = stage(bytes.clone()).await; + let fo = footer_only(&bytes); + let pred_cols = resolve_predicate_parquet_columns(&schema, &fo, &["price".to_string()]); + + // Predicate on price (0); project qty (1). Offset set = {0,1}. + let aug = load_scoped_page_index_cols(&store, &loc, &fo, &pred_cols, &[1]) + .await + .unwrap(); + + let selection = RowSelection::from(vec![RowSelector::skip(16), RowSelector::select(16)]); + let scoped_vals = read_selected_column(&bytes, &aug, 1, selection.clone()) + .expect("projected non-predicate read must succeed with col-scoped metadata"); + let full = full_index(&bytes); + let full_vals = read_selected_column(&bytes, &full, 1, selection).unwrap(); + let expected: Vec = (116..132).collect(); + assert_eq!(scoped_vals, expected); + assert_eq!(scoped_vals, full_vals); + } + + /// Column-scoping the OffsetIndex shrinks the cached bytes vs all-columns. + #[tokio::test] + async fn col_scoping_reduces_cached_bytes() { + let _g = CACHE_TEST_GUARD.lock().unwrap(); + clear_scoped_cache_for_test(); + let (bytes, schema) = wide4_parquet(); + let (store, loc) = stage(bytes.clone()).await; + let fo = footer_only(&bytes); + let pred_cols = resolve_predicate_parquet_columns(&schema, &fo, &["n1".to_string()]); + + // All-columns OffsetIndex (Step 1). + let _ = load_scoped_page_index(&store, &loc, &fo, &pred_cols).await.unwrap(); + let all_cols_bytes = scoped_cache_bytes_for_test(); + clear_scoped_cache_for_test(); + + // Col-scoped to {0,1,2} (drops col 3's OffsetIndex). + let _ = load_scoped_page_index_cols(&store, &loc, &fo, &pred_cols, &[2]).await.unwrap(); + let scoped_bytes = scoped_cache_bytes_for_test(); + + assert!( + scoped_bytes < all_cols_bytes, + "col-scoped bytes ({scoped_bytes}) must be < all-columns bytes ({all_cols_bytes})" + ); + clear_scoped_cache_for_test(); + } + + /// If the offset-column set covers every column, the entry collapses to the + /// Step-1 all-columns artifact (shares its cache slot, no separate entry). + #[tokio::test] + async fn col_scoping_full_coverage_collapses_to_all_columns_entry() { + let _g = CACHE_TEST_GUARD.lock().unwrap(); + clear_scoped_cache_for_test(); + let (bytes, schema) = two_col_parquet(); + let (store, loc) = stage(bytes.clone()).await; + let fo = footer_only(&bytes); + let pred_cols = resolve_predicate_parquet_columns(&schema, &fo, &["price".to_string()]); + + // Step-1 all-columns entry. + let _ = load_scoped_page_index(&store, &loc, &fo, &pred_cols).await.unwrap(); + assert_eq!(scoped_cache_len_for_test(), 1); + + // Col-scoped to {1} → union {0,1} = all 2 columns → collapses to the same + // key → HIT, no new entry. + let _ = load_scoped_page_index_cols(&store, &loc, &fo, &pred_cols, &[1]).await.unwrap(); + assert_eq!(scoped_cache_len_for_test(), 1, "full coverage must reuse the all-columns entry"); + assert!(scoped_cache_stats().hits >= 1); + + clear_scoped_cache_for_test(); + } + + /// 4 columns (2 int `n0`,`n1` + 2 wide string `s0`,`s1`), 1 RG, multiple pages. + fn wide4_parquet() -> (Bytes, SchemaRef) { + use arrow::array::StringArray; + let schema = Arc::new(Schema::new(vec![ + Field::new("n0", DataType::Int32, false), + Field::new("n1", DataType::Int32, false), + Field::new("s0", DataType::Utf8, false), + Field::new("s1", DataType::Utf8, false), + ])); + const ROWS: i32 = 256; + let n0: Vec = (0..ROWS).collect(); + let n1: Vec = (0..ROWS).collect(); + let s0: Vec = (0..ROWS).map(|r| format!("s0_{r:05}_padpadpad")).collect(); + let s1: Vec = (0..ROWS).map(|r| format!("s1_{r:05}_padpadpad")).collect(); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(n0)), + Arc::new(Int32Array::from(n1)), + Arc::new(StringArray::from(s0)), + Arc::new(StringArray::from(s1)), + ], + ) + .unwrap(); + let props = WriterProperties::builder() + .set_max_row_group_size(ROWS as usize) + .set_data_page_row_count_limit(32) + .set_write_batch_size(32) + .set_statistics_enabled(EnabledStatistics::Page) + .build(); + let mut buf: Vec = Vec::new(); + let mut w = ArrowWriter::try_new(&mut buf, schema.clone(), Some(props)).unwrap(); + w.write(&batch).unwrap(); + w.close().unwrap(); + (Bytes::from(buf), schema) + } } From 90dc109547b183608a7169b3c931c3dc6d3c3ae9 Mon Sep 17 00:00:00 2001 From: G Date: Wed, 17 Jun 2026 00:48:49 +0530 Subject: [PATCH 13/18] adding stats Signed-off-by: G --- .../rust/src/ffm.rs | 27 +- .../src/indexed_table/page_index_loader.rs | 1696 ++++++++--------- .../rust/src/stats.rs | 89 +- .../be/datafusion/DataFusionPlugin.java | 8 +- .../be/datafusion/DatafusionSettings.java | 3 +- .../RestClearScopedPageIndexCacheAction.java | 39 +- .../be/datafusion/cache/CacheSettings.java | 36 +- .../be/datafusion/cache/CacheUtils.java | 21 +- .../be/datafusion/nativelib/NativeBridge.java | 30 +- .../be/datafusion/nativelib/StatsLayout.java | 53 +- .../be/datafusion/stats/CacheStats.java | 56 +- .../DataFusionPluginSettingsTests.java | 4 +- .../be/datafusion/DataFusionServiceTests.java | 2 + .../DatafusionCacheManagerTests.java | 2 + .../datafusion/DatafusionSettingsTests.java | 3 +- .../nativelib/StatsLayoutPropertyTests.java | 75 +- .../nativelib/StatsLayoutTests.java | 6 +- .../be/datafusion/stats/CacheStatsTests.java | 43 +- .../stats/DataFusionStatsPropertyTests.java | 7 +- .../stats/DataFusionStatsTests.java | 3 +- .../analytics/qa/ScopedPageIndexCacheIT.java | 466 +++-- 21 files changed, 1478 insertions(+), 1191 deletions(-) diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs index 9d9ad397c544d..771efe1f34e03 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs @@ -1040,23 +1040,32 @@ pub unsafe extern "C" fn df_cache_manager_get_total_memory(runtime_ptr: i64) -> Ok(manager.get_total_memory_consumed() as i64) } -/// Set the byte budget of the process-global scoped page-index cache. +/// Set the byte budget of the process-global scoped ColumnIndex cache. /// -/// Unlike the metadata/statistics caches, the scoped page-index cache is a -/// process-wide singleton (see [`crate::indexed_table::page_index_loader`]), so +/// Unlike the metadata/statistics caches, the scoped page-index caches are +/// process-wide singletons (see [`crate::indexed_table::page_index_loader`]), so /// there is no manager pointer — the limit applies globally. Negative values are /// rejected; zero is ignored (keeps the existing budget). Shrinking the limit /// evicts least-recently-used entries immediately. #[ffm_safe] #[no_mangle] -pub extern "C" fn df_set_scoped_page_index_cache_limit(size_limit: i64) -> i64 { +pub extern "C" fn df_set_column_index_cache_limit(size_limit: i64) -> i64 { if size_limit < 0 { - return Err(format!( - "df_set_scoped_page_index_cache_limit: negative limit {}", - size_limit - )); + return Err(format!("df_set_column_index_cache_limit: negative limit {}", size_limit)); + } + crate::indexed_table::page_index_loader::set_column_index_cache_limit(size_limit as usize); + Ok(0) +} + +/// Set the byte budget of the process-global scoped OffsetIndex cache. See +/// [`df_set_column_index_cache_limit`] for the singleton/eviction semantics. +#[ffm_safe] +#[no_mangle] +pub extern "C" fn df_set_offset_index_cache_limit(size_limit: i64) -> i64 { + if size_limit < 0 { + return Err(format!("df_set_offset_index_cache_limit: negative limit {}", size_limit)); } - crate::indexed_table::page_index_loader::set_scoped_cache_limit(size_limit as usize); + crate::indexed_table::page_index_loader::set_offset_index_cache_limit(size_limit as usize); Ok(0) } diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/page_index_loader.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/page_index_loader.rs index 579ac5757aa53..e75d8c336fd72 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/page_index_loader.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/page_index_loader.rs @@ -6,7 +6,7 @@ * compatible open source license. */ -//! Unified, scoped parquet page-index cache (Step 1a). +//! Scoped parquet page-index caches — TWO caches, by consumer. //! //! # Why this exists //! @@ -15,84 +15,67 @@ //! `OffsetIndex` (per-page byte offsets), for every column of every row group. //! On wide schemas (the production `textbench` index has 402 columns) this is //! ~82% of native heap and is re-decoded per query, even when the query filters -//! a single column. +//! a single column. The level-1 metadata cache is kept footer-only (see +//! [`crate::cache`]); this module rebuilds a *scoped* page index per query and +//! caches it, shared by both scan paths (the DataFusion `ListingTable` path and +//! the custom indexed-table executor). //! -//! This module loads the page index **scoped to the query's predicate columns** -//! and caches the result in **one** process-global cache that is shared by both -//! scan paths (the DataFusion `ListingTable` path and the custom indexed-table -//! executor). For the same `(file, predicate columns)` an entry built by one -//! path is reusable by the other. +//! # Two caches, because the two indexes have different drivers //! -//! ## What is scoped, and why each scope is safe +//! The `ColumnIndex` and `OffsetIndex` are consumed by different parts of +//! DataFusion 54 / parquet 58, with **different natural cache keys**. Forcing +//! them into one key makes the projection-driven OffsetIndex poison the +//! predicate-driven ColumnIndex's broad cross-path sharing (the failure mode of +//! the prior iteration). So they are split: //! -//! - **`ColumnIndex` → predicate columns only.** It is read only at *prune* -//! time, and the pruner (`StatisticsConverter`, and our [`super::page_pruner`]) -//! only ever dereferences predicate-column positions. Non-predicate positions -//! carry [`ColumnIndexMetaData::NONE`] placeholders, which keep absolute -//! `index[rg][col]` indexing valid and cost nothing. This scoping is where the -//! heap win comes from. -//! - **`OffsetIndex` → all columns (this step also: all row groups).** Unlike -//! the `ColumnIndex`, the `OffsetIndex` is read at *scan* time: when a -//! `RowSelection` is active, arrow-rs's `InMemoryRowGroup::fetch_ranges` -//! dereferences `offset_index[col].page_locations()` for every **projected** -//! column (by absolute index) to compute which page byte ranges to fetch. A -//! placeholder there for a projected column fetches zero ranges and the read -//! fails ("failed to skip rows, expected N, got 0"). Because this loader runs -//! before the projection is known, it keeps a real `OffsetIndex` for every -//! column. That is cheap relative to the `ColumnIndex`: fixed-width page -//! offsets/sizes, no per-page string stats. +//! - **ColumnIndex — predicate-driven.** Read only at *prune* time, and only for +//! the predicate column being evaluated +//! (`page_filter::PagesPruningStatistics`, `offset_index[rg][predicate_col]`). +//! Key: `(file, predicate_cols, surviving_rgs)`. Deterministic in the +//! *predicate* (independent of what you `SELECT`), so the same filter shares +//! its entry across scan paths **and** across queries with different +//! projections. This is the heavy index (string min/max) and the big heap win. +//! Scoped to predicate columns (`NONE` placeholders elsewhere) and, optionally, +//! to the row groups that pass footer-stats pruning ([`surviving_row_groups`]). //! -//! ## Why the cache stores the (ColumnIndex, OffsetIndex) PAIR, not metadata +//! - **OffsetIndex — projection-driven.** Read at *scan* time for **projected** +//! columns (`InMemoryRowGroup::fetch_ranges`, `projection.leaf_included(idx)`), +//! and at prune time for the predicate column, and at column 0 for the +//! page-skip metric. Key: `(file, offset_cols)` where +//! `offset_cols = predicate ∪ projection ∪ {0}`. This is the cheap, fixed-width +//! index (no per-page string stats). Built for **all row groups** (an empty +//! OffsetIndex on a row group DataFusion scans panics / breaks reads, and +//! DataFusion chooses the scanned set itself, after our load — see +//! HANDOFF_step2_rg_scoping.md §1e). //! -//! `ParquetMetaData` owns its `row_groups: Vec` (~5–6 MB on -//! the 402-column schema). Caching a full augmented `ParquetMetaData` per entry -//! would **duplicate that footer** in every entry (the prior iteration's -//! over-allocation bug: a row-group-pruned entry that read 0 row groups still -//! cost a 60 MB footer clone). Instead this cache stores only the decoded -//! `(ParquetColumnIndex, ParquetOffsetIndex)` pair plus a size estimate, and -//! **grafts** the pair onto the caller's already-resident footer at lookup via -//! [`ParquetMetaData::into_builder`] → `set_column_index`/`set_offset_index`. -//! The footer is never duplicated into the cache; the only footer copy is the -//! one transient graft handed to the scan (one deep clone per file per query, -//! the same order as a cold metadata fetch). +//! Each cache stores only its decoded vector (`ParquetColumnIndex` / +//! `ParquetOffsetIndex`) — never a full `ParquetMetaData` (no footer +//! duplication). On lookup the two are **grafted** onto the caller's +//! already-resident footer via [`ParquetMetaData::into_builder`] → +//! `set_column_index`/`set_offset_index`. //! -//! **Consequence for tests:** a cache hit returns a *fresh* `Arc`, so -//! `Arc::ptr_eq` is the wrong signal for "served from cache" — assert via the -//! [`ScopedCacheStats::hits`] counter instead. -//! -//! ## Key -//! -//! The cache key is `(file path, predicate parquet-column indices)` **only**. -//! Both scan paths resolve predicate columns the same way -//! ([`resolve_predicate_parquet_columns`]) and build the identical artifact -//! (all-RG, all-column `OffsetIndex`; predicate-scoped `ColumnIndex`), so an -//! entry is shared across paths. Row-group scoping is a *future* axis that would -//! extend the key — deliberately omitted here so Step 1 keeps one unified key. +//! **Consequence for tests:** a lookup returns a *fresh* `Arc`, so `Arc::ptr_eq` +//! is the wrong signal for "served from cache" — assert via the per-cache hit +//! counters ([`column_index_cache_stats`] / [`offset_index_cache_stats`]). //! //! ## Correctness / fallback //! -//! Any failure (file has no page index, a predicate column lacks an index range, -//! a decode/IO error) makes the load return `None`. The caller keeps its +//! Any failure (file has no page index, a column lacks an index range, a +//! decode/IO error) makes the load return `None`. The caller keeps its //! footer-only metadata and the pruner conservatively no-ops (scans the whole //! row group) — never a wrong result. //! //! ## Upstream note //! -//! arrow-rs is moving toward first-class selective metadata decoding: a unified -//! options mechanism that would include "decode column indexes only for -//! predicate columns" and "row-group selection" (apache/arrow-rs#8643, open), -//! and the `ParquetMetaDataOptions` / `ParquetStatisticsPolicy::skip_except` -//! pattern just merged for page *encoding* statistics (apache/arrow-rs#8797, -//! built on the selective-decode metadata "index" of #8714). None of these yet -//! expose a page-index column/row-group projection. Until one does, the only -//! **public** API that decodes a *subset* of columns is the deprecated -//! [`read_columns_indexes`]/[`read_offset_indexes`] (the per-column primitives -//! are `pub(crate)`; `ParquetMetaDataReader`/`PageIndexPolicy` is -//! all-columns-or-nothing). This module is the interim hand-rolled equivalent; -//! when a `ParquetMetaDataOptions` page-index projection lands, migrate -//! [`build_scoped_page_index`] to it and drop the `#[allow(deprecated)]`. +//! arrow-rs is moving toward first-class selective metadata decoding +//! (apache/arrow-rs#8643 open; the `ParquetStatisticsPolicy::skip_except` pattern +//! merged in #8797 / #8714 for encoding stats). None yet expose a page-index +//! column/row-group projection, so we hand-roll it with the deprecated +//! [`read_columns_indexes`]/[`read_offset_indexes`] (the only public subset +//! decoders). Migrate to `ParquetMetaDataOptions` when it grows a page-index knob. use std::collections::HashMap; +use std::hash::Hash; use std::ops::Range; use std::sync::{Arc, Mutex}; @@ -113,52 +96,18 @@ use datafusion::parquet::file::reader::{ChunkReader, Length}; use object_store::ObjectStore; use prost::bytes::{Buf, Bytes}; -/// Default byte budget for the scoped page-index cache, used until the caller -/// sets one from the runtime's configured limit (see [`set_scoped_cache_limit`]). -/// 64 MiB is generous for a predicate-column-only page index yet a tiny fraction -/// of the footer-only baseline the page-index strip buys back. +/// Default byte budget for EACH scoped cache, used until the caller sets one from +/// the runtime's configured limit (see [`set_column_index_cache_limit`] / +/// [`set_offset_index_cache_limit`]). The two caches are budgeted independently: +/// the ColumnIndex (per-page string min/max) is the heavy one and the OffsetIndex +/// (fixed-width page offsets) is tiny, so they get separate, separately-tunable +/// limits rather than sharing one number. const DEFAULT_SCOPED_CACHE_LIMIT: usize = 64 * 1024 * 1024; -// ── Cache types ────────────────────────────────────────────────────────── - -/// Cache key: object-store path + the sorted/deduped set of parquet column -/// indices the page index was built for + the sorted/deduped set of surviving -/// row groups the (predicate-column) ColumnIndex was built for. -/// -/// `surviving_rgs` is empty for the Step-1 all-row-group artifact -/// ([`load_scoped_page_index`]); non-empty for the Step-2 RG-scoped artifact -/// ([`load_scoped_page_index_rgs`]). It is part of the key because a ColumnIndex -/// built for a different RG set is a different artifact — but because both scan -/// paths derive `surviving_rgs` deterministically from the SAME `(file, -/// predicate)` via [`surviving_row_groups`], they compute the identical key, so -/// cross-path sharing is preserved. -#[derive(Clone, PartialEq, Eq, Hash, Debug)] -struct ScopedKey { - path: String, - parquet_cols: Vec, - surviving_rgs: Vec, - /// The (normalized) set of columns the `OffsetIndex` was built for. Empty = - /// all columns (Step 1). Non-empty = `predicate ∪ projection ∪ {0}` (Step 2 - /// column-scoping). Part of the key because an OffsetIndex built for a - /// different column set is a different artifact; both scan paths derive it - /// deterministically (predicate + projection are known per query), so the key - /// stays identical across paths for the same logical request. - offset_cols: Vec, -} - -/// One cached entry: the decoded page-index pair (full-width over all row groups -/// and all columns; `ColumnIndex` real only at predicate positions), a size -/// estimate, and a last-used tick for LRU ordering. Deliberately does NOT hold a -/// `ParquetMetaData` — see the module docs (no footer duplication). -struct ScopedEntry { - column_index: ParquetColumnIndex, - offset_index: ParquetOffsetIndex, - size: usize, - last_used: u64, -} +// ── Generic byte-bounded LRU ──────────────────────────────────────────────── -/// Snapshot of scoped page-index cache counters. Surfaced on node-stats and used -/// by tests to assert hits/misses without relying on `Arc::ptr_eq`. +/// Snapshot of one scoped cache's counters plus occupancy. Surfaced on +/// node-stats and used by tests to assert hits/misses without `Arc::ptr_eq`. #[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] pub struct ScopedCacheStats { pub hits: u64, @@ -169,15 +118,18 @@ pub struct ScopedCacheStats { pub limit_bytes: usize, } -/// Byte-bounded LRU over scoped page-index pairs, keyed by -/// `(file, predicate-column-set)`. -/// -/// Evicts the least-recently-used entries once `used > limit`, so it always -/// serves cached data within a memory budget rather than silently degrading to -/// "decode every query" when full. Entry size is a deterministic structural -/// estimate (see [`scoped_page_index_size`]). -struct ScopedLru { - map: HashMap, +struct LruEntry { + value: V, + size: usize, + last_used: u64, +} + +/// Byte-bounded LRU. Evicts the least-recently-used entries once `used > limit`, +/// so it always serves cached data within a memory budget rather than silently +/// degrading to "decode every query" when full. `V` is cloned on hit (cheap: a +/// predicate-scoped `ColumnIndex` or a fixed-width `OffsetIndex`, never a footer). +struct Lru { + map: HashMap>, used: usize, limit: usize, /// Monotonic clock for LRU ordering (avoids `Instant`; fine single-process). @@ -187,7 +139,7 @@ struct ScopedLru { evictions: u64, } -impl ScopedLru { +impl Lru { fn new(limit: usize) -> Self { Self { map: HashMap::new(), @@ -205,15 +157,13 @@ impl ScopedLru { self.tick } - /// On hit, bump recency and return a clone of the pair (cheap relative to a - /// footer: predicate-scoped `ColumnIndex` + fixed-width `OffsetIndex`). - fn get(&mut self, key: &ScopedKey) -> Option<(ParquetColumnIndex, ParquetOffsetIndex)> { + fn get(&mut self, key: &K) -> Option { let t = self.next_tick(); match self.map.get_mut(key) { Some(entry) => { entry.last_used = t; self.hits += 1; - Some((entry.column_index.clone(), entry.offset_index.clone())) + Some(entry.value.clone()) } None => { self.misses += 1; @@ -222,28 +172,14 @@ impl ScopedLru { } } - fn insert( - &mut self, - key: ScopedKey, - column_index: ParquetColumnIndex, - offset_index: ParquetOffsetIndex, - size: usize, - ) { + fn insert(&mut self, key: K, value: V, size: usize) { // An entry larger than the whole budget can never be retained; skip it // rather than evicting everything else for something we'd drop anyway. if size > self.limit { return; } let t = self.next_tick(); - if let Some(old) = self.map.insert( - key, - ScopedEntry { - column_index, - offset_index, - size, - last_used: t, - }, - ) { + if let Some(old) = self.map.insert(key, LruEntry { value, size, last_used: t }) { self.used -= old.size; } self.used += size; @@ -282,51 +218,116 @@ impl ScopedLru { self.limit = limit; self.evict(); } + + fn clear_keep_limit(&mut self) { + self.map.clear(); + self.used = 0; + self.tick = 0; + self.hits = 0; + self.misses = 0; + self.evictions = 0; + } +} + +// ── Cache keys + the two global caches ────────────────────────────────────── + +/// ColumnIndex cache key — one decoded `ColumnIndexMetaData` **cell** per +/// `(file, column, row-group)`. The page index for a given column+RG is an +/// intrinsic property of the file: it is identical no matter which *other* +/// columns a query filters on, or which literal a predicate uses. Keying at the +/// cell granularity means a column's per-page string min/max is decoded and +/// stored **once per file**, then reused by every query whose predicate touches +/// that column — regardless of the predicate-column *combination* or the +/// surviving-row-group *set*. (The prior set-keyed design re-decoded and +/// re-stored a column for every distinct predicate/RG combination — storage grew +/// with query diversity, not schema width.) +/// +/// Both scan paths resolve the same `(file, col, rg)` for the same logical +/// request, so cells are shared across paths → cross-path sharing. +#[derive(Clone, PartialEq, Eq, Hash, Debug)] +struct CiCellKey { + path: Arc, + col: usize, + rg: usize, +} + +/// OffsetIndex cache key — one decoded value per `(file, column)`, where the +/// value is that column's `OffsetIndexMetaData` for **every** row group (a +/// `Vec` indexed by RG). Unlike the ColumnIndex, the OffsetIndex is read at scan +/// time for any RG DataFusion chooses to scan — and DataFusion picks that set +/// itself, after our load — so a column's OffsetIndex must always cover all RGs +/// (an empty entry on a scanned RG panics / breaks reads). RG can therefore never +/// be a key axis here; the cell is the whole-column, all-RG offset index. Keyed +/// only on `(file, col)`, so any query that reads a column reuses its offset +/// index irrespective of projection or predicate. +#[derive(Clone, PartialEq, Eq, Hash, Debug)] +struct OiCellKey { + path: Arc, + col: usize, } -/// Process-wide scoped page-index cache. -static SCOPED_CACHE: Lazy> = - Lazy::new(|| Mutex::new(ScopedLru::new(DEFAULT_SCOPED_CACHE_LIMIT))); +/// One column's OffsetIndex across all row groups (indexed by RG). The value type +/// of [`OFFSET_INDEX_CACHE`]. +type OiColumn = Vec; + +static COLUMN_INDEX_CACHE: Lazy>> = + Lazy::new(|| Mutex::new(Lru::new(DEFAULT_SCOPED_CACHE_LIMIT))); + +static OFFSET_INDEX_CACHE: Lazy>> = + Lazy::new(|| Mutex::new(Lru::new(DEFAULT_SCOPED_CACHE_LIMIT))); -/// Set the scoped cache's byte budget. Called from startup wiring with the -/// configured limit. Idempotent; shrinking evicts immediately. Zero is ignored -/// (keeps the existing budget). -pub fn set_scoped_cache_limit(limit: usize) { +// ── Limits / stats / clear ────────────────────────────────────────────────── + +/// Set the ColumnIndex cache's byte budget. Called from startup wiring with the +/// configured limit. Idempotent; shrinking evicts immediately. Zero ignored. +pub fn set_column_index_cache_limit(limit: usize) { if limit == 0 { return; } - if let Ok(mut c) = SCOPED_CACHE.lock() { + if let Ok(mut c) = COLUMN_INDEX_CACHE.lock() { c.set_limit(limit); } } -/// Snapshot of the scoped page-index cache counters plus occupancy. For -/// node-stats / observability and tests. -pub fn scoped_cache_stats() -> ScopedCacheStats { - SCOPED_CACHE.lock().map(|c| c.stats()).unwrap_or_default() +/// Set the OffsetIndex cache's byte budget. Called from startup wiring with the +/// configured limit. Idempotent; shrinking evicts immediately. Zero ignored. +pub fn set_offset_index_cache_limit(limit: usize) { + if limit == 0 { + return; + } + if let Ok(mut c) = OFFSET_INDEX_CACHE.lock() { + c.set_limit(limit); + } +} + +/// Counters + occupancy of the ColumnIndex (predicate-driven) cache. +pub fn column_index_cache_stats() -> ScopedCacheStats { + COLUMN_INDEX_CACHE.lock().map(|c| c.stats()).unwrap_or_default() } -/// Drop all cached entries and reset the hit/miss/eviction counters, keeping the -/// configured byte budget. Exposed for operational testing (clear the cache and -/// re-measure without a cluster restart). The next query repopulates lazily. +/// Counters + occupancy of the OffsetIndex (projection-driven) cache. +pub fn offset_index_cache_stats() -> ScopedCacheStats { + OFFSET_INDEX_CACHE.lock().map(|c| c.stats()).unwrap_or_default() +} + +/// Drop all entries and reset counters in BOTH caches, keeping the budgets. For +/// operational testing — reset and re-measure without a cluster restart. pub fn clear_scoped_cache() { - if let Ok(mut c) = SCOPED_CACHE.lock() { - c.map.clear(); - c.used = 0; - c.tick = 0; - c.hits = 0; - c.misses = 0; - c.evictions = 0; + if let Ok(mut c) = COLUMN_INDEX_CACHE.lock() { + c.clear_keep_limit(); + } + if let Ok(mut c) = OFFSET_INDEX_CACHE.lock() { + c.clear_keep_limit(); } } -// ── Public API ───────────────────────────────────────────────────────────── +// ── Public API ────────────────────────────────────────────────────────────── /// Map the query's arrow predicate-column names to this file's parquet column /// indices, using the same resolution the pruner uses -/// (`StatisticsConverter::parquet_column_index`). Columns absent from the -/// parquet file (schema evolution) are skipped. Returns a sorted, deduped set so -/// both scan paths produce an identical key for the same logical predicate. +/// (`StatisticsConverter::parquet_column_index`). Columns absent from the parquet +/// file (schema evolution) are skipped. Returns a sorted, deduped set so both +/// scan paths produce an identical key for the same logical predicate. pub fn resolve_predicate_parquet_columns( arrow_schema: &SchemaRef, metadata: &ParquetMetaData, @@ -344,51 +345,34 @@ pub fn resolve_predicate_parquet_columns( set.into_iter().collect() } -/// Load (or build + cache) the scoped page index for `parquet_cols` and graft it -/// onto `footer_meta`, returning fresh metadata that carries: -/// - a real `ColumnIndex` at the predicate-column positions (NONE elsewhere), -/// - a real `OffsetIndex` for every column, across all row groups. -/// -/// Returns `None` (caller keeps footer-only metadata) on any condition that -/// would make page pruning unsafe or impossible: empty column set, a file -/// without a page index, an out-of-range predicate column, or a decode/IO error. -/// -/// Consults and populates the shared `(file, predicate-cols)` cache. +/// Load + graft a scoped page index: ColumnIndex for `parquet_cols` (all RGs), +/// OffsetIndex for all columns/all RGs. The Step-1 baseline. pub async fn load_scoped_page_index( store: &Arc, location: &object_store::path::Path, footer_meta: &Arc, parquet_cols: &[usize], ) -> Option> { - load_scoped_page_index_inner(store, location, footer_meta, parquet_cols, None, None).await + load_combined(store, location, footer_meta, parquet_cols, None, None).await } -/// Like [`load_scoped_page_index`], but the `OffsetIndex` is decoded only for the -/// columns in `offset_cols` (others get an empty placeholder), instead of all -/// columns. -/// -/// # Why this is safe (and which columns must be included) -/// -/// The `OffsetIndex` is dereferenced in exactly three places in DataFusion 54 / -/// parquet 58: -/// - **Read** (`InMemoryRowGroup::fetch_ranges`): only for **projected** columns -/// (`projection.leaf_included(idx)`), and only when a `RowSelection` is active. -/// - **Prune** (`page_filter::PagesPruningStatistics`): only for the **predicate** -/// column being pruned (`offset_index[rg][parquet_column_index]`). -/// - **Metric** (`page_filter` `total_pages_in_group`): reads **column 0**'s page -/// count to seed the `page_index_rows_pruned` counter. This does NOT affect the -/// produced `RowSelection` (correctness), only the skipped-pages metric. -/// -/// So a real `OffsetIndex` is required for `predicate ∪ projection`; including -/// **column 0** additionally keeps the page-skip metric identical to stock -/// DataFusion. This fn always unions in the predicate columns and column 0 -/// defensively, so callers only need to pass the projection (∪ predicate). Any -/// column NOT in the final set gets an empty placeholder — safe because it is -/// never projected, never the predicate, and never column 0, so it is never -/// dereferenced. -/// -/// `offset_cols` joins the cache key. Callers MUST derive it deterministically -/// from `(predicate, projection)` so both scan paths agree and share entries. +/// Like [`load_scoped_page_index`], but the ColumnIndex is built only for the row +/// groups in `surviving_rgs` (footer-stats survivors — [`surviving_row_groups`]); +/// other RGs get a `NONE` ColumnIndex placeholder. OffsetIndex stays all-columns. +pub async fn load_scoped_page_index_rgs( + store: &Arc, + location: &object_store::path::Path, + footer_meta: &Arc, + parquet_cols: &[usize], + surviving_rgs: &[usize], +) -> Option> { + load_combined(store, location, footer_meta, parquet_cols, Some(surviving_rgs), None).await +} + +/// Like [`load_scoped_page_index`], but the OffsetIndex is built only for +/// `offset_cols` (the loader unions in the predicate columns + column 0 +/// defensively); other columns get an empty placeholder. ColumnIndex stays +/// all-RG. See [`OiKey`] for which columns must be real and why. pub async fn load_scoped_page_index_cols( store: &Arc, location: &object_store::path::Path, @@ -396,47 +380,32 @@ pub async fn load_scoped_page_index_cols( parquet_cols: &[usize], offset_cols: &[usize], ) -> Option> { - load_scoped_page_index_inner(store, location, footer_meta, parquet_cols, None, Some(offset_cols)) - .await + load_combined(store, location, footer_meta, parquet_cols, None, Some(offset_cols)).await } -/// Like [`load_scoped_page_index`], but the predicate-column `ColumnIndex` is -/// decoded ONLY for the row groups in `surviving_rgs` (those that pass footer -/// RG-stats pruning for the predicate — see [`surviving_row_groups`]); all other -/// row groups get a `NONE` ColumnIndex placeholder. The `OffsetIndex` is STILL -/// built for every column of every row group — it is read at scan time for any -/// row group DataFusion chooses to scan, and DataFusion picks that set itself -/// (after, and independently of, our page-index load), so an empty OffsetIndex on -/// a scanned RG would panic / break reads (see HANDOFF_step2_rg_scoping.md §1e). -/// -/// Safety of the ColumnIndex RG-scoping: `surviving_rgs` is a **superset** of the -/// RGs DataFusion actually scans (DF further applies bloom/range/limit pruning), -/// so every RG DF page-prunes has a real ColumnIndex. RGs we placeholdered are -/// exactly those DF also prunes by footer stats, so the `NONE` is never -/// dereferenced — and even if it were, `NONE` is handled gracefully ("no stats, -/// keep pages") as long as the OffsetIndex stays valid, which it does. -/// -/// `surviving_rgs` becomes part of the cache key, but since both scan paths derive -/// it deterministically from the same `(file, predicate)`, they share entries. -pub async fn load_scoped_page_index_rgs( +/// Fully scoped: ColumnIndex RG-scoped to `surviving_rgs`, OffsetIndex +/// column-scoped to `offset_cols` (∪ predicate ∪ {0}). The Step-2 target both +/// scan paths call once they know their surviving-RG set and projection. +pub async fn load_scoped_page_index_scoped( store: &Arc, location: &object_store::path::Path, footer_meta: &Arc, parquet_cols: &[usize], surviving_rgs: &[usize], + offset_cols: &[usize], ) -> Option> { - load_scoped_page_index_inner( + load_combined( store, location, footer_meta, parquet_cols, Some(surviving_rgs), - None, + Some(offset_cols), ) .await } -async fn load_scoped_page_index_inner( +async fn load_combined( store: &Arc, location: &object_store::path::Path, footer_meta: &Arc, @@ -447,89 +416,16 @@ async fn load_scoped_page_index_inner( if parquet_cols.is_empty() { return None; } - let num_cols = footer_meta.file_metadata().schema_descr().num_columns(); - - // Normalize the surviving-RG set into the key (sorted/deduped). `None` (all - // RGs, Step 1) is encoded as an empty vec, distinct from a real subset. - let key_rgs: Vec = match surviving_rgs { - Some(rgs) => { - let mut v = rgs.to_vec(); - v.sort_unstable(); - v.dedup(); - v - } - None => Vec::new(), - }; - - // Normalize the OffsetIndex column set. `None` (all columns, Step 1) → empty - // sentinel. `Some(cols)` → cols ∪ predicate-cols ∪ {0}, clamped to in-range - // and sorted/deduped (the defensive union guarantees prune + metric safety - // regardless of what the caller passed). If the union covers every column we - // collapse to the empty "all columns" sentinel so it shares the Step-1 entry. - let key_offset_cols: Vec = match offset_cols { - None => Vec::new(), - Some(cols) => { - let mut set: std::collections::BTreeSet = std::collections::BTreeSet::new(); - set.insert(0); // metric reads column 0 - for &c in parquet_cols { - set.insert(c); - } // prune reads predicate cols - for &c in cols { - if c < num_cols { - set.insert(c); - } - } // read needs projected cols - let v: Vec = set.into_iter().filter(|&c| c < num_cols).collect(); - if v.len() == num_cols { - Vec::new() // full coverage → reuse the all-columns artifact - } else { - v - } - } - }; - - let key = ScopedKey { - path: location.as_ref().to_string(), - parquet_cols: parquet_cols.to_vec(), - surviving_rgs: key_rgs, - offset_cols: key_offset_cols.clone(), - }; - - // Cache hit: graft the cached pair onto the caller's footer. No I/O, no - // decode. (Returns a fresh Arc — assert hits via the counter, not ptr_eq.) - if let Ok(mut cache) = SCOPED_CACHE.lock() { - if let Some((ci, oi)) = cache.get(&key) { - return Some(graft(footer_meta, ci, oi)); - } - } - - // Pass the normalized offset-col set (empty = all columns) to the builder. - let offset_cols_for_build: Option<&[usize]> = - if key_offset_cols.is_empty() { None } else { Some(&key_offset_cols) }; - - // Miss: range-read + decode the scoped page index. - let (column_index, offset_index, size) = build_scoped_page_index( - store, - location, - footer_meta, - parquet_cols, - surviving_rgs, - offset_cols_for_build, - ) - .await?; - - // Graft a copy for the caller; store the originals in the cache. - let grafted = graft(footer_meta, column_index.clone(), offset_index.clone()); - if let Ok(mut cache) = SCOPED_CACHE.lock() { - cache.insert(key, column_index, offset_index, size); - } - Some(grafted) + let column_index = + get_or_build_column_index(store, location, footer_meta, parquet_cols, surviving_rgs).await?; + let offset_index = + get_or_build_offset_index(store, location, footer_meta, parquet_cols, offset_cols).await?; + Some(graft(footer_meta, column_index, offset_index)) } -/// Build a fresh `ParquetMetaData` = `footer` with the scoped page-index pair -/// grafted on. This deep-clones the footer (the builder consumes an owned -/// `ParquetMetaData`); that transient clone is the only footer copy and is never -/// stored in the cache. +/// Build a fresh `ParquetMetaData` = `footer` with the page-index pair grafted +/// on. Deep-clones the footer (the builder consumes an owned `ParquetMetaData`); +/// that transient clone is the only footer copy and is never cached. fn graft( footer_meta: &Arc, column_index: ParquetColumnIndex, @@ -544,212 +440,310 @@ fn graft( Arc::new(rebuilt) } -// ── Build ──────────────────────────────────────────────────────────────── +// ── ColumnIndex cache lookup + build (per `(file, col, rg)` cell) ──────────── -/// Range-read and decode the page index scoped to `parquet_cols`, returning the -/// full-width `(ColumnIndex, OffsetIndex)` pair plus a size estimate. +/// Assemble the full-width `[rg][col]` `ColumnIndex` matrix (real cells only at +/// `parquet_cols` × built RGs; `NONE` everywhere else) by looking up each +/// `(file, col, rg)` cell in the cache and decoding only the cells that miss. /// -/// - `OffsetIndex`: real for the columns in `offset_cols` (the caller-normalized -/// `predicate ∪ projection ∪ {0}`), empty placeholder elsewhere. `None` → all -/// columns (Step 1 / no projection known). Built for EVERY row group (an -/// OffsetIndex omitted on a scanned RG would break reads — DataFusion chooses -/// the scanned set itself, after our load). -/// - `ColumnIndex`: real at the predicate-column positions, `NONE` elsewhere. -/// When `surviving_rgs` is `Some(set)`, the (heavy) predicate-column ColumnIndex -/// is decoded ONLY for row groups in `set`; other RGs get `NONE` (Step-2 -/// RG-scoping). When `None`, it is decoded for all RGs (Step-1). -/// -/// `None` on any unsafe/impossible condition. -async fn build_scoped_page_index( +/// `surviving_rgs == None` builds every RG; `Some(set)` restricts the built RGs +/// to footer-stats survivors ([`surviving_row_groups`]). Either way a cell is +/// keyed solely on `(file, col, rg)`, so it is decoded once per file and reused +/// across every predicate combination and surviving-RG set that touches it. +async fn get_or_build_column_index( store: &Arc, location: &object_store::path::Path, footer_meta: &Arc, parquet_cols: &[usize], surviving_rgs: Option<&[usize]>, - offset_cols: Option<&[usize]>, -) -> Option<(ParquetColumnIndex, ParquetOffsetIndex, usize)> { +) -> Option { let num_rgs = footer_meta.num_row_groups(); if num_rgs == 0 { return None; } let num_cols = footer_meta.file_metadata().schema_descr().num_columns(); - // Predicate indices index into the file schema (shared across RGs). if parquet_cols.iter().any(|&i| i >= num_cols) { return None; } - // Per-RG mask: build the predicate-column ColumnIndex for this RG? - // `None` → all RGs (Step 1). `Some(set)` → only RGs in the set (Step 2); - // out-of-range entries are ignored. The OffsetIndex is built for all RGs - // regardless (see fn docs). - let build_ci: Vec = match surviving_rgs { - None => vec![true; num_rgs], + // Which RGs to build the (heavy) predicate-column ColumnIndex for. + let build_rgs: Vec = match surviving_rgs { + None => (0..num_rgs).collect(), Some(set) => { - let mut v = vec![false; num_rgs]; - for &r in set { - if r < num_rgs { - v[r] = true; - } - } + let mut v: Vec = set.iter().copied().filter(|&r| r < num_rgs).collect(); + v.sort_unstable(); + v.dedup(); v } }; + if build_rgs.is_empty() { + // Nothing to build (e.g. an empty survivor set) → footer-only fallback. + return None; + } + + let path: Arc = Arc::from(location.as_ref()); + let mut matrix: ParquetColumnIndex = (0..num_rgs) + .map(|_| (0..num_cols).map(|_| ColumnIndexMetaData::NONE).collect()) + .collect(); + + // Phase 1: serve every needed cell that is already cached; collect misses. + let mut missing: Vec<(usize, usize)> = Vec::new(); // (col, rg) + if let Ok(mut cache) = COLUMN_INDEX_CACHE.lock() { + for &rg in &build_rgs { + for &col in parquet_cols { + let key = CiCellKey { path: path.clone(), col, rg }; + match cache.get(&key) { + Some(cell) => matrix[rg][col] = cell, + None => missing.push((col, rg)), + } + } + } + } else { + return None; + } + + // Phase 2: decode the missing cells (vectored fetch grouped by RG), place + // them in the matrix, and populate the cache. + if !missing.is_empty() { + let built = build_column_index_cells(store, location, footer_meta, &missing).await?; + if let Ok(mut cache) = COLUMN_INDEX_CACHE.lock() { + for (col, rg, cell, size) in built { + matrix[rg][col] = cell.clone(); + cache.insert(CiCellKey { path: path.clone(), col, rg }, cell, size); + } + } + } + + Some(matrix) +} + +/// Range-read + decode the requested `(col, rg)` ColumnIndex cells, grouping by +/// row group so each RG's columns share one vectored fetch + decode. Returns one +/// `(col, rg, ColumnIndexMetaData, size)` per requested cell. `None` if any +/// requested column lacks a column-index range (→ footer-only fallback). +async fn build_column_index_cells( + store: &Arc, + location: &object_store::path::Path, + footer_meta: &Arc, + missing: &[(usize, usize)], +) -> Option> { + use std::collections::BTreeMap; + let mut by_rg: BTreeMap> = BTreeMap::new(); + for &(col, rg) in missing { + by_rg.entry(rg).or_default().push(col); + } + + struct RgPlan { + rg: usize, + cols: Vec, + chunks: Vec, + range: Range, + } + let mut plans: Vec = Vec::with_capacity(by_rg.len()); + let mut fetch_ranges: Vec> = Vec::with_capacity(by_rg.len()); + for (rg, mut cols) in by_rg { + cols.sort_unstable(); + cols.dedup(); + let rgm = footer_meta.row_group(rg); + let chunks: Vec = cols.iter().map(|&i| rgm.column(i).clone()).collect(); + let range = column_index_union(&chunks)?; + fetch_ranges.push(range.clone()); + plans.push(RgPlan { rg, cols, chunks, range }); + } + if plans.is_empty() { + return None; + } + + let buffers = store.get_ranges(location, &fetch_ranges).await.ok()?; + if buffers.len() != fetch_ranges.len() { + return None; + } + + let mut out: Vec<(usize, usize, ColumnIndexMetaData, usize)> = Vec::with_capacity(missing.len()); + for (plan, buf) in plans.iter().zip(buffers.iter()) { + let reader = BufferChunkReader { base: plan.range.start, bytes: buf.clone() }; + // Deprecated but the only PUBLIC column-subset decoder (arrow-rs#8643). + #[allow(deprecated)] + let decoded = read_columns_indexes(&reader, &plan.chunks).ok()??; + if decoded.len() != plan.cols.len() { + return None; + } + let rgm = footer_meta.row_group(plan.rg); + for (k, &col) in plan.cols.iter().enumerate() { + let size = rgm.column(col).column_index_length().unwrap_or(0).max(0) as usize; + out.push((col, plan.rg, decoded[k].clone(), size)); + } + } + Some(out) +} + +// ── OffsetIndex cache lookup + build (per `(file, col)` cell, all RGs) ─────── + +/// Assemble the full-width `[rg][col]` `OffsetIndex` matrix (real entries only at +/// the resolved offset columns; empty placeholders elsewhere) from per-`(file, +/// col)` cells, decoding only the columns that miss. +/// +/// The resolved offset-column set is `predicate ∪ projection ∪ {0}` (`offset_cols +/// == None` → all columns); see [`OiCellKey`] for why each must be real. Each +/// cached cell is a column's OffsetIndex across **all** row groups, keyed only on +/// `(file, col)`, so it is decoded once per file and reused across every query +/// that reads that column irrespective of projection or predicate. +async fn get_or_build_offset_index( + store: &Arc, + location: &object_store::path::Path, + footer_meta: &Arc, + parquet_cols: &[usize], + offset_cols: Option<&[usize]>, +) -> Option { + let num_rgs = footer_meta.num_row_groups(); + if num_rgs == 0 { + return None; + } + let num_cols = footer_meta.file_metadata().schema_descr().num_columns(); - // Which columns to decode the OffsetIndex for. `None` → all columns. The - // caller already normalized `offset_cols` to include predicate cols + col 0, - // so this is just clamp/sort/dedup. Same indices for every RG (the column set - // is per-file). + // Resolve which columns need a real OffsetIndex: predicate ∪ projection ∪ {0}, + // clamped. `None` → all columns. let off_cols: Vec = match offset_cols { None => (0..num_cols).collect(), Some(cols) => { - let mut v: Vec = cols.iter().copied().filter(|&c| c < num_cols).collect(); - v.sort_unstable(); - v.dedup(); - v + let mut set: std::collections::BTreeSet = std::collections::BTreeSet::new(); + set.insert(0); // metric reads column 0 + for &c in parquet_cols { + set.insert(c); // prune reads predicate cols + } + for &c in cols { + set.insert(c); // read needs projected cols + } + set.into_iter().filter(|&c| c < num_cols).collect() } }; if off_cols.is_empty() { return None; } - // Phase 1: per RG, gather the OffsetIndex chunks for `off_cols` and — for RGs - // we're building the ColumnIndex for — the predicate columns' ColumnIndex - // chunks. Compute union byte ranges for a single vectored fetch. Bail to - // footer-only if any required index range is missing. + let path: Arc = Arc::from(location.as_ref()); + let mut matrix: ParquetOffsetIndex = (0..num_rgs) + .map(|_| (0..num_cols).map(|_| OffsetIndexBuilder::new().build()).collect()) + .collect(); + + // Phase 1: serve cached columns; collect misses. + let mut missing: Vec = Vec::new(); + if let Ok(mut cache) = OFFSET_INDEX_CACHE.lock() { + for &col in &off_cols { + let key = OiCellKey { path: path.clone(), col }; + match cache.get(&key) { + Some(column) => scatter_offset_column(&mut matrix, col, &column), + None => missing.push(col), + } + } + } else { + return None; + } + + // Phase 2: decode the missing columns (each spanning all RGs), scatter into + // the matrix, and populate the cache. + if !missing.is_empty() { + let built = build_offset_index_columns(store, location, footer_meta, &missing, num_rgs).await?; + if let Ok(mut cache) = OFFSET_INDEX_CACHE.lock() { + for (col, column, size) in built { + scatter_offset_column(&mut matrix, col, &column); + cache.insert(OiCellKey { path: path.clone(), col }, column, size); + } + } + } + + Some(matrix) +} + +/// Place a column's all-RG OffsetIndex (indexed by RG) into the matrix at `col`. +fn scatter_offset_column(matrix: &mut ParquetOffsetIndex, col: usize, column: &OiColumn) { + for (rg, entry) in column.iter().enumerate() { + if rg < matrix.len() { + matrix[rg][col] = entry.clone(); + } + } +} + +/// Range-read + decode the OffsetIndex for each requested column across **every** +/// row group (read-time safety — see [`OiCellKey`]). Returns one `(col, all-RG +/// OffsetIndex, size)` per requested column. `None` if any column lacks an +/// offset-index range (→ footer-only fallback). +async fn build_offset_index_columns( + store: &Arc, + location: &object_store::path::Path, + footer_meta: &Arc, + cols: &[usize], + num_rgs: usize, +) -> Option> { + // Group by RG: one vectored fetch + decode per RG over the requested columns. struct RgPlan { rg_idx: usize, - build_ci: bool, - pred_chunks: Vec, - off_chunks: Vec, - col_range: Option>, - off_range: Range, + chunks: Vec, + range: Range, } let mut plans: Vec = Vec::with_capacity(num_rgs); - let mut fetch_ranges: Vec> = Vec::with_capacity(num_rgs * 2); - // (off_buf_idx, Option) per plan. - let mut buf_idx: Vec<(usize, Option)> = Vec::with_capacity(num_rgs); + let mut fetch_ranges: Vec> = Vec::with_capacity(num_rgs); for rg_idx in 0..num_rgs { let rg = footer_meta.row_group(rg_idx); - let off_chunks: Vec = - off_cols.iter().map(|&i| rg.column(i).clone()).collect(); - let off_range = offset_index_union(&off_chunks)?; - - let off_i = fetch_ranges.len(); - fetch_ranges.push(off_range.clone()); - - let (pred_chunks, col_range, col_i) = if build_ci[rg_idx] { - let pred_chunks: Vec = - parquet_cols.iter().map(|&i| rg.column(i).clone()).collect(); - let col_range = column_index_union(&pred_chunks)?; - let col_i = fetch_ranges.len(); - fetch_ranges.push(col_range.clone()); - (pred_chunks, Some(col_range), Some(col_i)) - } else { - (Vec::new(), None, None) - }; - - buf_idx.push((off_i, col_i)); - plans.push(RgPlan { - rg_idx, - build_ci: build_ci[rg_idx], - pred_chunks, - off_chunks, - col_range, - off_range, - }); + let chunks: Vec = cols.iter().map(|&i| rg.column(i).clone()).collect(); + let range = offset_index_union(&chunks)?; + fetch_ranges.push(range.clone()); + plans.push(RgPlan { rg_idx, chunks, range }); } if plans.is_empty() { return None; } - // Phase 2: one vectored fetch of all required index byte ranges. let buffers = store.get_ranges(location, &fetch_ranges).await.ok()?; if buffers.len() != fetch_ranges.len() { return None; } - // Phase 3: decode and scatter into full-width per-RG vectors. Pre-fill with - // placeholders so absolute `index[rg][col]` indexing is always valid; only the - // columns in `off_cols` (OffsetIndex) / predicate cols (ColumnIndex) are - // overwritten with real entries. - let mut column_index: ParquetColumnIndex = (0..num_rgs) - .map(|_| (0..num_cols).map(|_| ColumnIndexMetaData::NONE).collect()) + // Per-column accumulator, one slot per RG (filled in RG order below). + let mut columns: Vec = cols + .iter() + .map(|_| Vec::with_capacity(num_rgs)) .collect(); - let mut offset_index: ParquetOffsetIndex = (0..num_rgs) - .map(|_| { - (0..num_cols) - .map(|_| OffsetIndexBuilder::new().build()) - .collect() - }) - .collect(); - - for (plan, &(off_i, col_i)) in plans.iter().zip(buf_idx.iter()) { - // OffsetIndex — for the scoped columns of this RG, scattered to absolute - // column positions. - let off_reader = BufferChunkReader { - base: plan.off_range.start, - bytes: buffers[off_i].clone(), - }; + for (plan, buf) in plans.iter().zip(buffers.iter()) { + let reader = BufferChunkReader { base: plan.range.start, bytes: buf.clone() }; #[allow(deprecated)] - let decoded_offs = read_offset_indexes(&off_reader, &plan.off_chunks).ok()??; - if decoded_offs.len() != off_cols.len() { + let decoded = read_offset_indexes(&reader, &plan.chunks).ok()??; + if decoded.len() != cols.len() { return None; } - let off_row = &mut offset_index[plan.rg_idx]; - for (k, &col) in off_cols.iter().enumerate() { - off_row[col] = decoded_offs[k].clone(); + for (k, entry) in decoded.into_iter().enumerate() { + columns[k].push(entry); } + } - // ColumnIndex — only for RGs we're building it for. - if plan.build_ci { - let col_range = plan.col_range.clone()?; - let col_i = col_i?; - let col_reader = BufferChunkReader { - base: col_range.start, - bytes: buffers[col_i].clone(), - }; - // `read_columns_indexes` / `read_offset_indexes` are deprecated in - // arrow-rs but are the only PUBLIC API that decodes a *column subset*. - // See module docs + apache/arrow-rs#8643. - #[allow(deprecated)] - let decoded_cols = read_columns_indexes(&col_reader, &plan.pred_chunks).ok()??; - if decoded_cols.len() != parquet_cols.len() { - return None; - } - let col_row = &mut column_index[plan.rg_idx]; - for (k, &parquet_col) in parquet_cols.iter().enumerate() { - col_row[parquet_col] = decoded_cols[k].clone(); - } + let mut out: Vec<(usize, OiColumn, usize)> = Vec::with_capacity(cols.len()); + for (k, &col) in cols.iter().enumerate() { + let mut size = 0usize; + for rg in footer_meta.row_groups() { + size += rg.column(col).offset_index_length().unwrap_or(0).max(0) as usize; } + out.push((col, std::mem::take(&mut columns[k]), size)); } - - let size = scoped_page_index_size(footer_meta, parquet_cols, surviving_rgs, &off_cols); - Some((column_index, offset_index, size)) + Some(out) } +// ── Surviving-RG computation (footer-stats prune; superset of DF's set) ────── + /// Compute the row groups that pass footer RG-statistics pruning for `predicate`. /// -/// This is a **superset** of the row groups DataFusion will actually scan: -/// DataFusion applies the same footer-stats pruning (`RowGroupAccessPlanFilter`) -/// plus bloom-filter / range / limit pruning, all of which only remove MORE row -/// groups. So scoping the predicate-column ColumnIndex to this set is safe — every -/// RG DataFusion page-prunes or scans has a real ColumnIndex. -/// -/// Returns all row groups (no scoping benefit) if the predicate can't be lowered -/// to a `PruningPredicate` or any RG lacks the stats to evaluate it — never fewer -/// than DataFusion would scan. Deterministic in `(footer_meta, schema, -/// predicate)`, so both scan paths compute the identical set (→ identical cache -/// key → cross-path sharing). +/// A **superset** of the row groups DataFusion will scan (DataFusion applies the +/// same footer-stats pruning plus bloom/range/limit, which only remove more), so +/// scoping the predicate-column ColumnIndex to this set is safe. Returns all row +/// groups if the predicate can't be lowered or stats are missing. Deterministic +/// in `(footer_meta, schema, predicate)` → both scan paths agree. pub fn surviving_row_groups( footer_meta: &ParquetMetaData, arrow_schema: &SchemaRef, predicate: &Arc, ) -> Vec { - use datafusion::parquet::arrow::arrow_reader::statistics::StatisticsConverter; + use arrow::array::{ArrayRef, BooleanArray, UInt64Array}; use datafusion::physical_optimizer::pruning::{PruningPredicate, PruningStatistics}; use datafusion::scalar::ScalarValue; - use arrow::array::{ArrayRef, BooleanArray, UInt64Array}; - use arrow::datatypes::FieldRef; let num_rgs = footer_meta.num_row_groups(); let all: Vec = (0..num_rgs).collect(); @@ -757,14 +751,10 @@ pub fn surviving_row_groups( return all; } - // Lower the predicate to a PruningPredicate; on failure, keep all RGs. let Ok(pp) = PruningPredicate::try_new(Arc::clone(predicate), Arc::clone(arrow_schema)) else { return all; }; - // Footer RG-stats source: min/max/null_count/row_count per RG per column, - // mirroring DataFusion's `RowGroupPruningStatistics` (footer stats only — no - // page index), via the public `StatisticsConverter`. struct RgStats<'a> { meta: &'a ParquetMetaData, schema: &'a SchemaRef, @@ -772,39 +762,33 @@ pub fn surviving_row_groups( } impl<'a> RgStats<'a> { fn conv(&self, col: &str) -> Option> { - StatisticsConverter::try_new( - col, - self.schema, - self.meta.file_metadata().schema_descr(), - ) - .ok() + StatisticsConverter::try_new(col, self.schema, self.meta.file_metadata().schema_descr()) + .ok() } } impl<'a> PruningStatistics for RgStats<'a> { fn min_values(&self, column: &datafusion::common::Column) -> Option { - let c = self.conv(&column.name)?; - c.row_group_mins(self.meta.row_groups().iter()).ok() + self.conv(&column.name)? + .row_group_mins(self.meta.row_groups().iter()) + .ok() } fn max_values(&self, column: &datafusion::common::Column) -> Option { - let c = self.conv(&column.name)?; - c.row_group_maxes(self.meta.row_groups().iter()).ok() + self.conv(&column.name)? + .row_group_maxes(self.meta.row_groups().iter()) + .ok() } fn num_containers(&self) -> usize { self.num_rgs } fn null_counts(&self, column: &datafusion::common::Column) -> Option { - let c = self.conv(&column.name)?; - c.row_group_null_counts(self.meta.row_groups().iter()) + self.conv(&column.name)? + .row_group_null_counts(self.meta.row_groups().iter()) .ok() .map(|a| Arc::new(a) as ArrayRef) } fn row_counts(&self) -> Option { - let counts: Vec = self - .meta - .row_groups() - .iter() - .map(|rg| rg.num_rows() as u64) - .collect(); + let counts: Vec = + self.meta.row_groups().iter().map(|rg| rg.num_rows() as u64).collect(); Some(Arc::new(UInt64Array::from(counts)) as ArrayRef) } fn contained( @@ -815,71 +799,23 @@ pub fn surviving_row_groups( None } } - // Suppress unused-import lint for FieldRef on toolchains where it's not needed. - let _ = std::marker::PhantomData::; - let stats = RgStats { - meta: footer_meta, - schema: arrow_schema, - num_rgs, - }; + let stats = RgStats { meta: footer_meta, schema: arrow_schema, num_rgs }; match pp.prune(&stats) { - Ok(mask) => { - let survivors: Vec = mask - .iter() - .enumerate() - .filter_map(|(i, keep)| if *keep { Some(i) } else { None }) - .collect(); - survivors - } - // Pruning evaluation failed → conservatively keep all RGs. + Ok(mask) => mask + .iter() + .enumerate() + .filter_map(|(i, keep)| if *keep { Some(i) } else { None }) + .collect(), Err(_) => all, } } -/// Deterministic size estimate for one cached pair, in bytes. -/// -/// Uses the **on-disk serialized lengths** from the footer -/// (`column_index_length` for predicate columns + `offset_index_length` for all -/// columns, summed over all row groups). This is a robust, monotonic proxy for -/// the decoded heap: it is never zero for a real page index, grows with the -/// number of predicate columns / row groups / pages, and — unlike -/// `ParquetMetaData::memory_size()` (which includes the footer we deliberately -/// don't store and rounds to ~0 on small files) — does not depend on private -/// `HeapSize` internals or on arrow-rs decoded-struct layout. The decoded form -/// is a small constant multiple of this; the configured limit is interpreted in -/// these units. -fn scoped_page_index_size( - footer_meta: &ParquetMetaData, - parquet_cols: &[usize], - surviving_rgs: Option<&[usize]>, - off_cols: &[usize], -) -> usize { - // ColumnIndex is only built for surviving RGs (Step 2); OffsetIndex only for - // off_cols (the predicate ∪ projection ∪ {0} set). - let build_ci = |rg_idx: usize| -> bool { - match surviving_rgs { - None => true, - Some(set) => set.contains(&rg_idx), - } - }; - let mut total = 0usize; - for (rg_idx, rg) in footer_meta.row_groups().iter().enumerate() { - if build_ci(rg_idx) { - for &pc in parquet_cols { - total += rg.column(pc).column_index_length().unwrap_or(0).max(0) as usize; - } - } - for &c in off_cols { - total += rg.column(c).offset_index_length().unwrap_or(0).max(0) as usize; - } - } - total -} +// ── Byte-range helpers + in-memory ChunkReader ────────────────────────────── /// Union of `column_index` byte ranges across the given column chunks. `None` if /// any chunk lacks a column index (we require all predicate columns to have one, -/// else we fall back to footer-only). +/// else fall back to footer-only). fn column_index_union(chunks: &[ColumnChunkMetaData]) -> Option> { range_union(chunks, |c| { let off = u64::try_from(c.column_index_offset()?).ok()?; @@ -914,10 +850,8 @@ fn range_union( /// A [`ChunkReader`] over an in-memory byte buffer representing the file region /// `[base, base + bytes.len())`. The arrow-rs page-index readers call -/// `get_bytes(absolute_offset, len)`; we translate the absolute file offset into -/// the buffer. +/// `get_bytes(absolute_offset, len)`; we translate into the buffer. struct BufferChunkReader { - /// Absolute file offset of `bytes[0]`. base: u64, bytes: Bytes, } @@ -943,8 +877,6 @@ impl ChunkReader for BufferChunkReader { } impl BufferChunkReader { - /// Translate an absolute file offset (and a length to validate) into a - /// buffer-relative start, erroring if it falls outside the prefetched span. fn rel(&self, start: u64, length: usize) -> ParquetResult { let rel = start.checked_sub(self.base).ok_or_else(|| { ParquetError::General(format!( @@ -966,43 +898,48 @@ impl BufferChunkReader { } } -// ── Test-only helpers (the cache is process-global; see SCOPED_CACHE_TEST_GUARD) ── +// ── Test-only helpers ──────────────────────────────────────────────────────── -/// Crate-wide guard so every test that touches the global scoped cache mutually -/// excludes — distinct fixtures alone aren't enough because the `InMemory` path -/// is always "data.parquet". Shared (not per-module) so future cache users -/// (`shard_scoped_reader`, the optimizer) serialize against the same lock. +/// Crate-wide guard so every test that touches the process-global caches mutually +/// excludes (distinct fixtures alone aren't enough — the `InMemory` path is always +/// "data.parquet"). Shared (not per-module) so all cache users serialize. #[cfg(test)] pub(crate) static SCOPED_CACHE_TEST_GUARD: std::sync::Mutex<()> = std::sync::Mutex::new(()); -/// Reset entries + counters AND restore the default limit. +/// Clear both caches AND restore the default limit on each. #[cfg(test)] pub(crate) fn clear_scoped_cache_for_test() { - if let Ok(mut c) = SCOPED_CACHE.lock() { - c.map.clear(); - c.used = 0; - c.tick = 0; + if let Ok(mut c) = COLUMN_INDEX_CACHE.lock() { + c.clear_keep_limit(); + c.limit = DEFAULT_SCOPED_CACHE_LIMIT; + } + if let Ok(mut c) = OFFSET_INDEX_CACHE.lock() { + c.clear_keep_limit(); c.limit = DEFAULT_SCOPED_CACHE_LIMIT; - c.hits = 0; - c.misses = 0; - c.evictions = 0; } } #[cfg(test)] -pub(crate) fn scoped_cache_len_for_test() -> usize { - SCOPED_CACHE.lock().map(|c| c.map.len()).unwrap_or(0) -} - -#[cfg(test)] -pub(crate) fn scoped_cache_bytes_for_test() -> usize { - SCOPED_CACHE.lock().map(|c| c.used).unwrap_or(0) +pub(crate) fn set_column_index_cache_limit_for_test(limit: usize) { + if let Ok(mut c) = COLUMN_INDEX_CACHE.lock() { + c.set_limit(limit); + } } +/// Combined view (sum of both caches) — test-only convenience for assertions that +/// only need "is the scoped machinery doing anything". Production code reads the +/// two caches separately ([`column_index_cache_stats`] / [`offset_index_cache_stats`]). #[cfg(test)] -pub(crate) fn set_scoped_cache_limit_for_test(limit: usize) { - if let Ok(mut c) = SCOPED_CACHE.lock() { - c.set_limit(limit); +pub(crate) fn scoped_cache_stats() -> ScopedCacheStats { + let a = column_index_cache_stats(); + let b = offset_index_cache_stats(); + ScopedCacheStats { + hits: a.hits + b.hits, + misses: a.misses + b.misses, + evictions: a.evictions + b.evictions, + entries: a.entries + b.entries, + used_bytes: a.used_bytes + b.used_bytes, + limit_bytes: a.limit_bytes.max(b.limit_bytes), } } @@ -1025,7 +962,6 @@ mod tests { use object_store::path::Path as ObjPath; use object_store::{ObjectStoreExt, PutPayload}; - // Aliased so every cache-touching test serializes on one lock. use super::SCOPED_CACHE_TEST_GUARD as CACHE_TEST_GUARD; // ── fixtures + expr helpers ────────────────────────────────────────── @@ -1040,10 +976,7 @@ mod tests { let qtys: Vec = (100..132).collect(); let batch = RecordBatch::try_new( schema.clone(), - vec![ - Arc::new(Int32Array::from(prices)), - Arc::new(Int32Array::from(qtys)), - ], + vec![Arc::new(Int32Array::from(prices)), Arc::new(Int32Array::from(qtys))], ) .unwrap(); let props = WriterProperties::builder() @@ -1059,6 +992,69 @@ mod tests { (Bytes::from(buf), schema) } + /// 4 row groups of 10 rows (`id` 0..40, `v` = id*2), page size 5. + fn four_rg_parquet() -> (Bytes, SchemaRef) { + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("v", DataType::Int32, false), + ])); + let ids: Vec = (0..40).collect(); + let vs: Vec = (0..40).map(|x| x * 2).collect(); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(ids)), Arc::new(Int32Array::from(vs))], + ) + .unwrap(); + let props = WriterProperties::builder() + .set_max_row_group_size(10) + .set_data_page_row_count_limit(5) + .set_write_batch_size(5) + .set_statistics_enabled(EnabledStatistics::Page) + .build(); + let mut buf: Vec = Vec::new(); + let mut w = ArrowWriter::try_new(&mut buf, schema.clone(), Some(props)).unwrap(); + w.write(&batch).unwrap(); + w.close().unwrap(); + (Bytes::from(buf), schema) + } + + /// 4 columns (2 int `n0`,`n1` + 2 wide string `s0`,`s1`), 1 RG, multiple pages. + fn wide4_parquet() -> (Bytes, SchemaRef) { + use arrow::array::StringArray; + let schema = Arc::new(Schema::new(vec![ + Field::new("n0", DataType::Int32, false), + Field::new("n1", DataType::Int32, false), + Field::new("s0", DataType::Utf8, false), + Field::new("s1", DataType::Utf8, false), + ])); + const ROWS: i32 = 256; + let n0: Vec = (0..ROWS).collect(); + let n1: Vec = (0..ROWS).collect(); + let s0: Vec = (0..ROWS).map(|r| format!("s0_{r:05}_padpadpad")).collect(); + let s1: Vec = (0..ROWS).map(|r| format!("s1_{r:05}_padpadpad")).collect(); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(n0)), + Arc::new(Int32Array::from(n1)), + Arc::new(StringArray::from(s0)), + Arc::new(StringArray::from(s1)), + ], + ) + .unwrap(); + let props = WriterProperties::builder() + .set_max_row_group_size(ROWS as usize) + .set_data_page_row_count_limit(32) + .set_write_batch_size(32) + .set_statistics_enabled(EnabledStatistics::Page) + .build(); + let mut buf: Vec = Vec::new(); + let mut w = ArrowWriter::try_new(&mut buf, schema.clone(), Some(props)).unwrap(); + w.write(&batch).unwrap(); + w.close().unwrap(); + (Bytes::from(buf), schema) + } + async fn stage(bytes: Bytes) -> (Arc, ObjPath) { let store: Arc = Arc::new(InMemory::new()); let loc = ObjPath::from("data.parquet"); @@ -1067,23 +1063,17 @@ mod tests { } fn footer_only(bytes: &Bytes) -> Arc { - ArrowReaderMetadata::load( - &bytes.clone(), - ArrowReaderOptions::new().with_page_index(false), - ) - .unwrap() - .metadata() - .clone() + ArrowReaderMetadata::load(&bytes.clone(), ArrowReaderOptions::new().with_page_index(false)) + .unwrap() + .metadata() + .clone() } fn full_index(bytes: &Bytes) -> Arc { - ArrowReaderMetadata::load( - &bytes.clone(), - ArrowReaderOptions::new().with_page_index(true), - ) - .unwrap() - .metadata() - .clone() + ArrowReaderMetadata::load(&bytes.clone(), ArrowReaderOptions::new().with_page_index(true)) + .unwrap() + .metadata() + .clone() } fn col(name: &str, idx: usize) -> Arc { @@ -1098,10 +1088,13 @@ mod tests { fn kept(sel: &RowSelection) -> usize { sel.iter().filter(|s| !s.skip).map(|s| s.row_count).sum() } + fn ci() -> ScopedCacheStats { + column_index_cache_stats() + } + fn oi() -> ScopedCacheStats { + offset_index_cache_stats() + } - /// Read the rows kept by `selection` for one projected leaf column — drives - /// the offset-index read path so we can prove non-predicate projected reads - /// succeed with scoped metadata. fn read_selected_column( bytes: &Bytes, meta: &Arc, @@ -1121,7 +1114,6 @@ mod tests { .with_row_selection(selection) .build() .map_err(|e| format!("build reader: {e}"))?; - let mut out = Vec::new(); while let Some(next) = reader.next() { let batch = next.map_err(|e| format!("read batch: {e}"))?; @@ -1137,9 +1129,8 @@ mod tests { Ok(out) } - // ── tests ──────────────────────────────────────────────────────────── + // ── baseline / correctness ──────────────────────────────────────────── - /// Footer-only load carries no page index — the baseline the cache restores. #[tokio::test] async fn footer_only_has_no_page_index() { let (bytes, _schema) = two_col_parquet(); @@ -1148,7 +1139,6 @@ mod tests { assert!(fo.offset_index().is_none()); } - /// Empty predicate-column set never builds an entry. #[tokio::test] async fn empty_column_set_returns_none() { let _g = CACHE_TEST_GUARD.lock().unwrap(); @@ -1157,11 +1147,10 @@ mod tests { let (store, loc) = stage(bytes.clone()).await; let fo = footer_only(&bytes); assert!(load_scoped_page_index(&store, &loc, &fo, &[]).await.is_none()); - assert_eq!(scoped_cache_len_for_test(), 0); + assert_eq!(ci().entries, 0); + assert_eq!(oi().entries, 0); } - /// The grafted metadata has a real ColumnIndex only at predicate positions - /// (NONE elsewhere) and a real OffsetIndex for every column. #[tokio::test] async fn scoped_index_is_predicate_scoped_for_column_index() { let _g = CACHE_TEST_GUARD.lock().unwrap(); @@ -1173,27 +1162,17 @@ mod tests { let cols = resolve_predicate_parquet_columns(&schema, &fo, &["price".to_string()]); assert_eq!(cols, vec![0]); - let aug = load_scoped_page_index(&store, &loc, &fo, &cols) - .await - .expect("augmentation must succeed"); - let ci = aug.column_index().expect("has column index"); - let oi = aug.offset_index().expect("has offset index"); - - assert!( - !matches!(ci[0][0], ColumnIndexMetaData::NONE), - "predicate col (price) must have a real ColumnIndex" - ); - assert!( - matches!(ci[0][1], ColumnIndexMetaData::NONE), - "non-predicate col (qty) ColumnIndex must be a NONE placeholder" - ); + let aug = load_scoped_page_index(&store, &loc, &fo, &cols).await.unwrap(); + let c = aug.column_index().unwrap(); + let o = aug.offset_index().unwrap(); + assert!(!matches!(c[0][0], ColumnIndexMetaData::NONE), "predicate col has real CI"); + assert!(matches!(c[0][1], ColumnIndexMetaData::NONE), "non-predicate col CI is NONE"); assert!( - !oi[0][0].page_locations().is_empty() && !oi[0][1].page_locations().is_empty(), - "OffsetIndex must be real for every column" + !o[0][0].page_locations().is_empty() && !o[0][1].page_locations().is_empty(), + "OffsetIndex real for every column (all-col default)" ); } - /// Scoped pruning must match full-index pruning on the predicate column. #[tokio::test] async fn scoped_pruning_matches_full_index() { let _g = CACHE_TEST_GUARD.lock().unwrap(); @@ -1201,23 +1180,16 @@ mod tests { let (bytes, schema) = two_col_parquet(); let (store, loc) = stage(bytes.clone()).await; let fo = footer_only(&bytes); - let cols = resolve_predicate_parquet_columns(&schema, &fo, &["price".to_string()]); let aug = load_scoped_page_index(&store, &loc, &fo, &cols).await.unwrap(); let full = full_index(&bytes); - - // price pages: 0..8,8..16,16..24,24..32; `price >= 20` keeps the last two. - let pp = build_pruning_predicate(&pred("price", 0, Operator::GtEq, 20), schema.clone()) - .unwrap(); + let pp = build_pruning_predicate(&pred("price", 0, Operator::GtEq, 20), schema.clone()).unwrap(); let s = PagePruner::new(&schema, Arc::clone(&aug)).prune_rg(&pp, 0, None); let f = PagePruner::new(&schema, full).prune_rg(&pp, 0, None); assert_eq!(s.as_ref().map(kept), f.as_ref().map(kept)); assert_eq!(s.as_ref().map(kept), Some(16)); } - /// A non-predicate but PROJECTED column must read correctly through the - /// scoped metadata — its OffsetIndex is real, so the offset-index read path - /// fetches the right page ranges (regression for "failed to skip rows"). #[tokio::test] async fn scoped_index_reads_non_predicate_projected_column() { let _g = CACHE_TEST_GUARD.lock().unwrap(); @@ -1225,25 +1197,22 @@ mod tests { let (bytes, schema) = two_col_parquet(); let (store, loc) = stage(bytes.clone()).await; let fo = footer_only(&bytes); - let cols = resolve_predicate_parquet_columns(&schema, &fo, &["price".to_string()]); let aug = load_scoped_page_index(&store, &loc, &fo, &cols).await.unwrap(); - let selection = RowSelection::from(vec![RowSelector::skip(16), RowSelector::select(16)]); - // Project the NON-predicate column (qty = leaf 1). - let scoped_vals = read_selected_column(&bytes, &aug, 1, selection.clone()) - .expect("non-predicate projected read must succeed with scoped metadata"); + let scoped_vals = read_selected_column(&bytes, &aug, 1, selection.clone()).unwrap(); let full = full_index(&bytes); let full_vals = read_selected_column(&bytes, &full, 1, selection).unwrap(); - let expected: Vec = (116..132).collect(); assert_eq!(scoped_vals, expected); assert_eq!(scoped_vals, full_vals); } - /// First load of a (file, column-set) is a miss; the next identical load is a - /// HIT served from cache — asserted via the counter (graft returns a fresh - /// Arc, so ptr_eq would be wrong), with no new entry and no byte growth. + // ── cache behavior: hits, independence, eviction ────────────────────── + + /// Second identical load is a pure hit in BOTH caches; no new cells/bytes. + /// Cells: predicate `price` → 1 CI cell `(col0,rg0)`; all-column OffsetIndex + /// (the default) → 2 OI cells (one per column). #[tokio::test] async fn second_load_is_cache_hit() { let _g = CACHE_TEST_GUARD.lock().unwrap(); @@ -1253,54 +1222,95 @@ mod tests { let fo = footer_only(&bytes); let cols = resolve_predicate_parquet_columns(&schema, &fo, &["price".to_string()]); - let _first = load_scoped_page_index(&store, &loc, &fo, &cols).await.unwrap(); - let s1 = scoped_cache_stats(); - assert_eq!((s1.hits, s1.misses, s1.entries), (0, 1, 1)); - let bytes_after_first = s1.used_bytes; - assert!(bytes_after_first > 0, "a real page index must cost > 0 bytes"); + let _ = load_scoped_page_index(&store, &loc, &fo, &cols).await.unwrap(); + let (c1, o1) = (ci(), oi()); + assert_eq!((c1.hits, c1.misses, c1.entries), (0, 1, 1), "1 CI cell (price,rg0)"); + assert_eq!((o1.hits, o1.misses, o1.entries), (0, 2, 2), "2 OI cells (col0,col1)"); + assert!(c1.used_bytes > 0 && o1.used_bytes > 0); - let _second = load_scoped_page_index(&store, &loc, &fo, &cols).await.unwrap(); - let s2 = scoped_cache_stats(); - assert_eq!( - (s2.hits, s2.misses, s2.entries, s2.used_bytes), - (1, 1, 1, bytes_after_first), - "second identical load: pure hit, no new entry, no byte growth" - ); + let _ = load_scoped_page_index(&store, &loc, &fo, &cols).await.unwrap(); + let (c2, o2) = (ci(), oi()); + assert_eq!((c2.hits, c2.misses, c2.entries, c2.used_bytes), (1, 1, 1, c1.used_bytes)); + assert_eq!((o2.hits, o2.misses, o2.entries, o2.used_bytes), (2, 2, 2, o1.used_bytes)); } - /// Distinct predicate-column sets are cached independently. + /// Distinct predicate columns → distinct CI cells, but the OffsetIndex column + /// cells are SHARED. Both loads default to the all-column OffsetIndex, so the + /// second load re-reads the SAME 2 OI cells from cache (no new cells). This is + /// the whole point of cell-keying: a column's index is stored once per file. #[tokio::test] - async fn distinct_column_sets_cached_independently() { + async fn distinct_predicates_share_offset_index() { let _g = CACHE_TEST_GUARD.lock().unwrap(); clear_scoped_cache_for_test(); let (bytes, schema) = two_col_parquet(); let (store, loc) = stage(bytes.clone()).await; let fo = footer_only(&bytes); - let c_price = resolve_predicate_parquet_columns(&schema, &fo, &["price".to_string()]); let c_qty = resolve_predicate_parquet_columns(&schema, &fo, &["qty".to_string()]); - assert_eq!(c_price, vec![0]); - assert_eq!(c_qty, vec![1]); - let a_price = load_scoped_page_index(&store, &loc, &fo, &c_price).await.unwrap(); - let _a_qty = load_scoped_page_index(&store, &loc, &fo, &c_qty).await.unwrap(); - assert_eq!(scoped_cache_len_for_test(), 2); + let _ = load_scoped_page_index(&store, &loc, &fo, &c_price).await.unwrap(); + let _ = load_scoped_page_index(&store, &loc, &fo, &c_qty).await.unwrap(); + + assert_eq!(ci().entries, 2, "distinct predicate cells: (price,rg0) + (qty,rg0)"); + assert_eq!(oi().entries, 2, "all-column OffsetIndex: 2 column cells, shared"); + // Second (qty) load re-read the same 2 OI cells from cache. + assert_eq!(oi().hits, 2); + } + + /// The cell-keying payoff: a predicate that ADDS a column reuses the cell the + /// first predicate already decoded, instead of re-decoding it inside a new + /// set-keyed entry. `price` then `{price, qty}` → `price`'s cell is a HIT; only + /// `qty`'s cell is freshly decoded. (Under the old set-keyed cache this was a + /// full miss that re-decoded `price`.) + #[tokio::test] + async fn adding_predicate_column_reuses_existing_cell() { + let _g = CACHE_TEST_GUARD.lock().unwrap(); + clear_scoped_cache_for_test(); + let (bytes, schema) = two_col_parquet(); + let (store, loc) = stage(bytes.clone()).await; + let fo = footer_only(&bytes); + let c_price = resolve_predicate_parquet_columns(&schema, &fo, &["price".to_string()]); + let c_both = resolve_predicate_parquet_columns( + &schema, + &fo, + &["price".to_string(), "qty".to_string()], + ); + + let _ = load_scoped_page_index(&store, &loc, &fo, &c_price).await.unwrap(); + assert_eq!((ci().hits, ci().misses, ci().entries), (0, 1, 1), "price cell decoded"); + + // Predicate now covers {price, qty}: price's cell hits, qty's cell misses. + let _ = load_scoped_page_index(&store, &loc, &fo, &c_both).await.unwrap(); + assert_eq!( + (ci().hits, ci().misses, ci().entries), + (1, 2, 2), + "price cell reused (hit); only qty cell freshly decoded" + ); + clear_scoped_cache_for_test(); + } - // qty-scoped index prunes a qty predicate; price-scoped prunes a price one. - let pp_qty = - build_pruning_predicate(&pred("qty", 1, Operator::GtEq, 120), schema.clone()).unwrap(); - // The qty entry must be usable; reload (hit) then prune. - let a_qty2 = load_scoped_page_index(&store, &loc, &fo, &c_qty).await.unwrap(); - let sel_qty = PagePruner::new(&schema, a_qty2).prune_rg(&pp_qty, 0, None).unwrap(); - assert_eq!(kept(&sel_qty), 16); + /// Two predicates on the SAME column with DIFFERENT literals resolve to the + /// same `(file, col)` parquet column, so they share the one CI cell — predicate + /// *value* never multiplies cache entries. (`status>=400` vs `status>=100`.) + #[tokio::test] + async fn different_literals_same_column_share_cell() { + let _g = CACHE_TEST_GUARD.lock().unwrap(); + clear_scoped_cache_for_test(); + let (bytes, schema) = two_col_parquet(); + let (store, loc) = stage(bytes.clone()).await; + let fo = footer_only(&bytes); + // Both predicates are on `price` (col 0) — only the literal differs, which + // never enters the cache key. + let cols = resolve_predicate_parquet_columns(&schema, &fo, &["price".to_string()]); - let pp_price = - build_pruning_predicate(&pred("price", 0, Operator::GtEq, 20), schema.clone()).unwrap(); - let sel_price = PagePruner::new(&schema, a_price).prune_rg(&pp_price, 0, None).unwrap(); - assert_eq!(kept(&sel_price), 16); + let _ = load_scoped_page_index(&store, &loc, &fo, &cols).await.unwrap(); + let _ = load_scoped_page_index(&store, &loc, &fo, &cols).await.unwrap(); + assert_eq!(ci().entries, 1, "same column → one cell regardless of literal"); + assert_eq!(ci().hits, 1); + clear_scoped_cache_for_test(); } - /// Hit/miss accounting across two distinct column-sets. + /// CI hit/miss accounting across two predicate-column sets. #[tokio::test] async fn stats_count_hits_and_misses() { let _g = CACHE_TEST_GUARD.lock().unwrap(); @@ -1312,28 +1322,21 @@ mod tests { let c_qty = resolve_predicate_parquet_columns(&schema, &fo, &["qty".to_string()]); let _ = load_scoped_page_index(&store, &loc, &fo, &c_price).await.unwrap(); - assert_eq!( - (scoped_cache_stats().hits, scoped_cache_stats().misses), - (0, 1) - ); + assert_eq!((ci().hits, ci().misses), (0, 1)); let _ = load_scoped_page_index(&store, &loc, &fo, &c_price).await.unwrap(); - assert_eq!( - (scoped_cache_stats().hits, scoped_cache_stats().misses), - (1, 1) - ); + assert_eq!((ci().hits, ci().misses), (1, 1)); let _ = load_scoped_page_index(&store, &loc, &fo, &c_qty).await.unwrap(); - assert_eq!( - (scoped_cache_stats().hits, scoped_cache_stats().misses), - (1, 2) - ); + assert_eq!((ci().hits, ci().misses), (1, 2)); let _ = load_scoped_page_index(&store, &loc, &fo, &c_price).await.unwrap(); let _ = load_scoped_page_index(&store, &loc, &fo, &c_qty).await.unwrap(); - let s = scoped_cache_stats(); + let s = ci(); assert_eq!((s.hits, s.misses, s.entries, s.evictions), (3, 2, 2, 0)); } - /// Byte-bounded LRU evicts the least-recently-used over budget, never - /// exceeding the limit, and never degrades to "cache nothing". + /// Byte-bounded LRU on the (now cell-keyed) ColumnIndex cache: with the budget + /// sized to hold ~1.5 cells, loading two distinct column cells evicts the LRU + /// one; the cache never exceeds its limit and never degrades to "cache + /// nothing"; the most-recently-used cell survives. #[tokio::test] async fn lru_evicts_over_byte_budget() { let _g = CACHE_TEST_GUARD.lock().unwrap(); @@ -1343,100 +1346,44 @@ mod tests { let fo = footer_only(&bytes); let c_price = resolve_predicate_parquet_columns(&schema, &fo, &["price".to_string()]); let c_qty = resolve_predicate_parquet_columns(&schema, &fo, &["qty".to_string()]); - let c_both = resolve_predicate_parquet_columns( - &schema, - &fo, - &["price".to_string(), "qty".to_string()], - ); - // Size one entry, then set the budget to ~1.5 entries so only one fits. + // Measure one CI cell (predicate `price` = col0 at the single RG), then set + // a budget of ~1.5 cells so a second distinct cell forces an eviction. let _ = load_scoped_page_index(&store, &loc, &fo, &c_price).await.unwrap(); - let one_size = scoped_cache_bytes_for_test(); - assert!(one_size > 0); - let budget = one_size + one_size / 2; - // Clear, set the budget, and fill from empty. + let one_cell = ci().used_bytes; + assert!(one_cell > 0); + let budget = one_cell + one_cell / 2; clear_scoped_cache_for_test(); - set_scoped_cache_limit_for_test(budget); + set_column_index_cache_limit_for_test(budget); - let _ = load_scoped_page_index(&store, &loc, &fo, &c_price).await.unwrap(); - let _ = load_scoped_page_index(&store, &loc, &fo, &c_qty).await.unwrap(); - let _ = load_scoped_page_index(&store, &loc, &fo, &c_both).await.unwrap(); + let _ = load_scoped_page_index(&store, &loc, &fo, &c_price).await.unwrap(); // cell (col0) + let _ = load_scoped_page_index(&store, &loc, &fo, &c_qty).await.unwrap(); // cell (col1) → evicts col0 - assert!( - scoped_cache_bytes_for_test() <= budget, - "cache bytes {} must stay within budget {}", - scoped_cache_bytes_for_test(), - budget - ); - assert!( - scoped_cache_len_for_test() >= 1, - "LRU must retain at least the most-recent entry" - ); - assert!(scoped_cache_stats().evictions >= 1, "something must have evicted"); + assert!(ci().used_bytes <= budget, "CI bytes {} must stay within {}", ci().used_bytes, budget); + assert_eq!(ci().entries, 1, "only the most-recent cell fits"); + assert!(ci().evictions >= 1, "the LRU cell must have evicted"); - // The most-recently-used (c_both) must still be a hit. - let hits_before = scoped_cache_stats().hits; - let _ = load_scoped_page_index(&store, &loc, &fo, &c_both).await.unwrap(); - assert_eq!( - scoped_cache_stats().hits, - hits_before + 1, - "most-recently-used entry must remain cached" - ); + // The most-recently-used cell (qty/col1) must still be a hit. + let hits_before = ci().hits; + let _ = load_scoped_page_index(&store, &loc, &fo, &c_qty).await.unwrap(); + assert_eq!(ci().hits, hits_before + 1, "MRU cell must remain cached"); clear_scoped_cache_for_test(); } - // ── Step 2: row-group scoping the ColumnIndex ───────────────────────── + // ── Step 2: RG-scoping the ColumnIndex ──────────────────────────────── - /// 4 row groups of 10 rows (`id` 0..40, `v` = id*2), page size 5 → real page - /// index, RG min/max well separated for footer-stats pruning. - fn four_rg_parquet() -> (Bytes, SchemaRef) { - let schema = Arc::new(Schema::new(vec![ - Field::new("id", DataType::Int32, false), - Field::new("v", DataType::Int32, false), - ])); - let ids: Vec = (0..40).collect(); - let vs: Vec = (0..40).map(|x| x * 2).collect(); - let batch = RecordBatch::try_new( - schema.clone(), - vec![Arc::new(Int32Array::from(ids)), Arc::new(Int32Array::from(vs))], - ) - .unwrap(); - let props = WriterProperties::builder() - .set_max_row_group_size(10) - .set_data_page_row_count_limit(5) - .set_write_batch_size(5) - .set_statistics_enabled(EnabledStatistics::Page) - .build(); - let mut buf: Vec = Vec::new(); - let mut w = ArrowWriter::try_new(&mut buf, schema.clone(), Some(props)).unwrap(); - w.write(&batch).unwrap(); - w.close().unwrap(); - (Bytes::from(buf), schema) - } - - /// `surviving_row_groups` returns exactly the RGs whose footer min/max overlap - /// the predicate — a superset of what DataFusion would scan, never fewer. #[tokio::test] async fn surviving_row_groups_matches_footer_stats_prune() { let (bytes, schema) = four_rg_parquet(); let fo = footer_only(&bytes); assert_eq!(fo.num_row_groups(), 4); - - // RG0 ids 0..10, RG1 10..20, RG2 20..30, RG3 30..40. - // `id >= 25` → RG2 (20..30, overlaps) + RG3 (30..40) survive; RG0,RG1 pruned. let p = pred("id", 0, Operator::GtEq, 25); - let survivors = surviving_row_groups(&fo, &schema, &p); - assert_eq!(survivors, vec![2, 3], "id>=25 must keep only RG2,RG3"); - - // `id < 12` → RG0 (0..10) + RG1 (10..20, overlaps) survive. + assert_eq!(surviving_row_groups(&fo, &schema, &p), vec![2, 3]); let p2 = pred("id", 0, Operator::Lt, 12); assert_eq!(surviving_row_groups(&fo, &schema, &p2), vec![0, 1]); } - /// RG-scoped load: ColumnIndex real ONLY on surviving RGs (NONE on pruned), - /// OffsetIndex real on ALL RGs (the §1e safety invariant), and pruning on a - /// surviving RG matches the full index. #[tokio::test] async fn rg_scoped_load_builds_column_index_only_for_survivors() { let _g = CACHE_TEST_GUARD.lock().unwrap(); @@ -1444,54 +1391,34 @@ mod tests { let (bytes, schema) = four_rg_parquet(); let (store, loc) = stage(bytes.clone()).await; let fo = footer_only(&bytes); - let cols = resolve_predicate_parquet_columns(&schema, &fo, &["id".to_string()]); - assert_eq!(cols, vec![0]); let surviving = vec![2usize, 3usize]; - let aug = load_scoped_page_index_rgs(&store, &loc, &fo, &cols, &surviving) - .await - .expect("rg-scoped load must succeed"); - let ci = aug.column_index().expect("has column index"); - let oi = aug.offset_index().expect("has offset index"); - assert_eq!(ci.len(), 4); - + let aug = load_scoped_page_index_rgs(&store, &loc, &fo, &cols, &surviving).await.unwrap(); + let c = aug.column_index().unwrap(); + let o = aug.offset_index().unwrap(); + assert_eq!(c.len(), 4); for &rg in &surviving { - assert!( - !matches!(ci[rg][0], ColumnIndexMetaData::NONE), - "surviving RG {rg} must have a real ColumnIndex for `id`" - ); + assert!(!matches!(c[rg][0], ColumnIndexMetaData::NONE), "survivor RG {rg} real CI"); } for &rg in &[0usize, 1usize] { - assert!( - matches!(ci[rg][0], ColumnIndexMetaData::NONE), - "pruned RG {rg} ColumnIndex must be NONE" - ); + assert!(matches!(c[rg][0], ColumnIndexMetaData::NONE), "pruned RG {rg} NONE CI"); } - // OffsetIndex must be REAL for EVERY RG and EVERY column (safety). for rg in 0..4 { - for c in 0..2 { - assert!( - !oi[rg][c].page_locations().is_empty(), - "OffsetIndex must be real for all RGs/cols (rg={rg} col={c})" - ); + for cc in 0..2 { + assert!(!o[rg][cc].page_locations().is_empty(), "OI real for all rg/col"); } } - - // Pruning on a surviving RG matches the full index. let full = full_index(&bytes); let pp = build_pruning_predicate(&pred("id", 0, Operator::GtEq, 25), schema.clone()).unwrap(); let s = PagePruner::new(&schema, Arc::clone(&aug)).prune_rg(&pp, 2, None); let f = PagePruner::new(&schema, full).prune_rg(&pp, 2, None); assert_eq!(s.as_ref().map(kept), f.as_ref().map(kept)); - clear_scoped_cache_for_test(); } - /// RG-scoping shrinks the cached ColumnIndex: a 2-of-4-RG scoped entry uses - /// fewer bytes than the all-RG (Step 1) entry for the same predicate column. #[tokio::test] - async fn rg_scoping_reduces_cached_bytes() { + async fn rg_scoping_reduces_column_index_bytes() { let _g = CACHE_TEST_GUARD.lock().unwrap(); clear_scoped_cache_for_test(); let (bytes, schema) = four_rg_parquet(); @@ -1499,25 +1426,18 @@ mod tests { let fo = footer_only(&bytes); let cols = resolve_predicate_parquet_columns(&schema, &fo, &["id".to_string()]); - // All-RG (Step 1). let _ = load_scoped_page_index(&store, &loc, &fo, &cols).await.unwrap(); - let all_rg_bytes = scoped_cache_bytes_for_test(); + let all_rg = ci().used_bytes; clear_scoped_cache_for_test(); - - // RG-scoped to 2 of 4 RGs (Step 2). let _ = load_scoped_page_index_rgs(&store, &loc, &fo, &cols, &[2, 3]).await.unwrap(); - let scoped_bytes = scoped_cache_bytes_for_test(); - - assert!( - scoped_bytes < all_rg_bytes, - "RG-scoped bytes ({scoped_bytes}) must be < all-RG bytes ({all_rg_bytes})" - ); + assert!(ci().used_bytes < all_rg, "RG-scoped CI bytes {} < all-RG {}", ci().used_bytes, all_rg); clear_scoped_cache_for_test(); } - /// The RG-scoped entry is keyed by its surviving-RG set: the same `(file, - /// cols, surviving_rgs)` hits; a different surviving-RG set is a distinct - /// entry (so two paths agree only when they compute the same survivor set). + /// CI cells are keyed per `(col, rg)`. Loading survivors {2,3} caches cells + /// (id,rg2) + (id,rg3); reloading the same survivor set hits both; a different + /// survivor set {0,1} adds two fresh cells. So a column's per-RG index is + /// reused across overlapping survivor sets instead of re-decoded per set. #[tokio::test] async fn rg_scoped_key_includes_surviving_rgs() { let _g = CACHE_TEST_GUARD.lock().unwrap(); @@ -1528,74 +1448,135 @@ mod tests { let cols = resolve_predicate_parquet_columns(&schema, &fo, &["id".to_string()]); let _ = load_scoped_page_index_rgs(&store, &loc, &fo, &cols, &[2, 3]).await.unwrap(); - assert_eq!((scoped_cache_stats().misses, scoped_cache_stats().entries), (1, 1)); + assert_eq!((ci().misses, ci().entries), (2, 2), "cells (id,rg2)+(id,rg3)"); + let _ = load_scoped_page_index_rgs(&store, &loc, &fo, &cols, &[2, 3]).await.unwrap(); + assert_eq!((ci().hits, ci().entries), (2, 2), "same survivors → both cells hit"); + let _ = load_scoped_page_index_rgs(&store, &loc, &fo, &cols, &[0, 1]).await.unwrap(); + assert_eq!((ci().misses, ci().entries), (4, 4), "new survivors → 2 fresh cells"); + // OI stayed all-columns across all three → 2 column cells, shared. + assert_eq!(oi().entries, 2); + clear_scoped_cache_for_test(); + } + + /// Partial-overlap survivor sets only decode the NEW row groups. Load + /// survivors {2,3} (cells rg2,rg3), then {1,2,3}: rg2+rg3 hit, only rg1 is + /// freshly decoded. Proves RG-scoping reuses per-RG cells across overlapping + /// survivor sets rather than re-decoding the whole set. + #[tokio::test] + async fn overlapping_survivor_sets_decode_only_new_rgs() { + let _g = CACHE_TEST_GUARD.lock().unwrap(); + clear_scoped_cache_for_test(); + let (bytes, schema) = four_rg_parquet(); + let (store, loc) = stage(bytes.clone()).await; + let fo = footer_only(&bytes); + let cols = resolve_predicate_parquet_columns(&schema, &fo, &["id".to_string()]); - // Same survivor set → hit, no new entry. let _ = load_scoped_page_index_rgs(&store, &loc, &fo, &cols, &[2, 3]).await.unwrap(); - assert_eq!((scoped_cache_stats().hits, scoped_cache_stats().entries), (1, 1)); + assert_eq!((ci().hits, ci().misses, ci().entries), (0, 2, 2), "cells (id,rg2)+(id,rg3)"); - // Different survivor set → distinct entry (miss). - let _ = load_scoped_page_index_rgs(&store, &loc, &fo, &cols, &[0, 1]).await.unwrap(); - assert_eq!((scoped_cache_stats().misses, scoped_cache_stats().entries), (2, 2)); + // {1,2,3}: rg2 & rg3 are cached (2 hits); only rg1 is new (1 miss). + let _ = load_scoped_page_index_rgs(&store, &loc, &fo, &cols, &[1, 2, 3]).await.unwrap(); + assert_eq!( + (ci().hits, ci().misses, ci().entries), + (2, 3, 3), + "rg2+rg3 reused (2 hits); only rg1 freshly decoded" + ); + clear_scoped_cache_for_test(); + } + /// The combined payoff across BOTH axes: a second query that adds a new + /// predicate column AND scans a wider RG set decodes only the genuinely new + /// `(col, rg)` cells. Uses `wide4` (1 RG) for the column axis and asserts CI + /// cell-level hit/miss deltas. + #[tokio::test] + async fn new_column_combination_caches_only_new_column_cells() { + let _g = CACHE_TEST_GUARD.lock().unwrap(); + clear_scoped_cache_for_test(); + let (bytes, schema) = wide4_parquet(); // n0,n1,s0,s1 — 1 RG + let (store, loc) = stage(bytes.clone()).await; + let fo = footer_only(&bytes); + let c_n0 = resolve_predicate_parquet_columns(&schema, &fo, &["n0".to_string()]); + let c_n0_n1 = resolve_predicate_parquet_columns( + &schema, + &fo, + &["n0".to_string(), "n1".to_string()], + ); + let c_n1_s0 = resolve_predicate_parquet_columns( + &schema, + &fo, + &["n1".to_string(), "s0".to_string()], + ); + + // {n0}: 1 new cell. + let _ = load_scoped_page_index(&store, &loc, &fo, &c_n0).await.unwrap(); + assert_eq!((ci().hits, ci().misses, ci().entries), (0, 1, 1)); + // {n0,n1}: n0 hits, n1 new. + let _ = load_scoped_page_index(&store, &loc, &fo, &c_n0_n1).await.unwrap(); + assert_eq!((ci().hits, ci().misses, ci().entries), (1, 2, 2), "n0 reused; n1 new"); + // {n1,s0}: n1 hits, s0 new. + let _ = load_scoped_page_index(&store, &loc, &fo, &c_n1_s0).await.unwrap(); + assert_eq!((ci().hits, ci().misses, ci().entries), (2, 3, 3), "n1 reused; s0 new"); + clear_scoped_cache_for_test(); + } + + /// OffsetIndex equivalent: different projections cache only the new column + /// cells. Project {s0} (offset cols n1∪s0∪{0}), then {s1} (offset cols + /// n1∪s1∪{0}) — the shared cols (0, n1) hit; only the genuinely new projected + /// column is decoded. + #[tokio::test] + async fn different_projections_cache_only_new_offset_columns() { + let _g = CACHE_TEST_GUARD.lock().unwrap(); + clear_scoped_cache_for_test(); + let (bytes, schema) = wide4_parquet(); // n0=0,n1=1,s0=2,s1=3 — 1 RG + let (store, loc) = stage(bytes.clone()).await; + let fo = footer_only(&bytes); + let pred_cols = resolve_predicate_parquet_columns(&schema, &fo, &["n1".to_string()]); + + // Project s0 (col 2): offset cols = {0, 1(n1), 2(s0)} → 3 new cells. + let _ = load_scoped_page_index_cols(&store, &loc, &fo, &pred_cols, &[2]).await.unwrap(); + assert_eq!((oi().hits, oi().misses, oi().entries), (0, 3, 3), "cols 0,1,2"); + + // Project s1 (col 3): offset cols = {0, 1, 3}. Cols 0 & 1 hit; col 3 new. + let _ = load_scoped_page_index_cols(&store, &loc, &fo, &pred_cols, &[3]).await.unwrap(); + assert_eq!( + (oi().hits, oi().misses, oi().entries), + (2, 4, 4), + "cols 0,1 reused (2 hits); only col 3 freshly decoded" + ); clear_scoped_cache_for_test(); } // ── Step 2: column-scoping the OffsetIndex ──────────────────────────── - /// Column-scoped load: OffsetIndex real ONLY for the scoped columns (the - /// caller's set ∪ predicate ∪ {0}); empty placeholder for the rest. ColumnIndex - /// still predicate-scoped. #[tokio::test] async fn col_scoped_offset_index_only_for_requested_columns() { let _g = CACHE_TEST_GUARD.lock().unwrap(); clear_scoped_cache_for_test(); - // 4 columns: predicate on n1, project s0 (leaf 2). Offset set should be - // {0 (metric), 1 (predicate), 2 (projection)} — NOT column 3 (s1). let (bytes, schema) = wide4_parquet(); let (store, loc) = stage(bytes.clone()).await; let fo = footer_only(&bytes); let pred_cols = resolve_predicate_parquet_columns(&schema, &fo, &["n1".to_string()]); assert_eq!(pred_cols, vec![1]); - // Caller passes projection = {2}; fn unions in predicate {1} + {0}. - let aug = load_scoped_page_index_cols(&store, &loc, &fo, &pred_cols, &[2]) - .await - .expect("col-scoped load must succeed"); - let oi = aug.offset_index().expect("has offset index"); - + let aug = load_scoped_page_index_cols(&store, &loc, &fo, &pred_cols, &[2]).await.unwrap(); + let o = aug.offset_index().unwrap(); for &c in &[0usize, 1, 2] { - assert!( - !oi[0][c].page_locations().is_empty(), - "col {c} (predicate/projection/metric) must have a real OffsetIndex" - ); + assert!(!o[0][c].page_locations().is_empty(), "col {c} (pred/proj/metric) real OI"); } - assert!( - oi[0][3].page_locations().is_empty(), - "col 3 (not projected/predicate/0) OffsetIndex must be an empty placeholder" - ); + assert!(o[0][3].page_locations().is_empty(), "col 3 OI is empty placeholder"); } - /// A projected non-predicate column reads correctly through the col-scoped - /// metadata (its OffsetIndex is real because it's in the projection set) — the - /// read-path safety invariant. #[tokio::test] async fn col_scoped_reads_projected_non_predicate_column() { let _g = CACHE_TEST_GUARD.lock().unwrap(); clear_scoped_cache_for_test(); - let (bytes, schema) = two_col_parquet(); // price (0), qty (1) + let (bytes, schema) = two_col_parquet(); let (store, loc) = stage(bytes.clone()).await; let fo = footer_only(&bytes); let pred_cols = resolve_predicate_parquet_columns(&schema, &fo, &["price".to_string()]); - - // Predicate on price (0); project qty (1). Offset set = {0,1}. - let aug = load_scoped_page_index_cols(&store, &loc, &fo, &pred_cols, &[1]) - .await - .unwrap(); - + let aug = load_scoped_page_index_cols(&store, &loc, &fo, &pred_cols, &[1]).await.unwrap(); let selection = RowSelection::from(vec![RowSelector::skip(16), RowSelector::select(16)]); - let scoped_vals = read_selected_column(&bytes, &aug, 1, selection.clone()) - .expect("projected non-predicate read must succeed with col-scoped metadata"); + let scoped_vals = read_selected_column(&bytes, &aug, 1, selection.clone()).unwrap(); let full = full_index(&bytes); let full_vals = read_selected_column(&bytes, &full, 1, selection).unwrap(); let expected: Vec = (116..132).collect(); @@ -1603,9 +1584,8 @@ mod tests { assert_eq!(scoped_vals, full_vals); } - /// Column-scoping the OffsetIndex shrinks the cached bytes vs all-columns. #[tokio::test] - async fn col_scoping_reduces_cached_bytes() { + async fn col_scoping_reduces_offset_index_bytes() { let _g = CACHE_TEST_GUARD.lock().unwrap(); clear_scoped_cache_for_test(); let (bytes, schema) = wide4_parquet(); @@ -1613,24 +1593,18 @@ mod tests { let fo = footer_only(&bytes); let pred_cols = resolve_predicate_parquet_columns(&schema, &fo, &["n1".to_string()]); - // All-columns OffsetIndex (Step 1). let _ = load_scoped_page_index(&store, &loc, &fo, &pred_cols).await.unwrap(); - let all_cols_bytes = scoped_cache_bytes_for_test(); + let all_cols = oi().used_bytes; clear_scoped_cache_for_test(); - - // Col-scoped to {0,1,2} (drops col 3's OffsetIndex). let _ = load_scoped_page_index_cols(&store, &loc, &fo, &pred_cols, &[2]).await.unwrap(); - let scoped_bytes = scoped_cache_bytes_for_test(); - - assert!( - scoped_bytes < all_cols_bytes, - "col-scoped bytes ({scoped_bytes}) must be < all-columns bytes ({all_cols_bytes})" - ); + assert!(oi().used_bytes < all_cols, "col-scoped OI {} < all-col {}", oi().used_bytes, all_cols); clear_scoped_cache_for_test(); } - /// If the offset-column set covers every column, the entry collapses to the - /// Step-1 all-columns artifact (shares its cache slot, no separate entry). + /// Cell-keying makes OffsetIndex reuse automatic: an all-columns load caches + /// per-column cells, and a later column-scoped load whose set is covered by + /// those cells hits them — no new entries, no special "collapse to all-columns + /// sentinel" needed (the prior set-keyed design's mechanism). #[tokio::test] async fn col_scoping_full_coverage_collapses_to_all_columns_entry() { let _g = CACHE_TEST_GUARD.lock().unwrap(); @@ -1640,53 +1614,35 @@ mod tests { let fo = footer_only(&bytes); let pred_cols = resolve_predicate_parquet_columns(&schema, &fo, &["price".to_string()]); - // Step-1 all-columns entry. let _ = load_scoped_page_index(&store, &loc, &fo, &pred_cols).await.unwrap(); - assert_eq!(scoped_cache_len_for_test(), 1); - - // Col-scoped to {1} → union {0,1} = all 2 columns → collapses to the same - // key → HIT, no new entry. + assert_eq!(oi().entries, 2, "all-columns load caches 2 column cells"); + // Project {1}; union {0,1} = both columns, both already cached → 2 hits. let _ = load_scoped_page_index_cols(&store, &loc, &fo, &pred_cols, &[1]).await.unwrap(); - assert_eq!(scoped_cache_len_for_test(), 1, "full coverage must reuse the all-columns entry"); - assert!(scoped_cache_stats().hits >= 1); - + assert_eq!(oi().entries, 2, "covered columns reuse their cells, no new entries"); + assert_eq!(oi().hits, 2); clear_scoped_cache_for_test(); } - /// 4 columns (2 int `n0`,`n1` + 2 wide string `s0`,`s1`), 1 RG, multiple pages. - fn wide4_parquet() -> (Bytes, SchemaRef) { - use arrow::array::StringArray; - let schema = Arc::new(Schema::new(vec![ - Field::new("n0", DataType::Int32, false), - Field::new("n1", DataType::Int32, false), - Field::new("s0", DataType::Utf8, false), - Field::new("s1", DataType::Utf8, false), - ])); - const ROWS: i32 = 256; - let n0: Vec = (0..ROWS).collect(); - let n1: Vec = (0..ROWS).collect(); - let s0: Vec = (0..ROWS).map(|r| format!("s0_{r:05}_padpadpad")).collect(); - let s1: Vec = (0..ROWS).map(|r| format!("s1_{r:05}_padpadpad")).collect(); - let batch = RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(Int32Array::from(n0)), - Arc::new(Int32Array::from(n1)), - Arc::new(StringArray::from(s0)), - Arc::new(StringArray::from(s1)), - ], - ) - .unwrap(); - let props = WriterProperties::builder() - .set_max_row_group_size(ROWS as usize) - .set_data_page_row_count_limit(32) - .set_write_batch_size(32) - .set_statistics_enabled(EnabledStatistics::Page) - .build(); - let mut buf: Vec = Vec::new(); - let mut w = ArrowWriter::try_new(&mut buf, schema.clone(), Some(props)).unwrap(); - w.write(&batch).unwrap(); - w.close().unwrap(); - (Bytes::from(buf), schema) + /// The fully-scoped entry point: CI RG-scoped + OI column-scoped together. + #[tokio::test] + async fn fully_scoped_load_combines_both_axes() { + let _g = CACHE_TEST_GUARD.lock().unwrap(); + clear_scoped_cache_for_test(); + let (bytes, schema) = four_rg_parquet(); + let (store, loc) = stage(bytes.clone()).await; + let fo = footer_only(&bytes); + let cols = resolve_predicate_parquet_columns(&schema, &fo, &["id".to_string()]); + + // CI scoped to RGs {2,3}; OI scoped to {0,1} = all 2 cols (collapses to all). + let aug = load_scoped_page_index_scoped(&store, &loc, &fo, &cols, &[2, 3], &[1]) + .await + .unwrap(); + let c = aug.column_index().unwrap(); + assert!(matches!(c[0][0], ColumnIndexMetaData::NONE), "RG0 pruned → NONE CI"); + assert!(!matches!(c[2][0], ColumnIndexMetaData::NONE), "RG2 survivor → real CI"); + // CI cells (id,rg2)+(id,rg3) = 2; OI cells (col0)+(col1) = 2. + assert_eq!(ci().entries, 2); + assert_eq!(oi().entries, 2); + clear_scoped_cache_for_test(); } } diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/stats.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/stats.rs index 58f83cafaa09d..84c146b051bed 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/stats.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/stats.rs @@ -20,7 +20,7 @@ //! | `plan_setup` | `TaskMonitorRepr` | 5 × i64 | //! | `fragment_executor_gate` | `PartitionGateRepr` | 8 × i64 | //! | `adaptive_budget` | `AdaptiveBudgetRepr` | 2 × i64 | -//! | `cache_stats` | `CacheStatsRepr` | 15 × i64 (3 × 5: metadata + statistics + scoped-page-index caches) | +//! | `cache_stats` | `CacheStatsRepr` | 20 × i64 (4 × 5: metadata + statistics + scoped column-index + scoped offset-index caches) | //! | `search_stats` | `SearchStatsRepr` | 17 × i64 | use tokio::runtime::Handle; @@ -92,12 +92,14 @@ pub struct CacheGroupRepr { pub struct CacheStatsRepr { pub metadata_cache: CacheGroupRepr, pub statistics_cache: CacheGroupRepr, - /// Process-wide scoped page-index cache (see - /// [`crate::indexed_table::page_index_loader`]). `hit_count`/`miss_count` - /// are cumulative since process start; `entry_count`/`memory_bytes`/ - /// `size_limit_bytes` are point-in-time. (Evictions are tracked internally - /// but not surfaced in this 5-field group.) - pub scoped_page_index_cache: CacheGroupRepr, + /// The predicate-driven scoped ColumnIndex cache, keyed per `(file, col, rg)` + /// cell (see [`crate::indexed_table::page_index_loader`]). `hit_count`/ + /// `miss_count` are cumulative since process start; `entry_count` (cells) / + /// `memory_bytes` / `size_limit_bytes` are point-in-time. + pub column_index_cache: CacheGroupRepr, + /// The projection-driven scoped OffsetIndex cache, keyed per `(file, col)` + /// cell (the value spans all row groups). + pub offset_index_cache: CacheGroupRepr, } impl Default for CacheGroupRepr { @@ -117,7 +119,8 @@ impl Default for CacheStatsRepr { Self { metadata_cache: CacheGroupRepr::default(), statistics_cache: CacheGroupRepr::default(), - scoped_page_index_cache: CacheGroupRepr::default(), + column_index_cache: CacheGroupRepr::default(), + offset_index_cache: CacheGroupRepr::default(), } } } @@ -168,14 +171,14 @@ const _: () = assert!(std::mem::size_of::() == 5 * 8); const _: () = assert!(std::mem::size_of::() == 8 * 8); const _: () = assert!(std::mem::size_of::() == 2 * 8); const _: () = assert!(std::mem::size_of::() == 5 * 8); -const _: () = assert!(std::mem::size_of::() == 15 * 8); +const _: () = assert!(std::mem::size_of::() == 20 * 8); const _: () = assert!(std::mem::size_of::() == 17 * 8); -const _: () = assert!(std::mem::size_of::() == 80 * 8); +const _: () = assert!(std::mem::size_of::() == 85 * 8); pub mod layout { use super::*; pub const BUFFER_BYTE_SIZE: usize = std::mem::size_of::(); - const _: () = assert!(BUFFER_BYTE_SIZE == 640); + const _: () = assert!(BUFFER_BYTE_SIZE == 680); } /// Snapshot a `RuntimeMonitor` and return a populated `RuntimeMetricsRepr`. @@ -311,9 +314,19 @@ pub fn pack_cache_stats(mgr: &CustomCacheManager) -> CacheStatsRepr { memory_bytes: statistics_memory, size_limit_bytes: mgr.statistics_cache_size_limit() as i64, }, - scoped_page_index_cache: { - // Process-global cache, not owned by the manager — read it directly. - let s = crate::indexed_table::page_index_loader::scoped_cache_stats(); + // Process-global caches, not owned by the manager — read them directly. + column_index_cache: { + let s = crate::indexed_table::page_index_loader::column_index_cache_stats(); + CacheGroupRepr { + hit_count: s.hits as i64, + miss_count: s.misses as i64, + entry_count: s.entries as i64, + memory_bytes: s.used_bytes as i64, + size_limit_bytes: s.limit_bytes as i64, + } + }, + offset_index_cache: { + let s = crate::indexed_table::page_index_loader::offset_index_cache_stats(); CacheGroupRepr { hit_count: s.hits as i64, miss_count: s.misses as i64, @@ -414,7 +427,7 @@ mod tests { search_stats: crate::search_stats::snapshot(), }; - assert_eq!(layout::BUFFER_BYTE_SIZE, 640); + assert_eq!(layout::BUFFER_BYTE_SIZE, 680); assert!(buf.io_runtime.workers_count > 0, "IO runtime workers_count should be > 0, got {}", buf.io_runtime.workers_count); assert!(buf.fragment_executor_gate.max_permits > 0, "fragment_executor_gate max_permits should be > 0, got {}", buf.fragment_executor_gate.max_permits); @@ -429,9 +442,9 @@ mod tests { #[test] fn test_df_stats_buffer_too_small() { // Verify that the buffer size assertion holds - assert_eq!(std::mem::size_of::(), 640); - assert_eq!(layout::BUFFER_BYTE_SIZE, 640); - // A buffer smaller than 640 bytes should be rejected by df_stats. + assert_eq!(std::mem::size_of::(), 680); + assert_eq!(layout::BUFFER_BYTE_SIZE, 680); + // A buffer smaller than 680 bytes should be rejected by df_stats. // We can't call df_stats directly without a runtime manager, // but we verify the constant is correct. assert!(layout::BUFFER_BYTE_SIZE > 0); @@ -449,33 +462,43 @@ mod tests { assert_eq!(g.memory_bytes, 0); assert_eq!(g.size_limit_bytes, 0); } - // The scoped page-index cache is process-global, not manager-owned, so it - // is NOT necessarily zero here (other tests may have populated it). It - // must always advertise a positive byte budget, though. - assert!(repr.scoped_page_index_cache.size_limit_bytes > 0); + // The scoped page-index caches are process-global, not manager-owned, so + // they are NOT necessarily zero here (other tests may have populated them). + // Each must always advertise a positive byte budget, though. + assert!(repr.column_index_cache.size_limit_bytes > 0); + assert!(repr.offset_index_cache.size_limit_bytes > 0); } #[test] fn test_pack_cache_stats_reflects_scoped_page_index_counters() { use crate::custom_cache_manager::CustomCacheManager; use crate::indexed_table::page_index_loader::{ - scoped_cache_stats, set_scoped_cache_limit, SCOPED_CACHE_TEST_GUARD, + column_index_cache_stats, offset_index_cache_stats, set_column_index_cache_limit, + set_offset_index_cache_limit, SCOPED_CACHE_TEST_GUARD, }; - // Serialize against the process-global scoped cache. + // Serialize against the process-global scoped caches. let _g = SCOPED_CACHE_TEST_GUARD.lock().unwrap(); - // Set a known budget and read it back through the packed repr. - set_scoped_cache_limit(123 * 1024 * 1024); + // Set known (distinct) budgets and read them back through the packed repr. + set_column_index_cache_limit(123 * 1024 * 1024); + set_offset_index_cache_limit(45 * 1024 * 1024); let mgr = CustomCacheManager::new(); let repr = pack_cache_stats(&mgr); - let live = scoped_cache_stats(); - - assert_eq!(repr.scoped_page_index_cache.size_limit_bytes, 123 * 1024 * 1024); - assert_eq!(repr.scoped_page_index_cache.hit_count, live.hits as i64); - assert_eq!(repr.scoped_page_index_cache.miss_count, live.misses as i64); - assert_eq!(repr.scoped_page_index_cache.entry_count, live.entries as i64); - assert_eq!(repr.scoped_page_index_cache.memory_bytes, live.used_bytes as i64); + let ci = column_index_cache_stats(); + let oi = offset_index_cache_stats(); + + assert_eq!(repr.column_index_cache.size_limit_bytes, 123 * 1024 * 1024); + assert_eq!(repr.column_index_cache.hit_count, ci.hits as i64); + assert_eq!(repr.column_index_cache.miss_count, ci.misses as i64); + assert_eq!(repr.column_index_cache.entry_count, ci.entries as i64); + assert_eq!(repr.column_index_cache.memory_bytes, ci.used_bytes as i64); + + assert_eq!(repr.offset_index_cache.size_limit_bytes, 45 * 1024 * 1024); + assert_eq!(repr.offset_index_cache.hit_count, oi.hits as i64); + assert_eq!(repr.offset_index_cache.miss_count, oi.misses as i64); + assert_eq!(repr.offset_index_cache.entry_count, oi.entries as i64); + assert_eq!(repr.offset_index_cache.memory_bytes, oi.used_bytes as i64); } #[test] diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionPlugin.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionPlugin.java index 7422149267f31..4ac7534bbf39b 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionPlugin.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionPlugin.java @@ -748,7 +748,13 @@ public List getSupportedFormats() { @Override public List> getActions() { - return List.of(new ActionHandler<>(DataFusionStatsActionType.INSTANCE, TransportDataFusionStatsAction.class)); + return List.of( + new ActionHandler<>(DataFusionStatsActionType.INSTANCE, TransportDataFusionStatsAction.class), + new ActionHandler<>( + org.opensearch.be.datafusion.action.stats.ClearScopedPageIndexCacheActionType.INSTANCE, + org.opensearch.be.datafusion.action.stats.TransportClearScopedPageIndexCacheAction.class + ) + ); } @Override diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionSettings.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionSettings.java index 417bd55d6469e..6ac8d791fd666 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionSettings.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionSettings.java @@ -298,7 +298,8 @@ static String derivePoolMinDefault(Settings settings, int percent) { CacheSettings.STATISTICS_CACHE_EVICTION_TYPE, CacheSettings.METADATA_CACHE_ENABLED, CacheSettings.STATISTICS_CACHE_ENABLED, - CacheSettings.SCOPED_PAGE_INDEX_CACHE_SIZE_LIMIT, + CacheSettings.COLUMN_INDEX_CACHE_SIZE_LIMIT, + CacheSettings.OFFSET_INDEX_CACHE_SIZE_LIMIT, // Concurrency gate settings CONCURRENCY_DATANODE_MULTIPLIER, diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/action/stats/RestClearScopedPageIndexCacheAction.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/action/stats/RestClearScopedPageIndexCacheAction.java index 81fbb9de09402..cab44a7ce84d8 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/action/stats/RestClearScopedPageIndexCacheAction.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/action/stats/RestClearScopedPageIndexCacheAction.java @@ -8,13 +8,10 @@ package org.opensearch.be.datafusion.action.stats; -import org.opensearch.be.datafusion.nativelib.NativeBridge; -import org.opensearch.transport.client.node.NodeClient; -import org.opensearch.core.rest.RestStatus; -import org.opensearch.core.xcontent.XContentBuilder; import org.opensearch.rest.BaseRestHandler; -import org.opensearch.rest.BytesRestResponse; import org.opensearch.rest.RestRequest; +import org.opensearch.rest.action.RestActions.NodesResponseRestListener; +import org.opensearch.transport.client.node.NodeClient; import java.io.IOException; import java.util.List; @@ -23,18 +20,18 @@ import static org.opensearch.rest.RestRequest.Method.POST; /** - * Clears the process-global scoped page-index cache on the local node (drops - * entries + resets counters, keeps the configured budget). + * Clears the process-global scoped page-index caches (ColumnIndex + OffsetIndex) + * across ALL nodes in the cluster (drops entries + resets counters, keeps the + * configured budgets). * - *

Operational/testing convenience: reset the cache and re-measure via the - * stats endpoint without a cluster restart. + *

Operational/testing convenience: reset the caches and re-measure via the + * stats endpoint without a cluster restart. Because the caches are per-node + * process-global singletons, this broadcasts to every node via + * {@link ClearScopedPageIndexCacheActionType} — clearing only the receiving node + * would leave the cluster-aggregated stats non-zero. * * 

POST /_plugins/_analytics_backend_datafusion/cache/scoped_page_index/_clear
* - *

Note: this acts on the node that receives the request only. For a single-node - * benchmark cluster that is exactly what's wanted; for multi-node clusters, send - * it to each node (or restart). - * * @opensearch.internal */ public class RestClearScopedPageIndexCacheAction extends BaseRestHandler { @@ -53,14 +50,12 @@ public List routes() { @Override protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) throws IOException { - return channel -> { - NativeBridge.clearScopedPageIndexCache(); - XContentBuilder builder = channel.newBuilder(); - builder.startObject(); - builder.field("acknowledged", true); - builder.field("cleared", "scoped_page_index_cache"); - builder.endObject(); - channel.sendResponse(new BytesRestResponse(RestStatus.OK, builder)); - }; + // Empty node-ids → broadcast to all nodes. + ClearScopedPageIndexCacheNodesRequest nodesRequest = new ClearScopedPageIndexCacheNodesRequest(); + return channel -> client.execute( + ClearScopedPageIndexCacheActionType.INSTANCE, + nodesRequest, + new NodesResponseRestListener<>(channel) + ); } } diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/cache/CacheSettings.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/cache/CacheSettings.java index c849df35480a6..1ae82712db2c0 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/cache/CacheSettings.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/cache/CacheSettings.java @@ -20,7 +20,8 @@ public class CacheSettings { public static final String METADATA_CACHE_SIZE_LIMIT_KEY = "datafusion.metadata.cache.size.limit"; public static final String STATISTICS_CACHE_SIZE_LIMIT_KEY = "datafusion.statistics.cache.size.limit"; - public static final String SCOPED_PAGE_INDEX_CACHE_SIZE_LIMIT_KEY = "datafusion.scoped_page_index.cache.size.limit"; + public static final String COLUMN_INDEX_CACHE_SIZE_LIMIT_KEY = "datafusion.column_index.cache.size.limit"; + public static final String OFFSET_INDEX_CACHE_SIZE_LIMIT_KEY = "datafusion.offset_index.cache.size.limit"; public static final Setting METADATA_CACHE_SIZE_LIMIT = new Setting<>( METADATA_CACHE_SIZE_LIMIT_KEY, "250mb", @@ -38,16 +39,30 @@ public class CacheSettings { ); /** - * Byte budget for the process-global scoped page-index cache (column index - * scoped to predicate columns, offset index for all columns). Pushed to - * native at startup via {@link CacheUtils#createCacheConfig}. Unlike the - * metadata/statistics caches this is a process-wide singleton, not owned by - * the cache manager. + * Byte budget for the process-global scoped ColumnIndex cache — the heavy, + * predicate-driven page index (per-page string min/max), keyed per + * {@code (file, col, rg)} cell. Pushed to native at startup via + * {@link CacheUtils#createCacheConfig}. Unlike the metadata/statistics caches + * this is a process-wide singleton, not owned by the cache manager. */ - public static final Setting SCOPED_PAGE_INDEX_CACHE_SIZE_LIMIT = new Setting<>( - SCOPED_PAGE_INDEX_CACHE_SIZE_LIMIT_KEY, + public static final Setting COLUMN_INDEX_CACHE_SIZE_LIMIT = new Setting<>( + COLUMN_INDEX_CACHE_SIZE_LIMIT_KEY, "64mb", - (s) -> ByteSizeValue.parseBytesSizeValue(s, new ByteSizeValue(0, ByteSizeUnit.KB), SCOPED_PAGE_INDEX_CACHE_SIZE_LIMIT_KEY), + (s) -> ByteSizeValue.parseBytesSizeValue(s, new ByteSizeValue(0, ByteSizeUnit.KB), COLUMN_INDEX_CACHE_SIZE_LIMIT_KEY), + Setting.Property.NodeScope, + Setting.Property.Dynamic + ); + + /** + * Byte budget for the process-global scoped OffsetIndex cache — the cheap, + * projection-driven page index (fixed-width page offsets), keyed per + * {@code (file, col)} cell. Far smaller per entry than the ColumnIndex, so it + * gets its own (smaller) default budget. + */ + public static final Setting OFFSET_INDEX_CACHE_SIZE_LIMIT = new Setting<>( + OFFSET_INDEX_CACHE_SIZE_LIMIT_KEY, + "16mb", + (s) -> ByteSizeValue.parseBytesSizeValue(s, new ByteSizeValue(0, ByteSizeUnit.KB), OFFSET_INDEX_CACHE_SIZE_LIMIT_KEY), Setting.Property.NodeScope, Setting.Property.Dynamic ); @@ -91,7 +106,8 @@ public class CacheSettings { STATISTICS_CACHE_ENABLED, STATISTICS_CACHE_SIZE_LIMIT, STATISTICS_CACHE_EVICTION_TYPE, - SCOPED_PAGE_INDEX_CACHE_SIZE_LIMIT + COLUMN_INDEX_CACHE_SIZE_LIMIT, + OFFSET_INDEX_CACHE_SIZE_LIMIT ); private static String validateEvictionType(String value) { diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/cache/CacheUtils.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/cache/CacheUtils.java index 2b55bb514728a..38d00a4f70288 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/cache/CacheUtils.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/cache/CacheUtils.java @@ -118,14 +118,21 @@ public static NativeCacheManagerHandle createCacheConfig(ClusterSettings cluster } } - // The scoped page-index cache is a process-global singleton (not owned by - // the cache manager), so its limit is pushed straight to native rather - // than created on the manager. Startup-only for now; a settings-update - // consumer can call NativeBridge.setScopedPageIndexCacheLimit to make it + // The scoped page-index caches are process-global singletons (not owned by + // the cache manager), so their limits are pushed straight to native rather + // than created on the manager. The heavy ColumnIndex and the cheap + // OffsetIndex are budgeted independently. Startup-only for now; a + // settings-update consumer can call the NativeBridge setters to make them // dynamic. - long scopedLimit = clusterSettings.get(CacheSettings.SCOPED_PAGE_INDEX_CACHE_SIZE_LIMIT).getBytes(); - logger.info("Configuring scoped page-index cache: size={} bytes", scopedLimit); - NativeBridge.setScopedPageIndexCacheLimit(scopedLimit); + long columnIndexLimit = clusterSettings.get(CacheSettings.COLUMN_INDEX_CACHE_SIZE_LIMIT).getBytes(); + long offsetIndexLimit = clusterSettings.get(CacheSettings.OFFSET_INDEX_CACHE_SIZE_LIMIT).getBytes(); + logger.info( + "Configuring scoped page-index caches: columnIndex={} bytes, offsetIndex={} bytes", + columnIndexLimit, + offsetIndexLimit + ); + NativeBridge.setColumnIndexCacheLimit(columnIndexLimit); + NativeBridge.setOffsetIndexCacheLimit(offsetIndexLimit); logger.info("Cache configuration completed"); return handle; diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/nativelib/NativeBridge.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/nativelib/NativeBridge.java index 195fa919f7bb2..d649fcef9c68f 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/nativelib/NativeBridge.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/nativelib/NativeBridge.java @@ -130,7 +130,8 @@ private static RuntimeException rethrowConverted(RuntimeException e) { private static final MethodHandle CACHE_MANAGER_GET_MEMORY_BY_TYPE; private static final MethodHandle CACHE_MANAGER_GET_TOTAL_MEMORY; private static final MethodHandle CACHE_MANAGER_CONTAINS_BY_TYPE; - private static final MethodHandle SET_SCOPED_PAGE_INDEX_CACHE_LIMIT; + private static final MethodHandle SET_COLUMN_INDEX_CACHE_LIMIT; + private static final MethodHandle SET_OFFSET_INDEX_CACHE_LIMIT; private static final MethodHandle CLEAR_SCOPED_PAGE_INDEX_CACHE; private static final MethodHandle CREATE_SESSION_CONTEXT; private static final MethodHandle CREATE_SESSION_CONTEXT_INDEXED; @@ -503,8 +504,13 @@ private static RuntimeException rethrowConverted(RuntimeException e) { ) ); - SET_SCOPED_PAGE_INDEX_CACHE_LIMIT = linker.downcallHandle( - lib.find("df_set_scoped_page_index_cache_limit").orElseThrow(), + SET_COLUMN_INDEX_CACHE_LIMIT = linker.downcallHandle( + lib.find("df_set_column_index_cache_limit").orElseThrow(), + FunctionDescriptor.of(ValueLayout.JAVA_LONG, ValueLayout.JAVA_LONG) + ); + + SET_OFFSET_INDEX_CACHE_LIMIT = linker.downcallHandle( + lib.find("df_set_offset_index_cache_limit").orElseThrow(), FunctionDescriptor.of(ValueLayout.JAVA_LONG, ValueLayout.JAVA_LONG) ); @@ -775,7 +781,7 @@ public static void setMemoryPoolLimit(long runtimePtr, long newLimitBytes) { } /** - * Sets the byte budget of the process-global scoped page-index cache. + * Sets the byte budget of the process-global scoped ColumnIndex cache. * *

This cache is a process-wide singleton (not owned by any cache manager), * so there is no runtime/manager pointer — the limit applies globally. @@ -783,9 +789,21 @@ public static void setMemoryPoolLimit(long runtimePtr, long newLimitBytes) { * * @param sizeLimitBytes the new byte budget (zero is ignored; keeps the current budget) */ - public static void setScopedPageIndexCacheLimit(long sizeLimitBytes) { + public static void setColumnIndexCacheLimit(long sizeLimitBytes) { + try (var call = new NativeCall()) { + call.invoke(SET_COLUMN_INDEX_CACHE_LIMIT, sizeLimitBytes); + } + } + + /** + * Sets the byte budget of the process-global scoped OffsetIndex cache. See + * {@link #setColumnIndexCacheLimit(long)} for the singleton/eviction semantics. + * + * @param sizeLimitBytes the new byte budget (zero is ignored; keeps the current budget) + */ + public static void setOffsetIndexCacheLimit(long sizeLimitBytes) { try (var call = new NativeCall()) { - call.invoke(SET_SCOPED_PAGE_INDEX_CACHE_LIMIT, sizeLimitBytes); + call.invoke(SET_OFFSET_INDEX_CACHE_LIMIT, sizeLimitBytes); } } diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/nativelib/StatsLayout.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/nativelib/StatsLayout.java index cd1821fc224f4..9218cf6a8df11 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/nativelib/StatsLayout.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/nativelib/StatsLayout.java @@ -28,9 +28,9 @@ * and provides {@link VarHandle} accessors for each field via layout path navigation. * *

The layout contains 10 named groups (2 runtime × 9 fields + 4 task monitor × 5 fields - * + 1 partition gate × 8 fields + 1 adaptive budget × 2 fields + 1 cache stats × 15 fields - * + 1 search stats × 17 fields = 80 longs = 640 bytes). The cache stats group holds three - * sub-caches × 5 fields each: metadata, statistics, and scoped page-index. + * + 1 partition gate × 8 fields + 1 adaptive budget × 2 fields + 1 cache stats × 20 fields + * + 1 search stats × 17 fields = 85 longs = 680 bytes). The cache stats group holds four + * sub-caches × 5 fields each: metadata, statistics, scoped column-index, and scoped offset-index. */ public final class StatsLayout { @@ -100,8 +100,8 @@ public final class StatsLayout { ); static { - if (LAYOUT.byteSize() != 80 * Long.BYTES) { - throw new AssertionError("StatsLayout size mismatch: expected " + (80 * Long.BYTES) + " but got " + LAYOUT.byteSize()); + if (LAYOUT.byteSize() != 85 * Long.BYTES) { + throw new AssertionError("StatsLayout size mismatch: expected " + (85 * Long.BYTES) + " but got " + LAYOUT.byteSize()); } } @@ -183,12 +183,19 @@ public final class StatsLayout { private static final VarHandle CACHE_STATS_MEMORY_BYTES = cacheHandle("statistics_cache", "memory_bytes"); private static final VarHandle CACHE_STATS_SIZE_LIMIT_BYTES = cacheHandle("statistics_cache", "size_limit_bytes"); - // ---- VarHandles for cache_stats.scoped_page_index_cache fields ---- - private static final VarHandle CACHE_SPI_HIT_COUNT = cacheHandle("scoped_page_index_cache", "hit_count"); - private static final VarHandle CACHE_SPI_MISS_COUNT = cacheHandle("scoped_page_index_cache", "miss_count"); - private static final VarHandle CACHE_SPI_ENTRY_COUNT = cacheHandle("scoped_page_index_cache", "entry_count"); - private static final VarHandle CACHE_SPI_MEMORY_BYTES = cacheHandle("scoped_page_index_cache", "memory_bytes"); - private static final VarHandle CACHE_SPI_SIZE_LIMIT_BYTES = cacheHandle("scoped_page_index_cache", "size_limit_bytes"); + // ---- VarHandles for cache_stats.column_index_cache fields ---- + private static final VarHandle CACHE_CI_HIT_COUNT = cacheHandle("column_index_cache", "hit_count"); + private static final VarHandle CACHE_CI_MISS_COUNT = cacheHandle("column_index_cache", "miss_count"); + private static final VarHandle CACHE_CI_ENTRY_COUNT = cacheHandle("column_index_cache", "entry_count"); + private static final VarHandle CACHE_CI_MEMORY_BYTES = cacheHandle("column_index_cache", "memory_bytes"); + private static final VarHandle CACHE_CI_SIZE_LIMIT_BYTES = cacheHandle("column_index_cache", "size_limit_bytes"); + + // ---- VarHandles for cache_stats.offset_index_cache fields ---- + private static final VarHandle CACHE_OI_HIT_COUNT = cacheHandle("offset_index_cache", "hit_count"); + private static final VarHandle CACHE_OI_MISS_COUNT = cacheHandle("offset_index_cache", "miss_count"); + private static final VarHandle CACHE_OI_ENTRY_COUNT = cacheHandle("offset_index_cache", "entry_count"); + private static final VarHandle CACHE_OI_MEMORY_BYTES = cacheHandle("offset_index_cache", "memory_bytes"); + private static final VarHandle CACHE_OI_SIZE_LIMIT_BYTES = cacheHandle("offset_index_cache", "size_limit_bytes"); // ---- VarHandles for search_stats fields ---- private static final VarHandle SS_LISTING_TABLE_SCAN = handle("search_stats", "listing_table_scan"); @@ -318,14 +325,21 @@ public static CacheStats readCacheStats(MemorySegment seg) { (long) CACHE_STATS_MEMORY_BYTES.get(seg, 0L), (long) CACHE_STATS_SIZE_LIMIT_BYTES.get(seg, 0L) ); - CacheGroupStats scopedPageIndex = new CacheGroupStats( - (long) CACHE_SPI_HIT_COUNT.get(seg, 0L), - (long) CACHE_SPI_MISS_COUNT.get(seg, 0L), - (long) CACHE_SPI_ENTRY_COUNT.get(seg, 0L), - (long) CACHE_SPI_MEMORY_BYTES.get(seg, 0L), - (long) CACHE_SPI_SIZE_LIMIT_BYTES.get(seg, 0L) + CacheGroupStats columnIndex = new CacheGroupStats( + (long) CACHE_CI_HIT_COUNT.get(seg, 0L), + (long) CACHE_CI_MISS_COUNT.get(seg, 0L), + (long) CACHE_CI_ENTRY_COUNT.get(seg, 0L), + (long) CACHE_CI_MEMORY_BYTES.get(seg, 0L), + (long) CACHE_CI_SIZE_LIMIT_BYTES.get(seg, 0L) + ); + CacheGroupStats offsetIndex = new CacheGroupStats( + (long) CACHE_OI_HIT_COUNT.get(seg, 0L), + (long) CACHE_OI_MISS_COUNT.get(seg, 0L), + (long) CACHE_OI_ENTRY_COUNT.get(seg, 0L), + (long) CACHE_OI_MEMORY_BYTES.get(seg, 0L), + (long) CACHE_OI_SIZE_LIMIT_BYTES.get(seg, 0L) ); - return new CacheStats(metadata, statistics, scopedPageIndex); + return new CacheStats(metadata, statistics, columnIndex, offsetIndex); } /** @@ -419,7 +433,8 @@ private static StructLayout cacheStatsGroup(String name) { return MemoryLayout.structLayout( cacheGroup("metadata_cache"), cacheGroup("statistics_cache"), - cacheGroup("scoped_page_index_cache") + cacheGroup("column_index_cache"), + cacheGroup("offset_index_cache") ).withName(name); } diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/stats/CacheStats.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/stats/CacheStats.java index f62ce980b21bd..553df644abe1b 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/stats/CacheStats.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/stats/CacheStats.java @@ -18,8 +18,11 @@ import java.util.Objects; /** - * Stats for the two parquet caches owned by {@code CustomCacheManager}: - * the parquet metadata (footer) cache and the column-statistics cache. + * Stats for the parquet caches: the {@code CustomCacheManager}-owned metadata + * (footer) cache and column-statistics cache, plus the two process-global scoped + * page-index caches — the predicate-driven {@code column_index_cache} (keyed per + * {@code (file, col, rg)} cell) and the projection-driven {@code offset_index_cache} + * (keyed per {@code (file, col)} cell). * *

Each sub-group reports five fields. Disabled caches surface as all-zero * groups (in particular {@code size_limit_bytes == 0}); the JSON shape stays @@ -29,19 +32,27 @@ public class CacheStats implements Writeable, ToXContentFragment { private final CacheGroupStats metadataCache; private final CacheGroupStats statisticsCache; - private final CacheGroupStats scopedPageIndexCache; + private final CacheGroupStats columnIndexCache; + private final CacheGroupStats offsetIndexCache; /** * Construct from individual sub-group stats. * - * @param metadataCache metadata cache counters (must not be null) - * @param statisticsCache statistics cache counters (must not be null) - * @param scopedPageIndexCache scoped page-index cache counters (must not be null) + * @param metadataCache metadata cache counters (must not be null) + * @param statisticsCache statistics cache counters (must not be null) + * @param columnIndexCache scoped ColumnIndex cache counters (must not be null) + * @param offsetIndexCache scoped OffsetIndex cache counters (must not be null) */ - public CacheStats(CacheGroupStats metadataCache, CacheGroupStats statisticsCache, CacheGroupStats scopedPageIndexCache) { + public CacheStats( + CacheGroupStats metadataCache, + CacheGroupStats statisticsCache, + CacheGroupStats columnIndexCache, + CacheGroupStats offsetIndexCache + ) { this.metadataCache = Objects.requireNonNull(metadataCache); this.statisticsCache = Objects.requireNonNull(statisticsCache); - this.scopedPageIndexCache = Objects.requireNonNull(scopedPageIndexCache); + this.columnIndexCache = Objects.requireNonNull(columnIndexCache); + this.offsetIndexCache = Objects.requireNonNull(offsetIndexCache); } /** @@ -53,14 +64,16 @@ public CacheStats(CacheGroupStats metadataCache, CacheGroupStats statisticsCache public CacheStats(StreamInput in) throws IOException { this.metadataCache = new CacheGroupStats(in); this.statisticsCache = new CacheGroupStats(in); - this.scopedPageIndexCache = new CacheGroupStats(in); + this.columnIndexCache = new CacheGroupStats(in); + this.offsetIndexCache = new CacheGroupStats(in); } @Override public void writeTo(StreamOutput out) throws IOException { metadataCache.writeTo(out); statisticsCache.writeTo(out); - scopedPageIndexCache.writeTo(out); + columnIndexCache.writeTo(out); + offsetIndexCache.writeTo(out); } @Override @@ -72,8 +85,11 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws builder.startObject("statistics_cache"); statisticsCache.toXContent(builder); builder.endObject(); - builder.startObject("scoped_page_index_cache"); - scopedPageIndexCache.toXContent(builder); + builder.startObject("column_index_cache"); + columnIndexCache.toXContent(builder); + builder.endObject(); + builder.startObject("offset_index_cache"); + offsetIndexCache.toXContent(builder); builder.endObject(); builder.endObject(); return builder; @@ -89,9 +105,14 @@ public CacheGroupStats getStatisticsCache() { return statisticsCache; } - /** Returns the scoped page-index cache counters. */ - public CacheGroupStats getScopedPageIndexCache() { - return scopedPageIndexCache; + /** Returns the scoped ColumnIndex cache counters. */ + public CacheGroupStats getColumnIndexCache() { + return columnIndexCache; + } + + /** Returns the scoped OffsetIndex cache counters. */ + public CacheGroupStats getOffsetIndexCache() { + return offsetIndexCache; } @Override @@ -101,11 +122,12 @@ public boolean equals(Object o) { CacheStats that = (CacheStats) o; return Objects.equals(metadataCache, that.metadataCache) && Objects.equals(statisticsCache, that.statisticsCache) - && Objects.equals(scopedPageIndexCache, that.scopedPageIndexCache); + && Objects.equals(columnIndexCache, that.columnIndexCache) + && Objects.equals(offsetIndexCache, that.offsetIndexCache); } @Override public int hashCode() { - return Objects.hash(metadataCache, statisticsCache, scopedPageIndexCache); + return Objects.hash(metadataCache, statisticsCache, columnIndexCache, offsetIndexCache); } } diff --git a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DataFusionPluginSettingsTests.java b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DataFusionPluginSettingsTests.java index 43b89b3ccee86..63825d606e4c5 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DataFusionPluginSettingsTests.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DataFusionPluginSettingsTests.java @@ -115,7 +115,9 @@ public void testGetSettingsReturnsAllIndexedSettings() { public void testGetSettingsReturnsTotalExpectedCount() { try (DataFusionPlugin plugin = new DataFusionPlugin()) { List> settings = plugin.getSettings(); - assertEquals(28, settings.size()); + // Split of the single scoped page-index limit into two (column_index + + // offset_index) added one net setting (29 → 30). + assertEquals(30, settings.size()); } catch (Exception e) { throw new AssertionError(e); } diff --git a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DataFusionServiceTests.java b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DataFusionServiceTests.java index 002e8e2e1e976..6337aa587c4a2 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DataFusionServiceTests.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DataFusionServiceTests.java @@ -228,6 +228,8 @@ private ClusterSettings createCacheClusterSettings(Settings settings) { all.add(CacheSettings.STATISTICS_CACHE_ENABLED); all.add(CacheSettings.STATISTICS_CACHE_SIZE_LIMIT); all.add(CacheSettings.STATISTICS_CACHE_EVICTION_TYPE); + all.add(CacheSettings.COLUMN_INDEX_CACHE_SIZE_LIMIT); + all.add(CacheSettings.OFFSET_INDEX_CACHE_SIZE_LIMIT); all.add(DataFusionPlugin.DATAFUSION_MEMORY_POOL_LIMIT); all.add(DataFusionPlugin.DATAFUSION_SPILL_MEMORY_LIMIT); return new ClusterSettings(settings, all); diff --git a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionCacheManagerTests.java b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionCacheManagerTests.java index f09497c72564c..2563ce8997b7a 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionCacheManagerTests.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionCacheManagerTests.java @@ -41,6 +41,8 @@ private void setup() { clusterSettingsToAdd.add(CacheSettings.STATISTICS_CACHE_ENABLED); clusterSettingsToAdd.add(CacheSettings.STATISTICS_CACHE_SIZE_LIMIT); clusterSettingsToAdd.add(CacheSettings.STATISTICS_CACHE_EVICTION_TYPE); + clusterSettingsToAdd.add(CacheSettings.COLUMN_INDEX_CACHE_SIZE_LIMIT); + clusterSettingsToAdd.add(CacheSettings.OFFSET_INDEX_CACHE_SIZE_LIMIT); clusterSettingsToAdd.add(DataFusionPlugin.DATAFUSION_MEMORY_POOL_LIMIT); clusterSettingsToAdd.add(DataFusionPlugin.DATAFUSION_SPILL_MEMORY_LIMIT); diff --git a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionSettingsTests.java b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionSettingsTests.java index 24c7303192f2c..e6c73ee335fec 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionSettingsTests.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionSettingsTests.java @@ -69,7 +69,8 @@ public void testMaxCollectorParallelismSettingDefinition() { } public void testAllSettingsContainsAllExpectedSettings() { - assertEquals(28, DatafusionSettings.ALL_SETTINGS.size()); + // Split of the scoped page-index limit into column_index + offset_index added one (29 → 30). + assertEquals(30, DatafusionSettings.ALL_SETTINGS.size()); assertTrue(DatafusionSettings.ALL_SETTINGS.contains(DataFusionPlugin.DATAFUSION_REDUCE_TARGET_PARTITIONS)); assertTrue(DatafusionSettings.ALL_SETTINGS.contains(DataFusionPlugin.DATAFUSION_SPILL_DIRECTORY)); assertTrue(DatafusionSettings.ALL_SETTINGS.contains(DatafusionSettings.INDEXED_BATCH_SIZE)); diff --git a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/nativelib/StatsLayoutPropertyTests.java b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/nativelib/StatsLayoutPropertyTests.java index bb1c0afea1229..3f9a23e71030a 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/nativelib/StatsLayoutPropertyTests.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/nativelib/StatsLayoutPropertyTests.java @@ -35,7 +35,7 @@ public class StatsLayoutPropertyTests extends OpenSearchTestCase { private static final int TRIES = 100; - private static final int FIELD_COUNT = 80; + private static final int FIELD_COUNT = 85; // ---- Generators ---- @@ -175,7 +175,7 @@ public void testPackThenDecodeRoundTripWithCpu() { assertEquals(values[46], bs.fallbacks); assertEquals(values[47], bs.rejections); - // Cache stats (offsets 48-62; 3 sub-caches × 5) + // Cache stats (offsets 48-67; 4 sub-caches × 5) var cs = StatsLayout.readCacheStats(seg); assertEquals(values[48], cs.getMetadataCache().hitCount); assertEquals(values[49], cs.getMetadataCache().missCount); @@ -187,31 +187,36 @@ public void testPackThenDecodeRoundTripWithCpu() { assertEquals(values[55], cs.getStatisticsCache().entryCount); assertEquals(values[56], cs.getStatisticsCache().memoryBytes); assertEquals(values[57], cs.getStatisticsCache().sizeLimitBytes); - assertEquals(values[58], cs.getScopedPageIndexCache().hitCount); - assertEquals(values[59], cs.getScopedPageIndexCache().missCount); - assertEquals(values[60], cs.getScopedPageIndexCache().entryCount); - assertEquals(values[61], cs.getScopedPageIndexCache().memoryBytes); - assertEquals(values[62], cs.getScopedPageIndexCache().sizeLimitBytes); - - // Search stats (offsets 63-79) + assertEquals(values[58], cs.getColumnIndexCache().hitCount); + assertEquals(values[59], cs.getColumnIndexCache().missCount); + assertEquals(values[60], cs.getColumnIndexCache().entryCount); + assertEquals(values[61], cs.getColumnIndexCache().memoryBytes); + assertEquals(values[62], cs.getColumnIndexCache().sizeLimitBytes); + assertEquals(values[63], cs.getOffsetIndexCache().hitCount); + assertEquals(values[64], cs.getOffsetIndexCache().missCount); + assertEquals(values[65], cs.getOffsetIndexCache().entryCount); + assertEquals(values[66], cs.getOffsetIndexCache().memoryBytes); + assertEquals(values[67], cs.getOffsetIndexCache().sizeLimitBytes); + + // Search stats (offsets 68-84) var ss = StatsLayout.readSearchStats(seg); - assertEquals(values[63], ss.listingTableScan); - assertEquals(values[64], ss.singleCollectorScan); - assertEquals(values[65], ss.bitmapTreeScan); - assertEquals(values[66], ss.delegationCalls); - assertEquals(values[67], ss.rgProcessed); - assertEquals(values[68], ss.rgSkipped); - assertEquals(values[69], ss.parquetScanTotalTimeMs); - assertEquals(values[70], ss.parquetScanUntilDataTimeMs); - assertEquals(values[71], ss.parquetProcessingTimeMs); - assertEquals(values[72], ss.parquetBytesScanned); - assertEquals(values[73], ss.prefetchWaitTimeMs); - assertEquals(values[74], ss.prefetchWaitCount); - assertEquals(values[75], ss.elapsedComputeMs); - assertEquals(values[76], ss.buildMaskTimeMs); - assertEquals(values[77], ss.onBatchMaskTimeMs); - assertEquals(values[78], ss.filterRecordBatchTimeMs); - assertEquals(values[79], ss.objectStoreReadTimeMs); + assertEquals(values[68], ss.listingTableScan); + assertEquals(values[69], ss.singleCollectorScan); + assertEquals(values[70], ss.bitmapTreeScan); + assertEquals(values[71], ss.delegationCalls); + assertEquals(values[72], ss.rgProcessed); + assertEquals(values[73], ss.rgSkipped); + assertEquals(values[74], ss.parquetScanTotalTimeMs); + assertEquals(values[75], ss.parquetScanUntilDataTimeMs); + assertEquals(values[76], ss.parquetProcessingTimeMs); + assertEquals(values[77], ss.parquetBytesScanned); + assertEquals(values[78], ss.prefetchWaitTimeMs); + assertEquals(values[79], ss.prefetchWaitCount); + assertEquals(values[80], ss.elapsedComputeMs); + assertEquals(values[81], ss.buildMaskTimeMs); + assertEquals(values[82], ss.onBatchMaskTimeMs); + assertEquals(values[83], ss.filterRecordBatchTimeMs); + assertEquals(values[84], ss.objectStoreReadTimeMs); } } } @@ -333,12 +338,18 @@ public void testDecodeThenReencodeIdentity() { cs.getStatisticsCache().entryCount, cs.getStatisticsCache().memoryBytes, cs.getStatisticsCache().sizeLimitBytes, - // cache_stats.scoped_page_index_cache (5) - cs.getScopedPageIndexCache().hitCount, - cs.getScopedPageIndexCache().missCount, - cs.getScopedPageIndexCache().entryCount, - cs.getScopedPageIndexCache().memoryBytes, - cs.getScopedPageIndexCache().sizeLimitBytes, + // cache_stats.column_index_cache (5) + cs.getColumnIndexCache().hitCount, + cs.getColumnIndexCache().missCount, + cs.getColumnIndexCache().entryCount, + cs.getColumnIndexCache().memoryBytes, + cs.getColumnIndexCache().sizeLimitBytes, + // cache_stats.offset_index_cache (5) + cs.getOffsetIndexCache().hitCount, + cs.getOffsetIndexCache().missCount, + cs.getOffsetIndexCache().entryCount, + cs.getOffsetIndexCache().memoryBytes, + cs.getOffsetIndexCache().sizeLimitBytes, // search_stats (17) ss.listingTableScan, ss.singleCollectorScan, diff --git a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/nativelib/StatsLayoutTests.java b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/nativelib/StatsLayoutTests.java index b938e1fd17f43..468e235538662 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/nativelib/StatsLayoutTests.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/nativelib/StatsLayoutTests.java @@ -22,10 +22,10 @@ */ public class StatsLayoutTests extends OpenSearchTestCase { - /** 7.1: Layout byte size must be 640 (80 × 8). */ + /** 7.1: Layout byte size must be 680 (85 × 8). */ public void testLayoutByteSize() { - assertEquals(640L, StatsLayout.LAYOUT.byteSize()); - assertEquals(80 * Long.BYTES, (int) StatsLayout.LAYOUT.byteSize()); + assertEquals(680L, StatsLayout.LAYOUT.byteSize()); + assertEquals(85 * Long.BYTES, (int) StatsLayout.LAYOUT.byteSize()); } /** 7.2: readRuntimeMetrics decodes 9 known values from io_runtime group. */ diff --git a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/stats/CacheStatsTests.java b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/stats/CacheStatsTests.java index 3b51f80648697..c430a6ad0ae4f 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/stats/CacheStatsTests.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/stats/CacheStatsTests.java @@ -91,27 +91,31 @@ public void testCacheGroupStatsEqualsAndHashCode() { public void testCacheStatsConstructorRejectsNullSubGroups() { CacheGroupStats g = new CacheGroupStats(0, 0, 0, 0, 0); - expectThrows(NullPointerException.class, () -> new CacheStats(null, g, g)); - expectThrows(NullPointerException.class, () -> new CacheStats(g, null, g)); - expectThrows(NullPointerException.class, () -> new CacheStats(g, g, null)); + expectThrows(NullPointerException.class, () -> new CacheStats(null, g, g, g)); + expectThrows(NullPointerException.class, () -> new CacheStats(g, null, g, g)); + expectThrows(NullPointerException.class, () -> new CacheStats(g, g, null, g)); + expectThrows(NullPointerException.class, () -> new CacheStats(g, g, g, null)); } public void testCacheStatsAccessors() { CacheGroupStats meta = new CacheGroupStats(1, 2, 3, 4, 5); CacheGroupStats stats = new CacheGroupStats(6, 7, 8, 9, 10); - CacheGroupStats scoped = new CacheGroupStats(11, 12, 13, 14, 15); - CacheStats c = new CacheStats(meta, stats, scoped); + CacheGroupStats columnIndex = new CacheGroupStats(11, 12, 13, 14, 15); + CacheGroupStats offsetIndex = new CacheGroupStats(16, 17, 18, 19, 20); + CacheStats c = new CacheStats(meta, stats, columnIndex, offsetIndex); assertSame(meta, c.getMetadataCache()); assertSame(stats, c.getStatisticsCache()); - assertSame(scoped, c.getScopedPageIndexCache()); + assertSame(columnIndex, c.getColumnIndexCache()); + assertSame(offsetIndex, c.getOffsetIndexCache()); } public void testCacheStatsWriteableRoundTrip() throws IOException { CacheStats original = new CacheStats( new CacheGroupStats(11, 12, 13, 14, 15), new CacheGroupStats(21, 22, 23, 24, 25), - new CacheGroupStats(31, 32, 33, 34, 35) + new CacheGroupStats(31, 32, 33, 34, 35), + new CacheGroupStats(41, 42, 43, 44, 45) ); BytesStreamOutput out = new BytesStreamOutput(); original.writeTo(out); @@ -124,7 +128,8 @@ public void testCacheStatsToXContentShape() throws IOException { CacheStats c = new CacheStats( new CacheGroupStats(11, 0, 3, 1024, 250_000_000), new CacheGroupStats(0, 7, 0, 0, 100_000_000), - new CacheGroupStats(5, 1, 2, 512, 64_000_000) + new CacheGroupStats(5, 1, 2, 512, 64_000_000), + new CacheGroupStats(4, 1, 2, 256, 16_000_000) ); XContentBuilder builder = XContentFactory.jsonBuilder(); builder.startObject(); @@ -134,10 +139,11 @@ public void testCacheStatsToXContentShape() throws IOException { // Top-level wrapper assertTrue("expected cache_stats wrapper, got: " + json, json.contains("\"cache_stats\"")); - // All three sub-groups present + // All four sub-groups present assertTrue(json.contains("\"metadata_cache\"")); assertTrue(json.contains("\"statistics_cache\"")); - assertTrue("expected scoped_page_index_cache group, got: " + json, json.contains("\"scoped_page_index_cache\"")); + assertTrue("expected column_index_cache group, got: " + json, json.contains("\"column_index_cache\"")); + assertTrue("expected offset_index_cache group, got: " + json, json.contains("\"offset_index_cache\"")); // Per-group fields assertTrue(json.contains("\"hit_count\":11")); assertTrue(json.contains("\"miss_count\":7")); @@ -146,10 +152,12 @@ public void testCacheStatsToXContentShape() throws IOException { assertTrue(json.contains("\"size_limit_bytes\":250000000")); assertTrue(json.contains("\"size_limit_bytes\":100000000")); assertTrue(json.contains("\"size_limit_bytes\":64000000")); + assertTrue(json.contains("\"size_limit_bytes\":16000000")); } public void testCacheStatsZeroedRendersAllZeros() throws IOException { CacheStats zero = new CacheStats( + new CacheGroupStats(0, 0, 0, 0, 0), new CacheGroupStats(0, 0, 0, 0, 0), new CacheGroupStats(0, 0, 0, 0, 0), new CacheGroupStats(0, 0, 0, 0, 0) @@ -160,9 +168,9 @@ public void testCacheStatsZeroedRendersAllZeros() throws IOException { builder.endObject(); String json = builder.toString(); - // Disabled-cache sentinel: all three size_limit_bytes are 0 + // Disabled-cache sentinel: all four size_limit_bytes are 0 long sizeLimitOccurrences = json.split("\"size_limit_bytes\":0", -1).length - 1; - assertEquals("expected size_limit_bytes:0 to appear three times (one per sub-cache)", 3, sizeLimitOccurrences); + assertEquals("expected size_limit_bytes:0 to appear four times (one per sub-cache)", 4, sizeLimitOccurrences); // hit_rate must not be NaN assertFalse("hit_rate must not be NaN: " + json, json.contains("NaN")); } @@ -171,18 +179,21 @@ public void testCacheStatsEqualsAndHashCode() { CacheStats a = new CacheStats( new CacheGroupStats(1, 2, 3, 4, 5), new CacheGroupStats(6, 7, 8, 9, 10), - new CacheGroupStats(11, 12, 13, 14, 15) + new CacheGroupStats(11, 12, 13, 14, 15), + new CacheGroupStats(16, 17, 18, 19, 20) ); CacheStats b = new CacheStats( new CacheGroupStats(1, 2, 3, 4, 5), new CacheGroupStats(6, 7, 8, 9, 10), - new CacheGroupStats(11, 12, 13, 14, 15) + new CacheGroupStats(11, 12, 13, 14, 15), + new CacheGroupStats(16, 17, 18, 19, 20) ); - // Differs only in the scoped group → must be unequal. + // Differs only in the offset-index group → must be unequal. CacheStats c = new CacheStats( new CacheGroupStats(1, 2, 3, 4, 5), new CacheGroupStats(6, 7, 8, 9, 10), - new CacheGroupStats(11, 12, 13, 14, 99) + new CacheGroupStats(11, 12, 13, 14, 15), + new CacheGroupStats(16, 17, 18, 19, 99) ); assertEquals(a, b); diff --git a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/stats/DataFusionStatsPropertyTests.java b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/stats/DataFusionStatsPropertyTests.java index 9943a0d993780..7251185f198d2 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/stats/DataFusionStatsPropertyTests.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/stats/DataFusionStatsPropertyTests.java @@ -108,7 +108,12 @@ private CacheGroupStats randomCacheGroupStats() { } private CacheStats randomCacheStats() { - return new CacheStats(randomCacheGroupStats(), randomCacheGroupStats(), randomCacheGroupStats()); + return new CacheStats( + randomCacheGroupStats(), + randomCacheGroupStats(), + randomCacheGroupStats(), + randomCacheGroupStats() + ); } private SearchStats randomSearchStats() { diff --git a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/stats/DataFusionStatsTests.java b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/stats/DataFusionStatsTests.java index a7c0d93104dc6..a22299278c455 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/stats/DataFusionStatsTests.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/stats/DataFusionStatsTests.java @@ -279,7 +279,8 @@ public void testCacheStatsPresentRendersIntoJson() throws IOException { CacheStats cache = new CacheStats( new CacheGroupStats(100, 5, 25, 4096, 250_000_000), new CacheGroupStats(50, 0, 25, 2048, 100_000_000), - new CacheGroupStats(20, 3, 8, 1024, 64_000_000) + new CacheGroupStats(20, 3, 8, 1024, 64_000_000), + new CacheGroupStats(12, 2, 6, 512, 16_000_000) ); DataFusionStats stats = new DataFusionStats( new NativeExecutorsStats(io, null, taskMonitors), diff --git a/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/ScopedPageIndexCacheIT.java b/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/ScopedPageIndexCacheIT.java index 416547da41566..60f2bc11aac75 100644 --- a/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/ScopedPageIndexCacheIT.java +++ b/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/ScopedPageIndexCacheIT.java @@ -17,27 +17,35 @@ import java.util.Map; /** - * End-to-end integration test for the unified scoped parquet page-index cache and - * the footer-only level-1 metadata cache, verified through the node-stats API + * End-to-end integration test for the cell-keyed scoped parquet page-index caches + * and the footer-only level-1 metadata cache, verified through the node-stats API * ({@code GET /_plugins/_analytics_backend_datafusion/stats}) — the - * {@code cache_stats.metadata_cache} and {@code cache_stats.scoped_page_index_cache} - * groups. + * {@code cache_stats.metadata_cache}, {@code cache_stats.column_index_cache}, and + * {@code cache_stats.offset_index_cache} groups. * - *

This suite is deliberately broad: it checks query correctness (not - * just cache counters), that the level-1 metadata cache still works after - * the page-index strip, the full scoped-cache hit/miss/size story, and that - * a spread of query shapes (aggregations, multi-column filters, full-text - * {@code match}) all still execute. The goal is to prove the page-index changes - * broke nothing. + *

The two scoped caches

* - *

Why assertions are same-method, twice-run deltas

+ * The scoped page index is split into two process-global caches, each keyed at + * cell granularity so an index is decoded and stored once per file and + * reused across query shapes: + *
    + *
  • {@code column_index_cache} — the heavy, predicate-driven ColumnIndex + * (per-page string min/max), keyed per {@code (file, col, rg)} cell. Adding a + * column to a predicate, or changing a literal, never re-decodes a cell that + * is already cached; only genuinely new {@code (col, rg)} cells are read.
    • + *
  • {@code offset_index_cache} — the cheap, projection-driven OffsetIndex + * (fixed-width page offsets), keyed per {@code (file, col)} cell (the value + * spans all row groups). Different projections reuse shared column cells.
    • + *
* - * The scoped cache is a process-global singleton that persists for the life - * of the node; the cluster is preserved across methods which run in randomized - * order; there is no reset endpoint; and provisioning triggers a refresh that - * warms the (footer-only) metadata cache. So a method must NOT assume a cold - * cache. The robust signal is measured within one method: run a query, snapshot, - * run the SAME query again, snapshot — the second run must be a pure hit. + *

Determinism: the clear endpoint + same-method deltas

+ * + * Most assertions clear the scoped caches first via + * {@code POST /_plugins/_analytics_backend_datafusion/cache/scoped_page_index/_clear} + * so a method starts from a known-empty state, then measure deltas across queries + * within the one method. (The caches are process-global singletons that persist for + * the life of the node and the cluster is preserved across randomly-ordered methods, + * so a method must never assume a globally cold cache — it clears explicitly.) * *

Run (fast): {@code ./gradlew :sandbox:qa:analytics-engine-rest:integTest * --tests "*ScopedPageIndexCacheIT" -Dsandbox.enabled=true -PrustDebug}. @@ -92,14 +100,26 @@ private CacheGroup fetchGroup(String groupName) throws IOException { return new CacheGroup(hits, misses, entries, memory, limit); } - private CacheGroup scoped() throws IOException { - return fetchGroup("scoped_page_index_cache"); + /** Scoped ColumnIndex cache (predicate-driven, per {@code (file, col, rg)} cell). */ + private CacheGroup columnIndex() throws IOException { + return fetchGroup("column_index_cache"); + } + + /** Scoped OffsetIndex cache (projection-driven, per {@code (file, col)} cell). */ + private CacheGroup offsetIndex() throws IOException { + return fetchGroup("offset_index_cache"); } private CacheGroup metadata() throws IOException { return fetchGroup("metadata_cache"); } + /** Drop all entries + reset counters in BOTH scoped caches (testing convenience). */ + private void clearScopedCaches() throws IOException { + Request request = new Request("POST", "/_plugins/_analytics_backend_datafusion/cache/scoped_page_index/_clear"); + assertOkAndParse(client().performRequest(request), "clear scoped page-index cache"); + } + private static long num(Map obj, String key) { Object v = obj.get(key); assertNotNull(key + " missing", v); @@ -127,52 +147,53 @@ private long scalarAgg(String ppl) throws IOException { return ((Number) first.get(0)).longValue(); } + private String src() { + return "source=" + DATASET.indexName; + } + // ---- exposure + correctness ----------------------------------------- /** - * The scoped page-index cache group must always be present with a positive - * byte budget (its configured 64mb limit), even before any query runs. + * Both scoped cache groups must always be present with a positive byte budget + * (their configured limits), even before any query runs. */ - public void testScopedCacheGroupIsExposedWithBudget() throws IOException { - CacheGroup snap = scoped(); - assertTrue( - "scoped_page_index_cache must advertise a positive size_limit_bytes, got " + snap.sizeLimitBytes(), - snap.sizeLimitBytes() > 0 - ); - assertTrue("hit_count >= 0", snap.hits() >= 0); - assertTrue("miss_count >= 0", snap.misses() >= 0); - assertTrue("entry_count >= 0", snap.entries() >= 0); + public void testScopedCacheGroupsAreExposedWithBudgets() throws IOException { + for (CacheGroup snap : new CacheGroup[] { columnIndex(), offsetIndex() }) { + assertTrue( + "scoped cache must advertise a positive size_limit_bytes, got " + snap.sizeLimitBytes(), + snap.sizeLimitBytes() > 0 + ); + assertTrue("hit_count >= 0", snap.hits() >= 0); + assertTrue("miss_count >= 0", snap.misses() >= 0); + assertTrue("entry_count >= 0", snap.entries() >= 0); + } } /** - * The page-index changes must not change query answers. Exact-count - * assertions over the listing path (numeric + keyword predicates) against - * known dataset cardinalities. + * The page-index changes must not change query answers. Exact-count assertions + * over the listing path (numeric + keyword predicates) against known dataset + * cardinalities. */ public void testListingQueryCorrectnessUnchanged() throws IOException { - assertEquals( - "total doc count must be exact", - TOTAL_DOCS, - scalarAgg("source=" + DATASET.indexName + " | stats count()") - ); + assertEquals("total doc count must be exact", TOTAL_DOCS, scalarAgg(src() + " | stats count()")); assertEquals( "status >= 400 count must be exact (numeric listing predicate)", STATUS_GE_400, - scalarAgg("source=" + DATASET.indexName + " | where status >= 400 | stats count()") + scalarAgg(src() + " | where status >= 400 | stats count()") ); assertEquals( "log_level = 'ERROR' count must be exact (keyword predicate)", LEVEL_ERROR, - scalarAgg("source=" + DATASET.indexName + " | where log_level = 'ERROR' | stats count()") + scalarAgg(src() + " | where log_level = 'ERROR' | stats count()") ); } /** * The same correctness must hold when the SAME query is re-run (served partly - * from the scoped cache) — a cached page index must never change the answer. + * from the scoped caches) — a cached page index must never change the answer. */ public void testCorrectnessIsStableAcrossCachedReRuns() throws IOException { - String q = "source=" + DATASET.indexName + " | where status >= 400 | stats count()"; + String q = src() + " | where status >= 400 | stats count()"; long first = scalarAgg(q); long second = scalarAgg(q); assertEquals("cold and warm runs must agree", first, second); @@ -184,186 +205,349 @@ public void testCorrectnessIsStableAcrossCachedReRuns() throws IOException { /** * The footer-only level-1 metadata cache must still function: after repeated * queries it holds entries and registers hits (footers are reused, not - * re-read). This guards against the page-index strip accidentally breaking - * normal footer caching. + * re-read). Guards against the page-index strip breaking normal footer caching. */ public void testMetadataCacheStillServesFooters() throws IOException { - String q = "source=" + DATASET.indexName + " | where status >= 200 | stats count() by service_name"; - // Warm. - executePpl(q); + String q = src() + " | where status >= 200 | stats count() by service_name"; + executePpl(q); // warm CacheGroup before = metadata(); - // Several more runs — footers must come from cache, driving hits up. for (int i = 0; i < 5; i++) { executePpl(q); } CacheGroup after = metadata(); - assertTrue( - "metadata cache must hold at least one footer entry, got " + after.entries(), - after.entries() >= 1 - ); + assertTrue("metadata cache must hold at least one footer entry, got " + after.entries(), after.entries() >= 1); assertTrue( String.format(Locale.ROOT, "metadata cache must register hits across repeated queries (before=%d after=%d)", before.hits(), after.hits()), after.hits() > before.hits() ); - assertTrue( - "metadata cache must advertise its configured byte budget", - after.sizeLimitBytes() > 0 - ); + assertTrue("metadata cache must advertise its configured byte budget", after.sizeLimitBytes() > 0); } - // ---- scoped cache: populate, hit, bounded --------------------------- + // ---- populate, hit, bounded ----------------------------------------- /** - * A filtered listing query populates the scoped cache (entry + bytes > 0) - * at query time, and re-running the IDENTICAL query is a pure hit: hits up; - * misses, entries, and memory_bytes flat. Measures hits, misses, AND size. + * A filtered listing query populates the scoped caches (entries + bytes > 0) + * at query time, and re-running the IDENTICAL query is a pure hit in BOTH + * caches: hits up; misses, entries, and memory_bytes flat. */ public void testSameListingQueryReRunIsPureCacheHit() throws IOException { - String query = "source=" + DATASET.indexName + " | where status >= 400 | stats count() by service_name"; + clearScopedCaches(); + String query = src() + " | where status >= 400 | stats count() by service_name"; executePpl(query); - CacheGroup afterFirst = scoped(); - assertTrue("scoped cache must hold >= 1 entry after a listing query", afterFirst.entries() >= 1); - assertTrue("populated scoped cache must consume memory_bytes", afterFirst.memoryBytes() > 0); + CacheGroup ci1 = columnIndex(); + CacheGroup oi1 = offsetIndex(); + assertTrue("ColumnIndex cache must hold >= 1 cell after a filtered query", ci1.entries() >= 1); + assertTrue("ColumnIndex cache must consume memory_bytes", ci1.memoryBytes() > 0); + assertTrue("OffsetIndex cache must hold >= 1 cell after a query that reads columns", oi1.entries() >= 1); executePpl(query); - CacheGroup afterSecond = scoped(); + CacheGroup ci2 = columnIndex(); + CacheGroup oi2 = offsetIndex(); assertTrue( - String.format(Locale.ROOT, "re-run must register hits (h1=%d m1=%d h2=%d m2=%d)", - afterFirst.hits(), afterFirst.misses(), afterSecond.hits(), afterSecond.misses()), - afterSecond.hits() > afterFirst.hits() + String.format(Locale.ROOT, "CI re-run must register hits (h1=%d h2=%d)", ci1.hits(), ci2.hits()), + ci2.hits() > ci1.hits() ); - assertEquals("re-run must NOT add misses", afterFirst.misses(), afterSecond.misses()); - assertEquals("re-run must NOT add entries (no duplication)", afterFirst.entries(), afterSecond.entries()); - assertEquals("re-run must NOT grow memory_bytes (no over-alloc)", afterFirst.memoryBytes(), afterSecond.memoryBytes()); + assertEquals("CI re-run must NOT add misses", ci1.misses(), ci2.misses()); + assertEquals("CI re-run must NOT add cells (no duplication)", ci1.entries(), ci2.entries()); + assertEquals("CI re-run must NOT grow memory_bytes", ci1.memoryBytes(), ci2.memoryBytes()); + + assertTrue("OI re-run must register hits", oi2.hits() > oi1.hits()); + assertEquals("OI re-run must NOT add misses", oi1.misses(), oi2.misses()); + assertEquals("OI re-run must NOT add cells", oi1.entries(), oi2.entries()); + assertEquals("OI re-run must NOT grow memory_bytes", oi1.memoryBytes(), oi2.memoryBytes()); } /** - * Repeated identical queries keep the cache bounded: entries and bytes do not + * Repeated identical queries keep the caches bounded: cells and bytes do not * grow after the first populating run; the re-runs all register hits. */ public void testRepeatedListingQueriesDoNotGrowTheCache() throws IOException { - String query = "source=" + DATASET.indexName + " | where status >= 200 | stats count()"; + clearScopedCaches(); + String query = src() + " | where status >= 200 | stats count()"; executePpl(query); - CacheGroup baseline = scoped(); - assertTrue("entry must exist after first run", baseline.entries() >= 1); + CacheGroup ciBase = columnIndex(); + CacheGroup oiBase = offsetIndex(); + assertTrue("CI cell must exist after first run", ciBase.entries() >= 1); for (int i = 0; i < 9; i++) { executePpl(query); } - CacheGroup after = scoped(); - - assertEquals("repeated queries must not add entries", baseline.entries(), after.entries()); - assertEquals("repeated queries must not grow memory_bytes", baseline.memoryBytes(), after.memoryBytes()); - assertEquals("repeated queries must not add misses", baseline.misses(), after.misses()); + CacheGroup ciAfter = columnIndex(); + CacheGroup oiAfter = offsetIndex(); + + assertEquals("repeated queries must not add CI cells", ciBase.entries(), ciAfter.entries()); + assertEquals("repeated queries must not grow CI memory_bytes", ciBase.memoryBytes(), ciAfter.memoryBytes()); + assertEquals("repeated queries must not add CI misses", ciBase.misses(), ciAfter.misses()); + assertEquals("repeated queries must not add OI cells", oiBase.entries(), oiAfter.entries()); + assertEquals("repeated queries must not add OI misses", oiBase.misses(), oiAfter.misses()); assertTrue( - String.format(Locale.ROOT, "the 9 re-runs must register hits (baseline=%d after=%d)", - baseline.hits(), after.hits()), - after.hits() >= baseline.hits() + 9 + String.format(Locale.ROOT, "the 9 CI re-runs must register hits (base=%d after=%d)", + ciBase.hits(), ciAfter.hits()), + ciAfter.hits() >= ciBase.hits() + 9 ); } /** - * The scoped cache must never exceed its configured byte budget — a basic + * Neither scoped cache may exceed its configured byte budget — a basic * "no over-allocation" invariant readable from the stats API. */ - public void testScopedCacheStaysWithinBudget() throws IOException { - // Exercise a few distinct predicate shapes to build whatever entries the - // listing path produces, then assert occupancy <= budget. - executePpl("source=" + DATASET.indexName + " | where status >= 400 | stats count()"); - executePpl("source=" + DATASET.indexName + " | where status < 300 | stats count()"); - CacheGroup snap = scoped(); + public void testScopedCachesStayWithinBudget() throws IOException { + executePpl(src() + " | where status >= 400 | stats count()"); + executePpl(src() + " | where status < 300 | stats count()"); + executePpl(src() + " | where log_level = 'ERROR' | stats count()"); + for (CacheGroup snap : new CacheGroup[] { columnIndex(), offsetIndex() }) { + assertTrue( + String.format(Locale.ROOT, "scoped cache memory_bytes (%d) must stay within size_limit_bytes (%d)", + snap.memoryBytes(), snap.sizeLimitBytes()), + snap.memoryBytes() <= snap.sizeLimitBytes() + ); + } + } + + // ---- cell reuse: ColumnIndex (predicate-driven) --------------------- + + /** + * Adding a column to a predicate must reuse the cells the first predicate + * already decoded — only the genuinely new column's cells are read. Filter + * {@code status}, then {@code status AND log_level}: the {@code status} cells + * are reused (CI hits strictly increase) and no fewer cells exist than before + * (the {@code log_level} cells are added, never replacing {@code status}). + */ + public void testAddingPredicateColumnReusesExistingCells() throws IOException { + clearScopedCaches(); + + executePpl(src() + " | where status >= 400 | stats count()"); + CacheGroup afterStatus = columnIndex(); + assertTrue("first predicate must populate CI cells", afterStatus.entries() >= 1); + long missesAfterStatus = afterStatus.misses(); + + // Predicate now also covers log_level: status cells reused, log_level new. + executePpl(src() + " | where status >= 400 and log_level = 'ERROR' | stats count()"); + CacheGroup afterBoth = columnIndex(); + + assertTrue( + String.format(Locale.ROOT, "adding a column must REUSE the status cells (hits %d -> %d)", + afterStatus.hits(), afterBoth.hits()), + afterBoth.hits() > afterStatus.hits() + ); assertTrue( - String.format(Locale.ROOT, "scoped cache memory_bytes (%d) must stay within size_limit_bytes (%d)", - snap.memoryBytes(), snap.sizeLimitBytes()), - snap.memoryBytes() <= snap.sizeLimitBytes() + "adding a column must keep all prior cells and add the new column's cells", + afterBoth.entries() >= afterStatus.entries() + ); + // The only NEW misses are for log_level's cells — status was never re-decoded. + assertTrue( + "new misses must be bounded by the newly added column's cells (status not re-decoded)", + afterBoth.misses() > missesAfterStatus ); } - // ---- no-breakage query sweep ---------------------------------------- - /** - * A spread of query shapes must all execute successfully (HTTP 200, parseable - * datarows) with the page-index changes in place: plain projection, - * aggregation, multi-column filter, full-text match (Lucene-delegated path), - * and a mixed predicate. This is the broad "nothing is broken" guard. + * Two predicates on the SAME column with DIFFERENT literals must share the + * cells — the predicate VALUE never enters the cache key, so changing it adds + * no cells and the second query is a pure hit on the first's cells. */ - public void testVariedQueryShapesAllExecute() throws IOException { - String idx = DATASET.indexName; + public void testDifferentLiteralsSameColumnShareCells() throws IOException { + clearScopedCaches(); - // Plain projection (no predicate). - assertEquals("plain projection returns all docs", TOTAL_DOCS, rowCount("source=" + idx + " | fields service_name, status")); + executePpl(src() + " | where status >= 400 | stats count()"); + CacheGroup first = columnIndex(); + assertTrue("first literal must populate CI cells", first.entries() >= 1); - // Aggregation with grouping. + executePpl(src() + " | where status >= 100 | stats count()"); + CacheGroup second = columnIndex(); + + assertEquals("a different literal on the same column must add NO cells", first.entries(), second.entries()); + assertEquals("a different literal must add NO misses (cells reused)", first.misses(), second.misses()); assertTrue( - "grouped aggregation must return at least one bucket", - rowCount("source=" + idx + " | stats count() by service_name") >= 1 + String.format(Locale.ROOT, "the second literal must HIT the existing cells (hits %d -> %d)", + first.hits(), second.hits()), + second.hits() > first.hits() ); + } - // Multi-column native predicate (listing path). - assertEquals( - "multi-column filter must be exact", - scalarAgg("source=" + idx + " | where status >= 400 and log_level = 'ERROR' | stats count()"), - scalarAgg("source=" + idx + " | where log_level = 'ERROR' and status >= 400 | stats count()") + /** + * A predicate that is a SUBSET of an already-cached predicate's columns reuses + * the relevant cells without decoding anything new. Cache a two-column + * predicate ({@code status AND log_level}), then run a one-column predicate + * ({@code status}) — {@code status}'s cells are a pure hit, adding no cells and + * no misses. (We use a compound predicate to bring the keyword {@code log_level} + * onto the native page-index path; a standalone keyword equality is fully + * Lucene-delegated and builds no page index — see the handoff notes.) + */ + public void testSubsetPredicateReusesCachedCells() throws IOException { + clearScopedCaches(); + + // Compound predicate caches cells for BOTH status and log_level. + executePpl(src() + " | where status >= 400 and log_level = 'ERROR' | stats count()"); + CacheGroup afterBoth = columnIndex(); + assertTrue("compound predicate must populate CI cells for both columns", afterBoth.entries() >= 2); + + // Subset predicate (status only): its cells are already cached → pure hit. + executePpl(src() + " | where status >= 400 | stats count()"); + CacheGroup afterSubset = columnIndex(); + assertEquals("a subset predicate must add NO new cells", afterBoth.entries(), afterSubset.entries()); + assertEquals("a subset predicate must add NO misses (cells reused)", afterBoth.misses(), afterSubset.misses()); + assertTrue( + String.format(Locale.ROOT, "a subset predicate must HIT the cached cells (hits %d -> %d)", + afterBoth.hits(), afterSubset.hits()), + afterSubset.hits() > afterBoth.hits() ); + } - // Full-text match — exercises the Lucene-delegated path; must not error. - // (Count is data-dependent; we only assert it executes and is non-negative.) - long matchCount = scalarAgg("source=" + idx + " | where match(message, 'timeout') | stats count()"); - assertTrue("match() query must execute and return a non-negative count", matchCount >= 0); + // ---- cell reuse: OffsetIndex (projection-driven) -------------------- - // Mixed: native predicate + full-text, then aggregate. - long mixed = scalarAgg("source=" + idx + " | where status >= 400 or match(message, 'error') | stats count()"); - assertTrue("mixed predicate query must execute", mixed >= 0); + /** + * Different projections on the same predicate must reuse the shared OffsetIndex + * column cells and only decode the newly projected column. Project a small set + * of fields, then a different set sharing some columns: OI hits increase + * (shared columns reused) while only the genuinely new column adds a cell. + */ + public void testDifferentProjectionsReuseOffsetIndexCells() throws IOException { + clearScopedCaches(); + + // First projection. + executePpl(src() + " | where status >= 400 | fields status, service_name"); + CacheGroup first = offsetIndex(); + assertTrue("first projection must populate OI cells", first.entries() >= 1); + + // Overlapping projection (shares status; adds log_level). + executePpl(src() + " | where status >= 400 | fields status, log_level"); + CacheGroup second = offsetIndex(); + + assertTrue( + String.format(Locale.ROOT, "overlapping projection must REUSE shared OI column cells (hits %d -> %d)", + first.hits(), second.hits()), + second.hits() > first.hits() + ); + assertTrue("overlapping projection must keep prior cells", second.entries() >= first.entries()); } - // ---- cross-path sharing (constraint #1) ------------------------------ + /** + * The OffsetIndex cache is keyed only on {@code (file, col)} — independent of + * the predicate. Two queries with DIFFERENT predicates but the SAME projection + * must reuse the same OffsetIndex column cells (predicate changes don't + * multiply OI cells). + */ + public void testOffsetIndexIndependentOfPredicate() throws IOException { + clearScopedCaches(); + + executePpl(src() + " | where status >= 400 | fields status, service_name"); + CacheGroup first = offsetIndex(); + + // Different predicate, same projected columns → same OI cells. + executePpl(src() + " | where log_level = 'ERROR' | fields status, service_name"); + CacheGroup second = offsetIndex(); + + assertEquals( + "same projection under a different predicate must add NO new OI cells", + first.entries(), + second.entries() + ); + assertTrue("the shared OI cells must be hit", second.hits() > first.hits()); + } + + // ---- cross-path sharing (one cache, both scan paths) ---------------- /** - * The unified cache must be shared across scan paths: an entry built for a - * predicate column on one path is reused by the other for the same column. - * - *

Here a listing query filters {@code status}, then an indexed query + * The scoped caches must be shared across scan paths: a cell built for a + * predicate column on the listing path is reused by the indexed path for the + * same column. A listing query filters {@code status}; then an indexed query * (forced onto the indexed path by a {@code match(message, ...)} full-text - * filter) ALSO filters {@code status}. Because both paths resolve the same - * file + predicate column to the same scoped-cache key, the indexed query - * must HIT the entry the listing query built: hits increase while entries do - * not grow (no second, path-specific entry). + * filter) ALSO filters {@code status}, so it resolves to the SAME + * {@code (file, status, rg)} cells and HITS them — CI hits increase while the + * cell count does not grow (no second, path-specific cells). */ public void testCrossPathSharingListingThenIndexed() throws IOException { - String idx = DATASET.indexName; - // Listing path: numeric predicate on `status` populates the scoped entry. - String listingQuery = "source=" + idx + " | where status >= 400 | stats count()"; - executePpl(listingQuery); - CacheGroup afterListing = scoped(); - assertTrue("listing query must populate a scoped entry", afterListing.entries() >= 1); + clearScopedCaches(); + + // Listing path: numeric predicate on `status` populates CI cells. + executePpl(src() + " | where status >= 400 | stats count()"); + CacheGroup afterListing = columnIndex(); + assertTrue("listing query must populate scoped CI cells", afterListing.entries() >= 1); // Indexed path: a match() filter forces indexed routing; it also filters - // `status`, so it resolves to the SAME (file, status) scoped-cache key. - String indexedQuery = - "source=" + idx + " | where match(message, 'timeout') and status >= 400 | stats count()"; - executePpl(indexedQuery); - CacheGroup afterIndexed = scoped(); + // `status`, so it resolves to the SAME (file, status, rg) cells. + executePpl(src() + " | where match(message, 'timeout') and status >= 400 | stats count()"); + CacheGroup afterIndexed = columnIndex(); assertTrue( String.format(Locale.ROOT, - "indexed query on the same predicate column must HIT the listing entry (listing hits=%d indexed hits=%d)", + "indexed query on the same predicate column must HIT the listing cells (listing hits=%d indexed hits=%d)", afterListing.hits(), afterIndexed.hits()), afterIndexed.hits() > afterListing.hits() ); assertEquals( - "cross-path reuse must NOT create a second entry for the same (file, predicate column)", + "cross-path reuse must NOT create new cells for the same (file, predicate column)", afterListing.entries(), afterIndexed.entries() ); assertEquals( - "cross-path reuse must NOT grow memory_bytes", + "cross-path reuse must NOT grow CI memory_bytes", afterListing.memoryBytes(), afterIndexed.memoryBytes() ); } + + // ---- clear endpoint -------------------------------------------------- + + /** + * The clear endpoint must drop all cells and reset counters in BOTH scoped + * caches: after a populating query and a clear, both groups read zero entries, + * zero hits, and zero misses, while keeping their configured budgets. + */ + public void testClearEndpointResetsBothCaches() throws IOException { + executePpl(src() + " | where status >= 400 | fields status, service_name"); + // Something must be cached before we clear. + assertTrue("a filtered+projected query must populate CI cells", columnIndex().entries() >= 1); + assertTrue("a filtered+projected query must populate OI cells", offsetIndex().entries() >= 1); + + clearScopedCaches(); + + CacheGroup ci = columnIndex(); + CacheGroup oi = offsetIndex(); + assertEquals("clear must reset CI cells", 0, ci.entries()); + assertEquals("clear must reset CI hits", 0, ci.hits()); + assertEquals("clear must reset CI misses", 0, ci.misses()); + assertEquals("clear must reset CI memory_bytes", 0, ci.memoryBytes()); + assertTrue("clear must keep the CI budget", ci.sizeLimitBytes() > 0); + assertEquals("clear must reset OI cells", 0, oi.entries()); + assertEquals("clear must reset OI hits", 0, oi.hits()); + assertEquals("clear must reset OI misses", 0, oi.misses()); + assertTrue("clear must keep the OI budget", oi.sizeLimitBytes() > 0); + } + + // ---- no-breakage query sweep ---------------------------------------- + + /** + * A spread of query shapes must all execute successfully with the cell-keyed + * caches in place: plain projection, aggregation, multi-column filter, + * full-text match (Lucene-delegated path), and a mixed predicate. + */ + public void testVariedQueryShapesAllExecute() throws IOException { + String idx = DATASET.indexName; + + assertEquals("plain projection returns all docs", TOTAL_DOCS, rowCount("source=" + idx + " | fields service_name, status")); + + assertTrue( + "grouped aggregation must return at least one bucket", + rowCount("source=" + idx + " | stats count() by service_name") >= 1 + ); + + assertEquals( + "multi-column filter must be order-independent", + scalarAgg("source=" + idx + " | where status >= 400 and log_level = 'ERROR' | stats count()"), + scalarAgg("source=" + idx + " | where log_level = 'ERROR' and status >= 400 | stats count()") + ); + + long matchCount = scalarAgg("source=" + idx + " | where match(message, 'timeout') | stats count()"); + assertTrue("match() query must execute and return a non-negative count", matchCount >= 0); + + long mixed = scalarAgg("source=" + idx + " | where status >= 400 or match(message, 'error') | stats count()"); + assertTrue("mixed predicate query must execute", mixed >= 0); + } } From 86ebdfdac55774ea1882d2575baa8a56f5c3b6dc Mon Sep 17 00:00:00 2001 From: G Date: Wed, 17 Jun 2026 08:28:51 +0530 Subject: [PATCH 14/18] fixes and debug logging Signed-off-by: G --- .../rust/src/indexed_executor.rs | 54 +++++++++++- .../src/indexed_table/page_index_loader.rs | 54 +++++++++++- .../rust/src/scoped_page_index_optimizer.rs | 17 ++++ .../rust/src/shard_scoped_reader.rs | 36 ++++++-- .../ClearScopedPageIndexCacheActionType.java | 31 +++++++ .../ClearScopedPageIndexCacheNodeRequest.java | 28 +++++++ ...ClearScopedPageIndexCacheNodeResponse.java | 38 +++++++++ ...ClearScopedPageIndexCacheNodesRequest.java | 32 ++++++++ ...learScopedPageIndexCacheNodesResponse.java | 67 +++++++++++++++ ...nsportClearScopedPageIndexCacheAction.java | 82 +++++++++++++++++++ 10 files changed, 432 insertions(+), 7 deletions(-) create mode 100644 sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/action/stats/ClearScopedPageIndexCacheActionType.java create mode 100644 sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/action/stats/ClearScopedPageIndexCacheNodeRequest.java create mode 100644 sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/action/stats/ClearScopedPageIndexCacheNodeResponse.java create mode 100644 sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/action/stats/ClearScopedPageIndexCacheNodesRequest.java create mode 100644 sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/action/stats/ClearScopedPageIndexCacheNodesResponse.java create mode 100644 sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/action/stats/TransportClearScopedPageIndexCacheAction.java diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_executor.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_executor.rs index 65e9804b1a7b4..f480d075c12e9 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_executor.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_executor.rs @@ -343,6 +343,40 @@ fn collect_predicate_column_names( } names.into_iter().collect() } +/// Collect the names of every column the logical plan references anywhere +/// (output projection, filter, sort, aggregation, …). This is a deliberate +/// **superset** of the columns the scan will actually read, used to scope the +/// OffsetIndex (`predicate ∪ projection`). A superset is safe: an OffsetIndex +/// built for more columns than are read just caches a little extra and never +/// breaks a read; building for *fewer* would panic the parquet reader (a +/// projected column with an empty page-locations list). Walking the whole plan +/// rather than just the top projection guarantees we never undershoot. +fn collect_plan_column_names(plan: &datafusion::logical_expr::LogicalPlan) -> Vec { + use datafusion::logical_expr::LogicalPlan; + let mut names = BTreeSet::new(); + let _ = plan.apply(|node| { + // Every expression on this node, recursed into for Column refs. + let _ = node.apply_expressions(|expr| { + let _ = expr.apply(|e| { + if let datafusion::logical_expr::Expr::Column(col) = e { + names.insert(col.name().to_string()); + } + Ok(TreeNodeRecursion::Continue) + }); + Ok(TreeNodeRecursion::Continue) + }); + // A bare TableScan with a projection carries its read columns in its + // schema rather than in expressions — fold those in too. + if let LogicalPlan::TableScan(scan) = node { + for f in scan.projected_schema.fields() { + names.insert(f.name().to_string()); + } + } + Ok(TreeNodeRecursion::Continue) + }); + names.into_iter().collect() +} + /// For a tree classified as `SingleCollector`, walk it to find the single /// Collector leaf and return its query bytes. fn single_collector_id(tree: &BoolNode) -> Option { @@ -964,6 +998,15 @@ async unsafe fn execute_indexed_with_context_inner( // metadata and pruning no-ops. let predicate_column_names = collect_predicate_column_names(extraction.as_ref(), &schema); if !predicate_column_names.is_empty() { + // Scope the OffsetIndex to the columns this query actually touches + // (projection ∪ predicate), instead of decoding+caching it for all + // columns. On the 402-column textbench schema the all-columns OffsetIndex + // was ~256 MB/file and was rebuilt every query regardless of the + // predicate; scoping it to the read columns is the difference between a + // few MB and 256 MB, and makes the cache actually fit its budget. + // `collect_plan_column_names` is a superset of the read set (safe — see + // its docs); the loader additionally unions in the predicate cols + col 0. + let projection_column_names = collect_plan_column_names(&logical_plan); for segment in segments.iter_mut() { let parquet_cols = crate::indexed_table::page_index_loader::resolve_predicate_parquet_columns( &schema, @@ -973,11 +1016,20 @@ async unsafe fn execute_indexed_with_context_inner( if parquet_cols.is_empty() { continue; } - if let Some(augmented) = crate::indexed_table::page_index_loader::load_scoped_page_index( + // Projection (∪ predicate ∪ {0}) leaf indices for THIS file, for the + // OffsetIndex column scoping. Resolved the same way as predicate cols + // so both scan paths agree on the cache key. + let offset_cols = crate::indexed_table::page_index_loader::resolve_predicate_parquet_columns( + &schema, + &segment.metadata, + &projection_column_names, + ); + if let Some(augmented) = crate::indexed_table::page_index_loader::load_scoped_page_index_cols( &store, &segment.object_path, &segment.metadata, &parquet_cols, + &offset_cols, ) .await { diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/page_index_loader.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/page_index_loader.rs index e75d8c336fd72..9fd7bd0450e6c 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/page_index_loader.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/page_index_loader.rs @@ -416,11 +416,30 @@ async fn load_combined( if parquet_cols.is_empty() { return None; } + let t_total = std::time::Instant::now(); + let t_ci = std::time::Instant::now(); let column_index = get_or_build_column_index(store, location, footer_meta, parquet_cols, surviving_rgs).await?; + let ci_ms = t_ci.elapsed(); + let t_oi = std::time::Instant::now(); let offset_index = get_or_build_offset_index(store, location, footer_meta, parquet_cols, offset_cols).await?; - Some(graft(footer_meta, column_index, offset_index)) + let oi_ms = t_oi.elapsed(); + // The graft deep-clones the whole footer (row_groups + column chunk metadata + // are owned, not Arc) — paid on EVERY query, hit or miss. Time it separately + // so we can see if this per-query clone (Problem B) dominates the scoped-cache + // path vs the cell decode/lookup. + let t_graft = std::time::Instant::now(); + let grafted = graft(footer_meta, column_index, offset_index); + log::debug!( + "scoped-load {}: total={:?} (colidx={:?} offidx={:?} graft={:?})", + location, + t_total.elapsed(), + ci_ms, + oi_ms, + t_graft.elapsed(), + ); + Some(grafted) } /// Build a fresh `ParquetMetaData` = `footer` with the page-index pair grafted @@ -501,10 +520,21 @@ async fn get_or_build_column_index( } else { return None; } + let ci_hits = build_rgs.len() * parquet_cols.len() - missing.len(); + log::debug!( + "scoped-colidx {}: cells {}rgs×{}cols, hit={} miss={} (rg_scoped={})", + location, + build_rgs.len(), + parquet_cols.len(), + ci_hits, + missing.len(), + surviving_rgs.is_some(), + ); // Phase 2: decode the missing cells (vectored fetch grouped by RG), place // them in the matrix, and populate the cache. if !missing.is_empty() { + let t_decode = std::time::Instant::now(); let built = build_column_index_cells(store, location, footer_meta, &missing).await?; if let Ok(mut cache) = COLUMN_INDEX_CACHE.lock() { for (col, rg, cell, size) in built { @@ -512,6 +542,12 @@ async fn get_or_build_column_index( cache.insert(CiCellKey { path: path.clone(), col, rg }, cell, size); } } + log::debug!( + "scoped-colidx {}: decoded {} miss cells in {:?}", + location, + missing.len(), + t_decode.elapsed(), + ); } Some(matrix) @@ -639,11 +675,27 @@ async fn get_or_build_offset_index( } else { return None; } + log::debug!( + "scoped-offidx {}: off_cols={}/{} (scoped={}) hit={} miss={}", + location, + off_cols.len(), + num_cols, + offset_cols.is_some(), + off_cols.len() - missing.len(), + missing.len(), + ); // Phase 2: decode the missing columns (each spanning all RGs), scatter into // the matrix, and populate the cache. if !missing.is_empty() { + let t_decode = std::time::Instant::now(); let built = build_offset_index_columns(store, location, footer_meta, &missing, num_rgs).await?; + log::debug!( + "scoped-offidx {}: decoded {} miss cols in {:?}", + location, + missing.len(), + t_decode.elapsed(), + ); if let Ok(mut cache) = OFFSET_INDEX_CACHE.lock() { for (col, column, size) in built { scatter_offset_column(&mut matrix, col, &column); diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/scoped_page_index_optimizer.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/scoped_page_index_optimizer.rs index 643a2070052fd..346b0dcaba2c3 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/scoped_page_index_optimizer.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/scoped_page_index_optimizer.rs @@ -109,6 +109,21 @@ impl PhysicalOptimizerRule for ScopedPageIndexOptimizer { return Ok(Transformed::no(node)); } + // Projected column NAMES (the columns this scan actually reads), for + // scoping the OffsetIndex to `predicate ∪ projection` instead of all + // columns. `projected_schema()` reflects the projection pushed into + // the scan; fall back to the full file schema if it can't be derived + // (safe — the loader then builds all-column offsets, the old behavior). + let projection_names: Vec = match config.projected_schema() { + Ok(ps) => ps + .fields() + .iter() + .map(|f| f.name().to_string()) + .filter(|n| file_schema.index_of(n).is_ok()) + .collect(), + Err(_) => Vec::new(), + }; + // Build the scoped factory and reinstall the source. The predicate is // retained for parity but not used for RG scoping (Step 1 builds an // all-row-group, column-scoped page index — see the reader's docs). @@ -116,6 +131,7 @@ impl PhysicalOptimizerRule for ScopedPageIndexOptimizer { Arc::clone(&self.store), Arc::clone(&self.metadata_cache), names, + projection_names, Some(Arc::clone(predicate)), Arc::clone(file_schema), )); @@ -233,6 +249,7 @@ mod tests { Arc::clone(&store), Arc::clone(&cache), vec!["a".to_string()], + vec!["a".to_string()], None, sch.clone(), )); diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/shard_scoped_reader.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/shard_scoped_reader.rs index 94222622d766b..f41292b30998b 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/shard_scoped_reader.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/shard_scoped_reader.rs @@ -73,7 +73,7 @@ use object_store::{ObjectStore, ObjectStoreExt}; use prost::bytes::Bytes; use crate::indexed_table::page_index_loader::{ - load_scoped_page_index, resolve_predicate_parquet_columns, + load_scoped_page_index_cols, resolve_predicate_parquet_columns, }; use crate::indexed_table::parquet_bridge::load_parquet_metadata; @@ -93,6 +93,10 @@ pub struct ScopedPageIndexReaderFactory { /// scoping" — `get_metadata` returns footer-only and the opener loads the /// page index on demand as usual. predicate_column_names: Arc>, + /// File-column names this scan PROJECTS (reads). Used to scope the + /// OffsetIndex to `predicate ∪ projection` instead of all columns. Empty = + /// fall back to all-column offsets (old behavior). + projection_column_names: Arc>, /// The physical predicate (if any). Retained in the constructor signature for /// parity with the indexed path, but intentionally NOT used for RG scoping /// here (see module docs). @@ -107,6 +111,7 @@ impl ScopedPageIndexReaderFactory { store: Arc, metadata_cache: Arc, predicate_column_names: Vec, + projection_column_names: Vec, predicate: Option>, file_schema: SchemaRef, ) -> Self { @@ -114,6 +119,7 @@ impl ScopedPageIndexReaderFactory { store, metadata_cache, predicate_column_names: Arc::new(predicate_column_names), + projection_column_names: Arc::new(projection_column_names), predicate, file_schema, } @@ -134,6 +140,7 @@ impl ParquetFileReaderFactory for ScopedPageIndexReaderFactory { store: Arc::clone(&self.store), metadata_cache: Arc::clone(&self.metadata_cache), predicate_column_names: Arc::clone(&self.predicate_column_names), + projection_column_names: Arc::clone(&self.projection_column_names), file_schema: Arc::clone(&self.file_schema), location: file.object_meta.location.clone(), metrics: file_metrics, @@ -145,6 +152,7 @@ struct ScopedPageIndexReader { store: Arc, metadata_cache: Arc, predicate_column_names: Arc>, + projection_column_names: Arc>, file_schema: SchemaRef, location: object_store::path::Path, metrics: ParquetFileMetrics, @@ -193,6 +201,7 @@ impl AsyncFileReader for ScopedPageIndexReader { let store = Arc::clone(&self.store); let metadata_cache = Arc::clone(&self.metadata_cache); let predicate_names = Arc::clone(&self.predicate_column_names); + let projection_names = Arc::clone(&self.projection_column_names); let file_schema = Arc::clone(&self.file_schema); let location = self.location.clone(); async move { @@ -213,14 +222,27 @@ impl AsyncFileReader for ScopedPageIndexReader { if !predicate_names.is_empty() { let parquet_cols = resolve_predicate_parquet_columns(&file_schema, &footer, &predicate_names); - if let Some(augmented) = - load_scoped_page_index(&store, &location, &footer, &parquet_cols).await + // Projected leaf indices for THIS file, for OffsetIndex column + // scoping (predicate ∪ projection ∪ {0}). Empty projection_names + // (couldn't derive) → empty offset_cols → loader unions in + // predicate + col0 only (still far smaller than all columns). + let offset_cols = + resolve_predicate_parquet_columns(&file_schema, &footer, &projection_names); + if let Some(augmented) = load_scoped_page_index_cols( + &store, + &location, + &footer, + &parquet_cols, + &offset_cols, + ) + .await { log::debug!( - "scoped-pageidx[listing]: {} rgs={} cols={}", + "scoped-pageidx[listing]: {} rgs={} pred_cols={} offset_cols={}", location, footer.num_row_groups(), - parquet_cols.len() + parquet_cols.len(), + offset_cols.len() ); return Ok(augmented); } @@ -310,6 +332,9 @@ mod tests { Arc::clone(&store), fresh_cache(), vec!["price".to_string()], + // Project both columns so the OffsetIndex is built for both (this test + // asserts a real OffsetIndex for every column). + vec!["price".to_string(), "qty".to_string()], None, schema, ); @@ -359,6 +384,7 @@ mod tests { Arc::clone(&store), fresh_cache(), vec![], + vec![], None, schema, ); diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/action/stats/ClearScopedPageIndexCacheActionType.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/action/stats/ClearScopedPageIndexCacheActionType.java new file mode 100644 index 0000000000000..65a8cea6ad577 --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/action/stats/ClearScopedPageIndexCacheActionType.java @@ -0,0 +1,31 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.be.datafusion.action.stats; + +import org.opensearch.action.ActionType; + +/** + * Action type for the broadcast "clear scoped page-index caches" transport action. + * + *

The scoped ColumnIndex and OffsetIndex caches are process-global singletons, + * one per node. Clearing them must fan out to every node, otherwise a + * single-node REST handler would leave the other nodes' caches populated — which + * is both surprising operationally and breaks cluster-aggregated stats assertions. + * + * @opensearch.internal + */ +public class ClearScopedPageIndexCacheActionType extends ActionType { + + public static final String NAME = "cluster:admin/_analytics_backend_datafusion/cache/scoped_page_index/clear"; + public static final ClearScopedPageIndexCacheActionType INSTANCE = new ClearScopedPageIndexCacheActionType(); + + private ClearScopedPageIndexCacheActionType() { + super(NAME, ClearScopedPageIndexCacheNodesResponse::new); + } +} diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/action/stats/ClearScopedPageIndexCacheNodeRequest.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/action/stats/ClearScopedPageIndexCacheNodeRequest.java new file mode 100644 index 0000000000000..3c7c3a6521d8a --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/action/stats/ClearScopedPageIndexCacheNodeRequest.java @@ -0,0 +1,28 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.be.datafusion.action.stats; + +import org.opensearch.core.common.io.stream.StreamInput; +import org.opensearch.transport.TransportRequest; + +import java.io.IOException; + +/** + * Per-node request in the clear-scoped-cache fan-out. Carries no payload. + * + * @opensearch.internal + */ +public class ClearScopedPageIndexCacheNodeRequest extends TransportRequest { + + public ClearScopedPageIndexCacheNodeRequest() {} + + public ClearScopedPageIndexCacheNodeRequest(StreamInput in) throws IOException { + super(in); + } +} diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/action/stats/ClearScopedPageIndexCacheNodeResponse.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/action/stats/ClearScopedPageIndexCacheNodeResponse.java new file mode 100644 index 0000000000000..0159040f82041 --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/action/stats/ClearScopedPageIndexCacheNodeResponse.java @@ -0,0 +1,38 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.be.datafusion.action.stats; + +import org.opensearch.action.support.nodes.BaseNodeResponse; +import org.opensearch.cluster.node.DiscoveryNode; +import org.opensearch.core.common.io.stream.StreamInput; +import org.opensearch.core.common.io.stream.StreamOutput; + +import java.io.IOException; + +/** + * Per-node response confirming the scoped page-index caches were cleared on this + * node. Carries only the node identity (the operation has no per-node payload). + * + * @opensearch.internal + */ +public class ClearScopedPageIndexCacheNodeResponse extends BaseNodeResponse { + + public ClearScopedPageIndexCacheNodeResponse(DiscoveryNode node) { + super(node); + } + + public ClearScopedPageIndexCacheNodeResponse(StreamInput in) throws IOException { + super(in); + } + + @Override + public void writeTo(StreamOutput out) throws IOException { + super.writeTo(out); + } +} diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/action/stats/ClearScopedPageIndexCacheNodesRequest.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/action/stats/ClearScopedPageIndexCacheNodesRequest.java new file mode 100644 index 0000000000000..f63a659e6d20d --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/action/stats/ClearScopedPageIndexCacheNodesRequest.java @@ -0,0 +1,32 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.be.datafusion.action.stats; + +import org.opensearch.action.support.nodes.BaseNodesRequest; +import org.opensearch.core.common.io.stream.StreamInput; + +import java.io.IOException; + +/** + * Cluster-level request to clear the scoped page-index caches on the target nodes + * (empty {@code nodesIds} means all nodes). Carries no payload — clearing is + * unconditional. + * + * @opensearch.internal + */ +public class ClearScopedPageIndexCacheNodesRequest extends BaseNodesRequest { + + public ClearScopedPageIndexCacheNodesRequest(String... nodesIds) { + super(nodesIds); + } + + public ClearScopedPageIndexCacheNodesRequest(StreamInput in) throws IOException { + super(in); + } +} diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/action/stats/ClearScopedPageIndexCacheNodesResponse.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/action/stats/ClearScopedPageIndexCacheNodesResponse.java new file mode 100644 index 0000000000000..d072a0aef0082 --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/action/stats/ClearScopedPageIndexCacheNodesResponse.java @@ -0,0 +1,67 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.be.datafusion.action.stats; + +import org.opensearch.action.FailedNodeException; +import org.opensearch.action.support.nodes.BaseNodesResponse; +import org.opensearch.cluster.ClusterName; +import org.opensearch.core.common.io.stream.StreamInput; +import org.opensearch.core.common.io.stream.StreamOutput; +import org.opensearch.core.xcontent.ToXContentFragment; +import org.opensearch.core.xcontent.XContentBuilder; + +import java.io.IOException; +import java.util.List; + +/** + * Aggregated cluster-wide response for the clear-scoped-cache broadcast. The + * {@code NodesResponseRestListener} wrapper supplies the {@code _nodes} header and + * {@code cluster_name}; this fragment adds only {@code acknowledged} and the list + * of node IDs that cleared. + * + * @opensearch.internal + */ +public class ClearScopedPageIndexCacheNodesResponse extends BaseNodesResponse + implements + ToXContentFragment { + + public ClearScopedPageIndexCacheNodesResponse( + ClusterName clusterName, + List nodes, + List failures + ) { + super(clusterName, nodes, failures); + } + + public ClearScopedPageIndexCacheNodesResponse(StreamInput in) throws IOException { + super(in); + } + + @Override + protected List readNodesFrom(StreamInput in) throws IOException { + return in.readList(ClearScopedPageIndexCacheNodeResponse::new); + } + + @Override + protected void writeNodesTo(StreamOutput out, List nodes) throws IOException { + out.writeList(nodes); + } + + @Override + public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException { + builder.field("acknowledged", failures() == null || failures().isEmpty()); + builder.field("cleared", "scoped_page_index_cache"); + builder.startArray("cleared_nodes"); + for (ClearScopedPageIndexCacheNodeResponse node : getNodes()) { + builder.value(node.getNode().getId()); + } + builder.endArray(); + return builder; + } +} diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/action/stats/TransportClearScopedPageIndexCacheAction.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/action/stats/TransportClearScopedPageIndexCacheAction.java new file mode 100644 index 0000000000000..4cdbc499908ff --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/action/stats/TransportClearScopedPageIndexCacheAction.java @@ -0,0 +1,82 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.be.datafusion.action.stats; + +import org.opensearch.action.FailedNodeException; +import org.opensearch.action.support.ActionFilters; +import org.opensearch.action.support.nodes.TransportNodesAction; +import org.opensearch.be.datafusion.nativelib.NativeBridge; +import org.opensearch.cluster.service.ClusterService; +import org.opensearch.common.inject.Inject; +import org.opensearch.core.common.io.stream.StreamInput; +import org.opensearch.threadpool.ThreadPool; +import org.opensearch.transport.TransportService; + +import java.io.IOException; +import java.util.List; + +/** + * Broadcast transport action that clears the process-global scoped page-index + * caches (ColumnIndex + OffsetIndex) on every target node. Mirrors + * {@link TransportDataFusionStatsAction}'s node fan-out so the operation reaches + * all nodes — not just the one that received the REST request. + * + * @opensearch.internal + */ +public class TransportClearScopedPageIndexCacheAction extends TransportNodesAction< + ClearScopedPageIndexCacheNodesRequest, + ClearScopedPageIndexCacheNodesResponse, + ClearScopedPageIndexCacheNodeRequest, + ClearScopedPageIndexCacheNodeResponse> { + + @Inject + public TransportClearScopedPageIndexCacheAction( + ThreadPool threadPool, + ClusterService clusterService, + TransportService transportService, + ActionFilters actionFilters + ) { + super( + ClearScopedPageIndexCacheActionType.NAME, + threadPool, + clusterService, + transportService, + actionFilters, + ClearScopedPageIndexCacheNodesRequest::new, + ClearScopedPageIndexCacheNodeRequest::new, + ThreadPool.Names.MANAGEMENT, + ClearScopedPageIndexCacheNodeResponse.class + ); + } + + @Override + protected ClearScopedPageIndexCacheNodesResponse newResponse( + ClearScopedPageIndexCacheNodesRequest request, + List responses, + List failures + ) { + return new ClearScopedPageIndexCacheNodesResponse(clusterService.getClusterName(), responses, failures); + } + + @Override + protected ClearScopedPageIndexCacheNodeRequest newNodeRequest(ClearScopedPageIndexCacheNodesRequest request) { + return new ClearScopedPageIndexCacheNodeRequest(); + } + + @Override + protected ClearScopedPageIndexCacheNodeResponse newNodeResponse(StreamInput in) throws IOException { + return new ClearScopedPageIndexCacheNodeResponse(in); + } + + @Override + protected ClearScopedPageIndexCacheNodeResponse nodeOperation(ClearScopedPageIndexCacheNodeRequest request) { + NativeBridge.clearScopedPageIndexCache(); + return new ClearScopedPageIndexCacheNodeResponse(clusterService.localNode()); + } +} From 1b5bacec9b9cf53e5c59bff4e88991c6c0bdcc4a Mon Sep 17 00:00:00 2001 From: G Date: Wed, 17 Jun 2026 12:19:11 +0530 Subject: [PATCH 15/18] page offset fixes and logger fixes Signed-off-by: G --- .../src/indexed_table/page_index_loader.rs | 106 +++++++++++++----- .../rust/src/shard_scoped_reader.rs | 8 ++ 2 files changed, 85 insertions(+), 29 deletions(-) diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/page_index_loader.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/page_index_loader.rs index 9fd7bd0450e6c..0bcf9533d3a58 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/page_index_loader.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/page_index_loader.rs @@ -323,12 +323,53 @@ pub fn clear_scoped_cache() { // ── Public API ────────────────────────────────────────────────────────────── -/// Map the query's arrow predicate-column names to this file's parquet column -/// indices, using the same resolution the pruner uses -/// (`StatisticsConverter::parquet_column_index`). Columns absent from the parquet -/// file (schema evolution) are skipped. Returns a sorted, deduped set so both -/// scan paths produce an identical key for the same logical predicate. +/// Map the query's predicate-column names to **this file's** parquet leaf +/// indices, resolving against the file's OWN schema so the indices are correct +/// even when the file is missing columns (schema evolution). +/// +/// # Why the file's own schema, not the shared table schema +/// +/// `StatisticsConverter`/`parquet_column` map a column by finding its position in +/// the supplied arrow schema and then matching that position to a parquet leaf +/// (`get_column_root_idx`). The textbench table schema is the **union** of all +/// files' columns (410 fields); a given file may physically contain fewer (e.g. +/// the merged file has 370 leaves — the absent columns are all-null and not +/// written). Resolving against the 410-field union therefore maps a column to the +/// WRONG leaf in a 370-leaf file (observed: `SeverityNumber` → leaf 291 against +/// the union, but its true leaf in the merged file is 149). We would then build +/// the scoped ColumnIndex/OffsetIndex at the wrong leaf and leave the real one an +/// empty placeholder — and DataFusion's pruner, which resolves against the file's +/// physical schema, reads the real leaf and panics on the empty `page_locations` +/// (`statistics.rs` `page_locations.last().unwrap()`). +/// +/// Deriving the arrow schema from the file footer (`parquet_to_arrow_schema`) +/// gives a 1:1 field↔leaf correspondence for that file, so the resolved index +/// matches what DataFusion dereferences. Columns absent from the file are skipped. pub fn resolve_predicate_parquet_columns( + _arrow_schema: &SchemaRef, + metadata: &ParquetMetaData, + predicate_column_names: &[String], +) -> Vec { + let parquet_schema = metadata.file_metadata().schema_descr(); + // Per-file arrow schema: 1:1 with this file's parquet leaves, so a column's + // arrow position maps to its true leaf. (The passed `_arrow_schema` is the + // union table schema and is intentionally NOT used for index resolution — + // see the doc comment.) + let file_arrow_schema = match datafusion::parquet::arrow::parquet_to_arrow_schema( + parquet_schema, + metadata.file_metadata().key_value_metadata(), + ) { + Ok(s) => Arc::new(s), + // If we can't derive the file schema, fall back to the union schema; the + // caller still falls back to footer-only on any downstream mismatch. + Err(_) => return resolve_with_schema(_arrow_schema, metadata, predicate_column_names), + }; + resolve_with_schema(&file_arrow_schema, metadata, predicate_column_names) +} + +/// Resolve predicate column names → parquet leaf indices against a specific arrow +/// schema, via the same `StatisticsConverter` mapping DataFusion's pruner uses. +fn resolve_with_schema( arrow_schema: &SchemaRef, metadata: &ParquetMetaData, predicate_column_names: &[String], @@ -431,7 +472,7 @@ async fn load_combined( // path vs the cell decode/lookup. let t_graft = std::time::Instant::now(); let grafted = graft(footer_meta, column_index, offset_index); - log::debug!( + native_bridge_common::log_debug!( "scoped-load {}: total={:?} (colidx={:?} offidx={:?} graft={:?})", location, t_total.elapsed(), @@ -520,21 +561,9 @@ async fn get_or_build_column_index( } else { return None; } - let ci_hits = build_rgs.len() * parquet_cols.len() - missing.len(); - log::debug!( - "scoped-colidx {}: cells {}rgs×{}cols, hit={} miss={} (rg_scoped={})", - location, - build_rgs.len(), - parquet_cols.len(), - ci_hits, - missing.len(), - surviving_rgs.is_some(), - ); - // Phase 2: decode the missing cells (vectored fetch grouped by RG), place // them in the matrix, and populate the cache. if !missing.is_empty() { - let t_decode = std::time::Instant::now(); let built = build_column_index_cells(store, location, footer_meta, &missing).await?; if let Ok(mut cache) = COLUMN_INDEX_CACHE.lock() { for (col, rg, cell, size) in built { @@ -542,12 +571,6 @@ async fn get_or_build_column_index( cache.insert(CiCellKey { path: path.clone(), col, rg }, cell, size); } } - log::debug!( - "scoped-colidx {}: decoded {} miss cells in {:?}", - location, - missing.len(), - t_decode.elapsed(), - ); } Some(matrix) @@ -658,8 +681,24 @@ async fn get_or_build_offset_index( } let path: Arc = Arc::from(location.as_ref()); + // Placeholder for columns we don't build: a SINGLE page spanning the whole row + // group, NOT an empty page-locations list. A scoped OffsetIndex is grafted as a + // full-width `[rg][col]` matrix; consumers (DataFusion's page pruner, arrow's + // reader, our indexed pruner) index it by absolute column and dereference + // `page_locations` (`.last()`, `[0]`, `windows(2)`). An EMPTY placeholder + // panics those (`page_locations.last().unwrap()` etc.) if any path touches a + // column we scoped out — which is hard to predict across every query shape + // (count/agg, SingleCollector prefetch, schema-evolved files). A one-page + // placeholder is always safe to dereference and makes pruning conservatively + // keep the whole RG (1 page = all rows → can't prune), never a wrong result. + let placeholder_for = |rg_idx: usize| -> datafusion::parquet::file::page_index::offset_index::OffsetIndexMetaData { + let mut b = OffsetIndexBuilder::new(); + b.append_offset_and_size(0, 0); + b.append_row_count(footer_meta.row_group(rg_idx).num_rows()); + b.build() + }; let mut matrix: ParquetOffsetIndex = (0..num_rgs) - .map(|_| (0..num_cols).map(|_| OffsetIndexBuilder::new().build()).collect()) + .map(|rg| (0..num_cols).map(|_| placeholder_for(rg)).collect()) .collect(); // Phase 1: serve cached columns; collect misses. @@ -675,12 +714,13 @@ async fn get_or_build_offset_index( } else { return None; } - log::debug!( - "scoped-offidx {}: off_cols={}/{} (scoped={}) hit={} miss={}", + native_bridge_common::log_debug!( + "scoped-offidx {}: off_cols={:?}/{} (scoped={}) pred={:?} hit={} miss={}", location, - off_cols.len(), + off_cols, num_cols, offset_cols.is_some(), + parquet_cols, off_cols.len() - missing.len(), missing.len(), ); @@ -690,7 +730,7 @@ async fn get_or_build_offset_index( if !missing.is_empty() { let t_decode = std::time::Instant::now(); let built = build_offset_index_columns(store, location, footer_meta, &missing, num_rgs).await?; - log::debug!( + native_bridge_common::log_debug!( "scoped-offidx {}: decoded {} miss cols in {:?}", location, missing.len(), @@ -763,6 +803,14 @@ async fn build_offset_index_columns( if decoded.len() != cols.len() { return None; } + // DIAG: report per-column page_locations length for this RG. + let plens: Vec = decoded.iter().map(|o| o.page_locations().len()).collect(); + native_bridge_common::log_debug!( + "scoped-offidx-decode rg={} cols={:?} page_locs_len={:?}", + plan.rg_idx, + cols, + plens, + ); for (k, entry) in decoded.into_iter().enumerate() { columns[k].push(entry); } diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/shard_scoped_reader.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/shard_scoped_reader.rs index f41292b30998b..38c4d231e6e1d 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/shard_scoped_reader.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/shard_scoped_reader.rs @@ -222,6 +222,14 @@ impl AsyncFileReader for ScopedPageIndexReader { if !predicate_names.is_empty() { let parquet_cols = resolve_predicate_parquet_columns(&file_schema, &footer, &predicate_names); + native_bridge_common::log_debug!( + "scoped-RESOLVE {}: names={:?} -> leaves={:?} (file_schema_fields={}, parquet_leaves={})", + location, + predicate_names.as_ref(), + parquet_cols, + file_schema.fields().len(), + footer.file_metadata().schema_descr().num_columns(), + ); // Projected leaf indices for THIS file, for OffsetIndex column // scoping (predicate ∪ projection ∪ {0}). Empty projection_names // (couldn't derive) → empty offset_cols → loader unions in From ce10cd9b456c03fdb8b889364f1e50c5c33f180f Mon Sep 17 00:00:00 2001 From: G Date: Thu, 18 Jun 2026 00:02:50 +0530 Subject: [PATCH 16/18] fixes for union schema - local schema mismatch + offset index columnar fixes Signed-off-by: G --- .../rust/src/indexed_executor.rs | 66 ++++- .../src/indexed_table/page_index_loader.rs | 144 ++++++++- .../rust/src/indexed_table/page_pruner.rs | 192 +++++++++++- .../rust/src/indexed_table/parquet_bridge.rs | 192 ++++++++++++ .../rust/src/indexed_table/tests_e2e/mod.rs | 196 +++++++++++++ .../indexed_table/tests_e2e/schema_drift.rs | 275 ++++++++++++++++++ 6 files changed, 1043 insertions(+), 22 deletions(-) diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_executor.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_executor.rs index f480d075c12e9..9459284e8e93b 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_executor.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_executor.rs @@ -343,16 +343,26 @@ fn collect_predicate_column_names( } names.into_iter().collect() } -/// Collect the names of every column the logical plan references anywhere -/// (output projection, filter, sort, aggregation, …). This is a deliberate -/// **superset** of the columns the scan will actually read, used to scope the -/// OffsetIndex (`predicate ∪ projection`). A superset is safe: an OffsetIndex -/// built for more columns than are read just caches a little extra and never -/// breaks a read; building for *fewer* would panic the parquet reader (a -/// projected column with an empty page-locations list). Walking the whole plan -/// rather than just the top projection guarantees we never undershoot. +/// Collect the names of every column the query actually reads — the columns +/// referenced by expressions anywhere in the plan (filter, residual, sort, +/// aggregation, output projection) plus a `TableScan`'s **explicit** projection. +/// Used to scope the OffsetIndex to `predicate ∪ projection`. +/// +/// # Why a `TableScan` with `projection == None` must NOT fold its full schema +/// +/// A `TableScan`'s `projected_schema` equals the **whole table schema** when +/// `projection` is `None` (a full column scan in the plan). On the indexed path +/// the substrait plan for e.g. `match(Body,..) AND SeverityNumber>=N | stats +/// count()` carries an unprojected `TableScan` (count reads no column values; +/// `match` is Lucene-delegated; only the residual `SeverityNumber` is read for +/// masking) — yet `projected_schema` lists all ~402 columns. Folding those made +/// `offset_cols` the entire schema (observed: offset cache 2MB→251MB, all +/// columns), which both bloated memory and, combined with per-file schema +/// evolution, produced a wrong count. The columns truly read are exactly the +/// `Expr::Column` references the walk already collects, so we only add a scan's +/// projection when it is **explicit** (`Some`). Still a safe superset of the read +/// set — never an undershoot that would panic the reader. fn collect_plan_column_names(plan: &datafusion::logical_expr::LogicalPlan) -> Vec { - use datafusion::logical_expr::LogicalPlan; let mut names = BTreeSet::new(); let _ = plan.apply(|node| { // Every expression on this node, recursed into for Column refs. @@ -365,13 +375,15 @@ fn collect_plan_column_names(plan: &datafusion::logical_expr::LogicalPlan) -> Ve }); Ok(TreeNodeRecursion::Continue) }); - // A bare TableScan with a projection carries its read columns in its - // schema rather than in expressions — fold those in too. - if let LogicalPlan::TableScan(scan) = node { - for f in scan.projected_schema.fields() { - names.insert(f.name().to_string()); - } - } + // NOTE: deliberately do NOT fold a `TableScan`'s declared projection. On + // the indexed path the substrait `TableScan` over-declares all columns + // (observed: projection = Some([0..401]) for `match() AND num | count()`), + // which scoped the OffsetIndex to all 402 columns — the bug. The columns + // actually read are captured by the expression walk above (filter / + // residual / sort / aggregate / output projection). A column referenced + // ONLY via a `TableScan` projection (not by any expression) is not read by + // value in a way that needs its real offset index for correctness; the + // one-page placeholder keeps any incidental dereference panic-safe. Ok(TreeNodeRecursion::Continue) }); names.into_iter().collect() @@ -620,6 +632,22 @@ mod tests { assert!(analyze_top_sort(&plan).is_none()); } + #[test] + fn collect_plan_column_names_includes_projection_and_predicate() { + // DIAG: what does collect_plan_column_names return for an explicit + // projection + predicate? `off_cols` is derived from this; if a SELECTed + // column is missing, it gets placeholdered and read through → the bug. + let plan = build_logical_plan("SELECT id, v FROM t WHERE ts > 100"); + let mut names = collect_plan_column_names(&plan); + names.sort(); + eprintln!("collect_plan_column_names => {:?}", names); + // The read set must include BOTH projected columns (id, v) AND the + // predicate column (ts). Anything read but absent here is placeholdered. + assert!(names.contains(&"id".to_string()), "missing projected id: {names:?}"); + assert!(names.contains(&"v".to_string()), "missing projected v: {names:?}"); + assert!(names.contains(&"ts".to_string()), "missing predicate ts: {names:?}"); + } + #[test] fn analyze_top_sort_returns_none_for_function_sort_key() { // `ORDER BY abs(v)` — leading sort key isn't a plain column. Catalog monotonicity @@ -1007,6 +1035,12 @@ async unsafe fn execute_indexed_with_context_inner( // `collect_plan_column_names` is a superset of the read set (safe — see // its docs); the loader additionally unions in the predicate cols + col 0. let projection_column_names = collect_plan_column_names(&logical_plan); + native_bridge_common::log_debug!( + "scoped-PROJCOLS predicate_names={:?} projection_names(n={})={:?}", + predicate_column_names, + projection_column_names.len(), + projection_column_names, + ); for segment in segments.iter_mut() { let parquet_cols = crate::indexed_table::page_index_loader::resolve_predicate_parquet_columns( &schema, diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/page_index_loader.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/page_index_loader.rs index 0bcf9533d3a58..e34fa15f1e76a 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/page_index_loader.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/page_index_loader.rs @@ -1290,6 +1290,135 @@ mod tests { assert_eq!(s.as_ref().map(kept), Some(16)); } + /// Schema-evolution fixture: a file whose physical layout is `[extra, price]` + /// (so `price` is parquet leaf **1**), 32 rows / 4 pages. Used to prove the + /// predicate→leaf resolution does NOT depend on a column's position in a wider + /// *union* schema. + fn evolved_extra_price_parquet() -> Bytes { + let schema = Arc::new(Schema::new(vec![ + Field::new("extra", DataType::Int32, false), + Field::new("price", DataType::Int32, false), + ])); + let extra: Vec = (1000..1032).collect(); + let prices: Vec = (0..32).collect(); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(extra)), Arc::new(Int32Array::from(prices))], + ) + .unwrap(); + let props = WriterProperties::builder() + .set_max_row_group_size(32) + .set_data_page_row_count_limit(8) + .set_write_batch_size(8) + .set_statistics_enabled(EnabledStatistics::Page) + .build(); + let mut buf: Vec = Vec::new(); + let mut w = ArrowWriter::try_new(&mut buf, schema.clone(), Some(props)).unwrap(); + w.write(&batch).unwrap(); + w.close().unwrap(); + Bytes::from(buf) + } + + /// Regression for the schema-evolution wrong-count bug: when the query's + /// (union) schema lists `price` at a DIFFERENT position than the file's + /// physical layout, the predicate must still resolve to the file's TRUE leaf + /// and scoped pruning must match the full-index pruning. Previously the + /// resolver used the union-schema position, scoped the page index at the wrong + /// leaf, and the residual mis-pruned → over-count. + #[tokio::test] + async fn scoped_resolution_is_per_file_under_schema_evolution() { + let _g = CACHE_TEST_GUARD.lock().unwrap(); + clear_scoped_cache_for_test(); + let bytes = evolved_extra_price_parquet(); + let (store, loc) = stage(bytes.clone()).await; + let fo = footer_only(&bytes); + + // The UNION/table schema the query carries: `price` at position 0 — but in + // THIS file `price` is physically leaf 1 (after `extra`). + let union_schema: SchemaRef = Arc::new(Schema::new(vec![ + Field::new("price", DataType::Int32, false), + Field::new("qty", DataType::Int32, false), + Field::new("extra", DataType::Int32, false), + ])); + + // Must resolve to the file's TRUE leaf for `price` = 1, NOT the union + // position 0 (which is `extra` in this file). + let cols = resolve_predicate_parquet_columns(&union_schema, &fo, &["price".to_string()]); + assert_eq!(cols, vec![1], "price must resolve to its per-file leaf (1), not union pos 0"); + + let aug = load_scoped_page_index(&store, &loc, &fo, &cols).await.unwrap(); + let full = full_index(&bytes); + // `price` pages: 0..8,8..16,16..24,24..32; `price >= 20` keeps the last two + // pages (rows 16..32 = 16 rows). Build the pruning predicate against the + // FILE schema (price at index 1) so the converter matches the data. + let file_schema: SchemaRef = Arc::new(Schema::new(vec![ + Field::new("extra", DataType::Int32, false), + Field::new("price", DataType::Int32, false), + ])); + let pp = build_pruning_predicate(&pred("price", 1, Operator::GtEq, 20), file_schema.clone()).unwrap(); + let s = PagePruner::new(&file_schema, Arc::clone(&aug)).prune_rg(&pp, 0, None); + let f = PagePruner::new(&file_schema, full).prune_rg(&pp, 0, None); + assert_eq!(s.as_ref().map(kept), f.as_ref().map(kept), "scoped pruning must match full index"); + assert_eq!(s.as_ref().map(kept), Some(16)); + clear_scoped_cache_for_test(); + } + + /// Page pruning over a column-scoped index must be a SAFE SUPERSET of the + /// full-index pruning — never drop a row the full index would keep (that + /// would be an under-count / lost result). It MAY keep extra rows (the + /// residual mask drops them post-decode), so equality is NOT required and is + /// the wrong invariant. + /// + /// The hazard this guards: with the page index scoped to `price` only, `qty` + /// is a one-page non-panicking OffsetIndex placeholder with a `NONE` + /// ColumnIndex. If the pruner TRUSTED that placeholder's (absent) stats it + /// would build a bogus single-page grid for `qty` and could mis-prune. The + /// `page_pruner` fix treats a `NONE`-ColumnIndex column as "no usable stats" + /// (like a schema-evolution-absent column) → it contributes "unknown" and + /// never prunes on `qty`, so the scoped result stays a conservative superset. + #[tokio::test] + async fn scoped_pruning_is_safe_superset_with_placeholdered_residual_col() { + let _g = CACHE_TEST_GUARD.lock().unwrap(); + clear_scoped_cache_for_test(); + // price 0..32 (pages 0..8,8..16,16..24,24..32); qty 100..132 (pages + // 100..108,108..116,116..124,124..132). 1 RG, 4 pages each. + let (bytes, schema) = two_col_parquet(); + let (store, loc) = stage(bytes.clone()).await; + let fo = footer_only(&bytes); + + // Scope the page index to `price` ONLY (mimics the predicate-scoped indexed + // path). `qty` therefore gets the one-page placeholder + NONE ColumnIndex. + let cols = resolve_predicate_parquet_columns(&schema, &fo, &["price".to_string()]); + assert_eq!(cols, vec![0]); + let aug = load_scoped_page_index(&store, &loc, &fo, &cols).await.unwrap(); + let full = full_index(&bytes); + + // Residual references BOTH columns: price >= 16 (full keeps pages 2,3) AND + // qty <= 115 (full keeps pages 0,1) → full intersection prunes to 0 rows. + let price_ge = pred("price", 0, Operator::GtEq, 16); + let qty_le = pred("qty", 1, Operator::LtEq, 115); + let residual: Arc = + Arc::new(BinaryExpr::new(price_ge, Operator::And, qty_le)); + let pp = build_pruning_predicate(&residual, schema.clone()).unwrap(); + + let s_kept = PagePruner::new(&schema, Arc::clone(&aug)).prune_rg(&pp, 0, None).map(|s| kept(&s)); + let f_kept = PagePruner::new(&schema, full).prune_rg(&pp, 0, None).map(|s| kept(&s)); + // Superset invariant: scoped must keep AT LEAST what full keeps (never + // fewer). It keeps more here (16 vs 0) because it correctly cannot prune + // the placeholdered `qty` — that's safe; the residual mask removes the + // extras post-decode. A scoped result SMALLER than full would be the real + // bug (lost rows). `None` = "kept everything" (no pruning) = the maximal + // superset, also safe. + let s = s_kept.unwrap_or(usize::MAX); + let f = f_kept.unwrap_or(usize::MAX); + assert!( + s >= f, + "scoped page pruning must be a safe superset of full ({} kept) but kept fewer ({})", + f, s + ); + clear_scoped_cache_for_test(); + } + #[tokio::test] async fn scoped_index_reads_non_predicate_projected_column() { let _g = CACHE_TEST_GUARD.lock().unwrap(); @@ -1660,10 +1789,21 @@ mod tests { let aug = load_scoped_page_index_cols(&store, &loc, &fo, &pred_cols, &[2]).await.unwrap(); let o = aug.offset_index().unwrap(); + // wide4 has 256 rows / page-size 32 → real columns have multiple pages. + let full = full_index(&bytes); + let real_pages = full.offset_index().unwrap()[0][0].page_locations().len(); + assert!(real_pages > 1, "fixture should have multi-page columns"); + // Scoped columns (predicate 1 ∪ projection 2 ∪ metric 0) carry the REAL + // page index; the rest carry a single whole-RG placeholder page (non-empty + // so any consumer dereference is safe — never empty, which would panic). for &c in &[0usize, 1, 2] { - assert!(!o[0][c].page_locations().is_empty(), "col {c} (pred/proj/metric) real OI"); + assert_eq!(o[0][c].page_locations().len(), real_pages, "col {c} (pred/proj/metric) real OI"); } - assert!(o[0][3].page_locations().is_empty(), "col 3 OI is empty placeholder"); + assert_eq!( + o[0][3].page_locations().len(), + 1, + "col 3 (scoped out) OI is a single-page placeholder, not real and not empty" + ); } #[tokio::test] diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/page_pruner.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/page_pruner.rs index 19f8ea7037baa..f416f05cb9a33 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/page_pruner.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/page_pruner.rs @@ -47,6 +47,7 @@ use datafusion::common::{Column, ScalarValue}; use datafusion::parquet::arrow::arrow_reader::statistics::StatisticsConverter; use datafusion::parquet::arrow::arrow_reader::{RowSelection, RowSelector}; use datafusion::parquet::file::metadata::ParquetMetaData; +use datafusion::parquet::file::page_index::column_index::ColumnIndexMetaData; #[cfg(test)] use datafusion::physical_expr::expressions::{BinaryExpr, Column as PhysColumn, Literal}; use datafusion::physical_expr::PhysicalExpr; @@ -92,12 +93,29 @@ impl PagePruner { // Early-exit if the file lacks a column/page index; without // both we can't produce page stats for pruning. - self.metadata.column_index()?; + let column_index = self.metadata.column_index()?; + let rg_column_index = column_index.get(rg_idx); let offset_index = self.metadata.offset_index()?; let rg_offsets = offset_index.get(rg_idx)?; let rg_meta = self.metadata.row_groups().get(rg_idx)?; let num_rows = rg_meta.num_rows() as usize; + // Per-file arrow schema (1:1 field<->leaf). Resolving StatisticsConverter + // against the shared/union table schema maps a column to the WRONG leaf + // in a schema-evolved/merged file (SeverityNumber -> union leaf 291, true + // leaf 149), so page stats come from the wrong column and nearly all + // pages are pruned -> the indexed match() AND row loss. Mirror + // eval_leaf / the scoped page-index resolver. + let prune_file_schema = datafusion::parquet::arrow::parquet_to_arrow_schema( + self.metadata.file_metadata().schema_descr(), + self.metadata.file_metadata().key_value_metadata(), + ) + .map(std::sync::Arc::new); + let prune_schema: &SchemaRef = match prune_file_schema.as_ref() { + Ok(s) => s, + Err(_) => &self.schema, + }; + // Build a common page grid from the union of all referenced // columns' page boundaries. Each grid cell is a row range that // falls within a single page of every column. @@ -119,7 +137,7 @@ impl PagePruner { for col in &columns { let converter = match StatisticsConverter::try_new( col.name(), - &self.schema, + prune_schema, self.metadata.file_metadata().schema_descr(), ) { Ok(c) => c, @@ -140,6 +158,24 @@ impl PagePruner { continue; } }; + // Scoped-page-index guard: with a column-scoped page index, only the + // query's scoped columns carry a real `ColumnIndex`; every other column + // is a `NONE` ColumnIndex cell paired with a synthetic one-page + // `OffsetIndex` placeholder. Trusting that placeholder's (absent) stats + // would build a bogus single-page grid for the column and admit rows the + // real index would prune — the indexed `match() AND ` + // over-count. Treat a `NONE` cell exactly like an absent column: push + // `(col, None)` so it contributes "unknown" (never prunes) while other + // columns still prune. This mirrors how DataFusion's listing pruner + // degrades on a column it has no usable stats for. + let has_real_column_index = rg_column_index + .and_then(|rg_ci| rg_ci.get(parquet_col_idx)) + .map(|cell| !matches!(cell, ColumnIndexMetaData::NONE)) + .unwrap_or(false); + if !has_real_column_index { + col_converters.push((col, None)); + continue; + } let col_locs = match rg_offsets.get(parquet_col_idx) { Some(oi) => oi.page_locations(), None => { @@ -177,7 +213,7 @@ impl PagePruner { &boundaries, num_grid_cells, &col_converters, - &self.schema, + prune_schema, &self.metadata, rg_idx, )?; @@ -525,7 +561,30 @@ fn eval_leaf( if columns.is_empty() { return vec![true; num]; } - let arrow_schema = schema.as_ref(); + // Resolve column statistics against THIS FILE's own arrow schema, not the + // shared/union table schema. `StatisticsConverter`/`parquet_column` map a + // column by its position in the supplied arrow schema, then match that + // position to a parquet leaf. The union table schema (e.g. textbench's 410 + // fields) maps a column to the WRONG leaf in a schema-evolved/merged file + // that physically has fewer columns (observed: `SeverityNumber` → leaf 291 + // against the union, but its true leaf in the merged file is 149). Reading + // the wrong leaf yields all-null min/max with `null_count == row_count`, so + // `PruningPredicate` evaluates `null_count != row_count AND …` to false and + // PRUNES the row group — silently dropping every row of RGs that actually + // satisfy the predicate (the indexed `match() AND ` ⅔ row loss: + // ~544 RGs / 557M rows skipped before the collector runs). Deriving the + // arrow schema from the file footer gives a 1:1 field↔leaf correspondence so + // the converter reads the real leaf. Mirrors + // `page_index_loader::resolve_predicate_parquet_columns`. + let file_arrow_schema = datafusion::parquet::arrow::parquet_to_arrow_schema( + metadata.file_metadata().schema_descr(), + metadata.file_metadata().key_value_metadata(), + ) + .map(Arc::new); + let arrow_schema: &datafusion::arrow::datatypes::Schema = match file_arrow_schema.as_ref() { + Ok(s) => s.as_ref(), + Err(_) => schema.as_ref(), + }; let rg_metas: Vec<_> = rg_indices .iter() .filter_map(|&idx| metadata.row_groups().get(idx)) @@ -544,6 +603,17 @@ fn eval_leaf( Ok(c) => c, Err(_) => continue, }; + // Schema-evolution guard: when the column can't be mapped to a parquet + // leaf in THIS file (`parquet_column_index == None`), the converter + // produces all-NULL min/max arrays. Inserting those makes + // `PruningPredicate` treat the RG as "cannot match" (null compares as + // false) and PRUNE it — wrongly dropping every row of a row group whose + // values would satisfy the predicate. This is the indexed + // `match() AND ` silent row loss: on schema-evolved / merged + // files the union-schema column index doesn't resolve to the file's leaf, + // so stats come back null and ~⅔ of row groups are skipped before the + // collector runs. Treat an unresolved column as "no usable stats" — skip + // it so other columns still prune but this RG is conservatively kept. let min_arr = match converter.row_group_mins(rg_metas.iter().copied()) { Ok(arr) => arr, Err(_) => continue, @@ -1001,6 +1071,120 @@ mod tests { assert!(pp.is_none()); } + /// REPRO for the cluster ⅔-drop at the RG level: `eval_leaf` (used by + /// `StatsPruneTree`) must NOT prune a row group whose Int8 `sev` values all + /// satisfy `sev >= 0`. On the cluster, StatsPruneTree skipped 544 RGs + /// (~557M rows) for `match AND sev>=0` — `sev >= 0` is a tautology so it must + /// prune ZERO. If `eval_leaf` mis-handles the Int8-vs-Int32 stats comparison + /// it returns `false` for can-match → the RG is wrongly skipped before the + /// collector runs → the silent ⅔ row loss. + #[test] + fn eval_leaf_int8_ge_zero_keeps_all_row_groups() { + use datafusion::arrow::array::{Int8Array, RecordBatch}; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::physical_expr::expressions::{BinaryExpr, Column as PhysCol, Literal}; + + let schema = Arc::new(Schema::new(vec![Field::new("sev", DataType::Int8, false)])); + // 3 row groups, all values >= 0. + let tmp = NamedTempFile::new().unwrap(); + let props = WriterProperties::builder() + .set_max_row_group_size(8) + .set_statistics_enabled(EnabledStatistics::Chunk) + .build(); + let mut w = + ArrowWriter::try_new(tmp.reopen().unwrap(), schema.clone(), Some(props)).unwrap(); + for base in [0i8, 8, 16] { + let vals: Vec = (base..base + 8).collect(); + let batch = + RecordBatch::try_new(schema.clone(), vec![Arc::new(Int8Array::from(vals))]).unwrap(); + w.write(&batch).unwrap(); + } + w.close().unwrap(); + let meta = ArrowReaderMetadata::load( + &tmp.reopen().unwrap(), + ArrowReaderOptions::new(), + ) + .unwrap(); + let arc_meta = meta.metadata().clone(); + let num_rgs = arc_meta.num_row_groups(); + assert!(num_rgs >= 2, "need multiple RGs, got {num_rgs}"); + let rg_indices: Vec = (0..num_rgs).collect(); + + // Residual `sev >= 0` with an Int32 literal (column is Int8) — as the + // planner emits. Build a PruningPredicate and run eval_leaf on every RG. + let expr: Arc = Arc::new(BinaryExpr::new( + Arc::new(PhysCol::new("sev", 0)), + Operator::GtEq, + Arc::new(Literal::new(ScalarValue::Int32(Some(0)))), + )); + let pp = build_pruning_predicate(&expr, schema.clone()) + .expect("`sev >= 0` should build a PruningPredicate"); + let keep = eval_leaf(&pp, &arc_meta, &schema, &rg_indices); + assert_eq!( + keep, + vec![true; num_rgs], + "every RG has sev>=0 → eval_leaf must keep ALL row groups; a `false` is \ + the Int8 stats mis-prune that drops ~557M rows on the cluster" + ); + } + + /// REPRO for the cluster ⅔-drop: an Int8 (`tinyint`) column with a residual + /// `sev >= 0` whose literal is Int32 (as the planner emits for an int + /// literal). `0` is ≤ every value in the page, so the page must be KEPT. + /// If `prune_rg` mis-handles the Int8-vs-Int32 stats comparison it will + /// wrongly prune (drop) the page → the silent row loss. + #[test] + fn int8_column_ge_int32_literal_keeps_all_pages() { + use datafusion::arrow::array::{Int8Array, RecordBatch}; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + let schema = Arc::new(Schema::new(vec![Field::new("sev", DataType::Int8, false)])); + // 32 rows, values 0..=21 cycling — all >= 0. + let vals: Vec = (0..32).map(|i| (i % 22) as i8).collect(); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int8Array::from(vals))], + ) + .unwrap(); + let tmp = NamedTempFile::new().unwrap(); + let props = WriterProperties::builder() + .set_max_row_group_size(32) + .set_data_page_row_count_limit(8) + .set_write_batch_size(8) + .set_statistics_enabled(EnabledStatistics::Page) + .build(); + let mut w = + ArrowWriter::try_new(tmp.reopen().unwrap(), schema.clone(), Some(props)).unwrap(); + w.write(&batch).unwrap(); + w.close().unwrap(); + let meta = ArrowReaderMetadata::load( + &tmp.reopen().unwrap(), + ArrowReaderOptions::new().with_page_index(true), + ) + .unwrap(); + let pruner = PagePruner::new(&schema, meta.metadata().clone()); + + // `sev >= 0` with an Int32 literal (column is Int8). Tautology → keep all 32. + let expr = bin(col("sev", 0), Operator::GtEq, lit_int(0)); + match build_pruning_predicate(&expr, schema) { + // always_true (no pruning) is also correct — nothing dropped. + None => {} + Some(pp) => { + // `prune_rg` → None means "can't prune → keep whole RG" (safe). + // `Some(sel)` must keep all 32 rows. Either is correct; a Some + // selection that drops rows is the tinyint mis-prune. + match pruner.prune_rg(&pp, 0, None) { + None => {} + Some(sel) => assert_eq!( + count_rows_kept(&sel), + 32, + "Int8 column `sev >= 0` (Int32 literal) must keep every row; \ + a smaller count is the tinyint stats-comparison mis-prune" + ), + } + } + } + } + // ───────────────────────────────────────────────────────────────── // Row-selection shape (adjacent merging, whole-RG, empty). // ───────────────────────────────────────────────────────────────── diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/parquet_bridge.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/parquet_bridge.rs index 72fcb80546dfb..25f968bf45825 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/parquet_bridge.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/parquet_bridge.rs @@ -349,3 +349,195 @@ impl AsyncFileReader for CachedMetadataReader { async move { Ok(metadata) }.boxed() } } + +#[cfg(test)] +mod schema_evolution_tests { + //! Reproduces the indexed-path schema-evolution residual bug: when the + //! physical parquet file's column layout differs from the union/table + //! schema (a column at a different leaf position than its union index), + //! a residual predicate referencing the column by its UNION index must + //! still read the correct leaf. The vanilla listing path does this via + //! per-file schema adaptation; this test checks the indexed bridge does + //! the same (and FAILS today, dropping rows / reading the wrong leaf). + + use super::*; + use datafusion::arrow::array::{Int32Array, RecordBatch, StringArray}; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::parquet::arrow::arrow_reader::{ArrowReaderMetadata, ArrowReaderOptions}; + use datafusion::parquet::arrow::ArrowWriter; + use datafusion::physical_expr::expressions::{BinaryExpr, Column, Literal}; + use datafusion::logical_expr::Operator; + use datafusion::scalar::ScalarValue; + use futures::StreamExt; + use std::sync::Arc; + use tempfile::NamedTempFile; + + /// Write a parquet file whose PHYSICAL schema is `physical_schema`. + fn write_parquet(physical_schema: SchemaRef, columns: Vec>) -> NamedTempFile { + let batch = RecordBatch::try_new(physical_schema.clone(), columns).unwrap(); + let tmp = NamedTempFile::new().unwrap(); + let props = datafusion::parquet::file::properties::WriterProperties::builder() + .set_statistics_enabled(datafusion::parquet::file::properties::EnabledStatistics::Page) + .build(); + let mut w = ArrowWriter::try_new(tmp.reopen().unwrap(), physical_schema, Some(props)).unwrap(); + w.write(&batch).unwrap(); + w.close().unwrap(); + tmp + } + + /// The bug: residual on a numeric column that sits at a DIFFERENT leaf in + /// the physical file than its position in the union/table schema. With the + /// union schema fed to `ParquetSource` and the predicate referencing the + /// union index, reading must still hit the right physical leaf. + /// + /// Physical file columns: `[brand:Utf8, sev:Int32]` (sev = leaf 1) + /// Union/table schema: `[brand:Utf8, inserted:Int32, sev:Int32]` (sev = union index 2) + /// Residual: `sev >= 0` → references union index 2. Every row has sev>=0, + /// so NO row may be dropped. If the path reads union-leaf 2 against a file + /// with only 2 leaves (0,1), it mis-reads / nulls → drops rows. + #[tokio::test] + async fn residual_on_shifted_leaf_reads_correct_column() { + // Physical schema (what's actually in the file): brand, sev. + let physical_schema: SchemaRef = Arc::new(Schema::new(vec![ + Field::new("brand", DataType::Utf8, true), + Field::new("sev", DataType::Int32, true), + ])); + let brands = StringArray::from(vec!["a", "b", "c", "d"]); + let sevs = Int32Array::from(vec![0, 17, 9, 21]); // all >= 0 + let tmp = write_parquet( + physical_schema.clone(), + vec![Arc::new(brands), Arc::new(sevs)], + ); + let path = tmp.path().to_path_buf(); + let size = std::fs::metadata(&path).unwrap().len(); + + let file = std::fs::File::open(&path).unwrap(); + let meta = + ArrowReaderMetadata::load(&file, ArrowReaderOptions::new().with_page_index(true)).unwrap(); + let parquet_meta = meta.metadata().clone(); + + // Union/table schema: an extra column `inserted` BEFORE sev, so sev is + // at union index 2 (but physical leaf 1). This simulates schema + // evolution / merged files where the column position shifted. + let union_schema: SchemaRef = Arc::new(Schema::new(vec![ + Field::new("brand", DataType::Utf8, true), + Field::new("inserted", DataType::Int32, true), + Field::new("sev", DataType::Int32, true), + ])); + + // Residual `sev >= 0`, with the Column referencing the UNION index (2), + // exactly as the indexed path builds it (bound to the union DFSchema). + let residual: Arc = Arc::new(BinaryExpr::new( + Arc::new(Column::new("sev", 2)), + Operator::GtEq, + Arc::new(Literal::new(ScalarValue::Int32(Some(0)))), + )); + + // Projection: read brand (0) and sev (2) by union index. + let projection = vec![0usize, 2usize]; + + let store: Arc = Arc::new(object_store::local::LocalFileSystem::new()); + let store_url = ObjectStoreUrl::local_filesystem(); + let config = RowGroupStreamConfig { + file_path: path.to_string_lossy().to_string(), + file_size: size, + store, + store_url, + full_schema: union_schema, + metadata: Arc::clone(&parquet_meta), + projection: Some(projection), + predicate: Some(residual), + io_stats: Arc::new(ReadIoStats::default()), + }; + + // Row-granular pushdown ON (push_predicate=true) — the cluster's + // low-selectivity regime that pushes the residual into parquet decode. + let (mut stream, _exec) = + create_row_selection_stream(&config, 0, full_rg_selection(&parquet_meta, 0), true) + .expect("stream builds"); + + let mut total_rows = 0usize; + while let Some(batch) = stream.next().await { + let b = batch.expect("batch decodes without error"); + total_rows += b.num_rows(); + } + + assert_eq!( + total_rows, 4, + "all 4 rows have sev>=0 and must survive the residual; got {total_rows} \ + (indexed path read the wrong leaf for the shifted column)" + ); + } + + /// Variant: the residual column is ABSENT from the physical file entirely + /// (pure schema evolution — column added after this file was written). The + /// listing path null-fills it; `sev >= 0` over null → null → row excluded, + /// which is CORRECT. But `sev IS NULL` must keep all rows. This pins whether + /// the indexed path null-fills absent columns the same way the listing path + /// does, rather than erroring or mis-reading. + #[tokio::test] + async fn residual_on_absent_column_null_fills() { + // Physical file has ONLY brand — no sev column at all. + let physical_schema: SchemaRef = Arc::new(Schema::new(vec![ + Field::new("brand", DataType::Utf8, true), + ])); + let brands = StringArray::from(vec!["a", "b", "c", "d"]); + let tmp = write_parquet(physical_schema.clone(), vec![Arc::new(brands)]); + let path = tmp.path().to_path_buf(); + let size = std::fs::metadata(&path).unwrap().len(); + + let file = std::fs::File::open(&path).unwrap(); + let meta = + ArrowReaderMetadata::load(&file, ArrowReaderOptions::new().with_page_index(true)).unwrap(); + let parquet_meta = meta.metadata().clone(); + + // Union schema adds `sev` at index 1 — absent from the physical file. + let union_schema: SchemaRef = Arc::new(Schema::new(vec![ + Field::new("brand", DataType::Utf8, true), + Field::new("sev", DataType::Int32, true), + ])); + + // Residual `sev IS NOT NULL` — over a null-filled absent column this is + // FALSE for all rows → 0 rows. That's correct. The point of the test is + // it must NOT error and must agree with listing-path null semantics. + // We assert the dual: `sev IS NULL` keeps all 4 rows. + use datafusion::physical_expr::expressions::IsNullExpr; + let residual: Arc = + Arc::new(IsNullExpr::new(Arc::new(Column::new("sev", 1)))); + + let projection = vec![0usize, 1usize]; + let store: Arc = Arc::new(object_store::local::LocalFileSystem::new()); + let store_url = ObjectStoreUrl::local_filesystem(); + let config = RowGroupStreamConfig { + file_path: path.to_string_lossy().to_string(), + file_size: size, + store, + store_url, + full_schema: union_schema, + metadata: Arc::clone(&parquet_meta), + projection: Some(projection), + predicate: Some(residual), + io_stats: Arc::new(ReadIoStats::default()), + }; + + let (mut stream, _exec) = + create_row_selection_stream(&config, 0, full_rg_selection(&parquet_meta, 0), true) + .expect("stream builds"); + let mut total_rows = 0usize; + while let Some(batch) = stream.next().await { + let b = batch.expect("batch decodes without error"); + total_rows += b.num_rows(); + } + assert_eq!( + total_rows, 4, + "sev IS NULL over an absent (null-filled) column must keep all 4 rows; got {total_rows}" + ); + } + + /// Build a RowSelection that selects every row of `rg_index`. + fn full_rg_selection(meta: &ParquetMetaData, rg_index: usize) -> RowSelection { + use datafusion::parquet::arrow::arrow_reader::RowSelector; + let n = meta.row_group(rg_index).num_rows() as usize; + RowSelection::from(vec![RowSelector::select(n)]) + } +} diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/mod.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/mod.rs index 662b71accee58..3612241c9aff2 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/mod.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/mod.rs @@ -335,6 +335,161 @@ async fn run_tree_and_plan( (rows, plan) } +/// Full-path run with a **column-scoped** page index grafted onto the segment +/// (mimics the indexed scan after `load_scoped_page_index_cols`), so we can prove +/// the end-to-end result (page-prune → decode → mask → rows) matches the +/// full-index path even when the page index is scoped to only `scope_cols` and +/// the boolean tree references a column that got the one-page placeholder + +/// `NONE` ColumnIndex. Uses the same `TreeBitsetSource` evaluator as +/// [`run_tree_and_plan`]; the only difference is the grafted scoped metadata. +/// +/// `scope_cols` = parquet leaf indices kept REAL in the scoped ColumnIndex +/// (everything else becomes a placeholder). Returns the SELECTed rows, sorted. +async fn run_tree_with_scoped_index( + tree: BoolNode, + scope_cols: &[usize], +) -> Vec<(String, i32, String, String)> { + let tmp = write_fixture_parquet(); + let path = tmp.path().to_path_buf(); + let size = std::fs::metadata(&path).unwrap().len(); + let file = std::fs::File::open(&path).unwrap(); + let meta = + ArrowReaderMetadata::load(&file, ArrowReaderOptions::new().with_page_index(false)).unwrap(); + let schema = meta.schema().clone(); + let footer = meta.metadata().clone(); + + // Graft a page index scoped to `scope_cols` (offset scoped likewise) — the + // exact artifact the indexed path builds. Non-scoped columns get NONE + // ColumnIndex + one-page placeholder OffsetIndex. + let store_local: Arc = + Arc::new(object_store::local::LocalFileSystem::new()); + let obj_loc = object_store::path::Path::from(path.to_string_lossy().as_ref()); + // ColumnIndex is scoped to `scope_cols` (the predicate columns); the + // OffsetIndex must cover every column the query READS (its projection), + // exactly as the indexed executor does via `collect_plan_column_names`. + // Passing `scope_cols` for the OffsetIndex too would leave projected-but- + // unscoped columns on the one-page placeholder and the reader would deref a + // fake (0,0) byte range. Scope the OffsetIndex to ALL columns here so the + // test mirrors production (ColumnIndex scoped, OffsetIndex covers reads). + let num_cols = footer.file_metadata().schema_descr().num_columns(); + let offset_cols: Vec = (0..num_cols).collect(); + let scoped_meta = crate::indexed_table::page_index_loader::load_scoped_page_index_cols( + &store_local, + &obj_loc, + &footer, + scope_cols, + &offset_cols, + ) + .await + .expect("scoped page index must build"); + + let mut rgs = Vec::new(); + let mut offset = 0i64; + for i in 0..scoped_meta.num_row_groups() { + let n = scoped_meta.row_group(i).num_rows(); + rgs.push(RowGroupInfo { index: i, first_row: offset, num_rows: n }); + offset += n; + } + let object_path = object_store::path::Path::from(path.to_string_lossy().as_ref()); + let segment = SegmentFileInfo { + writer_generation: 0, + max_doc: 16, + object_path, + parquet_size: size, + row_groups: rgs, + metadata: Arc::clone(&scoped_meta), + global_base: 0, + sort_min: None, + sort_max: None, + }; + + let tree = tree.push_not_down(); + let collectors = wire_collectors(&tree); + let per_leaf: Vec<(i32, Arc)> = collectors + .into_iter() + .enumerate() + .map(|(i, c)| (i as i32, c)) + .collect(); + let tree = Arc::new(tree); + let factory: super::table_provider::EvaluatorFactory = { + let per_leaf = per_leaf.clone(); + let tree = Arc::clone(&tree); + let schema = schema.clone(); + Arc::new(move |segment, _chunk, _stream_metrics, _stats_prune_tree| { + let resolved = tree.resolve(&per_leaf)?; + let pruner = Arc::new(PagePruner::new(&schema, Arc::clone(&segment.metadata))); + let eval: Arc = Arc::new(TreeBitsetSource { + tree: Arc::new(resolved), + evaluator: Arc::new(BitmapTreeEvaluator), + leaves: Arc::new(CollectorLeafBitmaps { + ffm_collector_calls: _stream_metrics.ffm_collector_calls.clone(), + }), + page_pruner: pruner, + cost_predicate: 1, + cost_collector: 10, + max_collector_parallelism: 1, + pruning_predicates: std::sync::Arc::new(std::collections::HashMap::new()), + page_prune_metrics: Some( + crate::indexed_table::page_pruner::PagePruneMetrics::from_stream_metrics( + _stream_metrics, + ), + ), + collector_strategy: + crate::indexed_table::eval::CollectorCallStrategy::TightenOuterBounds, + stats_prune_tree: None, + }); + Ok(eval) + }) + }; + + let store: Arc = + Arc::new(object_store::local::LocalFileSystem::new()); + let store_url = datafusion::execution::object_store::ObjectStoreUrl::local_filesystem(); + let qc = crate::datafusion_query_config::DatafusionQueryConfig::builder() + .target_partitions(1) + .force_strategy(Some(FilterStrategy::BooleanMask)) + .force_pushdown(Some(false)) + .build(); + let provider = Arc::new(IndexedTableProvider::new(IndexedTableConfig { + schema: schema.clone(), + segments: vec![segment], + store, + store_url, + evaluator_factory: factory, + pushdown_predicate: None, + query_config: std::sync::Arc::new(qc), + predicate_columns: vec![], + emit_row_ids: false, + prune_tree_config: None, + sort_fields: vec![], + sort_orders: vec![], + })); + let ctx = SessionContext::new(); + ctx.register_table("t", provider).unwrap(); + let df = ctx.sql("SELECT brand, price, status, category FROM t").await.unwrap(); + let plan = df.create_physical_plan().await.unwrap(); + let mut stream = + datafusion::physical_plan::execute_stream(plan, ctx.task_ctx()).unwrap(); + let mut rows: Vec<(String, i32, String, String)> = Vec::new(); + while let Some(batch) = stream.next().await { + let b = batch.unwrap(); + let brand = b.column(0).as_any().downcast_ref::().unwrap(); + let price = b.column(1).as_any().downcast_ref::().unwrap(); + let status = b.column(2).as_any().downcast_ref::().unwrap(); + let cat = b.column(3).as_any().downcast_ref::().unwrap(); + for i in 0..b.num_rows() { + rows.push(( + brand.value(i).to_string(), + price.value(i), + status.value(i).to_string(), + cat.value(i).to_string(), + )); + } + } + rows.sort(); + rows +} + // ── Tree-building helpers ────────────────────────────────────────── // // Test-only leaf encoding: `index_leaf(tag)` puts a single-byte tag into @@ -405,3 +560,44 @@ fn wire(node: &BoolNode, out: &mut Vec>) { BoolNode::DelegationPossible { .. } => {} } } + +// ── Full-path scoped-page-index correctness (task #6 regression) ────────────── +// +// The indexed path scopes the page index to the predicate columns; other columns +// become a one-page placeholder + NONE ColumnIndex. These tests drive the FULL +// path (page-prune → decode → mask → rows) over a scoped index and assert the +// rows are IDENTICAL to the full-index path — the real correctness contract +// (page-prune conservatism alone is harmless; what matters is the final rows). + +/// Tree `Collector(amazon) AND price>=100 AND status='archived'`, page index +/// scoped to `price` only (col 1). `status` (col 2) is therefore placeholdered. +/// The full-path result MUST equal the full-index path. (brand col 0 also +/// placeholdered, but the Collector — not the page index — handles brand.) +#[tokio::test] +async fn scoped_index_fullpath_matches_full_index_residual_on_placeholdered_col() { + let _g = crate::indexed_table::page_index_loader::SCOPED_CACHE_TEST_GUARD.lock().unwrap(); + crate::indexed_table::page_index_loader::clear_scoped_cache_for_test(); + + let make_tree = || { + BoolNode::And(vec![ + index_leaf(0), // Collector: brand == amazon + pred_int("price", Operator::GtEq, 100), // residual on price (scoped) + pred_str("status", Operator::Eq, "archived"), // residual on status (placeholdered) + ]) + }; + + // Full index (every column real) — the reference. + let full_rows = run_tree(make_tree()).await; + // Scoped to `price` only (col 1) — `status` is placeholdered. + let scoped_rows = run_tree_with_scoped_index(make_tree(), &[1]).await; + + // Ground truth: amazon rows with price>=100 AND status=archived → row 1 + // (amazon,150,archived) only. (row 12 amazon,30,archived fails price>=100.) + let mut expected = full_rows.clone(); + expected.sort(); + assert_eq!( + scoped_rows, expected, + "scoped-page-index full-path rows must equal full-index rows; a mismatch is the \ + indexed residual-on-placeholdered-column over-count (task #6)" + ); +} diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/schema_drift.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/schema_drift.rs index 665f5f0e878f4..39f1679fa6fb0 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/schema_drift.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/schema_drift.rs @@ -539,3 +539,278 @@ async fn utf8view_schema_with_utf8_file_does_not_panic() { .await; assert_eq!(count, 4); } + +// ══════════════════════════════════════════════════════════════════════ +// SingleCollector path + schema-evolved file + numeric residual. +// +// Reproduces the cluster bug: `match(Body,'connection') AND SeverityNumber>=0 +// | stats count()` drops ~2/3 of rows on schema-evolved/merged files, but the +// vanilla listing path (`LIKE ... AND sev>=0`) keeps all of them. +// +// The distinguishing ingredients vs the passing `reorder_columns_aligned_and_filtered` +// test above: +// - runs through `SingleCollectorEvaluator` (not `TreeBitsetSource`) +// - sets `pushdown_predicate: Some(residual)` + `predicate_columns` (union indices) +// - the file's PHYSICAL column order differs from the table/union schema, so +// a residual `Column` bound to a UNION index points at a different physical +// leaf in the file. +// +// The collector matches ALL rows; the residual `sev >= 0` is a tautology (every +// sev is non-negative). So every row MUST survive. If the indexed path reads the +// residual column from the wrong leaf (union-index, no per-file name remap), it +// sees null/garbage → `>= 0` is null/false → rows dropped. +// ══════════════════════════════════════════════════════════════════════ + +use crate::indexed_table::eval::single_collector::{ + CollectorCallStrategy, FfmDelegatedBackendCollectorFactory, SingleCollectorEvaluator, +}; + +/// Collector that matches every doc in `[min_doc, max_doc)`. +#[derive(Debug)] +struct MatchAllCollector; +impl RowGroupDocsCollector for MatchAllCollector { + fn collect_packed_u64_bitset(&self, min_doc: i32, max_doc: i32) -> Result, String> { + let span = (max_doc - min_doc).max(0) as usize; + let mut out = vec![0u64; span.div_ceil(64)]; + for rel in 0..span { + out[rel / 64] |= 1u64 << (rel % 64); + } + Ok(out) + } +} + +/// Run a single-collector query (collector matches all rows) with a residual +/// ` >= 0` referencing `residual_union_idx` in `table_schema`, +/// over a file whose physical schema is `file_schema`. Returns the row count. +async fn run_single_collector_residual( + file_schema: SchemaRef, + table_schema: SchemaRef, + batch: &RecordBatch, + residual_col: &str, + residual_union_idx: usize, +) -> usize { + use datafusion::physical_expr::expressions::{BinaryExpr, Column, Literal}; + // Residual ` >= 0`, Column bound to its UNION index (exactly + // how the indexed path builds it from the substrait/union DFSchema), with a + // RAW (uncoerced) Int32 literal — mirrors a plan that didn't coerce. + let residual: Arc = Arc::new(BinaryExpr::new( + Arc::new(Column::new(residual_col, residual_union_idx)), + Operator::GtEq, + Arc::new(Literal::new(ScalarValue::Int32(Some(0)))), + )); + run_single_collector_residual_with_expr( + file_schema, + table_schema, + batch, + residual, + vec![residual_union_idx], + ) + .await +} + +/// Same as `run_single_collector_residual` but takes a pre-built residual +/// PhysicalExpr (so tests can supply a production-coerced expr). +async fn run_single_collector_residual_with_expr( + file_schema: SchemaRef, + table_schema: SchemaRef, + batch: &RecordBatch, + residual: Arc, + predicate_columns: Vec, +) -> usize { + use std::collections::HashMap; + use std::sync::OnceLock; + + let tmp = NamedTempFile::new().unwrap(); + let (file, path) = tmp.keep().unwrap(); + let props = datafusion::parquet::file::properties::WriterProperties::builder() + .set_max_row_group_size(256) + .set_statistics_enabled(datafusion::parquet::file::properties::EnabledStatistics::Page) + .build(); + let mut w = ArrowWriter::try_new(file, file_schema, Some(props)).unwrap(); + w.write(batch).unwrap(); + w.close().unwrap(); + + let size = std::fs::metadata(&path).unwrap().len(); + let rfile = std::fs::File::open(&path).unwrap(); + let meta = + ArrowReaderMetadata::load(&rfile, ArrowReaderOptions::new().with_page_index(true)).unwrap(); + let parquet_meta = meta.metadata().clone(); + let mut rgs = Vec::new(); + let mut offset = 0i64; + for i in 0..parquet_meta.num_row_groups() { + let n = parquet_meta.row_group(i).num_rows(); + rgs.push(RowGroupInfo { index: i, first_row: offset, num_rows: n }); + offset += n; + } + let segment = SegmentFileInfo { + writer_generation: 0, + max_doc: batch.num_rows() as i64, + object_path: object_store::path::Path::from(path.to_string_lossy().as_ref()), + parquet_size: size, + row_groups: rgs, + metadata: Arc::clone(&parquet_meta), + global_base: 0, + sort_min: None, + sort_max: None, + }; + + let eval_factory: super::super::table_provider::EvaluatorFactory = { + let residual = Arc::clone(&residual); + let schema = table_schema.clone(); + Arc::new(move |segment, _chunk, stream_metrics, _stats_prune_tree| { + let pruner = Arc::new(PagePruner::new(&schema, Arc::clone(&segment.metadata))); + let eval: Arc = Arc::new(SingleCollectorEvaluator::new( + Some(Arc::new(MatchAllCollector) as Arc), + pruner, + None, + Some(Arc::clone(&residual)), + None, + stream_metrics.ffm_collector_calls.clone(), + CollectorCallStrategy::FullRange, + Arc::new(HashMap::>>::new()), + segment.writer_generation, + Arc::new(FfmDelegatedBackendCollectorFactory), + 0, + None, + None, + )); + Ok(eval) + }) + }; + + let qc = crate::datafusion_query_config::DatafusionQueryConfig::builder() + .target_partitions(1) + .force_pushdown(Some(true)) + // Force row-granular (min_skip_run=1 via a selectivity threshold of 1.0) + // so the residual is pushed into parquet's RowFilter (`with_predicate`) — + // the cluster's regime for a non-selective collector. + .min_skip_run_default(1) + .min_skip_run_selectivity_threshold(1.0) + .build(); + let provider = Arc::new(IndexedTableProvider::new(IndexedTableConfig { + schema: table_schema, + segments: vec![segment], + store: Arc::new(object_store::local::LocalFileSystem::new()) + as Arc, + store_url: datafusion::execution::object_store::ObjectStoreUrl::local_filesystem(), + evaluator_factory: eval_factory, + pushdown_predicate: Some(residual), + query_config: std::sync::Arc::new(qc), + predicate_columns, + emit_row_ids: false, + prune_tree_config: None, + sort_fields: vec![], + sort_orders: vec![], + })); + let ctx = SessionContext::new(); + ctx.register_table("t", provider).unwrap(); + let df = ctx.sql("SELECT * FROM t").await.unwrap(); + let mut stream = df.execute_stream().await.unwrap(); + let mut count = 0; + while let Some(b) = stream.next().await { + count += b.unwrap().num_rows(); + } + count +} + +/// Production-faithful Int8 residual: build `sev >= 0` through +/// `SessionState::create_physical_expr` (which runs the TypeCoercion analyzer, +/// exactly like `substrait_to_tree::convert_expr` does on the cluster), so the +/// resulting expr has the CAST DataFusion would insert. Then run it through the +/// SingleCollector path over an Int8 column. This distinguishes: +/// - coercion fixes it (test passes) → the fix is "always coerce the residual" +/// - still drops/errors (test fails) → the cast doesn't survive evaluate_residual +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn single_collector_residual_int8_coerced_keeps_all() { + use datafusion::arrow::array::Int8Array; + use datafusion::execution::context::SessionContext; + use datafusion::common::DFSchema; + use datafusion::prelude::{col, lit}; + + let schema = Arc::new(Schema::new(vec![ + Field::new("brand", DataType::Utf8, false), + Field::new("sev", DataType::Int8, false), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(StringArray::from(vec!["a", "b", "c", "d"])), + Arc::new(Int8Array::from(vec![0i8, 17, 9, 21])), + ], + ) + .unwrap(); + + // Build the coerced physical residual `sev >= 0` the production way. + let ctx = SessionContext::new(); + let df_schema = DFSchema::try_from(schema.as_ref().clone()).unwrap(); + let logical = col("sev").gt_eq(lit(0i32)); + let residual = ctx + .state() + .create_physical_expr(logical, &df_schema) + .expect("create_physical_expr (with TypeCoercion)"); + + let count = run_single_collector_residual_with_expr( + schema.clone(), + schema, + &batch, + residual, + vec![1], + ) + .await; + assert_eq!( + count, 4, + "coerced Int8 residual sev>=0 must keep all rows; got {count}" + ); +} + +/// Aligned baseline: file order == table order. Must keep all rows (sanity). +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn single_collector_residual_aligned_keeps_all() { + let schema = Arc::new(Schema::new(vec![ + Field::new("brand", DataType::Utf8, false), + Field::new("sev", DataType::Int32, false), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(StringArray::from(vec!["a", "b", "c", "d"])), + Arc::new(Int32Array::from(vec![0, 17, 9, 21])), + ], + ) + .unwrap(); + // sev is at union index 1 and physical leaf 1 — aligned. + let count = run_single_collector_residual(schema.clone(), schema, &batch, "sev", 1).await; + assert_eq!(count, 4, "aligned: all rows have sev>=0, none may drop"); +} + +/// THE BUG: file physical order != table/union order, so the residual's union +/// index points at the wrong physical leaf. Collector matches all; `sev >= 0` +/// is a tautology → all 4 rows must survive. Reproduces the cluster ⅔ drop. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn single_collector_residual_evolved_file_keeps_all() { + // Physical file order: [sev, brand] (sev = physical leaf 0) + let file_schema = Arc::new(Schema::new(vec![ + Field::new("sev", DataType::Int32, false), + Field::new("brand", DataType::Utf8, false), + ])); + // Table/union order: [brand, sev] (sev = union index 1, but physical leaf 0) + let table_schema = Arc::new(Schema::new(vec![ + Field::new("brand", DataType::Utf8, false), + Field::new("sev", DataType::Int32, false), + ])); + let batch = RecordBatch::try_new( + file_schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![0, 17, 9, 21])), + Arc::new(StringArray::from(vec!["a", "b", "c", "d"])), + ], + ) + .unwrap(); + // Residual references sev at UNION index 1; physical leaf is 0. + let count = run_single_collector_residual(file_schema, table_schema, &batch, "sev", 1).await; + assert_eq!( + count, 4, + "evolved file: all 4 rows have sev>=0 and must survive; a smaller count \ + means the residual read the wrong physical leaf (union-index, no per-file remap)" + ); +} From 76db1adef5030fd18ae2a5633112f21f7400cd4b Mon Sep 17 00:00:00 2001 From: G Date: Thu, 18 Jun 2026 19:32:51 +0530 Subject: [PATCH 17/18] metadata clone fixes Signed-off-by: G --- .../libs/dataformat-native/rust/Cargo.lock | 45 ++++------- .../libs/dataformat-native/rust/Cargo.toml | 25 ++++++ .../rust/src/indexed_executor.rs | 24 +++--- .../src/indexed_table/page_index_loader.rs | 77 +++++++++++++++++-- .../rust/src/shard_scoped_reader.rs | 25 +++--- 5 files changed, 131 insertions(+), 65 deletions(-) diff --git a/sandbox/libs/dataformat-native/rust/Cargo.lock b/sandbox/libs/dataformat-native/rust/Cargo.lock index 78827944527d7..86db668e38362 100644 --- a/sandbox/libs/dataformat-native/rust/Cargo.lock +++ b/sandbox/libs/dataformat-native/rust/Cargo.lock @@ -103,8 +103,7 @@ checksum = "7c02d123df017efcdfbd739ef81735b36c5ba83ec3c59c80a9d7ecc718f92e50" [[package]] name = "arrow" version = "58.3.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "378530e55cd479eda3c14eb345310799717e6f76d0c332041e8487022166b471" +source = "git+https://github.com/bharath-techie/arrow-rs.git?branch=opensearch-fixes#2755989eb982ee43966deb256bd8a48f76a7c76f" dependencies = [ "arrow-arith", "arrow-array", @@ -124,8 +123,7 @@ dependencies = [ [[package]] name = "arrow-arith" version = "58.3.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a0ab212d2c1886e802f51c5212d78ebbcbb0bec980fff9dadc1eb8d45cd0b738" +source = "git+https://github.com/bharath-techie/arrow-rs.git?branch=opensearch-fixes#2755989eb982ee43966deb256bd8a48f76a7c76f" dependencies = [ "arrow-array", "arrow-buffer", @@ -138,8 +136,7 @@ dependencies = [ [[package]] name = "arrow-array" version = "58.3.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cfd33d3e92f207444098c75b42de99d329562be0cf686b307b097cc52b4e999e" +source = "git+https://github.com/bharath-techie/arrow-rs.git?branch=opensearch-fixes#2755989eb982ee43966deb256bd8a48f76a7c76f" dependencies = [ "ahash", "arrow-buffer", @@ -157,8 +154,7 @@ dependencies = [ [[package]] name = "arrow-buffer" version = "58.3.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0c6cd424c2693bcdbc150d843dc9d4d137dd2de4782ce6df491ad11a3a0416c0" +source = "git+https://github.com/bharath-techie/arrow-rs.git?branch=opensearch-fixes#2755989eb982ee43966deb256bd8a48f76a7c76f" dependencies = [ "bytes", "half", @@ -169,8 +165,7 @@ dependencies = [ [[package]] name = "arrow-cast" version = "58.3.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4c5aefb56a2c02e9e2b30746241058b85f8983f0fcff2ba0c6d09006e1cded7f" +source = "git+https://github.com/bharath-techie/arrow-rs.git?branch=opensearch-fixes#2755989eb982ee43966deb256bd8a48f76a7c76f" dependencies = [ "arrow-array", "arrow-buffer", @@ -191,8 +186,7 @@ dependencies = [ [[package]] name = "arrow-csv" version = "58.3.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e94e8cf7e517657a52b91ea1263acf38c4ca62a84655d72458a3359b12ab97de" +source = "git+https://github.com/bharath-techie/arrow-rs.git?branch=opensearch-fixes#2755989eb982ee43966deb256bd8a48f76a7c76f" dependencies = [ "arrow-array", "arrow-cast", @@ -206,8 +200,7 @@ dependencies = [ [[package]] name = "arrow-data" version = "58.3.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3c88210023a2bfee1896af366309a3028fc3bcbd6515fa29a7990ee1baa08ee0" +source = "git+https://github.com/bharath-techie/arrow-rs.git?branch=opensearch-fixes#2755989eb982ee43966deb256bd8a48f76a7c76f" dependencies = [ "arrow-buffer", "arrow-schema", @@ -219,8 +212,7 @@ dependencies = [ [[package]] name = "arrow-ipc" version = "58.3.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "238438f0834483703d88896db6fe5a7138b2230debc31b34c0336c2996e3c64f" +source = "git+https://github.com/bharath-techie/arrow-rs.git?branch=opensearch-fixes#2755989eb982ee43966deb256bd8a48f76a7c76f" dependencies = [ "arrow-array", "arrow-buffer", @@ -235,8 +227,7 @@ dependencies = [ [[package]] name = "arrow-json" version = "58.3.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "205ca2119e6d679d5c133c6f30e68f027738d95ed948cf77677ea69c7800036b" +source = "git+https://github.com/bharath-techie/arrow-rs.git?branch=opensearch-fixes#2755989eb982ee43966deb256bd8a48f76a7c76f" dependencies = [ "arrow-array", "arrow-buffer", @@ -260,8 +251,7 @@ dependencies = [ [[package]] name = "arrow-ord" version = "58.3.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1bffd8fd2579286a5d63bac898159873e5094a79009940bcb42bbfce4f19f1d0" +source = "git+https://github.com/bharath-techie/arrow-rs.git?branch=opensearch-fixes#2755989eb982ee43966deb256bd8a48f76a7c76f" dependencies = [ "arrow-array", "arrow-buffer", @@ -273,8 +263,7 @@ dependencies = [ [[package]] name = "arrow-row" version = "58.3.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bab5994731204603c73ba69267616c50f80780774c6bb0476f1f830625115e0c" +source = "git+https://github.com/bharath-techie/arrow-rs.git?branch=opensearch-fixes#2755989eb982ee43966deb256bd8a48f76a7c76f" dependencies = [ "arrow-array", "arrow-buffer", @@ -286,8 +275,7 @@ dependencies = [ [[package]] name = "arrow-schema" version = "58.3.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f633dbfdf39c039ada1bf9e34c694816eb71fbb7dc78f613993b7245e078a1ed" +source = "git+https://github.com/bharath-techie/arrow-rs.git?branch=opensearch-fixes#2755989eb982ee43966deb256bd8a48f76a7c76f" dependencies = [ "bitflags", "serde_core", @@ -297,8 +285,7 @@ dependencies = [ [[package]] name = "arrow-select" version = "58.3.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8cd065c54172ac787cf3f2f8d4107e0d3fdc26edba76fdf4f4cc170258942222" +source = "git+https://github.com/bharath-techie/arrow-rs.git?branch=opensearch-fixes#2755989eb982ee43966deb256bd8a48f76a7c76f" dependencies = [ "ahash", "arrow-array", @@ -311,8 +298,7 @@ dependencies = [ [[package]] name = "arrow-string" version = "58.3.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "29dd7cda3ab9692f43a2e4acc444d760cc17b12bb6d8232ddf64e9bab7c06b42" +source = "git+https://github.com/bharath-techie/arrow-rs.git?branch=opensearch-fixes#2755989eb982ee43966deb256bd8a48f76a7c76f" dependencies = [ "arrow-array", "arrow-buffer", @@ -3012,8 +2998,7 @@ dependencies = [ [[package]] name = "parquet" version = "58.3.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5dafa7d01085b62a47dd0c1829550a0a36710ea9c4fe358a05a85477cec8a908" +source = "git+https://github.com/bharath-techie/arrow-rs.git?branch=opensearch-fixes#2755989eb982ee43966deb256bd8a48f76a7c76f" dependencies = [ "ahash", "arrow-array", diff --git a/sandbox/libs/dataformat-native/rust/Cargo.toml b/sandbox/libs/dataformat-native/rust/Cargo.toml index d632a5c550f4d..51c096575966d 100644 --- a/sandbox/libs/dataformat-native/rust/Cargo.toml +++ b/sandbox/libs/dataformat-native/rust/Cargo.toml @@ -96,3 +96,28 @@ codegen-units = 16 incremental = true debug = "full" strip = false + +# Redirect arrow-rs crates to the OpenSearch fork, which stores +# ParquetMetaData.row_groups behind an Arc so cloning the footer (to graft a +# scoped page index) is a refcount bump instead of a deep copy of every +# RowGroupMetaData/ColumnChunkMetaData. On wide, many-row-group footers +# (textbench: ~403 cols × ~64 RGs) that deep clone was ~60ms/query. +# Patch ALL arrow crates in the graph so a single arrow version is used +# (mixing patched + crates.io arrow → type mismatches). +# Fork: https://github.com/bharath-techie/arrow-rs branch: opensearch-fixes +[patch.crates-io] +arrow = { git = "https://github.com/bharath-techie/arrow-rs.git", branch = "opensearch-fixes" } +arrow-arith = { git = "https://github.com/bharath-techie/arrow-rs.git", branch = "opensearch-fixes" } +arrow-array = { git = "https://github.com/bharath-techie/arrow-rs.git", branch = "opensearch-fixes" } +arrow-buffer = { git = "https://github.com/bharath-techie/arrow-rs.git", branch = "opensearch-fixes" } +arrow-cast = { git = "https://github.com/bharath-techie/arrow-rs.git", branch = "opensearch-fixes" } +arrow-csv = { git = "https://github.com/bharath-techie/arrow-rs.git", branch = "opensearch-fixes" } +arrow-data = { git = "https://github.com/bharath-techie/arrow-rs.git", branch = "opensearch-fixes" } +arrow-ipc = { git = "https://github.com/bharath-techie/arrow-rs.git", branch = "opensearch-fixes" } +arrow-json = { git = "https://github.com/bharath-techie/arrow-rs.git", branch = "opensearch-fixes" } +arrow-ord = { git = "https://github.com/bharath-techie/arrow-rs.git", branch = "opensearch-fixes" } +arrow-row = { git = "https://github.com/bharath-techie/arrow-rs.git", branch = "opensearch-fixes" } +arrow-schema = { git = "https://github.com/bharath-techie/arrow-rs.git", branch = "opensearch-fixes" } +arrow-select = { git = "https://github.com/bharath-techie/arrow-rs.git", branch = "opensearch-fixes" } +arrow-string = { git = "https://github.com/bharath-techie/arrow-rs.git", branch = "opensearch-fixes" } +parquet = { git = "https://github.com/bharath-techie/arrow-rs.git", branch = "opensearch-fixes" } diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_executor.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_executor.rs index 9459284e8e93b..5d50a5ab21e2f 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_executor.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_executor.rs @@ -1042,22 +1042,20 @@ async unsafe fn execute_indexed_with_context_inner( projection_column_names, ); for segment in segments.iter_mut() { - let parquet_cols = crate::indexed_table::page_index_loader::resolve_predicate_parquet_columns( - &schema, - &segment.metadata, - &predicate_column_names, - ); + // Resolve predicate + projection leaves in ONE pass — both share the + // per-file arrow schema derivation, which is the dominant cost on wide + // schemas. (offset_cols = projection; predicate ∪ {0} is unioned inside + // the loader.) + let (parquet_cols, offset_cols) = + crate::indexed_table::page_index_loader::resolve_predicate_parquet_columns_pair( + &schema, + &segment.metadata, + &predicate_column_names, + &projection_column_names, + ); if parquet_cols.is_empty() { continue; } - // Projection (∪ predicate ∪ {0}) leaf indices for THIS file, for the - // OffsetIndex column scoping. Resolved the same way as predicate cols - // so both scan paths agree on the cache key. - let offset_cols = crate::indexed_table::page_index_loader::resolve_predicate_parquet_columns( - &schema, - &segment.metadata, - &projection_column_names, - ); if let Some(augmented) = crate::indexed_table::page_index_loader::load_scoped_page_index_cols( &store, &segment.object_path, diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/page_index_loader.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/page_index_loader.rs index e34fa15f1e76a..8452bbe404641 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/page_index_loader.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/page_index_loader.rs @@ -367,6 +367,39 @@ pub fn resolve_predicate_parquet_columns( resolve_with_schema(&file_arrow_schema, metadata, predicate_column_names) } +/// Resolve TWO name-sets (e.g. predicate columns and projection columns) against +/// the same file in one pass. Deriving the per-file arrow schema +/// (`parquet_to_arrow_schema`) is the dominant cost of name→leaf resolution on +/// wide schemas (it rebuilds the whole file's Schema); the two callers in the +/// indexed setup loop previously each rebuilt it, so doing it once here removes a +/// full schema reconstruction per file per query. Pure refactor — each returned +/// Vec is identical to calling `resolve_predicate_parquet_columns` separately. +pub fn resolve_predicate_parquet_columns_pair( + union_schema: &SchemaRef, + metadata: &ParquetMetaData, + names_a: &[String], + names_b: &[String], +) -> (Vec, Vec) { + let parquet_schema = metadata.file_metadata().schema_descr(); + match datafusion::parquet::arrow::parquet_to_arrow_schema( + parquet_schema, + metadata.file_metadata().key_value_metadata(), + ) { + Ok(s) => { + let file_arrow_schema = Arc::new(s); + ( + resolve_with_schema(&file_arrow_schema, metadata, names_a), + resolve_with_schema(&file_arrow_schema, metadata, names_b), + ) + } + // Same fallback as the single-name path: resolve against the union schema. + Err(_) => ( + resolve_with_schema(union_schema, metadata, names_a), + resolve_with_schema(union_schema, metadata, names_b), + ), + } +} + /// Resolve predicate column names → parquet leaf indices against a specific arrow /// schema, via the same `StatisticsConverter` mapping DataFusion's pruner uses. fn resolve_with_schema( @@ -468,24 +501,28 @@ async fn load_combined( let oi_ms = t_oi.elapsed(); // The graft deep-clones the whole footer (row_groups + column chunk metadata // are owned, not Arc) — paid on EVERY query, hit or miss. Time it separately - // so we can see if this per-query clone (Problem B) dominates the scoped-cache - // path vs the cell decode/lookup. + // so the debug log can attribute the per-query clone vs the cell decode/lookup. let t_graft = std::time::Instant::now(); let grafted = graft(footer_meta, column_index, offset_index); + let graft_ms = t_graft.elapsed(); native_bridge_common::log_debug!( "scoped-load {}: total={:?} (colidx={:?} offidx={:?} graft={:?})", location, t_total.elapsed(), ci_ms, oi_ms, - t_graft.elapsed(), + graft_ms, ); Some(grafted) } /// Build a fresh `ParquetMetaData` = `footer` with the page-index pair grafted -/// on. Deep-clones the footer (the builder consumes an owned `ParquetMetaData`); -/// that transient clone is the only footer copy and is never cached. +/// on. Clones the footer to get an owned value for the builder — but with +/// `ParquetMetaData.row_groups` held behind an `Arc` (see the arrow-rs change), +/// that clone is a refcount bump, not a deep copy of every row group's +/// column-chunk metadata. On wide / many-row-group files (e.g. textbench's +/// ~403-col, ~64-RG footers) that deep copy was ~60ms/query; sharing makes the +/// graft effectively free. fn graft( footer_meta: &Arc, column_index: ParquetColumnIndex, @@ -697,8 +734,17 @@ async fn get_or_build_offset_index( b.append_row_count(footer_meta.row_group(rg_idx).num_rows()); b.build() }; + // The placeholder is identical for every column within a row group (it only + // depends on the RG's row count). Build it ONCE per RG and clone it across the + // columns, instead of constructing `num_cols` identical OffsetIndexMetaData + // (each a heap alloc) per RG. On wide schemas (clickbench ~105 cols) this is the + // bulk of the per-file warm scoped-load cost — only the few scoped columns get + // real data scattered in afterward; the rest stay as clones of this placeholder. let mut matrix: ParquetOffsetIndex = (0..num_rgs) - .map(|rg| (0..num_cols).map(|_| placeholder_for(rg)).collect()) + .map(|rg| { + let ph = placeholder_for(rg); + vec![ph; num_cols] + }) .collect(); // Phase 1: serve cached columns; collect misses. @@ -1273,6 +1319,25 @@ mod tests { ); } + // The pair-resolve must return exactly what two separate single-name resolves + // return — it only shares the per-file arrow-schema derivation, nothing else. + #[tokio::test] + async fn resolve_pair_equals_two_single_resolves() { + let (bytes, schema) = two_col_parquet(); + let fo = footer_only(&bytes); + let names_a = vec!["price".to_string()]; + let names_b = vec!["qty".to_string(), "price".to_string()]; + let single_a = resolve_predicate_parquet_columns(&schema, &fo, &names_a); + let single_b = resolve_predicate_parquet_columns(&schema, &fo, &names_b); + let (pair_a, pair_b) = + resolve_predicate_parquet_columns_pair(&schema, &fo, &names_a, &names_b); + assert_eq!(pair_a, single_a, "pair predicate result must match single"); + assert_eq!(pair_b, single_b, "pair projection result must match single"); + // And the values are the expected leaves (price=0, qty=1). + assert_eq!(pair_a, vec![0]); + assert_eq!(pair_b, vec![0, 1]); + } + #[tokio::test] async fn scoped_pruning_matches_full_index() { let _g = CACHE_TEST_GUARD.lock().unwrap(); diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/shard_scoped_reader.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/shard_scoped_reader.rs index 38c4d231e6e1d..dd16391576f7a 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/shard_scoped_reader.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/shard_scoped_reader.rs @@ -73,7 +73,7 @@ use object_store::{ObjectStore, ObjectStoreExt}; use prost::bytes::Bytes; use crate::indexed_table::page_index_loader::{ - load_scoped_page_index_cols, resolve_predicate_parquet_columns, + load_scoped_page_index_cols, }; use crate::indexed_table::parquet_bridge::load_parquet_metadata; @@ -220,8 +220,12 @@ impl AsyncFileReader for ScopedPageIndexReader { // (per-file, so schema evolution across files is handled), then // augment with an all-RG, column-scoped page index. if !predicate_names.is_empty() { - let parquet_cols = - resolve_predicate_parquet_columns(&file_schema, &footer, &predicate_names); + // Resolve predicate + projection in one pass (shared per-file arrow + // schema derivation — the dominant cost on wide schemas). + let (parquet_cols, offset_cols) = + crate::indexed_table::page_index_loader::resolve_predicate_parquet_columns_pair( + &file_schema, &footer, &predicate_names, &projection_names, + ); native_bridge_common::log_debug!( "scoped-RESOLVE {}: names={:?} -> leaves={:?} (file_schema_fields={}, parquet_leaves={})", location, @@ -230,12 +234,8 @@ impl AsyncFileReader for ScopedPageIndexReader { file_schema.fields().len(), footer.file_metadata().schema_descr().num_columns(), ); - // Projected leaf indices for THIS file, for OffsetIndex column - // scoping (predicate ∪ projection ∪ {0}). Empty projection_names - // (couldn't derive) → empty offset_cols → loader unions in - // predicate + col0 only (still far smaller than all columns). - let offset_cols = - resolve_predicate_parquet_columns(&file_schema, &footer, &projection_names); + // offset_cols (projected leaves for OffsetIndex column scoping, + // predicate ∪ projection ∪ {0}) came from the pair resolve above. if let Some(augmented) = load_scoped_page_index_cols( &store, &location, @@ -245,13 +245,6 @@ impl AsyncFileReader for ScopedPageIndexReader { ) .await { - log::debug!( - "scoped-pageidx[listing]: {} rgs={} pred_cols={} offset_cols={}", - location, - footer.num_row_groups(), - parquet_cols.len(), - offset_cols.len() - ); return Ok(augmented); } } From 7f40d315c6cacfce3fcff82aceb755a47180ecec Mon Sep 17 00:00:00 2001 From: G Date: Fri, 19 Jun 2026 16:05:15 +0530 Subject: [PATCH 18/18] empty predicate offset index fixes Signed-off-by: G --- .../rust/src/indexed_executor.rs | 23 ++++---- .../src/indexed_table/page_index_loader.rs | 55 +++++++++++++++++-- 2 files changed, 62 insertions(+), 16 deletions(-) diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_executor.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_executor.rs index 5d50a5ab21e2f..2b14d7f76754d 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_executor.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_executor.rs @@ -1025,16 +1025,15 @@ async unsafe fn execute_indexed_with_context_inner( // (no RG scoping in Step 1). On any failure the segment keeps footer-only // metadata and pruning no-ops. let predicate_column_names = collect_predicate_column_names(extraction.as_ref(), &schema); - if !predicate_column_names.is_empty() { - // Scope the OffsetIndex to the columns this query actually touches - // (projection ∪ predicate), instead of decoding+caching it for all - // columns. On the 402-column textbench schema the all-columns OffsetIndex - // was ~256 MB/file and was rebuilt every query regardless of the - // predicate; scoping it to the read columns is the difference between a - // few MB and 256 MB, and makes the cache actually fit its budget. - // `collect_plan_column_names` is a superset of the read set (safe — see - // its docs); the loader additionally unions in the predicate cols + col 0. - let projection_column_names = collect_plan_column_names(&logical_plan); + // Projected columns this query reads (superset of the read set; safe). A + // match()-only query has NO residual predicate columns but DOES project real + // columns (e.g. `... | stats count() by URL` projects URL) — those columns + // still need a scoped OffsetIndex so the parquet reader fetches only the + // matched rows' pages instead of whole column chunks. Gating the scoped build + // on the predicate alone (the old behavior) skipped the OffsetIndex for these + // queries, making them read the full column (measured 2.5× more bytes_scanned). + let projection_column_names = collect_plan_column_names(&logical_plan); + if !predicate_column_names.is_empty() || !projection_column_names.is_empty() { native_bridge_common::log_debug!( "scoped-PROJCOLS predicate_names={:?} projection_names(n={})={:?}", predicate_column_names, @@ -1053,7 +1052,9 @@ async unsafe fn execute_indexed_with_context_inner( &predicate_column_names, &projection_column_names, ); - if parquet_cols.is_empty() { + // Proceed if EITHER a predicate (→ ColumnIndex for pruning) OR a + // projection (→ OffsetIndex for page-level IO) resolved to real leaves. + if parquet_cols.is_empty() && offset_cols.is_empty() { continue; } if let Some(augmented) = crate::indexed_table::page_index_loader::load_scoped_page_index_cols( diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/page_index_loader.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/page_index_loader.rs index 8452bbe404641..db426fd8fbe1b 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/page_index_loader.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/page_index_loader.rs @@ -487,13 +487,27 @@ async fn load_combined( surviving_rgs: Option<&[usize]>, offset_cols: Option<&[usize]>, ) -> Option> { - if parquet_cols.is_empty() { + // Nothing to build only when there is NEITHER a predicate (ColumnIndex for + // pruning) NOR an explicit projection (OffsetIndex for page-level IO). A + // match-only query has empty `parquet_cols` but a real, non-empty + // `offset_cols` (the projected columns) — it still needs the OffsetIndex so + // the parquet reader can fetch just the matched rows' pages instead of whole + // column chunks. `offset_cols == None` (all-columns) with no predicate is the + // legacy "nothing requested" case → return None as before. + let has_offset_work = offset_cols.map(|c| !c.is_empty()).unwrap_or(false); + if parquet_cols.is_empty() && !has_offset_work { return None; } let t_total = std::time::Instant::now(); let t_ci = std::time::Instant::now(); - let column_index = - get_or_build_column_index(store, location, footer_meta, parquet_cols, surviving_rgs).await?; + // ColumnIndex (predicate-driven, for page pruning). Skipped when there's no + // predicate — a `None` matrix grafts cleanly (consumers treat absent CI as + // "no usable stats", same as footer-only). + let column_index = if parquet_cols.is_empty() { + None + } else { + Some(get_or_build_column_index(store, location, footer_meta, parquet_cols, surviving_rgs).await?) + }; let ci_ms = t_ci.elapsed(); let t_oi = std::time::Instant::now(); let offset_index = @@ -525,13 +539,13 @@ async fn load_combined( /// graft effectively free. fn graft( footer_meta: &Arc, - column_index: ParquetColumnIndex, + column_index: Option, offset_index: ParquetOffsetIndex, ) -> Arc { let base = ParquetMetaData::clone(footer_meta); let rebuilt = base .into_builder() - .set_column_index(Some(column_index)) + .set_column_index(column_index) .set_offset_index(Some(offset_index)) .build(); Arc::new(rebuilt) @@ -1338,6 +1352,37 @@ mod tests { assert_eq!(pair_b, vec![0, 1]); } + /// Regression: a match()-only query has NO residual predicate columns + /// (`parquet_cols` empty) but DOES project columns (`offset_cols` non-empty). + /// The scoped load must still build the OffsetIndex for the projected column + /// so the parquet reader fetches only matched-row pages, not whole chunks. + /// (Bug: the load short-circuited on empty `parquet_cols`, skipping the + /// OffsetIndex → reader over-read ~2.5× the bytes on `... | stats ... by URL`.) + #[tokio::test] + async fn projection_only_builds_offset_index_without_predicate() { + let _g = CACHE_TEST_GUARD.lock().unwrap(); + clear_scoped_cache_for_test(); + let (bytes, _schema) = two_col_parquet(); // price=col0, qty=col1 + let (store, loc) = stage(bytes.clone()).await; + let fo = footer_only(&bytes); + + // No predicate columns; project qty (col 1). + let parquet_cols: Vec = vec![]; + let offset_cols: Vec = vec![1]; + let aug = load_scoped_page_index_cols(&store, &loc, &fo, &parquet_cols, &offset_cols) + .await + .expect("projection-only load must produce grafted metadata (not None)"); + + // No predicate → no ColumnIndex grafted. + assert!(aug.column_index().is_none(), "no predicate → ColumnIndex absent"); + + // OffsetIndex must be real for the projected col (1) AND col 0 (loader + // always unions in {0}); other behavior unchanged. + let o = aug.offset_index().expect("OffsetIndex must be grafted"); + assert!(!o[0][1].page_locations().is_empty(), "projected col qty has real OffsetIndex"); + assert!(!o[0][0].page_locations().is_empty(), "col 0 OffsetIndex real (always unioned)"); + } + #[tokio::test] async fn scoped_pruning_matches_full_index() { let _g = CACHE_TEST_GUARD.lock().unwrap();