Skip to content

Commit 0426b66

Browse files
GnaneshGnaniGnaneshGnani
authored
Sketchlib CountMin Support (#215)
* Integrate sketchlib-rust for Count-Min Sketch * Restore per-backend default constants, global default Legacy * Use per-backend defaults in fidelity, configurable impl_mode in UDF templates * report: scope to CMS only for PR 3 * UDFs: use same impl mode as QueryEngine (sketch_cms_impl, etc.) * Simplify UDF impl mode, default CMS to sketchlib * Fix black formatting in arroyo.py * Set deafult mode to sketchlib
1 parent 71100f9 commit 0426b66

15 files changed

Lines changed: 698 additions & 186 deletions

File tree

‎Cargo.lock‎

Lines changed: 27 additions & 28 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎asap-common/sketch-core/Cargo.toml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ serde.workspace = true
99
rmp-serde = "1.1"
1010
xxhash-rust = { version = "0.8", features = ["xxh32"] }
1111
dsrs = { git = "https://github.com/ProjectASAP/datasketches-rs", rev = "d748ec75c80fff21f7b24897244dd1c895df2e9a" }
12+
sketchlib-rust = { git = "https://github.com/ProjectASAP/sketchlib-rust", rev = "440427438fdaf3ac2298b53ee148f9e12a64ffcc" }
1213
clap = { version = "4.0", features = ["derive"] }
1314

1415
[dev-dependencies]

‎asap-common/sketch-core/report.md‎

Lines changed: 32 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,23 +1,17 @@
11
# Sketchlib Fidelity Report
22

3-
Compares the **legacy** sketch implementations in `sketch-core` vs the new **sketchlib-rust** backends for:
3+
Compares the **legacy** Count-Min Sketch implementation in `sketch-core` vs the new **sketchlib-rust** backend.
44

5-
- `CountMinSketch`
6-
- `CountMinSketchWithHeap` (Count-Min portion)
7-
- `KllSketch`
8-
- `HydraKllSketch` (via `KllSketch`)
5+
## Fidelity harness
96

10-
## Running Fidelity Tests
7+
The fidelity binary selects backends via CLI flags.
118

12-
The fidelity binary selects backends via CLI flags instead of environment variables.
9+
| Goal | Command |
10+
|-------------|---------------------------------------------------------------|
11+
| CMS sketchlib | `cargo run -p sketch-core --bin sketchlib_fidelity -- --cms-impl sketchlib` |
12+
| CMS legacy | `cargo run -p sketch-core --bin sketchlib_fidelity -- --cms-impl legacy` |
1313

14-
| Goal | Command |
15-
|--------------------------|--------------------------------------------------------------------------------------------------------------|
16-
| Default (all sketchlib) | `cargo run -p sketch-core --bin sketchlib_fidelity` |
17-
| All legacy | `cargo run -p sketch-core --bin sketchlib_fidelity -- --cms-impl legacy --kll-impl legacy --cmwh-impl legacy` |
18-
| Legacy KLL only | `cargo run -p sketch-core --bin sketchlib_fidelity -- --cms-impl sketchlib --kll-impl legacy --cmwh-impl sketchlib` |
19-
20-
## Unit Tests
14+
## Unit tests
2115

2216
Unit tests always run with **legacy** backends enabled (the test ctor calls
2317
`force_legacy_mode_for_tests()`), so you only need:
@@ -28,4 +22,27 @@ cargo test -p sketch-core
2822

2923
## Results
3024

31-
Fidelity results will be added as sketch implementations are integrated in subsequent PRs.
25+
### CountMinSketch (accuracy vs exact counts)
26+
27+
#### depth=3
28+
29+
| width | n | domain | Mode | Pearson corr | MAPE (%) | RMSE (%) |
30+
|-------|--------|--------|----------------|----------------|----------|----------|
31+
| 1024 | 100000 | 1000 | Legacy | 0.9998451189 | 24.48 | 52.76 |
32+
| 1024 | 100000 | 1000 | sketchlib-rust | 0.9998387103 | 24.36 | 54.11 |
33+
34+
#### depth=5
35+
36+
| width | n | domain | Mode | Pearson corr | MAPE (%) | RMSE (%) |
37+
|-------|--------|--------|----------------|----------------|----------|----------|
38+
| 2048 | 200000 | 2000 | Legacy | 0.9999733814 | 8.75 | 29.94 |
39+
| 2048 | 200000 | 2000 | sketchlib-rust | 0.9999744627 | 8.37 | 28.84 |
40+
| 2048 | 50000 | 500 | Legacy | 1.0000000000 | 0.00 | 0.00 |
41+
| 2048 | 50000 | 500 | sketchlib-rust | 1.0000000000 | 0.00 | 0.00 |
42+
43+
#### depth=7
44+
45+
| width | n | domain | Mode | Pearson corr | MAPE (%) | RMSE (%) |
46+
|-------|--------|--------|----------------|----------------|----------|----------|
47+
| 4096 | 200000 | 2000 | Legacy | 0.9999993694 | 0.20 | 3.69 |
48+
| 4096 | 200000 | 2000 | sketchlib-rust | 0.9999993499 | 0.21 | 4.27 |

‎asap-common/sketch-core/src/bin/sketchlib_fidelity.rs‎

Lines changed: 89 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,9 @@
1-
// Scaffold for fidelity benchmarks; helpers used in later PRs when sketch types are integrated.
1+
// Fidelity benchmarks comparing legacy vs sketchlib implementations across sketch types.
22
#![allow(dead_code)]
33

44
use clap::Parser;
55
use sketch_core::config::{self, ImplMode};
6+
use sketch_core::count_min::CountMinSketch;
67

78
#[derive(Clone)]
89
struct Lcg64 {
@@ -97,13 +98,58 @@ fn rank_fraction(sorted: &[f64], x: f64) -> f64 {
9798
(idx as f64) / (sorted.len() as f64)
9899
}
99100

101+
// --- CountMinSketch parameter sets and runner ---
102+
103+
struct CmsParams {
104+
depth: usize,
105+
width: usize,
106+
n: usize,
107+
domain: usize,
108+
}
109+
110+
struct CmsResult {
111+
pearson: f64,
112+
mape: f64,
113+
rmse: f64,
114+
}
115+
116+
fn run_countmin_once(seed: u64, p: &CmsParams) -> CmsResult {
117+
let mut rng = Lcg64::new(seed);
118+
let mut exact: Vec<f64> = vec![0.0; p.domain];
119+
let mut cms = CountMinSketch::new(p.depth, p.width);
120+
121+
for _ in 0..p.n {
122+
let r = rng.next_u64();
123+
let key_id = if (r & 0xFF) < 200 {
124+
(r as usize) % 20
125+
} else {
126+
(r as usize) % p.domain
127+
};
128+
let key = format!("k{key_id}");
129+
cms.update(&key, 1.0);
130+
exact[key_id] += 1.0;
131+
}
132+
133+
let mut est: Vec<f64> = Vec::with_capacity(p.domain);
134+
for key_id in 0..p.domain {
135+
let key = format!("k{key_id}");
136+
est.push(cms.query_key(&key));
137+
}
138+
139+
CmsResult {
140+
pearson: pearson_corr(&exact, &est),
141+
mape: mape(&exact, &est),
142+
rmse: rmse_percentage(&exact, &est),
143+
}
144+
}
145+
100146
#[derive(Parser)]
101147
struct Args {
102-
#[arg(long, value_enum, default_value_t = sketch_core::config::DEFAULT_IMPL_MODE)]
148+
#[arg(long, value_enum, default_value_t = sketch_core::config::DEFAULT_CMS_IMPL)]
103149
cms_impl: ImplMode,
104-
#[arg(long, value_enum, default_value_t = sketch_core::config::DEFAULT_IMPL_MODE)]
150+
#[arg(long, value_enum, default_value_t = sketch_core::config::DEFAULT_KLL_IMPL)]
105151
kll_impl: ImplMode,
106-
#[arg(long, value_enum, default_value_t = sketch_core::config::DEFAULT_IMPL_MODE)]
152+
#[arg(long, value_enum, default_value_t = sketch_core::config::DEFAULT_CMWH_IMPL)]
107153
cmwh_impl: ImplMode,
108154
}
109155

@@ -112,6 +158,7 @@ fn main() {
112158
config::configure(args.cms_impl, args.kll_impl, args.cmwh_impl)
113159
.expect("sketch backend already initialised");
114160

161+
let seed = 0xC0FFEE_u64;
115162
let mode = if matches!(args.cms_impl, ImplMode::Legacy)
116163
|| matches!(args.kll_impl, ImplMode::Legacy)
117164
|| matches!(args.cmwh_impl, ImplMode::Legacy)
@@ -121,7 +168,42 @@ fn main() {
121168
"sketchlib-rust"
122169
};
123170

124-
println!("# Sketchlib Fidelity Report ({})", mode);
125-
println!();
126-
println!("Fidelity tests will be added as sketch implementations are integrated.");
171+
// CountMinSketch: multiple (depth, width, n, domain)
172+
let cms_param_sets: Vec<CmsParams> = vec![
173+
CmsParams {
174+
depth: 3,
175+
width: 1024,
176+
n: 100_000,
177+
domain: 1000,
178+
},
179+
CmsParams {
180+
depth: 5,
181+
width: 2048,
182+
n: 200_000,
183+
domain: 2000,
184+
},
185+
CmsParams {
186+
depth: 7,
187+
width: 4096,
188+
n: 200_000,
189+
domain: 2000,
190+
},
191+
CmsParams {
192+
depth: 5,
193+
width: 2048,
194+
n: 50_000,
195+
domain: 500,
196+
},
197+
];
198+
199+
println!("## CountMinSketch ({mode})");
200+
println!("| depth | width | n_updates | domain | Pearson corr | MAPE (%) | RMSE (%) |");
201+
println!("|-------|-------|------------|--------|--------------|----------|----------|");
202+
for p in &cms_param_sets {
203+
let r = run_countmin_once(seed, p);
204+
println!(
205+
"| {} | {} | {} | {} | {:.10} | {:.6} | {:.6} |",
206+
p.depth, p.width, p.n, p.domain, r.pearson, r.mape, r.rmse
207+
);
208+
}
127209
}

‎asap-common/sketch-core/src/config.rs‎

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -9,30 +9,34 @@ pub enum ImplMode {
99
Sketchlib,
1010
}
1111

12-
/// Default backend when not explicitly configured (e.g. for binaries that don't pass CLI args).
13-
/// Single source of truth for sketch backend defaults.
12+
/// Global default when impl mode is not explicitly configured (e.g. env var parsing).
1413
pub const DEFAULT_IMPL_MODE: ImplMode = ImplMode::Legacy;
1514

15+
/// Per-backend defaults. Used when configure() has not been called.
16+
pub const DEFAULT_CMS_IMPL: ImplMode = ImplMode::Sketchlib;
17+
pub const DEFAULT_KLL_IMPL: ImplMode = ImplMode::Legacy;
18+
pub const DEFAULT_CMWH_IMPL: ImplMode = ImplMode::Legacy;
19+
1620
static COUNTMIN_MODE: OnceLock<ImplMode> = OnceLock::new();
1721

1822
/// Returns true if Count-Min operations should use sketchlib-rust internally.
1923
pub fn use_sketchlib_for_count_min() -> bool {
20-
*COUNTMIN_MODE.get_or_init(|| DEFAULT_IMPL_MODE) == ImplMode::Sketchlib
24+
*COUNTMIN_MODE.get_or_init(|| DEFAULT_CMS_IMPL) == ImplMode::Sketchlib
2125
}
2226

2327
static KLL_MODE: OnceLock<ImplMode> = OnceLock::new();
2428

2529
/// Returns true if KLL operations should use sketchlib-rust internally.
2630
pub fn use_sketchlib_for_kll() -> bool {
27-
*KLL_MODE.get_or_init(|| DEFAULT_IMPL_MODE) == ImplMode::Sketchlib
31+
*KLL_MODE.get_or_init(|| DEFAULT_KLL_IMPL) == ImplMode::Sketchlib
2832
}
2933

3034
static COUNTMIN_WITH_HEAP_MODE: OnceLock<ImplMode> = OnceLock::new();
3135

3236
/// Returns true if Count-Min-With-Heap operations should use sketchlib-rust internally for the
3337
/// Count-Min portion.
3438
pub fn use_sketchlib_for_count_min_with_heap() -> bool {
35-
*COUNTMIN_WITH_HEAP_MODE.get_or_init(|| DEFAULT_IMPL_MODE) == ImplMode::Sketchlib
39+
*COUNTMIN_WITH_HEAP_MODE.get_or_init(|| DEFAULT_CMWH_IMPL) == ImplMode::Sketchlib
3640
}
3741

3842
/// Set backend modes for all sketch types. Call once at process startup,

0 commit comments

Comments
 (0)