Fix multi-column ColumnIndex rebuild ordering via per-column radix sort + merge

CHANGELOG.md

@@ -8,6 +8,7 @@
 - Replace the global Rayon thread pool with an `egglog-concurrency` scoped `ThreadPool`; configure parallelism per `EGraph` via `with_num_threads` / `set_num_threads`.
 - Report full source file paths in egglog span and error messages.
 - Fix seminaive matching after nested containers rebuild in place by propagating dirty container ids through parent containers.
+- Fix multi-column secondary index rebuilds so each value's rows come back sorted by row id, and make all rebuild paths (serial, parallel, and bulk) record a row once even when its value repeats across covered columns (#914).
 - Render nullary AST calls without a trailing space, e.g. (foo) instead of (foo ).
 - Add a BigRat to-i64 primitive for integral rationals.
 - Add f64 exp, log, and sqrt primitives.

core-relations/Cargo.toml

@@ -42,6 +42,10 @@ rayon = { workspace = true }
 name = "writes"
 harness = false
 
+[[bench]]
+name = "build_index"
+harness = false
+
 [features]
 default = []
 debug-val-trace = []

core-relations/benches/build_index.rs

@@ -0,0 +1,88 @@
+//! Microbenchmarks for building `ColumnIndex`es.
+//!
+//! Two questions motivate these:
+//!   * `sort_merge_*`: the rebuild path sorts and merges `(Value, RowId)` pairs with hand-written
+//!     radix sort + tournament merge instead of the standard library. These compare the two so we
+//!     can track the speedup rather than assume it.
+//!   * `build_*`: end-to-end index construction, serially vs. in parallel, across subset sizes.
+
+use divan::{Bencher, counter::ItemsCount};
+use egglog_core_relations::bench_support::{
+    IndexInput, gen_blocks, sort_merge_custom, sort_merge_std, with_threads,
+};
+
+fn main() {
+    divan::main()
+}
+
+/// Values are drawn from `rows / DISTINCT_DIVISOR` distinct keys, giving each value several rows.
+const DISTINCT_DIVISOR: u32 = 8;
+
+fn distinct_for(rows: usize) -> u32 {
+    ((rows as u32) / DISTINCT_DIVISOR).max(1)
+}
+
+// -- Sort + merge primitives: custom vs. standard library ---------------------------------------
+
+/// Single column: only the per-block sort runs (no merge). Custom = radix sort.
+#[divan::bench(consts = [4096, 65_536, 1 << 20])]
+fn sort_single_custom<const N: usize>(bench: Bencher) {
+    let (pairs, bounds) = gen_blocks(N, 1, distinct_for(N), 1);
+    bench
+        .with_inputs(|| pairs.clone())
+        .input_counter(|p| ItemsCount::new(p.len()))
+        .bench_values(|p| sort_merge_custom(p, &bounds));
+}
+
+/// Single column, standard `sort_unstable`. No dedup: with one pair per row every `RowId` is
+/// distinct, so no `(Value, RowId)` pair repeats -- matching what the custom path does here.
+#[divan::bench(consts = [4096, 65_536, 1 << 20])]
+fn sort_single_std<const N: usize>(bench: Bencher) {
+    let (pairs, _bounds) = gen_blocks(N, 1, distinct_for(N), 1);
+    bench
+        .with_inputs(|| pairs.clone())
+        .input_counter(|p| ItemsCount::new(p.len()))
+        .bench_values(|p| sort_merge_std(p, false));
+}
+
+/// Four columns: per-block radix sort followed by the tournament merge with dedup.
+#[divan::bench(consts = [4096, 65_536, 1 << 20])]
+fn merge_multi_custom<const N: usize>(bench: Bencher) {
+    let (pairs, bounds) = gen_blocks(N, 4, distinct_for(N), 1);
+    bench
+        .with_inputs(|| pairs.clone())
+        .input_counter(|p| ItemsCount::new(p.len()))
+        .bench_values(|p| sort_merge_custom(p, &bounds));
+}
+
+/// Four columns, standard `sort_unstable` + `dedup` over the whole concatenation (a value can
+/// repeat across a row's columns, so duplicate pairs are possible and must be dropped).
+#[divan::bench(consts = [4096, 65_536, 1 << 20])]
+fn merge_multi_std<const N: usize>(bench: Bencher) {
+    let (pairs, _bounds) = gen_blocks(N, 4, distinct_for(N), 1);
+    bench
+        .with_inputs(|| pairs.clone())
+        .input_counter(|p| ItemsCount::new(p.len()))
+        .bench_values(|p| sort_merge_std(p, true));
+}
+
+// -- End-to-end index construction: serial vs. parallel -----------------------------------------
+
+const N_VAL_COLS: usize = 3;
+
+#[divan::bench(consts = [4096, 65_536, 1 << 20])]
+fn build_serial<const N: usize>(bench: Bencher) {
+    bench
+        .with_inputs(|| IndexInput::random(N, N_VAL_COLS, distinct_for(N), 1))
+        .input_counter(|_| ItemsCount::new(N * N_VAL_COLS))
+        .bench_refs(|input| input.build_serial());
+}
+
+#[divan::bench(consts = [4096, 65_536, 1 << 20])]
+fn build_parallel<const N: usize>(bench: Bencher) {
+    let threads = std::thread::available_parallelism().map_or(1, |n| n.get());
+    bench
+        .with_inputs(|| IndexInput::random(N, N_VAL_COLS, distinct_for(N), 1))
+        .input_counter(|_| ItemsCount::new(N * N_VAL_COLS))
+        .bench_refs(|input| with_threads(threads, || input.build_parallel()));
+}

core-relations/src/hash_index/bench_support.rs

@@ -0,0 +1,138 @@
+//! Benchmark-only entry points into the internal index-construction code.
+//!
+//! Exposed via `#[doc(hidden)] pub` solely so `benches/build_index.rs` can measure the custom
+//! radix-sort + tournament-merge against the standard library and time serial vs parallel
+//! `ColumnIndex` construction. Not a stable API; signatures may change without notice.
+
+use egglog_concurrency::ThreadPool;
+use rand::{Rng, SeedableRng, rngs::StdRng};
+
+use crate::{
+    common::Value,
+    free_join::Database,
+    numeric_id::NumericId,
+    offsets::RowId,
+    table::SortedWritesTable,
+    table_spec::{ColumnId, Table, WrappedTable},
+};
+
+use super::{ColumnIndex, IndexBase, merge_sorted_blocks_dedup, radix_sort_slice_by_value};
+
+/// Generate `n_cols` column blocks of `(Value, RowId)` pairs. Each block is RowId-ascending
+/// (matching a table scan) with values drawn from `0..distinct`, so equal-value runs occur.
+///
+/// Returns the concatenated pairs and the block bounds: block `b` is `[bounds[b], bounds[b + 1])`.
+pub fn gen_blocks(
+    n_rows: usize,
+    n_cols: usize,
+    distinct: u32,
+    seed: u64,
+) -> (Vec<(Value, RowId)>, Vec<usize>) {
+    let mut rng = StdRng::seed_from_u64(seed);
+    let mut pairs = Vec::with_capacity(n_rows * n_cols);
+    let mut bounds = vec![0usize];
+    for _ in 0..n_cols {
+        for r in 0..n_rows {
+            pairs.push((
+                Value::new(rng.random_range(0..distinct)),
+                RowId::from_usize(r),
+            ));
+        }
+        bounds.push(pairs.len());
+    }
+    (pairs, bounds)
+}
+
+/// The custom sort/merge from `rebuild_full`: radix-sort each block by value, then merge the
+/// blocks with dedup (a single block skips the merge).
+pub fn sort_merge_custom(mut pairs: Vec<(Value, RowId)>, bounds: &[usize]) -> Vec<(Value, RowId)> {
+    let widest = bounds.windows(2).map(|w| w[1] - w[0]).max().unwrap_or(0);
+    let mut scratch = vec![(Value::new_const(0), RowId::new_const(0)); widest];
+    for b in 0..bounds.len() - 1 {
+        radix_sort_slice_by_value(&mut pairs[bounds[b]..bounds[b + 1]], &mut scratch);
+    }
+    if bounds.len() <= 2 {
+        pairs
+    } else {
+        merge_sorted_blocks_dedup(pairs, bounds)
+    }
+}
+
+/// The standard-library analogue of [`sort_merge_custom`]: sort the whole concatenation by
+/// `(Value, RowId)`, then drop duplicates iff `dedup` is set.
+///
+/// Pass `dedup = false` for a single column and `true` for several, mirroring what
+/// `sort_merge_custom` does: a single block has one pair per row, so every `RowId` is distinct
+/// and no `(Value, RowId)` pair can repeat -- deduping there is a wasted pass that would unfairly
+/// handicap this side of the comparison. Only multiple columns can put the same value (hence the
+/// same pair) on one row.
+pub fn sort_merge_std(mut pairs: Vec<(Value, RowId)>, dedup: bool) -> Vec<(Value, RowId)> {
+    pairs.sort_unstable();
+    if dedup {
+        pairs.dedup();
+    }
+    pairs
+}
+
+/// Run `f` with an index thread pool of `n` threads installed, so parallel construction actually
+/// forks: `ColumnIndex` derives both its shard count and worker-pool size from the ambient pool.
+pub fn with_threads<R>(n: usize, f: impl FnOnce() -> R) -> R {
+    ThreadPool::new(n.max(1)).install(f)
+}
+
+/// A random table set up to (re)build a `ColumnIndex` over its value columns (`1..=n_val_cols`;
+/// column 0 is a unique key so no rows merge).
+pub struct IndexInput {
+    table: WrappedTable,
+    cols: Vec<ColumnId>,
+}
+
+impl IndexInput {
+    pub fn random(n_rows: usize, n_val_cols: usize, distinct: u32, seed: u64) -> Self {
+        let mut rng = StdRng::seed_from_u64(seed);
+        let n_cols = n_val_cols + 1;
+        let mut table = SortedWritesTable::new(
+            1,
+            n_cols,
+            None,
+            vec![],
+            Box::new(|_, _old, new, out: &mut Vec<Value>| {
+                out.extend_from_slice(new);
+                false
+            }),
+        );
+        {
+            let mut buf = table.new_buffer();
+            let mut row = vec![Value::new(0); n_cols];
+            for i in 0..n_rows {
+                row[0] = Value::from_usize(i);
+                for cell in row.iter_mut().skip(1) {
+                    *cell = Value::new(rng.random_range(0..distinct));
+                }
+                buf.stage_insert(&row);
+            }
+        }
+        let db = Database::default();
+        db.with_execution_state(|es| {
+            table.merge(es);
+        });
+        IndexInput {
+            table: WrappedTable::new(table),
+            cols: (1..n_cols).map(ColumnId::from_usize).collect(),
+        }
+    }
+
+    /// Build the index with the serial radix-sort + merge path.
+    pub fn build_serial(&self) {
+        let mut ci = ColumnIndex::new();
+        ci.rebuild_full(&self.cols, self.table.as_ref(), self.table.all().as_ref());
+        std::hint::black_box(&ci);
+    }
+
+    /// Build the index with the parallel sharded path. Call inside [`with_threads`].
+    pub fn build_parallel(&self) {
+        let mut ci = ColumnIndex::new();
+        ci.merge_parallel(&self.cols, self.table.as_ref(), self.table.all().as_ref());
+        std::hint::black_box(&ci);
+    }
+}