From a8e1a681fce8cb58ade3ab04b8ff6fdc3c34cd9a Mon Sep 17 00:00:00 2001 From: Eli Rosenthal Date: Sun, 19 Jul 2026 22:59:37 +0200 Subject: [PATCH 1/5] Scale generic joins with locality-aware scheduling Add private worker queues and scheduler metrics, partition generic joins by physical index shard, and retain prepared index handles for recursive probes. --- concurrency/src/lib.rs | 4 +- concurrency/src/threadpool/mod.rs | 657 ++++++++- concurrency/src/threadpool/tests.rs | 264 +++- core-relations/src/free_join/execute.rs | 1638 +++++++++++++++++++-- core-relations/src/hash_index/mod.rs | 72 + core-relations/src/hash_index/tests.rs | 60 + core-relations/src/parallel_heuristics.rs | 77 + core-relations/src/pool/mod.rs | 4 +- core-relations/src/tests.rs | 585 +++++++- 9 files changed, 3208 insertions(+), 153 deletions(-) diff --git a/concurrency/src/lib.rs b/concurrency/src/lib.rs index 322d2d67e..9d5993e4e 100644 --- a/concurrency/src/lib.rs +++ b/concurrency/src/lib.rs @@ -17,7 +17,9 @@ pub use notification_list::NotificationList; pub use parallel_writer::ParallelVecWriter; pub use resettable_oncelock::ResettableOnceLock; pub use shared_arena::{Handle, SharedArena, SharedRef}; -pub use threadpool::{Scope, ThreadPool, current_num_threads, scope, without_current_pool}; +pub use threadpool::{ + SchedulerMetrics, Scope, ThreadPool, current_num_threads, scope, without_current_pool, +}; #[cfg(test)] mod tests; diff --git a/concurrency/src/threadpool/mod.rs b/concurrency/src/threadpool/mod.rs index 427b7a559..1d7b3ae9b 100644 --- a/concurrency/src/threadpool/mod.rs +++ b/concurrency/src/threadpool/mod.rs @@ -1,12 +1,15 @@ -//! A small scoped thread pool backed by a shared crossbeam channel. +//! A small scoped thread pool with global and worker-local work queues. //! //! [`ThreadPool`] owns a fixed set of primary worker threads. Each worker -//! receives boxed `'static` jobs from a shared channel and exits when the -//! channel is closed. A primary worker that blocks waiting for nested scoped -//! work helps drain queued jobs until that nested scope completes, so the pool -//! does not lose a worker while work it needs may still be queued. [`Scope`] -//! provides the safe scoped API on top of those `'static` jobs by erasing task -//! lifetimes when work is sent to the channel, then waiting for all work in the +//! owns a private deque in addition to receiving boxed `'static` jobs from a +//! shared channel. Local work is pushed and popped at the back, giving the +//! owner depth-first execution; when another worker is stalled, half of the +//! oldest local work is donated from the front to the shared queue. A primary +//! worker that blocks waiting for nested scoped work helps drain its local +//! deque first and then the shared queue until that nested scope completes, so +//! the pool does not lose a worker while work it needs may still be queued. +//! [`Scope`] provides the safe scoped API on top of those `'static` jobs by +//! erasing task lifetimes when work is queued, then waiting for all work in the //! scope before returning to the caller. //! //! The root callback runs on the caller thread, and spawned callbacks run on @@ -95,22 +98,22 @@ use std::{ any::Any, - cell::Cell, + cell::{Cell, UnsafeCell}, + collections::VecDeque, marker::PhantomData, mem, panic::{self, AssertUnwindSafe}, ptr::{self, NonNull}, sync::{ Mutex, - atomic::{AtomicU64, Ordering}, + atomic::{AtomicU64, AtomicUsize, Ordering}, }, thread::{self, JoinHandle}, + time::{Duration, Instant}, }; -#[cfg(test)] -use std::sync::atomic::AtomicUsize; - use crossbeam::channel::{Receiver, Sender, bounded, select_biased, unbounded}; +use crossbeam::utils::CachePadded; const EXPECTED_SHIFT: u32 = 32; const COMPLETED_MASK: u64 = u32::MAX as u64; @@ -124,10 +127,69 @@ type PanicPayload = Box; thread_local! { static CURRENT_POOL: Cell<*const ThreadPoolState> = const { Cell::new(ptr::null()) }; + static CURRENT_WORKER: Cell<*const WorkerContext> = const { Cell::new(ptr::null()) }; static IS_BACKGROUND_WORKER: Cell = const { Cell::new(false) }; static INLINE_SCOPE_HELP_DEPTH: Cell = const { Cell::new(0) }; } +/// A cumulative snapshot of a thread pool's scheduler instrumentation. +/// +/// Counter fields and `stalled_time` are cumulative since the pool was created +/// or [`ThreadPool::reset_scheduler_metrics`] was called. Queue-depth fields +/// are high-water marks over that same period. `stalled_workers` is a gauge at +/// the instant of the snapshot rather than a cumulative value. +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +pub struct SchedulerMetrics { + /// Jobs pushed onto the shared global queue, including donations. + pub global_pushes: u64, + /// Jobs popped from the shared global queue. + pub global_pops: u64, + /// Jobs successfully pushed onto worker-private queues. + pub local_pushes: u64, + /// Jobs popped by their owning worker from a private queue. + pub local_pops: u64, + /// Jobs moved from private queues to the global queue. + pub donated_jobs: u64, + /// Aggregate wall-clock time for which primary workers were stalled. + /// + /// This includes the elapsed portion of stalls that are still in progress + /// when the snapshot is taken, which makes before/after snapshots useful + /// even when the pool is idle at both endpoints. + pub stalled_time: Duration, + /// Number of primary workers currently waiting for global work. + pub stalled_workers: usize, + /// Largest number of jobs observed in any one worker-private queue. + pub max_local_queue_depth: usize, + /// Largest number of jobs observed in the global queue. + pub max_global_queue_depth: usize, +} + +impl SchedulerMetrics { + /// Return the cumulative counter and stalled-time difference from `earlier`. + /// + /// Queue high-water marks cannot in general be subtracted. The returned + /// value therefore retains the later snapshot's high-water marks and its + /// current `stalled_workers` gauge. Call + /// [`ThreadPool::reset_scheduler_metrics`] immediately before an isolated + /// measurement when interval-specific queue high-water marks are needed. + pub fn delta_since(self, earlier: Self) -> Self { + Self { + global_pushes: self.global_pushes.saturating_sub(earlier.global_pushes), + global_pops: self.global_pops.saturating_sub(earlier.global_pops), + local_pushes: self.local_pushes.saturating_sub(earlier.local_pushes), + local_pops: self.local_pops.saturating_sub(earlier.local_pops), + donated_jobs: self.donated_jobs.saturating_sub(earlier.donated_jobs), + stalled_time: self + .stalled_time + .checked_sub(earlier.stalled_time) + .unwrap_or_default(), + stalled_workers: self.stalled_workers, + max_local_queue_depth: self.max_local_queue_depth, + max_global_queue_depth: self.max_global_queue_depth, + } + } +} + /// Return the number of threads in the currently installed thread pool. /// /// Returns `1` when called outside [`ThreadPool::install`], [`ThreadPool::scope`], @@ -222,11 +284,11 @@ where }) } -/// A thread pool using a single shared work queue. +/// A thread pool with a shared queue and a private deque per primary worker. /// /// The pool owns a fixed number of primary workers. A primary worker that -/// blocks waiting for nested scoped work helps drain the shared queue until the -/// nested scope completes. +/// blocks waiting for nested scoped work helps drain its private deque before +/// the shared queue until the nested scope completes. /// /// # Examples /// @@ -268,7 +330,7 @@ impl ThreadPool { let state = Box::new(ThreadPoolState::new(sender, receiver.clone(), thread_count)); let state_ptr = ThreadPoolStatePtr::new(&state); let workers = (0..thread_count) - .map(|_| spawn_worker(receiver.clone(), state_ptr)) + .map(|worker_id| spawn_worker(receiver.clone(), state_ptr, worker_id)) .collect(); Self { state, workers } @@ -288,6 +350,26 @@ impl ThreadPool { self.state.thread_count() } + /// Return a cumulative snapshot of scheduler instrumentation. + /// + /// A snapshot includes the elapsed portion of stalls that are still active + /// at the instant it is taken. Consequently, subtracting a before snapshot + /// from an after snapshot with [`SchedulerMetrics::delta_since`] does not + /// attribute idle time before the first snapshot to the measured interval. + pub fn scheduler_metrics(&self) -> SchedulerMetrics { + self.state.instrumentation.snapshot() + } + + /// Reset cumulative scheduler instrumentation and queue high-water marks. + /// + /// Active worker stalls are split at the reset instant, so idle time before + /// the reset is not charged to later snapshots. For meaningful queue + /// high-water marks, call this while no callbacks are running or being + /// scheduled; this method does not pause the pool. + pub fn reset_scheduler_metrics(&self) { + self.state.instrumentation.reset() + } + /// Run `f` with this thread pool installed as the current pool. /// /// The callback runs on the caller thread. While it is running, @@ -396,10 +478,228 @@ impl Drop for ThreadPool { } } +#[derive(Default)] +struct StallTimerState { + completed: Duration, + started: Option, +} + +struct WorkerInstrumentation { + local_pushes: AtomicU64, + local_pops: AtomicU64, + donated_jobs: AtomicU64, + max_local_queue_depth: AtomicUsize, + stall_timer: Mutex, +} + +impl WorkerInstrumentation { + fn new() -> Self { + Self { + local_pushes: AtomicU64::new(0), + local_pops: AtomicU64::new(0), + donated_jobs: AtomicU64::new(0), + max_local_queue_depth: AtomicUsize::new(0), + stall_timer: Mutex::new(StallTimerState::default()), + } + } +} + +struct SchedulerInstrumentation { + global_pushes: AtomicU64, + global_pops: AtomicU64, + stalled_workers: AtomicUsize, + workers: Vec>, + global_queue_depth: AtomicUsize, + max_global_queue_depth: AtomicUsize, +} + +impl SchedulerInstrumentation { + fn new(thread_count: usize) -> Self { + Self { + global_pushes: AtomicU64::new(0), + global_pops: AtomicU64::new(0), + stalled_workers: AtomicUsize::new(0), + workers: (0..thread_count) + .map(|_| CachePadded::new(WorkerInstrumentation::new())) + .collect(), + global_queue_depth: AtomicUsize::new(0), + max_global_queue_depth: AtomicUsize::new(0), + } + } + + fn snapshot(&self) -> SchedulerMetrics { + let now = Instant::now(); + let stalled_time = self + .workers + .iter() + .map(|worker| { + let timer = worker + .stall_timer + .lock() + .unwrap_or_else(|error| error.into_inner()); + timer.completed.saturating_add( + timer + .started + .map_or(Duration::ZERO, |started| now.duration_since(started)), + ) + }) + .fold(Duration::ZERO, Duration::saturating_add); + let local_pushes = self + .workers + .iter() + .map(|worker| worker.local_pushes.load(Ordering::Relaxed)) + .sum(); + let local_pops = self + .workers + .iter() + .map(|worker| worker.local_pops.load(Ordering::Relaxed)) + .sum(); + let donated_jobs = self + .workers + .iter() + .map(|worker| worker.donated_jobs.load(Ordering::Relaxed)) + .sum(); + let max_local_queue_depth = self + .workers + .iter() + .map(|worker| worker.max_local_queue_depth.load(Ordering::Relaxed)) + .max() + .unwrap_or(0); + + SchedulerMetrics { + global_pushes: self.global_pushes.load(Ordering::Relaxed), + global_pops: self.global_pops.load(Ordering::Relaxed), + local_pushes, + local_pops, + donated_jobs, + stalled_time, + stalled_workers: self.stalled_workers.load(Ordering::Acquire), + max_local_queue_depth, + max_global_queue_depth: self.max_global_queue_depth.load(Ordering::Relaxed), + } + } + + fn reset(&self) { + let now = Instant::now(); + for worker in &self.workers { + let mut timer = worker + .stall_timer + .lock() + .unwrap_or_else(|error| error.into_inner()); + timer.completed = Duration::ZERO; + if timer.started.is_some() { + timer.started = Some(now); + } + drop(timer); + + worker.local_pushes.store(0, Ordering::Relaxed); + worker.local_pops.store(0, Ordering::Relaxed); + worker.donated_jobs.store(0, Ordering::Relaxed); + worker.max_local_queue_depth.store(0, Ordering::Relaxed); + } + + self.global_pushes.store(0, Ordering::Relaxed); + self.global_pops.store(0, Ordering::Relaxed); + self.max_global_queue_depth.store( + self.global_queue_depth.load(Ordering::Acquire), + Ordering::Relaxed, + ); + } + + fn record_global_push(&self) { + self.global_pushes.fetch_add(1, Ordering::Relaxed); + let depth = self.global_queue_depth.fetch_add(1, Ordering::AcqRel) + 1; + self.max_global_queue_depth + .fetch_max(depth, Ordering::Relaxed); + } + + fn cancel_global_push(&self) { + self.global_pushes.fetch_sub(1, Ordering::Relaxed); + self.global_queue_depth.fetch_sub(1, Ordering::AcqRel); + } + + fn record_global_pop(&self) { + self.global_pops.fetch_add(1, Ordering::Relaxed); + self.global_queue_depth.fetch_sub(1, Ordering::AcqRel); + } + + fn record_local_push(&self, worker_id: usize, depth: usize) { + let worker = &self.workers[worker_id]; + worker.local_pushes.fetch_add(1, Ordering::Relaxed); + worker + .max_local_queue_depth + .fetch_max(depth, Ordering::Relaxed); + } + + fn record_local_pop(&self, worker_id: usize) { + self.workers[worker_id] + .local_pops + .fetch_add(1, Ordering::Relaxed); + } + + fn record_donation(&self, worker_id: usize, jobs: usize) { + self.workers[worker_id] + .donated_jobs + .fetch_add(jobs as u64, Ordering::Relaxed); + } + + fn has_stalled_workers(&self) -> bool { + self.stalled_workers.load(Ordering::Acquire) != 0 + } + + fn begin_stall(&self, worker_id: usize) { + let mut timer = self.workers[worker_id] + .stall_timer + .lock() + .unwrap_or_else(|error| error.into_inner()); + debug_assert!(timer.started.is_none()); + timer.started = Some(Instant::now()); + drop(timer); + self.stalled_workers.fetch_add(1, Ordering::AcqRel); + } + + fn end_stall(&self, worker_id: usize) { + let now = Instant::now(); + let mut timer = self.workers[worker_id] + .stall_timer + .lock() + .unwrap_or_else(|error| error.into_inner()); + let started = timer + .started + .take() + .expect("primary worker ended a stall it did not begin"); + timer.completed += now.duration_since(started); + drop(timer); + self.stalled_workers.fetch_sub(1, Ordering::AcqRel); + } +} + +struct StalledWorkerGuard<'pool> { + instrumentation: &'pool SchedulerInstrumentation, + worker_id: usize, +} + +impl<'pool> StalledWorkerGuard<'pool> { + fn new(instrumentation: &'pool SchedulerInstrumentation, worker_id: usize) -> Self { + instrumentation.begin_stall(worker_id); + Self { + instrumentation, + worker_id, + } + } +} + +impl Drop for StalledWorkerGuard<'_> { + fn drop(&mut self) { + self.instrumentation.end_stall(self.worker_id); + } +} + struct ThreadPoolState { sender: Option>, receiver: Receiver, thread_count: usize, + instrumentation: SchedulerInstrumentation, #[cfg(test)] backup_workers_spawned: AtomicUsize, #[cfg(test)] @@ -412,6 +712,7 @@ impl ThreadPoolState { sender: Some(sender), receiver, thread_count, + instrumentation: SchedulerInstrumentation::new(thread_count), #[cfg(test)] backup_workers_spawned: AtomicUsize::new(0), #[cfg(test)] @@ -423,6 +724,38 @@ impl ThreadPoolState { self.thread_count } + fn enqueue_global(&self, job: Job) { + let Some(sender) = self.sender.as_ref() else { + job(); + return; + }; + + self.instrumentation.record_global_push(); + if let Err(error) = sender.send(job) { + self.instrumentation.cancel_global_push(); + error.0(); + } + } + + fn try_recv_global(&self) -> Result { + let result = self.receiver.try_recv(); + if result.is_ok() { + self.instrumentation.record_global_pop(); + } + result + } + + fn recv_global(&self, worker_id: usize) -> Result { + let result = { + let _stalled = StalledWorkerGuard::new(&self.instrumentation, worker_id); + self.receiver.recv() + }; + if result.is_ok() { + self.instrumentation.record_global_pop(); + } + result + } + fn scope<'scope, F, R>(&self, f: F) -> R where F: FnOnce(&Scope<'scope>) -> R, @@ -450,12 +783,15 @@ impl ThreadPoolState { } let Some(_guard) = InlineScopeHelpGuard::try_enter() else { + let worker_id = current_worker_id(self); + donate_all_from_current_worker(self); let _backup = BackupWorker::spawn(self); + let _stalled = worker_id + .map(|worker_id| StalledWorkerGuard::new(&self.instrumentation, worker_id)); receive_scope_completion(done); return; }; - let receiver = self.receiver.clone(); loop { match done.try_recv() { Ok(()) => break, @@ -465,7 +801,12 @@ impl ThreadPoolState { } } - match receiver.try_recv() { + if let Some(job) = pop_current_worker_local(self) { + job(); + continue; + } + + match self.try_recv_global() { Ok(job) => { job(); continue; @@ -477,13 +818,31 @@ impl ThreadPoolState { } } - select_biased! { - recv(done) -> message => { + enum WaitEvent { + Scope(Result<(), crossbeam::channel::RecvError>), + Global(Result), + } + + let worker_id = current_worker_id(self); + let event = { + let _stalled = worker_id + .map(|worker_id| StalledWorkerGuard::new(&self.instrumentation, worker_id)); + select_biased! { + recv(done) -> message => WaitEvent::Scope(message), + recv(self.receiver) -> message => WaitEvent::Global(message), + } + }; + + match event { + WaitEvent::Scope(message) => { expect_scope_completion(message); break; } - recv(receiver) -> message => match message { - Ok(job) => job(), + WaitEvent::Global(message) => match message { + Ok(job) => { + self.instrumentation.record_global_pop(); + job(); + } Err(_) => { receive_scope_completion(done); break; @@ -530,6 +889,161 @@ unsafe impl Send for ThreadPoolStatePtr {} // boxed state remains alive until all worker threads have been joined. unsafe impl Sync for ThreadPoolStatePtr {} +/// State owned and accessed exclusively by one primary worker thread. +/// +/// The deque uses interior mutability so no Rust borrow of the queue is held +/// while a job is running. Jobs may recursively call `spawn_local`, which then +/// reaches the same deque through `CURRENT_WORKER` on that thread. +struct WorkerContext { + pool: ThreadPoolStatePtr, + worker_id: usize, + local: UnsafeCell>, +} + +impl WorkerContext { + fn new(pool: ThreadPoolStatePtr, worker_id: usize) -> Self { + Self { + pool, + worker_id, + local: UnsafeCell::new(VecDeque::new()), + } + } + + fn belongs_to(&self, pool: &ThreadPoolState) -> bool { + ptr::eq(self.pool.0, pool) + } + + fn push_local(&self, pool: &ThreadPoolState, job: Job) { + // SAFETY: a `WorkerContext` is installed only on its owning worker and + // `CURRENT_WORKER` is never made available to another thread. + let depth = unsafe { + let queue = &mut *self.local.get(); + queue.push_back(job); + queue.len() + }; + pool.instrumentation + .record_local_push(self.worker_id, depth); + self.donate_half_if_stalled(pool); + } + + fn pop_local(&self, pool: &ThreadPoolState) -> Option { + self.donate_half_if_stalled(pool); + // The owner takes the newest work, retaining depth-first locality. + // SAFETY: see `push_local`. + let job = unsafe { (&mut *self.local.get()).pop_back() }; + if job.is_some() { + pool.instrumentation.record_local_pop(self.worker_id); + } + job + } + + fn donate_half_if_stalled(&self, pool: &ThreadPoolState) { + if !pool.instrumentation.has_stalled_workers() { + return; + } + + // Donation takes the opposite end from the owner: the oldest siblings + // have the weakest connection to the owner's current depth-first path. + // SAFETY: see `push_local`. + let donated = unsafe { + let queue = &mut *self.local.get(); + let count = queue.len() / 2; + queue.drain(..count).collect::>() + }; + self.publish_donation(pool, donated); + } + + fn donate_all(&self, pool: &ThreadPoolState) { + // The inline-help depth escape hatch blocks this worker and starts a + // backup consumer. Publish everything so private work remains live. + // SAFETY: see `push_local`. + let donated = unsafe { (&mut *self.local.get()).drain(..).collect::>() }; + self.publish_donation(pool, donated); + } + + fn publish_donation(&self, pool: &ThreadPoolState, donated: Vec) { + if donated.is_empty() { + return; + } + + pool.instrumentation + .record_donation(self.worker_id, donated.len()); + for job in donated { + pool.enqueue_global(job); + } + } +} + +fn install_worker_context(worker: &WorkerContext, f: F) -> R +where + F: FnOnce() -> R, +{ + CURRENT_WORKER.with(|current| { + let previous = current.replace(worker); + let _guard = CurrentWorkerGuard { current, previous }; + f() + }) +} + +struct CurrentWorkerGuard<'a> { + current: &'a Cell<*const WorkerContext>, + previous: *const WorkerContext, +} + +impl Drop for CurrentWorkerGuard<'_> { + fn drop(&mut self) { + self.current.set(self.previous); + } +} + +fn with_current_worker(f: F) -> R +where + F: FnOnce(Option<&WorkerContext>) -> R, +{ + CURRENT_WORKER.with(|current| { + let worker = current.get(); + // SAFETY: the pointer is installed from a live `WorkerContext` on this + // same thread and is restored before that context is dropped. + f(unsafe { worker.as_ref() }) + }) +} + +fn current_worker_id(pool: &ThreadPoolState) -> Option { + with_current_worker(|worker| { + worker + .filter(|worker| worker.belongs_to(pool)) + .map(|worker| worker.worker_id) + }) +} + +fn push_current_worker_local(pool: &ThreadPoolState, job: Job) -> Result<(), Job> { + with_current_worker( + |worker| match worker.filter(|worker| worker.belongs_to(pool)) { + Some(worker) => { + worker.push_local(pool, job); + Ok(()) + } + None => Err(job), + }, + ) +} + +fn pop_current_worker_local(pool: &ThreadPoolState) -> Option { + with_current_worker(|worker| { + worker + .filter(|worker| worker.belongs_to(pool)) + .and_then(|worker| worker.pop_local(pool)) + }) +} + +fn donate_all_from_current_worker(pool: &ThreadPoolState) { + with_current_worker(|worker| { + if let Some(worker) = worker.filter(|worker| worker.belongs_to(pool)) { + worker.donate_all(pool); + } + }); +} + fn install_pool(pool: &ThreadPoolState, f: F) -> R where F: FnOnce() -> R, @@ -648,7 +1162,6 @@ where /// assert_eq!(counter.load(Ordering::Relaxed), 1); /// ``` pub struct Scope<'scope> { - sender: Sender, pool: ThreadPoolStatePtr, state: ScopeState, _scope: PhantomData &'scope ()>, @@ -657,18 +1170,18 @@ pub struct Scope<'scope> { impl<'scope> Scope<'scope> { fn new(pool: &ThreadPoolState) -> Self { Self { - sender: pool.sender.as_ref().unwrap().clone(), pool: ThreadPoolStatePtr::new(pool), state: ScopeState::new(), _scope: PhantomData, } } - /// Spawn a callback onto the pool. + /// Spawn a callback onto the pool's global queue. /// /// The callback receives the current scope and may add nested work. It may /// borrow values with the scope lifetime. The scope waits for the callback - /// before returning, so those borrows cannot outlive their owners. + /// before returning, so those borrows cannot outlive their owners. This is + /// retained as a compatibility alias for [`Scope::spawn_global`]. /// /// # Examples /// @@ -691,6 +1204,46 @@ impl<'scope> Scope<'scope> { /// assert_eq!(counter.load(Ordering::Relaxed), 2); /// ``` pub fn spawn(&self, f: F) + where + F: FnOnce(&Scope<'scope>) + Send + 'scope, + { + self.spawn_global(f); + } + + /// Spawn a callback onto the pool's global queue. + /// + /// Global work is visible to every worker immediately and is appropriate + /// for coarse top-level partitioning. + pub fn spawn_global(&self, f: F) + where + F: FnOnce(&Scope<'scope>) + Send + 'scope, + { + let job = self.prepare_job(f); + // SAFETY: this scope was created from a live pool, and the pool cannot + // be dropped while its borrowed `scope` call is active. + unsafe { self.pool.as_ref() }.enqueue_global(job); + } + + /// Spawn a callback onto the current worker's private deque. + /// + /// The owning worker executes local jobs depth-first before consulting the + /// global queue. When another worker is stalled, half of the oldest local + /// jobs are donated to the global queue. Calls made outside a primary + /// worker of this same pool—including calls from the root callback and + /// from a worker of a nested, different pool—fall back to the global queue. + pub fn spawn_local(&self, f: F) + where + F: FnOnce(&Scope<'scope>) + Send + 'scope, + { + let job = self.prepare_job(f); + // SAFETY: see `spawn_global`. + let pool = unsafe { self.pool.as_ref() }; + if let Err(job) = push_current_worker_local(pool, job) { + pool.enqueue_global(job); + } + } + + fn prepare_job(&self, f: F) -> Job where F: FnOnce(&Scope<'scope>) + Send + 'scope, { @@ -717,7 +1270,7 @@ impl<'scope> Scope<'scope> { // SAFETY: every erased job records completion in the scope state, and // `Scope::complete_root_and_wait` waits for all expected completions // before `ThreadPool::scope` returns. - enqueue(&self.sender, unsafe { erase_job_lifetime(job) }); + unsafe { erase_job_lifetime(job) } } fn complete_root_and_wait(&self) { @@ -898,7 +1451,10 @@ impl BackupWorker { select_biased! { recv(shutdown_receiver) -> _ => break, recv(receiver) -> message => match message { - Ok(job) => job(), + Ok(job) => { + pool.instrumentation.record_global_pop(); + job(); + } Err(_) => break, }, } @@ -947,30 +1503,45 @@ impl Drop for BackupWorkerLiveGuard<'_> { } } -fn spawn_worker(receiver: Receiver, pool: ThreadPoolStatePtr) -> JoinHandle<()> { +fn spawn_worker( + receiver: Receiver, + pool: ThreadPoolStatePtr, + worker_id: usize, +) -> JoinHandle<()> { thread::spawn(move || { // SAFETY: each worker is joined before the boxed pool state is dropped. let pool = unsafe { pool.as_ref() }; + let worker = WorkerContext::new(ThreadPoolStatePtr::new(pool), worker_id); install_pool(pool, || { - install_background_worker(|| { - for job in receiver { - job(); - } + install_worker_context(&worker, || { + install_background_worker(|| { + loop { + if let Some(job) = worker.pop_local(pool) { + job(); + continue; + } + + match receiver.try_recv() { + Ok(job) => { + pool.instrumentation.record_global_pop(); + job(); + continue; + } + Err(crossbeam::channel::TryRecvError::Empty) => {} + Err(crossbeam::channel::TryRecvError::Disconnected) => break, + } + + let Ok(job) = pool.recv_global(worker_id) else { + break; + }; + job(); + } + }) }); }); }) } -fn enqueue(sender: &Sender, job: Job) { - match sender.send(job) { - Ok(()) => {} - Err(error) => { - let job = error.0; - job(); - } - } -} - fn completed(value: u64) -> u32 { (value & COMPLETED_MASK) as u32 } diff --git a/concurrency/src/threadpool/tests.rs b/concurrency/src/threadpool/tests.rs index 4b43edf8e..36186df5e 100644 --- a/concurrency/src/threadpool/tests.rs +++ b/concurrency/src/threadpool/tests.rs @@ -5,7 +5,7 @@ use std::{ atomic::{AtomicBool, AtomicUsize, Ordering}, }, thread, - time::Duration, + time::{Duration, Instant}, }; use crate::{ @@ -13,7 +13,10 @@ use crate::{ without_current_pool, }; -use super::{MAX_INLINE_SCOPE_HELP_DEPTH, is_background_worker_thread}; +use super::{ + MAX_INLINE_SCOPE_HELP_DEPTH, ThreadPoolState, ThreadPoolStatePtr, WorkerContext, + is_background_worker_thread, +}; #[test] fn scope_runs_spawned_work_and_waits() { @@ -31,6 +34,213 @@ fn scope_runs_spawned_work_and_waits() { assert_eq!(counter.load(Ordering::Relaxed), 128); } +#[test] +fn spawn_remains_a_global_queue_compatibility_alias() { + let pool = ThreadPool::new(1); + pool.reset_scheduler_metrics(); + + pool.scope(|scope| scope.spawn(|_| {})); + + let metrics = pool.scheduler_metrics(); + assert_eq!(metrics.global_pushes, 1); + assert_eq!(metrics.global_pops, 1); + assert_eq!(metrics.local_pushes, 0); + assert_eq!(metrics.local_pops, 0); +} + +#[test] +fn spawn_local_from_root_falls_back_to_global_queue() { + let pool = ThreadPool::new(1); + pool.reset_scheduler_metrics(); + + pool.scope(|scope| scope.spawn_local(|_| {})); + + let metrics = pool.scheduler_metrics(); + assert_eq!(metrics.global_pushes, 1); + assert_eq!(metrics.global_pops, 1); + assert_eq!(metrics.local_pushes, 0); +} + +#[test] +fn owner_runs_local_work_depth_first_before_global_work() { + let pool = ThreadPool::new(1); + let first_started = Notification::new(); + let release_first = Notification::new(); + let order = Mutex::new(Vec::new()); + + pool.scope(|scope| { + scope.spawn_global(|scope| { + order.lock().unwrap().push("parent"); + first_started.notify(); + release_first.wait(); + scope.spawn_local(|_| order.lock().unwrap().push("local-oldest")); + scope.spawn_local(|_| order.lock().unwrap().push("local-newest")); + }); + + first_started.wait(); + scope.spawn_global(|_| order.lock().unwrap().push("global")); + release_first.notify(); + }); + + assert_eq!( + order.into_inner().unwrap(), + ["parent", "local-newest", "local-oldest", "global"] + ); + let metrics = pool.scheduler_metrics(); + assert_eq!(metrics.local_pushes, 2); + assert_eq!(metrics.local_pops, 2); + assert_eq!(metrics.donated_jobs, 0); +} + +#[test] +fn nested_scope_help_takes_local_work_before_global_work() { + let pool = ThreadPool::new(1); + let parent_started = Notification::new(); + let release_parent = Notification::new(); + let order = Mutex::new(Vec::new()); + + pool.scope(|scope| { + scope.spawn_global(|_| { + order.lock().unwrap().push("parent-start"); + parent_started.notify(); + release_parent.wait(); + + threadpool_scope(|nested| { + nested.spawn_local(|_| order.lock().unwrap().push("nested-local")); + }); + order.lock().unwrap().push("parent-after-nested"); + }); + + parent_started.wait(); + scope.spawn_global(|_| order.lock().unwrap().push("global")); + release_parent.notify(); + }); + + assert_eq!( + order.into_inner().unwrap(), + [ + "parent-start", + "nested-local", + "parent-after-nested", + "global" + ] + ); +} + +#[test] +fn donation_moves_half_of_the_oldest_local_jobs() { + let (sender, receiver) = super::unbounded(); + let state = ThreadPoolState::new(sender, receiver, 1); + let worker = WorkerContext::new(ThreadPoolStatePtr::new(&state), 0); + let order = Arc::new(Mutex::new(Vec::new())); + + // Build the queue directly so one donation decision observes all six jobs. + // This isolates deque-end and half-transfer policy from worker timing. + unsafe { + let queue = &mut *worker.local.get(); + for value in 0..6 { + let order = order.clone(); + queue.push_back(Box::new(move || order.lock().unwrap().push(value))); + } + } + state + .instrumentation + .stalled_workers + .store(1, Ordering::Release); + + worker.donate_half_if_stalled(&state); + state + .instrumentation + .stalled_workers + .store(0, Ordering::Release); + + for _ in 0..3 { + state.try_recv_global().unwrap()(); + } + while let Some(job) = worker.pop_local(&state) { + job(); + } + + assert_eq!(*order.lock().unwrap(), [0, 1, 2, 5, 4, 3]); + let metrics = state.instrumentation.snapshot(); + assert_eq!(metrics.donated_jobs, 3); + assert_eq!(metrics.global_pushes, 3); + assert_eq!(metrics.global_pops, 3); + assert_eq!(metrics.local_pops, 3); +} + +#[test] +fn stalled_worker_signal_causes_public_local_work_to_be_donated() { + let pool = ThreadPool::new(2); + let deadline = Instant::now() + Duration::from_secs(1); + while pool.scheduler_metrics().stalled_workers != 2 { + assert!(Instant::now() < deadline, "workers did not become stalled"); + thread::yield_now(); + } + pool.reset_scheduler_metrics(); + + pool.scope(|scope| { + scope.spawn_global(|scope| { + for _ in 0..16 { + scope.spawn_local(|_| {}); + } + }); + }); + + let metrics = pool.scheduler_metrics(); + assert_eq!(metrics.local_pushes, 16); + assert!(metrics.donated_jobs > 0); + assert_eq!( + metrics.local_pushes, + metrics.local_pops + metrics.donated_jobs + ); + assert_eq!(metrics.global_pushes, 1 + metrics.donated_jobs); + assert_eq!(metrics.global_pushes, metrics.global_pops); + assert!(metrics.max_local_queue_depth >= 2); +} + +#[test] +fn stalled_time_snapshots_include_stalls_in_progress() { + let pool = ThreadPool::new(2); + let deadline = Instant::now() + Duration::from_secs(1); + while pool.scheduler_metrics().stalled_workers != 2 { + assert!(Instant::now() < deadline, "workers did not become stalled"); + thread::yield_now(); + } + + let before = pool.scheduler_metrics(); + thread::sleep(Duration::from_millis(2)); + let delta = pool.scheduler_metrics().delta_since(before); + assert!(delta.stalled_time >= Duration::from_millis(2)); + assert_eq!(delta.stalled_workers, 2); + + pool.reset_scheduler_metrics(); + let after_reset = pool.scheduler_metrics(); + thread::sleep(Duration::from_millis(2)); + let reset_delta = pool.scheduler_metrics().delta_since(after_reset); + assert!(reset_delta.stalled_time >= Duration::from_millis(2)); +} + +#[test] +fn panic_from_local_work_is_propagated_after_the_scope_drains() { + let pool = ThreadPool::new(1); + let sibling_completed = AtomicBool::new(false); + + let result = panic::catch_unwind(panic::AssertUnwindSafe(|| { + pool.scope(|scope| { + scope.spawn_global(|scope| { + scope.spawn_local(|_| sibling_completed.store(true, Ordering::Release)); + // Local jobs are LIFO, so the panic runs first. Its job wrapper + // must record the panic and let the worker drain the sibling. + scope.spawn_local(|_| panic!("local worker panic")); + }); + }); + })); + + assert!(result.is_err()); + assert!(sibling_completed.load(Ordering::Acquire)); +} + #[test] fn scope_allows_borrowing_stack_data() { let pool = ThreadPool::new(2); @@ -377,6 +587,28 @@ fn install_sets_and_restores_current_pool() { assert_eq!(current_num_threads(), 1); } +#[test] +fn spawn_local_does_not_use_a_different_pools_worker_queue() { + let outer = ThreadPool::new(1); + let inner = ThreadPool::new(1); + let completed = AtomicBool::new(false); + inner.reset_scheduler_metrics(); + + outer.scope(|scope| { + scope.spawn_global(|_| { + inner.scope(|scope| { + scope.spawn_local(|_| completed.store(true, Ordering::Release)); + }); + }); + }); + + assert!(completed.load(Ordering::Acquire)); + let metrics = inner.scheduler_metrics(); + assert_eq!(metrics.local_pushes, 0); + assert_eq!(metrics.global_pushes, 1); + assert_eq!(metrics.global_pops, 1); +} + #[test] fn without_current_pool_clears_and_restores_current_pool() { let pool = ThreadPool::new(2); @@ -601,3 +833,31 @@ fn deeply_nested_background_waits_fall_back_to_backup_worker() { assert!(pool.state.backup_workers_spawned() >= 1); assert_eq!(pool.state.backup_workers_live(), 0); } + +#[test] +fn deeply_nested_local_work_is_donated_to_backup_worker() { + fn recurse(depth: usize, completed: &AtomicUsize) { + if depth == 0 { + completed.fetch_add(1, Ordering::Release); + return; + } + + threadpool_scope(|scope| { + scope.spawn_local(move |_| recurse(depth - 1, completed)); + }); + } + + let pool = ThreadPool::new(1); + let completed = AtomicUsize::new(0); + + pool.scope(|scope| { + scope.spawn_global(|_| { + recurse(MAX_INLINE_SCOPE_HELP_DEPTH + 8, &completed); + }); + }); + + assert_eq!(completed.load(Ordering::Acquire), 1); + assert!(pool.state.backup_workers_spawned() >= 1); + assert_eq!(pool.state.backup_workers_live(), 0); + assert!(pool.scheduler_metrics().donated_jobs >= 1); +} diff --git a/core-relations/src/free_join/execute.rs b/core-relations/src/free_join/execute.rs index 39d4b0f2b..75d1afa51 100644 --- a/core-relations/src/free_join/execute.rs +++ b/core-relations/src/free_join/execute.rs @@ -27,14 +27,18 @@ use web_time::Instant; use crate::{ Constraint, OffsetRange, Pool, SubsetRef, action::{Bindings, ExecutionState}, - common::{DashMap, Value}, + common::{DashMap, ShardData, ShardId, Value}, free_join::{ frame_update::{FrameUpdates, UpdateInstr}, get_index_from_tableinfo, }, - hash_index::{IndexBase, TupleIndex}, + hash_index::{ColumnIndex, Index, IndexBase, TupleIndex}, offsets::{Offsets, RowId, SortedOffsetSlice, SortedOffsetVector, Subset}, - parallel_heuristics::{action_batch_size, free_join_fork_depth, parallelize_db_level_op}, + parallel_heuristics::{ + GjChildCache, GjTopPromotion, action_batch_size, free_join_fork_depth, gj_child_cache, + gj_local_depth, gj_local_morsel, gj_local_queue_limit, gj_min_keys_per_worker, + gj_top_index_sharding, gj_top_promotion, parallelize_db_level_op, + }, pool::Pooled, query::RuleSet, row_buffer::TaggedRowBuffer, @@ -50,6 +54,52 @@ use super::{ const SMALL_RESIDUAL: usize = 8; +fn top_index_shape_is_eligible( + workers: usize, + leader_keys: usize, + nonempty_shards: usize, + min_keys_per_worker: usize, +) -> bool { + workers > 1 + && leader_keys >= min_keys_per_worker.saturating_mul(workers) + && nonempty_shards >= workers +} + +/// Select one already-evaluated eligible candidate. The iterator remains lazy, +/// so `Current` evaluates only position zero and `Eligible` stops at its first +/// usable fallback; only `Largest` evaluates the full +/// leading prefix. +fn select_top_candidate( + policy: GjTopPromotion, + current_estimate: usize, + min_threshold: usize, + mut candidates: impl Iterator>, + leader_keys: impl Fn(&T) -> usize, +) -> Option { + let current = candidates.next().flatten(); + match policy { + GjTopPromotion::Current => current, + GjTopPromotion::Eligible => { + if current.is_some() { + return current; + } + let promotion_limit = current_estimate.saturating_mul(4).max(min_threshold); + candidates + .flatten() + .find(|candidate| leader_keys(candidate) <= promotion_limit) + } + GjTopPromotion::Largest => { + current + .into_iter() + .chain(candidates.flatten()) + .fold(None, |best, candidate| match best { + Some(best) if leader_keys(&best) >= leader_keys(&candidate) => Some(best), + _ => Some(candidate), + }) + } + } +} + struct SparseColumnIndex { n_keys: usize, n_subsets: usize, @@ -244,18 +294,123 @@ impl SortedColumnIndex { } } -enum DynamicIndex { +/// A table-index slot retained for one `run_rule_set` call. +/// +/// The slot lazily acquires its Arc through the existing fully-refreshing +/// catalog helper on first cached use. Keeping that Arc in an execution-scoped +/// sidecar removes catalog lookups and refcount traffic from recursive join +/// execution without constructing indexes for plan accesses that choose a +/// residual-local strategy at runtime. Initialized slots are dropped before +/// the database resets its indexes during `merge_all`. +enum PreparedIndexSlot { + Tuple(OnceLock), + Column(OnceLock), + /// The table specification forbids a global cache for at least one key + /// column, so execution must use its existing dynamic-index path. + Uncacheable, +} + +fn columns_are_cacheable(info: &TableInfo, cols: &[ColumnId]) -> bool { + cols.iter().all(|col| { + !info + .spec + .uncacheable_columns + .get(*col) + .copied() + .unwrap_or(false) + }) +} + +/// Index handles for one immutable [`JoinStages`] value, positionally aligned +/// with `JoinStages::instrs` and with each stage's scans. +struct PreparedJoinIndexes { + stages: Box<[SmallVec<[PreparedIndexSlot; 4]>]>, +} + +impl PreparedJoinIndexes { + fn new(db: &Database, atoms: &Arc>, stages: &JoinStages) -> Self { + let stages = stages + .instrs + .iter() + .map(|stage| { + let mut handles = SmallVec::new(); + match stage { + JoinStage::Intersect { scans, .. } => { + handles.extend(scans.iter().map(|scan| { + let info = &db.tables[atoms[scan.atom].table]; + if !columns_are_cacheable(info, &[scan.column]) { + PreparedIndexSlot::Uncacheable + } else { + PreparedIndexSlot::Column(OnceLock::new()) + } + })); + } + JoinStage::FusedIntersect { to_intersect, .. } + | JoinStage::FusedIntersectMat { to_intersect, .. } => { + handles.extend(to_intersect.iter().map(|(scan, _)| { + let cols = scan.to_index.vars.as_slice(); + let info = &db.tables[atoms[scan.to_index.atom].table]; + if !columns_are_cacheable(info, cols) { + PreparedIndexSlot::Uncacheable + } else if cols.len() == 1 { + PreparedIndexSlot::Column(OnceLock::new()) + } else { + PreparedIndexSlot::Tuple(OnceLock::new()) + } + })); + } + } + handles + }) + .collect(); + Self { stages } + } + + fn stage(&self, index: usize) -> &[PreparedIndexSlot] { + &self.stages[index] + } +} + +/// Execution-scoped index sidecar mirroring the shape of a logical [`Plan`]. +enum PreparedPlanIndexes { + Single(PreparedJoinIndexes), + Decomposed { + blocks: Vec, + result: PreparedJoinIndexes, + }, +} + +impl PreparedPlanIndexes { + fn new(db: &Database, plan: &Plan) -> Self { + match plan { + Plan::SinglePlan(plan) => { + Self::Single(PreparedJoinIndexes::new(db, &plan.atoms, &plan.stages)) + } + Plan::DecomposedPlan(plan) => Self::Decomposed { + blocks: plan + .stages + .blocks + .iter() + .map(|(stages, _)| PreparedJoinIndexes::new(db, &plan.atoms, stages)) + .collect(), + result: PreparedJoinIndexes::new(db, &plan.atoms, &plan.result_block), + }, + } + } +} + +enum DynamicIndex<'plan> { Cached { /// When Some(range), intersect each subset from the index with this dense range. /// The range is the Dense outer subset known at Prober construction time. intersect_outer: Option, - table: HashIndex, + table: &'plan Index, }, CachedColumn { /// When Some(range), intersect each subset from the index with this dense range. /// The range is the Dense outer subset known at Prober construction time. intersect_outer: Option, - table: HashColumnIndex, + table: &'plan Index, }, Dynamic(TupleIndex), DynamicColumn(Arc), @@ -323,19 +478,19 @@ fn intersect_with_dense_ref<'a>(v: SubsetRef<'a>, range: OffsetRange) -> Option< } } -struct Prober { +struct Prober<'plan> { node: Arc, - ix: DynamicIndex, + ix: DynamicIndex<'plan>, } -impl Prober { +impl Prober<'_> { fn get_subset<'a>(&'a self, key: &'a [Value]) -> Option>> { match &self.ix { DynamicIndex::Cached { intersect_outer, table, } => { - let subset_ref = table.get().unwrap().get_subset(key)?; + let subset_ref = table.get_subset(key)?; let subset = if let Some(range) = intersect_outer { intersect_with_dense_ref(subset_ref, *range)? } else { @@ -348,7 +503,7 @@ impl Prober { table, } => { debug_assert_eq!(key.len(), 1); - let subset_ref = table.get().unwrap().get_subset(&key[0])?; + let subset_ref = table.get_subset(&key[0])?; let subset = if let Some(range) = intersect_outer { intersect_with_dense_ref(subset_ref, *range)? } else { @@ -373,7 +528,7 @@ impl Prober { table, } => { let range = *range; - table.get().unwrap().for_each(|k, v| { + table.for_each(|k, v| { if let Some(res) = intersect_with_dense_ref(v, range) { f(k, PotentiallyStale::maybe_stale(res)) } @@ -382,16 +537,13 @@ impl Prober { DynamicIndex::Cached { intersect_outer: None, table, - } => table - .get() - .unwrap() - .for_each(|k, v| f(k, PotentiallyStale::maybe_stale(v))), + } => table.for_each(|k, v| f(k, PotentiallyStale::maybe_stale(v))), DynamicIndex::CachedColumn { intersect_outer: Some(range), table, } => { let range = *range; - table.get().unwrap().for_each(|k, v| { + table.for_each(|k, v| { if let Some(res) = intersect_with_dense_ref(v, range) { f(&[*k], PotentiallyStale::maybe_stale(res)) } @@ -401,10 +553,7 @@ impl Prober { intersect_outer: None, table, } => { - table - .get() - .unwrap() - .for_each(|k, v| f(&[*k], PotentiallyStale::maybe_stale(v))); + table.for_each(|k, v| f(&[*k], PotentiallyStale::maybe_stale(v))); } DynamicIndex::Dynamic(tab) => { tab.for_each(|k, v| f(k, PotentiallyStale::not_stale(v))); @@ -418,10 +567,98 @@ impl Prober { } } + /// Enumerate one physical shard of a cached table index. + /// + /// Dynamic indexes are deliberately excluded: their storage is private to + /// this prober and has no coarse, pre-existing partition worth scheduling. + fn for_each_shard( + &self, + shard: usize, + mut f: impl FnMut(&[Value], PotentiallyStale), + ) { + match &self.ix { + DynamicIndex::Cached { + intersect_outer, + table, + } => { + let intersect_outer = *intersect_outer; + table.for_each_shard(shard, |key, subset| { + let subset = if let Some(range) = intersect_outer { + let Some(subset) = intersect_with_dense_ref(subset, range) else { + return; + }; + subset + } else { + subset + }; + f(key, PotentiallyStale::maybe_stale(subset)); + }); + } + DynamicIndex::CachedColumn { + intersect_outer, + table, + } => { + let intersect_outer = *intersect_outer; + table.for_each_shard(shard, |key, subset| { + let subset = if let Some(range) = intersect_outer { + let Some(subset) = intersect_with_dense_ref(subset, range) else { + return; + }; + subset + } else { + subset + }; + f(&[*key], PotentiallyStale::maybe_stale(subset)); + }); + } + DynamicIndex::Dynamic(_) + | DynamicIndex::DynamicColumn(_) + | DynamicIndex::SparseColumn(_) => { + unreachable!("only cached indexes expose physical shards") + } + } + } + + /// Return the number of coarse physical partitions when this prober is + /// backed by a cached hash index. + fn shard_count(&self) -> Option { + match &self.ix { + DynamicIndex::Cached { table, .. } => Some(table.shard_count()), + DynamicIndex::CachedColumn { table, .. } => Some(table.shard_count()), + DynamicIndex::Dynamic(_) + | DynamicIndex::DynamicColumn(_) + | DynamicIndex::SparseColumn(_) => None, + } + } + + fn shard_len(&self, shard: usize) -> Option { + match &self.ix { + DynamicIndex::Cached { + intersect_outer: None, + table, + } => Some(table.shard_len(shard)), + DynamicIndex::CachedColumn { + intersect_outer: None, + table, + } => Some(table.shard_len(shard)), + DynamicIndex::Cached { + intersect_outer: Some(_), + .. + } + | DynamicIndex::CachedColumn { + intersect_outer: Some(_), + .. + } + | DynamicIndex::Dynamic(_) + | DynamicIndex::DynamicColumn(_) + | DynamicIndex::SparseColumn(_) => None, + } + } + fn len(&self) -> usize { match &self.ix { - DynamicIndex::Cached { table, .. } => table.get().unwrap().len(), - DynamicIndex::CachedColumn { table, .. } => table.get().unwrap().len(), + DynamicIndex::Cached { table, .. } => table.len(), + DynamicIndex::CachedColumn { table, .. } => table.len(), DynamicIndex::Dynamic(tab) => tab.len(), DynamicIndex::DynamicColumn(tab) => tab.len(), DynamicIndex::SparseColumn(tab) => tab.len(), @@ -459,6 +696,15 @@ impl Database { }; let search_and_apply_timer = Instant::now(); + // Prepare one lazy slot for every table-index access named by the + // immutable plans before any worker starts. A slot acquires and refreshes + // its Arc through the regular catalog helper on first cached use. The + // sidecars are dropped before `merge_all` resets catalog entries. + let prepared_plans = rule_set + .plans + .values() + .map(|(plan, _, _)| PreparedPlanIndexes::new(self, plan)) + .collect::>(); // let mut rule_reports: HashMap>; let mut rule_reports: HashMap, Vec>; let exec_state = ExecutionState::new(self.read_only_view(), Default::default()); @@ -467,7 +713,9 @@ impl Database { Arc::new(DashMap::default()); let db: &Database = self; egglog_concurrency::scope(|scope| { - for (plan, desc, symbol_map) in rule_set.plans.values() { + for ((plan, desc, symbol_map), prepared_plan) in + rule_set.plans.values().zip(&prepared_plans) + { // TODO: add stats let report_plan = match report_level { ReportLevel::TimeOnly => None, @@ -498,17 +746,24 @@ impl Database { } } - match plan { - Plan::SinglePlan(plan) => { + match (plan, prepared_plan) { + (Plan::SinglePlan(plan), PreparedPlanIndexes::Single(prepared)) => { join_state.run_join_stages( &plan.stages, + prepared, &plan.atoms, plan.actions, &mut binding_info, &mut action_buf, ); } - Plan::DecomposedPlan(plan) => { + ( + Plan::DecomposedPlan(plan), + PreparedPlanIndexes::Decomposed { + blocks: prepared_blocks, + result: prepared_result, + }, + ) => { let mut materializations: DenseIdMap< MatId, Arc, RowBuffer>>, @@ -531,8 +786,8 @@ impl Database { ); let mut materializations = Arc::new(materializations); - for (mat_id, stage_block) in - plan.stages.blocks.iter().enumerate() + for (mat_id, (stage_block, prepared_block)) in + plan.stages.blocks.iter().zip(prepared_blocks).enumerate() { let mat_id = MatId::from_usize(mat_id); egglog_concurrency::scope(|stage_scope| { @@ -545,6 +800,7 @@ impl Database { }; join_state.run_join_stages( &stage_block.0, + prepared_block, &plan.atoms, mat_id, &mut binding_info, @@ -570,12 +826,14 @@ impl Database { } join_state.run_join_stages( &plan.result_block, + prepared_result, &plan.atoms, plan.actions, &mut binding_info, &mut action_buf, ); } + _ => unreachable!("prepared plan shape must match logical plan"), } } let search_and_apply_time = search_and_apply_timer.elapsed(); @@ -606,7 +864,9 @@ impl Database { match_counter: match_counter.as_ref(), batches: Default::default(), }; - for (plan, desc, symbol_map) in rule_set.plans.values() { + for ((plan, desc, symbol_map), prepared_plan) in + rule_set.plans.values().zip(&prepared_plans) + { let report_plan = match report_level { ReportLevel::TimeOnly => None, ReportLevel::WithPlan | ReportLevel::StageInfo => { @@ -625,17 +885,24 @@ impl Database { None => break 'eval, } } - match plan { - Plan::SinglePlan(plan) => { + match (plan, prepared_plan) { + (Plan::SinglePlan(plan), PreparedPlanIndexes::Single(prepared)) => { join_state.run_join_stages( &plan.stages, + prepared, &plan.atoms, plan.actions, &mut binding_info, &mut action_buf, ); } - Plan::DecomposedPlan(plan) => { + ( + Plan::DecomposedPlan(plan), + PreparedPlanIndexes::Decomposed { + blocks: prepared_blocks, + result: prepared_result, + }, + ) => { let mut materializations = DenseIdMap::with_capacity(plan.stages.blocks.len()); for i in 0..plan.stages.blocks.len() { @@ -654,10 +921,13 @@ impl Database { scratch_val: Default::default(), }; - for (mat_id, stage_block) in plan.stages.blocks.iter().enumerate() { + for (mat_id, (stage_block, prepared_block)) in + plan.stages.blocks.iter().zip(prepared_blocks).enumerate() + { let mat_id = MatId::from_usize(mat_id); join_state.run_join_stages( &stage_block.0, + prepared_block, &plan.atoms, mat_id, &mut binding_info, @@ -673,12 +943,14 @@ impl Database { } join_state.run_join_stages( &plan.result_block, + prepared_result, &plan.atoms, plan.actions, &mut binding_info, &mut action_buf, ); } + _ => unreachable!("prepared plan shape must match logical plan"), } } let search_and_apply_time = search_and_apply_timer.elapsed(); @@ -708,6 +980,10 @@ impl Database { // caused by individual queries. reports[i].num_matches = match_counter.read_matches(plan.actions()); } + // `merge_all` requires each catalog Arc to be uniquely owned so it can + // reset the corresponding ResettableOnceLock. No scoped worker can + // survive to this point. + drop(prepared_plans); let search_and_apply_time = search_and_apply_timer.elapsed(); let merge_timer = Instant::now(); @@ -752,15 +1028,95 @@ struct JoinState<'a> { /// Per-column indexes on a trie node's subset, lazily initialized on first access per column. type ColumnIndexes = IdVec>>; -// Each TrieNode is probed with exactly one column in practice, so we store a single -// (ColumnId, map) pair instead of a per-column IdVec of Mutexes. -// // The child cache (see [`TrieNode::get_cached_trie_node`]): keyed by the bound // value, storing the child node and the edge constraints used to build it. The // stored constraints guard against distinct scans reaching the same // (node, col, value) with different slow constraints; they are almost always // empty, in which case the guard is a cheap length check. -type ChildrenMaps = IdVec, Box<[Constraint]>)>>>; +type ChildEntry = (Arc, Box<[Constraint]>); +type ChildLock = RwLock>; +type ChildrenMaps = IdVec>; + +/// Cache of child trie nodes for one `(TrieNode, ColumnId)` pair. +/// +/// Descendant nodes lazily initialize [`ChildMap::Single`]. An eligible coarse +/// top-index partition may initialize a root column as [`ChildMap::Sharded`] +/// before its jobs start, using the same [`ShardData`] layout as the leader's +/// column index. There is deliberately no promotion from `Single` to `Sharded`: +/// an already-used cache remains valid and simply falls back to the old layout. +pub(crate) enum ChildMap { + Single(ChildLock), + Sharded { + shard_data: ShardData, + shards: IdVec>, + }, +} + +impl ChildMap { + fn single() -> Self { + Self::Single(RwLock::new(HashMap::default())) + } + + fn sharded(shard_count: usize) -> Self { + let shard_data = ShardData::new(shard_count); + debug_assert_eq!(shard_data.n_shards(), shard_count); + let mut shards = IdVec::with_capacity(shard_data.n_shards()); + shards.resize_with(shard_data.n_shards(), || { + CachePadded::new(RwLock::new(HashMap::default())) + }); + Self::Sharded { shard_data, shards } + } + + fn is_aligned_to(&self, shard_count: usize) -> bool { + matches!( + self, + Self::Sharded { shard_data, .. } if shard_data.n_shards() == shard_count + ) + } + + /// Select the lock with the same `Value` hashing and high-bit shard mapping + /// used by [`crate::hash_index::ColumnIndex`]. + fn lock_for(&self, value: Value) -> &ChildLock { + match self { + Self::Single(lock) => lock, + Self::Sharded { shard_data, shards } => { + shard_data.get_shard(&value, shards) as &ChildLock + } + } + } + + fn get_or_insert( + &self, + value: Value, + edge_cs: &[Constraint], + sub: impl FnOnce() -> Subset, + ) -> Arc { + let map = self.lock_for(value); + // Optimistic read path: most calls are cache hits, so try a shared lock + // first. A hit is only valid when the edge constraints match. + { + let guard = map.read().unwrap(); + if let Some((node, stored_cs)) = guard.get(&value) + && &**stored_cs == edge_cs + { + return node.clone(); + } + } + // Cache miss (or constraint mismatch): acquire the write lock and insert. + let mut guard = map.write().unwrap(); + if let Some((node, stored_cs)) = guard.get(&value) + && &**stored_cs == edge_cs + { + return node.clone(); + } + // Child nodes intentionally start with no cache layout. If they are + // probed later, they lazily choose `Single` and never inherit root + // sharding from this map. + let new_node = Arc::new(TrieNode::new(sub())); + guard.insert(value, (new_node.clone(), Box::from(edge_cs))); + new_node + } +} /// Canonical signature of a trie root: the table plus its sorted header (fast) /// constraints. Distinct signatures get distinct base ids from [`TrieCache`]. @@ -880,11 +1236,12 @@ pub(crate) struct TrieNode { subset: Subset, /// Any cached indexes on this subset. cached_subsets: OnceLock>, - /// Cached child trie nodes, keyed by value. In practice each TrieNode is - /// only ever probed with a single column, so we store one (col, map) pair - /// instead of an IdVec across all columns. When this node is a shared root - /// (or reachable from one), this cache is shared across plans too, so - /// children are shared without any global lookup. + /// Cached child trie nodes, keyed first by column and then by value. Each + /// column lazily chooses a single lock, except an eligible top-level root + /// may preinitialize a physically aligned sharded map before partition jobs + /// start. When this node is a shared root (or reachable from one), this + /// cache is shared across plans too, so children are shared without any + /// global lookup. cached_children: OnceLock>, } @@ -925,6 +1282,35 @@ impl TrieNode { .clone() } + fn child_map_slot(&self, col: ColumnId, arity: usize) -> &OnceLock { + &self.cached_children.get_or_init(|| { + let mut vec: Pooled = with_pool_set(|ps| ps.get()); + vec.resize_with(arity, OnceLock::new); + vec + })[col] + } + + fn child_map(&self, col: ColumnId, arity: usize) -> &ChildMap { + self.child_map_slot(col, arity) + .get_or_init(ChildMap::single) + } + + /// Initialize this node's child cache for `col` with the physical partition + /// used by the top index. If the column was already initialized, keep that + /// representation: in particular, do not promote an existing `Single` map. + /// The return value says whether the installed map has the requested + /// alignment. + fn initialize_aligned_child_map( + &self, + col: ColumnId, + arity: usize, + shard_count: usize, + ) -> bool { + self.child_map_slot(col, arity) + .get_or_init(|| ChildMap::sharded(shard_count)) + .is_aligned_to(shard_count) + } + /// Return the child node reached by additionally constraining `col = value` /// (and applying `edge_cs`). `sub` computes the child subset and is only /// called on a cache miss. @@ -943,31 +1329,8 @@ impl TrieNode { info: &TableInfo, sub: impl FnOnce() -> Subset, ) -> Arc { - let map = &self.cached_children.get_or_init(|| { - let mut vec: Pooled = with_pool_set(|ps| ps.get()); - vec.resize_with(info.spec.arity(), || RwLock::new(HashMap::default())); - vec - })[col]; - // Optimistic read path: most calls are cache hits, so try a shared lock - // first. A hit is only valid when the edge constraints match. - { - let guard = map.read().unwrap(); - if let Some((node, stored_cs)) = guard.get(&value) - && &**stored_cs == edge_cs - { - return node.clone(); - } - } - // Cache miss (or constraint mismatch): acquire the write lock and insert. - let mut guard = map.write().unwrap(); - if let Some((node, stored_cs)) = guard.get(&value) - && &**stored_cs == edge_cs - { - return node.clone(); - } - let new_node = Arc::new(TrieNode::new(sub())); - guard.insert(value, (new_node.clone(), Box::from(edge_cs))); - new_node + self.child_map(col, info.spec.arity()) + .get_or_insert(value, edge_cs, sub) } } @@ -1012,7 +1375,7 @@ impl BindingInfo { /// Probers returned from [`JoinState::get_index`] will move atom-related state out of the /// [`BindingInfo`]. Once the caller is done using a prober, this method moves it back. - fn move_back(&mut self, atom: AtomId, prober: Prober) { + fn move_back(&mut self, atom: AtomId, prober: Prober<'_>) { self.subsets.insert(atom, prober.node); } @@ -1029,6 +1392,44 @@ impl BindingInfo { } } +struct TopChildCacheTarget { + node: Arc, + column: ColumnId, + arity: usize, +} + +struct TopIndexCandidate { + order_position: usize, + leader_keys: usize, + shard_count: usize, + shards: Vec, + child_cache_targets: Vec, +} + +#[derive(Clone, Copy)] +struct TopIndexEligibility { + workers: usize, + min_keys_per_worker: usize, +} + +impl TopIndexCandidate { + /// Install an aligned cache layout only after this candidate has won top + /// variable selection. Candidate discovery may initialize column indexes, + /// but must not commit cache layouts for rejected stages. + fn configure_child_cache(&self) { + if gj_child_cache() != GjChildCache::Aligned { + return; + } + for target in &self.child_cache_targets { + let _ = target.node.initialize_aligned_child_map( + target.column, + target.arity, + self.shard_count, + ); + } + } +} + impl<'a> JoinState<'a> { fn new( db: &'a Database, @@ -1112,13 +1513,14 @@ impl<'a> JoinState<'a> { Some(subset) } - fn get_index( + fn get_index<'plan>( &self, atoms: &Arc>, atom: AtomId, binding_info: &mut BindingInfo, cols: impl Iterator, - ) -> Prober { + prepared: &'plan PreparedIndexSlot, + ) -> Prober<'plan> { let cols = SmallVec::<[ColumnId; 4]>::from_iter(cols); let trie_node = binding_info.subsets.unwrap_val(atom); let subset = &trie_node.subset; @@ -1132,14 +1534,7 @@ impl<'a> JoinState<'a> { cols[0], )) } else { - let all_cacheable = cols.iter().all(|col| { - !info - .spec - .uncacheable_columns - .get(*col) - .copied() - .unwrap_or(false) - }); + let all_cacheable = columns_are_cacheable(info, &cols); let whole_table = info.table.all(); if let Subset::Dense(range) = subset && all_cacheable @@ -1156,14 +1551,28 @@ impl<'a> JoinState<'a> { // large _or_ it is most of the table, or we already have a cached // index for it, then return it. if cols.len() != 1 { + let PreparedIndexSlot::Tuple(index) = prepared else { + unreachable!("multi-column scan must have a prepared tuple index") + }; + let index = + index.get_or_init(|| get_index_from_tableinfo(info, cols.as_slice())); DynamicIndex::Cached { intersect_outer, - table: get_index_from_tableinfo(info, &cols), + table: index + .get() + .expect("prepared tuple index must already be refreshed"), } } else { + let PreparedIndexSlot::Column(index) = prepared else { + unreachable!("single-column scan must have a prepared column index") + }; + let index = + index.get_or_init(|| get_column_index_from_tableinfo(info, cols[0])); DynamicIndex::CachedColumn { intersect_outer, - table: get_column_index_from_tableinfo(info, cols[0]).clone(), + table: index + .get() + .expect("prepared column index must already be refreshed"), } } } else if cols.len() != 1 { @@ -1178,14 +1587,162 @@ impl<'a> JoinState<'a> { ix: dyn_index, } } - fn get_column_index( + fn get_column_index<'plan>( &self, atoms: &Arc>, binding_info: &mut BindingInfo, atom: AtomId, col: ColumnId, - ) -> Prober { - self.get_index(atoms, atom, binding_info, iter::once(col)) + prepared: &'plan PreparedIndexSlot, + ) -> Prober<'plan> { + self.get_index(atoms, atom, binding_info, iter::once(col), prepared) + } + + /// Describe an eligible cached index that could drive a top-level + /// generic-join intersection. + /// + /// This mirrors the leader selection in [`Self::run_plan`]. Probers move + /// trie nodes out of `binding_info`, so restore every node before returning. + /// Every scan must cover its whole table through an unfiltered cached + /// hash/column index. Requiring this before constructing any prober avoids + /// building a dynamic nonleader index during discovery and then rebuilding + /// it in every shard job. Tiny or badly skewed leaders stay on the existing + /// path as well. + fn top_index_candidate( + &self, + order_position: usize, + stage: &JoinStage, + prepared: &[PreparedIndexSlot], + atoms: &Arc>, + binding_info: &mut BindingInfo, + eligibility: TopIndexEligibility, + ) -> Option { + let JoinStage::Intersect { scans, .. } = stage else { + return None; + }; + if scans.is_empty() || eligibility.workers <= 1 { + return None; + } + debug_assert_eq!(scans.len(), prepared.len()); + + // Mirror the cached/unfiltered branch in `get_index` without actually + // constructing any dynamic fallback indexes. + for scan in scans { + let node = &binding_info.subsets[scan.atom]; + let info = &self.db.tables[atoms[scan.atom].table]; + if node.subset.size() <= SMALL_RESIDUAL || !columns_are_cacheable(info, &[scan.column]) + { + return None; + } + let Subset::Dense(subset) = &node.subset else { + return None; + }; + let Subset::Dense(whole_table) = info.table.all() else { + return None; + }; + if subset != &whole_table { + return None; + } + } + + let mut leader = 0; + let mut leader_size = usize::MAX; + let mut probers = Vec::with_capacity(scans.len()); + for (i, (scan, prepared)) in scans.iter().zip(prepared).enumerate() { + let prober = + self.get_column_index(atoms, binding_info, scan.atom, scan.column, prepared); + let size = prober.len(); + // The two-way hot path historically breaks ties in favor of the + // second scan. Match it so partition discovery and execution use + // the same physical index. + if size < leader_size || (scans.len() == 2 && size == leader_size) { + leader = i; + leader_size = size; + } + probers.push(prober); + } + let candidate = probers[leader].shard_count().and_then(|shard_count| { + let shards = (0..shard_count) + .filter(|shard| probers[leader].shard_len(*shard).unwrap_or(0) != 0) + .collect::>(); + top_index_shape_is_eligible( + eligibility.workers, + leader_size, + shards.len(), + eligibility.min_keys_per_worker, + ) + .then(|| TopIndexCandidate { + order_position, + leader_keys: leader_size, + shard_count, + shards, + child_cache_targets: scans + .iter() + .zip(&probers) + .map(|(scan, prober)| TopChildCacheTarget { + node: prober.node.clone(), + column: scan.column, + arity: self.db.tables[atoms[scan.atom].table].spec.arity(), + }) + .collect(), + }) + }); + for (scan, prober) in scans.iter().zip(probers) { + binding_info.move_back(scan.atom, prober); + } + candidate + } + + /// Choose a top-index driver from the maximal leading run of plain + /// `Intersect` stages. Restricting promotion to that run is deliberately + /// conservative: every candidate is known to commute, so rotating one to + /// the front cannot cross a fused or materialization dependency. + fn select_top_index_candidate( + &self, + stages: &JoinStages, + prepared: &PreparedJoinIndexes, + atoms: &Arc>, + order: &InstrOrder, + binding_info: &mut BindingInfo, + ) -> Option { + let prefix_len = (0..order.len()) + .take_while(|position| { + matches!( + &stages.instrs[order.get(*position)], + JoinStage::Intersect { .. } + ) + }) + .count(); + if prefix_len == 0 { + return None; + } + + let eligibility = TopIndexEligibility { + workers: crate::parallel::current_num_threads(), + min_keys_per_worker: gj_min_keys_per_worker(), + }; + let current_estimate = estimate_size(&stages.instrs[order.get(0)], binding_info); + let min_threshold = eligibility + .min_keys_per_worker + .saturating_mul(eligibility.workers); + let candidates = (0..prefix_len).map(|order_position| { + let stage_index = order.get(order_position); + self.top_index_candidate( + order_position, + &stages.instrs[stage_index], + prepared.stage(stage_index), + atoms, + binding_info, + eligibility, + ) + }); + select_top_candidate( + gj_top_promotion(), + current_estimate, + min_threshold, + candidates, + |candidate| candidate.leader_keys, + ) } /// Runs the free join plan, starting with the header. @@ -1200,6 +1757,7 @@ impl<'a> JoinState<'a> { fn run_join_stages<'buf, A: NumericId + 'buf, BUF: ActionBuffer<'buf, A>>( &self, stages: &'buf JoinStages, + prepared: &'buf PreparedJoinIndexes, atoms: &'buf Arc>, action: A, binding_info: &mut BindingInfo, @@ -1218,13 +1776,82 @@ impl<'a> JoinState<'a> { let mut order = InstrOrder::from_iter(0..stages.instrs.len()); let mut leaf_scans: LeafScans = smallvec::smallvec![false; stages.instrs.len()]; sort_plan_by_size(&mut order, &mut leaf_scans, 0, &stages.instrs, binding_info); + + // A cached top index is already partitioned by hash. Schedule one + // coarse global job per nonempty shard. The shard may schedule one + // bounded level of coalesced local packets, but every packet's subtree + // remains serial. Besides avoiding an intermediate key copy, this keeps + // related nested probes on one worker by default. Buffers that cannot + // construct an independent partition (the in-place executors) decline + // this path. + if gj_top_index_sharding() + && !stages.instrs.is_empty() + && action_buf.supports_global_partition() + && let Some(candidate) = + self.select_top_index_candidate(stages, prepared, atoms, &order, binding_info) + { + if candidate.order_position != 0 { + order.promote_to_front(candidate.order_position); + recompute_leaf_scans(&order, &mut leaf_scans, &stages.instrs, 0); + } + // If configured, install the same value-hash partition in the + // root child cache as the selected leader before any shard job can + // touch these nodes. Rejected candidates never commit a cache + // layout. + candidate.configure_child_cache(); + let mut updates = FrameUpdates::with_capacity(0); + for shard in candidate.shards { + let db = self.db; + let exec_state_for_factory = self.exec_state.clone(); + let exec_state_for_work = self.exec_state.clone(); + let trie_cache = self.trie_cache.clone(); + action_buf.recur_global_serial( + BorrowedLocalState { + binding_info, + instr_order: &mut order, + leaf_scans: &mut leaf_scans, + updates: &mut updates, + }, + move || exec_state_for_factory.clone(), + move |BorrowedLocalState { + binding_info, + instr_order, + leaf_scans, + .. + }, + buf| { + JoinState { + db, + exec_state: exec_state_for_work, + pool: with_pool_set(|ps| ps.get_pool()), + trie_cache, + } + .run_plan( + stages, + prepared, + atoms, + action, + instr_order, + leaf_scans, + 0, + Some(shard), + binding_info, + buf, + ); + }, + ); + } + return; + } self.run_plan( stages, + prepared, atoms, action, &mut order, &mut leaf_scans, 0, + None, binding_info, action_buf, ); @@ -1242,11 +1869,13 @@ impl<'a> JoinState<'a> { fn run_plan<'buf, A: NumericId + 'buf, BUF: ActionBuffer<'buf, A>>( &self, stages: &'buf JoinStages, + prepared: &'buf PreparedJoinIndexes, atoms: &'buf Arc>, action: A, instr_order: &mut InstrOrder, leaf_scans: &mut LeafScans, cur: usize, + index_shard: Option, binding_info: &mut BindingInfo, action_buf: &mut BUF, ) where @@ -1282,7 +1911,13 @@ impl<'a> JoinState<'a> { } // TODO: `supports_parallel_drain`` is a hack because currently // `drain_updates_parallel!`` is a bit slower because of the additional ExecutionState clone. - if cur < free_join_fork_depth() && action_buf.supports_parallel_drain() { + if (index_shard.is_some() + && cur == 0 + && action_buf.supports_local_partition_drain()) + || (index_shard.is_none() + && cur < free_join_fork_depth() + && action_buf.supports_parallel_drain()) + { drain_updates_parallel!($updates) } else { $updates.drain(|update| match update { @@ -1310,11 +1945,13 @@ impl<'a> JoinState<'a> { } else { self.run_plan( stages, + prepared, atoms, action, instr_order, leaf_scans, cur + 1, + None, binding_info, action_buf, ); @@ -1370,11 +2007,13 @@ impl<'a> JoinState<'a> { } .run_plan( stages, + prepared, atoms, action, instr_order, leaf_scans, cur + 1, + None, binding_info, buf, ); @@ -1403,20 +2042,50 @@ impl<'a> JoinState<'a> { } } + // A sharded top-level job enumerates only its assigned physical index + // shard. Every recursive call clears `index_shard`, so this macro is + // used only by the leading prober of the initial `Intersect` stage. + macro_rules! for_each_leader { + ($prober:expr, $callback:expr) => { + if let Some(shard) = index_shard { + $prober.for_each_shard(shard, $callback) + } else { + $prober.for_each($callback) + } + }; + } + let pool = &self.pool; - match &stages.instrs[instr_order.get(cur)] { + let stage_index = instr_order.get(cur); + let stage = &stages.instrs[stage_index]; + let prepared_indexes = prepared.stage(stage_index); + debug_assert_eq!( + prepared_indexes.len(), + match stage { + JoinStage::Intersect { scans, .. } => scans.len(), + JoinStage::FusedIntersect { to_intersect, .. } + | JoinStage::FusedIntersectMat { to_intersect, .. } => to_intersect.len(), + } + ); + match stage { JoinStage::Intersect { var, scans } => match scans.as_slice() { [] => {} [a] => { if binding_info.has_empty_subset(a.atom) { return; } - let prober = self.get_column_index(atoms, binding_info, a.atom, a.column); + let prober = self.get_column_index( + atoms, + binding_info, + a.atom, + a.column, + &prepared_indexes[0], + ); let info = &self.db.tables[atoms[a.atom].table]; let table = info.table.as_ref(); let has_stale = table.has_stale_rows(); let mut updates = FrameUpdates::with_capacity(cmp::min(chunk_size, cur_size)); - prober.for_each(|val, x| { + for_each_leader!(prober, |val, x| { updates.push_binding(*var, val[0]); if x.size() <= 16 { let sub = refine_subset(x, &a.cs, &table, has_stale, pool); @@ -1448,9 +2117,21 @@ impl<'a> JoinState<'a> { binding_info.move_back(a.atom, prober); } [a, b] => { - let a_prober = self.get_column_index(atoms, binding_info, a.atom, a.column); - let b_prober = self.get_column_index(atoms, binding_info, b.atom, b.column); - + let a_prober = self.get_column_index( + atoms, + binding_info, + a.atom, + a.column, + &prepared_indexes[0], + ); + let b_prober = self.get_column_index( + atoms, + binding_info, + b.atom, + b.column, + &prepared_indexes[1], + ); + let ((smaller, smaller_scan), (larger, larger_scan)) = if a_prober.len() < b_prober.len() { ((&a_prober, a), (&b_prober, b)) @@ -1467,7 +2148,7 @@ impl<'a> JoinState<'a> { let small_table = small_info.table.as_ref(); let small_has_stale = small_table.has_stale_rows(); let mut updates = FrameUpdates::with_capacity(cmp::min(chunk_size, cur_size)); - smaller.for_each(|val, small_sub| { + for_each_leader!(smaller, |val, small_sub| { if let Some(large_sub) = larger.get_subset(val) { updates.push_binding(*var, val[0]); if small_sub.size() <= 16 { @@ -1555,9 +2236,14 @@ impl<'a> JoinState<'a> { let mut smallest = 0; let mut smallest_size = usize::MAX; let mut probers = Vec::with_capacity(rest.len()); - for (i, scan) in rest.iter().enumerate() { - let prober = - self.get_column_index(atoms, binding_info, scan.atom, scan.column); + for (i, (scan, prepared)) in rest.iter().zip(prepared_indexes).enumerate() { + let prober = self.get_column_index( + atoms, + binding_info, + scan.atom, + scan.column, + prepared, + ); let size = prober.len(); if size < smallest_size { smallest = i; @@ -1585,7 +2271,7 @@ impl<'a> JoinState<'a> { // Smallest leads the scan let mut updates = FrameUpdates::with_capacity(cmp::min(chunk_size, cur_size)); - probers[smallest].for_each(|key, sub| { + for_each_leader!(probers[smallest], |key, sub| { updates.push_binding(*var, key[0]); for (i, scan) in rest.iter().enumerate() { if i == smallest { @@ -1783,10 +2469,11 @@ impl<'a> JoinState<'a> { spec.to_index.atom, binding_info, spec.to_index.vars.iter().copied(), + &prepared_indexes[i], ), ) }) - .collect::>(); + .collect::); 4]>>(); // Pre-compute has_stale per prober to avoid vtable calls in the hot loop. let index_has_stale: SmallVec<[bool; 4]> = index_probers .iter() @@ -1962,15 +2649,17 @@ impl<'a> JoinState<'a> { let mut updates = FrameUpdates::with_capacity(cmp::min(chunk_size, cur_size)); let probers = to_intersect .iter() - .map(|(spec, _)| { + .zip(prepared_indexes) + .map(|((spec, _), prepared)| { self.get_index( atoms, spec.to_index.atom, binding_info, spec.to_index.vars.iter().copied(), + prepared, ) }) - .collect::>(); + .collect::; 4]>>(); // Pre-compute has_stale per prober to avoid vtable calls in the hot loop. let probers_has_stale: SmallVec<[bool; 4]> = to_intersect .iter() @@ -2125,6 +2814,9 @@ const LOCAL_ACTION_BATCH_SIZE: usize = 128; /// for serial and parallel modes. trait ActionBuffer<'state, A: NumericId>: Send { type AsLocal<'a>: ActionBuffer<'state, A> + where + 'state: 'a; + type AsGlobalSerial<'a>: ActionBuffer<'state, A> where 'state: 'a; @@ -2172,6 +2864,17 @@ trait ActionBuffer<'state, A: NumericId>: Send { work: impl for<'a> FnOnce(BorrowedLocalState<'a>, &mut Self::AsLocal<'a>) + Send + 'state, ); + /// Run one coarse partition as a global pool job, using a buffer whose + /// recursive work and action execution are both serial. Implementations + /// that return `false` from [`Self::supports_global_partition`] execute the + /// callback inline; the scoped implementations enqueue exactly one job. + fn recur_global_serial<'local>( + &mut self, + local: BorrowedLocalState<'local>, + to_exec_state: impl FnMut() -> ExecutionState<'state> + Send + 'state, + work: impl for<'a> FnOnce(BorrowedLocalState<'a>, &mut Self::AsGlobalSerial<'a>) + Send + 'state, + ); + /// The unit at which you should batch updates passed to calls to `recur`, /// potentially depending on the current level of recursion. /// @@ -2188,6 +2891,21 @@ trait ActionBuffer<'state, A: NumericId>: Send { fn supports_parallel_drain(&self) -> bool { true } + + /// Whether a coarse top-index partition may schedule its `cur = 0` + /// `FrameUpdates` packets on the current worker's local queue. + /// + /// This is separate from [`Self::supports_parallel_drain`] so a partition + /// can explicitly permit one local level while its packet buffer remains + /// fully serial. + fn supports_local_partition_drain(&self) -> bool { + false + } + + /// Whether this buffer can enqueue independent top-level partitions. + fn supports_global_partition(&self) -> bool { + false + } } /// The action buffer we use if we are executing in a single-threaded @@ -2203,6 +2921,10 @@ impl<'a, 'outer: 'a> ActionBuffer<'a, ActionId> for InPlaceActionBuffer<'outer> = Self where 'a: 'b; + type AsGlobalSerial<'b> + = Self + where + 'a: 'b; fn push_bindings( &mut self, @@ -2248,9 +2970,245 @@ impl<'a, 'outer: 'a> ActionBuffer<'a, ActionId> for InPlaceActionBuffer<'outer> work(local, self) } + fn recur_global_serial<'local>( + &mut self, + local: BorrowedLocalState<'local>, + _to_exec_state: impl FnMut() -> ExecutionState<'a> + Send + 'a, + work: impl for<'b> FnOnce(BorrowedLocalState<'b>, &mut Self) + Send + 'a, + ) { + work(local, self) + } + + fn supports_parallel_drain(&self) -> bool { + false + } +} + +#[derive(Clone)] +struct LocalPacketLimiter { + outstanding: Arc, + limit: usize, +} + +impl LocalPacketLimiter { + fn new() -> Self { + Self::with_limit(gj_local_queue_limit()) + } + + fn with_limit(limit: usize) -> Self { + Self { + outstanding: Arc::new(AtomicUsize::new(0)), + limit, + } + } + + fn try_acquire(&self) -> Option { + self.outstanding + .fetch_update(Ordering::AcqRel, Ordering::Acquire, |outstanding| { + (outstanding < self.limit).then_some(outstanding + 1) + }) + .ok() + .map(|_| LocalPacketPermit { + outstanding: self.outstanding.clone(), + }) + } +} + +struct LocalPacketPermit { + outstanding: Arc, +} + +impl Drop for LocalPacketPermit { + fn drop(&mut self) { + let previous = self.outstanding.fetch_sub(1, Ordering::AcqRel); + debug_assert!(previous > 0); + } +} + +/// Strictly serial action buffer used inside one locally scheduled packet. +/// It shares the rule-set match counter with sibling packets, but executes all +/// recursive join and action work inline and cannot schedule grandchildren. +struct SerialScopedActionBuffer<'scope> { + rule_set: &'scope RuleSet, + match_counter: Arc, + batches: DenseIdMap, +} + +impl<'scope> SerialScopedActionBuffer<'scope> { + fn new(rule_set: &'scope RuleSet, match_counter: Arc) -> Self { + Self { + rule_set, + match_counter, + batches: Default::default(), + } + } +} + +impl<'scope> ActionBuffer<'scope, ActionId> for SerialScopedActionBuffer<'scope> { + type AsLocal<'a> + = Self + where + 'scope: 'a; + type AsGlobalSerial<'a> + = Self + where + 'scope: 'a; + + fn push_bindings( + &mut self, + action: ActionId, + bindings: &DenseIdMap, + mut to_exec_state: impl FnMut() -> ExecutionState<'scope>, + ) { + let batch_size = action_batch_size(); + let action_state = self + .batches + .get_or_insert(action, || ActionState::new(batch_size)); + action_state.n_runs += 1; + action_state.len += 1; + let action_info = &self.rule_set.actions[action]; + // SAFETY: `used_vars` is constant for the rule and the bindings come + // from this rule's join plan. + unsafe { + action_state.bindings.push(bindings, &action_info.used_vars); + } + if action_state.len >= batch_size { + let succeeded = + to_exec_state().run_instrs(&action_info.instrs, &mut action_state.bindings); + action_state.bindings.clear(); + self.match_counter.inc_matches(action, succeeded); + action_state.len = 0; + } + } + + fn flush(&mut self, exec_state: &mut ExecutionState) { + flush_action_states( + exec_state, + &mut self.batches, + self.rule_set, + self.match_counter.as_ref(), + ); + } + + fn recur<'local>( + &mut self, + local: BorrowedLocalState<'local>, + _to_exec_state: impl FnMut() -> ExecutionState<'scope> + Send + 'scope, + work: impl for<'a> FnOnce(BorrowedLocalState<'a>, &mut Self) + Send + 'scope, + ) { + work(local, self); + } + + fn recur_global_serial<'local>( + &mut self, + local: BorrowedLocalState<'local>, + _to_exec_state: impl FnMut() -> ExecutionState<'scope> + Send + 'scope, + work: impl for<'a> FnOnce(BorrowedLocalState<'a>, &mut Self) + Send + 'scope, + ) { + work(local, self); + } + + fn supports_parallel_drain(&self) -> bool { + false + } +} + +/// Buffer owned by one globally scheduled top-index shard. It may fork +/// coalesced `cur = 0` packets onto that worker's local deque, but every packet +/// receives a [`SerialScopedActionBuffer`] and therefore runs the entire +/// `cur = 1` subtree and its actions serially. +struct TopPartitionActionBuffer<'inner, 'scope> { + scope: &'inner Scope<'scope>, + serial: SerialScopedActionBuffer<'scope>, + limiter: LocalPacketLimiter, +} + +impl<'inner, 'scope> TopPartitionActionBuffer<'inner, 'scope> { + fn new( + scope: &'inner Scope<'scope>, + rule_set: &'scope RuleSet, + match_counter: Arc, + ) -> Self { + Self { + scope, + serial: SerialScopedActionBuffer::new(rule_set, match_counter), + limiter: LocalPacketLimiter::new(), + } + } +} + +impl<'scope> ActionBuffer<'scope, ActionId> for TopPartitionActionBuffer<'_, 'scope> { + type AsLocal<'a> + = SerialScopedActionBuffer<'scope> + where + 'scope: 'a; + type AsGlobalSerial<'a> + = Self + where + 'scope: 'a; + + fn push_bindings( + &mut self, + action: ActionId, + bindings: &DenseIdMap, + to_exec_state: impl FnMut() -> ExecutionState<'scope>, + ) { + self.serial.push_bindings(action, bindings, to_exec_state); + } + + fn flush(&mut self, exec_state: &mut ExecutionState) { + self.serial.flush(exec_state); + } + + fn recur<'local>( + &mut self, + mut local: BorrowedLocalState<'local>, + mut to_exec_state: impl FnMut() -> ExecutionState<'scope> + Send + 'scope, + work: impl for<'a> FnOnce(BorrowedLocalState<'a>, &mut SerialScopedActionBuffer<'scope>) + + Send + + 'scope, + ) { + let rule_set = self.serial.rule_set; + let match_counter = self.serial.match_counter.clone(); + if let Some(permit) = self.limiter.try_acquire() { + let mut inner = local.clone_state(); + self.scope.spawn_local(move |_| { + let _permit = permit; + let mut buf = SerialScopedActionBuffer::new(rule_set, match_counter); + work(inner.borrow_mut(), &mut buf); + buf.flush(&mut to_exec_state()); + }); + } else { + let mut buf = SerialScopedActionBuffer::new(rule_set, match_counter); + work(local, &mut buf); + buf.flush(&mut to_exec_state()); + } + } + + fn recur_global_serial<'local>( + &mut self, + local: BorrowedLocalState<'local>, + _to_exec_state: impl FnMut() -> ExecutionState<'scope> + Send + 'scope, + work: impl for<'a> FnOnce(BorrowedLocalState<'a>, &mut Self) + Send + 'scope, + ) { + work(local, self); + } + + fn morsel_size(&mut self, level: usize, _total: usize) -> usize { + if level == 0 && gj_local_depth() == 1 { + gj_local_morsel() + } else { + 256 + } + } + fn supports_parallel_drain(&self) -> bool { false } + + fn supports_local_partition_drain(&self) -> bool { + gj_local_depth() == 1 + } } /// An action buffer that hands off batches of actions to scoped worker tasks. @@ -2283,6 +3241,10 @@ impl<'scope> ActionBuffer<'scope, ActionId> for ScopedActionBuffer<'_, 'scope> { = ScopedActionBuffer<'a, 'scope> where 'scope: 'a; + type AsGlobalSerial<'a> + = TopPartitionActionBuffer<'a, 'scope> + where + 'scope: 'a; fn push_bindings( &mut self, action: ActionId, @@ -2334,7 +3296,10 @@ impl<'scope> ActionBuffer<'scope, ActionId> for ScopedActionBuffer<'_, 'scope> { let rule_set = self.rule_set; let match_counter = self.match_counter.clone(); let mut inner = local.clone_state(); - self.scope.spawn(move |scope| { + // Keep recursive join work on the current worker when possible. If + // coarse top-index partitioning was unavailable, stalled workers cause + // the private deque to donate older siblings to the global queue. + self.scope.spawn_local(move |scope| { let mut buf: ScopedActionBuffer<'_, 'scope> = ScopedActionBuffer { scope, rule_set, @@ -2354,6 +3319,24 @@ impl<'scope> ActionBuffer<'scope, ActionId> for ScopedActionBuffer<'_, 'scope> { }); } + fn recur_global_serial<'local>( + &mut self, + mut local: BorrowedLocalState<'local>, + mut to_exec_state: impl FnMut() -> ExecutionState<'scope> + Send + 'scope, + work: impl for<'a> FnOnce(BorrowedLocalState<'a>, &mut TopPartitionActionBuffer<'a, 'scope>) + + Send + + 'scope, + ) { + let rule_set = self.rule_set; + let match_counter = self.match_counter.clone(); + let mut inner = local.clone_state(); + self.scope.spawn_global(move |scope| { + let mut buf = TopPartitionActionBuffer::new(scope, rule_set, match_counter); + work(inner.borrow_mut(), &mut buf); + buf.flush(&mut to_exec_state()); + }); + } + fn morsel_size(&mut self, _level: usize, _total: usize) -> usize { // Lower morsel size to increase parallelism. match _level { @@ -2361,6 +3344,10 @@ impl<'scope> ActionBuffer<'scope, ActionId> for ScopedActionBuffer<'_, 'scope> { _ => 256, } } + + fn supports_global_partition(&self) -> bool { + true + } } fn expand_binding_sets<'state, A: NumericId, BUF: ActionBuffer<'state, A> + ?Sized>( @@ -2435,6 +3422,10 @@ impl<'a> ActionBuffer<'a, MatId> for InPlaceMaterializer<'a> { = Self where 'a: 'b; + type AsGlobalSerial<'b> + = Self + where + 'a: 'b; fn push_bindings( &mut self, @@ -2480,11 +3471,201 @@ impl<'a> ActionBuffer<'a, MatId> for InPlaceMaterializer<'a> { work(local, self) } + fn recur_global_serial<'local>( + &mut self, + local: BorrowedLocalState<'local>, + _to_exec_state: impl FnMut() -> ExecutionState<'a> + Send + 'a, + work: impl for<'b> FnOnce(BorrowedLocalState<'b>, &mut Self) + Send + 'a, + ) { + work(local, self) + } + + fn supports_parallel_drain(&self) -> bool { + false + } +} + +/// Serial recursive view of a scoped materializer. Sibling top-level shards +/// still share the concurrent output maps, but this buffer never schedules +/// deeper work itself. +struct SerialScopedMaterializer { + specs: Arc>, + materializations: Arc, RowBuffer>>>>, + scratch_key: Vec, + scratch_val: Vec, +} + +impl SerialScopedMaterializer { + fn new( + specs: Arc>, + materializations: Arc, RowBuffer>>>>, + ) -> Self { + Self { + specs, + materializations, + scratch_key: Vec::new(), + scratch_val: Vec::new(), + } + } +} + +impl<'scope> ActionBuffer<'scope, MatId> for SerialScopedMaterializer { + type AsLocal<'a> + = Self + where + 'scope: 'a; + type AsGlobalSerial<'a> + = Self + where + 'scope: 'a; + + fn push_bindings( + &mut self, + mat_id: MatId, + bindings: &DenseIdMap, + _to_exec_state: impl FnMut() -> ExecutionState<'scope>, + ) { + let mat = self.materializations.get(mat_id).expect("invalid mat id"); + let spec = self.specs.get(mat_id).expect("invalid mat id"); + self.scratch_key.clear(); + self.scratch_key + .extend(spec.msg_vars.iter().map(|var| bindings[*var])); + self.scratch_val.clear(); + self.scratch_val + .extend(spec.val_vars.iter().map(|var| bindings[*var])); + if self.scratch_val.is_empty() { + self.scratch_val.push(Value::stale()); + } + match mat.entry(self.scratch_key.clone()) { + Entry::Occupied(mut occupied) => { + occupied.get_mut().add_row(&self.scratch_val); + } + Entry::Vacant(vacant) => { + let mut buffer = RowBuffer::new(usize::max(spec.val_vars.len(), 1)); + buffer.add_row(&self.scratch_val); + vacant.insert(buffer); + } + } + } + + fn flush(&mut self, _exec_state: &mut ExecutionState) {} + + fn recur<'local>( + &mut self, + local: BorrowedLocalState<'local>, + _to_exec_state: impl FnMut() -> ExecutionState<'scope> + Send + 'scope, + work: impl for<'a> FnOnce(BorrowedLocalState<'a>, &mut Self) + Send + 'scope, + ) { + work(local, self); + } + + fn recur_global_serial<'local>( + &mut self, + local: BorrowedLocalState<'local>, + _to_exec_state: impl FnMut() -> ExecutionState<'scope> + Send + 'scope, + work: impl for<'a> FnOnce(BorrowedLocalState<'a>, &mut Self) + Send + 'scope, + ) { + work(local, self); + } + fn supports_parallel_drain(&self) -> bool { false } } +/// Materialization counterpart to [`TopPartitionActionBuffer`]. Only the +/// `cur = 0` packet boundary may use the worker-local queue; every packet gets +/// a serial materializer and cannot spawn deeper work. +struct TopPartitionMaterializer<'inner, 'scope> { + scope: &'inner Scope<'scope>, + serial: SerialScopedMaterializer, + limiter: LocalPacketLimiter, +} + +impl<'inner, 'scope> TopPartitionMaterializer<'inner, 'scope> { + fn new( + scope: &'inner Scope<'scope>, + specs: Arc>, + materializations: Arc, RowBuffer>>>>, + ) -> Self { + Self { + scope, + serial: SerialScopedMaterializer::new(specs, materializations), + limiter: LocalPacketLimiter::new(), + } + } +} + +impl<'scope> ActionBuffer<'scope, MatId> for TopPartitionMaterializer<'_, 'scope> { + type AsLocal<'a> + = SerialScopedMaterializer + where + 'scope: 'a; + type AsGlobalSerial<'a> + = Self + where + 'scope: 'a; + + fn push_bindings( + &mut self, + mat_id: MatId, + bindings: &DenseIdMap, + to_exec_state: impl FnMut() -> ExecutionState<'scope>, + ) { + self.serial.push_bindings(mat_id, bindings, to_exec_state); + } + + fn flush(&mut self, exec_state: &mut ExecutionState) { + self.serial.flush(exec_state); + } + + fn recur<'local>( + &mut self, + mut local: BorrowedLocalState<'local>, + _to_exec_state: impl FnMut() -> ExecutionState<'scope> + Send + 'scope, + work: impl for<'a> FnOnce(BorrowedLocalState<'a>, &mut SerialScopedMaterializer) + Send + 'scope, + ) { + let specs = self.serial.specs.clone(); + let materializations = self.serial.materializations.clone(); + if let Some(permit) = self.limiter.try_acquire() { + let mut inner = local.clone_state(); + self.scope.spawn_local(move |_| { + let _permit = permit; + let mut buf = SerialScopedMaterializer::new(specs, materializations); + work(inner.borrow_mut(), &mut buf); + }); + } else { + let mut buf = SerialScopedMaterializer::new(specs, materializations); + work(local, &mut buf); + } + } + + fn recur_global_serial<'local>( + &mut self, + local: BorrowedLocalState<'local>, + _to_exec_state: impl FnMut() -> ExecutionState<'scope> + Send + 'scope, + work: impl for<'a> FnOnce(BorrowedLocalState<'a>, &mut Self) + Send + 'scope, + ) { + work(local, self); + } + + fn morsel_size(&mut self, level: usize, _total: usize) -> usize { + if level == 0 && gj_local_depth() == 1 { + gj_local_morsel() + } else { + 256 + } + } + + fn supports_parallel_drain(&self) -> bool { + false + } + + fn supports_local_partition_drain(&self) -> bool { + gj_local_depth() == 1 + } +} + struct ScopedMaterializer<'inner, 'scope> { scope: &'inner Scope<'scope>, specs: Arc>, @@ -2497,6 +3678,10 @@ impl<'scope> ActionBuffer<'scope, MatId> for ScopedMaterializer<'_, 'scope> { = ScopedMaterializer<'a, 'scope> where 'scope: 'a; + type AsGlobalSerial<'a> + = TopPartitionMaterializer<'a, 'scope> + where + 'scope: 'a; fn push_bindings( &mut self, @@ -2546,7 +3731,9 @@ impl<'scope> ActionBuffer<'scope, MatId> for ScopedMaterializer<'_, 'scope> { let specs = self.specs.clone(); let materializations = self.materializations.clone(); let mut inner = local.clone_state(); - scope.spawn(move |scope| { + // Match the action path: preserve locality by default and rely on + // private-deque donation when other workers need fallback work. + scope.spawn_local(move |scope| { let mut buf: ScopedMaterializer<'_, 'scope> = ScopedMaterializer { scope, specs, @@ -2557,6 +3744,27 @@ impl<'scope> ActionBuffer<'scope, MatId> for ScopedMaterializer<'_, 'scope> { work(inner.borrow_mut(), &mut buf); }); } + + fn recur_global_serial<'local>( + &mut self, + mut local: BorrowedLocalState<'local>, + _to_exec_state: impl FnMut() -> ExecutionState<'scope> + Send + 'scope, + work: impl for<'a> FnOnce(BorrowedLocalState<'a>, &mut TopPartitionMaterializer<'a, 'scope>) + + Send + + 'scope, + ) { + let specs = self.specs.clone(); + let materializations = self.materializations.clone(); + let mut inner = local.clone_state(); + self.scope.spawn_global(move |scope| { + let mut buf = TopPartitionMaterializer::new(scope, specs, materializations); + work(inner.borrow_mut(), &mut buf); + }); + } + + fn supports_global_partition(&self) -> bool { + true + } } struct MatchCounter { @@ -2608,7 +3816,7 @@ fn sort_plan_by_size( let mut last_pos = start; for i in start..instrs.len() { if matches!( - &instrs[i], + &instrs[order.get(i)], // These nodes don't commute JoinStage::FusedIntersectMat { mode: MatScanMode::Lookup(_) | MatScanMode::Value(_) | MatScanMode::Full, @@ -2725,8 +3933,8 @@ fn sort_plan_by_size_inner( let mut times_refined = with_pool_set(|ps| ps.get::>()); // Count how many times each atom has been refined so far. - for ins in instrs[..range.start].iter() { - match ins { + for position in 0..range.start { + match &instrs[order.get(position)] { JoinStage::Intersect { scans, .. } => scans.iter().for_each(|scan| { *times_refined.get_or_default(scan.atom) += 1; }), @@ -2844,6 +4052,12 @@ impl InstrOrder { fn len(&self) -> usize { self.data.len() } + + /// Move the stage at `position` to the front while preserving the relative + /// order of every stage it passes. + fn promote_to_front(&mut self, position: usize) { + self.data[..=position].rotate_right(1); + } } /// Per-position leaf-scan flags. `leaf_scans[i] == true` means the stage currently scheduled at @@ -2886,3 +4100,219 @@ impl LocalState { } } } + +#[cfg(test)] +mod local_packet_tests { + use std::{ + cell::Cell, + ptr, + sync::{ + Arc, Barrier, + atomic::{AtomicUsize, Ordering}, + }, + thread, + }; + + use crate::{ + common::Value, numeric_id::NumericId, offsets::Subset, parallel_heuristics::GjTopPromotion, + table_spec::ColumnId, + }; + + use super::{ + ChildMap, InstrOrder, LocalPacketLimiter, TrieNode, select_top_candidate, + top_index_shape_is_eligible, + }; + + #[test] + fn top_index_partitioning_rejects_serial_tiny_and_skewed_shapes() { + assert!(!top_index_shape_is_eligible(1, 10_000, 8, 64)); + assert!(!top_index_shape_is_eligible(4, 255, 8, 64)); + assert!(!top_index_shape_is_eligible(4, 10_000, 3, 64)); + assert!(top_index_shape_is_eligible(4, 256, 4, 64)); + assert!(top_index_shape_is_eligible(4, 40, 4, 10)); + } + + #[test] + fn current_top_policy_probes_only_sorted_position_zero() { + let probes = Cell::new(0); + let candidates = [None, Some((1, 512))].into_iter().inspect(|_| { + probes.set(probes.get() + 1); + }); + let selected = + select_top_candidate(GjTopPromotion::Current, 100, 256, candidates, |candidate| { + candidate.1 + }); + assert_eq!(selected, None); + assert_eq!(probes.get(), 1); + } + + #[test] + fn eligible_top_policy_keeps_current_or_uses_first_guarded_fallback() { + let probes = Cell::new(0); + let candidates = [Some((0, 256)), Some((1, 512))].into_iter().inspect(|_| { + probes.set(probes.get() + 1); + }); + assert_eq!( + select_top_candidate( + GjTopPromotion::Eligible, + 100, + 256, + candidates, + |candidate| candidate.1, + ), + Some((0, 256)) + ); + assert_eq!(probes.get(), 1); + + // max(4 * 100, 256) = 400: skip the first eligible-but-too-large + // fallback and preserve the first later candidate within the guard. + assert_eq!( + select_top_candidate( + GjTopPromotion::Eligible, + 100, + 256, + [None, Some((1, 401)), Some((2, 300))].into_iter(), + |candidate| candidate.1, + ), + Some((2, 300)) + ); + } + + #[test] + fn largest_top_policy_uses_actual_keys_and_keeps_earliest_tie() { + assert_eq!( + select_top_candidate( + GjTopPromotion::Largest, + 1, + 256, + [Some((0, 256)), None, Some((2, 1024)), Some((3, 1024))].into_iter(), + |candidate| candidate.1, + ), + Some((2, 1024)) + ); + } + + #[test] + fn top_promotion_preserves_passed_stage_order() { + let mut order = InstrOrder::from_iter([3, 1, 4, 2].into_iter()); + order.promote_to_front(2); + assert_eq!( + (0..order.len()).map(|i| order.get(i)).collect::>(), + vec![4, 3, 1, 2] + ); + } + + #[test] + fn local_packet_limiter_bounds_outstanding_work_and_reuses_slots() { + let limiter = LocalPacketLimiter::with_limit(2); + let first = limiter.try_acquire().expect("first packet should fit"); + let second = limiter.try_acquire().expect("second packet should fit"); + assert!(limiter.try_acquire().is_none()); + + drop(first); + let replacement = limiter + .try_acquire() + .expect("completed packet should release its slot"); + assert!(limiter.try_acquire().is_none()); + + drop((second, replacement)); + assert!(limiter.try_acquire().is_some()); + } + + #[test] + fn zero_local_packet_limit_always_executes_inline() { + let limiter = LocalPacketLimiter::with_limit(0); + assert!(limiter.try_acquire().is_none()); + } + + #[test] + fn sharded_child_cache_returns_one_arc_for_racing_same_key() { + const THREADS: usize = 16; + let map = Arc::new(ChildMap::sharded(8)); + let barrier = Arc::new(Barrier::new(THREADS)); + let constructed = Arc::new(AtomicUsize::new(0)); + let value = Value::from_usize(17); + + let handles = (0..THREADS) + .map(|_| { + let map = map.clone(); + let barrier = barrier.clone(); + let constructed = constructed.clone(); + thread::spawn(move || { + barrier.wait(); + map.get_or_insert(value, &[], || { + constructed.fetch_add(1, Ordering::Relaxed); + Subset::empty() + }) + }) + }) + .collect::>(); + let nodes = handles + .into_iter() + .map(|handle| handle.join().unwrap()) + .collect::>(); + + assert_eq!(constructed.load(Ordering::Relaxed), 1); + assert!(nodes[1..].iter().all(|node| Arc::ptr_eq(&nodes[0], node))); + } + + #[test] + fn sharded_child_cache_maps_distinct_values_to_every_aligned_lock() { + let map = Arc::new(ChildMap::sharded(8)); + let ChildMap::Sharded { shard_data, shards } = map.as_ref() else { + unreachable!() + }; + let mut values = vec![None; shard_data.n_shards()]; + for raw in 0..100_000 { + let value = Value::from_usize(raw); + let selected = shard_data.get_shard(&value, shards); + let shard = shards + .iter() + .find_map(|(shard, lock)| ptr::eq(lock, selected).then_some(shard.index())) + .unwrap(); + values[shard].get_or_insert(value); + if values.iter().all(Option::is_some) { + break; + } + } + assert!(values.iter().all(Option::is_some)); + + let handles = values + .into_iter() + .map(|value| { + let map = map.clone(); + thread::spawn(move || { + map.get_or_insert(value.unwrap(), &[], Subset::empty); + }) + }) + .collect::>(); + for handle in handles { + handle.join().unwrap(); + } + + let ChildMap::Sharded { shards, .. } = map.as_ref() else { + unreachable!() + }; + for (_, shard) in shards.iter() { + assert_eq!(shard.read().unwrap().len(), 1); + } + } + + #[test] + fn aligned_root_children_default_to_single_and_single_is_not_promoted() { + let col = ColumnId::new(0); + let root = TrieNode::new(Subset::empty()); + assert!(root.initialize_aligned_child_map(col, 1, 8)); + assert!(matches!(root.child_map(col, 1), ChildMap::Sharded { .. })); + assert!(!root.initialize_aligned_child_map(col, 1, 16)); + + let child = root + .child_map(col, 1) + .get_or_insert(Value::from_usize(3), &[], Subset::empty); + assert!(child.cached_children.get().is_none()); + assert!(matches!(child.child_map(col, 1), ChildMap::Single(_))); + + assert!(!child.initialize_aligned_child_map(col, 1, 8)); + assert!(matches!(child.child_map(col, 1), ChildMap::Single(_))); + } +} diff --git a/core-relations/src/hash_index/mod.rs b/core-relations/src/hash_index/mod.rs index 839667fe5..0cd2f0cb3 100644 --- a/core-relations/src/hash_index/mod.rs +++ b/core-relations/src/hash_index/mod.rs @@ -1,4 +1,6 @@ //! Hash-based secondary indexes. +#[cfg(test)] +use std::sync::atomic::{AtomicUsize, Ordering}; use std::{ cmp, hash::{Hash, Hasher}, @@ -132,6 +134,25 @@ impl Index { self.table.for_each(f); } + /// 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); + } + + /// Return the number of physical shards in this index. + pub(crate) fn shard_count(&self) -> usize { + self.table.shard_count() + } + + /// Return the number of distinct keys stored in one physical shard. + pub(crate) fn shard_len(&self, shard: usize) -> usize { + self.table.shard_len(shard) + } + pub(crate) fn len(&self) -> usize { self.table.len() } @@ -181,6 +202,12 @@ pub(crate) trait IndexBase { fn merge_rows(&mut self, buf: &TaggedRowBuffer); /// Call `f` over the elements of the index. fn for_each(&self, f: impl FnMut(&Self::Key, SubsetRef)); + /// Call `f` over the elements of one physical index shard. + fn for_each_shard(&self, shard: usize, f: impl FnMut(&Self::Key, SubsetRef)); + /// The number of physical shards in the index. + fn shard_count(&self) -> usize; + /// The number of distinct keys in one physical shard. + fn shard_len(&self, shard: usize) -> usize; /// The number of keys in the index. fn len(&self) -> usize; @@ -281,6 +308,21 @@ impl IndexBase for ColumnIndex { } } + fn for_each_shard(&self, shard: usize, mut f: impl FnMut(&Self::Key, SubsetRef)) { + let shard = &self.shards[ShardId::from_usize(shard)]; + for (key, subset) in shard.table.iter() { + f(key, subset.as_ref(&shard.subsets)); + } + } + + fn shard_count(&self) -> usize { + self.shards.len() + } + + fn shard_len(&self, shard: usize) -> usize { + self.shards[ShardId::from_usize(shard)].table.len() + } + fn len(&self) -> usize { self.shards.iter().map(|(_, shard)| shard.table.len()).sum() } @@ -720,6 +762,23 @@ impl IndexBase for TupleIndex { } } + fn for_each_shard(&self, shard: usize, mut f: impl FnMut(&Self::Key, SubsetRef)) { + let shard = &self.shards[ShardId::from_usize(shard)]; + 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(key, entry.vals.as_ref(&shard.subsets)); + } + } + + fn shard_count(&self) -> usize { + self.shards.len() + } + + fn shard_len(&self, shard: usize) -> usize { + self.shards[ShardId::from_usize(shard)].table.hash.len() + } + fn len(&self) -> usize { self.shards .iter() @@ -830,6 +889,8 @@ fn hash_key(key: &[Value]) -> u64 { #[derive(Default)] pub struct IndexCatalog { data: ReadOptimizedLock>, + #[cfg(test)] + get_or_insert_calls: AtomicUsize, } impl IndexCatalog @@ -839,6 +900,8 @@ where pub fn new() -> Self { IndexCatalog { data: ReadOptimizedLock::new(Vec::new()), + #[cfg(test)] + get_or_insert_calls: AtomicUsize::new(0), } } @@ -846,6 +909,8 @@ where let vec = self.data.read().iter().map(f).collect(); IndexCatalog { data: ReadOptimizedLock::new(vec), + #[cfg(test)] + get_or_insert_calls: AtomicUsize::new(0), } } @@ -856,6 +921,8 @@ where } pub fn get_or_insert(&self, k: K, init: impl FnOnce() -> I) -> I { + #[cfg(test)] + self.get_or_insert_calls.fetch_add(1, Ordering::Relaxed); let data = self.data.read(); let entry = data.iter().find(|(k1, _)| k1 == &k); if let Some(entry) = entry { @@ -872,6 +939,11 @@ where } } } + + #[cfg(test)] + pub(crate) fn get_or_insert_calls(&self) -> usize { + self.get_or_insert_calls.load(Ordering::Relaxed) + } } define_id!(BufferIndex, u32, "an index into a subset buffer"); diff --git a/core-relations/src/hash_index/tests.rs b/core-relations/src/hash_index/tests.rs index 35afc1ab8..d80914367 100644 --- a/core-relations/src/hash_index/tests.rs +++ b/core-relations/src/hash_index/tests.rs @@ -226,3 +226,63 @@ fn column_index_rebuild_matches_oracle() { } } } + +#[test] +fn physical_shards_partition_index_iteration() { + ThreadPool::new(4).install(|| { + let rows = (0..257) + .map(|i| vec![v(i), v(i % 37), v(10_000 + i)]) + .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()); + assert_eq!(column.shard_count(), 8); + let mut whole_column = BTreeMap::new(); + column.for_each(|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| { + shard_keys += 1; + assert!( + by_column_shard.insert(key.rep(), subset.size()).is_none(), + "a column-index key appeared in more than one shard" + ); + }); + assert_eq!(column.shard_len(shard), shard_keys); + } + assert_eq!(by_column_shard, whole_column); + + let mut tuple = Index::new( + vec![ColumnId::new(1), ColumnId::new(2)], + crate::TupleIndex::new(2), + ); + tuple.refresh(table.as_ref()); + assert_eq!(tuple.shard_count(), 8); + let mut whole_tuple = BTreeMap::new(); + tuple.for_each(|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| { + shard_keys += 1; + assert!( + by_tuple_shard + .insert((key[0].rep(), key[1].rep()), subset.size()) + .is_none(), + "a tuple-index key appeared in more than one shard" + ); + }); + assert_eq!(tuple.shard_len(shard), shard_keys); + } + assert_eq!(by_tuple_shard, whole_tuple); + }); +} diff --git a/core-relations/src/parallel_heuristics.rs b/core-relations/src/parallel_heuristics.rs index d6861470a..ef2de53af 100644 --- a/core-relations/src/parallel_heuristics.rs +++ b/core-relations/src/parallel_heuristics.rs @@ -14,6 +14,45 @@ const DEFAULT_TABLE_OP_CUTOFF: usize = 400_000; const DEFAULT_FREE_JOIN_FORK_DEPTH: usize = 2; const DEFAULT_ACTION_BATCH_SIZE: usize = 8 * 1024; +// Fixed generic-join scheduling choices selected by cross-workload +// benchmarking. Coarse index partitioning is enabled. Recursive fallback work +// uses worker-local queues, but already-sharded partitions do not create an +// additional packet level. Aligned child caches and top-variable promotion +// remain disabled. +const GJ_TOP_INDEX_SHARDING: bool = true; +const GJ_LOCAL_DEPTH: usize = 0; +const GJ_LOCAL_MORSEL: usize = 64; +const GJ_LOCAL_QUEUE_LIMIT: usize = 2; +const GJ_CHILD_CACHE: GjChildCache = GjChildCache::Single; +const GJ_TOP_PROMOTION: GjTopPromotion = GjTopPromotion::Current; +const GJ_MIN_KEYS_PER_WORKER: usize = 16; + +/// Layout for root trie-node child caches during coarse generic-join +/// partitioning. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) enum GjChildCache { + /// One lock per cached column. + Single, + /// One lock per top-index shard, using the same value hash partition. + Aligned, +} + +/// Policy for replacing the initially sorted top generic-join variable with a +/// coarse-partitionable later variable. +// The non-selected policies remain available to the executor's deterministic +// policy tests even though production uses the fixed `Current` policy. +#[cfg_attr(not(test), allow(dead_code))] +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) enum GjTopPromotion { + /// Consider only the variable already sorted into position zero. + Current, + /// Keep an eligible current variable, otherwise choose the first guarded + /// eligible variable in the leading Intersect prefix. + Eligible, + /// Choose the eligible variable with the largest actual leader cardinality. + Largest, +} + static CUTOFFS: OnceLock = OnceLock::new(); struct Cutoffs { @@ -70,6 +109,44 @@ pub(crate) fn action_batch_size() -> usize { cutoffs().action_batch_size } +/// Whether a parallel generic join may schedule one coarse global job per +/// physical shard of its top cached index. +pub(crate) fn gj_top_index_sharding() -> bool { + GJ_TOP_INDEX_SHARDING +} + +/// Number of recursive levels below a coarse top-index shard that may enqueue +/// worker-local packets. The implementation currently accepts only zero or one. +pub(crate) fn gj_local_depth() -> usize { + GJ_LOCAL_DEPTH +} + +/// Number of top-level generic-join frames coalesced into one local packet. +pub(crate) fn gj_local_morsel() -> usize { + GJ_LOCAL_MORSEL +} + +/// Maximum number of outstanding local packets per coarse index shard. +pub(crate) fn gj_local_queue_limit() -> usize { + GJ_LOCAL_QUEUE_LIMIT +} + +/// Child-cache layout used by eligible coarse top-index partitions. +pub(crate) fn gj_child_cache() -> GjChildCache { + GJ_CHILD_CACHE +} + +/// Policy for promoting an eligible top generic-join variable. +pub(crate) fn gj_top_promotion() -> GjTopPromotion { + GJ_TOP_PROMOTION +} + +/// Minimum number of leader keys required per worker for coarse top-index +/// partitioning. +pub(crate) fn gj_min_keys_per_worker() -> usize { + GJ_MIN_KEYS_PER_WORKER +} + fn should_parallelize(len: usize, cutoff: usize) -> bool { len > cutoff && crate::parallel::current_num_threads() > 1 } diff --git a/core-relations/src/pool/mod.rs b/core-relations/src/pool/mod.rs index d62574490..40123c22f 100644 --- a/core-relations/src/pool/mod.rs +++ b/core-relations/src/pool/mod.rs @@ -12,7 +12,7 @@ use std::{ use crate::{ AtomId, - free_join::execute::TrieNode, + free_join::execute::ChildMap, numeric_id::{DenseIdMap, IdVec}, }; use fixedbitset::FixedBitSet; @@ -442,7 +442,7 @@ pool_set! { cached_subsets: IdVec>> [ 4 << 20 ], intersected_on: DenseIdMap [ 1 << 20 ], - cached_child: IdVec, Box<[Constraint]>)>>> [ 1 << 20 ], + cached_child: IdVec> [ 1 << 20 ], } } diff --git a/core-relations/src/tests.rs b/core-relations/src/tests.rs index ec6a14742..04c747221 100644 --- a/core-relations/src/tests.rs +++ b/core-relations/src/tests.rs @@ -12,7 +12,10 @@ use crate::{ PlanStrategy, action::WriteVal, common::Value, - free_join::{CounterId, Database, TableId}, + free_join::{ + CounterId, Database, TableId, + plan::{JoinStage, Plan}, + }, make_external_func, query::RuleSetBuilder, table::SortedWritesTable, @@ -210,6 +213,586 @@ fn line_graph_1_test(strat: PlanStrategy) { assert_eq!(expected, got); } +#[test] +fn prepared_plan_indexes_refresh_across_runs_and_clear() { + let mut db = Database::default(); + let make_table = |arity| { + SortedWritesTable::new( + arity, + arity, + None, + vec![], + Box::new(|_, old, new, _| { + assert_eq!(old, new, "test tables have unique keys"); + false + }), + ) + }; + let driver = db.add_table(make_table(2), iter::empty(), iter::empty()); + let target = db.add_table(make_table(3), iter::empty(), iter::empty()); + let allowed = db.add_table(make_table(1), iter::empty(), iter::empty()); + let output = db.add_table(make_table(1), iter::empty(), iter::empty()); + + const KEYS: usize = 4; + const VALUES_PER_KEY: usize = 32; + let populate_inputs = |db: &Database, value_base: usize| { + { + let mut buf = db.new_buffer(target); + for x in 0..KEYS { + for z in 0..VALUES_PER_KEY { + buf.stage_insert(&[v(x), v(10_000 + x), v(value_base + x * 100 + z)]); + } + } + } + { + let mut buf = db.new_buffer(allowed); + for x in 0..KEYS { + for z in 0..VALUES_PER_KEY { + buf.stage_insert(&[v(value_base + x * 100 + z)]); + } + } + } + }; + + { + let mut buf = db.new_buffer(driver); + for x in 0..KEYS { + buf.stage_insert(&[v(x), v(10_000 + x)]); + } + } + populate_inputs(&db, 20_000); + db.merge_all(); + + let mut rsb = db.new_rule_set(); + let mut query = rsb.new_rule(); + query.set_plan_strategy(PlanStrategy::PureSize); + query.set_no_decomp(true); + let x = query.new_var_named("x"); + let y = query.new_var_named("y"); + let z = query.new_var_named("z"); + query.add_atom(driver, &[x.into(), y.into()], &[]).unwrap(); + query + .add_atom(target, &[x.into(), y.into(), z.into()], &[]) + .unwrap(); + query.add_atom(allowed, &[z.into()], &[]).unwrap(); + let mut rule = query.build(); + rule.insert(output, &[z.into()]).unwrap(); + rule.build_with_description("prepared-index-lifecycle"); + let rules = rsb.build(); + + let (plan, _, _) = rules.plans.values().next().unwrap(); + let Plan::SinglePlan(plan) = plan else { + panic!("set_no_decomp must produce a single plan") + }; + assert!( + plan.stages.instrs.iter().any(|stage| matches!( + stage, + JoinStage::FusedIntersect { to_intersect, .. } + if to_intersect + .iter() + .any(|(scan, _)| scan.to_index.vars.len() > 1) + )), + "test must exercise a prepared tuple index" + ); + assert!( + plan.stages + .instrs + .iter() + .any(|stage| matches!(stage, JoinStage::Intersect { scans, .. } if !scans.is_empty())), + "test must exercise a prepared column index" + ); + + let planned_tuple_lookups = plan + .stages + .instrs + .iter() + .map(|stage| match stage { + JoinStage::Intersect { .. } => 0, + JoinStage::FusedIntersect { to_intersect, .. } + | JoinStage::FusedIntersectMat { to_intersect, .. } => to_intersect + .iter() + .filter(|(scan, _)| { + plan.atoms[scan.to_index.atom].table == target && scan.to_index.vars.len() != 1 + }) + .count(), + }) + .sum::(); + let planned_allowed_column_lookups = plan + .stages + .instrs + .iter() + .map(|stage| match stage { + JoinStage::Intersect { scans, .. } => scans + .iter() + .filter(|scan| plan.atoms[scan.atom].table == allowed) + .count(), + JoinStage::FusedIntersect { to_intersect, .. } + | JoinStage::FusedIntersectMat { to_intersect, .. } => to_intersect + .iter() + .filter(|(scan, _)| { + plan.atoms[scan.to_index.atom].table == allowed && scan.to_index.vars.len() == 1 + }) + .count(), + }) + .sum::(); + assert!(planned_tuple_lookups > 0); + assert!(planned_allowed_column_lookups > 0); + + let expected = KEYS * VALUES_PER_KEY; + let tuple_lookups_before = db.get_table_info(target).indexes.get_or_insert_calls(); + let allowed_column_lookups_before = db + .get_table_info(allowed) + .column_indexes + .get_or_insert_calls(); + let first = db.run_rule_set(&rules, ReportLevel::TimeOnly); + assert_eq!(first.num_matches("prepared-index-lifecycle"), expected); + assert_eq!(db.get_table(output).len(), expected); + assert_eq!( + db.get_table_info(target).indexes.get_or_insert_calls() - tuple_lookups_before, + planned_tuple_lookups, + "tuple-index catalog access must be plan-static, not recursion-static" + ); + assert_eq!( + db.get_table_info(allowed) + .column_indexes + .get_or_insert_calls() + - allowed_column_lookups_before, + planned_allowed_column_lookups, + "column-index catalog access must be plan-static, not recursion-static" + ); + + // Reusing the same logical plan must not retain catalog Arcs across the + // merge at the end of the previous run. + let second = db.run_rule_set(&rules, ReportLevel::TimeOnly); + assert_eq!(second.num_matches("prepared-index-lifecycle"), expected); + assert_eq!(db.get_table(output).len(), expected); + + db.clear_table(target); + db.clear_table(allowed); + db.clear_table(output); + populate_inputs(&db, 40_000); + db.merge_all(); + + // Clearing bumps the table generation, so every prepared global index must + // be refreshed before its borrowed reference is exposed to execution. + let after_clear = db.run_rule_set(&rules, ReportLevel::TimeOnly); + assert_eq!( + after_clear.num_matches("prepared-index-lifecycle"), + expected + ); + let output_table = db.get_table(output); + let output_rows = output_table.scan(output_table.all().as_ref()); + assert_eq!(output_rows.len(), expected); + assert!(output_rows.iter().all(|(_, row)| row[0].index() >= 40_000)); +} + +#[test] +fn prepared_plan_indexes_preserve_uncacheable_columns() { + let mut db = Database::default(); + let displaced = db.add_table(DisplacedTable::default(), iter::empty(), iter::empty()); + let make_table = || { + SortedWritesTable::new( + 1, + 1, + None, + vec![], + Box::new(|_, old, new, _| { + assert_eq!(old, new); + false + }), + ) + }; + let representatives = db.add_table(make_table(), iter::empty(), iter::empty()); + let output = db.add_table(make_table(), iter::empty(), iter::empty()); + { + let mut buf = db.new_buffer(displaced); + buf.stage_insert(&[v(0), v(10), v(0)]); + } + db.merge_all(); + let displaced_table = db.get_table(displaced); + let displaced_rows = displaced_table.scan(displaced_table.all().as_ref()); + let displaced_row = displaced_rows.iter().next().unwrap().1; + let displaced_key = displaced_row[0]; + let canonical = displaced_row[1]; + let displaced_timestamp = displaced_row[2]; + { + let mut buf = db.new_buffer(representatives); + buf.stage_insert(&[canonical]); + } + db.merge_all(); + + let mut rsb = db.new_rule_set(); + let mut query = rsb.new_rule(); + query.set_plan_strategy(PlanStrategy::Gj); + query.set_no_decomp(true); + let representative = query.new_var_named("representative"); + query + .add_atom( + displaced, + &[ + displaced_key.into(), + representative.into(), + displaced_timestamp.into(), + ], + &[], + ) + .unwrap(); + query + .add_atom(representatives, &[representative.into()], &[]) + .unwrap(); + let mut rule = query.build(); + rule.insert(output, &[representative.into()]).unwrap(); + rule.build_with_description("uncacheable-prepared-index"); + let rules = rsb.build(); + + let (plan, _, _) = rules.plans.values().next().unwrap(); + let Plan::SinglePlan(plan) = plan else { + panic!("set_no_decomp must produce a single plan") + }; + assert!(plan.stages.instrs.iter().any(|stage| matches!( + stage, + JoinStage::Intersect { scans, .. } + if scans.iter().any(|scan| { + plan.atoms[scan.atom].table == displaced && scan.column == ColumnId::new(1) + }) + ))); + + let catalog_calls_before = db + .get_table_info(displaced) + .column_indexes + .get_or_insert_calls(); + let report = db.run_rule_set(&rules, ReportLevel::TimeOnly); + assert_eq!(report.num_matches("uncacheable-prepared-index"), 1); + assert_eq!( + db.get_table_info(displaced) + .column_indexes + .get_or_insert_calls(), + catalog_calls_before, + "preparation must not create a global index for a dynamic column" + ); + assert_eq!(db.get_table(output).len(), 1); +} + +#[test] +fn gj_top_index_shards_preserve_count_and_materialization() { + assert!(crate::parallel_heuristics::gj_top_index_sharding()); + assert_eq!(crate::parallel_heuristics::gj_local_depth(), 0); + assert_eq!( + crate::parallel_heuristics::gj_child_cache(), + crate::parallel_heuristics::GjChildCache::Single + ); + assert_eq!( + crate::parallel_heuristics::gj_top_promotion(), + crate::parallel_heuristics::GjTopPromotion::Current + ); + assert_eq!(crate::parallel_heuristics::gj_min_keys_per_worker(), 16); + gj_top_index_shards_preserve_count_and_materialization_inner(); +} + +fn gj_top_index_shards_preserve_count_and_materialization_inner() { + let pool = egglog_concurrency::ThreadPool::new(4); + pool.install(|| { + let mut db = Database::default(); + let make_table = || { + SortedWritesTable::new( + 2, + 2, + None, + vec![], + Box::new(|_, old, new, _| { + assert_eq!(old, new, "test tables have unique keys"); + false + }), + ) + }; + let left = db.add_table(make_table(), iter::empty(), iter::empty()); + let right = db.add_table(make_table(), iter::empty(), iter::empty()); + let output = db.add_table( + SortedWritesTable::new( + 2, + 2, + None, + vec![], + Box::new(|_, old, new, _| { + assert_eq!(old, new, "materialized pairs are unique"); + false + }), + ), + iter::empty(), + iter::empty(), + ); + + const KEYS: usize = 513; + const DEGREE: usize = 17; + { + let mut buf = db.new_buffer(left); + for x in 0..KEYS { + for y in 0..DEGREE { + buf.stage_insert(&[v(x), v(100_000 + x * DEGREE + y)]); + } + } + } + { + let mut buf = db.new_buffer(right); + for x in 0..KEYS { + for z in 0..DEGREE { + buf.stage_insert(&[v(x), v(200_000 + x * DEGREE + z)]); + } + } + } + db.merge_all(); + + let mut rsb = RuleSetBuilder::new(&mut db); + let mut query = rsb.new_rule(); + query.set_plan_strategy(PlanStrategy::Gj); + let x = query.new_var_named("x"); + let y = query.new_var_named("y"); + let z = query.new_var_named("z"); + query.add_atom(left, &[x.into(), y.into()], &[]).unwrap(); + query.add_atom(right, &[x.into(), z.into()], &[]).unwrap(); + let mut rule = query.build(); + rule.insert(output, &[y.into(), z.into()]).unwrap(); + rule.build_with_description("sharded-gj"); + let rules = rsb.build(); + + pool.reset_scheduler_metrics(); + let report = db.run_rule_set(&rules, ReportLevel::TimeOnly); + let scheduler = pool.scheduler_metrics(); + // The root rule job plus one job for each of the cached index's + // 2 * worker-count shards. Merge work may add more global jobs. + assert!( + scheduler.global_pushes >= 9, + "top index sharding did not activate: {scheduler:?}" + ); + assert_eq!( + scheduler.local_pushes, 0, + "the fixed scheduler policy should keep each shard subtree serial" + ); + let expected = KEYS * DEGREE * DEGREE; + assert_eq!(report.num_matches("sharded-gj"), expected); + + let table = db.get_table(output); + let materialized = table.scan(table.all().as_ref()); + assert_eq!(materialized.len(), expected); + for (_, row) in materialized.iter() { + let y = row[0].index() - 100_000; + let z = row[1].index() - 200_000; + assert_eq!(y / DEGREE, z / DEGREE); + } + }); +} + +#[test] +fn gj_small_top_fallback_uses_local_queue() { + let pool = egglog_concurrency::ThreadPool::new(4); + pool.install(|| { + let mut db = Database::default(); + let make_table = || { + SortedWritesTable::new( + 3, + 3, + None, + vec![], + Box::new(|_, old, new, _| { + assert_eq!(old, new, "test tables have unique keys"); + false + }), + ) + }; + let left = db.add_table(make_table(), iter::empty(), iter::empty()); + let right = db.add_table(make_table(), iter::empty(), iter::empty()); + let output = db.add_table(make_table(), iter::empty(), iter::empty()); + + // The sorted top variable has only two keys and is therefore too small + // for coarse index-shard partitioning, while the database is large + // enough to enable parallel rule execution. + const ROWS: usize = 8192; + for table in [left, right] { + let mut buf = db.new_buffer(table); + for i in 0..ROWS { + buf.stage_insert(&[v(i % 2), v(10_000 + i), v(20_000 + i)]); + } + } + db.merge_all(); + + let mut rsb = RuleSetBuilder::new(&mut db); + let mut query = rsb.new_rule(); + query.set_plan_strategy(PlanStrategy::Gj); + query.set_no_decomp(true); + let x = query.new_var_named("x"); + let y = query.new_var_named("y"); + let z = query.new_var_named("z"); + query + .add_atom(left, &[x.into(), y.into(), z.into()], &[]) + .unwrap(); + query + .add_atom(right, &[x.into(), y.into(), z.into()], &[]) + .unwrap(); + let mut rule = query.build(); + rule.insert(output, &[x.into(), y.into(), z.into()]) + .unwrap(); + rule.build_with_description("local-fallback-gj"); + let rules = rsb.build(); + + pool.reset_scheduler_metrics(); + let report = db.run_rule_set(&rules, ReportLevel::TimeOnly); + let scheduler = pool.scheduler_metrics(); + assert!( + scheduler.local_pushes > 0, + "small-top fallback did not use worker-local tasks: {scheduler:?}" + ); + assert_eq!( + scheduler.local_pushes, + scheduler.local_pops + scheduler.donated_jobs, + "every local task must run on its owner or be donated" + ); + assert_eq!(report.num_matches("local-fallback-gj"), ROWS); + + let table = db.get_table(output); + let materialized = table.scan(table.all().as_ref()); + assert_eq!(materialized.len(), ROWS); + for (_, row) in materialized.iter() { + let i = row[2].index() - 20_000; + assert_eq!(row[0].index(), i % 2); + assert_eq!(row[1].index() - 10_000, i); + } + }); +} + +#[test] +fn gj_decomposed_small_top_materialization_uses_local_queue() { + let pool = egglog_concurrency::ThreadPool::new(4); + pool.install(|| { + let mut db = Database::default(); + let make_table = |arity| { + SortedWritesTable::new( + arity, + arity, + None, + vec![], + Box::new(|_, old, new, _| { + assert_eq!(old, new, "test tables have unique keys"); + false + }), + ) + }; + let left_xa = db.add_table(make_table(2), iter::empty(), iter::empty()); + let left_ab = db.add_table(make_table(2), iter::empty(), iter::empty()); + let left_bx = db.add_table(make_table(2), iter::empty(), iter::empty()); + let right_xc = db.add_table(make_table(2), iter::empty(), iter::empty()); + let right_cd = db.add_table(make_table(2), iter::empty(), iter::empty()); + let right_dx = db.add_table(make_table(2), iter::empty(), iter::empty()); + let output = db.add_table(make_table(5), iter::empty(), iter::empty()); + + // Two cyclic bags joined through x force a decomposed plan. The left + // bag is large, but x has only two keys, so its materialization block + // must use recursive fallback work instead of top-index sharding. + const X_KEYS: usize = 2; + const LEFT_ROWS_PER_KEY: usize = 4096; + for x in 0..X_KEYS { + let mut xa = db.new_buffer(left_xa); + let mut ab = db.new_buffer(left_ab); + let mut bx = db.new_buffer(left_bx); + for i in 0..LEFT_ROWS_PER_KEY { + let a = 10_000 + x * LEFT_ROWS_PER_KEY + i; + let b = 20_000 + x * LEFT_ROWS_PER_KEY + i; + xa.stage_insert(&[v(x), v(a)]); + ab.stage_insert(&[v(a), v(b)]); + bx.stage_insert(&[v(b), v(x)]); + } + } + { + let mut xc = db.new_buffer(right_xc); + let mut cd = db.new_buffer(right_cd); + let mut dx = db.new_buffer(right_dx); + for x in 0..X_KEYS { + let c = 100_000 + x; + let d = 200_000 + x; + xc.stage_insert(&[v(x), v(c)]); + cd.stage_insert(&[v(c), v(d)]); + dx.stage_insert(&[v(d), v(x)]); + } + } + db.merge_all(); + + let mut rsb = RuleSetBuilder::new(&mut db); + let mut query = rsb.new_rule(); + query.set_plan_strategy(PlanStrategy::Gj); + let x = query.new_var_named("x"); + let a = query.new_var_named("a"); + let b = query.new_var_named("b"); + let c = query.new_var_named("c"); + let d = query.new_var_named("d"); + query.add_atom(left_xa, &[x.into(), a.into()], &[]).unwrap(); + query.add_atom(left_ab, &[a.into(), b.into()], &[]).unwrap(); + query.add_atom(left_bx, &[b.into(), x.into()], &[]).unwrap(); + query + .add_atom(right_xc, &[x.into(), c.into()], &[]) + .unwrap(); + query + .add_atom(right_cd, &[c.into(), d.into()], &[]) + .unwrap(); + query + .add_atom(right_dx, &[d.into(), x.into()], &[]) + .unwrap(); + let mut rule = query.build(); + rule.insert(output, &[x.into(), a.into(), b.into(), c.into(), d.into()]) + .unwrap(); + rule.build_with_description("local-materialization-fallback-gj"); + let rules = rsb.build(); + + let (plan, _, _) = rules.plans.values().next().unwrap(); + let Plan::DecomposedPlan(plan) = plan else { + panic!("the two cyclic bags must produce a decomposed plan") + }; + assert!(plan.stages.blocks.len() >= 2); + let first_materialization = &plan.stages.blocks[0].0; + assert!(first_materialization.instrs.len() >= 3); + let Some(JoinStage::Intersect { var, scans }) = first_materialization.instrs.first() else { + panic!("the first materialization block must start by intersecting x") + }; + assert_eq!(*var, x); + assert!( + scans.iter().any(|scan| { + let table = plan.atoms[scan.atom].table; + table == right_xc || table == right_dx + }), + "the top x intersection must include a two-row index, making coarse sharding ineligible" + ); + + pool.reset_scheduler_metrics(); + let report = db.run_rule_set(&rules, ReportLevel::TimeOnly); + let scheduler = pool.scheduler_metrics(); + assert!( + scheduler.local_pushes > 0, + "small-top materialization fallback did not use worker-local tasks: {scheduler:?}" + ); + assert_eq!( + scheduler.local_pushes, + scheduler.local_pops + scheduler.donated_jobs, + "every local task must run on its owner or be donated" + ); + + let expected = X_KEYS * LEFT_ROWS_PER_KEY; + assert_eq!( + report.num_matches("local-materialization-fallback-gj"), + expected + ); + let table = db.get_table(output); + let materialized = table.scan(table.all().as_ref()); + assert_eq!(materialized.len(), expected); + for (_, row) in materialized.iter() { + let x = row[0].index(); + let i = row[1].index() - 10_000 - x * LEFT_ROWS_PER_KEY; + assert!(x < X_KEYS); + assert!(i < LEFT_ROWS_PER_KEY); + assert_eq!(row[2].index(), 20_000 + x * LEFT_ROWS_PER_KEY + i); + assert_eq!(row[3].index(), 100_000 + x); + assert_eq!(row[4].index(), 200_000 + x); + } + }); +} + #[test] fn line_graph_2_fj_puresize() { run_serial_and_parallel(|| line_graph_2_test(PlanStrategy::PureSize)); From 252e49376f52584ddfe0585afc5df4ebddd9f969 Mon Sep 17 00:00:00 2001 From: Eli Rosenthal Date: Mon, 20 Jul 2026 00:01:36 +0200 Subject: [PATCH 2/5] Remove retired generic-join tuning paths --- core-relations/src/free_join/execute.rs | 796 +++------------------- core-relations/src/parallel_heuristics.rs | 79 +-- core-relations/src/pool/mod.rs | 4 +- core-relations/src/tests.rs | 11 - 4 files changed, 93 insertions(+), 797 deletions(-) diff --git a/core-relations/src/free_join/execute.rs b/core-relations/src/free_join/execute.rs index 75d1afa51..8f32d5419 100644 --- a/core-relations/src/free_join/execute.rs +++ b/core-relations/src/free_join/execute.rs @@ -27,7 +27,7 @@ use web_time::Instant; use crate::{ Constraint, OffsetRange, Pool, SubsetRef, action::{Bindings, ExecutionState}, - common::{DashMap, ShardData, ShardId, Value}, + common::{DashMap, Value}, free_join::{ frame_update::{FrameUpdates, UpdateInstr}, get_index_from_tableinfo, @@ -35,9 +35,8 @@ use crate::{ hash_index::{ColumnIndex, Index, IndexBase, TupleIndex}, offsets::{Offsets, RowId, SortedOffsetSlice, SortedOffsetVector, Subset}, parallel_heuristics::{ - GjChildCache, GjTopPromotion, action_batch_size, free_join_fork_depth, gj_child_cache, - gj_local_depth, gj_local_morsel, gj_local_queue_limit, gj_min_keys_per_worker, - gj_top_index_sharding, gj_top_promotion, parallelize_db_level_op, + MIN_TOP_INDEX_KEYS_PER_WORKER, action_batch_size, free_join_fork_depth, + parallelize_db_level_op, }, pool::Pooled, query::RuleSet, @@ -65,41 +64,6 @@ fn top_index_shape_is_eligible( && nonempty_shards >= workers } -/// Select one already-evaluated eligible candidate. The iterator remains lazy, -/// so `Current` evaluates only position zero and `Eligible` stops at its first -/// usable fallback; only `Largest` evaluates the full -/// leading prefix. -fn select_top_candidate( - policy: GjTopPromotion, - current_estimate: usize, - min_threshold: usize, - mut candidates: impl Iterator>, - leader_keys: impl Fn(&T) -> usize, -) -> Option { - let current = candidates.next().flatten(); - match policy { - GjTopPromotion::Current => current, - GjTopPromotion::Eligible => { - if current.is_some() { - return current; - } - let promotion_limit = current_estimate.saturating_mul(4).max(min_threshold); - candidates - .flatten() - .find(|candidate| leader_keys(candidate) <= promotion_limit) - } - GjTopPromotion::Largest => { - current - .into_iter() - .chain(candidates.flatten()) - .fold(None, |best, candidate| match best { - Some(best) if leader_keys(&best) >= leader_keys(&candidate) => Some(best), - _ => Some(candidate), - }) - } - } -} - struct SparseColumnIndex { n_keys: usize, n_subsets: usize, @@ -1034,88 +998,35 @@ type ColumnIndexes = IdVec>>; // (node, col, value) with different slow constraints; they are almost always // empty, in which case the guard is a cheap length check. type ChildEntry = (Arc, Box<[Constraint]>); -type ChildLock = RwLock>; -type ChildrenMaps = IdVec>; - -/// Cache of child trie nodes for one `(TrieNode, ColumnId)` pair. -/// -/// Descendant nodes lazily initialize [`ChildMap::Single`]. An eligible coarse -/// top-index partition may initialize a root column as [`ChildMap::Sharded`] -/// before its jobs start, using the same [`ShardData`] layout as the leader's -/// column index. There is deliberately no promotion from `Single` to `Sharded`: -/// an already-used cache remains valid and simply falls back to the old layout. -pub(crate) enum ChildMap { - Single(ChildLock), - Sharded { - shard_data: ShardData, - shards: IdVec>, - }, -} - -impl ChildMap { - fn single() -> Self { - Self::Single(RwLock::new(HashMap::default())) - } - - fn sharded(shard_count: usize) -> Self { - let shard_data = ShardData::new(shard_count); - debug_assert_eq!(shard_data.n_shards(), shard_count); - let mut shards = IdVec::with_capacity(shard_data.n_shards()); - shards.resize_with(shard_data.n_shards(), || { - CachePadded::new(RwLock::new(HashMap::default())) - }); - Self::Sharded { shard_data, shards } - } - - fn is_aligned_to(&self, shard_count: usize) -> bool { - matches!( - self, - Self::Sharded { shard_data, .. } if shard_data.n_shards() == shard_count - ) - } - - /// Select the lock with the same `Value` hashing and high-bit shard mapping - /// used by [`crate::hash_index::ColumnIndex`]. - fn lock_for(&self, value: Value) -> &ChildLock { - match self { - Self::Single(lock) => lock, - Self::Sharded { shard_data, shards } => { - shard_data.get_shard(&value, shards) as &ChildLock - } - } - } - - fn get_or_insert( - &self, - value: Value, - edge_cs: &[Constraint], - sub: impl FnOnce() -> Subset, - ) -> Arc { - let map = self.lock_for(value); - // Optimistic read path: most calls are cache hits, so try a shared lock - // first. A hit is only valid when the edge constraints match. - { - let guard = map.read().unwrap(); - if let Some((node, stored_cs)) = guard.get(&value) - && &**stored_cs == edge_cs - { - return node.clone(); - } - } - // Cache miss (or constraint mismatch): acquire the write lock and insert. - let mut guard = map.write().unwrap(); +pub(crate) type ChildLock = RwLock>; +type ChildrenMaps = IdVec>; + +fn get_or_insert_child( + map: &ChildLock, + value: Value, + edge_cs: &[Constraint], + sub: impl FnOnce() -> Subset, +) -> Arc { + // Optimistic read path: most calls are cache hits, so try a shared lock + // first. A hit is only valid when the edge constraints match. + { + let guard = map.read().unwrap(); if let Some((node, stored_cs)) = guard.get(&value) && &**stored_cs == edge_cs { return node.clone(); } - // Child nodes intentionally start with no cache layout. If they are - // probed later, they lazily choose `Single` and never inherit root - // sharding from this map. - let new_node = Arc::new(TrieNode::new(sub())); - guard.insert(value, (new_node.clone(), Box::from(edge_cs))); - new_node } + // Cache miss (or constraint mismatch): acquire the write lock and insert. + let mut guard = map.write().unwrap(); + if let Some((node, stored_cs)) = guard.get(&value) + && &**stored_cs == edge_cs + { + return node.clone(); + } + let new_node = Arc::new(TrieNode::new(sub())); + guard.insert(value, (new_node.clone(), Box::from(edge_cs))); + new_node } /// Canonical signature of a trie root: the table plus its sorted header (fast) @@ -1237,11 +1148,9 @@ pub(crate) struct TrieNode { /// Any cached indexes on this subset. cached_subsets: OnceLock>, /// Cached child trie nodes, keyed first by column and then by value. Each - /// column lazily chooses a single lock, except an eligible top-level root - /// may preinitialize a physically aligned sharded map before partition jobs - /// start. When this node is a shared root (or reachable from one), this - /// cache is shared across plans too, so children are shared without any - /// global lookup. + /// column's lock is allocated lazily. When this node is a shared root (or + /// reachable from one), this cache is shared across plans too, so children + /// are shared without any global lookup. cached_children: OnceLock>, } @@ -1282,33 +1191,13 @@ impl TrieNode { .clone() } - fn child_map_slot(&self, col: ColumnId, arity: usize) -> &OnceLock { - &self.cached_children.get_or_init(|| { + fn child_map(&self, col: ColumnId, arity: usize) -> &ChildLock { + self.cached_children.get_or_init(|| { let mut vec: Pooled = with_pool_set(|ps| ps.get()); vec.resize_with(arity, OnceLock::new); vec })[col] - } - - fn child_map(&self, col: ColumnId, arity: usize) -> &ChildMap { - self.child_map_slot(col, arity) - .get_or_init(ChildMap::single) - } - - /// Initialize this node's child cache for `col` with the physical partition - /// used by the top index. If the column was already initialized, keep that - /// representation: in particular, do not promote an existing `Single` map. - /// The return value says whether the installed map has the requested - /// alignment. - fn initialize_aligned_child_map( - &self, - col: ColumnId, - arity: usize, - shard_count: usize, - ) -> bool { - self.child_map_slot(col, arity) - .get_or_init(|| ChildMap::sharded(shard_count)) - .is_aligned_to(shard_count) + .get_or_init(|| RwLock::new(HashMap::default())) } /// Return the child node reached by additionally constraining `col = value` @@ -1329,8 +1218,7 @@ impl TrieNode { info: &TableInfo, sub: impl FnOnce() -> Subset, ) -> Arc { - self.child_map(col, info.spec.arity()) - .get_or_insert(value, edge_cs, sub) + get_or_insert_child(self.child_map(col, info.spec.arity()), value, edge_cs, sub) } } @@ -1392,44 +1280,6 @@ impl BindingInfo { } } -struct TopChildCacheTarget { - node: Arc, - column: ColumnId, - arity: usize, -} - -struct TopIndexCandidate { - order_position: usize, - leader_keys: usize, - shard_count: usize, - shards: Vec, - child_cache_targets: Vec, -} - -#[derive(Clone, Copy)] -struct TopIndexEligibility { - workers: usize, - min_keys_per_worker: usize, -} - -impl TopIndexCandidate { - /// Install an aligned cache layout only after this candidate has won top - /// variable selection. Candidate discovery may initialize column indexes, - /// but must not commit cache layouts for rejected stages. - fn configure_child_cache(&self) { - if gj_child_cache() != GjChildCache::Aligned { - return; - } - for target in &self.child_cache_targets { - let _ = target.node.initialize_aligned_child_map( - target.column, - target.arity, - self.shard_count, - ); - } - } -} - impl<'a> JoinState<'a> { fn new( db: &'a Database, @@ -1608,19 +1458,18 @@ impl<'a> JoinState<'a> { /// building a dynamic nonleader index during discovery and then rebuilding /// it in every shard job. Tiny or badly skewed leaders stay on the existing /// path as well. - fn top_index_candidate( + fn top_index_shards( &self, - order_position: usize, stage: &JoinStage, prepared: &[PreparedIndexSlot], atoms: &Arc>, binding_info: &mut BindingInfo, - eligibility: TopIndexEligibility, - ) -> Option { + workers: usize, + ) -> Option> { let JoinStage::Intersect { scans, .. } = stage else { return None; }; - if scans.is_empty() || eligibility.workers <= 1 { + if scans.is_empty() || workers <= 1 { return None; } debug_assert_eq!(scans.len(), prepared.len()); @@ -1661,87 +1510,46 @@ impl<'a> JoinState<'a> { } probers.push(prober); } - let candidate = probers[leader].shard_count().and_then(|shard_count| { + let shards = probers[leader].shard_count().and_then(|shard_count| { let shards = (0..shard_count) .filter(|shard| probers[leader].shard_len(*shard).unwrap_or(0) != 0) .collect::>(); top_index_shape_is_eligible( - eligibility.workers, + workers, leader_size, shards.len(), - eligibility.min_keys_per_worker, + MIN_TOP_INDEX_KEYS_PER_WORKER, ) - .then(|| TopIndexCandidate { - order_position, - leader_keys: leader_size, - shard_count, - shards, - child_cache_targets: scans - .iter() - .zip(&probers) - .map(|(scan, prober)| TopChildCacheTarget { - node: prober.node.clone(), - column: scan.column, - arity: self.db.tables[atoms[scan.atom].table].spec.arity(), - }) - .collect(), - }) + .then_some(shards) }); for (scan, prober) in scans.iter().zip(probers) { binding_info.move_back(scan.atom, prober); } - candidate + shards } - /// Choose a top-index driver from the maximal leading run of plain - /// `Intersect` stages. Restricting promotion to that run is deliberately - /// conservative: every candidate is known to commute, so rotating one to - /// the front cannot cross a fused or materialization dependency. - fn select_top_index_candidate( + /// Return a coarse partition for the variable already sorted to the top of + /// the join. Later variables are deliberately not promoted: benchmarking + /// found that probing and reordering them was not broadly beneficial. + fn select_top_index_shards( &self, stages: &JoinStages, prepared: &PreparedJoinIndexes, atoms: &Arc>, order: &InstrOrder, binding_info: &mut BindingInfo, - ) -> Option { - let prefix_len = (0..order.len()) - .take_while(|position| { - matches!( - &stages.instrs[order.get(*position)], - JoinStage::Intersect { .. } - ) - }) - .count(); - if prefix_len == 0 { + ) -> Option> { + if order.len() == 0 { return None; } - let eligibility = TopIndexEligibility { - workers: crate::parallel::current_num_threads(), - min_keys_per_worker: gj_min_keys_per_worker(), - }; - let current_estimate = estimate_size(&stages.instrs[order.get(0)], binding_info); - let min_threshold = eligibility - .min_keys_per_worker - .saturating_mul(eligibility.workers); - let candidates = (0..prefix_len).map(|order_position| { - let stage_index = order.get(order_position); - self.top_index_candidate( - order_position, - &stages.instrs[stage_index], - prepared.stage(stage_index), - atoms, - binding_info, - eligibility, - ) - }); - select_top_candidate( - gj_top_promotion(), - current_estimate, - min_threshold, - candidates, - |candidate| candidate.leader_keys, + let stage_index = order.get(0); + self.top_index_shards( + &stages.instrs[stage_index], + prepared.stage(stage_index), + atoms, + binding_info, + crate::parallel::current_num_threads(), ) } @@ -1777,30 +1585,18 @@ impl<'a> JoinState<'a> { let mut leaf_scans: LeafScans = smallvec::smallvec![false; stages.instrs.len()]; sort_plan_by_size(&mut order, &mut leaf_scans, 0, &stages.instrs, binding_info); - // A cached top index is already partitioned by hash. Schedule one - // coarse global job per nonempty shard. The shard may schedule one - // bounded level of coalesced local packets, but every packet's subtree - // remains serial. Besides avoiding an intermediate key copy, this keeps - // related nested probes on one worker by default. Buffers that cannot - // construct an independent partition (the in-place executors) decline - // this path. - if gj_top_index_sharding() - && !stages.instrs.is_empty() + // A cached top index is already partitioned by hash. Schedule one + // coarse global job per nonempty shard and run each shard's subtree + // serially. Besides avoiding an intermediate key copy, this keeps + // related nested probes on one worker. Buffers that cannot construct an + // independent partition (the in-place executors) decline this path. + if !stages.instrs.is_empty() && action_buf.supports_global_partition() - && let Some(candidate) = - self.select_top_index_candidate(stages, prepared, atoms, &order, binding_info) + && let Some(shards) = + self.select_top_index_shards(stages, prepared, atoms, &order, binding_info) { - if candidate.order_position != 0 { - order.promote_to_front(candidate.order_position); - recompute_leaf_scans(&order, &mut leaf_scans, &stages.instrs, 0); - } - // If configured, install the same value-hash partition in the - // root child cache as the selected leader before any shard job can - // touch these nodes. Rejected candidates never commit a cache - // layout. - candidate.configure_child_cache(); let mut updates = FrameUpdates::with_capacity(0); - for shard in candidate.shards { + for shard in shards { let db = self.db; let exec_state_for_factory = self.exec_state.clone(); let exec_state_for_work = self.exec_state.clone(); @@ -1859,12 +1655,10 @@ impl<'a> JoinState<'a> { /// The core method for executing a free join plan. /// - /// This method takes the plan, mutable data-structures for variable binding and staging - /// actions, and two indexes: `cur` which is the current stage of the plan to run, and `level` - /// which is the current "fan-out" node we are in. The latter parameter is an experimental - /// index used to detect if we are at the "top" of a plan rather than the "bottom", and is - /// currently used as a heuristic to determine if we should increase parallelism more than the - /// default. + /// This method takes the plan, mutable data structures for variable binding + /// and staging actions, and `cur`, the current stage of the plan. A top-level + /// coarse partition also passes `index_shard`; recursive calls clear it so + /// only the first intersection is restricted to that physical shard. #[allow(clippy::too_many_arguments)] fn run_plan<'buf, A: NumericId + 'buf, BUF: ActionBuffer<'buf, A>>( &self, @@ -1911,12 +1705,9 @@ impl<'a> JoinState<'a> { } // TODO: `supports_parallel_drain`` is a hack because currently // `drain_updates_parallel!`` is a bit slower because of the additional ExecutionState clone. - if (index_shard.is_some() - && cur == 0 - && action_buf.supports_local_partition_drain()) - || (index_shard.is_none() - && cur < free_join_fork_depth() - && action_buf.supports_parallel_drain()) + if index_shard.is_none() + && cur < free_join_fork_depth() + && action_buf.supports_parallel_drain() { drain_updates_parallel!($updates) } else { @@ -2892,16 +2683,6 @@ trait ActionBuffer<'state, A: NumericId>: Send { true } - /// Whether a coarse top-index partition may schedule its `cur = 0` - /// `FrameUpdates` packets on the current worker's local queue. - /// - /// This is separate from [`Self::supports_parallel_drain`] so a partition - /// can explicitly permit one local level while its packet buffer remains - /// fully serial. - fn supports_local_partition_drain(&self) -> bool { - false - } - /// Whether this buffer can enqueue independent top-level partitions. fn supports_global_partition(&self) -> bool { false @@ -2984,50 +2765,9 @@ impl<'a, 'outer: 'a> ActionBuffer<'a, ActionId> for InPlaceActionBuffer<'outer> } } -#[derive(Clone)] -struct LocalPacketLimiter { - outstanding: Arc, - limit: usize, -} - -impl LocalPacketLimiter { - fn new() -> Self { - Self::with_limit(gj_local_queue_limit()) - } - - fn with_limit(limit: usize) -> Self { - Self { - outstanding: Arc::new(AtomicUsize::new(0)), - limit, - } - } - - fn try_acquire(&self) -> Option { - self.outstanding - .fetch_update(Ordering::AcqRel, Ordering::Acquire, |outstanding| { - (outstanding < self.limit).then_some(outstanding + 1) - }) - .ok() - .map(|_| LocalPacketPermit { - outstanding: self.outstanding.clone(), - }) - } -} - -struct LocalPacketPermit { - outstanding: Arc, -} - -impl Drop for LocalPacketPermit { - fn drop(&mut self) { - let previous = self.outstanding.fetch_sub(1, Ordering::AcqRel); - debug_assert!(previous > 0); - } -} - -/// Strictly serial action buffer used inside one locally scheduled packet. -/// It shares the rule-set match counter with sibling packets, but executes all -/// recursive join and action work inline and cannot schedule grandchildren. +/// Strictly serial action buffer used inside one globally scheduled top-index +/// shard. It shares the rule-set match counter with sibling shards, but executes +/// all recursive join and action work inline. struct SerialScopedActionBuffer<'scope> { rule_set: &'scope RuleSet, match_counter: Arc, @@ -3113,104 +2853,6 @@ impl<'scope> ActionBuffer<'scope, ActionId> for SerialScopedActionBuffer<'scope> } } -/// Buffer owned by one globally scheduled top-index shard. It may fork -/// coalesced `cur = 0` packets onto that worker's local deque, but every packet -/// receives a [`SerialScopedActionBuffer`] and therefore runs the entire -/// `cur = 1` subtree and its actions serially. -struct TopPartitionActionBuffer<'inner, 'scope> { - scope: &'inner Scope<'scope>, - serial: SerialScopedActionBuffer<'scope>, - limiter: LocalPacketLimiter, -} - -impl<'inner, 'scope> TopPartitionActionBuffer<'inner, 'scope> { - fn new( - scope: &'inner Scope<'scope>, - rule_set: &'scope RuleSet, - match_counter: Arc, - ) -> Self { - Self { - scope, - serial: SerialScopedActionBuffer::new(rule_set, match_counter), - limiter: LocalPacketLimiter::new(), - } - } -} - -impl<'scope> ActionBuffer<'scope, ActionId> for TopPartitionActionBuffer<'_, 'scope> { - type AsLocal<'a> - = SerialScopedActionBuffer<'scope> - where - 'scope: 'a; - type AsGlobalSerial<'a> - = Self - where - 'scope: 'a; - - fn push_bindings( - &mut self, - action: ActionId, - bindings: &DenseIdMap, - to_exec_state: impl FnMut() -> ExecutionState<'scope>, - ) { - self.serial.push_bindings(action, bindings, to_exec_state); - } - - fn flush(&mut self, exec_state: &mut ExecutionState) { - self.serial.flush(exec_state); - } - - fn recur<'local>( - &mut self, - mut local: BorrowedLocalState<'local>, - mut to_exec_state: impl FnMut() -> ExecutionState<'scope> + Send + 'scope, - work: impl for<'a> FnOnce(BorrowedLocalState<'a>, &mut SerialScopedActionBuffer<'scope>) - + Send - + 'scope, - ) { - let rule_set = self.serial.rule_set; - let match_counter = self.serial.match_counter.clone(); - if let Some(permit) = self.limiter.try_acquire() { - let mut inner = local.clone_state(); - self.scope.spawn_local(move |_| { - let _permit = permit; - let mut buf = SerialScopedActionBuffer::new(rule_set, match_counter); - work(inner.borrow_mut(), &mut buf); - buf.flush(&mut to_exec_state()); - }); - } else { - let mut buf = SerialScopedActionBuffer::new(rule_set, match_counter); - work(local, &mut buf); - buf.flush(&mut to_exec_state()); - } - } - - fn recur_global_serial<'local>( - &mut self, - local: BorrowedLocalState<'local>, - _to_exec_state: impl FnMut() -> ExecutionState<'scope> + Send + 'scope, - work: impl for<'a> FnOnce(BorrowedLocalState<'a>, &mut Self) + Send + 'scope, - ) { - work(local, self); - } - - fn morsel_size(&mut self, level: usize, _total: usize) -> usize { - if level == 0 && gj_local_depth() == 1 { - gj_local_morsel() - } else { - 256 - } - } - - fn supports_parallel_drain(&self) -> bool { - false - } - - fn supports_local_partition_drain(&self) -> bool { - gj_local_depth() == 1 - } -} - /// An action buffer that hands off batches of actions to scoped worker tasks. struct ScopedActionBuffer<'inner, 'scope> { scope: &'inner Scope<'scope>, @@ -3242,7 +2884,7 @@ impl<'scope> ActionBuffer<'scope, ActionId> for ScopedActionBuffer<'_, 'scope> { where 'scope: 'a; type AsGlobalSerial<'a> - = TopPartitionActionBuffer<'a, 'scope> + = SerialScopedActionBuffer<'scope> where 'scope: 'a; fn push_bindings( @@ -3323,15 +2965,15 @@ impl<'scope> ActionBuffer<'scope, ActionId> for ScopedActionBuffer<'_, 'scope> { &mut self, mut local: BorrowedLocalState<'local>, mut to_exec_state: impl FnMut() -> ExecutionState<'scope> + Send + 'scope, - work: impl for<'a> FnOnce(BorrowedLocalState<'a>, &mut TopPartitionActionBuffer<'a, 'scope>) + work: impl for<'a> FnOnce(BorrowedLocalState<'a>, &mut SerialScopedActionBuffer<'scope>) + Send + 'scope, ) { let rule_set = self.rule_set; let match_counter = self.match_counter.clone(); let mut inner = local.clone_state(); - self.scope.spawn_global(move |scope| { - let mut buf = TopPartitionActionBuffer::new(scope, rule_set, match_counter); + self.scope.spawn_global(move |_| { + let mut buf = SerialScopedActionBuffer::new(rule_set, match_counter); work(inner.borrow_mut(), &mut buf); buf.flush(&mut to_exec_state()); }); @@ -3573,99 +3215,6 @@ impl<'scope> ActionBuffer<'scope, MatId> for SerialScopedMaterializer { } } -/// Materialization counterpart to [`TopPartitionActionBuffer`]. Only the -/// `cur = 0` packet boundary may use the worker-local queue; every packet gets -/// a serial materializer and cannot spawn deeper work. -struct TopPartitionMaterializer<'inner, 'scope> { - scope: &'inner Scope<'scope>, - serial: SerialScopedMaterializer, - limiter: LocalPacketLimiter, -} - -impl<'inner, 'scope> TopPartitionMaterializer<'inner, 'scope> { - fn new( - scope: &'inner Scope<'scope>, - specs: Arc>, - materializations: Arc, RowBuffer>>>>, - ) -> Self { - Self { - scope, - serial: SerialScopedMaterializer::new(specs, materializations), - limiter: LocalPacketLimiter::new(), - } - } -} - -impl<'scope> ActionBuffer<'scope, MatId> for TopPartitionMaterializer<'_, 'scope> { - type AsLocal<'a> - = SerialScopedMaterializer - where - 'scope: 'a; - type AsGlobalSerial<'a> - = Self - where - 'scope: 'a; - - fn push_bindings( - &mut self, - mat_id: MatId, - bindings: &DenseIdMap, - to_exec_state: impl FnMut() -> ExecutionState<'scope>, - ) { - self.serial.push_bindings(mat_id, bindings, to_exec_state); - } - - fn flush(&mut self, exec_state: &mut ExecutionState) { - self.serial.flush(exec_state); - } - - fn recur<'local>( - &mut self, - mut local: BorrowedLocalState<'local>, - _to_exec_state: impl FnMut() -> ExecutionState<'scope> + Send + 'scope, - work: impl for<'a> FnOnce(BorrowedLocalState<'a>, &mut SerialScopedMaterializer) + Send + 'scope, - ) { - let specs = self.serial.specs.clone(); - let materializations = self.serial.materializations.clone(); - if let Some(permit) = self.limiter.try_acquire() { - let mut inner = local.clone_state(); - self.scope.spawn_local(move |_| { - let _permit = permit; - let mut buf = SerialScopedMaterializer::new(specs, materializations); - work(inner.borrow_mut(), &mut buf); - }); - } else { - let mut buf = SerialScopedMaterializer::new(specs, materializations); - work(local, &mut buf); - } - } - - fn recur_global_serial<'local>( - &mut self, - local: BorrowedLocalState<'local>, - _to_exec_state: impl FnMut() -> ExecutionState<'scope> + Send + 'scope, - work: impl for<'a> FnOnce(BorrowedLocalState<'a>, &mut Self) + Send + 'scope, - ) { - work(local, self); - } - - fn morsel_size(&mut self, level: usize, _total: usize) -> usize { - if level == 0 && gj_local_depth() == 1 { - gj_local_morsel() - } else { - 256 - } - } - - fn supports_parallel_drain(&self) -> bool { - false - } - - fn supports_local_partition_drain(&self) -> bool { - gj_local_depth() == 1 - } -} - struct ScopedMaterializer<'inner, 'scope> { scope: &'inner Scope<'scope>, specs: Arc>, @@ -3679,7 +3228,7 @@ impl<'scope> ActionBuffer<'scope, MatId> for ScopedMaterializer<'_, 'scope> { where 'scope: 'a; type AsGlobalSerial<'a> - = TopPartitionMaterializer<'a, 'scope> + = SerialScopedMaterializer where 'scope: 'a; @@ -3749,15 +3298,13 @@ impl<'scope> ActionBuffer<'scope, MatId> for ScopedMaterializer<'_, 'scope> { &mut self, mut local: BorrowedLocalState<'local>, _to_exec_state: impl FnMut() -> ExecutionState<'scope> + Send + 'scope, - work: impl for<'a> FnOnce(BorrowedLocalState<'a>, &mut TopPartitionMaterializer<'a, 'scope>) - + Send - + 'scope, + work: impl for<'a> FnOnce(BorrowedLocalState<'a>, &mut SerialScopedMaterializer) + Send + 'scope, ) { let specs = self.specs.clone(); let materializations = self.materializations.clone(); let mut inner = local.clone_state(); - self.scope.spawn_global(move |scope| { - let mut buf = TopPartitionMaterializer::new(scope, specs, materializations); + self.scope.spawn_global(move |_| { + let mut buf = SerialScopedMaterializer::new(specs, materializations); work(inner.borrow_mut(), &mut buf); }); } @@ -4052,12 +3599,6 @@ impl InstrOrder { fn len(&self) -> usize { self.data.len() } - - /// Move the stage at `position` to the front while preserving the relative - /// order of every stage it passes. - fn promote_to_front(&mut self, position: usize) { - self.data[..=position].rotate_right(1); - } } /// Per-position leaf-scan flags. `leaf_scans[i] == true` means the stage currently scheduled at @@ -4102,10 +3643,8 @@ impl LocalState { } #[cfg(test)] -mod local_packet_tests { +mod top_index_tests { use std::{ - cell::Cell, - ptr, sync::{ Arc, Barrier, atomic::{AtomicUsize, Ordering}, @@ -4113,15 +3652,9 @@ mod local_packet_tests { thread, }; - use crate::{ - common::Value, numeric_id::NumericId, offsets::Subset, parallel_heuristics::GjTopPromotion, - table_spec::ColumnId, - }; + use crate::{common::Value, numeric_id::NumericId, offsets::Subset}; - use super::{ - ChildMap, InstrOrder, LocalPacketLimiter, TrieNode, select_top_candidate, - top_index_shape_is_eligible, - }; + use super::{ChildLock, get_or_insert_child, top_index_shape_is_eligible}; #[test] fn top_index_partitioning_rejects_serial_tiny_and_skewed_shapes() { @@ -4133,102 +3666,9 @@ mod local_packet_tests { } #[test] - fn current_top_policy_probes_only_sorted_position_zero() { - let probes = Cell::new(0); - let candidates = [None, Some((1, 512))].into_iter().inspect(|_| { - probes.set(probes.get() + 1); - }); - let selected = - select_top_candidate(GjTopPromotion::Current, 100, 256, candidates, |candidate| { - candidate.1 - }); - assert_eq!(selected, None); - assert_eq!(probes.get(), 1); - } - - #[test] - fn eligible_top_policy_keeps_current_or_uses_first_guarded_fallback() { - let probes = Cell::new(0); - let candidates = [Some((0, 256)), Some((1, 512))].into_iter().inspect(|_| { - probes.set(probes.get() + 1); - }); - assert_eq!( - select_top_candidate( - GjTopPromotion::Eligible, - 100, - 256, - candidates, - |candidate| candidate.1, - ), - Some((0, 256)) - ); - assert_eq!(probes.get(), 1); - - // max(4 * 100, 256) = 400: skip the first eligible-but-too-large - // fallback and preserve the first later candidate within the guard. - assert_eq!( - select_top_candidate( - GjTopPromotion::Eligible, - 100, - 256, - [None, Some((1, 401)), Some((2, 300))].into_iter(), - |candidate| candidate.1, - ), - Some((2, 300)) - ); - } - - #[test] - fn largest_top_policy_uses_actual_keys_and_keeps_earliest_tie() { - assert_eq!( - select_top_candidate( - GjTopPromotion::Largest, - 1, - 256, - [Some((0, 256)), None, Some((2, 1024)), Some((3, 1024))].into_iter(), - |candidate| candidate.1, - ), - Some((2, 1024)) - ); - } - - #[test] - fn top_promotion_preserves_passed_stage_order() { - let mut order = InstrOrder::from_iter([3, 1, 4, 2].into_iter()); - order.promote_to_front(2); - assert_eq!( - (0..order.len()).map(|i| order.get(i)).collect::>(), - vec![4, 3, 1, 2] - ); - } - - #[test] - fn local_packet_limiter_bounds_outstanding_work_and_reuses_slots() { - let limiter = LocalPacketLimiter::with_limit(2); - let first = limiter.try_acquire().expect("first packet should fit"); - let second = limiter.try_acquire().expect("second packet should fit"); - assert!(limiter.try_acquire().is_none()); - - drop(first); - let replacement = limiter - .try_acquire() - .expect("completed packet should release its slot"); - assert!(limiter.try_acquire().is_none()); - - drop((second, replacement)); - assert!(limiter.try_acquire().is_some()); - } - - #[test] - fn zero_local_packet_limit_always_executes_inline() { - let limiter = LocalPacketLimiter::with_limit(0); - assert!(limiter.try_acquire().is_none()); - } - - #[test] - fn sharded_child_cache_returns_one_arc_for_racing_same_key() { + fn child_cache_returns_one_arc_for_racing_same_key() { const THREADS: usize = 16; - let map = Arc::new(ChildMap::sharded(8)); + let map = Arc::new(ChildLock::default()); let barrier = Arc::new(Barrier::new(THREADS)); let constructed = Arc::new(AtomicUsize::new(0)); let value = Value::from_usize(17); @@ -4240,7 +3680,7 @@ mod local_packet_tests { let constructed = constructed.clone(); thread::spawn(move || { barrier.wait(); - map.get_or_insert(value, &[], || { + get_or_insert_child(&map, value, &[], || { constructed.fetch_add(1, Ordering::Relaxed); Subset::empty() }) @@ -4255,64 +3695,4 @@ mod local_packet_tests { assert_eq!(constructed.load(Ordering::Relaxed), 1); assert!(nodes[1..].iter().all(|node| Arc::ptr_eq(&nodes[0], node))); } - - #[test] - fn sharded_child_cache_maps_distinct_values_to_every_aligned_lock() { - let map = Arc::new(ChildMap::sharded(8)); - let ChildMap::Sharded { shard_data, shards } = map.as_ref() else { - unreachable!() - }; - let mut values = vec![None; shard_data.n_shards()]; - for raw in 0..100_000 { - let value = Value::from_usize(raw); - let selected = shard_data.get_shard(&value, shards); - let shard = shards - .iter() - .find_map(|(shard, lock)| ptr::eq(lock, selected).then_some(shard.index())) - .unwrap(); - values[shard].get_or_insert(value); - if values.iter().all(Option::is_some) { - break; - } - } - assert!(values.iter().all(Option::is_some)); - - let handles = values - .into_iter() - .map(|value| { - let map = map.clone(); - thread::spawn(move || { - map.get_or_insert(value.unwrap(), &[], Subset::empty); - }) - }) - .collect::>(); - for handle in handles { - handle.join().unwrap(); - } - - let ChildMap::Sharded { shards, .. } = map.as_ref() else { - unreachable!() - }; - for (_, shard) in shards.iter() { - assert_eq!(shard.read().unwrap().len(), 1); - } - } - - #[test] - fn aligned_root_children_default_to_single_and_single_is_not_promoted() { - let col = ColumnId::new(0); - let root = TrieNode::new(Subset::empty()); - assert!(root.initialize_aligned_child_map(col, 1, 8)); - assert!(matches!(root.child_map(col, 1), ChildMap::Sharded { .. })); - assert!(!root.initialize_aligned_child_map(col, 1, 16)); - - let child = root - .child_map(col, 1) - .get_or_insert(Value::from_usize(3), &[], Subset::empty); - assert!(child.cached_children.get().is_none()); - assert!(matches!(child.child_map(col, 1), ChildMap::Single(_))); - - assert!(!child.initialize_aligned_child_map(col, 1, 8)); - assert!(matches!(child.child_map(col, 1), ChildMap::Single(_))); - } } diff --git a/core-relations/src/parallel_heuristics.rs b/core-relations/src/parallel_heuristics.rs index ef2de53af..e52bf3cf9 100644 --- a/core-relations/src/parallel_heuristics.rs +++ b/core-relations/src/parallel_heuristics.rs @@ -14,44 +14,9 @@ const DEFAULT_TABLE_OP_CUTOFF: usize = 400_000; const DEFAULT_FREE_JOIN_FORK_DEPTH: usize = 2; const DEFAULT_ACTION_BATCH_SIZE: usize = 8 * 1024; -// Fixed generic-join scheduling choices selected by cross-workload -// benchmarking. Coarse index partitioning is enabled. Recursive fallback work -// uses worker-local queues, but already-sharded partitions do not create an -// additional packet level. Aligned child caches and top-variable promotion -// remain disabled. -const GJ_TOP_INDEX_SHARDING: bool = true; -const GJ_LOCAL_DEPTH: usize = 0; -const GJ_LOCAL_MORSEL: usize = 64; -const GJ_LOCAL_QUEUE_LIMIT: usize = 2; -const GJ_CHILD_CACHE: GjChildCache = GjChildCache::Single; -const GJ_TOP_PROMOTION: GjTopPromotion = GjTopPromotion::Current; -const GJ_MIN_KEYS_PER_WORKER: usize = 16; - -/// Layout for root trie-node child caches during coarse generic-join -/// partitioning. -#[derive(Clone, Copy, Debug, Eq, PartialEq)] -pub(crate) enum GjChildCache { - /// One lock per cached column. - Single, - /// One lock per top-index shard, using the same value hash partition. - Aligned, -} - -/// Policy for replacing the initially sorted top generic-join variable with a -/// coarse-partitionable later variable. -// The non-selected policies remain available to the executor's deterministic -// policy tests even though production uses the fixed `Current` policy. -#[cfg_attr(not(test), allow(dead_code))] -#[derive(Clone, Copy, Debug, Eq, PartialEq)] -pub(crate) enum GjTopPromotion { - /// Consider only the variable already sorted into position zero. - Current, - /// Keep an eligible current variable, otherwise choose the first guarded - /// eligible variable in the leading Intersect prefix. - Eligible, - /// Choose the eligible variable with the largest actual leader cardinality. - Largest, -} +/// Minimum number of top-index leader keys required per worker before coarse +/// generic-join partitioning is worthwhile. +pub(crate) const MIN_TOP_INDEX_KEYS_PER_WORKER: usize = 16; static CUTOFFS: OnceLock = OnceLock::new(); @@ -109,44 +74,6 @@ pub(crate) fn action_batch_size() -> usize { cutoffs().action_batch_size } -/// Whether a parallel generic join may schedule one coarse global job per -/// physical shard of its top cached index. -pub(crate) fn gj_top_index_sharding() -> bool { - GJ_TOP_INDEX_SHARDING -} - -/// Number of recursive levels below a coarse top-index shard that may enqueue -/// worker-local packets. The implementation currently accepts only zero or one. -pub(crate) fn gj_local_depth() -> usize { - GJ_LOCAL_DEPTH -} - -/// Number of top-level generic-join frames coalesced into one local packet. -pub(crate) fn gj_local_morsel() -> usize { - GJ_LOCAL_MORSEL -} - -/// Maximum number of outstanding local packets per coarse index shard. -pub(crate) fn gj_local_queue_limit() -> usize { - GJ_LOCAL_QUEUE_LIMIT -} - -/// Child-cache layout used by eligible coarse top-index partitions. -pub(crate) fn gj_child_cache() -> GjChildCache { - GJ_CHILD_CACHE -} - -/// Policy for promoting an eligible top generic-join variable. -pub(crate) fn gj_top_promotion() -> GjTopPromotion { - GJ_TOP_PROMOTION -} - -/// Minimum number of leader keys required per worker for coarse top-index -/// partitioning. -pub(crate) fn gj_min_keys_per_worker() -> usize { - GJ_MIN_KEYS_PER_WORKER -} - fn should_parallelize(len: usize, cutoff: usize) -> bool { len > cutoff && crate::parallel::current_num_threads() > 1 } diff --git a/core-relations/src/pool/mod.rs b/core-relations/src/pool/mod.rs index 40123c22f..0695cdddf 100644 --- a/core-relations/src/pool/mod.rs +++ b/core-relations/src/pool/mod.rs @@ -12,7 +12,7 @@ use std::{ use crate::{ AtomId, - free_join::execute::ChildMap, + free_join::execute::ChildLock, numeric_id::{DenseIdMap, IdVec}, }; use fixedbitset::FixedBitSet; @@ -442,7 +442,7 @@ pool_set! { cached_subsets: IdVec>> [ 4 << 20 ], intersected_on: DenseIdMap [ 1 << 20 ], - cached_child: IdVec> [ 1 << 20 ], + cached_child: IdVec> [ 1 << 20 ], } } diff --git a/core-relations/src/tests.rs b/core-relations/src/tests.rs index 04c747221..9cb9bc6d7 100644 --- a/core-relations/src/tests.rs +++ b/core-relations/src/tests.rs @@ -475,17 +475,6 @@ fn prepared_plan_indexes_preserve_uncacheable_columns() { #[test] fn gj_top_index_shards_preserve_count_and_materialization() { - assert!(crate::parallel_heuristics::gj_top_index_sharding()); - assert_eq!(crate::parallel_heuristics::gj_local_depth(), 0); - assert_eq!( - crate::parallel_heuristics::gj_child_cache(), - crate::parallel_heuristics::GjChildCache::Single - ); - assert_eq!( - crate::parallel_heuristics::gj_top_promotion(), - crate::parallel_heuristics::GjTopPromotion::Current - ); - assert_eq!(crate::parallel_heuristics::gj_min_keys_per_worker(), 16); gj_top_index_shards_preserve_count_and_materialization_inner(); } From 2a8e4693b2154647bab725a91534c7e427801837 Mon Sep 17 00:00:00 2001 From: Eli Rosenthal Date: Thu, 30 Jul 2026 09:35:01 -0700 Subject: [PATCH 3/5] Make deletion hash casts explicit --- core-relations/src/table/mod.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/core-relations/src/table/mod.rs b/core-relations/src/table/mod.rs index 8a520f4ef..dbfaca9e0 100644 --- a/core-relations/src/table/mod.rs +++ b/core-relations/src/table/mod.rs @@ -627,7 +627,7 @@ impl SortedWritesTable { let (actual_shard, hc) = hash_code(shard_data, to_remove, n_keys); assert_eq!(actual_shard, shard_id); if let Ok(entry) = shard.find_entry(hc, |entry| { - entry.hashcode == (hc as _) + entry.hashcode == (hc as HashCode) && &data.get_row(entry.row)[0..n_keys] == to_remove }) { let (ent, _) = entry.remove(); @@ -665,7 +665,7 @@ impl SortedWritesTable { let (actual_shard, hc) = hash_code(shard_data, to_remove, self.n_keys); assert_eq!(actual_shard, shard_id); if let Ok(entry) = shard.find_entry(hc, |entry| { - entry.hashcode == (hc as _) + entry.hashcode == (hc as HashCode) && &self.data.get_row(entry.row).unwrap()[0..self.n_keys] == to_remove }) { From 0b0c54632046331b412115f8401e9d0fea4d22a3 Mon Sep 17 00:00:00 2001 From: Eli Rosenthal Date: Sun, 9 Aug 2026 22:45:34 -0700 Subject: [PATCH 4/5] Restore logical-order DVO heuristic --- core-relations/src/free_join/execute.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/core-relations/src/free_join/execute.rs b/core-relations/src/free_join/execute.rs index 8f32d5419..c6fa40026 100644 --- a/core-relations/src/free_join/execute.rs +++ b/core-relations/src/free_join/execute.rs @@ -3363,7 +3363,7 @@ fn sort_plan_by_size( let mut last_pos = start; for i in start..instrs.len() { if matches!( - &instrs[order.get(i)], + &instrs[i], // These nodes don't commute JoinStage::FusedIntersectMat { mode: MatScanMode::Lookup(_) | MatScanMode::Value(_) | MatScanMode::Full, @@ -3480,8 +3480,8 @@ fn sort_plan_by_size_inner( let mut times_refined = with_pool_set(|ps| ps.get::>()); // Count how many times each atom has been refined so far. - for position in 0..range.start { - match &instrs[order.get(position)] { + for ins in &instrs[..range.start] { + match ins { JoinStage::Intersect { scans, .. } => scans.iter().for_each(|scan| { *times_refined.get_or_default(scan.atom) += 1; }), From e4ebef25b5d35f54c52b81bd73a7ada817fa718e Mon Sep 17 00:00:00 2001 From: Eli Rosenthal Date: Sun, 9 Aug 2026 22:59:09 -0700 Subject: [PATCH 5/5] Address parallel join review feedback --- core-relations/src/free_join/execute.rs | 43 ++++++++++++++----------- 1 file changed, 25 insertions(+), 18 deletions(-) diff --git a/core-relations/src/free_join/execute.rs b/core-relations/src/free_join/execute.rs index c6fa40026..35b38873a 100644 --- a/core-relations/src/free_join/execute.rs +++ b/core-relations/src/free_join/execute.rs @@ -664,7 +664,7 @@ impl Database { // immutable plans before any worker starts. A slot acquires and refreshes // its Arc through the regular catalog helper on first cached use. The // sidecars are dropped before `merge_all` resets catalog entries. - let prepared_plans = rule_set + let prepared_indexes = rule_set .plans .values() .map(|(plan, _, _)| PreparedPlanIndexes::new(self, plan)) @@ -677,8 +677,8 @@ impl Database { Arc::new(DashMap::default()); let db: &Database = self; egglog_concurrency::scope(|scope| { - for ((plan, desc, symbol_map), prepared_plan) in - rule_set.plans.values().zip(&prepared_plans) + for ((plan, desc, symbol_map), prepared_index) in + rule_set.plans.values().zip(&prepared_indexes) { // TODO: add stats let report_plan = match report_level { @@ -710,7 +710,7 @@ impl Database { } } - match (plan, prepared_plan) { + match (plan, prepared_index) { (Plan::SinglePlan(plan), PreparedPlanIndexes::Single(prepared)) => { join_state.run_join_stages( &plan.stages, @@ -828,8 +828,8 @@ impl Database { match_counter: match_counter.as_ref(), batches: Default::default(), }; - for ((plan, desc, symbol_map), prepared_plan) in - rule_set.plans.values().zip(&prepared_plans) + for ((plan, desc, symbol_map), prepared_index) in + rule_set.plans.values().zip(&prepared_indexes) { let report_plan = match report_level { ReportLevel::TimeOnly => None, @@ -849,7 +849,7 @@ impl Database { None => break 'eval, } } - match (plan, prepared_plan) { + match (plan, prepared_index) { (Plan::SinglePlan(plan), PreparedPlanIndexes::Single(prepared)) => { join_state.run_join_stages( &plan.stages, @@ -947,7 +947,7 @@ impl Database { // `merge_all` requires each catalog Arc to be uniquely owned so it can // reset the corresponding ResettableOnceLock. No scoped worker can // survive to this point. - drop(prepared_plans); + drop(prepared_indexes); let search_and_apply_time = search_and_apply_timer.elapsed(); let merge_timer = Instant::now(); @@ -1501,10 +1501,7 @@ impl<'a> JoinState<'a> { let prober = self.get_column_index(atoms, binding_info, scan.atom, scan.column, prepared); let size = prober.len(); - // The two-way hot path historically breaks ties in favor of the - // second scan. Match it so partition discovery and execution use - // the same physical index. - if size < leader_size || (scans.len() == 2 && size == leader_size) { + if size < leader_size { leader = i; leader_size = size; } @@ -1588,8 +1585,11 @@ impl<'a> JoinState<'a> { // A cached top index is already partitioned by hash. Schedule one // coarse global job per nonempty shard and run each shard's subtree // serially. Besides avoiding an intermediate key copy, this keeps - // related nested probes on one worker. Buffers that cannot construct an - // independent partition (the in-place executors) decline this path. + // related nested probes on one worker. Selecting this path commits the + // whole subtree to coarse-only parallelism: `recur_global_serial` + // supplies a serial buffer whose recursive work stays inline. Buffers + // that cannot construct an independent partition (the in-place + // executors) decline this path. if !stages.instrs.is_empty() && action_buf.supports_global_partition() && let Some(shards) = @@ -1657,8 +1657,11 @@ impl<'a> JoinState<'a> { /// /// This method takes the plan, mutable data structures for variable binding /// and staging actions, and `cur`, the current stage of the plan. A top-level - /// coarse partition also passes `index_shard`; recursive calls clear it so - /// only the first intersection is restricted to that physical shard. + /// coarse partition also passes `index_shard`; `Some` is only supplied with + /// a serial action buffer, so no nested parallelism is available in that + /// subtree. Recursive calls clear the value so only the first intersection + /// is restricted to that physical shard, while the serial buffer continues + /// to keep later work inline. #[allow(clippy::too_many_arguments)] fn run_plan<'buf, A: NumericId + 'buf, BUF: ActionBuffer<'buf, A>>( &self, @@ -1924,7 +1927,7 @@ impl<'a> JoinState<'a> { ); let ((smaller, smaller_scan), (larger, larger_scan)) = - if a_prober.len() < b_prober.len() { + if a_prober.len() <= b_prober.len() { ((&a_prober, a), (&b_prober, b)) } else { ((&b_prober, b), (&a_prober, a)) @@ -2767,7 +2770,9 @@ impl<'a, 'outer: 'a> ActionBuffer<'a, ActionId> for InPlaceActionBuffer<'outer> /// Strictly serial action buffer used inside one globally scheduled top-index /// shard. It shares the rule-set match counter with sibling shards, but executes -/// all recursive join and action work inline. +/// all recursive join and action work inline. It deliberately has no +/// `needs_flush` flag: its owner unconditionally flushes it once after the shard +/// callback returns, and recursive calls reuse the same buffer synchronously. struct SerialScopedActionBuffer<'scope> { rule_set: &'scope RuleSet, match_counter: Arc, @@ -2975,6 +2980,8 @@ impl<'scope> ActionBuffer<'scope, ActionId> for ScopedActionBuffer<'_, 'scope> { self.scope.spawn_global(move |_| { let mut buf = SerialScopedActionBuffer::new(rule_set, match_counter); work(inner.borrow_mut(), &mut buf); + // Unlike a nested `ScopedActionBuffer`, this buffer is owned by the + // one shard job and is always flushed exactly once before it exits. buf.flush(&mut to_exec_state()); }); }