From aefddbff0c18a90b3ea95fd33ceb3459dbd8c9f7 Mon Sep 17 00:00:00 2001 From: Eli Rosenthal Date: Tue, 21 Jul 2026 23:50:58 +0200 Subject: [PATCH 1/3] Add packed arena trie allocation --- concurrency/src/lib.rs | 2 +- concurrency/src/shared_arena.rs | 157 +++++++- concurrency/src/tests.rs | 204 ++++++++++ core-relations/src/free_join/mod.rs | 3 + core-relations/src/free_join/packed_trie.rs | 409 ++++++++++++++++++++ 5 files changed, 772 insertions(+), 3 deletions(-) create mode 100644 core-relations/src/free_join/packed_trie.rs diff --git a/concurrency/src/lib.rs b/concurrency/src/lib.rs index 9d5993e4e..91af0e44c 100644 --- a/concurrency/src/lib.rs +++ b/concurrency/src/lib.rs @@ -16,7 +16,7 @@ pub use notification::Notification; pub use notification_list::NotificationList; pub use parallel_writer::ParallelVecWriter; pub use resettable_oncelock::ResettableOnceLock; -pub use shared_arena::{Handle, SharedArena, SharedRef}; +pub use shared_arena::{Handle, RawAllocation, SharedArena, SharedRef}; pub use threadpool::{ SchedulerMetrics, Scope, ThreadPool, current_num_threads, scope, without_current_pool, }; diff --git a/concurrency/src/shared_arena.rs b/concurrency/src/shared_arena.rs index bd1d816f0..6779003b4 100644 --- a/concurrency/src/shared_arena.rs +++ b/concurrency/src/shared_arena.rs @@ -11,10 +11,17 @@ //! [`SharedArena`] is dropped. Since the arena itself is `Send`, values must be //! `Send` when they are allocated: the arena may be dropped on a different //! thread from the one that performed the allocation. +//! +//! [`Handle::alloc_layout`] is the narrow escape hatch for packed allocations. +//! It returns an uninitialized, non-`Send` [`RawAllocation`]. The caller writes +//! a no-drop header and any trailing data, then explicitly publishes the +//! header as a [`SharedRef`]. Raw allocations never register destructors; this +//! makes abandoning a partially initialized allocation during unwinding safe, +//! provided every initialized tail value may also be safely forgotten. use std::{ - cell::RefCell, marker::PhantomData, mem, ops::Deref, pin::Pin, ptr::NonNull, rc::Rc, - sync::Mutex, + alloc::Layout, cell::RefCell, marker::PhantomData, mem, ops::Deref, pin::Pin, ptr::NonNull, + rc::Rc, sync::Mutex, }; use bumpalo::Bump; @@ -188,6 +195,136 @@ impl<'arena> Handle<'arena> { _lifetime: PhantomData, } } + + /// Allocate uninitialized storage with exactly `layout` in the parent + /// [`SharedArena`]. + /// + /// The returned [`RawAllocation`] is an initialization token. It exposes a + /// raw pointer for constructing a packed allocation, but no reference to + /// the storage exists until + /// [`RawAllocation::assume_init_no_drop`] publishes a header at the start + /// of the allocation. Dropping the token simply abandons the storage until + /// the arena itself is reclaimed. + /// + /// Raw allocations do not register destructors. This is useful for packed + /// data whose header and trailing values are safe to forget, but it is not + /// a replacement for [`Handle::alloc`] when values own resources. + /// + /// The token intentionally does not implement `Send` or `Sync`: + /// + /// ```compile_fail + /// use std::alloc::Layout; + /// use egglog_concurrency::SharedArena; + /// + /// let arena = SharedArena::new(); + /// let handle = arena.new_handle(); + /// let raw = handle.alloc_layout(Layout::new::()); + /// std::thread::scope(|scope| { + /// scope.spawn(move || drop(raw)); + /// }); + /// ``` + pub fn alloc_layout(&self, layout: Layout) -> RawAllocation<'arena> { + // SAFETY: `self.local` points to a boxed `LocalArena` stored in the + // parent `SharedArena`. Boxes are never removed from that vector before + // the `SharedArena` is dropped, and the raw allocation's lifetime keeps + // that arena borrowed until the initialization token is consumed or + // dropped. + let local = unsafe { self.local.as_ref() }; + RawAllocation { + ptr: local.alloc_layout(layout), + layout, + _arena: PhantomData, + _not_send_sync: PhantomData, + } + } +} + +/// Uninitialized storage owned by a [`SharedArena`]. +/// +/// A raw allocation is an opaque, single-use publication token returned by +/// [`Handle::alloc_layout`]. It is deliberately not `Send` or `Sync`: initialize +/// it on the same thread that owns the allocating [`Handle`]. Its storage stays +/// allocated until the parent arena is dropped, even if the token is abandoned +/// because initialization panics. +/// +/// No destructor is registered for any part of this allocation. The caller of +/// [`RawAllocation::assume_init_no_drop`] is responsible for ensuring that the +/// header and all initialized trailing values can safely be forgotten. +pub struct RawAllocation<'arena> { + ptr: NonNull, + layout: Layout, + _arena: PhantomData<&'arena SharedArena>, + _not_send_sync: PhantomData>, +} + +impl<'arena> RawAllocation<'arena> { + /// Return the allocation's exact requested layout. + pub fn layout(&self) -> Layout { + self.layout + } + + /// Return a pointer to the start of the uninitialized allocation. + /// + /// Dereferencing this pointer is unsafe. Although this method requires a + /// mutable borrow of the initialization token, raw pointers copied from it + /// are not tracked by Rust's borrow checker. + pub fn as_mut_ptr(&mut self) -> *mut u8 { + self.ptr.as_ptr() + } + + /// Publish an initialized, no-drop `T` header at the start of this + /// allocation. + /// + /// The allocation may be larger than `T`; this is the intended way to + /// publish a sized header followed by packed trailing arrays. This method + /// checks, before treating the storage as `T`, that the original layout is + /// large enough and sufficiently aligned for `T`, and that + /// `mem::needs_drop::()` is false. A failed check panics without + /// publishing a reference. + /// + /// # Safety + /// + /// If the checked layout requirements hold, the caller must ensure that: + /// + /// - a valid `T` has been completely initialized at [`Self::as_mut_ptr`]; + /// - every byte that `T`'s safe interface can read, including trailing + /// storage reached through offsets or pointers in `T`, is initialized; + /// - omitting destructors for every initialized trailing value is sound, + /// both after publication and if initialization unwinds partway through; + /// - the `Send` and `Sync` behavior of `T` accurately accounts for any + /// trailing values its interface can access; and + /// - no pointer retained from [`Self::as_mut_ptr`] is used to mutate the + /// allocation after publication, except through valid interior + /// mutability. + /// + /// Raw trailing storage is not described by Rust's type system, so these + /// invariants cannot be checked by this API. + pub unsafe fn assume_init_no_drop(self) -> SharedRef<'arena, T> + where + T: Send + 'arena, + { + assert!( + self.layout.size() >= mem::size_of::(), + "raw arena allocation of {} bytes is too small for a {}-byte header", + self.layout.size(), + mem::size_of::() + ); + assert!( + self.layout.align() >= mem::align_of::(), + "raw arena allocation alignment {} is insufficient for header alignment {}", + self.layout.align(), + mem::align_of::() + ); + assert!( + !mem::needs_drop::(), + "raw arena allocation headers must not require drop" + ); + + SharedRef { + ptr: self.ptr.cast(), + _lifetime: PhantomData, + } + } } /// An immutable reference to a value allocated in a [`SharedArena`]. @@ -220,6 +357,18 @@ impl Clone for SharedRef<'_, T> { } } +impl<'arena, T> SharedRef<'arena, T> { + /// Convert this copyable handle into a reference valid for the lifetime of + /// its parent [`SharedArena`]. + pub fn into_ref(self) -> &'arena T { + // SAFETY: `ptr` denotes a fully initialized value in a bump allocation + // owned by the parent arena. The lifetime marker prevents that arena + // from being dropped for `'arena`, and publication never exposes a + // mutable reference to this value. + unsafe { self.ptr.as_ref() } + } +} + impl Deref for SharedRef<'_, T> { type Target = T; @@ -270,6 +419,10 @@ impl LocalArena { ptr } + + fn alloc_layout(&self, layout: Layout) -> NonNull { + self.bump.alloc_layout(layout) + } } impl Drop for LocalArena { diff --git a/concurrency/src/tests.rs b/concurrency/src/tests.rs index 5cd44a5e2..e558c3be0 100644 --- a/concurrency/src/tests.rs +++ b/concurrency/src/tests.rs @@ -1,4 +1,5 @@ use std::{ + alloc::Layout, cell::Cell, mem, sync::{ @@ -688,6 +689,209 @@ fn shared_arena_allocates_from_scoped_threads() { assert_eq!(values, (0..N_THREADS * PER_THREAD).collect::>()); } +#[test] +fn shared_arena_raw_layout_honors_alignment() { + #[repr(C, align(128))] + struct AlignedHeader { + value: u64, + } + + let arena = SharedArena::new(); + let handle = arena.new_handle(); + let layout = Layout::new::(); + let mut raw = handle.alloc_layout(layout); + + assert_eq!(raw.layout(), layout); + assert_eq!((raw.as_mut_ptr() as usize) % layout.align(), 0); + // SAFETY: The raw allocation has exactly `AlignedHeader`'s layout. We + // initialize the complete no-drop header before publishing it and retain + // no pointer used for mutation afterward. + let published = unsafe { + raw.as_mut_ptr() + .cast::() + .write(AlignedHeader { value: 37 }); + raw.assume_init_no_drop::() + }; + + assert_eq!(published.value, 37); +} + +#[test] +fn shared_arena_raw_layout_publishes_header_with_trailing_data() { + #[repr(C)] + struct PackedHeader { + len: usize, + tail_offset: usize, + } + + let values = [2u32, 3, 5, 7, 11]; + let (layout, tail_offset) = Layout::new::() + .extend(Layout::array::(values.len()).unwrap()) + .unwrap(); + let layout = layout.pad_to_align(); + let arena = SharedArena::new(); + let handle = arena.new_handle(); + let mut raw = handle.alloc_layout(layout); + let base = raw.as_mut_ptr(); + + // SAFETY: `Layout::extend` produced both offsets. We initialize the full + // no-drop header and every trailing `u32` before publishing, and all writes + // end before the resulting shared reference is created. + let published = unsafe { + base.cast::().write(PackedHeader { + len: values.len(), + tail_offset, + }); + let tail = base.add(tail_offset).cast::(); + for (idx, value) in values.into_iter().enumerate() { + tail.add(idx).write(value); + } + raw.assume_init_no_drop::() + }; + let header = published.into_ref(); + + // SAFETY: The header and tail came from the checked packed layout above, + // and `header.len` elements were initialized before publication. + let got = unsafe { + std::slice::from_raw_parts( + (header as *const PackedHeader) + .cast::() + .add(header.tail_offset) + .cast::(), + header.len, + ) + }; + assert_eq!(got, &values); +} + +#[test] +fn shared_arena_raw_publications_outlive_parallel_handles() { + #[repr(transparent)] + struct ParallelHeader(usize); + + const N_THREADS: usize = 8; + let arena = SharedArena::new(); + let published = Mutex::new(Vec::new()); + + thread::scope(|scope| { + for value in 0..N_THREADS { + let arena = &arena; + let published = &published; + scope.spawn(move || { + let handle = arena.new_handle(); + let mut raw = handle.alloc_layout(Layout::new::()); + // SAFETY: The allocation has the header's exact layout and the + // complete no-drop header is initialized before publication. + let shared = unsafe { + raw.as_mut_ptr() + .cast::() + .write(ParallelHeader(value)); + raw.assume_init_no_drop::() + }; + published.lock().unwrap().push(shared); + }); + } + }); + + // Every allocating handle and worker is gone, but the arena still owns the + // published headers. + let mut values = published + .into_inner() + .unwrap() + .into_iter() + .map(|value| value.into_ref().0) + .collect::>(); + values.sort_unstable(); + assert_eq!(values, (0..N_THREADS).collect::>()); +} + +#[test] +fn shared_arena_raw_initialization_panic_does_not_publish() { + #[repr(transparent)] + struct Header(usize); + + fn build<'arena>( + handle: &crate::Handle<'arena>, + panic_before_publish: bool, + ) -> SharedRef<'arena, Header> { + let mut raw = handle.alloc_layout(Layout::new::
()); + // SAFETY: The allocation has `Header`'s exact layout and this writes the + // complete no-drop header. It is not yet published. + unsafe { + raw.as_mut_ptr().cast::
().write(Header(41)); + } + assert!(!panic_before_publish, "injected initialization panic"); + // SAFETY: The complete header was initialized above and no raw pointer + // is used after this publication point. + unsafe { raw.assume_init_no_drop::
() } + } + + let arena = SharedArena::new(); + let handle = arena.new_handle(); + let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + let _ = build(&handle, true); + })); + assert!(result.is_err()); + + // The abandoned allocation registered no destructor and published no + // reference; another allocation in the same local arena remains usable. + assert_eq!(build(&handle, false).into_ref().0, 41); +} + +#[test] +fn shared_arena_raw_publication_validates_header_requirements() { + #[repr(C)] + struct WideHeader([usize; 2]); + + #[repr(C, align(64))] + struct OverAlignedHeader([u8; 64]); + + struct DropHeader { + _value: u8, + } + impl Drop for DropHeader { + fn drop(&mut self) {} + } + + let arena = SharedArena::new(); + let handle = arena.new_handle(); + + let too_small = handle.alloc_layout(Layout::new::()); + let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + // SAFETY: `assume_init_no_drop` checks size before it relies on the + // initialization invariant, so this deliberately undersized token + // must panic without interpreting its storage as `WideHeader`. + let _ = unsafe { too_small.assume_init_no_drop::() }; + })); + assert!(result.is_err()); + + let under_aligned = handle + .alloc_layout(Layout::from_size_align(mem::size_of::(), 1).unwrap()); + let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + // SAFETY: The method checks alignment before relying on initialization, + // so this must panic without interpreting the under-aligned storage. + let _ = unsafe { under_aligned.assume_init_no_drop::() }; + })); + assert!(result.is_err()); + + let mut needs_drop = handle.alloc_layout(Layout::new::()); + // SAFETY: Initialize a valid header so the only rejected condition is its + // destructor. The raw API deliberately forgets the value after rejecting + // it; this test type's destructor owns no resources. + unsafe { + needs_drop + .as_mut_ptr() + .cast::() + .write(DropHeader { _value: 0 }); + } + let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + // SAFETY: The storage contains a valid `DropHeader`; the method must + // reject it before publication because raw headers may not need drop. + let _ = unsafe { needs_drop.assume_init_no_drop::() }; + })); + assert!(result.is_err()); +} + #[test] fn shared_arena_drops_values_lifo_within_a_handle() { #[derive(Debug)] diff --git a/core-relations/src/free_join/mod.rs b/core-relations/src/free_join/mod.rs index b45239da5..82dbe92ee 100644 --- a/core-relations/src/free_join/mod.rs +++ b/core-relations/src/free_join/mod.rs @@ -38,6 +38,9 @@ use crate::action::ExecutionState; pub(crate) mod execute; pub(crate) mod frame_update; +// The packed trie is exercised independently before it is wired into execution. +#[allow(dead_code)] +pub(crate) mod packed_trie; pub(crate) mod plan; define_id!( diff --git a/core-relations/src/free_join/packed_trie.rs b/core-relations/src/free_join/packed_trie.rs new file mode 100644 index 000000000..4db5f9fec --- /dev/null +++ b/core-relations/src/free_join/packed_trie.rs @@ -0,0 +1,409 @@ +//! Arena-backed scalar trie levels for one join execution. +//! +//! A node is one allocation containing an immutable sorted column index and, +//! when another level follows it, one child publication slot per distinct key. +//! The allocation is tied to a [`Handle`] and deliberately has no destructor: +//! its only non-`Copy` trailing values are `OnceLock`s containing references +//! into the same execution-scoped arena. + +use std::{alloc::Layout, marker::PhantomData, ptr, sync::OnceLock}; + +use egglog_concurrency::Handle; + +use crate::{ + OffsetRange, SubsetRef, Value, + numeric_id::NumericId, + offsets::{RowId, SortedOffsetSlice}, + table_spec::ColumnId, +}; + +const NO_CHILDREN: u32 = u32::MAX; + +/// One immutable, scalar trie level allocated for the lifetime of an execution. +/// +/// The dynamically sized portions immediately following this header are laid +/// out as keys, boundaries, grouped row ids, and (optionally) child locks. Use +/// the accessors rather than relying on those offsets outside this module. +#[repr(C)] +pub(crate) struct PackedTrieNode<'exec> { + key_len: u32, + row_len: u32, + column: ColumnId, + keys_offset: u32, + boundaries_offset: u32, + rows_offset: u32, + children_offset: u32, + allocation_bytes: u32, + _execution: PhantomData<&'exec ()>, +} + +struct PackedLayout { + allocation: Layout, + keys_offset: usize, + boundaries_offset: usize, + rows_offset: usize, + children_offset: Option, +} + +impl PackedLayout { + fn new(key_len: usize, row_len: usize, with_children: bool) -> Self { + // Lifetimes do not affect layout; use one concrete instantiation so + // callers do not need to manufacture a lifetime solely for arithmetic. + let (layout, keys_offset) = Layout::new::>() + .extend(Layout::array::(key_len).expect("packed trie key layout overflow")) + .expect("packed trie key offset overflow"); + let (layout, boundaries_offset) = layout + .extend( + Layout::array::( + key_len + .checked_add(1) + .expect("packed trie boundary count overflow"), + ) + .expect("packed trie boundary layout overflow"), + ) + .expect("packed trie boundary offset overflow"); + let (layout, rows_offset) = layout + .extend(Layout::array::(row_len).expect("packed trie row layout overflow")) + .expect("packed trie row offset overflow"); + let (layout, children_offset) = if with_children { + let (layout, offset) = layout + .extend( + Layout::array::>>(key_len) + .expect("packed trie child layout overflow"), + ) + .expect("packed trie child offset overflow"); + (layout, Some(offset)) + } else { + (layout, None) + }; + + Self { + allocation: layout.pad_to_align(), + keys_offset, + boundaries_offset, + rows_offset, + children_offset, + } + } + + fn checked_u32(value: usize, section: &str) -> u32 { + u32::try_from(value) + .unwrap_or_else(|_| panic!("packed trie {section} exceeds u32::MAX bytes")) + } +} + +impl<'exec> PackedTrieNode<'exec> { + /// Build a node from pairs sorted lexicographically by `(Value, RowId)`. + /// + /// Sorting by value groups equal keys. Sorting each group by row id is what + /// makes every range returned by [`Self::subset_at`] a valid `SubsetRef`. + pub(crate) fn build_from_sorted_pairs( + arena: &Handle<'exec>, + column: ColumnId, + pairs: &[(Value, RowId)], + with_children: bool, + ) -> &'exec Self { + assert!( + pairs.windows(2).all(|pair| pair[0] <= pair[1]), + "packed trie input must be sorted by (Value, RowId)" + ); + + let row_len = u32::try_from(pairs.len()) + .expect("a packed trie node cannot contain more than u32::MAX rows"); + let key_len_usize = pairs + .iter() + .enumerate() + .filter(|(index, pair)| *index == 0 || pairs[*index - 1].0 != pair.0) + .count(); + let key_len = u32::try_from(key_len_usize) + .expect("a packed trie node cannot contain more than u32::MAX keys"); + let layout = PackedLayout::new(key_len_usize, pairs.len(), with_children); + let keys_offset = PackedLayout::checked_u32(layout.keys_offset, "key offset"); + let boundaries_offset = + PackedLayout::checked_u32(layout.boundaries_offset, "boundary offset"); + let rows_offset = PackedLayout::checked_u32(layout.rows_offset, "row offset"); + let children_offset = layout + .children_offset + .map(|offset| PackedLayout::checked_u32(offset, "child offset")) + .unwrap_or(NO_CHILDREN); + let allocation_bytes = + PackedLayout::checked_u32(layout.allocation.size(), "allocation size"); + let mut allocation = arena.alloc_layout(layout.allocation); + let base = allocation.as_mut_ptr(); + + // SAFETY: every pointer below is derived from `Layout::extend` offsets + // within this allocation. Each element is written exactly once, no + // shared reference exists yet, and all lengths were checked above. + unsafe { + let keys = base.add(layout.keys_offset).cast::(); + let boundaries = base.add(layout.boundaries_offset).cast::(); + let rows = base.add(layout.rows_offset).cast::(); + let children = layout + .children_offset + .map(|offset| base.add(offset).cast::>()); + + let mut next_key = 0usize; + for (row_index, &(value, row_id)) in pairs.iter().enumerate() { + if row_index == 0 || pairs[row_index - 1].0 != value { + ptr::write(keys.add(next_key), value); + ptr::write(boundaries.add(next_key), row_index as u32); + if let Some(children) = children { + ptr::write(children.add(next_key), OnceLock::new()); + } + next_key += 1; + } + ptr::write(rows.add(row_index), row_id); + } + debug_assert_eq!(next_key, key_len_usize); + ptr::write(boundaries.add(key_len_usize), row_len); + ptr::write( + base.cast::(), + Self { + key_len, + row_len, + column, + keys_offset, + boundaries_offset, + rows_offset, + children_offset, + allocation_bytes, + _execution: PhantomData, + }, + ); + + // The complete header and every trailing element are initialized + // before this conversion publishes an immutable arena reference. + allocation.assume_init_no_drop::().into_ref() + } + } + + pub(crate) fn values(&self) -> &[Value] { + // SAFETY: construction initializes exactly `key_len` values at this + // layout-derived offset, and the arena outlives `self`. + unsafe { + std::slice::from_raw_parts( + (self as *const Self) + .cast::() + .add(self.keys_offset as usize) + .cast(), + self.key_len(), + ) + } + } + + pub(crate) fn boundaries(&self) -> &[u32] { + // SAFETY: construction initializes one boundary per key plus a final + // sentinel at this layout-derived offset. + unsafe { + std::slice::from_raw_parts( + (self as *const Self) + .cast::() + .add(self.boundaries_offset as usize) + .cast(), + self.key_len() + 1, + ) + } + } + + pub(crate) fn rows(&self) -> &[RowId] { + // SAFETY: construction initializes exactly `row_len` row ids at this + // layout-derived offset. + unsafe { + std::slice::from_raw_parts( + (self as *const Self) + .cast::() + .add(self.rows_offset as usize) + .cast(), + self.row_len(), + ) + } + } + + pub(crate) fn find(&self, value: Value) -> Option { + self.values().binary_search(&value).ok() + } + + pub(crate) fn column(&self) -> ColumnId { + self.column + } + + pub(crate) fn subset_at(&self, key_index: usize) -> SubsetRef<'_> { + assert!(key_index < self.key_len(), "packed trie key out of bounds"); + let boundaries = self.boundaries(); + let rows = &self.rows()[boundaries[key_index] as usize..boundaries[key_index + 1] as usize]; + debug_assert!(!rows.is_empty()); + let first = rows[0]; + let last = rows[rows.len() - 1]; + if last.index() - first.index() == rows.len() - 1 { + SubsetRef::Dense(OffsetRange::new(first, last.inc())) + } else { + // SAFETY: the constructor requires `(Value, RowId)` order, so the + // rows within each equal-value range are non-decreasing. + SubsetRef::Sparse(unsafe { SortedOffsetSlice::new_unchecked(rows) }) + } + } + + pub(crate) fn child_slot( + &self, + key_index: usize, + ) -> Option<&OnceLock<&'exec PackedTrieNode<'exec>>> { + assert!(key_index < self.key_len(), "packed trie key out of bounds"); + if self.children_offset == NO_CHILDREN { + return None; + } + // SAFETY: a non-leaf construction initializes exactly `key_len` child + // locks at this layout-derived offset. + Some(unsafe { + &*(self as *const Self) + .cast::() + .add(self.children_offset as usize) + .cast::>() + .add(key_index) + }) + } + + pub(crate) fn allocation_bytes(&self) -> usize { + self.allocation_bytes as usize + } + + fn key_len(&self) -> usize { + self.key_len as usize + } + + fn row_len(&self) -> usize { + self.row_len as usize + } +} + +#[cfg(test)] +mod tests { + use std::{ + mem::{align_of, size_of}, + sync::atomic::{AtomicUsize, Ordering}, + }; + + use egglog_concurrency::SharedArena; + + use super::*; + + fn value(value: usize) -> Value { + Value::from_usize(value) + } + + fn row(row: usize) -> RowId { + RowId::from_usize(row) + } + + fn column(column: usize) -> ColumnId { + ColumnId::from_usize(column) + } + + #[test] + fn packed_trie_builds_sorted_ranges() { + let arena = SharedArena::new(); + let handle = arena.new_handle(); + let pairs = [ + (value(1), row(1)), + (value(1), row(3)), + (value(2), row(2)), + (value(3), row(4)), + (value(3), row(5)), + ]; + let node = PackedTrieNode::build_from_sorted_pairs(&handle, column(2), &pairs, true); + + assert_eq!(node.column(), column(2)); + assert_eq!(node.values(), &[value(1), value(2), value(3)]); + assert_eq!(node.boundaries(), &[0, 2, 3, 5]); + assert_eq!(node.rows(), &[row(1), row(3), row(2), row(4), row(5)]); + assert_eq!(node.find(value(2)), Some(1)); + assert_eq!(node.find(value(9)), None); + + let SubsetRef::Sparse(one) = node.subset_at(0) else { + panic!("noncontiguous range should be sparse") + }; + assert_eq!(one.inner(), &[row(1), row(3)]); + assert!(matches!( + node.subset_at(1), + SubsetRef::Dense(range) if range == OffsetRange::new(row(2), row(3)) + )); + assert!(matches!( + node.subset_at(2), + SubsetRef::Dense(range) if range == OffsetRange::new(row(4), row(6)) + )); + } + + #[test] + fn packed_trie_leaf_omits_child_locks() { + let arena = SharedArena::new(); + let handle = arena.new_handle(); + let pairs = [(value(1), row(0)), (value(2), row(1))]; + let leaf = PackedTrieNode::build_from_sorted_pairs(&handle, column(0), &pairs, false); + let branch = PackedTrieNode::build_from_sorted_pairs(&handle, column(0), &pairs, true); + + assert!(leaf.child_slot(0).is_none()); + assert!(leaf.child_slot(1).is_none()); + assert!(branch.child_slot(0).is_some()); + assert!(branch.child_slot(1).is_some()); + assert!(leaf.allocation_bytes() < branch.allocation_bytes()); + } + + #[test] + fn packed_trie_sections_are_aligned() { + let arena = SharedArena::new(); + let handle = arena.new_handle(); + let pairs = [(value(1), row(0)), (value(2), row(2))]; + let node = PackedTrieNode::build_from_sorted_pairs(&handle, column(1), &pairs, true); + + assert_eq!( + node as *const _ as usize % align_of::>(), + 0 + ); + assert_eq!(node.values().as_ptr() as usize % align_of::(), 0); + assert_eq!(node.boundaries().as_ptr() as usize % align_of::(), 0); + assert_eq!(node.rows().as_ptr() as usize % align_of::(), 0); + assert_eq!( + node.child_slot(0).unwrap() as *const _ as usize + % align_of::>>(), + 0 + ); + assert!(node.allocation_bytes() >= size_of::>()); + } + + #[test] + fn packed_trie_child_is_published_once_under_race() { + let arena = SharedArena::new(); + let handle = arena.new_handle(); + let parent = PackedTrieNode::build_from_sorted_pairs( + &handle, + column(0), + &[(value(1), row(0))], + true, + ); + let slot = parent.child_slot(0).unwrap(); + let initializations = AtomicUsize::new(0); + + std::thread::scope(|scope| { + for _ in 0..8 { + let arena = &arena; + let initializations = &initializations; + scope.spawn(move || { + let child = slot.get_or_init(|| { + initializations.fetch_add(1, Ordering::Relaxed); + let handle = arena.new_handle(); + PackedTrieNode::build_from_sorted_pairs( + &handle, + column(1), + &[(value(7), row(9))], + false, + ) + }); + assert_eq!(child.values(), &[value(7)]); + }); + } + }); + + assert_eq!(initializations.load(Ordering::Relaxed), 1); + assert_eq!(slot.get().unwrap().column(), column(1)); + assert_eq!(slot.get().unwrap().values(), &[value(7)]); + } +} From f5636fee1f5f0f94349fa56c9703d3ed5101a9fe Mon Sep 17 00:00:00 2001 From: Eli Rosenthal Date: Thu, 30 Jul 2026 09:38:00 -0700 Subject: [PATCH 2/3] Add positioned indexes and packed trie layouts --- core-relations/src/free_join/packed_trie.rs | 1069 +++++++++++++++++-- core-relations/src/hash_index/mod.rs | 194 +++- core-relations/src/hash_index/tests.rs | 196 +++- 3 files changed, 1353 insertions(+), 106 deletions(-) diff --git a/core-relations/src/free_join/packed_trie.rs b/core-relations/src/free_join/packed_trie.rs index 4db5f9fec..9343fe427 100644 --- a/core-relations/src/free_join/packed_trie.rs +++ b/core-relations/src/free_join/packed_trie.rs @@ -6,7 +6,16 @@ //! its only non-`Copy` trailing values are `OnceLock`s containing references //! into the same execution-scoped arena. -use std::{alloc::Layout, marker::PhantomData, ptr, sync::OnceLock}; +use std::{ + alloc::Layout, + marker::PhantomData, + mem::{ManuallyDrop, align_of, size_of}, + ptr, + sync::{ + OnceLock, + atomic::{AtomicPtr, Ordering}, + }, +}; use egglog_concurrency::Handle; @@ -14,10 +23,48 @@ use crate::{ OffsetRange, SubsetRef, Value, numeric_id::NumericId, offsets::{RowId, SortedOffsetSlice}, - table_spec::ColumnId, + table_spec::{ColumnId, WrappedTableRef}, }; -const NO_CHILDREN: u32 = u32::MAX; +const HAS_CHILDREN: u32 = 1 << 31; +const DYNAMIC_CHILDREN: u32 = 1 << 30; +const KEY_LEN_MASK: u32 = DYNAMIC_CHILDREN - 1; + +type ChildSlot<'exec> = OnceLock<&'exec PackedTrieNode<'exec>>; + +/// How one packed trie level reaches its children. +/// +/// Tuple-index interior levels and fixed-order join levels use [`Self::Direct`] +/// and retain one inline child slot per key. A dynamically ordered join level +/// uses one lazily allocated child-slot family per eligible successor. The +/// family arrays themselves are raw arena allocations and therefore add no +/// registered destructors. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) enum ChildShape { + Leaf, + Direct, + Dynamic { families: usize }, +} + +impl ChildShape { + fn flags(self) -> u32 { + match self { + Self::Leaf => 0, + Self::Direct => HAS_CHILDREN, + Self::Dynamic { families } => { + assert!( + families > 0, + "a dynamic packed trie node needs at least one child family" + ); + assert!( + u32::try_from(families).is_ok(), + "a packed trie node cannot contain more than u32::MAX child families" + ); + HAS_CHILDREN | DYNAMIC_CHILDREN + } + } + } +} /// One immutable, scalar trie level allocated for the lifetime of an execution. /// @@ -25,16 +72,128 @@ const NO_CHILDREN: u32 = u32::MAX; /// out as keys, boundaries, grouped row ids, and (optionally) child locks. Use /// the accessors rather than relying on those offsets outside this module. #[repr(C)] +#[derive(Debug)] pub(crate) struct PackedTrieNode<'exec> { - key_len: u32, + key_len_and_flags: u32, row_len: u32, - column: ColumnId, - keys_offset: u32, - boundaries_offset: u32, - rows_offset: u32, - children_offset: u32, - allocation_bytes: u32, - _execution: PhantomData<&'exec ()>, + // The trailing child slots contain `OnceLock<&'exec PackedTrieNode<'exec>>`, + // which is invariant in `'exec` because `OnceLock` permits publication. + // Model that invariance in the sized header even though the slots live in + // trailing storage. A function argument plus return keeps the marker + // invariant without making the immutable node `!Sync`. + _execution: PhantomData &'exec ()>, +} + +/// A checked position within one packed scalar trie level. +/// +/// Keeping the ordinal beside the arena reference makes descendant binding +/// state `Copy`: the represented subset is the corresponding range in the +/// node's row array, rather than another owned trie node. +#[derive(Clone, Copy, Debug)] +pub(crate) struct PackedCursor<'node, 'exec> { + node: &'node PackedTrieNode<'exec>, + key_index: u32, +} + +impl<'node, 'exec> PackedCursor<'node, 'exec> +where + 'exec: 'node, +{ + pub(crate) fn new(node: &'node PackedTrieNode<'exec>, key_index: usize) -> Self { + assert!( + key_index < node.key_len(), + "packed trie cursor key out of bounds" + ); + Self { + node, + key_index: u32::try_from(key_index) + .expect("packed trie cursor key ordinal exceeds u32::MAX"), + } + } + + pub(crate) fn subset(self) -> SubsetRef<'node> { + self.node.subset_at(self.key_index()) + } + + pub(crate) fn size(self) -> usize { + let boundaries = self.node.boundaries(); + boundaries[self.key_index() + 1] as usize - boundaries[self.key_index()] as usize + } + + /// Return the next scalar index below this cursor, building and publishing + /// it once for the current key when necessary. + /// + /// A child slot belongs to a frozen execution plan: every use must request + /// the same next column and leaf shape. The compact node header records the + /// shape but deliberately does not repeat the plan's column, so requesting + /// the same column is a caller invariant. + pub(crate) fn child_index( + self, + arena: &Handle<'exec>, + table: WrappedTableRef<'_>, + column: ColumnId, + family: usize, + child_shape: ChildShape, + scratch: &mut Vec<(Value, RowId)>, + ) -> &'exec PackedTrieNode<'exec> { + self.child_index_with(arena, family, child_shape, || { + PackedTrieNode::build_from_subset( + arena, + table, + self.subset(), + column, + child_shape, + scratch, + ) + }) + } + + /// Like [`Self::child_index`], but publish an index over a subset that the + /// caller has already refined by the next scan's row constraints. + /// + /// The prefiltered subset is copied into the arena allocation during this + /// call, so it may borrow short-lived executor scratch. As with the column + /// and leaf shape, a frozen plan must supply equivalent filtering on every + /// use of this cursor's child slot. + #[allow(clippy::too_many_arguments)] + pub(crate) fn child_index_from_subset( + self, + arena: &Handle<'exec>, + table: WrappedTableRef<'_>, + subset: SubsetRef<'_>, + column: ColumnId, + family: usize, + child_shape: ChildShape, + scratch: &mut Vec<(Value, RowId)>, + ) -> &'exec PackedTrieNode<'exec> { + self.child_index_with(arena, family, child_shape, || { + PackedTrieNode::build_from_subset(arena, table, subset, column, child_shape, scratch) + }) + } + + /// Return this cursor's child, invoking `build` only when its slot is + /// empty. Callers that must copy or refine the cursor subset should put + /// that work in `build`, so a cache hit remains allocation-free. + pub(crate) fn child_index_with( + self, + arena: &Handle<'exec>, + family: usize, + child_shape: ChildShape, + build: impl FnOnce() -> &'exec PackedTrieNode<'exec>, + ) -> &'exec PackedTrieNode<'exec> { + let slot = self.node.child_slot(arena, self.key_index(), family); + let child = *slot.get_or_init(build); + assert_eq!( + child.child_shape(), + child_shape, + "packed trie child cache shape mismatch" + ); + child + } + + fn key_index(self) -> usize { + self.key_index as usize + } } struct PackedLayout { @@ -43,10 +202,12 @@ struct PackedLayout { boundaries_offset: usize, rows_offset: usize, children_offset: Option, + dynamic_family_count_offset: Option, } impl PackedLayout { - fn new(key_len: usize, row_len: usize, with_children: bool) -> Self { + fn new(key_len: usize, row_len: usize, child_shape: ChildShape) -> Self { + let _ = child_shape.flags(); // Lifetimes do not affect layout; use one concrete instantiation so // callers do not need to manufacture a lifetime solely for arithmetic. let (layout, keys_offset) = Layout::new::>() @@ -65,16 +226,29 @@ impl PackedLayout { let (layout, rows_offset) = layout .extend(Layout::array::(row_len).expect("packed trie row layout overflow")) .expect("packed trie row offset overflow"); - let (layout, children_offset) = if with_children { - let (layout, offset) = layout - .extend( - Layout::array::>>(key_len) - .expect("packed trie child layout overflow"), - ) - .expect("packed trie child offset overflow"); - (layout, Some(offset)) - } else { - (layout, None) + let (layout, children_offset, dynamic_family_count_offset) = match child_shape { + ChildShape::Leaf => (layout, None, None), + ChildShape::Direct => { + let (layout, offset) = layout + .extend( + Layout::array::>(key_len) + .expect("packed trie direct-child layout overflow"), + ) + .expect("packed trie direct-child offset overflow"); + (layout, Some(offset), None) + } + ChildShape::Dynamic { families } => { + let (layout, family_count_offset) = layout + .extend(Layout::new::()) + .expect("packed trie dynamic-family count offset overflow"); + let (layout, children_offset) = layout + .extend( + Layout::array::>>(families) + .expect("packed trie dynamic-family layout overflow"), + ) + .expect("packed trie dynamic-family offset overflow"); + (layout, Some(children_offset), Some(family_count_offset)) + } }; Self { @@ -83,27 +257,76 @@ impl PackedLayout { boundaries_offset, rows_offset, children_offset, + dynamic_family_count_offset, } } - - fn checked_u32(value: usize, section: &str) -> u32 { - u32::try_from(value) - .unwrap_or_else(|_| panic!("packed trie {section} exceeds u32::MAX bytes")) - } } impl<'exec> PackedTrieNode<'exec> { + /// Project one column of an already-filtered subset and build its packed + /// scalar trie level. + /// + /// `subset` must already reflect every row-local constraint that applies + /// at this point in the join. The caller owns `scratch` so recursive join + /// execution can reuse the projection and radix-sort allocation. On + /// return, `scratch` contains the sorted `(Value, RowId)` pairs used to + /// build the node; its capacity also retains the radix sort's ping-pong + /// storage for the next build. + pub(crate) fn build_from_subset( + arena: &Handle<'exec>, + table: WrappedTableRef<'_>, + subset: SubsetRef<'_>, + column: ColumnId, + child_shape: ChildShape, + scratch: &mut Vec<(Value, RowId)>, + ) -> &'exec Self { + scratch.clear(); + scratch.reserve(subset.size()); + table.for_each_col(subset, column, &mut |row_id, value| { + scratch.push((value, row_id)); + }); + // A SubsetRef is RowId-ordered, and for_each_col preserves that scan + // order. The repository's value-stable radix sort therefore produces + // full (Value, RowId) order without a separate RowId pass. A scalar + // projection has exactly one pair per input row, so retaining every + // pair (rather than deduplicating) preserves the subset exactly. Keep + // the sort's ping-pong half in the same caller-owned Vec so repeated + // descendant construction does not allocate another temporary buffer. + debug_assert!( + scratch.windows(2).all(|pair| pair[0].1 <= pair[1].1), + "packed trie subset projection must be RowId-ordered" + ); + let pair_len = scratch.len(); + if pair_len < 64 { + // This mirrors radix_sort_slice_by_value's small-input path, but + // avoids growing and initializing a second half that it would not + // use. Most lower-level trie subsets fall into this case. + scratch.sort_unstable(); + } else { + let storage_len = pair_len + .checked_mul(2) + .expect("packed trie sort scratch size overflow"); + scratch.resize(storage_len, (Value::new_const(0), RowId::new_const(0))); + { + let (pairs, sort_scratch) = scratch.split_at_mut(pair_len); + crate::hash_index::radix_sort_slice_by_value(pairs, sort_scratch); + } + scratch.truncate(pair_len); + } + + Self::build_from_sorted_pairs(arena, scratch, child_shape) + } + /// Build a node from pairs sorted lexicographically by `(Value, RowId)`. /// /// Sorting by value groups equal keys. Sorting each group by row id is what /// makes every range returned by [`Self::subset_at`] a valid `SubsetRef`. pub(crate) fn build_from_sorted_pairs( arena: &Handle<'exec>, - column: ColumnId, pairs: &[(Value, RowId)], - with_children: bool, + child_shape: ChildShape, ) -> &'exec Self { - assert!( + debug_assert!( pairs.windows(2).all(|pair| pair[0] <= pair[1]), "packed trie input must be sorted by (Value, RowId)" ); @@ -115,19 +338,43 @@ impl<'exec> PackedTrieNode<'exec> { .enumerate() .filter(|(index, pair)| *index == 0 || pairs[*index - 1].0 != pair.0) .count(); - let key_len = u32::try_from(key_len_usize) - .expect("a packed trie node cannot contain more than u32::MAX keys"); - let layout = PackedLayout::new(key_len_usize, pairs.len(), with_children); - let keys_offset = PackedLayout::checked_u32(layout.keys_offset, "key offset"); - let boundaries_offset = - PackedLayout::checked_u32(layout.boundaries_offset, "boundary offset"); - let rows_offset = PackedLayout::checked_u32(layout.rows_offset, "row offset"); - let children_offset = layout - .children_offset - .map(|offset| PackedLayout::checked_u32(offset, "child offset")) - .unwrap_or(NO_CHILDREN); - let allocation_bytes = - PackedLayout::checked_u32(layout.allocation.size(), "allocation size"); + let key_len_and_flags = Self::encode_key_len(key_len_usize, child_shape); + let layout = PackedLayout::new(key_len_usize, pairs.len(), child_shape); + debug_assert_eq!(layout.keys_offset, Self::keys_offset()); + debug_assert_eq!( + layout.boundaries_offset, + Self::boundaries_offset_for(key_len_usize) + ); + debug_assert_eq!(layout.rows_offset, Self::rows_offset_for(key_len_usize)); + match child_shape { + ChildShape::Leaf => { + debug_assert_eq!(layout.children_offset, None); + debug_assert_eq!(layout.dynamic_family_count_offset, None); + } + ChildShape::Direct => { + debug_assert_eq!( + layout.children_offset, + Some(Self::direct_children_offset_for(key_len_usize, pairs.len())) + ); + debug_assert_eq!(layout.dynamic_family_count_offset, None); + } + ChildShape::Dynamic { .. } => { + debug_assert_eq!( + layout.dynamic_family_count_offset, + Some(Self::dynamic_family_count_offset_for( + key_len_usize, + pairs.len() + )) + ); + debug_assert_eq!( + layout.children_offset, + Some(Self::dynamic_children_offset_for( + key_len_usize, + pairs.len() + )) + ); + } + } let mut allocation = arena.alloc_layout(layout.allocation); let base = allocation.as_mut_ptr(); @@ -138,16 +385,21 @@ impl<'exec> PackedTrieNode<'exec> { let keys = base.add(layout.keys_offset).cast::(); let boundaries = base.add(layout.boundaries_offset).cast::(); let rows = base.add(layout.rows_offset).cast::(); - let children = layout - .children_offset - .map(|offset| base.add(offset).cast::>()); + let direct_children = matches!(child_shape, ChildShape::Direct).then(|| { + base.add( + layout + .children_offset + .expect("direct children need an offset"), + ) + .cast::>() + }); let mut next_key = 0usize; for (row_index, &(value, row_id)) in pairs.iter().enumerate() { if row_index == 0 || pairs[row_index - 1].0 != value { ptr::write(keys.add(next_key), value); ptr::write(boundaries.add(next_key), row_index as u32); - if let Some(children) = children { + if let Some(children) = direct_children { ptr::write(children.add(next_key), OnceLock::new()); } next_key += 1; @@ -156,17 +408,35 @@ impl<'exec> PackedTrieNode<'exec> { } debug_assert_eq!(next_key, key_len_usize); ptr::write(boundaries.add(key_len_usize), row_len); + if let ChildShape::Dynamic { families } = child_shape { + let family_count = base + .add( + layout + .dynamic_family_count_offset + .expect("dynamic children need a family-count offset"), + ) + .cast::(); + ptr::write( + family_count, + u32::try_from(families) + .expect("a packed trie node cannot contain more than u32::MAX families"), + ); + let family_table = base + .add( + layout + .children_offset + .expect("dynamic children need a table offset"), + ) + .cast::>>(); + for family in 0..families { + ptr::write(family_table.add(family), AtomicPtr::new(ptr::null_mut())); + } + } ptr::write( base.cast::(), Self { - key_len, + key_len_and_flags, row_len, - column, - keys_offset, - boundaries_offset, - rows_offset, - children_offset, - allocation_bytes, _execution: PhantomData, }, ); @@ -184,7 +454,7 @@ impl<'exec> PackedTrieNode<'exec> { std::slice::from_raw_parts( (self as *const Self) .cast::() - .add(self.keys_offset as usize) + .add(Self::keys_offset()) .cast(), self.key_len(), ) @@ -198,7 +468,7 @@ impl<'exec> PackedTrieNode<'exec> { std::slice::from_raw_parts( (self as *const Self) .cast::() - .add(self.boundaries_offset as usize) + .add(Self::boundaries_offset_for(self.key_len())) .cast(), self.key_len() + 1, ) @@ -212,7 +482,7 @@ impl<'exec> PackedTrieNode<'exec> { std::slice::from_raw_parts( (self as *const Self) .cast::() - .add(self.rows_offset as usize) + .add(Self::rows_offset_for(self.key_len())) .cast(), self.row_len(), ) @@ -223,10 +493,6 @@ impl<'exec> PackedTrieNode<'exec> { self.values().binary_search(&value).ok() } - pub(crate) fn column(&self) -> ColumnId { - self.column - } - pub(crate) fn subset_at(&self, key_index: usize) -> SubsetRef<'_> { assert!(key_index < self.key_len(), "packed trie key out of bounds"); let boundaries = self.boundaries(); @@ -234,7 +500,10 @@ impl<'exec> PackedTrieNode<'exec> { debug_assert!(!rows.is_empty()); let first = rows[0]; let last = rows[rows.len() - 1]; - if last.index() - first.index() == rows.len() - 1 { + if rows + .windows(2) + .all(|pair| pair[0].index().checked_add(1) == Some(pair[1].index())) + { SubsetRef::Dense(OffsetRange::new(first, last.inc())) } else { // SAFETY: the constructor requires `(Value, RowId)` order, so the @@ -243,36 +512,234 @@ impl<'exec> PackedTrieNode<'exec> { } } - pub(crate) fn child_slot( - &self, - key_index: usize, - ) -> Option<&OnceLock<&'exec PackedTrieNode<'exec>>> { + /// Return the direct child slot for `key_index`. + /// + /// This accessor exists for fixed-order and tuple-index traversal. Dynamic + /// traversal must go through [`Self::child_slot`] so its family array is + /// allocated and published safely. + pub(crate) fn direct_child_slot(&self, key_index: usize) -> Option<&ChildSlot<'exec>> { assert!(key_index < self.key_len(), "packed trie key out of bounds"); - if self.children_offset == NO_CHILDREN { - return None; + match self.child_shape() { + ChildShape::Leaf => return None, + ChildShape::Direct => {} + ChildShape::Dynamic { .. } => { + panic!("a dynamic packed trie node has no direct child slots") + } } - // SAFETY: a non-leaf construction initializes exactly `key_len` child + let children_offset = Self::direct_children_offset_for(self.key_len(), self.row_len()); + // SAFETY: direct construction initializes exactly `key_len` child // locks at this layout-derived offset. Some(unsafe { &*(self as *const Self) .cast::() - .add(self.children_offset as usize) - .cast::>() + .add(children_offset) + .cast::>() .add(key_index) }) } - pub(crate) fn allocation_bytes(&self) -> usize { - self.allocation_bytes as usize + /// Get a key's child slot in `family`, lazily publishing the whole dynamic + /// family array when this is its first use. + /// + fn child_slot( + &self, + arena: &Handle<'exec>, + key_index: usize, + family: usize, + ) -> &ChildSlot<'exec> { + assert!(key_index < self.key_len(), "packed trie key out of bounds"); + match self.child_shape() { + ChildShape::Leaf => panic!("cannot build a child index below a packed trie leaf"), + ChildShape::Direct => { + // Direct nodes have exactly one deterministic successor. The + // executor passes its plan-local AccessId uniformly as + // `family`; it is intentionally ignored in this shape. + return self + .direct_child_slot(key_index) + .expect("direct packed trie node must have child slots"); + } + ChildShape::Dynamic { families } => { + assert!( + family < families, + "packed trie dynamic child family out of bounds" + ); + } + } + + let family_entry = unsafe { + // SAFETY: `child_shape` established that this is a dynamic node, + // and construction initialized `families` AtomicPtr entries at + // this layout-derived offset. The bounds check above selects one. + &*(self as *const Self) + .cast::() + .add(Self::dynamic_children_offset_for( + self.key_len(), + self.row_len(), + )) + .cast::>>() + .add(family) + }; + let mut slots = family_entry.load(Ordering::Acquire); + if slots.is_null() { + let candidate = Self::allocate_child_slot_family(arena, self.key_len()); + slots = match family_entry.compare_exchange( + ptr::null_mut(), + candidate, + Ordering::AcqRel, + Ordering::Acquire, + ) { + Ok(_) => candidate, + Err(winner) => { + debug_assert!( + !winner.is_null(), + "a failed null-to-candidate CAS must observe its winner" + ); + winner + } + }; + } + + // SAFETY: either the Acquire load observed a fully initialized family + // array, or the CAS published/observed one with a Release/Acquire + // synchronization edge. Every family array has exactly `key_len` + // slots, and `key_index` was checked above. Arena lifetime is `'exec`. + unsafe { &*slots.add(key_index) } + } + + fn allocate_child_slot_family(arena: &Handle<'exec>, key_len: usize) -> *mut ChildSlot<'exec> { + debug_assert!(key_len > 0, "a cursor cannot exist for an empty node"); + let layout = Layout::array::>(key_len) + .expect("packed trie dynamic child-slot family layout overflow"); + let mut allocation = arena.alloc_layout(layout); + let slots = allocation.as_mut_ptr().cast::>(); + // SAFETY: `slots` is aligned storage for exactly `key_len` ChildSlots. + // Each no-drop OnceLock is initialized once before the pointer can be + // published. Publishing the first slot as the raw allocation's header + // ties the returned reference (and therefore this copied pointer) to + // the arena lifetime; all trailing slots have the same lifetime. + unsafe { + for key_index in 0..key_len { + ptr::write(slots.add(key_index), OnceLock::new()); + } + allocation + // `OnceLock` has drop glue even when `T` is a reference. + // The arena intentionally forgets these reference-only locks, + // so publish a transparent ManuallyDrop header while retaining + // the ChildSlot pointer used by the initialized trailing array. + .assume_init_no_drop::>>() + .into_ref() as *const ManuallyDrop> + as *mut ChildSlot<'exec> + } + } + + fn dynamic_family_count(&self) -> usize { + assert!( + self.key_len_and_flags & DYNAMIC_CHILDREN != 0, + "only dynamic packed trie nodes store a family count" + ); + // SAFETY: dynamic construction writes one u32 at this derived offset + // before publishing the immutable node. + unsafe { + *(self as *const Self) + .cast::() + .add(Self::dynamic_family_count_offset_for( + self.key_len(), + self.row_len(), + )) + .cast::() as usize + } + } + + pub(crate) fn child_shape(&self) -> ChildShape { + if !self.has_children() { + return ChildShape::Leaf; + } + if self.key_len_and_flags & DYNAMIC_CHILDREN != 0 { + ChildShape::Dynamic { + families: self.dynamic_family_count(), + } + } else { + ChildShape::Direct + } + } + + #[cfg(test)] + fn allocation_bytes(&self) -> usize { + PackedLayout::new(self.key_len(), self.row_len(), self.child_shape()) + .allocation + .size() + } + + fn has_children(&self) -> bool { + self.key_len_and_flags & HAS_CHILDREN != 0 } fn key_len(&self) -> usize { - self.key_len as usize + (self.key_len_and_flags & KEY_LEN_MASK) as usize } fn row_len(&self) -> usize { self.row_len as usize } + + fn encode_key_len(key_len: usize, child_shape: ChildShape) -> u32 { + assert!( + key_len <= KEY_LEN_MASK as usize, + "a packed trie node cannot contain 2^30 or more keys" + ); + key_len as u32 | child_shape.flags() + } + + #[inline] + fn keys_offset() -> usize { + align_up(size_of::(), align_of::()) + } + + #[inline] + fn boundaries_offset_for(key_len: usize) -> usize { + align_up( + Self::keys_offset() + key_len * size_of::(), + align_of::(), + ) + } + + #[inline] + fn rows_offset_for(key_len: usize) -> usize { + align_up( + Self::boundaries_offset_for(key_len) + (key_len + 1) * size_of::(), + align_of::(), + ) + } + + #[inline] + fn direct_children_offset_for(key_len: usize, row_len: usize) -> usize { + align_up( + Self::rows_offset_for(key_len) + row_len * size_of::(), + align_of::>(), + ) + } + + #[inline] + fn dynamic_family_count_offset_for(key_len: usize, row_len: usize) -> usize { + align_up( + Self::rows_offset_for(key_len) + row_len * size_of::(), + align_of::(), + ) + } + + #[inline] + fn dynamic_children_offset_for(key_len: usize, row_len: usize) -> usize { + align_up( + Self::dynamic_family_count_offset_for(key_len, row_len) + size_of::(), + align_of::>>(), + ) + } +} + +#[inline] +fn align_up(offset: usize, alignment: usize) -> usize { + debug_assert!(alignment.is_power_of_two()); + (offset + alignment - 1) & !(alignment - 1) } #[cfg(test)] @@ -285,6 +752,10 @@ mod tests { use egglog_concurrency::SharedArena; use super::*; + use crate::{ + table_shortcuts::fill_table, + table_spec::{Constraint, Table, WrappedTableRef}, + }; fn value(value: usize) -> Value { Value::from_usize(value) @@ -298,6 +769,40 @@ mod tests { ColumnId::from_usize(column) } + #[test] + fn packed_trie_header_is_exactly_two_u32s() { + assert_eq!(size_of::>(), 8); + assert_eq!(align_of::>(), align_of::()); + assert_eq!(std::mem::offset_of!(PackedTrieNode<'static>, row_len), 4); + assert_eq!( + std::mem::offset_of!(PackedTrieNode<'static>, _execution), + 8, + "the invariant lifetime marker must not grow the header" + ); + + assert_eq!( + PackedTrieNode::encode_key_len(KEY_LEN_MASK as usize, ChildShape::Leaf), + KEY_LEN_MASK + ); + assert_eq!( + PackedTrieNode::encode_key_len(KEY_LEN_MASK as usize, ChildShape::Direct), + HAS_CHILDREN | KEY_LEN_MASK + ); + assert_eq!( + PackedTrieNode::encode_key_len( + KEY_LEN_MASK as usize, + ChildShape::Dynamic { families: 2 } + ), + u32::MAX, + ); + } + + #[test] + #[should_panic(expected = "cannot contain 2^30 or more keys")] + fn packed_trie_rejects_key_count_that_uses_shape_bit() { + PackedTrieNode::encode_key_len(DYNAMIC_CHILDREN as usize, ChildShape::Leaf); + } + #[test] fn packed_trie_builds_sorted_ranges() { let arena = SharedArena::new(); @@ -309,16 +814,17 @@ mod tests { (value(3), row(4)), (value(3), row(5)), ]; - let node = PackedTrieNode::build_from_sorted_pairs(&handle, column(2), &pairs, true); + let node = PackedTrieNode::build_from_sorted_pairs(&handle, &pairs, ChildShape::Direct); - assert_eq!(node.column(), column(2)); assert_eq!(node.values(), &[value(1), value(2), value(3)]); assert_eq!(node.boundaries(), &[0, 2, 3, 5]); assert_eq!(node.rows(), &[row(1), row(3), row(2), row(4), row(5)]); assert_eq!(node.find(value(2)), Some(1)); assert_eq!(node.find(value(9)), None); - let SubsetRef::Sparse(one) = node.subset_at(0) else { + let cursor = PackedCursor::new(node, 0); + assert_eq!(cursor.size(), 2); + let SubsetRef::Sparse(one) = cursor.subset() else { panic!("noncontiguous range should be sparse") }; assert_eq!(one.inner(), &[row(1), row(3)]); @@ -332,18 +838,135 @@ mod tests { )); } + #[test] + fn packed_trie_duplicate_rows_are_not_misclassified_as_dense() { + let arena = SharedArena::new(); + let handle = arena.new_handle(); + // The endpoints span exactly three row ids, but row 1 is repeated and + // row 2 is absent. An endpoint-only density check would incorrectly + // turn this into the range 1..4 and admit row 2. + let pairs = [(value(1), row(1)), (value(1), row(1)), (value(1), row(3))]; + let node = PackedTrieNode::build_from_sorted_pairs(&handle, &pairs, ChildShape::Leaf); + + let SubsetRef::Sparse(rows) = node.subset_at(0) else { + panic!("a non-contiguous row sequence with duplicates must stay sparse") + }; + assert_eq!(rows.inner(), &[row(1), row(1), row(3)]); + } + + #[test] + #[should_panic(expected = "packed trie cursor key out of bounds")] + fn packed_cursor_checks_key_ordinal() { + let arena = SharedArena::new(); + let handle = arena.new_handle(); + let node = PackedTrieNode::build_from_sorted_pairs( + &handle, + &[(value(1), row(0))], + ChildShape::Leaf, + ); + let _ = PackedCursor::new(node, 1); + } + + #[test] + fn packed_trie_builds_from_prefiltered_subset_with_reusable_scratch() { + // 65 selected rows exercises the radix path. The parity constraint + // also makes the input subset sparse, while its RowIds remain sorted. + let table = fill_table( + (0..130).map(|row| vec![value(row), value((row * 37) % 11), value(row % 2)]), + 1, + None, + |_, new| Some(new.to_vec()), + ); + let filtered = table.refine_one( + table.all(), + &Constraint::EqConst { + col: column(2), + val: value(0), + }, + ); + assert_eq!(filtered.as_ref().size(), 65); + assert!(matches!(filtered.as_ref(), SubsetRef::Sparse(_))); + + let arena = SharedArena::new(); + let handle = arena.new_handle(); + let mut scratch = vec![(value(999), row(999))]; + WrappedTableRef::with_wrapper(&table, |table| { + let node = PackedTrieNode::build_from_subset( + &handle, + table, + filtered.as_ref(), + column(1), + ChildShape::Direct, + &mut scratch, + ); + let mut expected: Vec<_> = (0..130) + .step_by(2) + .map(|row_id| (value((row_id * 37) % 11), row(row_id))) + .collect(); + expected.sort_unstable(); + + assert_eq!(scratch, expected); + assert_eq!(node.rows().len(), expected.len()); + let expected_values: Vec<_> = (0..11).map(value).collect(); + assert_eq!(node.values(), expected_values.as_slice()); + for key_index in 0..node.values().len() { + let lo = node.boundaries()[key_index] as usize; + let hi = node.boundaries()[key_index + 1] as usize; + let expected_rows: Vec<_> = + expected[lo..hi].iter().map(|&(_, row_id)| row_id).collect(); + assert_eq!(node.rows()[lo..hi], expected_rows); + } + }); + + // Reusing the same Vec for a much smaller subset must clear both old + // projected pairs and the retained radix ping-pong region. + let small = table.refine_one( + table.refine_one( + table.all(), + &Constraint::EqConst { + col: column(2), + val: value(1), + }, + ), + &Constraint::LtConst { + col: column(0), + val: value(20), + }, + ); + WrappedTableRef::with_wrapper(&table, |table| { + let node = PackedTrieNode::build_from_subset( + &handle, + table, + small.as_ref(), + column(1), + ChildShape::Leaf, + &mut scratch, + ); + let mut expected: Vec<_> = (1..20) + .step_by(2) + .map(|row_id| (value((row_id * 37) % 11), row(row_id))) + .collect(); + expected.sort_unstable(); + assert_eq!(scratch, expected); + assert_eq!(node.rows().len(), 10); + }); + } + #[test] fn packed_trie_leaf_omits_child_locks() { let arena = SharedArena::new(); let handle = arena.new_handle(); let pairs = [(value(1), row(0)), (value(2), row(1))]; - let leaf = PackedTrieNode::build_from_sorted_pairs(&handle, column(0), &pairs, false); - let branch = PackedTrieNode::build_from_sorted_pairs(&handle, column(0), &pairs, true); + let leaf = PackedTrieNode::build_from_sorted_pairs(&handle, &pairs, ChildShape::Leaf); + let branch = PackedTrieNode::build_from_sorted_pairs(&handle, &pairs, ChildShape::Direct); - assert!(leaf.child_slot(0).is_none()); - assert!(leaf.child_slot(1).is_none()); - assert!(branch.child_slot(0).is_some()); - assert!(branch.child_slot(1).is_some()); + assert_eq!(leaf.key_len_and_flags, 2); + assert_eq!(branch.key_len_and_flags, HAS_CHILDREN | 2); + assert_eq!(leaf.values(), branch.values()); + assert!(leaf.direct_child_slot(0).is_none()); + assert!(leaf.direct_child_slot(1).is_none()); + assert!(branch.direct_child_slot(0).is_some()); + assert!(branch.direct_child_slot(1).is_some()); assert!(leaf.allocation_bytes() < branch.allocation_bytes()); } @@ -352,7 +975,7 @@ mod tests { let arena = SharedArena::new(); let handle = arena.new_handle(); let pairs = [(value(1), row(0)), (value(2), row(2))]; - let node = PackedTrieNode::build_from_sorted_pairs(&handle, column(1), &pairs, true); + let node = PackedTrieNode::build_from_sorted_pairs(&handle, &pairs, ChildShape::Direct); assert_eq!( node as *const _ as usize % align_of::>(), @@ -362,11 +985,20 @@ mod tests { assert_eq!(node.boundaries().as_ptr() as usize % align_of::(), 0); assert_eq!(node.rows().as_ptr() as usize % align_of::(), 0); assert_eq!( - node.child_slot(0).unwrap() as *const _ as usize + node.direct_child_slot(0).unwrap() as *const _ as usize % align_of::>>(), 0 ); - assert!(node.allocation_bytes() >= size_of::>()); + assert_eq!( + node.values().as_ptr() as usize - node as *const _ as usize, + size_of::>() + ); + assert_eq!( + node.allocation_bytes(), + PackedLayout::new(2, 2, ChildShape::Direct) + .allocation + .size() + ); } #[test] @@ -375,11 +1007,10 @@ mod tests { let handle = arena.new_handle(); let parent = PackedTrieNode::build_from_sorted_pairs( &handle, - column(0), &[(value(1), row(0))], - true, + ChildShape::Direct, ); - let slot = parent.child_slot(0).unwrap(); + let slot = parent.direct_child_slot(0).unwrap(); let initializations = AtomicUsize::new(0); std::thread::scope(|scope| { @@ -392,9 +1023,8 @@ mod tests { let handle = arena.new_handle(); PackedTrieNode::build_from_sorted_pairs( &handle, - column(1), &[(value(7), row(9))], - false, + ChildShape::Leaf, ) }); assert_eq!(child.values(), &[value(7)]); @@ -403,7 +1033,262 @@ mod tests { }); assert_eq!(initializations.load(Ordering::Relaxed), 1); - assert_eq!(slot.get().unwrap().column(), column(1)); assert_eq!(slot.get().unwrap().values(), &[value(7)]); } + + #[test] + fn packed_cursor_child_index_reuses_cached_node() { + let table = fill_table( + (0..8).map(|row_id| vec![value(row_id), value(row_id % 3), value(7 - row_id)]), + 1, + None, + |_, new| Some(new.to_vec()), + ); + let arena = SharedArena::new(); + let handle = arena.new_handle(); + let parent_pairs: Vec<_> = (0..8).map(|row_id| (value(0), row(row_id))).collect(); + let parent = + PackedTrieNode::build_from_sorted_pairs(&handle, &parent_pairs, ChildShape::Direct); + let cursor = PackedCursor::new(parent, 0); + + WrappedTableRef::with_wrapper(&table, |table| { + let mut scratch = Vec::new(); + let first = + cursor.child_index(&handle, table, column(1), 7, ChildShape::Leaf, &mut scratch); + assert_eq!(first.values(), &[value(0), value(1), value(2)]); + + let sentinel = (value(999), row(999)); + scratch.clear(); + scratch.push(sentinel); + let second = + cursor.child_index(&handle, table, column(1), 7, ChildShape::Leaf, &mut scratch); + assert!(ptr::eq(first, second)); + assert_eq!(scratch, &[sentinel], "a cache hit must not use scratch"); + }); + } + + #[test] + fn packed_dynamic_children_keep_two_successor_families_distinct() { + let table = fill_table( + (0..12).map(|row_id| vec![value(row_id), value(row_id % 2), value((row_id * 5) % 3)]), + 1, + None, + |_, new| Some(new.to_vec()), + ); + let arena = SharedArena::new(); + { + let handle = arena.new_handle(); + let parent_pairs: Vec<_> = (0..12).map(|row_id| (value(0), row(row_id))).collect(); + let parent = PackedTrieNode::build_from_sorted_pairs( + &handle, + &parent_pairs, + ChildShape::Dynamic { families: 2 }, + ); + assert_eq!(parent.child_shape(), ChildShape::Dynamic { families: 2 }); + assert_eq!( + parent.key_len_and_flags, + HAS_CHILDREN | DYNAMIC_CHILDREN | 1 + ); + let cursor = PackedCursor::new(parent, 0); + + WrappedTableRef::with_wrapper(&table, |table| { + let mut scratch = Vec::new(); + let first_family = cursor.child_index( + &handle, + table, + column(1), + 0, + ChildShape::Leaf, + &mut scratch, + ); + assert_eq!(first_family.values(), &[value(0), value(1)]); + + let second_family = cursor.child_index( + &handle, + table, + column(2), + 1, + ChildShape::Leaf, + &mut scratch, + ); + assert_eq!(second_family.values(), &[value(0), value(1), value(2)]); + assert!(!ptr::eq(first_family, second_family)); + + let cached_first = cursor.child_index( + &handle, + table, + column(1), + 0, + ChildShape::Leaf, + &mut scratch, + ); + assert!(ptr::eq(first_family, cached_first)); + }); + } + } + + #[test] + fn packed_dynamic_same_family_converges_under_race() { + let table = fill_table( + (0..256).map(|row_id| vec![value(row_id), value((row_id * 17) % 31)]), + 1, + None, + |_, new| Some(new.to_vec()), + ); + let arena = SharedArena::new(); + let handle = arena.new_handle(); + let parent_pairs: Vec<_> = (0..256).map(|row_id| (value(0), row(row_id))).collect(); + let parent = PackedTrieNode::build_from_sorted_pairs( + &handle, + &parent_pairs, + ChildShape::Dynamic { families: 2 }, + ); + let cursor = PackedCursor::new(parent, 0); + let published = AtomicUsize::new(0); + let start = std::sync::Barrier::new(16); + + WrappedTableRef::with_wrapper(&table, |table| { + std::thread::scope(|scope| { + for _ in 0..16 { + let arena = &arena; + let published = &published; + let start = &start; + scope.spawn(move || { + let handle = arena.new_handle(); + let mut scratch = Vec::new(); + start.wait(); + let child = cursor.child_index( + &handle, + table, + column(1), + 1, + ChildShape::Leaf, + &mut scratch, + ); + assert_eq!(child.values().len(), 31); + let address = child as *const _ as usize; + let observed = published + .compare_exchange(0, address, Ordering::AcqRel, Ordering::Acquire) + .unwrap_or_else(|already_published| already_published); + assert!(observed == 0 || observed == address); + }); + } + }); + }); + + assert_ne!(published.load(Ordering::Acquire), 0); + } + + #[test] + fn packed_cursor_child_index_uses_prefiltered_subset() { + let table = fill_table( + (0..8).map(|row_id| vec![value(row_id), value(row_id % 3), value(7 - row_id)]), + 1, + None, + |_, new| Some(new.to_vec()), + ); + let filtered = table.refine_one( + table.all(), + &Constraint::LtConst { + col: column(0), + val: value(4), + }, + ); + let arena = SharedArena::new(); + let handle = arena.new_handle(); + let parent_pairs: Vec<_> = (0..8).map(|row_id| (value(0), row(row_id))).collect(); + let parent = + PackedTrieNode::build_from_sorted_pairs(&handle, &parent_pairs, ChildShape::Direct); + let cursor = PackedCursor::new(parent, 0); + + WrappedTableRef::with_wrapper(&table, |table| { + let mut scratch = Vec::new(); + let child = cursor.child_index_from_subset( + &handle, + table, + filtered.as_ref(), + column(1), + 0, + ChildShape::Leaf, + &mut scratch, + ); + assert_eq!(child.rows().len(), 4); + assert!(child.rows().iter().all(|row_id| row_id.index() < 4)); + }); + } + + #[test] + fn packed_cursor_child_index_is_published_once_under_race() { + let table = fill_table( + (0..96).map(|row_id| vec![value(row_id), value((row_id * 17) % 13), value(row_id % 5)]), + 1, + None, + |_, new| Some(new.to_vec()), + ); + let arena = SharedArena::new(); + let handle = arena.new_handle(); + let parent_pairs: Vec<_> = (0..96).map(|row_id| (value(0), row(row_id))).collect(); + let parent = + PackedTrieNode::build_from_sorted_pairs(&handle, &parent_pairs, ChildShape::Direct); + let cursor = PackedCursor::new(parent, 0); + let published = AtomicUsize::new(0); + + WrappedTableRef::with_wrapper(&table, |table| { + std::thread::scope(|scope| { + for _ in 0..8 { + let arena = &arena; + let published = &published; + scope.spawn(move || { + let handle = arena.new_handle(); + let mut scratch = Vec::new(); + let child = cursor.child_index( + &handle, + table, + column(1), + 0, + ChildShape::Leaf, + &mut scratch, + ); + assert_eq!(child.values().len(), 13); + let address = child as *const _ as usize; + let observed = published + .compare_exchange(0, address, Ordering::AcqRel, Ordering::Acquire) + .unwrap_or_else(|already_published| already_published); + assert!(observed == 0 || observed == address); + }); + } + }); + }); + assert_ne!(published.load(Ordering::Acquire), 0); + } + + #[test] + #[should_panic(expected = "packed trie child cache shape mismatch")] + fn packed_cursor_rejects_cached_child_with_another_shape() { + let table = fill_table( + (0..4).map(|row_id| vec![value(row_id), value(row_id % 2), value(row_id + 10)]), + 1, + None, + |_, new| Some(new.to_vec()), + ); + let arena = SharedArena::new(); + let handle = arena.new_handle(); + let parent_pairs: Vec<_> = (0..4).map(|row_id| (value(0), row(row_id))).collect(); + let parent = + PackedTrieNode::build_from_sorted_pairs(&handle, &parent_pairs, ChildShape::Direct); + let cursor = PackedCursor::new(parent, 0); + + WrappedTableRef::with_wrapper(&table, |table| { + let mut scratch = Vec::new(); + cursor.child_index(&handle, table, column(1), 0, ChildShape::Leaf, &mut scratch); + cursor.child_index( + &handle, + table, + column(1), + 0, + ChildShape::Direct, + &mut scratch, + ); + }); + } } diff --git a/core-relations/src/hash_index/mod.rs b/core-relations/src/hash_index/mod.rs index 0cd2f0cb3..ea15f45ee 100644 --- a/core-relations/src/hash_index/mod.rs +++ b/core-relations/src/hash_index/mod.rs @@ -51,6 +51,38 @@ pub(crate) struct Index { table: TI, } +/// A stable position for one key within a frozen index. +/// +/// Positions are scoped to the exact index that produced them. They may be +/// retained while that index is only read. Existing positions also survive an +/// append-only incremental refresh, but a clear/full rebuild invalidates every +/// position. The two components are deliberately private so callers cannot +/// manufacture positions belonging to another index. +#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)] +#[allow(dead_code)] // Consumed by the packed-executor integration layer. +pub(crate) struct IndexPosition { + shard: u32, + slot: u32, +} + +#[allow(dead_code)] // Consumed by the packed-executor integration layer. +impl IndexPosition { + fn new(shard: ShardId, slot: usize) -> Self { + Self { + shard: u32::try_from(shard.index()).expect("index shard exceeds u32::MAX"), + slot: u32::try_from(slot).expect("index shard contains more than u32::MAX keys"), + } + } + + pub(crate) fn shard(self) -> usize { + self.shard as usize + } + + pub(crate) fn slot(self) -> usize { + self.slot as usize + } +} + impl Index { pub(crate) fn new(key: Vec, table: TI) -> Self { Index { @@ -135,10 +167,6 @@ impl Index { } /// Call `f` over the elements stored in one physical index shard. - /// - /// The shards form a disjoint partition of the index keys. Exposing that - /// partition lets query execution create a small number of coarse tasks - /// without first copying keys into a separate work list. pub(crate) fn for_each_shard(&self, shard: usize, f: impl FnMut(&TI::Key, SubsetRef)) { self.table.for_each_shard(shard, f); } @@ -158,6 +186,34 @@ impl Index { } } +#[allow(dead_code)] // Consumed by the packed-executor integration layer. +impl Index { + /// Get a key's execution-stable position along with its nonempty subset. + pub(crate) fn get_subset_positioned<'a>( + &'a self, + key: &TI::Key, + ) -> Option<(IndexPosition, SubsetRef<'a>)> { + self.table.get_subset_positioned(key) + } + + /// Visit every key together with its execution-stable position. + pub(crate) fn for_each_positioned<'a>( + &'a self, + f: impl FnMut(IndexPosition, &'a TI::Key, SubsetRef<'a>), + ) { + self.table.for_each_positioned(f); + } + + /// Visit one physical shard together with each key's stable position. + pub(crate) fn for_each_shard_positioned<'a>( + &'a self, + shard: usize, + f: impl FnMut(IndexPosition, &'a TI::Key, SubsetRef<'a>), + ) { + self.table.for_each_shard_positioned(shard, f); + } +} + pub(crate) struct SubsetTable { keys: RowBuffer, hash: Pooled>>, @@ -232,6 +288,30 @@ pub(crate) trait IndexBase { } } +/// Read-only index operations that also expose a stable per-index key position. +/// +/// Implementations may choose different physical slot schemes. Positions only +/// need to remain stable while the index is frozen; callers must treat them as +/// opaque tokens belonging to the exact index that returned them. +#[allow(dead_code)] // Consumed by the packed-executor integration layer. +pub(crate) trait PositionedIndexBase: IndexBase { + fn get_subset_positioned<'a>( + &'a self, + key: &Self::Key, + ) -> Option<(IndexPosition, SubsetRef<'a>)>; + + fn for_each_positioned<'a>( + &'a self, + f: impl FnMut(IndexPosition, &'a Self::Key, SubsetRef<'a>), + ); + + fn for_each_shard_positioned<'a>( + &'a self, + shard: usize, + f: impl FnMut(IndexPosition, &'a Self::Key, SubsetRef<'a>), + ); +} + struct ColumnIndexShard { /// It's important that table is implemented using IndexMap instead of the more efficient /// HashMap because we want stable enumeration order. @@ -299,12 +379,12 @@ impl IndexBase for ColumnIndex { } fn for_each(&self, mut f: impl FnMut(&Self::Key, SubsetRef)) { - for (subsets, (k, v)) in self + for (subsets, (key, subset)) in self .shards .iter() - .flat_map(|(_, shard)| shard.table.iter().map(|x| (&shard.subsets, x))) + .flat_map(|(_, shard)| shard.table.iter().map(|entry| (&shard.subsets, entry))) { - f(k, v.as_ref(subsets)); + f(key, subset.as_ref(subsets)); } } @@ -446,6 +526,51 @@ impl IndexBase for ColumnIndex { } } +impl PositionedIndexBase for ColumnIndex { + fn get_subset_positioned<'a>(&'a self, key: &Value) -> Option<(IndexPosition, SubsetRef<'a>)> { + let mut hasher = FxHasher::default(); + key.hash(&mut hasher); + let shard_id = self.shard_data.shard_id(hasher.finish()); + let shard = &self.shards[shard_id]; + let (slot, _, subset) = shard.table.get_full(key)?; + Some(( + IndexPosition::new(shard_id, slot), + subset.as_ref(&shard.subsets), + )) + } + + fn for_each_positioned<'a>( + &'a self, + mut f: impl FnMut(IndexPosition, &'a Value, SubsetRef<'a>), + ) { + for (shard_id, shard) in self.shards.iter() { + for (slot, (key, subset)) in shard.table.iter().enumerate() { + f( + IndexPosition::new(shard_id, slot), + key, + subset.as_ref(&shard.subsets), + ); + } + } + } + + fn for_each_shard_positioned<'a>( + &'a self, + shard: usize, + mut f: impl FnMut(IndexPosition, &'a Value, SubsetRef<'a>), + ) { + let shard_id = ShardId::from_usize(shard); + let shard = &self.shards[shard_id]; + for (slot, (key, subset)) in shard.table.iter().enumerate() { + f( + IndexPosition::new(shard_id, slot), + key, + subset.as_ref(&shard.subsets), + ); + } + } +} + /// Number of 8-bit radix passes needed to cover values up to `max`. fn radix_passes_for(max: u32) -> u32 { if max < 256 { @@ -752,6 +877,7 @@ impl IndexBase for TupleIndex { self.add_row(key, src_id); } } + fn for_each(&self, mut f: impl FnMut(&Self::Key, SubsetRef)) { for (_, shard) in self.shards.iter() { for entry in shard.table.hash.iter() { @@ -874,6 +1000,60 @@ impl IndexBase for TupleIndex { } } +impl PositionedIndexBase for TupleIndex { + fn get_subset_positioned<'a>( + &'a self, + key: &[Value], + ) -> Option<(IndexPosition, SubsetRef<'a>)> { + let hash = hash_key(key); + let shard_id = self.shard_data.shard_id(hash); + let shard = &self.shards[shard_id]; + let entry = shard.table.hash.find(hash, |entry| { + // SAFETY: entry.key was stored by add_row, which returns a valid RowId. + entry.hash == hash && unsafe { shard.table.keys.get_row_unchecked(entry.key) } == key + })?; + Some(( + IndexPosition::new(shard_id, entry.key.index()), + entry.vals.as_ref(&shard.subsets), + )) + } + + fn for_each_positioned<'a>( + &'a self, + mut f: impl FnMut(IndexPosition, &'a [Value], SubsetRef<'a>), + ) { + for (shard_id, shard) in self.shards.iter() { + for entry in shard.table.hash.iter() { + // SAFETY: entry.key was stored by add_row, so it is always in-bounds. + let key = unsafe { shard.table.keys.get_row_unchecked(entry.key) }; + f( + IndexPosition::new(shard_id, entry.key.index()), + key, + entry.vals.as_ref(&shard.subsets), + ); + } + } + } + + fn for_each_shard_positioned<'a>( + &'a self, + shard: usize, + mut f: impl FnMut(IndexPosition, &'a [Value], SubsetRef<'a>), + ) { + let shard_id = ShardId::from_usize(shard); + let shard = &self.shards[shard_id]; + for entry in shard.table.hash.iter() { + // SAFETY: entry.key was stored by add_row, so it is always in-bounds. + let key = unsafe { shard.table.keys.get_row_unchecked(entry.key) }; + f( + IndexPosition::new(shard_id, entry.key.index()), + key, + entry.vals.as_ref(&shard.subsets), + ); + } + } +} + fn hash_key(key: &[Value]) -> u64 { let mut hasher = FxHasher::default(); key.hash(&mut hasher); diff --git a/core-relations/src/hash_index/tests.rs b/core-relations/src/hash_index/tests.rs index d80914367..317fd5fb2 100644 --- a/core-relations/src/hash_index/tests.rs +++ b/core-relations/src/hash_index/tests.rs @@ -3,7 +3,11 @@ use std::collections::BTreeMap; use egglog_concurrency::ThreadPool; use rand::{Rng, SeedableRng, rngs::StdRng}; -use crate::{common::Value, numeric_id::NumericId, offsets::Offsets}; +use crate::{ + common::{HashSet, Value}, + numeric_id::NumericId, + offsets::{Offsets, RowId, SubsetRef}, +}; use crate::{ TupleIndex, @@ -12,7 +16,13 @@ use crate::{ table_spec::{ColumnId, WrappedTable}, }; -use super::{Index, IndexBase}; +use super::{Index, IndexBase, IndexPosition, PositionedIndexBase}; + +fn subset_rows(subset: SubsetRef<'_>) -> Vec { + let mut rows = Vec::new(); + subset.offsets(|row| rows.push(row.index())); + rows +} #[test] fn basic_updates() { @@ -158,7 +168,7 @@ fn oracle(rows: &[Vec], cols: &[usize]) -> BTreeMap> { /// Collect a built [`ColumnIndex`] into `value -> row ids` for comparison against [`oracle`]. fn collect(index: &ColumnIndex) -> BTreeMap> { let mut got: BTreeMap> = BTreeMap::new(); - index.for_each(|val, subset| { + index.for_each_positioned(|_, val, subset| { let mut ids = Vec::new(); subset.offsets(|row_id| ids.push(row_id.index())); got.insert(val.rep(), ids); @@ -240,15 +250,16 @@ fn physical_shards_partition_index_iteration() { let mut column = Index::new(vec![ColumnId::new(1)], ColumnIndex::new()); column.refresh(table.as_ref()); + assert!(column.get_subset_positioned(&v(9_999_999)).is_none()); assert_eq!(column.shard_count(), 8); let mut whole_column = BTreeMap::new(); - column.for_each(|key, subset| { + column.for_each_positioned(|_, key, subset| { whole_column.insert(key.rep(), subset.size()); }); let mut by_column_shard = BTreeMap::new(); for shard in 0..column.shard_count() { let mut shard_keys = 0; - column.for_each_shard(shard, |key, subset| { + column.for_each_shard_positioned(shard, |_, key, subset| { shard_keys += 1; assert!( by_column_shard.insert(key.rep(), subset.size()).is_none(), @@ -266,13 +277,13 @@ fn physical_shards_partition_index_iteration() { tuple.refresh(table.as_ref()); assert_eq!(tuple.shard_count(), 8); let mut whole_tuple = BTreeMap::new(); - tuple.for_each(|key, subset| { + tuple.for_each_positioned(|_, key, subset| { whole_tuple.insert((key[0].rep(), key[1].rep()), subset.size()); }); let mut by_tuple_shard = BTreeMap::new(); for shard in 0..tuple.shard_count() { let mut shard_keys = 0; - tuple.for_each_shard(shard, |key, subset| { + tuple.for_each_shard_positioned(shard, |_, key, subset| { shard_keys += 1; assert!( by_tuple_shard @@ -286,3 +297,174 @@ fn physical_shards_partition_index_iteration() { assert_eq!(by_tuple_shard, whole_tuple); }); } + +#[test] +fn positioned_lookups_and_iteration_agree() { + ThreadPool::new(4).install(|| { + let rows = (0..257) + .map(|i| vec![v(i), v(i % 37), v(10_000 + i % 53)]) + .collect::>(); + let table = WrappedTable::new(fill_table(rows, 1, None, |old, new| { + assert_eq!(old, new, "unique keys, so no conflicts"); + None + })); + + let mut column = Index::new(vec![ColumnId::new(1)], ColumnIndex::new()); + column.refresh(table.as_ref()); + let mut positioned_column = BTreeMap::new(); + let mut column_positions = HashSet::default(); + column.for_each_positioned(|position, key, subset| { + let (lookup_position, lookup_subset) = + column.get_subset_positioned(key).expect("enumerated key"); + assert_eq!(lookup_position, position); + assert_eq!(subset_rows(lookup_subset), subset_rows(subset)); + assert!(position.slot() < column.shard_len(position.shard())); + assert!(column_positions.insert(position)); + positioned_column.insert(key.rep(), (position, subset_rows(subset))); + }); + assert_eq!(positioned_column.len(), column.len()); + + let mut legacy_column = BTreeMap::new(); + column.for_each_positioned(|_, key, subset| { + legacy_column.insert(key.rep(), subset_rows(subset)); + }); + assert_eq!( + legacy_column, + positioned_column + .iter() + .map(|(key, (_, rows))| (*key, rows.clone())) + .collect() + ); + + let mut column_by_shard = BTreeMap::new(); + for shard in 0..column.shard_count() { + let mut slots = Vec::new(); + column.for_each_shard_positioned(shard, |position, key, subset| { + assert_eq!(position.shard(), shard); + slots.push(position.slot()); + assert_eq!( + positioned_column.get(&key.rep()), + Some(&(position, subset_rows(subset))) + ); + assert!(column_by_shard.insert(key.rep(), position).is_none()); + }); + slots.sort_unstable(); + assert_eq!(slots, (0..column.shard_len(shard)).collect::>()); + } + assert_eq!(column_by_shard.len(), positioned_column.len()); + + let mut tuple = Index::new(vec![ColumnId::new(1), ColumnId::new(2)], TupleIndex::new(2)); + tuple.refresh(table.as_ref()); + assert!( + tuple + .get_subset_positioned(&[v(9_999_999), v(9_999_998)]) + .is_none() + ); + let mut positioned_tuple = BTreeMap::new(); + let mut tuple_positions = HashSet::default(); + tuple.for_each_positioned(|position, key, subset| { + let (lookup_position, lookup_subset) = + tuple.get_subset_positioned(key).expect("enumerated key"); + assert_eq!(lookup_position, position); + assert_eq!(subset_rows(lookup_subset), subset_rows(subset)); + assert!(position.slot() < tuple.shard_len(position.shard())); + assert!(tuple_positions.insert(position)); + positioned_tuple.insert( + (key[0].rep(), key[1].rep()), + (position, subset_rows(subset)), + ); + }); + assert_eq!(positioned_tuple.len(), tuple.len()); + + let mut legacy_tuple = BTreeMap::new(); + tuple.for_each_positioned(|_, key, subset| { + legacy_tuple.insert((key[0].rep(), key[1].rep()), subset_rows(subset)); + }); + assert_eq!( + legacy_tuple, + positioned_tuple + .iter() + .map(|(key, (_, rows))| (*key, rows.clone())) + .collect() + ); + + let mut tuple_by_shard = BTreeMap::new(); + for shard in 0..tuple.shard_count() { + let mut slots = Vec::new(); + tuple.for_each_shard_positioned(shard, |position, key, subset| { + assert_eq!(position.shard(), shard); + slots.push(position.slot()); + let key = (key[0].rep(), key[1].rep()); + assert_eq!( + positioned_tuple.get(&key), + Some(&(position, subset_rows(subset))) + ); + assert!(tuple_by_shard.insert(key, position).is_none()); + }); + slots.sort_unstable(); + assert_eq!(slots, (0..tuple.shard_len(shard)).collect::>()); + } + assert_eq!(tuple_by_shard.len(), positioned_tuple.len()); + }); +} + +#[test] +fn positions_survive_incremental_insertions() { + ThreadPool::new(4).install(|| { + let mut column = Index::new(vec![ColumnId::new(0)], ColumnIndex::new()); + column.table.add_row(&[v(10)], RowId::new(0)); + column.table.add_row(&[v(20)], RowId::new(1)); + column.table.add_row(&[v(10)], RowId::new(2)); + let old_column = [v(10), v(20)].map(|key| { + ( + key, + column + .get_subset_positioned(&key) + .expect("initial column key") + .0, + ) + }); + + column.table.add_row(&[v(10)], RowId::new(3)); + column.table.add_row(&[v(30)], RowId::new(4)); + for (key, old_position) in old_column { + let (position, _) = column + .get_subset_positioned(&key) + .expect("retained column key"); + assert_eq!(position, old_position); + } + let mut column_positions = HashSet::::default(); + column.for_each_positioned(|position, _, _| { + assert!(column_positions.insert(position)); + }); + assert_eq!(column_positions.len(), column.len()); + + let mut tuple = Index::new(vec![ColumnId::new(0), ColumnId::new(1)], TupleIndex::new(2)); + tuple.table.add_row(&[v(10), v(100)], RowId::new(0)); + tuple.table.add_row(&[v(20), v(200)], RowId::new(1)); + tuple.table.add_row(&[v(10), v(100)], RowId::new(2)); + let old_tuple = [[v(10), v(100)], [v(20), v(200)]].map(|key| { + ( + key, + tuple + .get_subset_positioned(&key) + .expect("initial tuple key") + .0, + ) + }); + + tuple.table.add_row(&[v(10), v(100)], RowId::new(3)); + tuple.table.add_row(&[v(30), v(300)], RowId::new(4)); + for (key, old_position) in old_tuple { + let (position, _) = tuple + .get_subset_positioned(&key) + .expect("retained tuple key"); + assert_eq!(position, old_position); + } + let mut tuple_positions = HashSet::::default(); + tuple.for_each_positioned(|position, _, _| { + assert!(tuple_positions.insert(position)); + }); + assert_eq!(tuple_positions.len(), tuple.len()); + }); +} From 51ba5b9a8a4286e3442f28292327a63ce7419741 Mon Sep 17 00:00:00 2001 From: Eli Rosenthal Date: Sun, 9 Aug 2026 22:46:17 -0700 Subject: [PATCH 3/3] Document packed trie arena layout --- core-relations/src/free_join/packed_trie.rs | 39 +++++++++++++++++++-- 1 file changed, 36 insertions(+), 3 deletions(-) diff --git a/core-relations/src/free_join/packed_trie.rs b/core-relations/src/free_join/packed_trie.rs index 9343fe427..e73d0561a 100644 --- a/core-relations/src/free_join/packed_trie.rs +++ b/core-relations/src/free_join/packed_trie.rs @@ -68,9 +68,42 @@ impl ChildShape { /// One immutable, scalar trie level allocated for the lifetime of an execution. /// -/// The dynamically sized portions immediately following this header are laid -/// out as keys, boundaries, grouped row ids, and (optionally) child locks. Use -/// the accessors rather than relying on those offsets outside this module. +/// Each node starts with the same packed index. `boundaries[i]..boundaries[i + 1]` +/// selects the rows belonging to `keys[i]`: +/// +/// ```text +/// +----------------------------+ +/// | PackedTrieNode header | +/// +----------------------------+ +/// | Value keys[K] | +/// +----------------------------+ +/// | u32 boundaries[K + 1] | +/// +----------------------------+ +/// | RowId rows[R] | +/// +----------------------------+ +/// ``` +/// +/// A leaf ends there. A direct node appends one inline child-publication slot +/// per key: +/// +/// ```text +/// | ChildSlot direct[K] | +/// +----------------------------+ +/// ``` +/// +/// A dynamically ordered node instead appends a family count and one atomic +/// pointer per eligible successor. Each pointer lazily publishes a separate +/// arena allocation containing one `ChildSlot[K]` array: +/// +/// ```text +/// | u32 family_count | family pointer +/// +----------------------------+ | +/// | AtomicPtr families[F] |-------------+--> ChildSlot family[K] +/// +----------------------------+ +/// ``` +/// +/// Alignment may insert padding between regions. Use the accessors rather +/// than relying on these offsets outside this module. #[repr(C)] #[derive(Debug)] pub(crate) struct PackedTrieNode<'exec> {