From 0b4b44c1de10c344c035274b1cbe7e33279a4e55 Mon Sep 17 00:00:00 2001 From: Yihong Zhang Date: Sat, 18 Jul 2026 13:49:36 -0700 Subject: [PATCH 1/2] Skip building per-subset join indexes for probers that will only be probed once. --- CHANGELOG.md | 1 + Cargo.toml | 2 + core-relations/Cargo.toml | 2 + core-relations/src/free_join/execute.rs | 341 +++++++++++++++++--- core-relations/src/free_join/index_stats.rs | 239 ++++++++++++++ core-relations/src/free_join/mod.rs | 1 + core-relations/src/lib.rs | 2 + src/cli.rs | 3 + 8 files changed, 544 insertions(+), 47 deletions(-) create mode 100644 core-relations/src/free_join/index_stats.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index f9087421e..622e2e70d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,7 @@ ## [Unreleased] - ReleaseDate +- Skip building per-subset join indexes for probers that will only be probed once. Join stages now pass a probe-count bound into index acquisition, and single-probe probers answer by scanning the subset directly (a per-trie-node scan budget caps the extra work when a node is probed again later). Most probe-side index builds served exactly one probe, so this removes 60–99% of on-the-fly index-build volume on LLM-compilation benchmarks (whisper −9%, qwen −9%, python_array_optimize −12% end-to-end). A new feature-gated `index-stats` instrumentation (`--features index-stats`) reports per-index build/usage histograms at CLI exit. - Speed up query evaluation by building on-the-fly per-subset column indexes as sorted arrays (`SortedColumnIndex`) instead of hash maps. These indexes are typically iterated once and probed a bounded number of times over high-cardinality columns, so skipping hash-table construction is a large win (e.g. ~33% faster on the `gemma` benchmark). - Add `make nightly` and `scripts/nightly_bench.py`, a hyperfine-based benchmark harness that measures every `tests/**/*.egg` program at 1/2/4/8 threads and (where supported) in proof-testing mode, caps each run at a 2-minute timeout, skips sub-50ms programs, and emits an HTML dashboard (one row per benchmark, one column per configuration) for nightly.cs.washington.edu. The dashboard uses [eval-live](https://github.com/oflatt/eval-live) for interactive filtering and sorting. - Add typed `EGraph` extension state that clones with `EGraph` and is restored by `push`/`pop`. diff --git a/Cargo.toml b/Cargo.toml index 7df6396e5..9e457ee66 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -130,6 +130,8 @@ bin = [ ] serde = ["egraph-serialize/serde"] graphviz = ["egraph-serialize/graphviz"] +# Report join index build/usage statistics at CLI exit. +index-stats = ["egglog-core-relations/index-stats"] [dependencies] clap = { workspace = true, features = ["derive"], optional = true } diff --git a/core-relations/Cargo.toml b/core-relations/Cargo.toml index e73e2b4cf..2b822a867 100644 --- a/core-relations/Cargo.toml +++ b/core-relations/Cargo.toml @@ -49,3 +49,5 @@ harness = false [features] default = [] debug-val-trace = [] +# Collect and report statistics on join index builds and their usage. +index-stats = [] diff --git a/core-relations/src/free_join/execute.rs b/core-relations/src/free_join/execute.rs index b58d45faa..fce7c1a58 100644 --- a/core-relations/src/free_join/execute.rs +++ b/core-relations/src/free_join/execute.rs @@ -8,6 +8,7 @@ use std::{ use crate::{ common::{HashMap, IndexMap}, + free_join::index_stats::{BuildStats, IndexKind}, free_join::plan::{JoinStages, MatId, MatScanMode, MatSpec}, numeric_id::{DenseIdMap, IdVec, NumericId}, query::Atom, @@ -47,12 +48,22 @@ use super::{ const SMALL_RESIDUAL: usize = 8; +/// Probers expecting at most this many probes answer each probe by scanning +/// the subset instead of building a single-column index (see +/// [`DynamicIndex::ScanColumn`]). A single scan is one pass with no +/// allocation, while a build is a collection pass plus a (cheap, for small +/// subsets) sort — so scanning only wins for probers probed exactly once. +/// Benchmarks confirm higher limits regress: most scan-eligible probers +/// receive one probe, and at two or more the build already pays off. +const SCAN_PROBE_LIMIT: usize = 1; + struct SparseColumnIndex { n_keys: usize, n_subsets: usize, keys: [Value; SMALL_RESIDUAL], offsets: [usize; SMALL_RESIDUAL], subset_ids: [RowId; SMALL_RESIDUAL], + stats: BuildStats, } /// Return a SubsetRef for the given range of rows in a SparseColumnIndex. @@ -121,6 +132,7 @@ impl SparseColumnIndex { keys, offsets, subset_ids, + stats: BuildStats::new(IndexKind::SparseColumn, n_subsets), } } @@ -187,6 +199,7 @@ pub(crate) struct SortedColumnIndex { keys: Vec<(Value, u32)>, /// Row ids grouped by key; each group is ascending. row_ids: Vec, + stats: BuildStats, } impl SortedColumnIndex { @@ -209,7 +222,11 @@ impl SortedColumnIndex { row_ids.push(row); } keys.push((Value::new_const(0), row_ids.len() as u32)); - SortedColumnIndex { keys, row_ids } + SortedColumnIndex { + keys, + row_ids, + stats: BuildStats::new(IndexKind::SortedColumn, n), + } } fn get_subset(&self, key: Value) -> Option> { @@ -242,7 +259,7 @@ impl SortedColumnIndex { } } -enum DynamicIndex { +enum DynamicIndex<'a> { Cached { /// When Some(range), intersect each subset from the index with this dense range. /// The range is the Dense outer subset known at Prober construction time. @@ -255,9 +272,21 @@ enum DynamicIndex { intersect_outer: Option, table: HashColumnIndex, }, - Dynamic(TupleIndex), + Dynamic { + table: TupleIndex, + stats: BuildStats, + }, DynamicColumn(Arc), SparseColumn(SparseColumnIndex), + /// No index at all: each probe answers by scanning the trie node's subset + /// for rows matching the probed value. Chosen by [`JoinState::get_index`] + /// when the caller's `probe_hint` shows the prober will be probed at most + /// [`SCAN_PROBE_LIMIT`] times, which makes an index build unprofitable. + ScanColumn { + info: &'a TableInfo, + col: ColumnId, + stats: BuildStats, + }, } /// This struct is used to mark subsets that can contain non-stale entries. @@ -298,6 +327,42 @@ impl PotentiallyStale> { } } +/// The result of a point probe on a [`Prober`]: either a view into an index's +/// group of rows, or a subset the prober had to materialize (e.g. by scanning). +enum ProbedSubset<'a> { + Borrowed(SubsetRef<'a>), + Owned(Subset), +} + +impl ProbedSubset<'_> { + fn size(&self) -> usize { + match self { + ProbedSubset::Borrowed(s) => s.size(), + ProbedSubset::Owned(s) => s.size(), + } + } + + fn into_owned(self, pool: &Pool) -> Subset { + match self { + ProbedSubset::Borrowed(s) => s.to_owned(pool), + ProbedSubset::Owned(s) => s, + } + } +} + +impl PotentiallyStale> { + fn size(&self) -> usize { + self.inner.size() + } + + fn into_owned(self, pool: &Pool) -> PotentiallyStale { + PotentiallyStale { + inner: self.inner.into_owned(pool), + can_be_stale: self.can_be_stale, + } + } +} + /// Intersect a `SubsetRef` with a dense `OffsetRange` and return the result as a /// borrowed `SubsetRef`, or `None` if the intersection is empty. /// @@ -328,13 +393,13 @@ fn intersect_with_dense_ref<'a>(v: SubsetRef<'a>, range: OffsetRange) -> Option< } } -struct Prober { +struct Prober<'a> { node: Arc, - ix: DynamicIndex, + ix: DynamicIndex<'a>, } -impl Prober { - fn get_subset<'a>(&'a self, key: &'a [Value]) -> Option>> { +impl Prober<'_> { + fn get_subset<'s>(&'s self, key: &'s [Value]) -> Option>> { match &self.ix { DynamicIndex::Cached { intersect_outer, @@ -346,7 +411,9 @@ impl Prober { } else { subset_ref }; - Some(PotentiallyStale::maybe_stale(subset)) + Some(PotentiallyStale::maybe_stale(ProbedSubset::Borrowed( + subset, + ))) } DynamicIndex::CachedColumn { intersect_outer, @@ -359,15 +426,44 @@ impl Prober { } else { subset_ref }; - Some(PotentiallyStale::maybe_stale(subset)) + Some(PotentiallyStale::maybe_stale(ProbedSubset::Borrowed( + subset, + ))) + } + DynamicIndex::Dynamic { table, stats } => { + stats.record_probe(); + let subset = table.get_subset(key)?; + Some(PotentiallyStale::not_stale(ProbedSubset::Borrowed(subset))) } - DynamicIndex::Dynamic(tab) => tab.get_subset(key).map(PotentiallyStale::not_stale), DynamicIndex::DynamicColumn(tab) => { - tab.get_subset(key[0]).map(PotentiallyStale::not_stale) + tab.stats.record_probe(); + let subset = tab.get_subset(key[0])?; + Some(PotentiallyStale::not_stale(ProbedSubset::Borrowed(subset))) } DynamicIndex::SparseColumn(tab) => { debug_assert_eq!(key.len(), 1); - tab.get_subset(key[0]).map(PotentiallyStale::not_stale) + tab.stats.record_probe(); + let subset = tab.get_subset(key[0])?; + Some(PotentiallyStale::not_stale(ProbedSubset::Borrowed(subset))) + } + DynamicIndex::ScanColumn { info, col, stats } => { + debug_assert_eq!(key.len(), 1); + stats.record_probe(); + let mut result = Subset::empty(); + info.table.as_ref().for_each_col( + self.node.subset.as_ref(), + *col, + &mut |row, val| { + if val == key[0] { + result.add_row_sorted(row); + } + }, + ); + if result.size() == 0 { + return None; + } + // Scans (like on-the-fly index builds) never see stale rows. + Some(PotentiallyStale::not_stale(ProbedSubset::Owned(result))) } } } @@ -411,15 +507,26 @@ impl Prober { .unwrap() .for_each(|k, v| f(&[*k], PotentiallyStale::maybe_stale(v))); } - DynamicIndex::Dynamic(tab) => { - tab.for_each(|k, v| f(k, PotentiallyStale::not_stale(v))); + DynamicIndex::Dynamic { table, stats } => { + stats.record_iter(); + table.for_each(|k, v| f(k, PotentiallyStale::not_stale(v))); + } + DynamicIndex::DynamicColumn(tab) => { + tab.stats.record_iter(); + tab.for_each(|k, v| { + f(&[k], PotentiallyStale::not_stale(v)); + }) } - DynamicIndex::DynamicColumn(tab) => tab.for_each(|k, v| { - f(&[k], PotentiallyStale::not_stale(v)); - }), DynamicIndex::SparseColumn(tab) => { + tab.stats.record_iter(); tab.for_each(|k, v| f(k, PotentiallyStale::not_stale(v))); } + DynamicIndex::ScanColumn { .. } => { + unreachable!( + "Iterating requires grouping rows by value, i.e. an actual index. + Callers that request a scan-based prober are not expected to iterate it" + ) + } } } @@ -427,9 +534,13 @@ impl Prober { match &self.ix { DynamicIndex::Cached { table, .. } => table.get().unwrap().len(), DynamicIndex::CachedColumn { table, .. } => table.get().unwrap().len(), - DynamicIndex::Dynamic(tab) => tab.len(), + DynamicIndex::Dynamic { table, .. } => table.len(), DynamicIndex::DynamicColumn(tab) => tab.len(), DynamicIndex::SparseColumn(tab) => tab.len(), + // See `for_each`: not expected for scan-based probers. + DynamicIndex::ScanColumn { .. } => { + unreachable!("Unexpected iteration of a scan-based prober") + } } } } @@ -772,6 +883,12 @@ pub(crate) struct TrieNode { /// only ever probed with a single column, so we store one (col, map) pair /// instead of an IdVec across all columns. cached_children: OnceLock>, + /// How many more index-free scan probes ([`DynamicIndex::ScanColumn`]) + /// this node's subset may serve. Starts at [`SCAN_PROBE_LIMIT`]; when a + /// reservation no longer fits, `get_index` builds a real index instead. + /// This bounds the total scan work per node even when the node is probed + /// again across many join frames. + scan_budget: AtomicUsize, } impl std::fmt::Debug for TrieNode { @@ -788,12 +905,24 @@ impl TrieNode { subset, cached_subsets: Default::default(), cached_children: Default::default(), + scan_budget: AtomicUsize::new(SCAN_PROBE_LIMIT), } } fn size(&self) -> usize { self.subset.size() } + + /// Try to reserve `probes` scan probes from this node's budget, returning + /// whether the reservation fits. See [`TrieNode::scan_budget`]. + fn try_reserve_scans(&self, probes: usize) -> bool { + use std::sync::atomic::Ordering; + self.scan_budget + .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |budget| { + budget.checked_sub(probes) + }) + .is_ok() + } fn get_cached_index(&self, col: ColumnId, info: &TableInfo) -> Arc { self.cached_subsets.get_or_init(|| { // Pre-size the vector so we do not need to borrow it mutably to initialize the index. @@ -811,6 +940,11 @@ impl TrieNode { .clone() } + /// The cached index for `col`, if one has already been built. + fn peek_cached_index(&self, col: ColumnId) -> Option> { + self.cached_subsets.get()?[col].get().cloned() + } + fn get_cached_trie_node( &self, col: ColumnId, @@ -872,6 +1006,7 @@ impl BindingInfo { node.cached_subsets.take(); node.cached_children.take(); node.subset = subset; + *node.scan_budget.get_mut() = SCAN_PROBE_LIMIT; return; } self.subsets.insert(atom, Arc::new(TrieNode::new(subset))); @@ -909,19 +1044,26 @@ impl<'a> JoinState<'a> { } } + /// Build a [`Prober`] for the given atom's current subset over `cols`. + /// + /// `probe_hint`, when known by the caller, is an upper bound on the number + /// of point probes ([`Prober::get_subset`] calls) this prober will + /// receive. A small bound lets us skip building an index entirely and + /// answer each probe with a scan. fn get_index( &self, atoms: &Arc>, atom: AtomId, binding_info: &mut BindingInfo, cols: impl Iterator, - ) -> Prober { + probe_hint: Option, + ) -> Prober<'a> { let cols = SmallVec::<[ColumnId; 4]>::from_iter(cols); let trie_node = binding_info.subsets.unwrap_val(atom); let subset = &trie_node.subset; let table_id = atoms[atom].table; - let info = &self.db.tables[table_id]; + let info: &'a TableInfo = &self.db.tables[table_id]; let dyn_index = if subset.size() <= SMALL_RESIDUAL && cols.len() == 1 { DynamicIndex::SparseColumn(SparseColumnIndex::new( info.table.as_ref(), @@ -965,7 +1107,23 @@ impl<'a> JoinState<'a> { } } else if cols.len() != 1 { // NB: we should have a caching strategy for non-column indexes. - DynamicIndex::Dynamic(info.table.group_by_key(subset.as_ref(), &cols)) + DynamicIndex::Dynamic { + stats: BuildStats::new(IndexKind::Tuple, subset.size()), + table: info.table.group_by_key(subset.as_ref(), &cols), + } + } else if let Some(index) = trie_node.peek_cached_index(cols[0]) { + // An index for this column already exists; probing it is free. + DynamicIndex::DynamicColumn(index) + } else if probe_hint.is_some_and(|probes| trie_node.try_reserve_scans(probes)) { + // Few enough probes are coming that scanning per probe beats + // building an index. (Reservations larger than the node's + // starting budget never succeed, so this also enforces + // `probes <= SCAN_PROBE_LIMIT`.) + DynamicIndex::ScanColumn { + info, + col: cols[0], + stats: BuildStats::new(IndexKind::Scan, subset.size()), + } } else { DynamicIndex::DynamicColumn(trie_node.get_cached_index(cols[0], info)) } @@ -981,8 +1139,9 @@ impl<'a> JoinState<'a> { binding_info: &mut BindingInfo, atom: AtomId, col: ColumnId, - ) -> Prober { - self.get_index(atoms, atom, binding_info, iter::once(col)) + probe_hint: Option, + ) -> Prober<'a> { + self.get_index(atoms, atom, binding_info, iter::once(col), probe_hint) } /// Runs the free join plan, starting with the header. @@ -1207,7 +1366,7 @@ impl<'a> JoinState<'a> { if binding_info.has_empty_subset(a.atom) { return; } - let prober = self.get_column_index(atoms, binding_info, a.atom, a.column); + let prober = self.get_column_index(atoms, binding_info, a.atom, a.column, None); let info = &self.db.tables[atoms[a.atom].table]; let table = info.table.as_ref(); let has_stale = table.has_stale_rows(); @@ -1243,14 +1402,43 @@ impl<'a> JoinState<'a> { binding_info.move_back(a.atom, prober); } [a, b] => { - let a_prober = self.get_column_index(atoms, binding_info, a.atom, a.column); - let b_prober = self.get_column_index(atoms, binding_info, b.atom, b.column); - - let ((smaller, smaller_scan), (larger, larger_scan)) = - if a_prober.len() < b_prober.len() { - ((&a_prober, a), (&b_prober, b)) + // Index the side with the smaller subset first: its key + // count is an exact bound on how many times the other + // side would be probed. When that bound is small, + // `get_index` skips building the other side's index and + // iterating the first side is trivially the right order; + // otherwise both sides have indexes and the iteration + // side is picked by exact key count, which minimizes + // probe misses. + let (first_scan, second_scan) = if binding_info.subsets[a.atom].size() + <= binding_info.subsets[b.atom].size() + { + (a, b) + } else { + (b, a) + }; + let first = self.get_column_index( + atoms, + binding_info, + first_scan.atom, + first_scan.column, + None, + ); + let first_len = first.len(); + let second = self.get_column_index( + atoms, + binding_info, + second_scan.atom, + second_scan.column, + Some(first_len), + ); + let (smaller, larger, smaller_scan, larger_scan) = + // We only consider swapping when `second` is not a ScanColumn, + // which is meant to be scanned. + if first_len > SCAN_PROBE_LIMIT && second.len() < first_len { + (second, first, second_scan, first_scan) } else { - ((&b_prober, b), (&a_prober, a)) + (first, second, first_scan, second_scan) }; let smaller_atom = smaller_scan.atom; @@ -1299,7 +1487,7 @@ impl<'a> JoinState<'a> { } if large_sub.size() <= 16 { let large_sub = refine_subset( - large_sub.to_owned(pool), + large_sub.into_owned(pool), &larger_scan.cs, &large_table, large_has_stale, @@ -1316,7 +1504,7 @@ impl<'a> JoinState<'a> { large_info, || { refine_subset( - large_sub.to_owned(pool), + large_sub.into_owned(pool), &larger_scan.cs, &large_table, large_has_stale, @@ -1337,23 +1525,60 @@ impl<'a> JoinState<'a> { }); drain_updates!(updates); - binding_info.move_back(a.atom, a_prober); - binding_info.move_back(b.atom, b_prober); + binding_info.move_back(smaller_atom, smaller); + binding_info.move_back(larger_atom, larger); } rest => { - let mut smallest = 0; - let mut smallest_size = usize::MAX; + // As in the two-scan case: index the atom with the + // smallest subset first. When its key count is small it + // leads the scan and bounds the number of probes every + // other atom receives; otherwise all atoms get indexes + // and the lead is picked by exact key count. + let mut min_subset_idx = 0; + let mut min_subset_size = usize::MAX; + for (i, scan) in rest.iter().enumerate() { + let size = binding_info.subsets[scan.atom].size(); + if size < min_subset_size { + min_subset_idx = i; + min_subset_size = size; + } + } + let lead_prober = self.get_column_index( + atoms, + binding_info, + rest[min_subset_idx].atom, + rest[min_subset_idx].column, + None, + ); + let lead_len = lead_prober.len(); + let probe_hint = (lead_len <= SCAN_PROBE_LIMIT).then_some(lead_len); let mut probers = Vec::with_capacity(rest.len()); + let mut smallest = min_subset_idx; + let mut smallest_size = lead_len; for (i, scan) in rest.iter().enumerate() { - let prober = - self.get_column_index(atoms, binding_info, scan.atom, scan.column); - let size = prober.len(); - if size < smallest_size { - smallest = i; - smallest_size = size; + if i == min_subset_idx { + continue; + } + let prober = self.get_column_index( + atoms, + binding_info, + scan.atom, + scan.column, + probe_hint, + ); + // Probers built without a hint have an index; re-pick + // the lead by exact key count among them. + // Only do this when hint is some to avoid calling `len()` on ScanColumn probers. + if probe_hint.is_none() { + let size = prober.len(); + if size < smallest_size { + smallest = i; + smallest_size = size; + } } probers.push(prober); } + probers.insert(min_subset_idx, lead_prober); let main_spec = &rest[smallest]; let main_spec_info = &self.db.tables[atoms[main_spec.atom].table]; @@ -1385,7 +1610,7 @@ impl<'a> JoinState<'a> { self.db.tables[atoms[rest[i].atom].table].table.as_ref(); if sub.size() <= 16 { let sub = refine_subset( - sub.to_owned(pool), + sub.into_owned(pool), &rest[i].cs, &table, rest_has_stale[i], @@ -1402,7 +1627,7 @@ impl<'a> JoinState<'a> { &self.db.tables[atoms[scan.atom].table], || { refine_subset( - sub.to_owned(pool), + sub.into_owned(pool), &rest[i].cs, &table, rest_has_stale[i], @@ -1555,6 +1780,9 @@ impl<'a> JoinState<'a> { if binding_info.has_empty_subset(cover_atom) { return; } + // Each index is probed once per cover row surviving the + // constraints, so the cover's subset size bounds the probes. + let cover_size = binding_info.subsets[cover_atom].size(); let index_probers = to_intersect .iter() .enumerate() @@ -1567,6 +1795,7 @@ impl<'a> JoinState<'a> { spec.to_index.atom, binding_info, spec.to_index.vars.iter().copied(), + Some(cover_size), ), ) }) @@ -1626,7 +1855,7 @@ impl<'a> JoinState<'a> { let table_info = &self.db.tables[atoms[*atom].table]; let cs = &to_intersect[*i].0.constraints; let subset = refine_subset( - subset.to_owned(pool), + subset.into_owned(pool), cs, &table_info.table.as_ref(), index_has_stale[prober_idx], @@ -1743,6 +1972,23 @@ impl<'a> JoinState<'a> { } => { let cover_mat = binding_info.materializations[*cover].clone(); let mut updates = FrameUpdates::with_capacity(cmp::min(chunk_size, cur_size)); + // Each index is probed once per enumerated frame of the + // materialization: one per group in `KeyOnly` mode, one per + // row in `Full` mode, and one per row of the group selected + // by the current bindings in `Value` mode. + let probe_hint = match mode { + MatScanMode::KeyOnly => Some(cover_mat.len()), + MatScanMode::Full => Some(cover_mat.values().map(RowBuffer::len).sum()), + MatScanMode::Value(index_vars) => { + let keys: Vec = index_vars + .iter() + .map(|var| binding_info.bindings[*var]) + .collect(); + Some(cover_mat.get(&keys).map_or(0, RowBuffer::len)) + } + // `Lookup` requires `to_intersect` to be empty: no probers. + MatScanMode::Lookup(_) => None, + }; let probers = to_intersect .iter() .map(|(spec, _)| { @@ -1751,6 +1997,7 @@ impl<'a> JoinState<'a> { spec.to_index.atom, binding_info, spec.to_index.vars.iter().copied(), + probe_hint, ) }) .collect::>(); @@ -1789,7 +2036,7 @@ impl<'a> JoinState<'a> { } if let Some(subset) = prober.get_subset(&key) { let subset = refine_subset( - subset.to_owned(pool), + subset.into_owned(pool), &spec.constraints, &self.db.tables[atoms[spec.to_index.atom].table] .table diff --git a/core-relations/src/free_join/index_stats.rs b/core-relations/src/free_join/index_stats.rs new file mode 100644 index 000000000..5b228324c --- /dev/null +++ b/core-relations/src/free_join/index_stats.rs @@ -0,0 +1,239 @@ +//! Feature-gated statistics on per-subset join indexes: how many rows each +//! index build covered and how many times the index was used before being +//! dropped. +//! +//! With the `index-stats` feature enabled, every dynamically-built join index +//! (see `DynamicIndex` in `execute.rs`) carries a [`BuildStats`] that counts +//! point probes (`get_subset`) and whole-index iterations (`for_each`). When +//! the index is dropped, the totals are folded into global histograms keyed by +//! index kind, usage class, build size, and use count; [`report`] renders +//! them. Without the feature, [`BuildStats`] is a zero-sized no-op. + +/// The kind of index a [`BuildStats`] describes. +#[derive(Copy, Clone, Debug)] +pub(crate) enum IndexKind { + /// `SortedColumnIndex`: single-column sort-based index, cached per trie node. + SortedColumn = 0, + /// `TupleIndex` from `group_by_key`: multi-column hash index, never cached. + Tuple = 1, + /// `SparseColumnIndex`: stack-allocated index over at most 8 rows. + SparseColumn = 2, + /// `ScanColumn` prober: no index at all; each probe scans the subset. + /// "Rows" counts the subset size once (each probe passes over it). + Scan = 3, +} + +#[cfg(feature = "index-stats")] +pub(crate) use imp::BuildStats; +#[cfg(feature = "index-stats")] +pub use imp::report; + +#[cfg(not(feature = "index-stats"))] +pub(crate) use noop::BuildStats; + +#[cfg(not(feature = "index-stats"))] +mod noop { + use super::IndexKind; + + /// No-op stand-in for the instrumented version; see the module docs. + pub(crate) struct BuildStats; + + impl BuildStats { + #[inline(always)] + pub(crate) fn new(_kind: IndexKind, _rows: usize) -> Self { + BuildStats + } + + #[inline(always)] + pub(crate) fn record_probe(&self) {} + + #[inline(always)] + pub(crate) fn record_iter(&self) {} + } +} + +#[cfg(feature = "index-stats")] +mod imp { + use std::fmt::Write; + use std::sync::atomic::{AtomicU64, Ordering}; + + use super::IndexKind; + + const KINDS: usize = 4; + const KIND_NAMES: [&str; KINDS] = ["SortedColumn", "Tuple", "SparseColumn", "Scan"]; + + /// How a build was used over its lifetime. + const CLASSES: usize = 4; + const CLASS_NAMES: [&str; CLASSES] = ["unused", "probe-only", "iter-only", "mixed"]; + + /// Build sizes bucketed by floor(log2(rows)): bucket `i` covers + /// `[2^i, 2^(i+1))` rows, with 0 and 1 sharing bucket 0. + const ROW_BUCKETS: usize = 24; + + /// Use counts: 0..=3 exactly, then power-of-two ranges. + const USE_BUCKETS: usize = 16; + const USE_LABELS: [&str; USE_BUCKETS] = [ + "0", "1", "2", "3", "4-7", "8-15", "16-31", "32-63", "64-127", "128-255", "256-511", + "512-1023", "1K-2K", "2K-4K", "4K-8K", ">=8K", + ]; + + fn row_bucket(rows: u64) -> usize { + (63 - rows.max(1).leading_zeros() as usize).min(ROW_BUCKETS - 1) + } + + fn use_bucket(uses: u64) -> usize { + if uses < 4 { + uses as usize + } else { + (2 + (63 - uses.leading_zeros() as usize)).min(USE_BUCKETS - 1) + } + } + + type Histogram = [[[[AtomicU64; USE_BUCKETS]; ROW_BUCKETS]; CLASSES]; KINDS]; + + // Interior mutability and size are both intentional: this is only the + // zero-initializer for the two static histograms below. + #[allow(clippy::declare_interior_mutable_const, clippy::large_const_arrays)] + const EMPTY_HISTOGRAM: Histogram = [const { + [const { [const { [const { AtomicU64::new(0) }; USE_BUCKETS] }; ROW_BUCKETS] }; CLASSES] + }; KINDS]; + + /// Histogram cell arrays, indexed by `[kind][class][row_bucket][use_bucket]`. + static BUILDS: Histogram = EMPTY_HISTOGRAM; + static ROWS: Histogram = EMPTY_HISTOGRAM; + + /// Exact per-(kind, class) totals: builds, rows, probes, iters. + static TOTALS: [[[AtomicU64; 4]; CLASSES]; KINDS] = + [const { [const { [const { AtomicU64::new(0) }; 4] }; CLASSES] }; KINDS]; + + /// Per-build usage counters attached to a single join index. Dropping the + /// value folds its counts into the global histograms. + pub(crate) struct BuildStats { + kind: IndexKind, + rows: u64, + probes: AtomicU64, + iters: AtomicU64, + } + + impl BuildStats { + pub(crate) fn new(kind: IndexKind, rows: usize) -> Self { + BuildStats { + kind, + rows: rows as u64, + probes: AtomicU64::new(0), + iters: AtomicU64::new(0), + } + } + + #[inline] + pub(crate) fn record_probe(&self) { + self.probes.fetch_add(1, Ordering::Relaxed); + } + + #[inline] + pub(crate) fn record_iter(&self) { + self.iters.fetch_add(1, Ordering::Relaxed); + } + } + + impl Drop for BuildStats { + fn drop(&mut self) { + let probes = *self.probes.get_mut(); + let iters = *self.iters.get_mut(); + let class = match (probes > 0, iters > 0) { + (false, false) => 0, + (true, false) => 1, + (false, true) => 2, + (true, true) => 3, + }; + let kind = self.kind as usize; + let rb = row_bucket(self.rows); + let ub = use_bucket(probes + iters); + BUILDS[kind][class][rb][ub].fetch_add(1, Ordering::Relaxed); + ROWS[kind][class][rb][ub].fetch_add(self.rows, Ordering::Relaxed); + let totals = &TOTALS[kind][class]; + totals[0].fetch_add(1, Ordering::Relaxed); + totals[1].fetch_add(self.rows, Ordering::Relaxed); + totals[2].fetch_add(probes, Ordering::Relaxed); + totals[3].fetch_add(iters, Ordering::Relaxed); + } + } + + fn fmt_count(n: u64) -> String { + if n >= 10_000_000_000 { + format!("{:.1}G", n as f64 / 1e9) + } else if n >= 10_000_000 { + format!("{:.1}M", n as f64 / 1e6) + } else if n >= 10_000 { + format!("{:.1}K", n as f64 / 1e3) + } else { + n.to_string() + } + } + + fn row_label(bucket: usize) -> String { + let lo = if bucket == 0 { 0u64 } else { 1 << bucket }; + if bucket == ROW_BUCKETS - 1 { + format!(">={}", fmt_count(lo)) + } else { + format!("{}-{}", fmt_count(lo), fmt_count((1 << (bucket + 1)) - 1)) + } + } + + fn render_table( + out: &mut String, + title: &str, + cells: &[[AtomicU64; USE_BUCKETS]; ROW_BUCKETS], + ) { + writeln!(out, " {title} (rows \\ uses):").unwrap(); + write!(out, " {:>12}", "").unwrap(); + for label in USE_LABELS { + write!(out, " {label:>8}").unwrap(); + } + writeln!(out).unwrap(); + for (rb, row) in cells.iter().enumerate() { + if row.iter().all(|c| c.load(Ordering::Relaxed) == 0) { + continue; + } + write!(out, " {:>12}", row_label(rb)).unwrap(); + for cell in row { + let v = cell.load(Ordering::Relaxed); + if v == 0 { + write!(out, " {:>8}", ".").unwrap(); + } else { + write!(out, " {:>8}", fmt_count(v)).unwrap(); + } + } + writeln!(out).unwrap(); + } + } + + /// Render the histograms collected so far. + pub fn report() -> String { + let mut out = String::new(); + writeln!(out, "=== join index build/usage stats ===").unwrap(); + for kind in 0..KINDS { + for class in 0..CLASSES { + let totals = &TOTALS[kind][class]; + let builds = totals[0].load(Ordering::Relaxed); + if builds == 0 { + continue; + } + writeln!( + out, + "\n{} / {}: {} builds, {} rows built, {} probes, {} iters", + KIND_NAMES[kind], + CLASS_NAMES[class], + fmt_count(builds), + fmt_count(totals[1].load(Ordering::Relaxed)), + fmt_count(totals[2].load(Ordering::Relaxed)), + fmt_count(totals[3].load(Ordering::Relaxed)), + ) + .unwrap(); + render_table(&mut out, "builds", &BUILDS[kind][class]); + render_table(&mut out, "rows built", &ROWS[kind][class]); + } + } + out + } +} diff --git a/core-relations/src/free_join/mod.rs b/core-relations/src/free_join/mod.rs index 97417988f..db993a8a5 100644 --- a/core-relations/src/free_join/mod.rs +++ b/core-relations/src/free_join/mod.rs @@ -38,6 +38,7 @@ use crate::action::ExecutionState; pub(crate) mod execute; pub(crate) mod frame_update; +pub(crate) mod index_stats; pub(crate) mod plan; define_id!( diff --git a/core-relations/src/lib.rs b/core-relations/src/lib.rs index 1adf0c094..ce4084536 100644 --- a/core-relations/src/lib.rs +++ b/core-relations/src/lib.rs @@ -27,6 +27,8 @@ pub use action::{ExecutionState, MergeVal, QueryEntry, WriteVal}; pub use base_values::{BaseValue, BaseValueId, BaseValuePrinter, BaseValues, Boxed}; pub use common::Value; pub use containers::{ContainerRebuildSummary, ContainerValue, ContainerValueId, ContainerValues}; +#[cfg(feature = "index-stats")] +pub use free_join::index_stats::report as index_stats_report; pub use free_join::{ AtomId, CounterId, Database, ExternalFunction, ExternalFunctionId, TableId, Variable, make_external_func, plan::PlanStrategy, diff --git a/src/cli.rs b/src/cli.rs index 4c75768aa..15c3bc51f 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -198,6 +198,9 @@ pub fn cli(mut egraph: EGraph) { log::info!("Saved report to {report_path:?}"); } + #[cfg(feature = "index-stats")] + eprintln!("{}", egglog_core_relations::index_stats_report()); + // no need to drop the egraph if we are going to exit std::mem::forget(egraph) } From 64e8fe3953d03e3219d1898ffdeee7dfc5d4c8d2 Mon Sep 17 00:00:00 2001 From: Yihong Zhang Date: Sat, 18 Jul 2026 15:33:59 -0700 Subject: [PATCH 2/2] interleave run --- scripts/bench.py | 250 ++++++++++++++++++++++++++++++----------------- 1 file changed, 159 insertions(+), 91 deletions(-) diff --git a/scripts/bench.py b/scripts/bench.py index 055fc65db..a279cb7af 100644 --- a/scripts/bench.py +++ b/scripts/bench.py @@ -1,4 +1,3 @@ - #!/usr/bin/env python3 """ bench.py — egglog benchmark harness @@ -8,17 +7,28 @@ run --compare-unstaged — benchmark HEAD vs current working-tree changes diff — print the most recently saved diff report +Each state is built in its own temporary git worktree (the unstaged state +builds in the main checkout), so the checkout you are working in is never +touched. The two binaries' runs are interleaved per benchmark — alternating +back-to-back, swapping which binary goes first every round — so slow drift in +machine state (thermals, cache pressure) hits both states equally instead of +biasing whichever was measured last. + +Temporary worktrees are removed on exit, including on SIGINT/SIGTERM; ones +orphaned by a hard kill (SIGKILL) are swept up by the next invocation. + Improvement policy: IMPROVEMENT iff at least one benchmark improved >=3% AND net average change is negative. """ -import json -import os -import shlex +import atexit import shutil +import signal +import statistics import subprocess import sys import tempfile +import time from datetime import datetime from pathlib import Path @@ -26,7 +36,6 @@ REPO_ROOT = SCRIPT_DIR.parent BENCH_DIR = REPO_ROOT / "benchmarks" TEST_DIR = REPO_ROOT / "tests" -EGGLOG = REPO_ROOT / "target" / "release" / "egglog" BENCHMARKS = [ "hardboiled_conv1d_32.egg", @@ -44,7 +53,7 @@ "whisper.egg", ] -HYPERFINE_RUNS = 15 +TIMED_RUNS = 15 WARMUP_RUNS = 3 IMPROVE_THRESHOLD = 3.0 UNCHANGED_BAND = 0.5 @@ -65,32 +74,109 @@ def _resolve(ref: str) -> tuple[str, str, str]: return full, short, subject -def _current_ref() -> str: - """Branch name if on a branch, else full hash (detached HEAD).""" - try: - return _git("symbolic-ref", "--short", "HEAD") - except subprocess.CalledProcessError: - return _git("rev-parse", "HEAD") +# ── temporary worktrees ────────────────────────────────────────────────────── +_WORKTREE_PREFIX = "egglog-bench-" +_temp_worktrees: list[Path] = [] -# ── benchmarking ───────────────────────────────────────────────────────────── + +def _remove_worktree(path: Path) -> None: + subprocess.run( + ["git", "-C", str(REPO_ROOT), "worktree", "remove", "--force", str(path)], + capture_output=True, + ) + # `git worktree remove` refuses in some states (e.g. a build was killed + # mid-write); make sure the directory is gone either way, then drop any + # dangling registration. + shutil.rmtree(path, ignore_errors=True) + subprocess.run( + ["git", "-C", str(REPO_ROOT), "worktree", "prune"], + capture_output=True, + ) -def benchmark_state(label: str) -> dict[str, tuple[float, float]]: - """Build the current working tree and hyperfine all benchmarks. +def _cleanup_worktrees() -> None: + while _temp_worktrees: + _remove_worktree(_temp_worktrees.pop()) - Returns {bench_name: (mean_s, stddev_s)}. - """ + +def _sweep_stale_worktrees() -> None: + """Remove bench worktrees left behind by a previous killed run.""" + out = _git("worktree", "list", "--porcelain") + for line in out.splitlines(): + if line.startswith("worktree "): + path = Path(line[len("worktree ") :]) + if path.name.startswith(_WORKTREE_PREFIX): + print(f" Removing stale bench worktree {path}") + _remove_worktree(path) + + +def _add_worktree(commit: str, label: str) -> Path: + """Check `commit` out into a fresh temporary worktree and return its path.""" + path = Path(tempfile.mkdtemp(prefix=f"{_WORKTREE_PREFIX}{label}-")) + path.rmdir() # `git worktree add` wants to create the directory itself. + _temp_worktrees.append(path) + _git("worktree", "add", "--detach", str(path), commit) + return path + + +def _exit_on_signal(signum: int, _frame) -> None: + # Turn the signal into a normal exit so `finally` blocks and the atexit + # worktree cleanup run. + sys.exit(128 + signum) + + +atexit.register(_cleanup_worktrees) +for _sig in (signal.SIGINT, signal.SIGTERM, getattr(signal, "SIGHUP", None)): + if _sig is not None: + signal.signal(_sig, _exit_on_signal) + + +# ── benchmarking ───────────────────────────────────────────────────────────── + + +def _build(manifest_dir: Path, label: str) -> Path: + """Build egglog (release) in `manifest_dir` and return the binary path.""" print(f"\n [{label}] Building egglog (release)...") subprocess.run( - ["cargo", "build", "--release", "--manifest-path", str(REPO_ROOT / "Cargo.toml")], + ["cargo", "build", "--release", "--manifest-path", str(manifest_dir / "Cargo.toml")], check=True, ) + return manifest_dir / "target" / "release" / "egglog" + - results: dict[str, tuple[float, float]] = {} - print(f"\n [{label}] Benchmarking ({HYPERFINE_RUNS} runs each)\n") - print(f" {'Benchmark':<42} {'Mean (s)':>9} {'± Stddev':>9}") - print(" " + "─" * 62) +def _time_run(binary: Path, src: Path) -> float | None: + """One timed run; returns wall-clock seconds, or None if the run failed.""" + start = time.perf_counter() + proc = subprocess.run( + [str(binary), str(src)], + cwd=REPO_ROOT, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + elapsed = time.perf_counter() - start + return elapsed if proc.returncode == 0 else None + + +def benchmark_interleaved( + label1: str, + bin1: Path, + label2: str, + bin2: Path, +) -> tuple[dict[str, tuple[float, float]], dict[str, tuple[float, float]]]: + """Time all benchmarks, interleaving runs of the two binaries. + + Both binaries run against the main checkout's `tests/` directory (same + inputs for both states) with the repo root as working directory. + + Returns ({bench: (mean_s, stddev_s)}, {bench: (mean_s, stddev_s)}). + """ + results1: dict[str, tuple[float, float]] = {} + results2: dict[str, tuple[float, float]] = {} + + print(f"\n Benchmarking ({TIMED_RUNS} interleaved runs each, {WARMUP_RUNS} warmup)\n") + print(f" {'Benchmark':<42} {label1:>15} {label2:>15}") + print(" " + "─" * 76) for bench in BENCHMARKS: src = TEST_DIR / bench @@ -98,43 +184,40 @@ def benchmark_state(label: str) -> dict[str, tuple[float, float]]: print(f" SKIP {bench}") continue - shell_cmd = shlex.join([str(EGGLOG), str(src)]) + " >/dev/null 2>&1" - with tempfile.NamedTemporaryFile(suffix=".json", delete=False) as tf: - tmp_json = tf.name - - try: - proc = subprocess.run( - [ - "hyperfine", - "--warmup", - str(WARMUP_RUNS), - "--runs", - str(HYPERFINE_RUNS), - "--export-json", - tmp_json, - shell_cmd, - ], - capture_output=True, - text=True, - ) - if proc.returncode != 0: - print(f" FAIL {bench}") - if proc.stderr: - print(" " + proc.stderr.splitlines()[0]) - continue - - with open(tmp_json) as f: - hf = json.load(f) + times1: list[float] = [] + times2: list[float] = [] + failed = False + + for _ in range(WARMUP_RUNS): + if _time_run(bin1, src) is None or _time_run(bin2, src) is None: + failed = True + break + + for i in range(TIMED_RUNS): + if failed: + break + # Swap the order every round so neither binary always runs first. + pair = [(bin1, times1), (bin2, times2)] + if i % 2: + pair.reverse() + for binary, acc in pair: + t = _time_run(binary, src) + if t is None: + failed = True + break + acc.append(t) + + if failed: + print(f" FAIL {bench}") + continue - r = hf["results"][0] - mean = r["mean"] - stddev = r.get("stddev", 0.0) - results[bench] = (mean, stddev) - print(f" {bench:<42} {mean:>9.3f} ±{stddev:>8.3f}") - finally: - os.unlink(tmp_json) + mean1, stddev1 = statistics.mean(times1), statistics.stdev(times1) + mean2, stddev2 = statistics.mean(times2), statistics.stdev(times2) + results1[bench] = (mean1, stddev1) + results2[bench] = (mean2, stddev2) + print(f" {bench:<42} {mean1:>7.3f} ±{stddev1:>6.3f} {mean2:>7.3f} ±{stddev2:>6.3f}") - return results + return results1, results2 # ── diff formatting ─────────────────────────────────────────────────────────── @@ -244,37 +327,27 @@ def _save_diff(report: str, slug: str, timestamp: str) -> Path: def cmd_run(argv: list[str]) -> None: - if not shutil.which("hyperfine"): - sys.exit("hyperfine not found — install with: cargo install hyperfine") - - if argv == ["--compare-unstaged"]: - _run_compare_unstaged() - elif len(argv) == 2: - _run_compare_commits(argv[0], argv[1]) - else: - sys.exit("Usage:\n" " bench.py run \n" " bench.py run --compare-unstaged") + _sweep_stale_worktrees() + try: + if argv == ["--compare-unstaged"]: + _run_compare_unstaged() + elif len(argv) == 2: + _run_compare_commits(argv[0], argv[1]) + else: + sys.exit("Usage:\n" " bench.py run \n" " bench.py run --compare-unstaged") + finally: + _cleanup_worktrees() def _run_compare_commits(ref1: str, ref2: str) -> None: - dirty = subprocess.run( - ["git", "-C", str(REPO_ROOT), "status", "--porcelain", "--untracked-files=no"], - capture_output=True, - text=True, - ).stdout.strip() - if dirty: - sys.exit("Working tree is dirty. Stash or commit your changes first.") - full1, short1, subject1 = _resolve(ref1) full2, short2, subject2 = _resolve(ref2) - orig = _current_ref() - try: - _git("checkout", full1) - times1 = benchmark_state(short1) - _git("checkout", full2) - times2 = benchmark_state(short2) - finally: - _git("checkout", orig) + wt1 = _add_worktree(full1, short1) + wt2 = _add_worktree(full2, short2) + bin1 = _build(wt1, short1) + bin2 = _build(wt2, short2) + times1, times2 = benchmark_interleaved(short1, bin1, short2, bin2) timestamp = datetime.now().strftime("%Y-%m-%dT%H:%M:%S") report = format_diff( @@ -302,14 +375,10 @@ def _run_compare_unstaged() -> None: full_head, short_head, subject_head = _resolve("HEAD") - # Benchmark working-tree state (with changes) first, then HEAD as baseline. - times_wt = benchmark_state("working-tree") - - _git("stash", "push", "--include-untracked", "-m", "bench.py temp stash") - try: - times_head = benchmark_state(short_head) - finally: - _git("stash", "pop") + wt_head = _add_worktree(full_head, short_head) + bin_head = _build(wt_head, short_head) + bin_wt = _build(REPO_ROOT, "working-tree") + times_head, times_wt = benchmark_interleaved(short_head, bin_head, "working-tree", bin_wt) timestamp = datetime.now().strftime("%Y-%m-%dT%H:%M:%S") report = format_diff( @@ -354,4 +423,3 @@ def cmd_diff() -> None: else: print(__doc__) sys.exit(1) -