diff --git a/Cargo.lock b/Cargo.lock index a1b0dbd..04baa67 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -204,6 +204,17 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b05b61dc5112cbb17e4b6cd61790d9845d13888356391624cbe7e41efeac1e75" +[[package]] +name = "core_affinity" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a034b3a7b624016c6e13f5df875747cc25f884156aad2abd12b6c46797971342" +dependencies = [ + "libc", + "num_cpus", + "winapi", +] + [[package]] name = "criterion" version = "0.5.1" @@ -240,6 +251,15 @@ dependencies = [ "itertools", ] +[[package]] +name = "crossbeam-channel" +version = "0.5.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "82b8f8f868b36967f9606790d1903570de9ceaf870a7bf9fbbd3016d636a2cb2" +dependencies = [ + "crossbeam-utils", +] + [[package]] name = "crossbeam-deque" version = "0.8.6" @@ -474,6 +494,16 @@ dependencies = [ "autocfg", ] +[[package]] +name = "num_cpus" +version = "1.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91df4bbde75afed763b708b7eee1e8e7651e02d97f6d5dd763e89367e957b23b" +dependencies = [ + "hermit-abi", + "libc", +] + [[package]] name = "once_cell" version = "1.21.3" @@ -850,7 +880,9 @@ version = "0.1.0" dependencies = [ "bytes", "clap", + "core_affinity", "criterion", + "crossbeam-channel", "pcap", "prost", "prost-build", diff --git a/Cargo.toml b/Cargo.toml index 2de61d3..4d35056 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -14,6 +14,8 @@ clap = { version = "4.0", features = ["derive"] } smallvec = "1.13.2" prost = "0.13" bytes = "1" +core_affinity = "0.8" +crossbeam-channel = "0.5" [dev-dependencies] criterion = "0.5" @@ -63,6 +65,11 @@ harness = false name = "median_bench" harness = false +[[bench]] +name = "octo_speedup" +harness = false +path = "benches/octo_speedup.rs" + #[profile.release] #debug = true #strip = false diff --git a/benches/octo_speedup.rs b/benches/octo_speedup.rs new file mode 100644 index 0000000..fcf7114 --- /dev/null +++ b/benches/octo_speedup.rs @@ -0,0 +1,595 @@ +use criterion::{BenchmarkId, Criterion, Throughput, black_box, criterion_group, criterion_main}; +use crossbeam_channel::{TryRecvError, bounded}; +use rand::{Rng, SeedableRng, rngs::StdRng}; +use sketchlib_rust::{ + CmDelta, Count, CountMin, CountOctoAggregator, CountOctoWorker, HllOctoAggregator, + HllOctoWorker, HyperLogLog, OctoAggregator, OctoConfig, OctoWorker, Regular, RegularPath, + SketchInput, Vector2D, +}; +use std::sync::{Once, mpsc}; +use std::thread; +use std::time::Duration; + +const RNG_SEED: u64 = 0x0c70_2026_5eed_1234; +const CM_COUNT_ROWS: usize = 3; +const CM_COUNT_COLS: usize = 65536; +const CM_COUNT_INPUTS: usize = 1_000_000; +const HLL_INPUTS: usize = 500_000; +const DOMAIN_MASK: u64 = (1 << 20) - 1; +const QUEUE_CAPACITY: usize = 65_536; +const WORKER_COUNTS: [usize; 4] = [1, 2, 4, 8]; +const MERGE_INTERVAL: usize = 10_000; + +static SANITY_ONCE: Once = Once::new(); + +struct BenchCmOctoWorker { + sketch: CountMin, RegularPath>, +} + +impl BenchCmOctoWorker { + fn new(rows: usize, cols: usize) -> Self { + Self { + sketch: CountMin::with_dimensions(rows, cols), + } + } +} + +impl OctoWorker for BenchCmOctoWorker { + type Delta = CmDelta; + + fn process(&mut self, input: &SketchInput, emit: &mut F) + where + F: FnMut(Self::Delta), + { + self.sketch.insert_emit_delta(input, emit); + } +} + +struct BenchCmOctoAggregator { + sketch: CountMin, RegularPath>, +} + +impl OctoAggregator for BenchCmOctoAggregator { + type Delta = CmDelta; + + fn apply(&mut self, delta: Self::Delta) { + self.sketch.apply_delta(delta); + } +} + +fn build_inputs(sample_count: usize) -> Vec> { + let mut rng = StdRng::seed_from_u64(RNG_SEED ^ (sample_count as u64)); + (0..sample_count) + .map(|_| SketchInput::U64(rng.random::() & DOMAIN_MASK)) + .collect() +} + +fn build_round_robin_shards( + inputs: &[SketchInput<'static>], + num_workers: usize, +) -> Vec>> { + let mut shards = vec![Vec::new(); num_workers]; + for (idx, input) in inputs.iter().enumerate() { + shards[idx % num_workers].push(input.clone()); + } + shards +} + +fn run_octo_sharded( + shards: &[Vec>], + config: &OctoConfig, + worker_factory: impl Fn(usize) -> W, + parent_factory: impl FnOnce() -> P + Send, +) -> P +where + W: OctoWorker + 'static, + P: OctoAggregator + Send + 'static, +{ + let num_workers = shards.len(); + thread::scope(|s| { + let queue_capacity = config.queue_capacity.max(1); + let mut worker_txs = Vec::with_capacity(num_workers); + let mut worker_rxs = Vec::with_capacity(num_workers); + for _ in 0..num_workers { + let (tx, rx) = bounded::(queue_capacity); + worker_txs.push(tx); + worker_rxs.push(Some(rx)); + } + + let aggregator = s.spawn(move || { + let mut parent = parent_factory(); + if config.pin_cores { + let _ = core_affinity::set_for_current(core_affinity::CoreId { id: num_workers }); + } + + let mut disconnected = 0usize; + while disconnected < num_workers { + let mut made_progress = false; + for rx_slot in &mut worker_rxs { + let Some(rx) = rx_slot else { + continue; + }; + match rx.try_recv() { + Ok(delta) => { + parent.apply(delta); + made_progress = true; + } + Err(TryRecvError::Empty) => {} + Err(TryRecvError::Disconnected) => { + *rx_slot = None; + disconnected += 1; + } + } + } + if !made_progress { + std::hint::spin_loop(); + } + } + parent + }); + + for (worker_id, (shard, delta_tx_worker)) in + shards.iter().zip(worker_txs.into_iter()).enumerate() + { + let mut worker = worker_factory(worker_id); + let pin_cores = config.pin_cores; + s.spawn(move || { + if pin_cores { + let _ = core_affinity::set_for_current(core_affinity::CoreId { id: worker_id }); + } + for input in shard { + worker.process(input, &mut |delta| { + delta_tx_worker + .send(delta) + .expect("octo sharded aggregator dropped unexpectedly"); + }); + } + }); + } + + aggregator + .join() + .expect("octo sharded aggregator thread panicked") + }) +} + +enum CountMinMergeMsg { + Snapshot(CountMin, RegularPath>), + Done, +} + +enum CountMergeMsg { + Snapshot(Count, RegularPath>), + Done, +} + +enum HllMergeMsg { + Snapshot(HyperLogLog), + Done, +} + +fn run_periodic_countmin_full_merge_sharded( + shards: &[Vec>], +) -> CountMin, RegularPath> { + let num_workers = shards.len(); + thread::scope(|s| { + let (merge_tx, merge_rx) = mpsc::channel::(); + let aggregator = s.spawn(move || { + let mut parent = CountMin::, RegularPath>::with_dimensions( + CM_COUNT_ROWS, + CM_COUNT_COLS, + ); + let mut finished_workers = 0usize; + while finished_workers < num_workers { + match merge_rx + .recv() + .expect("CountMin merge channel closed unexpectedly") + { + CountMinMergeMsg::Snapshot(snapshot) => parent.merge(&snapshot), + CountMinMergeMsg::Done => finished_workers += 1, + } + } + parent + }); + + for shard in shards { + let merge_tx_worker = merge_tx.clone(); + s.spawn(move || { + let mut child = CountMin::, RegularPath>::with_dimensions( + CM_COUNT_ROWS, + CM_COUNT_COLS, + ); + let mut local_inserts = 0usize; + for input in shard { + child.insert(input); + local_inserts += 1; + if local_inserts % MERGE_INTERVAL == 0 { + merge_tx_worker + .send(CountMinMergeMsg::Snapshot(child)) + .expect("CountMin aggregator dropped unexpectedly"); + child = CountMin::, RegularPath>::with_dimensions( + CM_COUNT_ROWS, + CM_COUNT_COLS, + ); + } + } + merge_tx_worker + .send(CountMinMergeMsg::Snapshot(child)) + .expect("CountMin aggregator dropped unexpectedly"); + merge_tx_worker + .send(CountMinMergeMsg::Done) + .expect("CountMin aggregator dropped unexpectedly"); + }); + } + drop(merge_tx); + + aggregator + .join() + .expect("CountMin periodic merge aggregator panicked") + }) +} + +fn run_periodic_count_full_merge_sharded( + shards: &[Vec>], +) -> Count, RegularPath> { + let num_workers = shards.len(); + thread::scope(|s| { + let (merge_tx, merge_rx) = mpsc::channel::(); + let aggregator = s.spawn(move || { + let mut parent = + Count::, RegularPath>::with_dimensions(CM_COUNT_ROWS, CM_COUNT_COLS); + let mut finished_workers = 0usize; + while finished_workers < num_workers { + match merge_rx + .recv() + .expect("Count merge channel closed unexpectedly") + { + CountMergeMsg::Snapshot(snapshot) => parent.merge(&snapshot), + CountMergeMsg::Done => finished_workers += 1, + } + } + parent + }); + + for shard in shards { + let merge_tx_worker = merge_tx.clone(); + s.spawn(move || { + let mut child = Count::, RegularPath>::with_dimensions( + CM_COUNT_ROWS, + CM_COUNT_COLS, + ); + let mut local_inserts = 0usize; + for input in shard { + child.insert(input); + local_inserts += 1; + if local_inserts % MERGE_INTERVAL == 0 { + merge_tx_worker + .send(CountMergeMsg::Snapshot(child)) + .expect("Count aggregator dropped unexpectedly"); + child = Count::, RegularPath>::with_dimensions( + CM_COUNT_ROWS, + CM_COUNT_COLS, + ); + } + } + merge_tx_worker + .send(CountMergeMsg::Snapshot(child)) + .expect("Count aggregator dropped unexpectedly"); + merge_tx_worker + .send(CountMergeMsg::Done) + .expect("Count aggregator dropped unexpectedly"); + }); + } + drop(merge_tx); + + aggregator + .join() + .expect("Count periodic merge aggregator panicked") + }) +} + +fn run_periodic_hll_full_merge_sharded( + shards: &[Vec>], +) -> HyperLogLog { + let num_workers = shards.len(); + thread::scope(|s| { + let (merge_tx, merge_rx) = mpsc::channel::(); + let aggregator = s.spawn(move || { + let mut parent = HyperLogLog::::default(); + let mut finished_workers = 0usize; + while finished_workers < num_workers { + match merge_rx + .recv() + .expect("HLL merge channel closed unexpectedly") + { + HllMergeMsg::Snapshot(snapshot) => parent.merge(&snapshot), + HllMergeMsg::Done => finished_workers += 1, + } + } + parent + }); + + for shard in shards { + let merge_tx_worker = merge_tx.clone(); + s.spawn(move || { + let mut child = HyperLogLog::::default(); + let mut local_inserts = 0usize; + for input in shard { + child.insert(input); + local_inserts += 1; + if local_inserts % MERGE_INTERVAL == 0 { + merge_tx_worker + .send(HllMergeMsg::Snapshot(child)) + .expect("HLL aggregator dropped unexpectedly"); + child = HyperLogLog::::default(); + } + } + merge_tx_worker + .send(HllMergeMsg::Snapshot(child)) + .expect("HLL aggregator dropped unexpectedly"); + merge_tx_worker + .send(HllMergeMsg::Done) + .expect("HLL aggregator dropped unexpectedly"); + }); + } + drop(merge_tx); + + aggregator + .join() + .expect("HLL periodic merge aggregator panicked") + }) +} + +fn sanity_check_periodic_baselines() { + SANITY_ONCE.call_once(|| { + let inputs = build_inputs(20_000); + let workers = 4; + let shards = build_round_robin_shards(&inputs, workers); + + let config = OctoConfig { + num_workers: workers, + pin_cores: false, + queue_capacity: QUEUE_CAPACITY, + }; + + let octo_cm = run_octo_sharded( + &shards, + &config, + |_| BenchCmOctoWorker::new(CM_COUNT_ROWS, CM_COUNT_COLS), + || BenchCmOctoAggregator { + sketch: CountMin::with_dimensions(CM_COUNT_ROWS, CM_COUNT_COLS), + }, + ) + .sketch; + let periodic_cm = run_periodic_countmin_full_merge_sharded(&shards); + for key in [1_u64, 7, 42, 4097] { + let input = SketchInput::U64(key); + assert!(octo_cm.estimate(&input) <= periodic_cm.estimate(&input)); + } + + let octo_count = run_octo_sharded( + &shards, + &config, + |_| CountOctoWorker::new(CM_COUNT_ROWS, CM_COUNT_COLS), + || CountOctoAggregator { + sketch: Count::with_dimensions(CM_COUNT_ROWS, CM_COUNT_COLS), + }, + ) + .sketch; + let periodic_count = run_periodic_count_full_merge_sharded(&shards); + for key in [1_u64, 7, 42, 4097] { + let input = SketchInput::U64(key); + let octo = octo_count.estimate(&input); + let full = periodic_count.estimate(&input); + assert!(octo.abs() <= full.abs() + 500.0); + } + + let octo_hll = run_octo_sharded( + &shards, + &config, + |_| HllOctoWorker::new(), + || HllOctoAggregator { + sketch: HyperLogLog::::default(), + }, + ) + .sketch; + let periodic_hll = run_periodic_hll_full_merge_sharded(&shards); + assert!(octo_hll.estimate() <= periodic_hll.estimate()); + }); +} + +fn bench_countmin_merge_compare(c: &mut Criterion) { + sanity_check_periodic_baselines(); + + let inputs = build_inputs(CM_COUNT_INPUTS); + let mut group = c.benchmark_group("countmin_merge_compare"); + group.sample_size(10); + group.warm_up_time(Duration::from_secs(2)); + group.measurement_time(Duration::from_secs(8)); + group.throughput(Throughput::Elements(inputs.len() as u64)); + + group.bench_function("single_thread_insert", |b| { + b.iter_with_setup( + || { + CountMin::, RegularPath>::with_dimensions( + CM_COUNT_ROWS, + CM_COUNT_COLS, + ) + }, + |mut sketch| { + for input in &inputs { + sketch.insert(input); + } + black_box(sketch); + }, + ); + }); + + for &workers in &WORKER_COUNTS { + let shards = build_round_robin_shards(&inputs, workers); + group.bench_with_input( + BenchmarkId::new("octo_style", workers), + &workers, + |b, &w| { + b.iter(|| { + let config = OctoConfig { + num_workers: w, + pin_cores: false, + queue_capacity: QUEUE_CAPACITY, + }; + let parent = run_octo_sharded( + &shards, + &config, + |_| BenchCmOctoWorker::new(CM_COUNT_ROWS, CM_COUNT_COLS), + || BenchCmOctoAggregator { + sketch: CountMin::with_dimensions(CM_COUNT_ROWS, CM_COUNT_COLS), + }, + ); + black_box(parent.sketch); + }); + }, + ); + + group.bench_with_input( + BenchmarkId::new("regular_periodic_full_merge", workers), + &workers, + |b, _| { + b.iter(|| { + let merged = run_periodic_countmin_full_merge_sharded(&shards); + black_box(merged); + }); + }, + ); + } + + group.finish(); +} + +fn bench_count_merge_compare(c: &mut Criterion) { + sanity_check_periodic_baselines(); + + let inputs = build_inputs(CM_COUNT_INPUTS); + let mut group = c.benchmark_group("count_merge_compare"); + group.sample_size(10); + group.warm_up_time(Duration::from_secs(2)); + group.measurement_time(Duration::from_secs(8)); + group.throughput(Throughput::Elements(inputs.len() as u64)); + + group.bench_function("single_thread_insert", |b| { + b.iter_with_setup( + || Count::, RegularPath>::with_dimensions(CM_COUNT_ROWS, CM_COUNT_COLS), + |mut sketch| { + for input in &inputs { + sketch.insert(input); + } + black_box(sketch); + }, + ); + }); + + for &workers in &WORKER_COUNTS { + let shards = build_round_robin_shards(&inputs, workers); + group.bench_with_input( + BenchmarkId::new("octo_style", workers), + &workers, + |b, &w| { + b.iter(|| { + let config = OctoConfig { + num_workers: w, + pin_cores: false, + queue_capacity: QUEUE_CAPACITY, + }; + let parent = run_octo_sharded( + &shards, + &config, + |_| CountOctoWorker::new(CM_COUNT_ROWS, CM_COUNT_COLS), + || CountOctoAggregator { + sketch: Count::with_dimensions(CM_COUNT_ROWS, CM_COUNT_COLS), + }, + ); + black_box(parent.sketch); + }); + }, + ); + + group.bench_with_input( + BenchmarkId::new("regular_periodic_full_merge", workers), + &workers, + |b, _| { + b.iter(|| { + let merged = run_periodic_count_full_merge_sharded(&shards); + black_box(merged); + }); + }, + ); + } + + group.finish(); +} + +fn bench_hll_merge_compare(c: &mut Criterion) { + sanity_check_periodic_baselines(); + + let inputs = build_inputs(HLL_INPUTS); + let mut group = c.benchmark_group("hll_merge_compare"); + group.sample_size(10); + group.warm_up_time(Duration::from_secs(2)); + group.measurement_time(Duration::from_secs(8)); + group.throughput(Throughput::Elements(inputs.len() as u64)); + + group.bench_function("single_thread_insert", |b| { + b.iter_with_setup(HyperLogLog::::default, |mut sketch| { + for input in &inputs { + sketch.insert(input); + } + black_box(sketch); + }); + }); + + for &workers in &WORKER_COUNTS { + let shards = build_round_robin_shards(&inputs, workers); + group.bench_with_input( + BenchmarkId::new("octo_style", workers), + &workers, + |b, &w| { + b.iter(|| { + let config = OctoConfig { + num_workers: w, + pin_cores: false, + queue_capacity: QUEUE_CAPACITY, + }; + let parent = run_octo_sharded( + &shards, + &config, + |_| HllOctoWorker::new(), + || HllOctoAggregator { + sketch: HyperLogLog::::default(), + }, + ); + black_box(parent.sketch); + }); + }, + ); + + group.bench_with_input( + BenchmarkId::new("regular_periodic_full_merge", workers), + &workers, + |b, _| { + b.iter(|| { + let merged = run_periodic_hll_full_merge_sharded(&shards); + black_box(merged); + }); + }, + ); + } + + group.finish(); +} + +criterion_group!( + octo_speedup_benches, + bench_countmin_merge_compare, + bench_count_merge_compare, + bench_hll_merge_compare +); +criterion_main!(octo_speedup_benches); diff --git a/src/sketch_framework/mod.rs b/src/sketch_framework/mod.rs index 62db908..63113fa 100644 --- a/src/sketch_framework/mod.rs +++ b/src/sketch_framework/mod.rs @@ -35,6 +35,12 @@ pub use nitro::{NitroBatch, NitroEstimate, NitroTarget}; pub mod eh_univ_optimized; pub use eh_univ_optimized::{EHMapBucket, EHUnivMonBucket, EHUnivOptimized, EHUnivQueryResult}; +pub mod octo; +pub use octo::{ + CountOctoAggregator, CountOctoWorker, HllOctoAggregator, HllOctoWorker, OctoAggregator, + OctoConfig, OctoReadHandle, OctoResult, OctoRuntime, OctoWorker, run_octo, +}; + pub mod tumbling; pub use tumbling::{ FoldCMSConfig, FoldCSConfig, KLLConfig, SketchPool, TumblingWindow, TumblingWindowSketch, diff --git a/src/sketch_framework/octo.rs b/src/sketch_framework/octo.rs new file mode 100644 index 0000000..17343e7 --- /dev/null +++ b/src/sketch_framework/octo.rs @@ -0,0 +1,1016 @@ +//! OctoSketch multi-threaded sketch framework. +//! +//! Implements the parent-child delta-promotion architecture from OctoSketch (NSDI 2024). +//! Worker threads maintain lightweight child sketches with small counters and emit +//! compact delta entries via an MPSC channel when counters overflow a promotion +//! threshold. An aggregator thread applies deltas to a full-precision parent sketch. + +use std::marker::PhantomData; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; +use std::sync::{Arc, RwLock, Weak}; +use std::thread; + +use crossbeam_channel::{Receiver, Sender, TryRecvError, bounded}; + +use crate::{ + CmDelta, Count, CountDelta, CountMin, HllDelta, HyperLogLog, Regular, RegularPath, SketchInput, + Vector2D, +}; + +/// Legacy queue capacity default retained for config compatibility. +const DEFAULT_QUEUE_CAPACITY: usize = 65536; + +// --------------------------------------------------------------------------- +// Traits +// --------------------------------------------------------------------------- + +/// Worker-side trait: processes inputs and emits deltas. +pub trait OctoWorker: Send { + type Delta: Copy + Send + 'static; + + /// Process one input and emit zero or more deltas. + fn process(&mut self, input: &SketchInput, emit: &mut F) + where + F: FnMut(Self::Delta); +} + +/// Aggregator-side trait: absorbs deltas into a full-precision sketch. +pub trait OctoAggregator: Send { + type Delta: Copy + Send + 'static; + + /// Apply a single delta to the parent sketch. + fn apply(&mut self, delta: Self::Delta); +} + +// --------------------------------------------------------------------------- +// Configuration +// --------------------------------------------------------------------------- + +/// Configuration for `run_octo`. +pub struct OctoConfig { + /// Number of worker threads (default: 4). + pub num_workers: usize, + /// Pin worker threads to cores (default: true). + /// Worker i is pinned to core i, aggregator to core num_workers. + /// Silently skipped if pinning fails. + pub pin_cores: bool, + /// Queue capacity for bounded worker-input and worker-delta channels (default: 65536). + pub queue_capacity: usize, +} + +impl Default for OctoConfig { + fn default() -> Self { + Self { + num_workers: 4, + pin_cores: true, + queue_capacity: DEFAULT_QUEUE_CAPACITY, + } + } +} + +/// Result of an `run_octo` execution. +pub struct OctoResult

{ + pub parent: P, +} + +enum WorkerMsg { + Data(SketchInput<'static>), + End, +} + +/// Extends a `SketchInput` lifetime to `'static` for cross-thread transport in +/// streaming mode. Caller must ensure all borrowed data outlives worker processing. +#[inline(always)] +unsafe fn assume_input_static(input: SketchInput<'_>) -> SketchInput<'static> { + // SAFETY: enforced by caller contract described above. + unsafe { std::mem::transmute::, SketchInput<'static>>(input) } +} + +/// Streaming Octo runtime that accepts incremental inserts and finalizes into a parent sketch. +pub struct OctoRuntime +where + W: OctoWorker + 'static, + P: OctoAggregator + Send + Sync + 'static, +{ + core: Option>, + _worker_marker: PhantomData, +} + +/// Read-only handle for querying the live aggregator state while runtime is active. +pub struct OctoReadHandle

{ + parent: Weak>, +} + +impl

Clone for OctoReadHandle

{ + fn clone(&self) -> Self { + Self { + parent: Weak::clone(&self.parent), + } + } +} + +impl

OctoReadHandle

{ + /// Executes a read-only closure over the live parent state. + pub fn with_parent(&self, f: impl FnOnce(&P) -> R) -> R { + let parent = self + .parent + .upgrade() + .expect("Octo runtime has been finished and parent state was dropped"); + let guard = parent.read().expect("parent lock poisoned"); + f(&guard) + } +} + +struct OctoCore

{ + worker_input_txs: Vec>, + next_worker: AtomicUsize, + worker_handles: Vec>, + aggregator_handle: Option>, + parent: Arc>, + closed: AtomicBool, +} + +impl

OctoCore

{ + fn read_handle(&self) -> OctoReadHandle

{ + OctoReadHandle { + parent: Arc::downgrade(&self.parent), + } + } + + fn close(&self) { + if self.closed.swap(true, Ordering::AcqRel) { + return; + } + for tx in &self.worker_input_txs { + let _ = tx.send(WorkerMsg::End); + } + } +} + +impl

OctoCore

{ + fn into_parent(mut self) -> P { + self.close(); + + for handle in self.worker_handles.drain(..) { + handle.join().expect("worker thread panicked during finish"); + } + + if let Some(aggregator) = self.aggregator_handle.take() { + aggregator + .join() + .expect("aggregator thread panicked during finish"); + } + + let parent_lock = match Arc::try_unwrap(self.parent) { + Ok(lock) => lock, + Err(_) => panic!("Octo parent still has external strong references at finish"), + }; + parent_lock.into_inner().expect("parent lock poisoned") + } +} + +impl

OctoCore

+where + P: Send + Sync + 'static, +{ + fn start( + workers: Vec, + parent: P, + num_workers: usize, + pin_cores: bool, + queue_capacity: usize, + ) -> Self + where + W: OctoWorker + 'static, + P: OctoAggregator, + { + assert_eq!(workers.len(), num_workers); + let queue_capacity = queue_capacity.max(1); + + let parent = Arc::new(RwLock::new(parent)); + let parent_for_aggregator = Arc::clone(&parent); + let mut delta_txs: Vec> = Vec::with_capacity(num_workers); + let mut delta_rxs: Vec>> = Vec::with_capacity(num_workers); + for _ in 0..num_workers { + let (tx, rx) = bounded::(queue_capacity); + delta_txs.push(tx); + delta_rxs.push(Some(rx)); + } + + let aggregator_handle = thread::spawn(move || { + if pin_cores { + let _ = core_affinity::set_for_current(core_affinity::CoreId { id: num_workers }); + } + + let mut disconnected_workers = 0usize; + while disconnected_workers < num_workers { + let mut made_progress = false; + for rx_slot in &mut delta_rxs { + let Some(rx) = rx_slot else { + continue; + }; + match rx.try_recv() { + Ok(delta) => { + let mut guard = parent_for_aggregator + .write() + .expect("parent lock poisoned in aggregator"); + guard.apply(delta); + made_progress = true; + } + Err(TryRecvError::Empty) => {} + Err(TryRecvError::Disconnected) => { + *rx_slot = None; + disconnected_workers += 1; + } + } + } + if !made_progress { + std::hint::spin_loop(); + } + } + }); + + let mut worker_input_txs = Vec::with_capacity(num_workers); + let mut worker_handles = Vec::with_capacity(num_workers); + for (worker_id, (mut worker, delta_tx_worker)) in + workers.into_iter().zip(delta_txs.into_iter()).enumerate() + { + let (worker_tx, worker_rx) = bounded::(queue_capacity); + worker_input_txs.push(worker_tx); + let pin_cores = pin_cores; + worker_handles.push(thread::spawn(move || { + if pin_cores { + let _ = core_affinity::set_for_current(core_affinity::CoreId { id: worker_id }); + } + while let Ok(msg) = worker_rx.recv() { + match msg { + WorkerMsg::Data(input) => worker.process(&input, &mut |delta| { + delta_tx_worker + .send(delta) + .expect("aggregator receiver dropped while workers still running"); + }), + WorkerMsg::End => break, + } + } + })); + } + + Self { + worker_input_txs, + next_worker: AtomicUsize::new(0), + worker_handles, + aggregator_handle: Some(aggregator_handle), + parent, + closed: AtomicBool::new(false), + } + } +} + +impl OctoRuntime +where + W: OctoWorker + 'static, + P: OctoAggregator + Send + Sync + 'static, +{ + pub fn new(config: &OctoConfig, worker_factory: F, parent_factory: PF) -> Self + where + F: Fn(usize) -> W, + PF: FnOnce() -> P, + { + let num_workers = config.num_workers.max(1); + let workers: Vec = (0..num_workers).map(worker_factory).collect(); + let parent = parent_factory(); + let core = OctoCore::start( + workers, + parent, + num_workers, + config.pin_cores, + config.queue_capacity, + ); + + Self { + core: Some(core), + _worker_marker: PhantomData, + } + } + + pub fn read_handle(&self) -> OctoReadHandle

{ + self.core + .as_ref() + .expect("runtime core missing") + .read_handle() + } + + pub fn close(&self) { + self.core.as_ref().expect("runtime core missing").close(); + } + + pub fn insert(&mut self, input: SketchInput<'_>) { + let core = self.core.as_ref().expect("runtime core missing"); + if core.closed.load(Ordering::Acquire) { + panic!("cannot insert after runtime has been closed"); + } + + let worker_id = + core.next_worker.fetch_add(1, Ordering::AcqRel) % core.worker_input_txs.len(); + // SAFETY: caller explicitly guarantees borrowed data lives long enough. + let static_input = unsafe { assume_input_static(input) }; + core.worker_input_txs[worker_id] + .send(WorkerMsg::Data(static_input)) + .expect("worker receiver dropped while runtime is active"); + } + + pub fn insert_batch(&mut self, inputs: &[SketchInput<'_>]) { + for input in inputs { + self.insert(input.clone()); + } + } + + pub fn finish(mut self) -> OctoResult

{ + let parent = self + .core + .take() + .expect("runtime core missing") + .into_parent(); + + OctoResult { parent } + } +} + +// --------------------------------------------------------------------------- +// Concrete worker/parent implementations +// --------------------------------------------------------------------------- + +// -- CountMin -- + +/// OctoSketch worker backed by `CountMin`. +pub struct CmOctoWorker { + sketch: CountMin, RegularPath>, +} + +impl CmOctoWorker { + pub fn new(rows: usize, cols: usize) -> Self { + Self { + sketch: CountMin::with_dimensions(rows, cols), + } + } +} + +impl OctoWorker for CmOctoWorker { + type Delta = CmDelta; + + #[inline(always)] + fn process(&mut self, input: &SketchInput, emit: &mut F) + where + F: FnMut(Self::Delta), + { + self.sketch.insert_emit_delta(input, emit); + } +} + +/// OctoSketch parent wrapping a full-precision `CountMin`. +pub struct CmOctoAggregator { + pub sketch: CountMin, RegularPath>, +} + +impl OctoAggregator for CmOctoAggregator { + type Delta = CmDelta; + + #[inline(always)] + fn apply(&mut self, delta: CmDelta) { + self.sketch.apply_delta(delta); + } +} + +// -- Count Sketch -- + +/// OctoSketch worker backed by `Count`. +pub struct CountOctoWorker { + child: Count, RegularPath>, +} + +impl CountOctoWorker { + pub fn new(rows: usize, cols: usize) -> Self { + Self { + child: Count::with_dimensions(rows, cols), + } + } +} + +impl OctoWorker for CountOctoWorker { + type Delta = CountDelta; + + #[inline(always)] + fn process(&mut self, input: &SketchInput, emit: &mut F) + where + F: FnMut(Self::Delta), + { + self.child.insert_emit_delta(input, emit); + } +} + +/// OctoSketch parent wrapping a full-precision `Count`. +pub struct CountOctoAggregator { + pub sketch: Count, RegularPath>, +} + +impl OctoAggregator for CountOctoAggregator { + type Delta = CountDelta; + + #[inline(always)] + fn apply(&mut self, delta: CountDelta) { + self.sketch.apply_delta(delta); + } +} + +// -- HyperLogLog -- + +/// OctoSketch worker backed by `HyperLogLog`. +pub struct HllOctoWorker { + child: HyperLogLog, +} + +impl HllOctoWorker { + pub fn new() -> Self { + Self { + child: HyperLogLog::default(), + } + } +} + +impl Default for HllOctoWorker { + fn default() -> Self { + Self::new() + } +} + +impl OctoWorker for HllOctoWorker { + type Delta = HllDelta; + + #[inline(always)] + fn process(&mut self, input: &SketchInput, emit: &mut F) + where + F: FnMut(Self::Delta), + { + self.child.insert_emit_delta(input, emit); + } +} + +/// OctoSketch parent wrapping a full-precision `HyperLogLog`. +pub struct HllOctoAggregator { + pub sketch: HyperLogLog, +} + +impl OctoAggregator for HllOctoAggregator { + type Delta = HllDelta; + + #[inline(always)] + fn apply(&mut self, delta: HllDelta) { + self.sketch.apply_delta(delta); + } +} + +// --------------------------------------------------------------------------- +// Core execution engine +// --------------------------------------------------------------------------- + +/// Runs the OctoSketch multi-threaded insert protocol. +/// +/// 1. Dispatches `inputs` across workers round-robin through per-worker channels. +/// 2. Each worker maintains a child sketch, emitting deltas via an MPSC channel. +/// 3. The aggregator blocks on the channel and applies deltas to the parent. +/// 4. Returns the fully-merged parent sketch. +pub fn run_octo( + inputs: &[SketchInput<'_>], + config: &OctoConfig, + worker_factory: impl Fn(usize) -> W, + parent_factory: impl FnOnce() -> P, +) -> OctoResult

+where + W: OctoWorker + 'static, + P: OctoAggregator + Send + Sync + 'static, +{ + let mut runtime = OctoRuntime::new(config, worker_factory, parent_factory); + for input in inputs { + runtime.insert(input.clone()); + } + runtime.finish() +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::SketchInput; + + // ----------------------------------------------------------------------- + // Layer 2 unit tests: child sketches + apply_delta + // ----------------------------------------------------------------------- + + #[test] + fn cm_insert_emit_delta_emits_at_threshold() { + let mut worker_sketch = CountMin::with_dimensions(3, 64); + let key = SketchInput::U64(42); + let mut deltas: Vec = Vec::new(); + + // Insert CM_PROMASK-1 times: no delta yet. + for _ in 0..(crate::CM_PROMASK - 1) { + worker_sketch.insert_emit_delta(&key, &mut |d| deltas.push(d)); + } + assert!(deltas.is_empty(), "should not emit before threshold"); + + // One more insert triggers the delta. + worker_sketch.insert_emit_delta(&key, &mut |d| deltas.push(d)); + assert_eq!(deltas.len(), 3, "should emit one delta per row (3 rows)"); + for d in &deltas { + assert_eq!(d.value, crate::CM_PROMASK); + } + } + + #[test] + fn cm_apply_delta_increments_parent() { + let mut parent = CountMin::, RegularPath>::with_dimensions(3, 64); + let delta = CmDelta { + row: 1, + col: 5, + value: 100, + }; + parent.apply_delta(delta); + assert_eq!(parent.as_storage().query_one_counter(1, 5), 100); + + parent.apply_delta(delta); + assert_eq!(parent.as_storage().query_one_counter(1, 5), 200); + } + + #[test] + fn count_child_emits_delta_at_threshold() { + let mut child = Count::, RegularPath>::with_dimensions(3, 64); + let key = SketchInput::U64(99); + let mut deltas: Vec = Vec::new(); + + // Insert enough times to trigger at least one delta. + for _ in 0..200 { + child.insert_emit_delta(&key, &mut |d| deltas.push(d)); + } + // With COUNT_PROMASK = 63 and 3 rows, after 200 inserts we + // should have emitted at least 3 deltas (one per row per threshold). + assert!( + deltas.len() >= 3, + "expected at least 3 deltas, got {}", + deltas.len() + ); + } + + #[test] + fn count_apply_delta_increments_parent() { + let mut parent = Count::, RegularPath>::with_dimensions(3, 64); + let delta = CountDelta { + row: 0, + col: 10, + value: -50, + }; + parent.apply_delta(delta); + assert_eq!(parent.as_storage().query_one_counter(0, 10), -50); + } + + #[test] + fn hll_child_emits_delta_on_improvement() { + let mut child = HyperLogLog::::default(); + let mut deltas: Vec = Vec::new(); + + // First insert should always emit (register goes from 0 to something > 0). + child.insert_emit_delta(&SketchInput::U64(1), &mut |d| deltas.push(d)); + assert_eq!(deltas.len(), 1, "first insert should emit a delta"); + + // Inserting the same value again should NOT emit (register unchanged). + let len_before = deltas.len(); + child.insert_emit_delta(&SketchInput::U64(1), &mut |d| deltas.push(d)); + assert_eq!(deltas.len(), len_before, "duplicate should not emit"); + } + + #[test] + fn hll_apply_delta_updates_register() { + let mut parent = HyperLogLog::::default(); + let delta = HllDelta { + pos: 100, + value: 10, + }; + parent.apply_delta(delta); + + // Apply a smaller value — should not change. + let smaller = HllDelta { pos: 100, value: 5 }; + parent.apply_delta(smaller); + + // Apply a larger value — should update. + let larger = HllDelta { + pos: 100, + value: 15, + }; + parent.apply_delta(larger); + } + + // ----------------------------------------------------------------------- + // Layer 3 integration tests: run_octo + // ----------------------------------------------------------------------- + + #[test] + fn run_octo_cm_matches_single_thread() { + let rows = 3; + let cols = 4096; + let n = 100_000u64; + + let inputs: Vec> = (0..n).map(|i| SketchInput::U64(i % 1024)).collect(); + + // Single-threaded reference. + let mut reference = CountMin::, RegularPath>::with_dimensions(rows, cols); + for input in &inputs { + reference.insert(input); + } + + // Multi-threaded via OctoSketch. + let config = OctoConfig { + num_workers: 4, + pin_cores: false, // don't pin in tests (CI may have few cores) + queue_capacity: 8192, + }; + let result = run_octo( + &inputs, + &config, + |_| CmOctoWorker::new(rows, cols), + || CmOctoAggregator { + sketch: CountMin::with_dimensions(rows, cols), + }, + ); + + // Verify estimates are close. + // OctoSketch child uses u8 counters, so there may be slight differences + // due to the promotion threshold batching, but the total counts should match. + for key_val in 0u64..1024 { + let key = SketchInput::U64(key_val); + let ref_est = reference.estimate(&key); + let octo_est = result.parent.sketch.estimate(&key); + // The octo estimate should be <= reference (some counts may be + // lost in the u8 remainder that didn't reach PROMASK). + assert!( + (ref_est - octo_est).abs() < 200, + "key {key_val}: ref={ref_est}, octo={octo_est}, diff={}", + (ref_est - octo_est).abs() + ); + } + } + + #[test] + fn run_octo_hll_cardinality() { + let n = 50_000u64; + let inputs: Vec> = (0..n).map(SketchInput::U64).collect(); + + let config = OctoConfig { + num_workers: 4, + pin_cores: false, + queue_capacity: 16384, + }; + + let result = run_octo( + &inputs, + &config, + |_| HllOctoWorker::new(), + || HllOctoAggregator { + sketch: HyperLogLog::::default(), + }, + ); + + let estimate = result.parent.sketch.estimate(); + let truth = n as f64; + let error = (estimate as f64 - truth).abs() / truth; + assert!( + error < 0.05, + "HLL cardinality error {error:.4} exceeded 5% (truth {truth}, estimate {estimate})" + ); + } + + #[test] + fn run_octo_count_sketch_basic() { + let rows = 3; + let cols = 4096; + let n = 50_000u64; + + let inputs: Vec> = (0..n).map(|i| SketchInput::U64(i % 512)).collect(); + + let config = OctoConfig { + num_workers: 2, + pin_cores: false, + queue_capacity: 8192, + }; + + let result = run_octo( + &inputs, + &config, + |_| CountOctoWorker::new(rows, cols), + || CountOctoAggregator { + sketch: Count::with_dimensions(rows, cols), + }, + ); + + // Each key appears ~97 times. Check a sample of keys. + let expected_per_key = (n / 512) as f64; + for key_val in [0u64, 100, 255, 511] { + let key = SketchInput::U64(key_val); + let est = result.parent.sketch.estimate(&key); + assert!( + (est - expected_per_key).abs() < expected_per_key * 0.5, + "key {key_val}: estimate={est}, expected~{expected_per_key}" + ); + } + } + + #[test] + fn run_octo_single_worker() { + // Edge case: single worker should still work. + let inputs: Vec> = (0..1000u64).map(SketchInput::U64).collect(); + + let config = OctoConfig { + num_workers: 1, + pin_cores: false, + queue_capacity: 4096, + }; + + let result = run_octo( + &inputs, + &config, + |_| HllOctoWorker::new(), + || HllOctoAggregator { + sketch: HyperLogLog::::default(), + }, + ); + + let estimate = result.parent.sketch.estimate(); + assert!( + estimate > 900 && estimate < 1100, + "single-worker HLL estimate {estimate} out of range for 1000 distinct keys" + ); + } + + #[test] + fn octo_runtime_streaming_cm_matches_batch_api() { + let rows = 3; + let cols = 4096; + let n = 30_000u64; + let inputs: Vec> = (0..n).map(|i| SketchInput::U64(i % 1024)).collect(); + let config = OctoConfig { + num_workers: 4, + pin_cores: false, + queue_capacity: 8192, + }; + + let batch_result = run_octo( + &inputs, + &config, + |_| CmOctoWorker::new(rows, cols), + || CmOctoAggregator { + sketch: CountMin::with_dimensions(rows, cols), + }, + ); + + let mut runtime = OctoRuntime::new( + &config, + move |_| CmOctoWorker::new(rows, cols), + move || CmOctoAggregator { + sketch: CountMin::with_dimensions(rows, cols), + }, + ); + for input in &inputs { + runtime.insert(input.clone()); + } + let streaming_result = runtime.finish(); + + for key_val in 0u64..128 { + let key = SketchInput::U64(key_val); + let batch_est = batch_result.parent.sketch.estimate(&key); + let stream_est = streaming_result.parent.sketch.estimate(&key); + assert_eq!(batch_est, stream_est, "key {key_val} mismatch"); + } + } + + #[test] + fn octo_runtime_mixed_insert_and_batch_hll() { + let config = OctoConfig { + num_workers: 2, + pin_cores: false, + queue_capacity: 4096, + }; + let mut runtime = OctoRuntime::new( + &config, + |_| HllOctoWorker::new(), + || HllOctoAggregator { + sketch: HyperLogLog::::default(), + }, + ); + + runtime.insert(SketchInput::U64(1)); + runtime.insert(SketchInput::U64(2)); + let batch: Vec> = (3..2000).map(SketchInput::U64).collect(); + runtime.insert_batch(&batch); + let result = runtime.finish(); + let estimate = result.parent.sketch.estimate(); + assert!( + estimate > 1700 && estimate < 2300, + "runtime mixed insert+batch estimate {estimate} is out of expected range" + ); + } + + #[test] + fn octo_runtime_live_read_handle_observes_progress() { + let config = OctoConfig { + num_workers: 2, + pin_cores: false, + queue_capacity: 4096, + }; + let mut runtime = OctoRuntime::new( + &config, + |_| CountingWorker, + || CountingAggregator { total: 0 }, + ); + let reader = runtime.read_handle(); + + for i in 0..64u64 { + runtime.insert(SketchInput::U64(i)); + } + std::thread::sleep(std::time::Duration::from_millis(5)); + let observed = reader.with_parent(|p| p.total); + assert!( + observed <= 64, + "live reader should observe a partial or complete total" + ); + + let result = runtime.finish(); + assert_eq!( + result.parent.total, 64, + "all inserted items should be accounted for" + ); + assert!( + result.parent.total >= observed, + "final total should not be less than live snapshot" + ); + } + + #[test] + fn octo_runtime_close_is_idempotent() { + let config = OctoConfig { + num_workers: 2, + pin_cores: false, + queue_capacity: 4096, + }; + let runtime = OctoRuntime::new( + &config, + |_| HllOctoWorker::new(), + || HllOctoAggregator { + sketch: HyperLogLog::::default(), + }, + ); + runtime.close(); + runtime.close(); + let result = runtime.finish(); + assert_eq!(result.parent.sketch.estimate(), 0); + } + + #[test] + #[should_panic(expected = "cannot insert after runtime has been closed")] + fn octo_runtime_insert_after_close_panics() { + let config = OctoConfig { + num_workers: 2, + pin_cores: false, + queue_capacity: 4096, + }; + let mut runtime = OctoRuntime::new( + &config, + |_| HllOctoWorker::new(), + || HllOctoAggregator { + sketch: HyperLogLog::::default(), + }, + ); + runtime.close(); + runtime.insert(SketchInput::U64(1)); + } + + #[test] + fn octo_runtime_empty_stream_finishes() { + let config = OctoConfig { + num_workers: 4, + pin_cores: false, + queue_capacity: 4096, + }; + let runtime = OctoRuntime::new( + &config, + |_| HllOctoWorker::new(), + || HllOctoAggregator { + sketch: HyperLogLog::::default(), + }, + ); + let result = runtime.finish(); + let estimate = result.parent.sketch.estimate(); + assert!( + estimate == 0, + "empty runtime should estimate 0 cardinality, got {estimate}" + ); + } + + struct CountingWorker; + + impl OctoWorker for CountingWorker { + type Delta = u64; + + fn process(&mut self, _input: &SketchInput, emit: &mut F) + where + F: FnMut(Self::Delta), + { + emit(1); + } + } + + struct CountingAggregator { + total: u64, + } + + impl OctoAggregator for CountingAggregator { + type Delta = u64; + + fn apply(&mut self, delta: Self::Delta) { + self.total += delta; + } + } + + #[test] + fn octo_runtime_close_preserves_queued_items_without_dispatcher() { + let config = OctoConfig { + num_workers: 4, + pin_cores: false, + queue_capacity: 4096, + }; + let n = 257usize; + let mut runtime = OctoRuntime::new( + &config, + |_| CountingWorker, + || CountingAggregator { total: 0 }, + ); + + for i in 0..n as u64 { + runtime.insert(SketchInput::U64(i + 42)); + } + runtime.close(); + let result = runtime.finish(); + assert_eq!( + result.parent.total, n as u64, + "runtime close should preserve already queued items" + ); + } + + struct WorkerIdEmitter { + worker_id: usize, + } + + impl OctoWorker for WorkerIdEmitter { + type Delta = usize; + + fn process(&mut self, _input: &SketchInput, emit: &mut F) + where + F: FnMut(Self::Delta), + { + emit(self.worker_id); + } + } + + struct WorkerLoadAggregator { + loads: Vec, + } + + impl OctoAggregator for WorkerLoadAggregator { + type Delta = usize; + + fn apply(&mut self, delta: Self::Delta) { + self.loads[delta] += 1; + } + } + + #[test] + fn octo_runtime_round_robin_selector_distributes_deterministically() { + let num_workers = 3; + let inserts = 10u64; + let config = OctoConfig { + num_workers, + pin_cores: false, + queue_capacity: 4096, + }; + let mut runtime = OctoRuntime::new( + &config, + |worker_id| WorkerIdEmitter { worker_id }, + || WorkerLoadAggregator { + loads: vec![0; num_workers], + }, + ); + + for i in 0..inserts { + runtime.insert(SketchInput::U64(i)); + } + let result = runtime.finish(); + + assert_eq!(result.parent.loads, vec![4, 3, 3]); + } +} diff --git a/src/sketches/count.rs b/src/sketches/count.rs index 2e0a0ad..af52c8b 100644 --- a/src/sketches/count.rs +++ b/src/sketches/count.rs @@ -2,7 +2,7 @@ use crate::{ DefaultMatrixI32, DefaultMatrixI64, DefaultMatrixI128, DefaultXxHasher, FastPath, FastPathHasher, FixedMatrix, MatrixFastHash, MatrixStorage, NitroTarget, QuickMatrixI64, QuickMatrixI128, RegularPath, SketchHasher, SketchInput, Vector1D, Vector2D, - compute_median_inline_f64, + compute_median_inline_f64, hash64_seeded, }; use rmp_serde::{ decode::Error as RmpDecodeError, encode::Error as RmpEncodeError, from_slice, to_vec_named, @@ -716,6 +716,49 @@ impl CountL2HH { } } +use crate::octo_delta::{COUNT_PROMASK, CountDelta}; + +/// Worker-side update for OctoSketch delta promotion. +impl Count, RegularPath> { + /// Insert a key with the Count sketch sign convention. + /// Emits `CountDelta` when |counter| >= `COUNT_PROMASK`, then resets the counter. + #[inline(always)] + pub fn insert_emit_delta(&mut self, value: &SketchInput, emit: &mut impl FnMut(CountDelta)) { + let rows = self.counts.rows(); + let cols = self.counts.cols(); + let data = self.counts.as_mut_slice(); + for r in 0..rows { + let hashed = hash64_seeded(r, value); + let col = ((hashed & LOWER_32_MASK) as usize) % cols; + let sign: i32 = if ((hashed >> 63) & 1) == 1 { 1 } else { -1 }; + let cell = &mut data[r * cols + col]; + *cell += sign; + if cell.unsigned_abs() >= COUNT_PROMASK as u32 { + emit(CountDelta { + row: r as u16, + col: col as u16, + value: *cell as i8, + }); + *cell = 0; + } + } + } +} + +/// Apply a `CountDelta` to a full-precision parent Count sketch. +impl Count +where + S::Counter: Copy + std::ops::AddAssign + From, +{ + pub fn apply_delta(&mut self, delta: CountDelta) { + self.counts.increment_by_row( + delta.row as usize, + delta.col as usize, + S::Counter::from(delta.value as i32), + ); + } +} + #[cfg(test)] mod tests { use super::*; @@ -725,6 +768,21 @@ mod tests { use crate::{SketchInput, hash64_seeded}; use std::collections::HashMap; + #[test] + fn count_child_insert_emits_at_threshold() { + let mut child = Count::, RegularPath>::with_dimensions(3, 64); + let key = SketchInput::U64(99); + let mut deltas: Vec = Vec::new(); + + for _ in 0..200 { + child.insert_emit_delta(&key, &mut |d| deltas.push(d)); + } + assert!( + deltas.len() >= 3, + "expected at least one promoted delta per row" + ); + } + fn counter_sign(row: usize, key: &SketchInput) -> i32 { let hash = hash64_seeded(row, key); if (hash >> 63) & 1 == 1 { 1 } else { -1 } diff --git a/src/sketches/countmin.rs b/src/sketches/countmin.rs index 30be6cd..7753f0d 100644 --- a/src/sketches/countmin.rs +++ b/src/sketches/countmin.rs @@ -4,10 +4,11 @@ use rmp_serde::{ use serde::{Deserialize, Serialize}; use std::marker::PhantomData; +use crate::octo_delta::{CM_PROMASK, CmDelta}; use crate::{ DefaultMatrixI32, DefaultMatrixI64, DefaultMatrixI128, DefaultXxHasher, FastPath, FastPathHasher, FixedMatrix, MatrixStorage, NitroTarget, QuickMatrixI64, QuickMatrixI128, - RegularPath, SketchHasher, SketchInput, Vector2D, + RegularPath, SketchHasher, SketchInput, Vector2D, hash64_seeded, }; const DEFAULT_ROW_NUM: usize = 3; @@ -327,6 +328,35 @@ where /// Count-Min sketch with floating-point counters (no integer rounding). pub type CountMinF64 = CountMin, RegularPath, H>; +impl CountMin, RegularPath> { + /// Inserts one key and emits a `CmDelta` whenever a touched counter reaches + /// a `CM_PROMASK` multiple. + #[inline(always)] + pub fn insert_emit_delta(&mut self, value: &SketchInput, emit: &mut impl FnMut(CmDelta)) { + let rows = self.counts.rows(); + let cols = self.counts.cols(); + // let threshold = CM_PROMASK as i32; + for r in 0..rows { + let hashed = hash64_seeded(r, value); + let col = ((hashed & LOWER_32_MASK) as usize) % cols; + self.counts.increment_by_row(r, col, 1); + let current = self.counts.query_one_counter(r, col); + if current % CM_PROMASK as i32 == 0 { + emit(CmDelta { + row: r as u16, + col: col as u16, + value: CM_PROMASK, + }); + } + } + } + /// Apply a `CmDelta` to a full-precision parent CountMin sketch. + pub fn apply_delta(&mut self, delta: CmDelta) { + self.counts + .increment_by_row(delta.row as usize, delta.col as usize, delta.value as i32); + } +} + // SketchInput adapters for the fast-path Count-Min update rule. // Fast-path CountMin operations using precomputed hashes. Uses PartialOrd for f64 support. impl CountMin @@ -479,6 +509,51 @@ mod tests { use core::f64; use std::collections::HashMap; + #[test] + fn countmin_insert_emit_delta_emits_at_threshold_and_resets_period() { + let mut sketch = CountMin::, RegularPath>::with_dimensions(3, 64); + let key = SketchInput::U64(42); + let mut deltas: Vec = Vec::new(); + + for _ in 0..(CM_PROMASK - 1) { + sketch.insert_emit_delta(&key, &mut |d| deltas.push(d)); + } + assert!( + deltas.is_empty(), + "regular CMS worker path should not emit before threshold" + ); + + sketch.insert_emit_delta(&key, &mut |d| deltas.push(d)); + assert_eq!( + deltas.len(), + 3, + "should emit one delta per row at threshold" + ); + assert!(deltas.iter().all(|d| d.value == CM_PROMASK)); + + for _ in 0..(CM_PROMASK - 1) { + sketch.insert_emit_delta(&key, &mut |d| deltas.push(d)); + } + assert_eq!(deltas.len(), 3, "no second emission before next threshold"); + sketch.insert_emit_delta(&key, &mut |d| deltas.push(d)); + assert_eq!(deltas.len(), 6, "should emit again on next threshold"); + } + + #[test] + fn countmin_apply_delta_increments_parent_counter() { + let mut parent = CountMin::, RegularPath>::with_dimensions(3, 64); + let delta = CmDelta { + row: 1, + col: 5, + value: CM_PROMASK, + }; + parent.apply_delta(delta); + assert_eq!( + parent.as_storage().query_one_counter(1, 5), + CM_PROMASK as i32 + ); + } + fn run_zipf_stream( rows: usize, cols: usize, diff --git a/src/sketches/hll.rs b/src/sketches/hll.rs index 2d0ab79..848c107 100644 --- a/src/sketches/hll.rs +++ b/src/sketches/hll.rs @@ -331,6 +331,45 @@ impl HyperLogLogHIP { } } } +use crate::octo_delta::HllDelta; + +impl HyperLogLog { + /// Insert using a pre-computed hash and emit `HllDelta` on register improvement. + #[inline(always)] + pub fn insert_emit_delta_with_hash( + &mut self, + hashed_val: u64, + emit: &mut impl FnMut(HllDelta), + ) { + let bucket_num = ((hashed_val >> HLL_Q) & HLL_P_MASK) as usize; + let leading_zero = ((hashed_val << HLL_P) + HLL_P_MASK).leading_zeros() as u8 + 1; + if leading_zero > self.registers[bucket_num] { + self.registers[bucket_num] = leading_zero; + emit(HllDelta { + pos: bucket_num as u16, + value: leading_zero, + }); + } + } + + /// Insert a key and emit `HllDelta` on register improvement. + #[inline(always)] + pub fn insert_emit_delta(&mut self, obj: &SketchInput, emit: &mut impl FnMut(HllDelta)) { + let hashed_val = hash64_seeded(CANONICAL_HASH_SEED, obj); + self.insert_emit_delta_with_hash(hashed_val, emit); + } +} + +/// Apply an `HllDelta` to a full-precision parent HyperLogLog sketch. +impl HyperLogLog { + pub fn apply_delta(&mut self, delta: HllDelta) { + let pos = delta.pos as usize; + if delta.value > self.registers[pos] { + self.registers[pos] = delta.value; + } + } +} + #[cfg(test)] mod tests { @@ -341,6 +380,19 @@ mod tests { const ERROR_TOLERANCE: f64 = 0.02; const SERDE_SAMPLE: usize = 100_000; + #[test] + fn hll_child_insert_emits_on_improvement() { + let mut child = HyperLogLog::::default(); + let mut deltas: Vec = Vec::new(); + + child.insert_emit_delta(&SketchInput::U64(1), &mut |d| deltas.push(d)); + assert_eq!(deltas.len(), 1, "first insert should improve one register"); + + let before = deltas.len(); + child.insert_emit_delta(&SketchInput::U64(1), &mut |d| deltas.push(d)); + assert_eq!(deltas.len(), before, "duplicate should not emit"); + } + trait HllEstimator: Default { fn push(&mut self, input: &SketchInput); fn insert_with_hash(&mut self, hashed: u64); diff --git a/src/sketches/mod.rs b/src/sketches/mod.rs index e504fa7..a119757 100644 --- a/src/sketches/mod.rs +++ b/src/sketches/mod.rs @@ -44,6 +44,9 @@ pub use cms_heap::CMSHeap; pub mod cs_heap; pub use cs_heap::CSHeap; +pub mod octo_delta; +pub use octo_delta::{CM_PROMASK, COUNT_PROMASK, CmDelta, CountDelta, HLL_PROMASK, HllDelta}; + pub mod fold_cms; pub use fold_cms::{FoldCMS, FoldCell, FoldEntry}; diff --git a/src/sketches/octo_delta.rs b/src/sketches/octo_delta.rs new file mode 100644 index 0000000..04a3fb6 --- /dev/null +++ b/src/sketches/octo_delta.rs @@ -0,0 +1,39 @@ +//! Delta entry types for OctoSketch-style multi-threaded sketch updates. +//! +//! Each delta represents an accumulated counter change emitted by a child +//! worker sketch when a local counter overflows the promotion threshold (PROMASK). + +/// CountMin promotion threshold: emit delta when u8 counter >= 127. +pub const CM_PROMASK: u8 = 0x7f; + +/// Count sketch promotion threshold: emit delta when |i8 counter| >= 63. +pub const COUNT_PROMASK: u8 = 0x3f; + +/// HLL promotion threshold: 0 means every register improvement is emitted. +pub const HLL_PROMASK: u8 = 0; + +/// Delta emitted by a CountMin child worker. +/// Represents an accumulated unsigned count for a single cell. +#[derive(Clone, Copy, Debug)] +pub struct CmDelta { + pub row: u16, + pub col: u16, + pub value: u8, +} + +/// Delta emitted by a Count sketch child worker. +/// Represents a signed accumulated count for a single cell. +#[derive(Clone, Copy, Debug)] +pub struct CountDelta { + pub row: u16, + pub col: u16, + pub value: i8, +} + +/// Delta emitted by an HLL child worker. +/// Represents a register improvement (max-register semantics). +#[derive(Clone, Copy, Debug)] +pub struct HllDelta { + pub pos: u16, + pub value: u8, +}