diff --git a/HANDOFF_step2_rg_scoping.md b/HANDOFF_step2_rg_scoping.md new file mode 100644 index 0000000000000..cadf57cd3eb2c --- /dev/null +++ b/HANDOFF_step2_rg_scoping.md @@ -0,0 +1,265 @@ +# 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. + +--- + +## 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 +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/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/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..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,19 +359,18 @@ 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 - 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/ffm.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs index 55e8719813e64..771efe1f34e03 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs @@ -1040,6 +1040,45 @@ 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 ColumnIndex cache. +/// +/// 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_column_index_cache_limit(size_limit: i64) -> i64 { + if size_limit < 0 { + 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_offset_index_cache_limit(size_limit as usize); + 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_executor.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_executor.rs index faf1615d0b3b0..2b14d7f76754d 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,80 @@ 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() +} +/// 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 { + 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) + }); + // 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() +} + /// 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 { @@ -558,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 @@ -921,6 +1011,66 @@ 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); + // 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, + projection_column_names.len(), + projection_column_names, + ); + for segment in segments.iter_mut() { + // 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, + ); + // 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( + &store, + &segment.object_path, + &segment.metadata, + &parquet_cols, + &offset_cols, + ) + .await + { + segment.metadata = augmented; + } + } + } + let factory: EvaluatorFactory = match classification { FilterClass::None => { // Predicate-only scan: page-pruned universe, residual applied in 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..db426fd8fbe1b --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/page_index_loader.rs @@ -0,0 +1,1998 @@ +/* + * 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 parquet page-index caches — TWO caches, by consumer. +//! +//! # 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. 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). +//! +//! # Two caches, because the two indexes have different drivers +//! +//! 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-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`]). +//! +//! - **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). +//! +//! 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 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 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 +//! (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}; + +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 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; + +// ── Generic byte-bounded LRU ──────────────────────────────────────────────── + +/// 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, + pub misses: u64, + pub evictions: u64, + pub entries: usize, + pub used_bytes: usize, + pub limit_bytes: usize, +} + +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). + tick: u64, + hits: u64, + misses: u64, + evictions: u64, +} + +impl Lru { + 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 + } + + 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.value.clone()) + } + None => { + self.misses += 1; + None + } + } + } + + 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, LruEntry { value, 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(); + } + + 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, +} + +/// 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))); + +// ── 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) = COLUMN_INDEX_CACHE.lock() { + c.set_limit(limit); + } +} + +/// 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() +} + +/// 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) = COLUMN_INDEX_CACHE.lock() { + c.clear_keep_limit(); + } + if let Ok(mut c) = OFFSET_INDEX_CACHE.lock() { + c.clear_keep_limit(); + } +} + +// ── Public API ────────────────────────────────────────────────────────────── + +/// 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 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( + 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 + 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_combined(store, location, footer_meta, parquet_cols, None, None).await +} + +/// 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, + footer_meta: &Arc, + parquet_cols: &[usize], + offset_cols: &[usize], +) -> Option> { + load_combined(store, location, footer_meta, parquet_cols, None, Some(offset_cols)).await +} + +/// 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_combined( + store, + location, + footer_meta, + parquet_cols, + Some(surviving_rgs), + Some(offset_cols), + ) + .await +} + +async fn load_combined( + store: &Arc, + location: &object_store::path::Path, + footer_meta: &Arc, + parquet_cols: &[usize], + surviving_rgs: Option<&[usize]>, + offset_cols: Option<&[usize]>, +) -> Option> { + // 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(); + // 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 = + get_or_build_offset_index(store, location, footer_meta, parquet_cols, offset_cols).await?; + 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 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, + graft_ms, + ); + Some(grafted) +} + +/// Build a fresh `ParquetMetaData` = `footer` with the page-index pair grafted +/// 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: Option, + offset_index: ParquetOffsetIndex, +) -> Arc { + let base = ParquetMetaData::clone(footer_meta); + let rebuilt = base + .into_builder() + .set_column_index(column_index) + .set_offset_index(Some(offset_index)) + .build(); + Arc::new(rebuilt) +} + +// ── ColumnIndex cache lookup + build (per `(file, col, rg)` cell) ──────────── + +/// 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. +/// +/// `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]>, +) -> 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(); + if parquet_cols.iter().any(|&i| i >= num_cols) { + return None; + } + + // 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 = 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(); + + // 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 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; + } + + 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() + }; + // 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| { + let ph = placeholder_for(rg); + vec![ph; num_cols] + }) + .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; + } + native_bridge_common::log_debug!( + "scoped-offidx {}: off_cols={:?}/{} (scoped={}) pred={:?} hit={} miss={}", + location, + off_cols, + num_cols, + offset_cols.is_some(), + parquet_cols, + 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?; + native_bridge_common::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); + 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, + chunks: Vec, + range: Range, + } + let mut plans: Vec = 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 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; + } + + let buffers = store.get_ranges(location, &fetch_ranges).await.ok()?; + if buffers.len() != fetch_ranges.len() { + return None; + } + + // 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(); + for (plan, buf) in plans.iter().zip(buffers.iter()) { + let reader = BufferChunkReader { base: plan.range.start, bytes: buf.clone() }; + #[allow(deprecated)] + let decoded = read_offset_indexes(&reader, &plan.chunks).ok()??; + 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); + } + } + + 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)); + } + 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`. +/// +/// 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 arrow::array::{ArrayRef, BooleanArray, UInt64Array}; + use datafusion::physical_optimizer::pruning::{PruningPredicate, PruningStatistics}; + use datafusion::scalar::ScalarValue; + + let num_rgs = footer_meta.num_row_groups(); + let all: Vec = (0..num_rgs).collect(); + if num_rgs == 0 { + return all; + } + + let Ok(pp) = PruningPredicate::try_new(Arc::clone(predicate), Arc::clone(arrow_schema)) else { + return all; + }; + + 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 { + self.conv(&column.name)? + .row_group_mins(self.meta.row_groups().iter()) + .ok() + } + fn max_values(&self, column: &datafusion::common::Column) -> Option { + 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 { + 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(); + Some(Arc::new(UInt64Array::from(counts)) as ArrayRef) + } + fn contained( + &self, + _column: &datafusion::common::Column, + _values: &std::collections::HashSet, + ) -> Option { + None + } + } + + let stats = RgStats { meta: footer_meta, schema: arrow_schema, num_rgs }; + match pp.prune(&stats) { + Ok(mask) => mask + .iter() + .enumerate() + .filter_map(|(i, keep)| if *keep { Some(i) } else { None }) + .collect(), + Err(_) => all, + } +} + +// ── 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 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 into the buffer. +struct BufferChunkReader { + 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 { + 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 ──────────────────────────────────────────────────────── + +/// 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(()); + +/// 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) = 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; + } +} + +#[cfg(test)] +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 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), + } +} + +#[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}; + + 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) + } + + /// 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"); + 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() + } + fn ci() -> ScopedCacheStats { + column_index_cache_stats() + } + fn oi() -> ScopedCacheStats { + offset_index_cache_stats() + } + + 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) + } + + // ── baseline / correctness ──────────────────────────────────────────── + + #[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()); + } + + #[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!(ci().entries, 0); + assert_eq!(oi().entries, 0); + } + + #[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.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!( + !o[0][0].page_locations().is_empty() && !o[0][1].page_locations().is_empty(), + "OffsetIndex real for every column (all-col default)" + ); + } + + // 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]); + } + + /// 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(); + 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); + 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)); + } + + /// 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(); + 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)]); + 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); + } + + // ── 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(); + 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 _ = 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 _ = 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 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_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()]); + + 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(); + } + + /// 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 _ = 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(); + } + + /// 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(); + 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!((ci().hits, ci().misses), (0, 1)); + let _ = load_scoped_page_index(&store, &loc, &fo, &c_price).await.unwrap(); + assert_eq!((ci().hits, ci().misses), (1, 1)); + let _ = load_scoped_page_index(&store, &loc, &fo, &c_qty).await.unwrap(); + 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 = ci(); + assert_eq!((s.hits, s.misses, s.entries, s.evictions), (3, 2, 2, 0)); + } + + /// 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(); + 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()]); + + // 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_cell = ci().used_bytes; + assert!(one_cell > 0); + let budget = one_cell + one_cell / 2; + clear_scoped_cache_for_test(); + set_column_index_cache_limit_for_test(budget); + + 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!(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 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: RG-scoping the ColumnIndex ──────────────────────────────── + + #[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); + let p = pred("id", 0, Operator::GtEq, 25); + 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]); + } + + #[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()]); + let surviving = vec![2usize, 3usize]; + + 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!(c[rg][0], ColumnIndexMetaData::NONE), "survivor RG {rg} real CI"); + } + for &rg in &[0usize, 1usize] { + assert!(matches!(c[rg][0], ColumnIndexMetaData::NONE), "pruned RG {rg} NONE CI"); + } + for rg in 0..4 { + for cc in 0..2 { + assert!(!o[rg][cc].page_locations().is_empty(), "OI real for all rg/col"); + } + } + 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(); + } + + #[tokio::test] + 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(); + 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(&store, &loc, &fo, &cols).await.unwrap(); + let all_rg = ci().used_bytes; + clear_scoped_cache_for_test(); + let _ = load_scoped_page_index_rgs(&store, &loc, &fo, &cols, &[2, 3]).await.unwrap(); + assert!(ci().used_bytes < all_rg, "RG-scoped CI bytes {} < all-RG {}", ci().used_bytes, all_rg); + clear_scoped_cache_for_test(); + } + + /// 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(); + 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!((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()]); + + let _ = load_scoped_page_index_rgs(&store, &loc, &fo, &cols, &[2, 3]).await.unwrap(); + assert_eq!((ci().hits, ci().misses, ci().entries), (0, 2, 2), "cells (id,rg2)+(id,rg3)"); + + // {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 ──────────────────────────── + + #[tokio::test] + async fn col_scoped_offset_index_only_for_requested_columns() { + 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()]); + assert_eq!(pred_cols, vec![1]); + + 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_eq!(o[0][c].page_locations().len(), real_pages, "col {c} (pred/proj/metric) real OI"); + } + 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] + 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(); + let (store, loc) = stage(bytes.clone()).await; + let fo = footer_only(&bytes); + let pred_cols = resolve_predicate_parquet_columns(&schema, &fo, &["price".to_string()]); + 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()).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); + } + + #[tokio::test] + 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(); + let (store, loc) = stage(bytes.clone()).await; + let fo = footer_only(&bytes); + let pred_cols = resolve_predicate_parquet_columns(&schema, &fo, &["n1".to_string()]); + + let _ = load_scoped_page_index(&store, &loc, &fo, &pred_cols).await.unwrap(); + let all_cols = oi().used_bytes; + clear_scoped_cache_for_test(); + let _ = load_scoped_page_index_cols(&store, &loc, &fo, &pred_cols, &[2]).await.unwrap(); + assert!(oi().used_bytes < all_cols, "col-scoped OI {} < all-col {}", oi().used_bytes, all_cols); + clear_scoped_cache_for_test(); + } + + /// 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(); + 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()]); + + let _ = load_scoped_page_index(&store, &loc, &fo, &pred_cols).await.unwrap(); + 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!(oi().entries, 2, "covered columns reuse their cells, no new entries"); + assert_eq!(oi().hits, 2); + clear_scoped_cache_for_test(); + } + + /// 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/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 5231187a694e2..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 @@ -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))?; @@ -297,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)" + ); +} 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..346b0dcaba2c3 --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/scoped_page_index_optimizer.rs @@ -0,0 +1,373 @@ +/* + * 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)); + } + + // 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). + let factory = Arc::new(ScopedPageIndexReaderFactory::new( + Arc::clone(&self.store), + Arc::clone(&self.metadata_cache), + names, + projection_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()], + 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..dd16391576f7a --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/shard_scoped_reader.rs @@ -0,0 +1,406 @@ +/* + * 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_cols, +}; +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>, + /// 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). + #[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, + projection_column_names: Vec, + predicate: Option>, + file_schema: SchemaRef, + ) -> Self { + Self { + store, + metadata_cache, + predicate_column_names: Arc::new(predicate_column_names), + projection_column_names: Arc::new(projection_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), + 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, + })) + } +} + +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, +} + +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 projection_names = Arc::clone(&self.projection_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() { + // 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, + predicate_names.as_ref(), + parquet_cols, + file_schema.fields().len(), + footer.file_metadata().schema_descr().num_columns(), + ); + // 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, + &footer, + &parquet_cols, + &offset_cols, + ) + .await + { + 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()], + // 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, + ); + 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![], + 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/rust/src/stats.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/stats.rs index 2fcec9ce4d86e..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` | 10 × i64 (2 × 5: metadata + statistics 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,6 +92,14 @@ pub struct CacheGroupRepr { pub struct CacheStatsRepr { pub metadata_cache: CacheGroupRepr, pub statistics_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 { @@ -111,6 +119,8 @@ impl Default for CacheStatsRepr { Self { metadata_cache: CacheGroupRepr::default(), statistics_cache: CacheGroupRepr::default(), + column_index_cache: CacheGroupRepr::default(), + offset_index_cache: CacheGroupRepr::default(), } } } @@ -161,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::() == 10 * 8); +const _: () = assert!(std::mem::size_of::() == 20 * 8); const _: () = assert!(std::mem::size_of::() == 17 * 8); -const _: () = assert!(std::mem::size_of::() == 75 * 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 == 600); + const _: () = assert!(BUFFER_BYTE_SIZE == 680); } /// Snapshot a `RuntimeMonitor` and return a populated `RuntimeMetricsRepr`. @@ -304,6 +314,27 @@ pub fn pack_cache_stats(mgr: &CustomCacheManager) -> CacheStatsRepr { memory_bytes: statistics_memory, size_limit_bytes: mgr.statistics_cache_size_limit() as i64, }, + // 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, + entry_count: s.entries as i64, + memory_bytes: s.used_bytes as i64, + size_limit_bytes: s.limit_bytes as i64, + } + }, } } @@ -396,7 +427,7 @@ mod tests { search_stats: crate::search_stats::snapshot(), }; - assert_eq!(layout::BUFFER_BYTE_SIZE, 600); + 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); @@ -411,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::(), 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::(), 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); @@ -431,6 +462,43 @@ mod tests { assert_eq!(g.memory_bytes, 0); assert_eq!(g.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::{ + 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 caches. + let _g = SCOPED_CACHE_TEST_GUARD.lock().unwrap(); + // 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 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 ab8892c23cbfd..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 @@ -764,7 +770,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/DatafusionSettings.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionSettings.java index 06de123d0a675..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,6 +298,8 @@ static String derivePoolMinDefault(Settings settings, int percent) { CacheSettings.STATISTICS_CACHE_EVICTION_TYPE, CacheSettings.METADATA_CACHE_ENABLED, CacheSettings.STATISTICS_CACHE_ENABLED, + 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/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/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..cab44a7ce84d8 --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/action/stats/RestClearScopedPageIndexCacheAction.java @@ -0,0 +1,61 @@ +/* + * 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.rest.BaseRestHandler; +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; + +import static java.util.Collections.singletonList; +import static org.opensearch.rest.RestRequest.Method.POST; + +/** + * 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 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
+ * + * @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 { + // 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/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()); + } +} 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..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,6 +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 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", @@ -36,6 +38,35 @@ public class CacheSettings { Setting.Property.Dynamic ); + /** + * 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 COLUMN_INDEX_CACHE_SIZE_LIMIT = new Setting<>( + COLUMN_INDEX_CACHE_SIZE_LIMIT_KEY, + "64mb", + (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 + ); + public static final Setting METADATA_CACHE_EVICTION_TYPE = new Setting( "datafusion.metadata.cache.eviction.type", "LRU", @@ -74,7 +105,9 @@ public class CacheSettings { METADATA_CACHE_EVICTION_TYPE, STATISTICS_CACHE_ENABLED, STATISTICS_CACHE_SIZE_LIMIT, - STATISTICS_CACHE_EVICTION_TYPE + STATISTICS_CACHE_EVICTION_TYPE, + 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 1b95478fd83aa..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 @@ -117,6 +117,23 @@ public static NativeCacheManagerHandle createCacheConfig(ClusterSettings cluster logger.debug("Cache type {} is disabled", type.getCacheTypeName()); } } + + // 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 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 475f8d228191b..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,6 +130,9 @@ 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_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; private static final MethodHandle CLOSE_SESSION_CONTEXT; @@ -501,6 +504,21 @@ private static RuntimeException rethrowConverted(RuntimeException e) { ) ); + 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) + ); + + 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( @@ -762,6 +780,44 @@ public static void setMemoryPoolLimit(long runtimePtr, long newLimitBytes) { } } + /** + * 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. + * 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 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_OFFSET_INDEX_CACHE_LIMIT, 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 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..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,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 × 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 { @@ -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() != 85 * Long.BYTES) { + throw new AssertionError("StatsLayout size mismatch: expected " + (85 * Long.BYTES) + " but got " + LAYOUT.byteSize()); } } @@ -182,6 +183,20 @@ 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.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"); private static final VarHandle SS_SINGLE_COLLECTOR_SCAN = handle("search_stats", "single_collector_scan"); @@ -290,11 +305,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 +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) ); - return new CacheStats(metadata, statistics); + 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, columnIndex, offsetIndex); } /** @@ -402,7 +430,12 @@ 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("column_index_cache"), + cacheGroup("offset_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..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,16 +32,27 @@ public class CacheStats implements Writeable, ToXContentFragment { private final CacheGroupStats metadataCache; private final CacheGroupStats statisticsCache; + 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 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) { + public CacheStats( + CacheGroupStats metadataCache, + CacheGroupStats statisticsCache, + CacheGroupStats columnIndexCache, + CacheGroupStats offsetIndexCache + ) { this.metadataCache = Objects.requireNonNull(metadataCache); this.statisticsCache = Objects.requireNonNull(statisticsCache); + this.columnIndexCache = Objects.requireNonNull(columnIndexCache); + this.offsetIndexCache = Objects.requireNonNull(offsetIndexCache); } /** @@ -50,12 +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.columnIndexCache = new CacheGroupStats(in); + this.offsetIndexCache = new CacheGroupStats(in); } @Override public void writeTo(StreamOutput out) throws IOException { metadataCache.writeTo(out); statisticsCache.writeTo(out); + columnIndexCache.writeTo(out); + offsetIndexCache.writeTo(out); } @Override @@ -67,6 +85,12 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws builder.startObject("statistics_cache"); statisticsCache.toXContent(builder); builder.endObject(); + builder.startObject("column_index_cache"); + columnIndexCache.toXContent(builder); + builder.endObject(); + builder.startObject("offset_index_cache"); + offsetIndexCache.toXContent(builder); + builder.endObject(); builder.endObject(); return builder; } @@ -81,16 +105,29 @@ public CacheGroupStats getStatisticsCache() { return statisticsCache; } + /** Returns the scoped ColumnIndex cache counters. */ + public CacheGroupStats getColumnIndexCache() { + return columnIndexCache; + } + + /** Returns the scoped OffsetIndex cache counters. */ + public CacheGroupStats getOffsetIndexCache() { + return offsetIndexCache; + } + @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(columnIndexCache, that.columnIndexCache) + && Objects.equals(offsetIndexCache, that.offsetIndexCache); } @Override public int hashCode() { - return Objects.hash(metadataCache, statisticsCache); + 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 4fd7679ce5b9b..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 = 75; + 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-57) + // 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,26 +187,36 @@ public void testPackThenDecodeRoundTripWithCpu() { assertEquals(values[55], cs.getStatisticsCache().entryCount); assertEquals(values[56], cs.getStatisticsCache().memoryBytes); assertEquals(values[57], cs.getStatisticsCache().sizeLimitBytes); - - // Search stats (offsets 58-74) + 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[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[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); } } } @@ -328,6 +338,18 @@ public void testDecodeThenReencodeIdentity() { cs.getStatisticsCache().entryCount, cs.getStatisticsCache().memoryBytes, cs.getStatisticsCache().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 9a72c2af21a26..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 600 (75 × 8). */ + /** 7.1: Layout byte size must be 680 (85 × 8). */ public void testLayoutByteSize() { - assertEquals(600L, StatsLayout.LAYOUT.byteSize()); - assertEquals(75 * 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 77a839c377b9c..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 @@ -90,21 +90,33 @@ 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, 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); - CacheStats c = new CacheStats(meta, stats); + 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(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)); + 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(41, 42, 43, 44, 45) + ); BytesStreamOutput out = new BytesStreamOutput(); original.writeTo(out); StreamInput in = out.bytes().streamInput(); @@ -113,7 +125,12 @@ 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), + new CacheGroupStats(4, 1, 2, 256, 16_000_000) + ); XContentBuilder builder = XContentFactory.jsonBuilder(); builder.startObject(); c.toXContent(builder, ToXContent.EMPTY_PARAMS); @@ -122,9 +139,11 @@ public void testCacheStatsToXContentShape() throws IOException { // Top-level wrapper assertTrue("expected cache_stats wrapper, got: " + json, json.contains("\"cache_stats\"")); - // Both sub-groups present + // All four sub-groups present assertTrue(json.contains("\"metadata_cache\"")); assertTrue(json.contains("\"statistics_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")); @@ -132,27 +151,50 @@ 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")); + 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)); + 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) + ); 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 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 twice (one per sub-cache)", 2, 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")); } 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), + 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(16, 17, 18, 19, 20) + ); + // 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, 15), + new CacheGroupStats(16, 17, 18, 19, 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..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()); + 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 075a48c5ebcc5..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 @@ -278,7 +278,9 @@ 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), + 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 new file mode 100644 index 0000000000000..60f2bc11aac75 --- /dev/null +++ b/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/ScopedPageIndexCacheIT.java @@ -0,0 +1,553 @@ +/* + * 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.List; +import java.util.Locale; +import java.util.Map; + +/** + * 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}, {@code cache_stats.column_index_cache}, and + * {@code cache_stats.offset_index_cache} groups. + * + *

The two scoped caches

+ * + * 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.
    • + *
+ * + *

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}. + */ +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) { + DatasetProvisioner.provision(client(), DATASET); + dataProvisioned = true; + } + } + + // ---- stats helpers --------------------------------------------------- + + /** 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 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"); + + 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 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 CacheGroup(hits, misses, entries, memory, limit); + } + + /** 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); + return ((Number) v).longValue(); + } + + /** 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(); + } + + private String src() { + return "source=" + DATASET.indexName; + } + + // ---- exposure + correctness ----------------------------------------- + + /** + * Both scoped cache groups must always be present with a positive byte budget + * (their configured limits), even before any query runs. + */ + 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. + */ + public void testListingQueryCorrectnessUnchanged() throws IOException { + 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(src() + " | where status >= 400 | stats count()") + ); + assertEquals( + "log_level = 'ERROR' count must be exact (keyword predicate)", + LEVEL_ERROR, + 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 caches) — a cached page index must never change the answer. + */ + public void testCorrectnessIsStableAcrossCachedReRuns() throws IOException { + 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); + 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). Guards against the page-index strip breaking normal footer caching. + */ + public void testMetadataCacheStillServesFooters() throws IOException { + String q = src() + " | where status >= 200 | stats count() by service_name"; + executePpl(q); // warm + CacheGroup before = metadata(); + 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( + 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); + } + + // ---- populate, hit, bounded ----------------------------------------- + + /** + * 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 { + clearScopedCaches(); + String query = src() + " | where status >= 400 | stats count() by service_name"; + + executePpl(query); + 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 ci2 = columnIndex(); + CacheGroup oi2 = offsetIndex(); + + assertTrue( + String.format(Locale.ROOT, "CI re-run must register hits (h1=%d h2=%d)", ci1.hits(), ci2.hits()), + ci2.hits() > ci1.hits() + ); + 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 caches bounded: cells and bytes do not + * grow after the first populating run; the re-runs all register hits. + */ + public void testRepeatedListingQueriesDoNotGrowTheCache() throws IOException { + clearScopedCaches(); + String query = src() + " | where status >= 200 | stats count()"; + + executePpl(query); + 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 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 CI re-runs must register hits (base=%d after=%d)", + ciBase.hits(), ciAfter.hits()), + ciAfter.hits() >= ciBase.hits() + 9 + ); + } + + /** + * Neither scoped cache may exceed its configured byte budget — a basic + * "no over-allocation" invariant readable from the stats API. + */ + 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( + "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 + ); + } + + /** + * 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 testDifferentLiteralsSameColumnShareCells() throws IOException { + clearScopedCaches(); + + executePpl(src() + " | where status >= 400 | stats count()"); + CacheGroup first = columnIndex(); + assertTrue("first literal must populate CI cells", first.entries() >= 1); + + 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( + String.format(Locale.ROOT, "the second literal must HIT the existing cells (hits %d -> %d)", + first.hits(), second.hits()), + second.hits() > first.hits() + ); + } + + /** + * 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() + ); + } + + // ---- cell reuse: OffsetIndex (projection-driven) -------------------- + + /** + * 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()); + } + + /** + * 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 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}, 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 { + 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, 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 cells (listing hits=%d indexed hits=%d)", + afterListing.hits(), afterIndexed.hits()), + afterIndexed.hits() > afterListing.hits() + ); + assertEquals( + "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 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); + } +}