From 07d30dbb3488780c3cd66cf8aebc8a7f94c25732 Mon Sep 17 00:00:00 2001 From: Zeying Zhu Date: Fri, 1 May 2026 14:44:55 -0400 Subject: [PATCH 1/2] refactor: retire sketch-core mirror --- Cargo.lock | 183 +----- Cargo.toml | 2 - asap-common/sketch-core/Cargo.toml | 16 - asap-common/sketch-core/report.md | 125 ---- .../sketch-core/src/bin/sketchlib_fidelity.rs | 496 --------------- asap-common/sketch-core/src/config.rs | 83 --- asap-common/sketch-core/src/count_min.rs | 548 ---------------- .../sketch-core/src/count_min_sketchlib.rs | 59 -- .../sketch-core/src/count_min_with_heap.rs | 597 ------------------ .../src/count_min_with_heap_sketchlib.rs | 109 ---- asap-common/sketch-core/src/count_sketch.rs | 274 -------- asap-common/sketch-core/src/dd_sketch.rs | 588 ----------------- .../sketch-core/src/delta_set_aggregator.rs | 71 --- asap-common/sketch-core/src/hll_sketch.rs | 296 --------- asap-common/sketch-core/src/hydra_kll.rs | 295 --------- asap-common/sketch-core/src/kll.rs | 366 ----------- asap-common/sketch-core/src/kll_sketchlib.rs | 36 -- asap-common/sketch-core/src/lib.rs | 19 - asap-common/sketch-core/src/set_aggregator.rs | 152 ----- asap-query-engine/Cargo.toml | 3 +- asap-query-engine/Dockerfile | 3 - asap-query-engine/src/drivers/ingest/otel.rs | 72 +-- asap-query-engine/src/lib.rs | 10 +- asap-query-engine/src/main.rs | 16 +- .../src/precompute_engine/ingest_handler.rs | 48 +- .../src/precompute_engine/worker.rs | 2 +- .../count_min_sketch_accumulator.rs | 20 +- .../count_min_sketch_with_heap_accumulator.rs | 9 +- .../count_sketch_accumulator.rs | 18 +- .../datasketches_kll_accumulator.rs | 7 +- .../dd_sketch_accumulator.rs | 14 +- .../delta_set_aggregator_accumulator.rs | 5 +- .../hll_sketch_accumulator.rs | 16 +- .../hydra_kll_accumulator.rs | 7 +- .../set_aggregator_accumulator.rs | 7 +- .../src/stores/promsketch_store/query.rs | 10 +- .../src/stores/promsketch_store/series.rs | 4 +- .../tests/test_utilities/engine_factories.rs | 2 +- .../tests/e2e_modified_otlp_sketch_path.rs | 5 +- .../tests/e2e_precompute_equivalence.rs | 2 +- 40 files changed, 127 insertions(+), 4468 deletions(-) delete mode 100644 asap-common/sketch-core/Cargo.toml delete mode 100644 asap-common/sketch-core/report.md delete mode 100644 asap-common/sketch-core/src/bin/sketchlib_fidelity.rs delete mode 100644 asap-common/sketch-core/src/config.rs delete mode 100644 asap-common/sketch-core/src/count_min.rs delete mode 100644 asap-common/sketch-core/src/count_min_sketchlib.rs delete mode 100644 asap-common/sketch-core/src/count_min_with_heap.rs delete mode 100644 asap-common/sketch-core/src/count_min_with_heap_sketchlib.rs delete mode 100644 asap-common/sketch-core/src/count_sketch.rs delete mode 100644 asap-common/sketch-core/src/dd_sketch.rs delete mode 100644 asap-common/sketch-core/src/delta_set_aggregator.rs delete mode 100644 asap-common/sketch-core/src/hll_sketch.rs delete mode 100644 asap-common/sketch-core/src/hydra_kll.rs delete mode 100644 asap-common/sketch-core/src/kll.rs delete mode 100644 asap-common/sketch-core/src/kll_sketchlib.rs delete mode 100644 asap-common/sketch-core/src/lib.rs delete mode 100644 asap-common/sketch-core/src/set_aggregator.rs diff --git a/Cargo.lock b/Cargo.lock index 87215c65..37af1546 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -423,21 +423,21 @@ dependencies = [ [[package]] name = "asap_sketchlib" version = "0.1.0" -source = "git+https://github.com/ProjectASAP/asap_sketchlib#2442d95ef8bc8470f9909728fadc72b5e082139d" +source = "git+https://github.com/ProjectASAP/asap_sketchlib?branch=refactor%2Fadopt-sketch-core-modules#63051a830e029679db9fe8da3cd4488773a18bbe" dependencies = [ "bytes", "clap 4.6.0", - "core_affinity", - "crossbeam-channel", - "pcap", + "dsrs", "prost", "prost-build", - "rand 0.9.2", + "protoc-bin-vendored", + "rand 0.9.4", "rmp-serde", "serde", "serde-big-array", "smallvec", "twox-hash 2.1.2", + "xxhash-rust", ] [[package]] @@ -971,17 +971,6 @@ version = "0.8.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b" -[[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 = "cpufeatures" version = "0.2.17" @@ -1534,7 +1523,7 @@ dependencies = [ "itertools 0.13.0", "log", "paste", - "petgraph 0.6.5", + "petgraph", ] [[package]] @@ -1714,17 +1703,6 @@ version = "1.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" -[[package]] -name = "errno" -version = "0.2.8" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f639046355ee4f37944e44f60642c6f3a7efa3cf6b78c78a0d989a8ce6c396a1" -dependencies = [ - "errno-dragonfly", - "libc", - "winapi", -] - [[package]] name = "errno" version = "0.3.14" @@ -1732,17 +1710,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.61.2", -] - -[[package]] -name = "errno-dragonfly" -version = "0.1.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "aa68f1b12764fab894d2755d2518754e71b4fd80ecfb822714a1206c2aab39bf" -dependencies = [ - "cc", - "libc", + "windows-sys 0.52.0", ] [[package]] @@ -1786,12 +1754,6 @@ version = "0.4.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0ce7134b9999ecaf8bcd65542e436736ef32ddca1b3e06094cb6ec5755203b80" -[[package]] -name = "fixedbitset" -version = "0.5.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1d674e81391d1e1ab681a28d99df07927c6d4aa5b027d7da16ba32d1d21ecd99" - [[package]] name = "flatbuffers" version = "24.12.23" @@ -2476,7 +2438,7 @@ checksum = "3640c1c38b8e4e43584d8df18be5fc6b0aa314ce6ebf51b53313d4306cca8e46" dependencies = [ "hermit-abi 0.5.2", "libc", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -2503,15 +2465,6 @@ dependencies = [ "either", ] -[[package]] -name = "itertools" -version = "0.14.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2b192c782037fadd9cfa75548310488aabdbf3d2da73885b31bd0abd03351285" -dependencies = [ - "either", -] - [[package]] name = "itoa" version = "1.0.18" @@ -2530,7 +2483,7 @@ dependencies = [ "portable-atomic", "portable-atomic-util", "serde_core", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -2656,16 +2609,6 @@ version = "0.2.184" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "48f5d2a454e16a5ea0f4ced81bd44e4cfc7bd3a507b61887c99fd3538b28e4af" -[[package]] -name = "libloading" -version = "0.6.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "351a32417a12d5f7e82c368a66781e307834dae04c6ce0cd4456d52989229883" -dependencies = [ - "cfg-if", - "winapi", -] - [[package]] name = "libm" version = "0.2.16" @@ -2913,7 +2856,7 @@ version = "0.50.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.59.0", ] [[package]] @@ -3180,21 +3123,6 @@ version = "1.0.15" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" -[[package]] -name = "pcap" -version = "1.3.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "99e935fc73d54a89fff576526c2ccd42bbf8247aae05b358693475b14fd4ff79" -dependencies = [ - "bitflags 1.3.2", - "errno 0.2.8", - "libc", - "libloading", - "pkg-config", - "regex", - "windows-sys 0.36.1", -] - [[package]] name = "percent-encoding" version = "2.3.2" @@ -3207,17 +3135,7 @@ version = "0.6.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b4c5cc86750666a3ed20bdaf5ca2a0344f9c67674cae0515bec2da16fbaa47db" dependencies = [ - "fixedbitset 0.4.2", - "indexmap 2.13.1", -] - -[[package]] -name = "petgraph" -version = "0.7.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3672b37090dbd86368a4145bc067582552b29c27377cad4e0a306c97f9bd7772" -dependencies = [ - "fixedbitset 0.5.7", + "fixedbitset", "indexmap 2.13.1", ] @@ -3479,11 +3397,11 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "be769465445e8c1474e9c5dac2018218498557af32d9ed057325ec9a41ae81bf" dependencies = [ "heck 0.5.0", - "itertools 0.14.0", + "itertools 0.10.5", "log", "multimap", "once_cell", - "petgraph 0.7.1", + "petgraph", "prettyplease", "prost", "prost-types", @@ -3499,7 +3417,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8a56d757972c98b346a9b766e3f02746cde6dd1cd1d1d563472929fdd74bec4d" dependencies = [ "anyhow", - "itertools 0.14.0", + "itertools 0.10.5", "proc-macro2", "quote", "syn 2.0.117", @@ -3638,7 +3556,6 @@ dependencies = [ "serde", "serde_json", "serde_yaml", - "sketch-core", "snap", "sql_utilities", "sqlparser 0.59.0", @@ -3691,9 +3608,9 @@ dependencies = [ [[package]] name = "rand" -version = "0.9.2" +version = "0.9.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6db2770f06117d490610c7488547d543617b21bfa07796d7a12f6f1bd53850d1" +checksum = "44c5af06bb1b7d3216d91932aed5265164bf384dc89cd6ba05cf59a35f5f76ea" dependencies = [ "rand_chacha 0.9.0", "rand_core 0.9.5", @@ -3964,10 +3881,10 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b6fe4565b9518b83ef4f91bb47ce29620ca828bd32cb7e408f0062e9930ba190" dependencies = [ "bitflags 2.11.0", - "errno 0.3.14", + "errno", "libc", "linux-raw-sys", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -4176,7 +4093,7 @@ version = "1.4.8" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c4db69cba1110affc0e9f7bcd48bbf87b3f4fc7c61fc9155afd4c469eb3d6c1b" dependencies = [ - "errno 0.3.14", + "errno", "libc", ] @@ -4192,19 +4109,6 @@ version = "1.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b2aa850e253778c88a04c3d7323b043aeda9d3e30d5971937c1855769763678e" -[[package]] -name = "sketch-core" -version = "0.1.0" -dependencies = [ - "asap_sketchlib", - "clap 4.6.0", - "ctor", - "dsrs", - "rmp-serde", - "serde", - "xxhash-rust", -] - [[package]] name = "slab" version = "0.4.12" @@ -4488,7 +4392,7 @@ dependencies = [ "getrandom 0.4.2", "once_cell", "rustix", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -4928,7 +4832,7 @@ version = "2.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9ea3136b675547379c4bd395ca6b938e5ad3c3d20fad76e7fe85f9e0d011419c" dependencies = [ - "rand 0.9.2", + "rand 0.9.4", ] [[package]] @@ -5235,7 +5139,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.48.0", ] [[package]] @@ -5303,19 +5207,6 @@ dependencies = [ "windows-link", ] -[[package]] -name = "windows-sys" -version = "0.36.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ea04155a16a59f9eab786fe12a4a450e75cdb175f9e0d80da1e17db09f55b8d2" -dependencies = [ - "windows_aarch64_msvc 0.36.1", - "windows_i686_gnu 0.36.1", - "windows_i686_msvc 0.36.1", - "windows_x86_64_gnu 0.36.1", - "windows_x86_64_msvc 0.36.1", -] - [[package]] name = "windows-sys" version = "0.48.0" @@ -5395,12 +5286,6 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" -[[package]] -name = "windows_aarch64_msvc" -version = "0.36.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9bb8c3fd39ade2d67e9874ac4f3db21f0d710bee00fe7cab16949ec184eeaa47" - [[package]] name = "windows_aarch64_msvc" version = "0.48.5" @@ -5413,12 +5298,6 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" -[[package]] -name = "windows_i686_gnu" -version = "0.36.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "180e6ccf01daf4c426b846dfc66db1fc518f074baa793aa7d9b9aaeffad6a3b6" - [[package]] name = "windows_i686_gnu" version = "0.48.5" @@ -5437,12 +5316,6 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" -[[package]] -name = "windows_i686_msvc" -version = "0.36.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e2e7917148b2812d1eeafaeb22a97e4813dfa60a3f8f78ebe204bcc88f12f024" - [[package]] name = "windows_i686_msvc" version = "0.48.5" @@ -5455,12 +5328,6 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" -[[package]] -name = "windows_x86_64_gnu" -version = "0.36.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4dcd171b8776c41b97521e5da127a2d86ad280114807d0b2ab1e462bc764d9e1" - [[package]] name = "windows_x86_64_gnu" version = "0.48.5" @@ -5485,12 +5352,6 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" -[[package]] -name = "windows_x86_64_msvc" -version = "0.36.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c811ca4a8c853ef420abd8592ba53ddbbac90410fab6903b3e79972a631f7680" - [[package]] name = "windows_x86_64_msvc" version = "0.48.5" diff --git a/Cargo.toml b/Cargo.toml index c47a9da0..3dc9e28f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,7 +1,6 @@ [workspace] resolver = "2" members = [ - "asap-common/sketch-core", "asap-common/dependencies/rs/promql_utilities", "asap-common/dependencies/rs/sql_utilities", "asap-common/dependencies/rs/elastic_dsl_utilities", @@ -34,7 +33,6 @@ tokio = { version = "1.0", features = ["full"] } arc-swap = "1.7" # Internal crates -sketch-core = { path = "asap-common/sketch-core" } promql_utilities = { path = "asap-common/dependencies/rs/promql_utilities" } sql_utilities = { path = "asap-common/dependencies/rs/sql_utilities" } asap_types = { path = "asap-common/dependencies/rs/asap_types" } diff --git a/asap-common/sketch-core/Cargo.toml b/asap-common/sketch-core/Cargo.toml deleted file mode 100644 index fe9cf24b..00000000 --- a/asap-common/sketch-core/Cargo.toml +++ /dev/null @@ -1,16 +0,0 @@ -[package] -name = "sketch-core" -version.workspace = true -edition.workspace = true - -[dependencies] -serde.workspace = true -# sketch-core-specific, keep pinned -rmp-serde = "1.1" -xxhash-rust = { version = "0.8", features = ["xxh32"] } -dsrs = { git = "https://github.com/ProjectASAP/datasketches-rs", rev = "d748ec75c80fff21f7b24897244dd1c895df2e9a" } -asap_sketchlib = { git = "https://github.com/ProjectASAP/asap_sketchlib" } -clap = { version = "4.0", features = ["derive"] } - -[dev-dependencies] -ctor = "0.2" diff --git a/asap-common/sketch-core/report.md b/asap-common/sketch-core/report.md deleted file mode 100644 index 9dfde749..00000000 --- a/asap-common/sketch-core/report.md +++ /dev/null @@ -1,125 +0,0 @@ -# Report - -Compares the **legacy** sketch implementations in `sketch-core` vs the **asap_sketchlib** backends (Count-Min Sketch, Count-Min-With-Heap, KLL, HydraKLL). - -## Fidelity harness - -The fidelity binary selects backends via CLI flags (`--cms-impl`, `--kll-impl`, `--cmwh-impl`). - -| Goal | Command | -|--------------------------|--------------------------------------------------------------------------------------------------------------| -| Default (all sketchlib) | `cargo run -p sketch-core --bin sketchlib_fidelity` | -| All legacy | `cargo run -p sketch-core --bin sketchlib_fidelity -- --cms-impl legacy --kll-impl legacy --cmwh-impl legacy` | -| Legacy KLL only | `cargo run -p sketch-core --bin sketchlib_fidelity -- --cms-impl sketchlib --kll-impl legacy --cmwh-impl sketchlib` | -| CMS sketchlib only | `cargo run -p sketch-core --bin sketchlib_fidelity -- --cms-impl sketchlib` | -| CMS legacy only | `cargo run -p sketch-core --bin sketchlib_fidelity -- --cms-impl legacy` | - -## Unit tests - -Unit tests always run with **legacy** backends enabled (the test ctor calls -`force_legacy_mode_for_tests()`), so you only need: - -```bash -cargo test -p sketch-core -``` - -## Results - -### CountMinSketch (accuracy vs exact counts) - -#### depth=3 - -| width | n | domain | Mode | Pearson corr | MAPE (%) | RMSE (%) | -|-------|--------|--------|----------------|----------------|----------|----------| -| 1024 | 100000 | 1000 | Legacy | 0.9998451189 | 24.48 | 52.76 | -| 1024 | 100000 | 1000 | asap_sketchlib | 0.9998387103 | 24.36 | 54.11 | - -#### depth=5 - -| width | n | domain | Mode | Pearson corr | MAPE (%) | RMSE (%) | -|-------|--------|--------|----------------|----------------|----------|----------| -| 2048 | 200000 | 2000 | Legacy | 0.9999733814 | 8.75 | 29.94 | -| 2048 | 200000 | 2000 | asap_sketchlib | 0.9999744627 | 8.37 | 28.84 | -| 2048 | 50000 | 500 | Legacy | 1.0000000000 | 0.00 | 0.00 | -| 2048 | 50000 | 500 | asap_sketchlib | 1.0000000000 | 0.00 | 0.00 | - -#### depth=7 - -| width | n | domain | Mode | Pearson corr | MAPE (%) | RMSE (%) | -|-------|--------|--------|----------------|----------------|----------|----------| -| 4096 | 200000 | 2000 | Legacy | 0.9999993694 | 0.20 | 3.69 | -| 4096 | 200000 | 2000 | asap_sketchlib | 0.9999993499 | 0.21 | 4.27 | - ---- - -### CountMinSketchWithHeap (top-k + CMS accuracy on exact top-k) - -The heap is maintained by local updates; recall is measured against the **true** top-k at the end of the stream. - -#### depth=3 - -| width | n | domain | heap_size | Mode | Top-k recall | Pearson (top-k) | MAPE (%) | RMSE (%) | -|-------|--------|--------|-----------|----------------|--------------|-----------------|----------|----------| -| 1024 | 100000 | 1000 | 10 | Legacy | 0.40 | 0.9571 | 0.174 | 0.319 | -| 1024 | 100000 | 1000 | 10 | asap_sketchlib | 0.80 | 1.0000 | 0.000 | 0.000 | - -#### depth=5 - -| width | n | domain | heap_size | Mode | Top-k recall | Pearson (top-k) | MAPE (%) | RMSE (%) | -|-------|--------|--------|-----------|----------------|--------------|-----------------|----------|----------| -| 2048 | 200000 | 2000 | 20 | Legacy | 0.60 | 0.9964 | 0.045 | 0.101 | -| 2048 | 200000 | 2000 | 20 | asap_sketchlib | 1.00 | 0.9982 | 0.021 | 0.067 | -| 2048 | 200000 | 2000 | 50 | Legacy | 0.40 | 0.9999983 | 5.60 | 16.49 | -| 2048 | 200000 | 2000 | 50 | asap_sketchlib | 0.48 | 0.9999990 | 3.90 | 12.95 | - ---- - -### KllSketch (quantiles, absolute rank error) - -For each quantile \(q\), we compute the sketch estimate `est_value`, then: -`abs_rank_error = |rank_fraction(exact_sorted_values, est_value) - q|`. - -#### k=20 - -| n_updates | Mode | q=0.5 | q=0.9 | q=0.99 | -|-----------|----------------|---------|---------|---------| -| 200000 | Legacy | 0.0104 | 0.0145 | 0.0028 | -| 200000 | asap_sketchlib | 0.0275 | 0.0470 | 0.0061 | -| 50000 | Legacy | 0.0131 | 0.0091 | 0.0054 | -| 50000 | asap_sketchlib | 0.0110 | 0.0116 | 0.0031 | - -#### k=50 - -| n_updates | Mode | q=0.5 | q=0.9 | q=0.99 | -|-----------|----------------|---------|---------|---------| -| 200000 | Legacy | 0.0013 | 0.0021 | 0.0012 | -| 200000 | asap_sketchlib | 0.0101 | 0.0044 | 0.0074 | - -#### k=200 - -| n_updates | Mode | q=0.5 | q=0.9 | q=0.99 | -|-----------|----------------|---------|---------|---------| -| 200000 | Legacy | 0.0021 | 0.0036 | 0.0000 | -| 200000 | asap_sketchlib | 0.0015 | 0.0001 | 0.0002 | - ---- - -### HydraKllSketch (per-key quantiles, mean/max absolute rank error across 50 keys) - -#### rows=2, cols=64 - -| k | n | domain | Mode | q=0.5 (mean / max) | q=0.9 (mean / max) | -|-----|--------|--------|----------------|--------------------|--------------------| -| 20 | 200000 | 200 | Legacy | 0.0170 / 0.0546 | 0.0165 / 0.0452 | -| 20 | 200000 | 200 | asap_sketchlib | 0.0254 / 0.0629 | 0.0546 / 0.0942 | - -#### rows=3, cols=128 - -| k | n | domain | Mode | q=0.5 (mean / max) | q=0.9 (mean / max) | -|-----|--------|--------|----------------|--------------------|--------------------| -| 20 | 200000 | 200 | Legacy | 0.0166 / 0.0591 | 0.0114 / 0.0304 | -| 20 | 200000 | 200 | asap_sketchlib | 0.0216 / 0.0534 | 0.0238 / 0.1087 | -| 50 | 200000 | 200 | Legacy | 0.0099 / 0.0352 | 0.0087 / 0.0330 | -| 50 | 200000 | 200 | asap_sketchlib | 0.0119 / 0.0458 | 0.0119 / 0.0296 | -| 20 | 100000 | 100 | Legacy | 0.0141 / 0.0574 | 0.0149 / 0.0471 | -| 20 | 100000 | 100 | asap_sketchlib | 0.0202 / 0.0621 | 0.0287 / 0.0779 | diff --git a/asap-common/sketch-core/src/bin/sketchlib_fidelity.rs b/asap-common/sketch-core/src/bin/sketchlib_fidelity.rs deleted file mode 100644 index ca95cb6f..00000000 --- a/asap-common/sketch-core/src/bin/sketchlib_fidelity.rs +++ /dev/null @@ -1,496 +0,0 @@ -// Fidelity benchmarks comparing legacy vs sketchlib implementations across sketch types. -#![allow(dead_code)] - -use std::collections::HashMap; - -use clap::Parser; -use sketch_core::config::{self, ImplMode}; -use sketch_core::count_min::CountMinSketch; -use sketch_core::count_min_with_heap::CountMinSketchWithHeap; -use sketch_core::hydra_kll::HydraKllSketch; -use sketch_core::kll::KllSketch; - -#[derive(Clone)] -struct Lcg64 { - state: u64, -} - -impl Lcg64 { - fn new(seed: u64) -> Self { - Self { state: seed } - } - - fn next_u64(&mut self) -> u64 { - self.state = self - .state - .wrapping_mul(6364136223846793005) - .wrapping_add(1442695040888963407); - self.state - } - - fn next_f64_0_1(&mut self) -> f64 { - let x = self.next_u64() >> 11; - (x as f64) / ((1u64 << 53) as f64) - } -} - -fn pearson_corr(exact: &[f64], est: &[f64]) -> f64 { - let n = exact.len().min(est.len()); - if n == 0 { - return f64::NAN; - } - let (mut sum_x, mut sum_y) = (0.0, 0.0); - for i in 0..n { - sum_x += exact[i]; - sum_y += est[i]; - } - let mean_x = sum_x / (n as f64); - let mean_y = sum_y / (n as f64); - let (mut num, mut den_x, mut den_y) = (0.0, 0.0, 0.0); - for i in 0..n { - let dx = exact[i] - mean_x; - let dy = est[i] - mean_y; - num += dx * dy; - den_x += dx * dx; - den_y += dy * dy; - } - if den_x == 0.0 || den_y == 0.0 { - return f64::NAN; - } - num / (den_x.sqrt() * den_y.sqrt()) -} - -fn mape(exact: &[f64], est: &[f64]) -> f64 { - let n = exact.len().min(est.len()); - let mut num = 0.0; - let mut denom = 0.0; - for i in 0..n { - if exact[i] == 0.0 { - continue; - } - num += ((exact[i] - est[i]) / exact[i]).abs(); - denom += 1.0; - } - if denom == 0.0 { - return if exact == est { 0.0 } else { f64::INFINITY }; - } - (num / denom) * 100.0 -} - -fn rmse_percentage(exact: &[f64], est: &[f64]) -> f64 { - let n = exact.len().min(est.len()); - let mut sum_sq = 0.0; - let mut denom = 0.0; - for i in 0..n { - if exact[i] == 0.0 { - continue; - } - let rel = (exact[i] - est[i]) / exact[i]; - sum_sq += rel * rel; - denom += 1.0; - } - if denom == 0.0 { - return if exact == est { 0.0 } else { f64::INFINITY }; - } - (sum_sq / denom).sqrt() * 100.0 -} - -#[derive(Parser)] -struct Args { - #[arg(long, value_enum, default_value_t = sketch_core::config::DEFAULT_CMS_IMPL)] - cms_impl: ImplMode, - #[arg(long, value_enum, default_value_t = sketch_core::config::DEFAULT_KLL_IMPL)] - kll_impl: ImplMode, - #[arg(long, value_enum, default_value_t = sketch_core::config::DEFAULT_CMWH_IMPL)] - cmwh_impl: ImplMode, -} - -fn rank_fraction(sorted: &[f64], x: f64) -> f64 { - if sorted.is_empty() { - return 0.0; - } - let idx = sorted.partition_point(|v| *v <= x); - (idx as f64) / (sorted.len() as f64) -} - -// --- CountMinSketch parameter sets and runner --- - -struct CmsParams { - depth: usize, - width: usize, - n: usize, - domain: usize, -} - -struct CmsResult { - pearson: f64, - mape: f64, - rmse: f64, -} - -fn run_countmin_once(seed: u64, p: &CmsParams) -> CmsResult { - let mut rng = Lcg64::new(seed); - let mut exact: Vec = vec![0.0; p.domain]; - let mut cms = CountMinSketch::new(p.depth, p.width); - - for _ in 0..p.n { - let r = rng.next_u64(); - let key_id = if (r & 0xFF) < 200 { - (r as usize) % 20 - } else { - (r as usize) % p.domain - }; - let key = format!("k{key_id}"); - cms.update(&key, 1.0); - exact[key_id] += 1.0; - } - - let mut est: Vec = Vec::with_capacity(p.domain); - for key_id in 0..p.domain { - let key = format!("k{key_id}"); - est.push(cms.query_key(&key)); - } - - CmsResult { - pearson: pearson_corr(&exact, &est), - mape: mape(&exact, &est), - rmse: rmse_percentage(&exact, &est), - } -} - -// --- CountMinSketchWithHeap --- - -struct CmwhParams { - depth: usize, - width: usize, - n: usize, - domain: usize, - heap_size: usize, -} - -struct CmwhResult { - topk_recall: f64, - pearson: f64, - mape: f64, - rmse: f64, -} - -fn run_countmin_with_heap_once(seed: u64, p: &CmwhParams) -> CmwhResult { - let mut rng = Lcg64::new(seed ^ 0xA5A5_A5A5); - let mut exact: Vec = vec![0.0; p.domain]; - let mut cms = CountMinSketchWithHeap::new(p.depth, p.width, p.heap_size); - - for _ in 0..p.n { - let r = rng.next_u64(); - let key_id = if (r & 0xFF) < 200 { - (r as usize) % 20 - } else { - (r as usize) % p.domain - }; - let key = format!("k{key_id}"); - cms.update(&key, 1.0); - exact[key_id] += 1.0; - } - - let mut exact_pairs: Vec<(usize, f64)> = exact.iter().copied().enumerate().collect(); - exact_pairs.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap()); - exact_pairs.truncate(p.heap_size); - - let exact_top: HashMap = exact_pairs - .into_iter() - .map(|(k, v)| (format!("k{k}"), v)) - .collect(); - - let mut est_vals = Vec::with_capacity(exact_top.len()); - let mut exact_vals = Vec::with_capacity(exact_top.len()); - let mut hit = 0usize; - for item in cms.topk_heap_items() { - if exact_top.contains_key(&item.key) { - hit += 1; - } - } - for (k, v) in &exact_top { - exact_vals.push(*v); - est_vals.push(cms.query_key(k)); - } - - CmwhResult { - topk_recall: (hit as f64) / (p.heap_size as f64), - pearson: pearson_corr(&exact_vals, &est_vals), - mape: mape(&exact_vals, &est_vals), - rmse: rmse_percentage(&exact_vals, &est_vals), - } -} - -// --- KllSketch --- - -struct KllParams { - k: u16, - n: usize, -} - -struct KllResult { - rank_err_50: f64, - rank_err_90: f64, - rank_err_99: f64, -} - -fn run_kll_once(seed: u64, p: &KllParams) -> KllResult { - let mut rng = Lcg64::new(seed ^ 0x1234_5678); - let mut values: Vec = Vec::with_capacity(p.n); - let mut sk = KllSketch::new(p.k); - - for _ in 0..p.n { - let v = rng.next_f64_0_1() * 1_000_000.0; - values.push(v); - sk.update(v); - } - - values.sort_by(f64::total_cmp); - let qs = [0.5, 0.9, 0.99]; - let rank_err = |q: f64| (rank_fraction(&values, sk.get_quantile(q)) - q).abs(); - - KllResult { - rank_err_50: rank_err(qs[0]), - rank_err_90: rank_err(qs[1]), - rank_err_99: rank_err(qs[2]), - } -} - -// --- HydraKllSketch --- - -struct HydraKllParams { - rows: usize, - cols: usize, - k: u16, - n: usize, - domain: usize, - eval_keys: usize, -} - -struct HydraKllResult { - mean_50: f64, - max_50: f64, - mean_90: f64, - max_90: f64, -} - -fn run_hydra_kll_once(seed: u64, p: &HydraKllParams) -> HydraKllResult { - let mut rng = Lcg64::new(seed ^ 0xDEAD_BEEF); - let mut hydra = HydraKllSketch::new(p.rows, p.cols, p.k); - let mut exact: HashMap> = HashMap::new(); - - for _ in 0..p.n { - let r = rng.next_u64(); - let key_id = if (r & 0xFF) < 200 { - (r as usize) % 20 - } else { - (r as usize) % p.domain - }; - let key = format!("k{key_id}"); - let v = rng.next_f64_0_1() * 1_000_000.0; - hydra.update(&key, v); - exact.entry(key).or_default().push(v); - } - - let mut keys: Vec = exact.keys().cloned().collect(); - keys.sort(); - keys.truncate(p.eval_keys); - - let mut mean_50 = 0.0f64; - let mut max_50 = 0.0f64; - let mut mean_90 = 0.0f64; - let mut max_90 = 0.0f64; - let nk = keys.len() as f64; - for key in &keys { - let mut vals = exact.get(key).cloned().unwrap_or_default(); - vals.sort_by(f64::total_cmp); - for (q, mean_ref, max_ref) in [ - (0.5, &mut mean_50, &mut max_50), - (0.9, &mut mean_90, &mut max_90), - ] { - let est = hydra.query(key, q); - let err = (rank_fraction(&vals, est) - q).abs(); - *mean_ref += err; - if err > *max_ref { - *max_ref = err; - } - } - } - mean_50 /= nk; - mean_90 /= nk; - - HydraKllResult { - mean_50, - max_50, - mean_90, - max_90, - } -} - -fn main() { - let args = Args::parse(); - config::configure(args.cms_impl, args.kll_impl, args.cmwh_impl) - .expect("sketch backend already initialised"); - - let seed = 0xC0FFEE_u64; - let cms_mode = if matches!(args.cms_impl, ImplMode::Legacy) { - "Legacy" - } else { - "asap_sketchlib" - }; - let cmwh_mode = if matches!(args.cmwh_impl, ImplMode::Legacy) { - "Legacy" - } else { - "asap_sketchlib" - }; - let kll_mode = if matches!(args.kll_impl, ImplMode::Legacy) { - "Legacy" - } else { - "asap_sketchlib" - }; - - // CountMinSketch: multiple (depth, width, n, domain) - let cms_param_sets: Vec = vec![ - CmsParams { - depth: 3, - width: 1024, - n: 100_000, - domain: 1000, - }, - CmsParams { - depth: 5, - width: 2048, - n: 200_000, - domain: 2000, - }, - CmsParams { - depth: 7, - width: 4096, - n: 200_000, - domain: 2000, - }, - CmsParams { - depth: 5, - width: 2048, - n: 50_000, - domain: 500, - }, - ]; - - println!("## CountMinSketch ({cms_mode})"); - println!("| depth | width | n_updates | domain | Pearson corr | MAPE (%) | RMSE (%) |"); - println!("|-------|-------|------------|--------|--------------|----------|----------|"); - for p in &cms_param_sets { - let r = run_countmin_once(seed, p); - println!( - "| {} | {} | {} | {} | {:.10} | {:.6} | {:.6} |", - p.depth, p.width, p.n, p.domain, r.pearson, r.mape, r.rmse - ); - } - - // CountMinSketchWithHeap - let cmwh_param_sets: Vec = vec![ - CmwhParams { - depth: 3, - width: 1024, - n: 100_000, - domain: 1000, - heap_size: 10, - }, - CmwhParams { - depth: 5, - width: 2048, - n: 200_000, - domain: 2000, - heap_size: 20, - }, - CmwhParams { - depth: 5, - width: 2048, - n: 200_000, - domain: 2000, - heap_size: 50, - }, - ]; - - println!("\n## CountMinSketchWithHeap ({cmwh_mode})"); - println!("| depth | width | n | domain | heap_size | Top-k recall | Pearson (top-k) | MAPE (%) | RMSE (%) |"); - println!("|-------|-------|-----|--------|-----------|--------------|-----------------|----------|----------|"); - for p in &cmwh_param_sets { - let r = run_countmin_with_heap_once(seed, p); - println!( - "| {} | {} | {} | {} | {} | {:.4} | {:.10} | {:.6} | {:.6} |", - p.depth, p.width, p.n, p.domain, p.heap_size, r.topk_recall, r.pearson, r.mape, r.rmse - ); - } - // KllSketch - let kll_param_sets: Vec = vec![ - KllParams { k: 20, n: 200_000 }, - KllParams { k: 50, n: 200_000 }, - KllParams { k: 200, n: 200_000 }, - KllParams { k: 20, n: 50_000 }, - ]; - - println!("\n## KllSketch ({kll_mode})"); - println!( - "| k | n_updates | q=0.5 abs_rank_error | q=0.9 abs_rank_error | q=0.99 abs_rank_error |" - ); - println!( - "|---|-----------|----------------------|----------------------|-----------------------|" - ); - for p in &kll_param_sets { - let r = run_kll_once(seed, p); - println!( - "| {} | {} | {:.6} | {:.6} | {:.6} |", - p.k, p.n, r.rank_err_50, r.rank_err_90, r.rank_err_99 - ); - } - - // HydraKllSketch - let hydra_param_sets: Vec = vec![ - HydraKllParams { - rows: 2, - cols: 64, - k: 20, - n: 200_000, - domain: 200, - eval_keys: 50, - }, - HydraKllParams { - rows: 3, - cols: 128, - k: 20, - n: 200_000, - domain: 200, - eval_keys: 50, - }, - HydraKllParams { - rows: 3, - cols: 128, - k: 50, - n: 200_000, - domain: 200, - eval_keys: 50, - }, - HydraKllParams { - rows: 3, - cols: 128, - k: 20, - n: 100_000, - domain: 100, - eval_keys: 50, - }, - ]; - - println!("\n## HydraKllSketch ({kll_mode})"); - println!("| rows | cols | k | n | domain | q=0.5 mean/max | q=0.9 mean/max |"); - println!("|------|------|---|-----|--------|----------------|----------------|"); - for p in &hydra_param_sets { - let r = run_hydra_kll_once(seed, p); - println!( - "| {} | {} | {} | {} | {} | {:.5} / {:.5} | {:.5} / {:.5} |", - p.rows, p.cols, p.k, p.n, p.domain, r.mean_50, r.max_50, r.mean_90, r.max_90 - ); - } -} diff --git a/asap-common/sketch-core/src/config.rs b/asap-common/sketch-core/src/config.rs deleted file mode 100644 index b7c6abc5..00000000 --- a/asap-common/sketch-core/src/config.rs +++ /dev/null @@ -1,83 +0,0 @@ -use std::sync::OnceLock; - -/// Implementation mode for sketch-core internals. -#[derive(Debug, Clone, Copy, PartialEq, Eq, clap::ValueEnum)] -pub enum ImplMode { - /// Use the original hand-written implementations. - Legacy, - /// Use asap_sketchlib backed implementations. - Sketchlib, -} - -/// Global default when impl mode is not explicitly configured (e.g. env var parsing). -pub const DEFAULT_IMPL_MODE: ImplMode = ImplMode::Sketchlib; - -/// Per-backend defaults. Used when configure() has not been called. -pub const DEFAULT_CMS_IMPL: ImplMode = ImplMode::Sketchlib; -pub const DEFAULT_KLL_IMPL: ImplMode = ImplMode::Sketchlib; -pub const DEFAULT_CMWH_IMPL: ImplMode = ImplMode::Sketchlib; - -static COUNTMIN_MODE: OnceLock = OnceLock::new(); - -/// Returns true if Count-Min operations should use asap_sketchlib internally. -pub fn use_sketchlib_for_count_min() -> bool { - *COUNTMIN_MODE.get_or_init(|| DEFAULT_CMS_IMPL) == ImplMode::Sketchlib -} - -static KLL_MODE: OnceLock = OnceLock::new(); - -/// Returns true if KLL operations should use asap_sketchlib internally. -pub fn use_sketchlib_for_kll() -> bool { - *KLL_MODE.get_or_init(|| DEFAULT_KLL_IMPL) == ImplMode::Sketchlib -} - -static COUNTMIN_WITH_HEAP_MODE: OnceLock = OnceLock::new(); - -/// Returns true if Count-Min-With-Heap operations should use asap_sketchlib internally for the -/// Count-Min portion. -pub fn use_sketchlib_for_count_min_with_heap() -> bool { - *COUNTMIN_WITH_HEAP_MODE.get_or_init(|| DEFAULT_CMWH_IMPL) == ImplMode::Sketchlib -} - -/// Set backend modes for all sketch types. Call once at process startup, -/// before any sketch operation. Returns Err if any OnceLock was already set. -pub fn configure(cms: ImplMode, kll: ImplMode, cmwh: ImplMode) -> Result<(), &'static str> { - let a = COUNTMIN_MODE.set(cms); - let b = KLL_MODE.set(kll); - let c = COUNTMIN_WITH_HEAP_MODE.set(cmwh); - if a.is_err() || b.is_err() || c.is_err() { - Err("configure() called after sketch backends were already initialised") - } else { - Ok(()) - } -} - -pub fn force_legacy_mode_for_tests() { - let _ = COUNTMIN_MODE.set(ImplMode::Legacy); - let _ = KLL_MODE.set(ImplMode::Legacy); - let _ = COUNTMIN_WITH_HEAP_MODE.set(ImplMode::Legacy); -} - -/// Helper used by UDF templates and documentation examples to parse implementation mode -/// from environment variables in a robust way. This is not used in the hot path. -pub fn parse_mode(var: Result) -> ImplMode { - match var { - Ok(v) => match v.to_ascii_lowercase().as_str() { - "legacy" => ImplMode::Legacy, - "sketchlib" => ImplMode::Sketchlib, - other => { - eprintln!( - "sketch-core: unrecognised IMPL value {other:?}, defaulting to {DEFAULT_IMPL_MODE:?}" - ); - DEFAULT_IMPL_MODE - } - }, - Err(std::env::VarError::NotPresent) => DEFAULT_IMPL_MODE, - Err(std::env::VarError::NotUnicode(v)) => { - eprintln!( - "sketch-core: IMPL env var has invalid UTF-8 ({v:?}), defaulting to {DEFAULT_IMPL_MODE:?}" - ); - DEFAULT_IMPL_MODE - } - } -} diff --git a/asap-common/sketch-core/src/count_min.rs b/asap-common/sketch-core/src/count_min.rs deleted file mode 100644 index a41f2d1c..00000000 --- a/asap-common/sketch-core/src/count_min.rs +++ /dev/null @@ -1,548 +0,0 @@ -// Adapted from QueryEngineRust/src/precompute_operators/count_min_sketch_accumulator.rs -// Changes: -// - Renamed CountMinSketchAccumulator -> CountMinSketch -// - _update(&KeyByLabelValues) -> pub update(&str) (caller does key-to-string conversion) -// - query_key(&KeyByLabelValues) -> query_key(&str) -// - serialize_to_bytes (trait) -> serialize_msgpack (inherent method) -// - deserialize_from_bytes_arroyo -> deserialize_msgpack -// - merge_accumulators -> merge -// - Removed: deserialize_from_json, deserialize_from_bytes (legacy QE formats, stay in QE) -// - Removed: merge_multiple (QE trait-object helper, stays in QE) -// - Removed: AggregateCore, SerializableToSink, MergeableAccumulator, MultipleSubpopulationAggregate impls -// - Added: aggregate_count() / aggregate_sum() one-shot helpers for Arroyo call pattern - -use serde::{Deserialize, Serialize}; -use xxhash_rust::xxh32::xxh32; - -use crate::config::use_sketchlib_for_count_min; -use crate::count_min_sketchlib::{ - matrix_from_sketchlib_cms, new_sketchlib_cms, sketchlib_cms_from_matrix, sketchlib_cms_query, - sketchlib_cms_update, SketchlibCms, -}; - -#[derive(Serialize, Deserialize)] -struct WireFormat { - sketch: Vec>, - row_num: usize, - col_num: usize, -} - -/// Backend implementation for Count-Min Sketch. Only one is active at a time. -#[derive(Debug, Clone)] -pub enum CountMinBackend { - /// Original hand-written matrix implementation. - Legacy(Vec>), - /// asap_sketchlib backed implementation. - Sketchlib(SketchlibCms), -} - -/// Count-Min Sketch probabilistic data structure for frequency counting. -/// Sparse delta between two consecutive CountMinSketch snapshots — -/// the input shape for [`CountMinSketch::apply_delta`]. Mirrors the -/// `CountMinDelta` proto in -/// `sketchlib-go/proto/countminsketch/countminsketch.proto` (packed -/// encoding only). -/// -/// Cells apply additively: `matrix[row][col] += d_count`. Per-row -/// L1 and L2 norm deltas are carried for downstream error-accounting -/// but are not consumed by `apply_delta` itself. -#[derive(Debug, Clone, Default)] -pub struct CountMinDelta { - pub rows: u32, - pub cols: u32, - pub cells: Vec<(u32, u32, i64)>, - pub l1: Vec, - pub l2: Vec, -} - -/// Provides approximate frequency counts with error bounds. -/// This is the canonical shared implementation; the msgpack wire format is the -/// contract between Arroyo UDAFs (producers) and QueryEngineRust (consumer). -#[derive(Debug, Clone)] -pub struct CountMinSketch { - pub row_num: usize, - pub col_num: usize, - pub backend: CountMinBackend, -} - -impl CountMinSketch { - pub fn new(row_num: usize, col_num: usize) -> Self { - let backend = if use_sketchlib_for_count_min() { - CountMinBackend::Sketchlib(new_sketchlib_cms(row_num, col_num)) - } else { - CountMinBackend::Legacy(vec![vec![0.0; col_num]; row_num]) - }; - Self { - row_num, - col_num, - backend, - } - } - - /// Returns the sketch matrix (for wire format, serialization, tests). - pub fn sketch(&self) -> Vec> { - match &self.backend { - CountMinBackend::Legacy(m) => m.clone(), - CountMinBackend::Sketchlib(s) => matrix_from_sketchlib_cms(s), - } - } - - /// Mutable access to the matrix. Only `Some` for Legacy backend. - pub fn sketch_mut(&mut self) -> Option<&mut Vec>> { - match &mut self.backend { - CountMinBackend::Legacy(m) => Some(m), - CountMinBackend::Sketchlib(_) => None, - } - } - - /// Construct from a legacy matrix (used by deserialization and query engine). - pub fn from_legacy_matrix(sketch: Vec>, row_num: usize, col_num: usize) -> Self { - let backend = if use_sketchlib_for_count_min() { - CountMinBackend::Sketchlib(sketchlib_cms_from_matrix(row_num, col_num, &sketch)) - } else { - CountMinBackend::Legacy(sketch) - }; - Self { - row_num, - col_num, - backend, - } - } - - pub fn update(&mut self, key: &str, value: f64) { - match &mut self.backend { - CountMinBackend::Legacy(sketch) => { - let key_bytes = key.as_bytes(); - for (i, row) in sketch.iter_mut().enumerate().take(self.row_num) { - let hash_value = xxh32(key_bytes, i as u32); - let col_index = (hash_value as usize) % self.col_num; - row[col_index] += value; - } - } - CountMinBackend::Sketchlib(s) => { - sketchlib_cms_update(s, key, value); - } - } - } - - pub fn query_key(&self, key: &str) -> f64 { - match &self.backend { - CountMinBackend::Legacy(sketch) => { - let key_bytes = key.as_bytes(); - let mut min_value = f64::MAX; - for (i, row) in sketch.iter().enumerate().take(self.row_num) { - let hash_value = xxh32(key_bytes, i as u32); - let col_index = (hash_value as usize) % self.col_num; - min_value = min_value.min(row[col_index]); - } - min_value - } - CountMinBackend::Sketchlib(s) => sketchlib_cms_query(s, key), - } - } - - pub fn merge( - accumulators: Vec, - ) -> Result> { - if accumulators.is_empty() { - return Err("No accumulators to merge".into()); - } - - if accumulators.len() == 1 { - return Ok(accumulators.into_iter().next().unwrap()); - } - - // Check that all accumulators have the same dimensions - let row_num = accumulators[0].row_num; - let col_num = accumulators[0].col_num; - - for acc in &accumulators { - if acc.row_num != row_num || acc.col_num != col_num { - return Err( - "Cannot merge CountMinSketch accumulators with different dimensions".into(), - ); - } - } - - if use_sketchlib_for_count_min() { - let mut sketchlib_inners: Vec = Vec::with_capacity(accumulators.len()); - for acc in accumulators { - let matrix = acc.sketch(); - let inner = sketchlib_cms_from_matrix(acc.row_num, acc.col_num, &matrix); - sketchlib_inners.push(inner); - } - let merged_sketchlib = sketchlib_inners - .into_iter() - .reduce(|mut lhs: SketchlibCms, rhs: SketchlibCms| { - lhs.merge(&rhs); - lhs - }) - .ok_or("No accumulators to merge")?; - - let sketch = matrix_from_sketchlib_cms(&merged_sketchlib); - let row_num = sketch.len(); - let col_num = sketch.first().map(|r| r.len()).unwrap_or(0); - - Ok(Self { - row_num, - col_num, - backend: CountMinBackend::Sketchlib(merged_sketchlib), - }) - } else { - let mut merged = accumulators[0].clone(); - for acc in &accumulators[1..] { - let acc_matrix = acc.sketch(); - if let CountMinBackend::Legacy(merged_matrix) = &mut merged.backend { - for (merged_row, acc_row) in merged_matrix.iter_mut().zip(acc_matrix.iter()) { - for (m_cell, a_cell) in merged_row.iter_mut().zip(acc_row.iter()) { - *m_cell += *a_cell; - } - } - } - } - Ok(merged) - } - } - - /// Merge from references, allocating only the output — no input clones. - pub fn merge_refs( - accumulators: &[&Self], - ) -> Result> { - if accumulators.is_empty() { - return Err("No accumulators to merge".into()); - } - - let row_num = accumulators[0].row_num; - let col_num = accumulators[0].col_num; - - for acc in accumulators { - if acc.row_num != row_num || acc.col_num != col_num { - return Err( - "Cannot merge CountMinSketch accumulators with different dimensions".into(), - ); - } - } - - if use_sketchlib_for_count_min() { - let mut sketchlib_inners: Vec = Vec::with_capacity(accumulators.len()); - for acc in accumulators { - let acc_matrix = acc.sketch(); - let matrix_has_values = acc_matrix - .iter() - .any(|row: &Vec| row.iter().any(|&v| v != 0.0)); - - let inner = if matrix_has_values { - sketchlib_cms_from_matrix(acc.row_num, acc.col_num, &acc_matrix) - } else if let CountMinBackend::Sketchlib(s) = &acc.backend { - s.clone() - } else { - sketchlib_cms_from_matrix(acc.row_num, acc.col_num, &acc_matrix) - }; - - sketchlib_inners.push(inner); - } - - let merged_sketchlib = sketchlib_inners - .into_iter() - .reduce(|mut lhs: SketchlibCms, rhs: SketchlibCms| { - lhs.merge(&rhs); - lhs - }) - .ok_or("No accumulators to merge")?; - - let sketch = matrix_from_sketchlib_cms(&merged_sketchlib); - let r = sketch.len(); - let c = sketch.first().map(|row| row.len()).unwrap_or(0); - - Ok(Self { - row_num: r, - col_num: c, - backend: CountMinBackend::Sketchlib(merged_sketchlib), - }) - } else { - let mut merged = Self::new(row_num, col_num); - if let CountMinBackend::Legacy(ref mut merged_sketch) = merged.backend { - for acc in accumulators { - let acc_matrix = acc.sketch(); - for (merged_row, acc_row) in merged_sketch.iter_mut().zip(acc_matrix.iter()) { - for (m_cell, a_cell) in merged_row.iter_mut().zip(acc_row.iter()) { - *m_cell += *a_cell; - } - } - } - } - Ok(merged) - } - } - - /// Apply a sparse delta in place. Matches the `ApplyDelta` - /// semantics in `sketchlib-go/sketches/CountMinSketch/delta.go`: - /// `matrix[row][col] += d_count` for each cell in the delta. - /// - /// Both backends (Legacy Vec-of-Vec and Sketchlib FFI) are - /// supported: for Legacy we mutate `sketch_mut()` in place; for - /// Sketchlib the FFI handle is opaque, so we snapshot the - /// matrix, apply cell updates, and rebuild the backend. The - /// rebuild is O(rows × cols) per delta and is acceptable for - /// ingest-side reconstitution — no delta should fire more than - /// once per window (10s–300s in the paper's B3 / B4 configs). - pub fn apply_delta( - &mut self, - delta: &CountMinDelta, - ) -> Result<(), Box> { - for (row, col, _) in &delta.cells { - let r = *row as usize; - let c = *col as usize; - if r >= self.row_num || c >= self.col_num { - return Err(format!( - "CountMinDelta cell ({r},{c}) out of range (matrix={}x{})", - self.row_num, self.col_num - ) - .into()); - } - } - match &mut self.backend { - CountMinBackend::Legacy(matrix) => { - for (row, col, d_count) in &delta.cells { - matrix[*row as usize][*col as usize] += *d_count as f64; - } - } - CountMinBackend::Sketchlib(_) => { - let mut matrix = self.sketch(); - for (row, col, d_count) in &delta.cells { - matrix[*row as usize][*col as usize] += *d_count as f64; - } - self.backend = CountMinBackend::Sketchlib(sketchlib_cms_from_matrix( - self.row_num, - self.col_num, - &matrix, - )); - } - } - Ok(()) - } - - /// Serialize to MessagePack — matches the Arroyo UDF wire format exactly. - pub fn serialize_msgpack(&self) -> Vec { - let sketch = self.sketch(); - let wire = WireFormat { - sketch, - row_num: self.row_num, - col_num: self.col_num, - }; - - let mut buf = Vec::new(); - wire.serialize(&mut rmp_serde::Serializer::new(&mut buf)) - .unwrap(); - buf - } - - /// Deserialize from MessagePack produced by the Arroyo UDF. - pub fn deserialize_msgpack(buffer: &[u8]) -> Result> { - let wire: WireFormat = - rmp_serde::from_slice(buffer).map_err(|e| -> Box { - format!("Failed to deserialize CountMinSketch from MessagePack: {e}").into() - })?; - - let backend = if use_sketchlib_for_count_min() { - CountMinBackend::Sketchlib(sketchlib_cms_from_matrix( - wire.row_num, - wire.col_num, - &wire.sketch, - )) - } else { - CountMinBackend::Legacy(wire.sketch) - }; - - Ok(Self { - row_num: wire.row_num, - col_num: wire.col_num, - backend, - }) - } - - /// One-shot aggregation for the Arroyo UDAF call pattern: build a sketch from - /// parallel key/value slices and return the msgpack bytes. - pub fn aggregate_count( - depth: usize, - width: usize, - keys: &[&str], - values: &[f64], - ) -> Option> { - if keys.is_empty() { - return None; - } - let mut sketch = Self::new(depth, width); - for (key, &value) in keys.iter().zip(values.iter()) { - sketch.update(key, value); - } - Some(sketch.serialize_msgpack()) - } - - /// Same as aggregate_count — CMS accumulates sums by construction. - pub fn aggregate_sum( - depth: usize, - width: usize, - keys: &[&str], - values: &[f64], - ) -> Option> { - Self::aggregate_count(depth, width, keys, values) - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_count_min_sketch_creation() { - let cms = CountMinSketch::new(4, 1000); - assert_eq!(cms.row_num, 4); - assert_eq!(cms.col_num, 1000); - let sketch = cms.sketch(); - assert_eq!(sketch.len(), 4); - assert_eq!(sketch[0].len(), 1000); - - // Check all values are initialized to 0 - for row in &sketch { - for &value in row { - assert_eq!(value, 0.0); - } - } - } - - #[test] - fn test_count_min_sketch_update() { - let mut cms = CountMinSketch::new(2, 10); - cms.update("key1", 1.0); - // Query should return at least the updated value - let result = cms.query_key("key1"); - assert!(result >= 1.0); - } - - #[test] - fn test_count_min_sketch_query_empty() { - let cms = CountMinSketch::new(2, 10); - assert_eq!(cms.query_key("anything"), 0.0); - } - - #[test] - fn test_count_min_sketch_merge() { - // Use from_legacy_matrix so the test works regardless of sketchlib/legacy config - let mut sketch1 = vec![vec![0.0; 3]; 2]; - sketch1[0][0] = 5.0; - sketch1[1][2] = 10.0; - let cms1 = CountMinSketch::from_legacy_matrix(sketch1, 2, 3); - - let mut sketch2 = vec![vec![0.0; 3]; 2]; - sketch2[0][0] = 3.0; - sketch2[0][1] = 7.0; - let cms2 = CountMinSketch::from_legacy_matrix(sketch2, 2, 3); - - let merged = CountMinSketch::merge(vec![cms1, cms2]).unwrap(); - let merged_sketch = merged.sketch(); - - assert_eq!(merged_sketch[0][0], 8.0); // 5 + 3 - assert_eq!(merged_sketch[0][1], 7.0); // 0 + 7 - assert_eq!(merged_sketch[1][2], 10.0); // 10 + 0 - } - - #[test] - fn test_count_min_sketch_merge_dimension_mismatch() { - let cms1 = CountMinSketch::new(2, 3); - let cms2 = CountMinSketch::new(3, 3); - assert!(CountMinSketch::merge(vec![cms1, cms2]).is_err()); - } - - #[test] - fn test_count_min_sketch_msgpack_round_trip() { - let mut cms = CountMinSketch::new(4, 256); - cms.update("apple", 5.0); - cms.update("banana", 3.0); - cms.update("apple", 2.0); // total "apple" = 7 - - let bytes = cms.serialize_msgpack(); - let deserialized = CountMinSketch::deserialize_msgpack(&bytes).unwrap(); - - assert_eq!(deserialized.row_num, 4); - assert_eq!(deserialized.col_num, 256); - assert!(deserialized.query_key("apple") >= 7.0); - assert!(deserialized.query_key("banana") >= 3.0); - } - - #[test] - fn test_aggregate_count() { - let keys = ["a", "b", "a"]; - let values = [1.0, 2.0, 3.0]; - let bytes = CountMinSketch::aggregate_count(4, 100, &keys, &values).unwrap(); - let cms = CountMinSketch::deserialize_msgpack(&bytes).unwrap(); - // "a" was updated twice (1.0 + 3.0 = 4.0), "b" once (2.0) - assert!(cms.query_key("a") >= 4.0); - assert!(cms.query_key("b") >= 2.0); - } - - #[test] - fn test_aggregate_count_empty() { - assert!(CountMinSketch::aggregate_count(4, 100, &[], &[]).is_none()); - } - - #[test] - fn test_apply_delta_additive() { - let mut cms = CountMinSketch::from_legacy_matrix( - vec![vec![1.0, 2.0, 3.0], vec![4.0, 5.0, 6.0]], - 2, - 3, - ); - let delta = CountMinDelta { - rows: 2, - cols: 3, - cells: vec![(0, 0, 10), (1, 2, 100)], - l1: vec![], - l2: vec![], - }; - cms.apply_delta(&delta).unwrap(); - assert_eq!( - cms.sketch(), - vec![vec![11.0, 2.0, 3.0], vec![4.0, 5.0, 106.0]] - ); - } - - #[test] - fn test_apply_delta_matches_full_merge() { - let base = CountMinSketch::from_legacy_matrix( - vec![vec![1.0, 2.0], vec![3.0, 4.0]], - 2, - 2, - ); - let addition = CountMinSketch::from_legacy_matrix( - vec![vec![10.0, 0.0], vec![0.0, 20.0]], - 2, - 2, - ); - let via_merge = CountMinSketch::merge(vec![base.clone(), addition]).unwrap(); - - let delta = CountMinDelta { - rows: 2, - cols: 2, - cells: vec![(0, 0, 10), (1, 1, 20)], - l1: vec![], - l2: vec![], - }; - let mut via_delta = base; - via_delta.apply_delta(&delta).unwrap(); - assert_eq!(via_delta.sketch(), via_merge.sketch()); - } - - #[test] - fn test_apply_delta_out_of_range() { - let mut cms = CountMinSketch::new(2, 3); - let delta = CountMinDelta { - rows: 2, - cols: 3, - cells: vec![(5, 0, 1)], - l1: vec![], - l2: vec![], - }; - assert!(cms.apply_delta(&delta).is_err()); - } -} diff --git a/asap-common/sketch-core/src/count_min_sketchlib.rs b/asap-common/sketch-core/src/count_min_sketchlib.rs deleted file mode 100644 index 6fb3924d..00000000 --- a/asap-common/sketch-core/src/count_min_sketchlib.rs +++ /dev/null @@ -1,59 +0,0 @@ -use asap_sketchlib::{CountMin, RegularPath, SketchInput, Vector2D}; - -/// Concrete Count-Min type from asap_sketchlib when sketchlib backend is enabled. -/// Uses f64 counters (Vector2D) for weighted updates without integer rounding. -pub type SketchlibCms = CountMin, RegularPath>; - -/// Creates a fresh sketchlib Count-Min sketch with the given dimensions. -pub fn new_sketchlib_cms(row_num: usize, col_num: usize) -> SketchlibCms { - SketchlibCms::with_dimensions(row_num, col_num) -} - -/// Builds a sketchlib Count-Min sketch from an existing `sketch` matrix. -pub fn sketchlib_cms_from_matrix( - row_num: usize, - col_num: usize, - sketch: &[Vec], -) -> SketchlibCms { - let matrix = Vector2D::from_fn(row_num, col_num, |r, c| { - sketch - .get(r) - .and_then(|row| row.get(c)) - .copied() - .unwrap_or(0.0) - }); - SketchlibCms::from_storage(matrix) -} - -/// Converts a sketchlib Count-Min sketch into the legacy `Vec>` matrix. -pub fn matrix_from_sketchlib_cms(inner: &SketchlibCms) -> Vec> { - let storage: &Vector2D = inner.as_storage(); - let rows = storage.rows(); - let cols = storage.cols(); - let mut sketch = vec![vec![0.0; cols]; rows]; - - for (r, row) in sketch.iter_mut().enumerate().take(rows) { - for (c, cell) in row.iter_mut().enumerate().take(cols) { - if let Some(v) = storage.get(r, c) { - *cell = *v; - } - } - } - - sketch -} - -/// Helper to update a sketchlib Count-Min with a weighted key. -pub fn sketchlib_cms_update(inner: &mut SketchlibCms, key: &str, value: f64) { - if value <= 0.0 { - return; - } - let input = SketchInput::String(key.to_owned()); - inner.insert_many(&input, value); -} - -/// Helper to query a sketchlib Count-Min for a key, returning f64. -pub fn sketchlib_cms_query(inner: &SketchlibCms, key: &str) -> f64 { - let input = SketchInput::String(key.to_owned()); - inner.estimate(&input) -} diff --git a/asap-common/sketch-core/src/count_min_with_heap.rs b/asap-common/sketch-core/src/count_min_with_heap.rs deleted file mode 100644 index f02efc70..00000000 --- a/asap-common/sketch-core/src/count_min_with_heap.rs +++ /dev/null @@ -1,597 +0,0 @@ -// Adapted from QueryEngineRust/src/precompute_operators/count_min_sketch_with_heap_accumulator.rs -// Changes: -// - Renamed CountMinSketchWithHeapAccumulator -> CountMinSketchWithHeap -// - Inner CmsData helper renamed to avoid name collision with count_min::CountMinSketch -// - update() takes &str instead of &KeyByLabelValues -// - query_key() takes &str -// - serialize_to_bytes (trait) -> serialize_msgpack (inherent method) -// - deserialize_from_bytes_arroyo -> deserialize_msgpack -// - merge_accumulators -> merge -// - Removed: deserialize_from_json, deserialize_from_bytes (legacy QE formats, stay in QE) -// - Removed: AggregateCore, SerializableToSink, MergeableAccumulator, MultipleSubpopulationAggregate impls -// - Removed: get_topk_keys (returns KeyByLabelValues — QE-specific) -// - Added: insert_or_update_heap helper, aggregate_topk() one-shot helper -// - Refactored to enum-based backend (Legacy vs Sketchlib) -// -// NOTE (bug, do not fix): QueryEngineRust uses xxhash-rust::xxh32; the Arroyo template uses -// twox-hash::XxHash32. Bucket assignments differ, so query results will be wrong until the -// hash crate mismatch is resolved. Tracked separately. - -use serde::{Deserialize, Serialize}; -use std::collections::HashSet; -use xxhash_rust::xxh32::xxh32; - -use crate::config::use_sketchlib_for_count_min_with_heap; -use crate::count_min_with_heap_sketchlib::{ - heap_to_wire, matrix_from_sketchlib_cms_heap, new_sketchlib_cms_heap, - sketchlib_cms_heap_from_matrix_and_heap, sketchlib_cms_heap_query, sketchlib_cms_heap_update, - SketchlibCMSHeap, WireHeapItem, -}; - -/// Item in the top-k heap representing a key-value pair. -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct HeapItem { - pub key: String, - pub value: f64, -} - -/// Helper struct matching Arroyo's nested serialization format (inner CMS). -#[derive(Debug, Clone, Serialize, Deserialize)] -struct CmsData { - sketch: Vec>, - row_num: usize, - col_num: usize, -} - -/// Helper struct matching Arroyo's serialization format (outer wrapper). -#[derive(Debug, Clone, Serialize, Deserialize)] -struct CountMinSketchWithHeapSerialized { - sketch: CmsData, - topk_heap: Vec, - heap_size: usize, -} - -/// Backend implementation for Count-Min Sketch with Heap. Only one is active at a time. -pub enum CountMinWithHeapBackend { - /// Legacy implementation: matrix + local heap. - Legacy { - sketch: Vec>, - heap: Vec, - }, - /// asap_sketchlib CMSHeap implementation. - Sketchlib(SketchlibCMSHeap), -} - -impl std::fmt::Debug for CountMinWithHeapBackend { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - CountMinWithHeapBackend::Legacy { sketch, heap } => f - .debug_struct("Legacy") - .field("sketch", sketch) - .field("heap", heap) - .finish(), - CountMinWithHeapBackend::Sketchlib(_) => write!(f, "Sketchlib(..)"), - } - } -} - -/// Count-Min Sketch with Heap for top-k tracking. -/// Combines probabilistic frequency counting with efficient top-k maintenance. -pub struct CountMinSketchWithHeap { - pub row_num: usize, - pub col_num: usize, - pub heap_size: usize, - pub backend: CountMinWithHeapBackend, -} - -impl std::fmt::Debug for CountMinSketchWithHeap { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("CountMinSketchWithHeap") - .field("row_num", &self.row_num) - .field("col_num", &self.col_num) - .field("heap_size", &self.heap_size) - .field("backend", &self.backend) - .finish() - } -} - -impl Clone for CountMinSketchWithHeap { - fn clone(&self) -> Self { - let backend = match &self.backend { - CountMinWithHeapBackend::Legacy { sketch, heap } => CountMinWithHeapBackend::Legacy { - sketch: sketch.clone(), - heap: heap.clone(), - }, - CountMinWithHeapBackend::Sketchlib(cms_heap) => { - let sketch = matrix_from_sketchlib_cms_heap(cms_heap); - let heap_items: Vec = heap_to_wire(cms_heap) - .into_iter() - .map(|w| HeapItem { - key: w.key, - value: w.value, - }) - .collect(); - let wire_ref: Vec = heap_items - .iter() - .map(|h| WireHeapItem { - key: h.key.clone(), - value: h.value, - }) - .collect(); - CountMinWithHeapBackend::Sketchlib(sketchlib_cms_heap_from_matrix_and_heap( - self.row_num, - self.col_num, - self.heap_size, - &sketch, - &wire_ref, - )) - } - }; - Self { - row_num: self.row_num, - col_num: self.col_num, - heap_size: self.heap_size, - backend, - } - } -} - -impl CountMinSketchWithHeap { - pub fn new(row_num: usize, col_num: usize, heap_size: usize) -> Self { - let backend = if use_sketchlib_for_count_min_with_heap() { - CountMinWithHeapBackend::Sketchlib(new_sketchlib_cms_heap(row_num, col_num, heap_size)) - } else { - CountMinWithHeapBackend::Legacy { - sketch: vec![vec![0.0; col_num]; row_num], - heap: Vec::new(), - } - }; - Self { - row_num, - col_num, - heap_size, - backend, - } - } - - /// Create from legacy matrix and heap (e.g. from JSON deserialization). - pub fn from_legacy_matrix( - sketch: Vec>, - topk_heap: Vec, - row_num: usize, - col_num: usize, - heap_size: usize, - ) -> Self { - Self { - row_num, - col_num, - heap_size, - backend: CountMinWithHeapBackend::Legacy { - sketch, - heap: topk_heap, - }, - } - } - - /// Mutable reference to the sketch matrix. Only valid for Legacy backend. - pub fn sketch_mut(&mut self) -> Option<&mut Vec>> { - match &mut self.backend { - CountMinWithHeapBackend::Legacy { sketch, .. } => Some(sketch), - CountMinWithHeapBackend::Sketchlib(_) => None, - } - } - - /// Get the top-k heap items (works for both backends). - pub fn topk_heap_items(&self) -> Vec { - match &self.backend { - CountMinWithHeapBackend::Legacy { heap, .. } => heap.clone(), - CountMinWithHeapBackend::Sketchlib(cms_heap) => heap_to_wire(cms_heap) - .into_iter() - .map(|w| HeapItem { - key: w.key, - value: w.value, - }) - .collect(), - } - } - - /// Get the sketch matrix (works for both backends). - pub fn sketch_matrix(&self) -> Vec> { - match &self.backend { - CountMinWithHeapBackend::Legacy { sketch, .. } => sketch.clone(), - CountMinWithHeapBackend::Sketchlib(cms_heap) => { - matrix_from_sketchlib_cms_heap(cms_heap) - } - } - } - - pub fn update(&mut self, key: &str, value: f64) { - match &mut self.backend { - CountMinWithHeapBackend::Legacy { sketch, heap } => { - let key_bytes = key.as_bytes(); - for (i, row) in sketch.iter_mut().enumerate().take(self.row_num) { - let hash_value = xxh32(key_bytes, i as u32); - let col_index = (hash_value as usize) % self.col_num; - row[col_index] += value; - } - Self::insert_or_update_heap_inline(heap, key, value, self.heap_size); - } - CountMinWithHeapBackend::Sketchlib(cms_heap) => { - sketchlib_cms_heap_update(cms_heap, key, value); - } - } - } - - fn insert_or_update_heap_inline( - heap: &mut Vec, - key: &str, - value: f64, - heap_size: usize, - ) { - if let Some(item) = heap.iter_mut().find(|i| i.key == key) { - item.value += value; - } else if heap.len() < heap_size { - heap.push(HeapItem { - key: key.to_string(), - value, - }); - } else if let Some(min_item) = heap.iter_mut().min_by(|a, b| { - a.value - .partial_cmp(&b.value) - .unwrap_or(std::cmp::Ordering::Equal) - }) { - if value > min_item.value { - *min_item = HeapItem { - key: key.to_string(), - value, - }; - } - } - } - - pub fn query_key(&self, key: &str) -> f64 { - match &self.backend { - CountMinWithHeapBackend::Legacy { sketch, .. } => { - let key_bytes = key.as_bytes(); - let mut min_value = f64::MAX; - for (i, row) in sketch.iter().enumerate().take(self.row_num) { - let hash_value = xxh32(key_bytes, i as u32); - let col_index = (hash_value as usize) % self.col_num; - min_value = min_value.min(row[col_index]); - } - min_value - } - CountMinWithHeapBackend::Sketchlib(cms_heap) => sketchlib_cms_heap_query(cms_heap, key), - } - } - - pub fn merge( - accumulators: Vec, - ) -> Result> { - if accumulators.is_empty() { - return Err("No accumulators to merge".into()); - } - - if accumulators.len() == 1 { - return Ok(accumulators.into_iter().next().unwrap()); - } - - let row_num = accumulators[0].row_num; - let col_num = accumulators[0].col_num; - - for acc in &accumulators { - if acc.row_num != row_num || acc.col_num != col_num { - return Err( - "Cannot merge CountMinSketchWithHeap accumulators with different dimensions" - .into(), - ); - } - } - - let min_heap_size = accumulators - .iter() - .map(|acc| acc.heap_size) - .min() - .unwrap_or(0); - - let mut all_keys: HashSet = HashSet::new(); - for acc in &accumulators { - for item in acc.topk_heap_items() { - all_keys.insert(item.key); - } - } - - match &accumulators[0].backend { - CountMinWithHeapBackend::Sketchlib(_) => { - let mut sketchlib_cms_heaps: Vec = - Vec::with_capacity(accumulators.len()); - for acc in accumulators { - let (sketch, heap) = match &acc.backend { - CountMinWithHeapBackend::Legacy { sketch, heap } => { - (sketch.clone(), heap.clone()) - } - CountMinWithHeapBackend::Sketchlib(cms_heap) => ( - matrix_from_sketchlib_cms_heap(cms_heap), - heap_to_wire(cms_heap) - .into_iter() - .map(|w| HeapItem { - key: w.key, - value: w.value, - }) - .collect(), - ), - }; - let wire_heap: Vec = heap - .iter() - .map(|h| WireHeapItem { - key: h.key.clone(), - value: h.value, - }) - .collect(); - sketchlib_cms_heaps.push(sketchlib_cms_heap_from_matrix_and_heap( - acc.row_num, - acc.col_num, - acc.heap_size, - &sketch, - &wire_heap, - )); - } - - let merged_sketchlib = sketchlib_cms_heaps - .into_iter() - .reduce(|mut lhs, rhs| { - lhs.merge(&rhs); - lhs - }) - .ok_or("No accumulators to merge")?; - - let _merged_sketch = matrix_from_sketchlib_cms_heap(&merged_sketchlib); - let _heap_items: Vec = heap_to_wire(&merged_sketchlib) - .into_iter() - .map(|w| HeapItem { - key: w.key, - value: w.value, - }) - .collect(); - - Ok(CountMinSketchWithHeap { - row_num, - col_num, - heap_size: min_heap_size, - backend: CountMinWithHeapBackend::Sketchlib(merged_sketchlib), - }) - } - CountMinWithHeapBackend::Legacy { .. } => { - let mut merged_sketch = vec![vec![0.0; col_num]; row_num]; - for acc in &accumulators { - let sketch = match &acc.backend { - CountMinWithHeapBackend::Legacy { sketch, .. } => sketch, - CountMinWithHeapBackend::Sketchlib(_) => { - return Err( - "Cannot mix Legacy and Sketchlib backends when merging".into() - ); - } - }; - for (i, row) in merged_sketch.iter_mut().enumerate() { - for (j, cell) in row.iter_mut().enumerate() { - *cell += sketch[i][j]; - } - } - } - - let temp_merged = Self::from_legacy_matrix( - merged_sketch.clone(), - Vec::new(), - row_num, - col_num, - min_heap_size, - ); - - let mut heap_items: Vec = all_keys - .into_iter() - .map(|key_str| { - let frequency = temp_merged.query_key(&key_str); - HeapItem { - key: key_str, - value: frequency, - } - }) - .collect(); - - heap_items.sort_by(|a, b| b.value.partial_cmp(&a.value).unwrap()); - heap_items.truncate(min_heap_size); - - Ok(CountMinSketchWithHeap { - row_num, - col_num, - heap_size: min_heap_size, - backend: CountMinWithHeapBackend::Legacy { - sketch: merged_sketch, - heap: heap_items, - }, - }) - } - } - } - - pub fn serialize_msgpack(&self) -> Vec { - let (sketch, topk_heap) = (self.sketch_matrix(), self.topk_heap_items()); - - let serialized = CountMinSketchWithHeapSerialized { - sketch: CmsData { - sketch, - row_num: self.row_num, - col_num: self.col_num, - }, - topk_heap, - heap_size: self.heap_size, - }; - - let mut buf = Vec::new(); - serialized - .serialize(&mut rmp_serde::Serializer::new(&mut buf)) - .unwrap(); - buf - } - - pub fn deserialize_msgpack(buffer: &[u8]) -> Result> { - let serialized: CountMinSketchWithHeapSerialized = - rmp_serde::from_slice(buffer).map_err(|e| { - format!("Failed to deserialize CountMinSketchWithHeap from MessagePack: {e}") - })?; - - let mut sorted_topk_heap = serialized.topk_heap; - sorted_topk_heap.sort_by(|a, b| b.value.partial_cmp(&a.value).unwrap()); - - let backend = if use_sketchlib_for_count_min_with_heap() { - let wire_heap: Vec = sorted_topk_heap - .iter() - .map(|h| WireHeapItem { - key: h.key.clone(), - value: h.value, - }) - .collect(); - CountMinWithHeapBackend::Sketchlib(sketchlib_cms_heap_from_matrix_and_heap( - serialized.sketch.row_num, - serialized.sketch.col_num, - serialized.heap_size, - &serialized.sketch.sketch, - &wire_heap, - )) - } else { - CountMinWithHeapBackend::Legacy { - sketch: serialized.sketch.sketch, - heap: sorted_topk_heap, - } - }; - - Ok(Self { - row_num: serialized.sketch.row_num, - col_num: serialized.sketch.col_num, - heap_size: serialized.heap_size, - backend, - }) - } - - pub fn aggregate_topk( - row_num: usize, - col_num: usize, - heap_size: usize, - keys: &[&str], - values: &[f64], - ) -> Option> { - if keys.is_empty() { - return None; - } - let mut sketch = Self::new(row_num, col_num, heap_size); - for (key, &value) in keys.iter().zip(values.iter()) { - sketch.update(key, value); - } - Some(sketch.serialize_msgpack()) - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_creation() { - let cms = CountMinSketchWithHeap::new(4, 1000, 20); - assert_eq!(cms.row_num, 4); - assert_eq!(cms.col_num, 1000); - assert_eq!(cms.heap_size, 20); - assert_eq!(cms.sketch_matrix().len(), 4); - assert_eq!(cms.sketch_matrix()[0].len(), 1000); - assert_eq!(cms.topk_heap_items().len(), 0); - } - - #[test] - fn test_query_empty() { - let cms = CountMinSketchWithHeap::new(2, 10, 5); - assert_eq!(cms.query_key("anything"), 0.0); - } - - #[test] - fn test_merge() { - let mut cms1 = CountMinSketchWithHeap::new(2, 10, 5); - let mut cms2 = CountMinSketchWithHeap::new(2, 10, 3); - - if let Some(sketch) = cms1.sketch_mut() { - sketch[0][0] = 10.0; - sketch[1][1] = 20.0; - } - if let Some(sketch) = cms2.sketch_mut() { - sketch[0][0] = 5.0; - sketch[1][1] = 15.0; - } - if let CountMinWithHeapBackend::Legacy { heap, .. } = &mut cms1.backend { - heap.push(HeapItem { - key: "key1".to_string(), - value: 100.0, - }); - heap.push(HeapItem { - key: "key2".to_string(), - value: 50.0, - }); - } - if let CountMinWithHeapBackend::Legacy { heap, .. } = &mut cms2.backend { - heap.push(HeapItem { - key: "key3".to_string(), - value: 75.0, - }); - heap.push(HeapItem { - key: "key1".to_string(), - value: 80.0, - }); - } - - let merged = CountMinSketchWithHeap::merge(vec![cms1, cms2]).unwrap(); - - assert_eq!(merged.sketch_matrix()[0][0], 15.0); - assert_eq!(merged.sketch_matrix()[1][1], 35.0); - assert_eq!(merged.heap_size, 3); - assert!(merged.topk_heap_items().len() <= 3); - } - - #[test] - fn test_merge_dimension_mismatch() { - let cms1 = CountMinSketchWithHeap::new(2, 10, 5); - let cms2 = CountMinSketchWithHeap::new(3, 10, 5); - assert!(CountMinSketchWithHeap::merge(vec![cms1, cms2]).is_err()); - } - - #[test] - fn test_msgpack_round_trip() { - let mut cms = CountMinSketchWithHeap::new(4, 128, 3); - cms.update("hot", 100.0); - cms.update("cold", 1.0); - - let bytes = cms.serialize_msgpack(); - let deserialized = CountMinSketchWithHeap::deserialize_msgpack(&bytes).unwrap(); - - assert_eq!(deserialized.row_num, 4); - assert_eq!(deserialized.col_num, 128); - assert_eq!(deserialized.heap_size, 3); - assert!(!deserialized.topk_heap_items().is_empty()); - assert_eq!(deserialized.topk_heap_items()[0].key, "hot"); - assert!(deserialized.topk_heap_items()[0].value >= 100.0); - assert!(deserialized.query_key("hot") >= 100.0); - assert!(deserialized.query_key("cold") >= 1.0); - } - - #[test] - fn test_aggregate_topk() { - let keys = ["a", "b", "a", "c"]; - let values = [1.0, 2.0, 3.0, 0.5]; - let bytes = CountMinSketchWithHeap::aggregate_topk(4, 100, 2, &keys, &values).unwrap(); - let cms = CountMinSketchWithHeap::deserialize_msgpack(&bytes).unwrap(); - assert_eq!(cms.heap_size, 2); - assert!(cms.topk_heap_items().len() <= 2); - } - - #[test] - fn test_aggregate_topk_empty() { - assert!(CountMinSketchWithHeap::aggregate_topk(4, 100, 10, &[], &[]).is_none()); - } -} diff --git a/asap-common/sketch-core/src/count_min_with_heap_sketchlib.rs b/asap-common/sketch-core/src/count_min_with_heap_sketchlib.rs deleted file mode 100644 index 451c5b87..00000000 --- a/asap-common/sketch-core/src/count_min_with_heap_sketchlib.rs +++ /dev/null @@ -1,109 +0,0 @@ -//! asap_sketchlib CMSHeap integration for CountMinSketchWithHeap. -//! -//! Uses CMSHeap (CountMin + HHHeap) from asap_sketchlib instead of CountMin + local heap, -//! providing automatic top-k tracking during insert and merge. - -use asap_sketchlib::RegularPath; -use asap_sketchlib::{CMSHeap, SketchInput, Vector2D}; - -/// Wire-format heap item (key, value) to avoid circular dependency with count_min_with_heap. -pub struct WireHeapItem { - pub key: String, - pub value: f64, -} - -/// Concrete Count-Min-with-Heap type from asap_sketchlib (CMS + HHHeap). -pub type SketchlibCMSHeap = CMSHeap, RegularPath>; - -/// Creates a fresh CMSHeap with the given dimensions and heap capacity. -pub fn new_sketchlib_cms_heap( - row_num: usize, - col_num: usize, - heap_size: usize, -) -> SketchlibCMSHeap { - CMSHeap::new(row_num, col_num, heap_size) -} - -/// Builds a CMSHeap from an existing sketch matrix and optional heap items. -/// Used when deserializing or when ensuring sketchlib from legacy state. -pub fn sketchlib_cms_heap_from_matrix_and_heap( - row_num: usize, - col_num: usize, - heap_size: usize, - sketch: &[Vec], - topk_heap: &[WireHeapItem], -) -> SketchlibCMSHeap { - let matrix = Vector2D::from_fn(row_num, col_num, |r, c| { - sketch - .get(r) - .and_then(|row| row.get(c)) - .copied() - .unwrap_or(0.0) - .round() as i64 - }); - let mut cms_heap = CMSHeap::from_storage(matrix, heap_size); - - // Populate the heap from wire-format topk_heap - for item in topk_heap { - let count = item.value.round() as i64; - if count > 0 { - let input = SketchInput::Str(&item.key); - cms_heap.heap_mut().update(&input, count); - } - } - - cms_heap -} - -/// Converts a CMSHeap's storage into the legacy `Vec>` matrix. -pub fn matrix_from_sketchlib_cms_heap(cms_heap: &SketchlibCMSHeap) -> Vec> { - let storage = cms_heap.cms().as_storage(); - let rows = storage.rows(); - let cols = storage.cols(); - let mut sketch = vec![vec![0.0; cols]; rows]; - - for (r, row) in sketch.iter_mut().enumerate().take(rows) { - for (c, cell) in row.iter_mut().enumerate().take(cols) { - if let Some(v) = storage.get(r, c) { - *cell = *v as f64; - } - } - } - - sketch -} - -/// Converts sketchlib HHHeap items to wire-format (key, value) pairs. -pub fn heap_to_wire(cms_heap: &SketchlibCMSHeap) -> Vec { - cms_heap - .heap() - .heap() - .iter() - .map(|hh_item| { - let key = match &hh_item.key { - asap_sketchlib::HeapItem::String(s) => s.clone(), - other => format!("{:?}", other), - }; - WireHeapItem { - key, - value: hh_item.count as f64, - } - }) - .collect() -} - -/// Updates a CMSHeap with a weighted key. Automatically updates the heap. -pub fn sketchlib_cms_heap_update(cms_heap: &mut SketchlibCMSHeap, key: &str, value: f64) { - let many = value.round() as i64; - if many <= 0 { - return; - } - let input = SketchInput::String(key.to_owned()); - cms_heap.insert_many(&input, many); -} - -/// Queries a CMSHeap for a key's frequency estimate. -pub fn sketchlib_cms_heap_query(cms_heap: &SketchlibCMSHeap, key: &str) -> f64 { - let input = SketchInput::String(key.to_owned()); - cms_heap.estimate(&input) as f64 -} diff --git a/asap-common/sketch-core/src/count_sketch.rs b/asap-common/sketch-core/src/count_sketch.rs deleted file mode 100644 index 01afc460..00000000 --- a/asap-common/sketch-core/src/count_sketch.rs +++ /dev/null @@ -1,274 +0,0 @@ -//! Count Sketch (a.k.a. Count-Min-style signed-counter sketch) — -//! element-wise mergeable frequency estimator. -//! -//! Parallel to `count_min::CountMinSketch` but with **signed** counters, -//! matching the `asap_sketchlib::proto::sketchlib::CountSketchState` wire -//! format that DataCollector's `countsketchprocessor` emits via the -//! modified OTLP `Metric.data = CountSketch{…}` variant. -//! -//! This is the minimal surface needed for PR C-CountSketch in the -//! modified-OTLP hot path: construct from a decoded proto state, merge -//! element-wise with another sketch, emit the matrix for queries and -//! serialization. The richer query semantics of Count Sketch (median- -//! of-estimators heavy-hitter tracking, `TopKState` integration, etc.) -//! are intentionally deferred to a follow-up — the wire format already -//! carries the matrix losslessly, so the merge/store round-trip works -//! with just a matrix today. - -use serde::{Deserialize, Serialize}; - -/// Sparse delta between two consecutive CountSketch snapshots — -/// the input shape for [`CountSketch::apply_delta`]. Mirrors the -/// `CountSketchDelta` proto in -/// `sketchlib-go/proto/countsketch/countsketch.proto` (packed -/// encoding only — the deprecated `cells_legacy` path + the -/// non-delta `topk` / `hh_keys` top-K carrier aren't modeled here). -/// -/// Cells apply additively: `matrix[row][col] += d_count` for each -/// `(row, col, d_count)` triple. Top-K on the delta path is a -/// separate follow-up (CS top-K is non-linear; merging deltas -/// would require re-querying the merged matrix). -#[derive(Debug, Clone, Default)] -pub struct CountSketchDelta { - pub rows: u32, - pub cols: u32, - /// `(row, col, d_count)` cell updates, additive on the CS matrix. - pub cells: Vec<(u32, u32, i64)>, - /// Per-row L2 norm deltas. Additive, one scalar per row of the - /// base sketch. Kept on the delta surface for downstream - /// error-accounting; `apply_delta` itself ignores L2. - pub l2: Vec, -} - -/// Minimal Count Sketch state — a flat `rows × cols` matrix of signed -/// counts. Element-wise mergeable (sum over aligned cells). -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct CountSketch { - pub row_num: usize, - pub col_num: usize, - /// Row-major matrix of signed counts. `matrix[r][c]` is the value of - /// hash row `r`, column `c`. - pub matrix: Vec>, -} - -impl CountSketch { - /// Construct an all-zero sketch with the given dimensions. - pub fn new(row_num: usize, col_num: usize) -> Self { - Self { - row_num, - col_num, - matrix: vec![vec![0.0; col_num]; row_num], - } - } - - /// Construct from a pre-built matrix (used by the modified-OTLP - /// proto-decode path). - pub fn from_legacy_matrix(matrix: Vec>, row_num: usize, col_num: usize) -> Self { - debug_assert_eq!(matrix.len(), row_num, "row count mismatch"); - debug_assert!( - matrix.iter().all(|r| r.len() == col_num), - "column count mismatch in at least one row" - ); - Self { - row_num, - col_num, - matrix, - } - } - - /// Borrow the inner matrix. - pub fn sketch(&self) -> &Vec> { - &self.matrix - } - - /// Merge one other sketch into self via element-wise addition. Both - /// operands must have identical dimensions. - pub fn merge( - &mut self, - other: &CountSketch, - ) -> Result<(), Box> { - if self.row_num != other.row_num || self.col_num != other.col_num { - return Err(format!( - "CountSketch dimension mismatch: self={}x{}, other={}x{}", - self.row_num, self.col_num, other.row_num, other.col_num - ) - .into()); - } - for r in 0..self.row_num { - for c in 0..self.col_num { - self.matrix[r][c] += other.matrix[r][c]; - } - } - Ok(()) - } - - /// Apply a sparse delta in place. Matches the `ApplyDelta` - /// semantics in `sketchlib-go/sketches/CountSketch/delta.go`: - /// `matrix[row][col] += d_count` for each cell in the delta. - /// Returns `Err` if any `(row, col)` is out of range — indicating - /// a dimension mismatch between the snapshot this sketch was - /// built from and the delta sender. - pub fn apply_delta( - &mut self, - delta: &CountSketchDelta, - ) -> Result<(), Box> { - for (row, col, d_count) in &delta.cells { - let r = *row as usize; - let c = *col as usize; - if r >= self.row_num || c >= self.col_num { - return Err(format!( - "CountSketchDelta cell ({r},{c}) out of range (matrix={}x{})", - self.row_num, self.col_num - ) - .into()); - } - // `d_count` is signed on the wire; CS counts are signed - // too (can go negative under adversarial keys). - self.matrix[r][c] += *d_count as f64; - } - Ok(()) - } - - /// Merge a slice of references into a single new sketch. All inputs - /// must share the same dimensions; returns `Err` on mismatch or an - /// empty input. - pub fn merge_refs( - inputs: &[&CountSketch], - ) -> Result> { - let first = inputs - .first() - .ok_or("CountSketch::merge_refs called with empty input")?; - let mut merged = CountSketch::new(first.row_num, first.col_num); - for cs in inputs { - merged.merge(cs)?; - } - Ok(merged) - } - - /// Serialize to MessagePack bytes (used by the legacy Arroyo path - /// and by PR I's `_ENCODING_MSGPACK` variant when that lands). - pub fn serialize_msgpack(&self) -> Vec { - rmp_serde::to_vec(self).unwrap_or_default() - } - - /// Deserialize from MessagePack bytes. - pub fn deserialize_msgpack( - buffer: &[u8], - ) -> Result> { - Ok(rmp_serde::from_slice(buffer)?) - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_new_empty() { - let cs = CountSketch::new(2, 3); - assert_eq!(cs.row_num, 2); - assert_eq!(cs.col_num, 3); - assert_eq!(cs.sketch(), &vec![vec![0.0, 0.0, 0.0], vec![0.0, 0.0, 0.0]]); - } - - #[test] - fn test_from_legacy_matrix() { - let m = vec![vec![1.0, -2.0, 3.0], vec![-4.0, 5.0, -6.0]]; - let cs = CountSketch::from_legacy_matrix(m.clone(), 2, 3); - assert_eq!(cs.sketch(), &m); - } - - #[test] - fn test_merge_element_wise() { - let mut a = CountSketch::from_legacy_matrix(vec![vec![1.0, 2.0], vec![3.0, 4.0]], 2, 2); - let b = CountSketch::from_legacy_matrix(vec![vec![-1.0, -2.0], vec![-3.0, -4.0]], 2, 2); - a.merge(&b).unwrap(); - assert_eq!(a.sketch(), &vec![vec![0.0, 0.0], vec![0.0, 0.0]]); - } - - #[test] - fn test_merge_dimension_mismatch() { - let mut a = CountSketch::new(2, 3); - let b = CountSketch::new(3, 3); - assert!(a.merge(&b).is_err()); - } - - #[test] - fn test_merge_refs() { - let a = CountSketch::from_legacy_matrix(vec![vec![1.0, 2.0]], 1, 2); - let b = CountSketch::from_legacy_matrix(vec![vec![3.0, 4.0]], 1, 2); - let c = CountSketch::from_legacy_matrix(vec![vec![5.0, 6.0]], 1, 2); - let merged = CountSketch::merge_refs(&[&a, &b, &c]).unwrap(); - assert_eq!(merged.sketch(), &vec![vec![9.0, 12.0]]); - } - - #[test] - fn test_apply_delta_additive() { - let mut cs = CountSketch::from_legacy_matrix( - vec![vec![1.0, -2.0, 3.0], vec![-4.0, 5.0, -6.0]], - 2, - 3, - ); - let delta = CountSketchDelta { - rows: 2, - cols: 3, - cells: vec![ - (0, 0, 10), // 1 + 10 = 11 - (0, 2, -3), // 3 - 3 = 0 - (1, 1, -15), // 5 - 15 = -10 - ], - l2: vec![], - }; - cs.apply_delta(&delta).unwrap(); - assert_eq!(cs.sketch(), &vec![vec![11.0, -2.0, 0.0], vec![-4.0, -10.0, -6.0]]); - } - - #[test] - fn test_apply_delta_matches_full_merge() { - let base = CountSketch::from_legacy_matrix( - vec![vec![1.0, 2.0], vec![3.0, 4.0]], - 2, - 2, - ); - let addition = CountSketch::from_legacy_matrix( - vec![vec![10.0, 0.0], vec![0.0, 20.0]], - 2, - 2, - ); - let mut via_merge = base.clone(); - via_merge.merge(&addition).unwrap(); - - let delta = CountSketchDelta { - rows: 2, - cols: 2, - cells: vec![(0, 0, 10), (1, 1, 20)], - l2: vec![], - }; - let mut via_delta = base; - via_delta.apply_delta(&delta).unwrap(); - assert_eq!(via_delta.sketch(), via_merge.sketch()); - } - - #[test] - fn test_apply_delta_out_of_range() { - let mut cs = CountSketch::new(2, 3); - let delta = CountSketchDelta { - rows: 2, - cols: 3, - cells: vec![(2, 0, 1)], // row 2 out of range for 2-row matrix - l2: vec![], - }; - assert!(cs.apply_delta(&delta).is_err()); - } - - #[test] - fn test_msgpack_round_trip() { - let original = - CountSketch::from_legacy_matrix(vec![vec![1.5, -2.5], vec![3.5, -4.5]], 2, 2); - let bytes = original.serialize_msgpack(); - let decoded = CountSketch::deserialize_msgpack(&bytes).unwrap(); - assert_eq!(decoded.sketch(), original.sketch()); - assert_eq!(decoded.row_num, original.row_num); - assert_eq!(decoded.col_num, original.col_num); - } -} diff --git a/asap-common/sketch-core/src/dd_sketch.rs b/asap-common/sketch-core/src/dd_sketch.rs deleted file mode 100644 index bccf5b8b..00000000 --- a/asap-common/sketch-core/src/dd_sketch.rs +++ /dev/null @@ -1,588 +0,0 @@ -//! DDSketch — log-bucketed quantile sketch, mergeable by store-index alignment. -//! -//! Parallel to `count_sketch::CountSketch`: the minimum viable surface -//! needed for the modified-OTLP `Metric.data = DDSketch{…}` hot path -//! (PR C-CountSketch follow-up). Holds the bucket counts, their -//! absolute-index base offset, and the aggregate `{count, sum, min, max}`. -//! -//! Merge semantics: two sketches with the same relative-accuracy -//! parameter `alpha` are merged by aligning bucket arrays along their -//! `store_offset` and summing counts element-wise, with `min`/`max` -//! combined via min/max and `count`/`sum` added. -//! -//! The wire format is the protobuf-encoded -//! `asap_sketchlib::proto::sketchlib::DDSketchState` emitted by -//! DataCollector's `ddsketchprocessor`. Quantile estimation against -//! stored data is intentionally deferred — queries currently return -//! a placeholder error and fall through to the §5.2 fallback. - -use serde::{Deserialize, Serialize}; - -/// Sparse delta between two consecutive DDSketch snapshots — the -/// input shape for [`DdSketch::apply_delta`]. Mirrors the -/// `DDSketchDelta` proto in `sketchlib-go/proto/ddsketch/ddsketch.proto` -/// (and its Rust bindings vendored in `asap_otel_proto::sketchlib::v1`). -/// Kept as a plain struct in sketch-core so the pure-math crate doesn't -/// need a tonic/prost dependency; proto decode lives in the accumulator. -#[derive(Debug, Clone, Default)] -pub struct DdSketchDelta { - /// `(absolute_bucket_index, Δcount)` pairs, additive. - pub buckets: Vec<(i32, u64)>, - /// Δ total count. May be negative (signed on the wire). - pub d_count: i64, - /// Δ sum. - pub d_sum: f64, - /// Whether `new_min` carries a meaningful value. Min can only - /// decrease; a delta that didn't lower min sends `false`. - pub min_changed: bool, - pub new_min: f64, - /// Whether `new_max` carries a meaningful value. Max can only - /// increase. - pub max_changed: bool, - pub new_max: f64, -} - -/// Minimal DDSketch state — bucket counts + alpha + aggregates. -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct DdSketch { - /// Relative accuracy parameter; must satisfy `0 < alpha < 1`. - pub alpha: f64, - /// Bucket counts in absolute-index order. The absolute index of - /// `store_counts[i]` is `i + store_offset`. - pub store_counts: Vec, - /// Absolute bucket index corresponding to `store_counts[0]`. May - /// be negative. - pub store_offset: i32, - pub count: u64, - pub sum: f64, - pub min: f64, - pub max: f64, -} - -impl DdSketch { - /// Construct an empty sketch. - pub fn new(alpha: f64) -> Self { - Self { - alpha, - store_counts: Vec::new(), - store_offset: 0, - count: 0, - sum: 0.0, - min: f64::INFINITY, - max: f64::NEG_INFINITY, - } - } - - /// Construct from the decoded wire fields. - #[allow(clippy::too_many_arguments)] - pub fn from_raw( - alpha: f64, - store_counts: Vec, - store_offset: i32, - count: u64, - sum: f64, - min: f64, - max: f64, - ) -> Self { - Self { - alpha, - store_counts, - store_offset, - count, - sum, - min, - max, - } - } - - /// Merge one other sketch into self by aligning bucket arrays on - /// absolute indices. Both operands must share the same `alpha`. - pub fn merge( - &mut self, - other: &DdSketch, - ) -> Result<(), Box> { - if (self.alpha - other.alpha).abs() > f64::EPSILON { - return Err(format!( - "DdSketch alpha mismatch: self={}, other={}", - self.alpha, other.alpha - ) - .into()); - } - - if other.store_counts.is_empty() { - self.count += other.count; - self.sum += other.sum; - if other.min < self.min { - self.min = other.min; - } - if other.max > self.max { - self.max = other.max; - } - return Ok(()); - } - if self.store_counts.is_empty() { - self.store_counts = other.store_counts.clone(); - self.store_offset = other.store_offset; - } else { - let self_start = self.store_offset as i64; - let self_end = self_start + self.store_counts.len() as i64; - let other_start = other.store_offset as i64; - let other_end = other_start + other.store_counts.len() as i64; - let new_start = self_start.min(other_start); - let new_end = self_end.max(other_end); - let new_len = (new_end - new_start) as usize; - let mut merged = vec![0u64; new_len]; - for (i, c) in self.store_counts.iter().enumerate() { - let idx = (self_start + i as i64 - new_start) as usize; - merged[idx] = merged[idx].saturating_add(*c); - } - for (i, c) in other.store_counts.iter().enumerate() { - let idx = (other_start + i as i64 - new_start) as usize; - merged[idx] = merged[idx].saturating_add(*c); - } - self.store_counts = merged; - self.store_offset = new_start as i32; - } - self.count += other.count; - self.sum += other.sum; - if other.min < self.min { - self.min = other.min; - } - if other.max > self.max { - self.max = other.max; - } - Ok(()) - } - - /// Apply a sparse delta to this sketch in place. Matches the - /// `ApplyDelta` logic in `sketchlib-go/sketches/DDSketch/delta.go`: - /// bucket counts add, total count + sum add, min can only decrease - /// and max can only increase. Used by the backend ingest path to - /// reconstitute a full sketch from a base snapshot + subsequent - /// delta-transmission frames (paper §6.2 B3 / B4 baselines). - pub fn apply_delta(&mut self, delta: &DdSketchDelta) { - for (abs_idx, d_count) in &delta.buckets { - if self.store_counts.is_empty() { - self.store_counts = vec![0u64; 1]; - self.store_offset = *abs_idx; - } - let cur_start = self.store_offset as i64; - let cur_end = cur_start + self.store_counts.len() as i64; - let k = *abs_idx as i64; - if k < cur_start { - // Prepend zeros. - let pad = (cur_start - k) as usize; - let mut buf = vec![0u64; pad]; - buf.append(&mut self.store_counts); - self.store_counts = buf; - self.store_offset = *abs_idx; - } else if k >= cur_end { - let pad = (k - cur_end + 1) as usize; - self.store_counts - .extend(std::iter::repeat(0u64).take(pad)); - } - let arr_idx = (k - self.store_offset as i64) as usize; - self.store_counts[arr_idx] = - self.store_counts[arr_idx].saturating_add(*d_count); - } - if delta.d_count >= 0 { - self.count = self.count.saturating_add(delta.d_count as u64); - } else { - self.count = self.count.saturating_sub((-delta.d_count) as u64); - } - self.sum += delta.d_sum; - if delta.min_changed && delta.new_min < self.min { - self.min = delta.new_min; - } - if delta.max_changed && delta.new_max > self.max { - self.max = delta.new_max; - } - } - - /// Merge a slice of references into a single new sketch. Returns - /// `Err` on alpha mismatch or an empty input. - pub fn merge_refs( - inputs: &[&DdSketch], - ) -> Result> { - let first = inputs - .first() - .ok_or("DdSketch::merge_refs called with empty input")?; - let mut merged = DdSketch::new(first.alpha); - for d in inputs { - merged.merge(d)?; - } - Ok(merged) - } - - /// Insert a single positive value. Updates count, sum, min/max - /// and increments the bucket for `floor(ln(v) / ln(gamma))` - /// where `gamma = (1+α)/(1-α)`. Provided primarily so tests can - /// build a ground-truth sketch to compare delta-apply output - /// against. - pub fn insert(&mut self, value: f64) { - if value <= 0.0 { - // DDSketch is defined for positive reals; the paper's - // sketchlib-go rejects non-positive values silently. - return; - } - let gamma = (1.0 + self.alpha) / (1.0 - self.alpha); - let ln_gamma = gamma.ln(); - let idx = (value.ln() / ln_gamma).floor() as i32; - if self.store_counts.is_empty() { - self.store_counts = vec![1]; - self.store_offset = idx; - } else { - let cur_start = self.store_offset as i64; - let cur_end = cur_start + self.store_counts.len() as i64; - let k = idx as i64; - if k < cur_start { - let pad = (cur_start - k) as usize; - let mut buf = vec![0u64; pad]; - buf.append(&mut self.store_counts); - self.store_counts = buf; - self.store_offset = idx; - } else if k >= cur_end { - let pad = (k - cur_end + 1) as usize; - self.store_counts.extend(std::iter::repeat(0u64).take(pad)); - } - let arr_idx = (k - self.store_offset as i64) as usize; - self.store_counts[arr_idx] = self.store_counts[arr_idx].saturating_add(1); - } - self.count = self.count.saturating_add(1); - self.sum += value; - if value < self.min { - self.min = value; - } - if value > self.max { - self.max = value; - } - } - - /// Estimate the quantile at rank `q` ∈ [0, 1]. Walks the bucket - /// array in ascending absolute-index order, accumulating counts - /// until the target rank; returns the bucket's representative - /// value `gamma^(k + 0.5)` where `k` is the bucket's absolute - /// index. Returns `None` if the sketch is empty. - /// - /// Accuracy: bounded by DDSketch's α parameter — the estimated - /// quantile value is within `(1+α)/(1-α)` relative error of the - /// true quantile. - pub fn quantile(&self, q: f64) -> Option { - if self.count == 0 || self.store_counts.is_empty() { - return None; - } - let target = (q * (self.count.saturating_sub(1)) as f64).floor() as u64; - let mut cumulative: u64 = 0; - let gamma = (1.0 + self.alpha) / (1.0 - self.alpha); - for (i, &c) in self.store_counts.iter().enumerate() { - cumulative = cumulative.saturating_add(c); - if cumulative > target { - let k = (self.store_offset as i64 + i as i64) as f64; - // Bucket midpoint: gamma^(k + 0.5) — centers the - // estimate in the logarithmic bucket. - return Some(gamma.powf(k + 0.5)); - } - } - // Numerical edge case: if we fall off the end, return max. - Some(self.max) - } - - /// Serialize to MessagePack bytes. - pub fn serialize_msgpack(&self) -> Vec { - rmp_serde::to_vec(self).unwrap_or_default() - } - - /// Deserialize from MessagePack bytes. - pub fn deserialize_msgpack( - buffer: &[u8], - ) -> Result> { - Ok(rmp_serde::from_slice(buffer)?) - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_new_empty() { - let d = DdSketch::new(0.01); - assert_eq!(d.count, 0); - assert!(d.store_counts.is_empty()); - assert_eq!(d.min, f64::INFINITY); - assert_eq!(d.max, f64::NEG_INFINITY); - } - - #[test] - fn test_merge_aligned_same_offset() { - let mut a = DdSketch::from_raw(0.01, vec![1, 2, 3], -1, 6, 30.0, 1.0, 5.0); - let b = DdSketch::from_raw(0.01, vec![10, 20, 30], -1, 60, 300.0, 0.5, 6.0); - a.merge(&b).unwrap(); - assert_eq!(a.store_counts, vec![11, 22, 33]); - assert_eq!(a.store_offset, -1); - assert_eq!(a.count, 66); - assert_eq!(a.sum, 330.0); - assert_eq!(a.min, 0.5); - assert_eq!(a.max, 6.0); - } - - #[test] - fn test_merge_overlapping_offsets() { - // a covers indices [-1, 0, 1]; b covers indices [0, 1, 2] - let mut a = DdSketch::from_raw(0.01, vec![1, 1, 1], -1, 3, 3.0, 1.0, 3.0); - let b = DdSketch::from_raw(0.01, vec![10, 10, 10], 0, 30, 30.0, 1.0, 3.0); - a.merge(&b).unwrap(); - // Merged window is [-1, 0, 1, 2] → [1, 11, 11, 10] - assert_eq!(a.store_counts, vec![1, 11, 11, 10]); - assert_eq!(a.store_offset, -1); - assert_eq!(a.count, 33); - } - - #[test] - fn test_merge_disjoint_offsets() { - let mut a = DdSketch::from_raw(0.01, vec![1, 2], 0, 3, 3.0, 1.0, 2.0); - let b = DdSketch::from_raw(0.01, vec![3, 4], 5, 7, 7.0, 5.0, 6.0); - a.merge(&b).unwrap(); - // Window [0..7) → [1,2,0,0,0,3,4] - assert_eq!(a.store_counts, vec![1, 2, 0, 0, 0, 3, 4]); - assert_eq!(a.store_offset, 0); - } - - #[test] - fn test_apply_delta_additive_inside_store() { - let mut base = DdSketch::from_raw(0.01, vec![1, 2, 3], -1, 6, 30.0, 1.0, 5.0); - let delta = DdSketchDelta { - buckets: vec![(-1, 4), (0, 8), (1, 12)], - d_count: 24, - d_sum: 120.0, - min_changed: false, - new_min: 0.0, - max_changed: true, - new_max: 9.0, - }; - base.apply_delta(&delta); - assert_eq!(base.store_counts, vec![5, 10, 15]); - assert_eq!(base.count, 30); - assert_eq!(base.sum, 150.0); - assert_eq!(base.min, 1.0); - assert_eq!(base.max, 9.0); - } - - #[test] - fn test_apply_delta_expands_store_on_new_bucket() { - // Base covers [0..2]; delta adds a bucket at absolute index 4. - let mut base = DdSketch::from_raw(0.01, vec![1, 2], 0, 3, 3.0, 1.0, 2.0); - let delta = DdSketchDelta { - buckets: vec![(4, 7)], - d_count: 7, - d_sum: 35.0, - min_changed: false, - new_min: 0.0, - max_changed: true, - new_max: 6.0, - }; - base.apply_delta(&delta); - assert_eq!(base.store_counts, vec![1, 2, 0, 0, 7]); - assert_eq!(base.store_offset, 0); - assert_eq!(base.count, 10); - assert_eq!(base.max, 6.0); - } - - #[test] - fn test_apply_delta_matches_full_merge() { - // Snapshot the sketch, add more samples via a merge, and confirm - // the delta+apply path lands at the same state. - let base = DdSketch::from_raw(0.01, vec![1, 2, 3], 0, 6, 12.0, 1.0, 3.0); - let addition = DdSketch::from_raw(0.01, vec![10, 0, 20], 0, 30, 70.0, 0.5, 5.0); - let mut via_merge = base.clone(); - via_merge.merge(&addition).unwrap(); - - let delta = DdSketchDelta { - buckets: vec![(0, 10), (2, 20)], - d_count: 30, - d_sum: 70.0, - min_changed: true, - new_min: 0.5, - max_changed: true, - new_max: 5.0, - }; - let mut via_delta = base; - via_delta.apply_delta(&delta); - - assert_eq!(via_delta.store_counts, via_merge.store_counts); - assert_eq!(via_delta.count, via_merge.count); - assert_eq!(via_delta.sum, via_merge.sum); - assert_eq!(via_delta.min, via_merge.min); - assert_eq!(via_delta.max, via_merge.max); - } - - #[test] - fn test_merge_alpha_mismatch() { - let mut a = DdSketch::new(0.01); - let b = DdSketch::new(0.02); - assert!(a.merge(&b).is_err()); - } - - #[test] - fn test_msgpack_round_trip() { - let original = DdSketch::from_raw(0.01, vec![1, 2, 3], -2, 6, 30.0, 1.0, 5.0); - let bytes = original.serialize_msgpack(); - let decoded = DdSketch::deserialize_msgpack(&bytes).unwrap(); - assert_eq!(decoded.store_counts, original.store_counts); - assert_eq!(decoded.store_offset, original.store_offset); - assert_eq!(decoded.count, original.count); - } - - #[test] - fn test_insert_and_quantile_lognormal() { - // Ground-truth: insert a large i.i.d. log-normal sample into - // a full sketch and sanity-check P50 / P90 / P99. - let mut gt = DdSketch::new(0.01); - let mut rng = 0xdead_beefu64; - let mut next = || { - // xorshift64* - rng ^= rng << 13; - rng ^= rng >> 7; - rng ^= rng << 17; - rng - }; - // Box-Muller normal → log-normal(mu=3, sigma=0.7). - let lognormal = |u: u64, v: u64| -> f64 { - let r1 = (u as f64) / (u64::MAX as f64).max(1.0); - let r2 = (v as f64) / (u64::MAX as f64).max(1.0); - let z = (-2.0 * r1.max(1e-12).ln()).sqrt() * (2.0 * std::f64::consts::PI * r2).cos(); - (3.0 + 0.7 * z).exp() - }; - for _ in 0..100_000 { - gt.insert(lognormal(next(), next())); - } - let p50 = gt.quantile(0.5).unwrap(); - let p99 = gt.quantile(0.99).unwrap(); - // Analytical P50 = exp(mu) = e^3 ≈ 20.09; - // P99 ≈ exp(mu + sigma * Φ⁻¹(0.99)) = e^(3 + 0.7×2.326) ≈ 102.4. - assert!((p50 / 20.09).ln().abs() < 0.05, "P50 {} not close to 20.09", p50); - assert!((p99 / 102.4).ln().abs() < 0.05, "P99 {} not close to 102.4", p99); - } - - /// Core accuracy claim for PRs #60-#63 end-to-end: building a - /// sketch via `base + apply_delta()` produces quantile estimates - /// within DDSketch's α bound of the ground-truth full-sketch - /// path. If this fails, the paper's delta-reconstitution story - /// is broken. - #[test] - fn test_delta_chain_preserves_quantile_accuracy() { - let alpha = 0.01; - let mut rng = 0xcafe_babeu64; - let mut next = || { - rng ^= rng << 13; - rng ^= rng >> 7; - rng ^= rng << 17; - rng - }; - let lognormal = |u: u64, v: u64| -> f64 { - let r1 = (u as f64) / (u64::MAX as f64).max(1.0); - let r2 = (v as f64) / (u64::MAX as f64).max(1.0); - let z = (-2.0 * r1.max(1e-12).ln()).sqrt() * (2.0 * std::f64::consts::PI * r2).cos(); - (3.0 + 0.7 * z).exp() - }; - - // Path A (ground truth): one sketch, 50k samples inserted - // directly. - let mut full = DdSketch::new(alpha); - // Path B (delta chain): a base sketch from the first 10k - // samples, then 4 incremental "flushes" of 10k samples each, - // each transmitted as a delta computed against the previous - // snapshot. - let mut reconstituted = DdSketch::new(alpha); - let mut prev_snapshot = DdSketch::new(alpha); // what the "receiver" has cached - - let batch = 10_000; - let batches = 5; - for b in 0..batches { - let mut this_batch = prev_snapshot.clone(); - for _ in 0..batch { - let v = lognormal(next(), next()); - full.insert(v); - this_batch.insert(v); - } - // Compute a "delta" = diff of this_batch vs prev_snapshot - // in our in-memory struct shape. Matches what - // `sketchlib-go/sketches/DDSketch/delta.go::ComputeDelta` - // would put on the wire. - let delta = compute_dd_delta(&prev_snapshot, &this_batch); - if b == 0 { - // First batch seeds the reconstituted sketch. - reconstituted = this_batch.clone(); - } else { - reconstituted.apply_delta(&delta); - } - prev_snapshot = this_batch; - } - - // Both paths should see the same total count + sum (exact). - assert_eq!(reconstituted.count, full.count, "count diverged"); - assert!( - (reconstituted.sum - full.sum).abs() < 1e-6, - "sum diverged: recon={} full={}", - reconstituted.sum, - full.sum - ); - - // And P50 / P90 / P99 should agree within α bound. - for q in [0.5, 0.9, 0.99] { - let got = reconstituted.quantile(q).unwrap(); - let want = full.quantile(q).unwrap(); - let rel_err = (got / want - 1.0).abs(); - assert!( - rel_err <= alpha, - "q={} rel_err={:.4} exceeds α={}: reconstituted={}, full={}", - q, rel_err, alpha, got, want, - ); - } - } - - /// Helper used only in the delta-chain test: computes a - /// `DdSketchDelta` from two snapshots. Mirrors the sketchlib-go - /// `ComputeDelta` logic so the test exercises the wire-format - /// path end-to-end. - fn compute_dd_delta(snapshot: &DdSketch, current: &DdSketch) -> DdSketchDelta { - let mut cells = Vec::new(); - if !current.store_counts.is_empty() { - for (i, &c) in current.store_counts.iter().enumerate() { - if c == 0 { - continue; - } - let k = current.store_offset + i as i32; - let snap_count: u64 = if !snapshot.store_counts.is_empty() { - let idx = k as i64 - snapshot.store_offset as i64; - if idx >= 0 && (idx as usize) < snapshot.store_counts.len() { - snapshot.store_counts[idx as usize] - } else { - 0 - } - } else { - 0 - }; - let dc = c.saturating_sub(snap_count); - if dc > 0 { - cells.push((k, dc)); - } - } - } - let d_count = current.count as i64 - snapshot.count as i64; - let d_sum = current.sum - snapshot.sum; - let min_changed = current.count > 0 - && (snapshot.count == 0 || current.min < snapshot.min); - let max_changed = current.count > 0 - && (snapshot.count == 0 || current.max > snapshot.max); - DdSketchDelta { - buckets: cells, - d_count, - d_sum, - min_changed, - new_min: current.min, - max_changed, - new_max: current.max, - } - } -} diff --git a/asap-common/sketch-core/src/delta_set_aggregator.rs b/asap-common/sketch-core/src/delta_set_aggregator.rs deleted file mode 100644 index c086e2a2..00000000 --- a/asap-common/sketch-core/src/delta_set_aggregator.rs +++ /dev/null @@ -1,71 +0,0 @@ -// Adapted from QueryEngineRust/src/precompute_operators/delta_set_aggregator_accumulator.rs -// Changes: -// - Only the wire format (DeltaResult) and its serialize/deserialize functions are extracted. -// - The Arroyo side uses lazy_static for stateful window tracking — that streaming logic -// stays in the Arroyo template and does NOT belong in sketch-core. -// - DeltaResult made pub (was private inline struct in QE). -// - serialize_msgpack / deserialize_msgpack are module-level free functions -// (not methods on DeltaSetAggregatorAccumulator, which stays in QE). -// - Removed: all QE accumulator struct/impls (stay in QE) - -use serde::{Deserialize, Serialize}; -use std::collections::HashSet; - -/// Wire format for the delta set aggregator — shared between Arroyo and QueryEngineRust. -/// Both sides agree on `{ added: HashSet, removed: HashSet }` in msgpack. -#[derive(Serialize, Deserialize, Clone, Debug)] -pub struct DeltaResult { - pub added: HashSet, - pub removed: HashSet, -} - -/// Serialize a delta result to MessagePack. -pub fn serialize_msgpack(added: &HashSet, removed: &HashSet) -> Vec { - let result = DeltaResult { - added: added.clone(), - removed: removed.clone(), - }; - let mut buf = Vec::new(); - rmp_serde::encode::write(&mut buf, &result).unwrap(); - buf -} - -/// Deserialize a delta result from MessagePack produced by the Arroyo UDF. -pub fn deserialize_msgpack(buffer: &[u8]) -> Result> { - rmp_serde::from_slice(buffer) - .map_err(|e| format!("Failed to deserialize DeltaResult from MessagePack: {e}").into()) -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_msgpack_round_trip() { - let mut added = HashSet::new(); - added.insert("web".to_string()); - added.insert("api".to_string()); - - let mut removed = HashSet::new(); - removed.insert("db".to_string()); - - let bytes = serialize_msgpack(&added, &removed); - let result = deserialize_msgpack(&bytes).unwrap(); - - assert_eq!(result.added.len(), 2); - assert!(result.added.contains("web")); - assert!(result.added.contains("api")); - assert_eq!(result.removed.len(), 1); - assert!(result.removed.contains("db")); - } - - #[test] - fn test_empty_sets() { - let added = HashSet::new(); - let removed = HashSet::new(); - let bytes = serialize_msgpack(&added, &removed); - let result = deserialize_msgpack(&bytes).unwrap(); - assert!(result.added.is_empty()); - assert!(result.removed.is_empty()); - } -} diff --git a/asap-common/sketch-core/src/hll_sketch.rs b/asap-common/sketch-core/src/hll_sketch.rs deleted file mode 100644 index db28f473..00000000 --- a/asap-common/sketch-core/src/hll_sketch.rs +++ /dev/null @@ -1,296 +0,0 @@ -//! HyperLogLog sketch — register-wise mergeable cardinality estimator. -//! -//! Parallel to `count_sketch::CountSketch`: the minimum viable surface -//! needed for the modified-OTLP `Metric.data = HLLSketch{…}` hot path -//! (PR C-CountSketch follow-up). Wraps a flat `Vec` of register -//! values (length = `2^precision`) and merges element-wise by taking -//! the maximum across aligned registers, which is the standard HLL -//! merge semantics. -//! -//! The wire format is the protobuf-encoded -//! `asap_sketchlib::proto::sketchlib::HyperLogLogState` emitted by -//! DataCollector's `hllprocessor`. This type carries the register -//! bytes and the variant/precision metadata losslessly, so the -//! merge + store round-trip works end-to-end. Cardinality estimation -//! against stored HLL data is intentionally deferred to a follow-up -//! — queries currently return a placeholder error and fall through -//! to the §5.2 fallback. - -use serde::{Deserialize, Serialize}; - -/// HLL estimator variant. Mirrors `asap_sketchlib::proto::sketchlib::HllVariant` -/// so the proto round-trip preserves the algorithm identity — the three -/// variants are not mutually compatible on register contents. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] -pub enum HllVariant { - Unspecified, - Regular, - Datafusion, - Hip, -} - -/// Sparse delta between two consecutive HLL snapshots — the input -/// shape for [`HllSketch::apply_delta`]. Mirrors the `HLLDelta` proto -/// in `sketchlib-go/proto/hll/hll.proto` (and its Rust bindings -/// vendored in `asap_otel_proto::sketchlib::v1`). HLL registers merge -/// with max semantics, so a delta carries only the register indices -/// whose value increased since the last snapshot. -#[derive(Debug, Clone, Default)] -pub struct HllDelta { - /// `(register_index, new_value)` pairs. `new_value` is the full - /// post-update register value; `apply_delta` does - /// `registers[i] = max(registers[i], new_value)`. - pub updates: Vec<(u32, u8)>, -} - -/// Minimal HLL state — registers + variant + precision. Register-wise -/// mergeable (max over aligned cells). -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct HllSketch { - pub variant: HllVariant, - pub precision: u32, - /// Flat register array, length = `2^precision`. - pub registers: Vec, - /// HIP accumulator components — populated only when `variant == Hip`. - pub hip_kxq0: f64, - pub hip_kxq1: f64, - pub hip_est: f64, -} - -impl HllSketch { - /// Construct an empty sketch at the given precision. - pub fn new(variant: HllVariant, precision: u32) -> Self { - let n = 1usize << precision; - Self { - variant, - precision, - registers: vec![0u8; n], - hip_kxq0: 0.0, - hip_kxq1: 0.0, - hip_est: 0.0, - } - } - - /// Construct from pre-built register bytes (used by the modified-OTLP - /// proto-decode path). - pub fn from_raw( - variant: HllVariant, - precision: u32, - registers: Vec, - hip_kxq0: f64, - hip_kxq1: f64, - hip_est: f64, - ) -> Self { - Self { - variant, - precision, - registers, - hip_kxq0, - hip_kxq1, - hip_est, - } - } - - /// Merge one other sketch into self via register-wise max. Both - /// operands must have identical variant and precision. - pub fn merge( - &mut self, - other: &HllSketch, - ) -> Result<(), Box> { - if self.variant != other.variant { - return Err(format!( - "HllSketch variant mismatch: self={:?}, other={:?}", - self.variant, other.variant - ) - .into()); - } - if self.precision != other.precision { - return Err(format!( - "HllSketch precision mismatch: self={}, other={}", - self.precision, other.precision - ) - .into()); - } - if self.registers.len() != other.registers.len() { - return Err(format!( - "HllSketch register-length mismatch: self={}, other={}", - self.registers.len(), - other.registers.len() - ) - .into()); - } - for (s, o) in self.registers.iter_mut().zip(other.registers.iter()) { - if *o > *s { - *s = *o; - } - } - // HIP accumulators add on merge (each source carried its own - // running estimate; merged state inherits the combined - // components). - if self.variant == HllVariant::Hip { - self.hip_kxq0 += other.hip_kxq0; - self.hip_kxq1 += other.hip_kxq1; - self.hip_est += other.hip_est; - } - Ok(()) - } - - /// Apply a sparse register delta in place. Matches the - /// `registers[i] = max(registers[i], new_value)` logic in - /// `sketchlib-go/sketches/HLL/delta.go::ApplyRegisterDelta`. Used - /// by the backend ingest path to reconstitute a full sketch from - /// a base snapshot + subsequent delta-transmission frames (paper - /// §6.2 B3 / B4 baselines). - /// - /// Returns `Err` if any delta index is out of range for the - /// sketch's precision — indicating a precision mismatch between - /// the snapshot this sketch was built from and the delta sender. - pub fn apply_delta( - &mut self, - delta: &HllDelta, - ) -> Result<(), Box> { - let n = self.registers.len(); - for (idx, new_val) in &delta.updates { - let i = *idx as usize; - if i >= n { - return Err(format!( - "HllDelta index {i} out of range (precision={} → {n} registers)", - self.precision - ) - .into()); - } - if *new_val > self.registers[i] { - self.registers[i] = *new_val; - } - } - Ok(()) - } - - /// Merge a slice of references into a single new sketch. All inputs - /// must share the same variant and precision; returns `Err` on - /// mismatch or an empty input. - pub fn merge_refs( - inputs: &[&HllSketch], - ) -> Result> { - let first = inputs - .first() - .ok_or("HllSketch::merge_refs called with empty input")?; - let mut merged = HllSketch::new(first.variant, first.precision); - for hll in inputs { - merged.merge(hll)?; - } - Ok(merged) - } - - /// Serialize to MessagePack bytes (used by the legacy Arroyo path - /// and by PR I's `_ENCODING_MSGPACK` variant when that lands). - pub fn serialize_msgpack(&self) -> Vec { - rmp_serde::to_vec(self).unwrap_or_default() - } - - /// Deserialize from MessagePack bytes. - pub fn deserialize_msgpack( - buffer: &[u8], - ) -> Result> { - Ok(rmp_serde::from_slice(buffer)?) - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_new_empty() { - let h = HllSketch::new(HllVariant::Regular, 4); - assert_eq!(h.registers.len(), 16); - assert!(h.registers.iter().all(|&r| r == 0)); - } - - #[test] - fn test_merge_register_wise_max() { - let mut a = HllSketch::from_raw(HllVariant::Regular, 2, vec![1, 5, 3, 7], 0.0, 0.0, 0.0); - let b = HllSketch::from_raw(HllVariant::Regular, 2, vec![4, 2, 6, 0], 0.0, 0.0, 0.0); - a.merge(&b).unwrap(); - assert_eq!(a.registers, vec![4, 5, 6, 7]); - } - - #[test] - fn test_apply_delta_max_semantics() { - let mut h = - HllSketch::from_raw(HllVariant::Regular, 2, vec![1, 5, 3, 7], 0.0, 0.0, 0.0); - let delta = HllDelta { - updates: vec![(0, 4), (1, 2), (2, 6), (3, 0)], - }; - h.apply_delta(&delta).unwrap(); - // reg[0]: max(1,4)=4, reg[1]: max(5,2)=5, reg[2]: max(3,6)=6, - // reg[3]: max(7,0)=7. - assert_eq!(h.registers, vec![4, 5, 6, 7]); - } - - #[test] - fn test_apply_delta_out_of_range() { - let mut h = HllSketch::new(HllVariant::Regular, 2); // 4 registers - let delta = HllDelta { - updates: vec![(7, 3)], - }; - assert!(h.apply_delta(&delta).is_err()); - } - - #[test] - fn test_apply_delta_matches_full_merge() { - let base = - HllSketch::from_raw(HllVariant::Regular, 2, vec![1, 5, 3, 7], 0.0, 0.0, 0.0); - let addition = - HllSketch::from_raw(HllVariant::Regular, 2, vec![4, 0, 6, 0], 0.0, 0.0, 0.0); - let mut via_merge = base.clone(); - via_merge.merge(&addition).unwrap(); - - let delta = HllDelta { - updates: vec![(0, 4), (2, 6)], - }; - let mut via_delta = base; - via_delta.apply_delta(&delta).unwrap(); - assert_eq!(via_delta.registers, via_merge.registers); - } - - #[test] - fn test_merge_variant_mismatch() { - let mut a = HllSketch::new(HllVariant::Regular, 4); - let b = HllSketch::new(HllVariant::Datafusion, 4); - assert!(a.merge(&b).is_err()); - } - - #[test] - fn test_merge_precision_mismatch() { - let mut a = HllSketch::new(HllVariant::Regular, 4); - let b = HllSketch::new(HllVariant::Regular, 5); - assert!(a.merge(&b).is_err()); - } - - #[test] - fn test_merge_refs() { - let a = HllSketch::from_raw(HllVariant::Regular, 1, vec![1, 0], 0.0, 0.0, 0.0); - let b = HllSketch::from_raw(HllVariant::Regular, 1, vec![0, 3], 0.0, 0.0, 0.0); - let c = HllSketch::from_raw(HllVariant::Regular, 1, vec![2, 2], 0.0, 0.0, 0.0); - let merged = HllSketch::merge_refs(&[&a, &b, &c]).unwrap(); - assert_eq!(merged.registers, vec![2, 3]); - } - - #[test] - fn test_msgpack_round_trip() { - let original = HllSketch::from_raw( - HllVariant::Hip, - 3, - vec![0, 1, 2, 3, 4, 5, 6, 7], - 1.0, - 2.0, - 3.0, - ); - let bytes = original.serialize_msgpack(); - let decoded = HllSketch::deserialize_msgpack(&bytes).unwrap(); - assert_eq!(decoded.registers, original.registers); - assert_eq!(decoded.precision, original.precision); - assert_eq!(decoded.hip_kxq0, 1.0); - } -} diff --git a/asap-common/sketch-core/src/hydra_kll.rs b/asap-common/sketch-core/src/hydra_kll.rs deleted file mode 100644 index e6888d55..00000000 --- a/asap-common/sketch-core/src/hydra_kll.rs +++ /dev/null @@ -1,295 +0,0 @@ -// Adapted from QueryEngineRust/src/precompute_operators/hydra_kll_accumulator.rs -// Changes: -// - Renamed HydraKllSketchAccumulator -> HydraKllSketch -// - KllSketchData import replaced by crate::kll::{KllSketch, KllSketchData} -// - Inner cells are KllSketch instead of DatasketchesKLLAccumulator -// - update() takes &str instead of &KeyByLabelValues -// - query_key() takes &str; renamed to query() -// - serialize_to_bytes (trait) -> serialize_msgpack (inherent method) -// - deserialize_from_bytes_arroyo -> deserialize_msgpack -// - merge_accumulators -> merge -// - Removed: deserialize_from_bytes (stub, stays in QE) -// - Removed: AggregateCore, SerializableToSink, MergeableAccumulator, MultipleSubpopulationAggregate impls -// - Removed: base64, serde_json imports (QE-specific) -// - Added: aggregate_hydrakll() one-shot helper - -use crate::kll::{KllSketch, KllSketchData}; -use serde::{Deserialize, Serialize}; -use std::cmp::Ordering; -use xxhash_rust::xxh32::xxh32; - -#[derive(Serialize, Deserialize)] -struct HydraKllSketchData { - row_num: usize, - col_num: usize, - sketches: Vec>, -} - -#[derive(Debug, Clone)] -pub struct HydraKllSketch { - pub sketch: Vec>, - pub row_num: usize, - pub col_num: usize, -} - -impl HydraKllSketch { - pub fn new(row_num: usize, col_num: usize, k: u16) -> Self { - let sketch = vec![vec![KllSketch::new(k); col_num]; row_num]; - Self { - sketch, - row_num, - col_num, - } - } - - pub fn update(&mut self, key: &str, value: f64) { - let key_bytes = key.as_bytes(); - // Update each row using different hash functions - for i in 0..self.row_num { - let hash_value = xxh32(key_bytes, i as u32); - let col_index = (hash_value as usize) % self.col_num; - self.sketch[i][col_index].update(value); - } - } - - pub fn query(&self, key: &str, quantile: f64) -> f64 { - let key_bytes = key.as_bytes(); - let mut quantiles = Vec::with_capacity(self.row_num); - - for i in 0..self.row_num { - let hash_value = xxh32(key_bytes, i as u32); - let col_index = (hash_value as usize) % self.col_num; - quantiles.push(self.sketch[i][col_index].get_quantile(quantile)); - } - - if quantiles.is_empty() { - return 0.0; - } - - quantiles.sort_by(|a, b| match a.partial_cmp(b) { - Some(ordering) => ordering, - None => Ordering::Equal, - }); - - let mid = quantiles.len() / 2; - if quantiles.len() % 2 == 0 { - (quantiles[mid - 1] + quantiles[mid]) / 2.0 - } else { - quantiles[mid] - } - } - - pub fn merge( - accumulators: Vec, - ) -> Result> { - if accumulators.is_empty() { - return Err("No accumulators to merge".into()); - } - - // Check dimensions match - let row_num = accumulators[0].row_num; - let col_num = accumulators[0].col_num; - for acc in &accumulators { - if acc.row_num != row_num || acc.col_num != col_num { - return Err( - "Cannot merge HydraKllSketch accumulators with different dimensions".into(), - ); - } - } - - // Transpose Vec into Vec>> indexed [row][col][acc], - // consuming the owned accumulators so no per-cell clones are needed. - let mut by_cell: Vec>> = (0..row_num) - .map(|_| (0..col_num).map(|_| Vec::new()).collect()) - .collect(); - for acc in accumulators { - for (i, row) in acc.sketch.into_iter().enumerate() { - for (j, cell) in row.into_iter().enumerate() { - by_cell[i][j].push(cell); - } - } - } - - // Merge each cell independently - let mut merged_sketch = Vec::with_capacity(row_num); - for row in by_cell { - let mut merged_row = Vec::with_capacity(col_num); - for cells in row { - merged_row.push(KllSketch::merge(cells)?); - } - merged_sketch.push(merged_row); - } - - Ok(HydraKllSketch { - sketch: merged_sketch, - row_num, - col_num, - }) - } - - /// Serialize to MessagePack — matches the Arroyo UDF wire format exactly. - pub fn serialize_msgpack(&self) -> Vec { - let mut sketches = Vec::with_capacity(self.row_num); - for row in &self.sketch { - let mut row_data = Vec::with_capacity(self.col_num); - for cell in row { - // Serialize each KllSketch to KllSketchData - let cell_bytes = cell.serialize_msgpack(); - let kll_data: KllSketchData = rmp_serde::from_slice(&cell_bytes) - .expect("Failed to deserialize KllSketchData from cell"); - row_data.push(kll_data); - } - sketches.push(row_data); - } - - let serialized = HydraKllSketchData { - row_num: self.row_num, - col_num: self.col_num, - sketches, - }; - - let mut buf = Vec::new(); - rmp_serde::encode::write(&mut buf, &serialized).unwrap(); - buf - } - - /// Deserialize from MessagePack produced by the Arroyo UDF. - pub fn deserialize_msgpack(buffer: &[u8]) -> Result> { - let deserialized_sketch_data: HydraKllSketchData = rmp_serde::from_slice(buffer) - .map_err(|e| format!("Failed to deserialize HydraKLL from MessagePack: {e}"))?; - - if deserialized_sketch_data.sketches.len() != deserialized_sketch_data.row_num { - return Err(format!( - "HydraKLL row count mismatch: expected {}, got {}", - deserialized_sketch_data.row_num, - deserialized_sketch_data.sketches.len() - ) - .into()); - } - - let mut sketch: Vec> = Vec::with_capacity(deserialized_sketch_data.row_num); - - for (row_idx, row) in deserialized_sketch_data.sketches.into_iter().enumerate() { - if row.len() != deserialized_sketch_data.col_num { - return Err(format!( - "HydraKLL column count mismatch in row {}: expected {}, got {}", - row_idx, - deserialized_sketch_data.col_num, - row.len() - ) - .into()); - } - - let mut accum_row: Vec = - Vec::with_capacity(deserialized_sketch_data.col_num); - for cell in row { - let cell_bytes = rmp_serde::to_vec(&cell) - .map_err(|e| format!("Failed to serialize nested KLL sketch: {e}"))?; - let kll = KllSketch::deserialize_msgpack(&cell_bytes)?; - accum_row.push(kll); - } - - sketch.push(accum_row); - } - - Ok(Self { - sketch, - row_num: deserialized_sketch_data.row_num, - col_num: deserialized_sketch_data.col_num, - }) - } - - /// One-shot aggregation for the Arroyo UDAF call pattern. - pub fn aggregate_hydrakll( - row_num: usize, - col_num: usize, - k: u16, - keys: &[&str], - values: &[f64], - ) -> Option> { - if keys.is_empty() { - return None; - } - let mut sketch = Self::new(row_num, col_num, k); - for (key, &value) in keys.iter().zip(values.iter()) { - sketch.update(key, value); - } - Some(sketch.serialize_msgpack()) - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_creation() { - let h = HydraKllSketch::new(2, 3, 200); - assert_eq!(h.row_num, 2); - assert_eq!(h.col_num, 3); - assert_eq!(h.sketch.len(), 2); - assert_eq!(h.sketch[0].len(), 3); - } - - #[test] - fn test_update_and_query() { - let mut h = HydraKllSketch::new(2, 10, 200); - h.update("key1", 5.0); - h.update("key1", 10.0); - // With 2 values, median quantile should be between them - let q = h.query("key1", 0.5); - assert!(q >= 0.0); - } - - #[test] - fn test_merge() { - let mut h1 = HydraKllSketch::new(2, 5, 200); - let mut h2 = HydraKllSketch::new(2, 5, 200); - - for i in 1..=5 { - h1.update("key1", i as f64); - } - for i in 6..=10 { - h2.update("key1", i as f64); - } - - let merged = HydraKllSketch::merge(vec![h1, h2]).unwrap(); - assert_eq!(merged.row_num, 2); - assert_eq!(merged.col_num, 5); - } - - #[test] - fn test_merge_dimension_mismatch() { - let h1 = HydraKllSketch::new(2, 5, 200); - let h2 = HydraKllSketch::new(3, 5, 200); - assert!(HydraKllSketch::merge(vec![h1, h2]).is_err()); - } - - #[test] - fn test_msgpack_round_trip() { - let mut h = HydraKllSketch::new(2, 3, 200); - h.update("key1", 5.0); - h.update("key2", 10.0); - - let bytes = h.serialize_msgpack(); - let deserialized = HydraKllSketch::deserialize_msgpack(&bytes).unwrap(); - - assert_eq!(deserialized.row_num, 2); - assert_eq!(deserialized.col_num, 3); - } - - #[test] - fn test_aggregate_hydrakll() { - let keys = ["a", "b", "a"]; - let values = [1.0, 2.0, 3.0]; - let bytes = HydraKllSketch::aggregate_hydrakll(2, 5, 200, &keys, &values).unwrap(); - let h = HydraKllSketch::deserialize_msgpack(&bytes).unwrap(); - assert_eq!(h.row_num, 2); - assert_eq!(h.col_num, 5); - } - - #[test] - fn test_aggregate_hydrakll_empty() { - assert!(HydraKllSketch::aggregate_hydrakll(2, 5, 200, &[], &[]).is_none()); - } -} diff --git a/asap-common/sketch-core/src/kll.rs b/asap-common/sketch-core/src/kll.rs deleted file mode 100644 index 5751cb23..00000000 --- a/asap-common/sketch-core/src/kll.rs +++ /dev/null @@ -1,366 +0,0 @@ -// Adapted from QueryEngineRust/src/precompute_operators/datasketches_kll_accumulator.rs -// Changes: -// - Renamed DatasketchesKLLAccumulator -> KllSketch -// - KllSketchData made pub (used by hydra_kll) -// - _update -> pub update -// - serialize_to_bytes (trait) -> serialize_msgpack (inherent method) -// - deserialize_from_bytes_arroyo -> deserialize_msgpack -// - merge_accumulators -> merge -// - Removed: deserialize_from_json, deserialize_from_bytes (legacy QE formats, stay in QE) -// - Removed: merge_multiple (QE trait-object helper, stays in QE) -// - Removed: AggregateCore, SerializableToSink, MergeableAccumulator, SingleSubpopulationAggregate impls -// - Removed: base64, serde_json, tracing imports (QE-specific) -// - Added: aggregate_kll() one-shot helper - -use core::panic; -use dsrs::KllDoubleSketch; -use serde::{Deserialize, Serialize}; - -use crate::config::use_sketchlib_for_kll; -use crate::kll_sketchlib::{ - bytes_from_sketchlib_kll, new_sketchlib_kll, sketchlib_kll_from_bytes, sketchlib_kll_merge, - sketchlib_kll_quantile, sketchlib_kll_update, SketchlibKll, -}; - -/// Wire format used in MessagePack serialization (matches Arroyo UDF output). -#[derive(Deserialize, Serialize)] -pub struct KllSketchData { - pub k: u16, - pub sketch_bytes: Vec, -} - -/// Backend implementation for KLL Sketch. Only one is active at a time. -pub enum KllBackend { - /// dsrs (DataSketches) implementation. - Legacy(KllDoubleSketch), - /// asap_sketchlib backed implementation. - Sketchlib(SketchlibKll), -} - -impl std::fmt::Debug for KllBackend { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - KllBackend::Legacy(_) => write!(f, "Legacy(..)"), - KllBackend::Sketchlib(_) => write!(f, "Sketchlib(..)"), - } - } -} - -impl Clone for KllBackend { - fn clone(&self) -> Self { - match self { - KllBackend::Legacy(s) => { - if s.get_n() == 0 { - KllBackend::Legacy(KllDoubleSketch::with_k(200)) // k will be overwritten by KllSketch - } else { - let bytes = s.serialize(); - KllBackend::Legacy(KllDoubleSketch::deserialize(bytes.as_ref()).unwrap()) - } - } - KllBackend::Sketchlib(s) => KllBackend::Sketchlib(s.clone()), - } - } -} - -pub struct KllSketch { - pub k: u16, - pub backend: KllBackend, -} - -impl KllSketch { - pub fn new(k: u16) -> Self { - let backend = if use_sketchlib_for_kll() { - KllBackend::Sketchlib(new_sketchlib_kll(k)) - } else { - KllBackend::Legacy(KllDoubleSketch::with_k(k)) - }; - Self { k, backend } - } - - /// Returns the raw sketch bytes (for JSON serialization, etc.). - pub fn sketch_bytes(&self) -> Vec { - match &self.backend { - KllBackend::Legacy(s) => s.serialize().as_ref().to_vec(), - KllBackend::Sketchlib(s) => bytes_from_sketchlib_kll(s), - } - } - - pub fn update(&mut self, value: f64) { - match &mut self.backend { - KllBackend::Legacy(s) => s.update(value), - KllBackend::Sketchlib(s) => sketchlib_kll_update(s, value), - } - } - - pub fn count(&self) -> u64 { - match &self.backend { - KllBackend::Legacy(s) => s.get_n(), - KllBackend::Sketchlib(s) => s.count() as u64, - } - } - - pub fn get_quantile(&self, quantile: f64) -> f64 { - if self.count() == 0 { - return 0.0; - } - match &self.backend { - KllBackend::Legacy(s) => s.get_quantile(quantile), - KllBackend::Sketchlib(s) => sketchlib_kll_quantile(s, quantile), - } - } - - pub fn merge( - accumulators: Vec, - ) -> Result> { - if accumulators.is_empty() { - return Err("No accumulators to merge".into()); - } - - let k = accumulators[0].k; - for acc in &accumulators { - if acc.k != k { - return Err("Cannot merge KllSketch with different k values".into()); - } - } - - let mut merged = KllSketch::new(k); - match &mut merged.backend { - KllBackend::Legacy(merged_legacy) => { - for acc in accumulators { - if let KllBackend::Legacy(acc_legacy) = acc.backend { - merged_legacy.merge(&acc_legacy); - } else { - return Err("Cannot merge Legacy with Sketchlib KLL".into()); - } - } - } - KllBackend::Sketchlib(merged_sketchlib) => { - for acc in accumulators { - if let KllBackend::Sketchlib(acc_sketchlib) = &acc.backend { - sketchlib_kll_merge(merged_sketchlib, acc_sketchlib); - } else { - return Err("Cannot merge Sketchlib with Legacy KLL".into()); - } - } - } - } - - Ok(merged) - } - - /// Serialize to MessagePack — matches the Arroyo UDF wire format exactly. - pub fn serialize_msgpack(&self) -> Vec { - let sketch_bytes = self.sketch_bytes(); - let serialized = KllSketchData { - k: self.k, - sketch_bytes, - }; - - let mut buf = Vec::new(); - match rmp_serde::encode::write(&mut buf, &serialized) { - Ok(_) => buf, - Err(_) => { - panic!("Failed to serialize KllSketchData to MessagePack"); - } - } - } - - /// Deserialize from MessagePack produced by the Arroyo UDF. - pub fn deserialize_msgpack(buffer: &[u8]) -> Result> { - let wire: KllSketchData = rmp_serde::from_slice(buffer) - .map_err(|e| format!("Failed to deserialize KllSketchData from MessagePack: {e}"))?; - - let backend = if use_sketchlib_for_kll() { - KllBackend::Sketchlib(sketchlib_kll_from_bytes(&wire.sketch_bytes)?) - } else { - KllBackend::Legacy( - KllDoubleSketch::deserialize(&wire.sketch_bytes) - .map_err(|e| format!("Failed to deserialize KLL sketch: {e}"))?, - ) - }; - - Ok(Self { k: wire.k, backend }) - } - - /// Merge from references without cloning. - pub fn merge_refs( - sketches: &[&Self], - ) -> Result> { - if sketches.is_empty() { - return Err("No sketches to merge".into()); - } - let k = sketches[0].k; - for s in sketches { - if s.k != k { - return Err("Cannot merge KllSketch with different k values".into()); - } - } - let mut merged = Self::new(k); - match &mut merged.backend { - KllBackend::Legacy(merged_legacy) => { - for s in sketches { - if let KllBackend::Legacy(s_legacy) = &s.backend { - merged_legacy.merge(s_legacy); - } else { - return Err("Cannot merge Legacy with Sketchlib KLL".into()); - } - } - } - KllBackend::Sketchlib(merged_sketchlib) => { - for s in sketches { - if let KllBackend::Sketchlib(s_sketchlib) = &s.backend { - sketchlib_kll_merge(merged_sketchlib, s_sketchlib); - } else { - return Err("Cannot merge Sketchlib with Legacy KLL".into()); - } - } - } - } - Ok(merged) - } - - /// Deserialize from a raw datasketches byte buffer (legacy Flink/FlinkSketch format). - pub fn from_dsrs_bytes(bytes: &[u8], k: u16) -> Result> { - let sketch = KllDoubleSketch::deserialize(bytes) - .map_err(|e| format!("Failed to deserialize KLL sketch from dsrs bytes: {e}"))?; - Ok(Self { - k, - backend: KllBackend::Legacy(sketch), - }) - } - - /// One-shot aggregation for the Arroyo UDAF call pattern. - pub fn aggregate_kll(k: u16, values: &[f64]) -> Option> { - if values.is_empty() { - return None; - } - let mut sketch = Self::new(k); - for &value in values { - sketch.update(value); - } - Some(sketch.serialize_msgpack()) - } -} - -// Manual trait implementations since the C++ and sketchlib types don't provide Clone -impl Clone for KllSketch { - fn clone(&self) -> Self { - let backend = match &self.backend { - KllBackend::Legacy(sketch) => { - let new_sketch = if sketch.get_n() == 0 { - KllDoubleSketch::with_k(self.k) - } else { - let bytes = sketch.serialize(); - KllDoubleSketch::deserialize(bytes.as_ref()).unwrap() - }; - KllBackend::Legacy(new_sketch) - } - KllBackend::Sketchlib(s) => { - let bytes = bytes_from_sketchlib_kll(s); - KllBackend::Sketchlib(sketchlib_kll_from_bytes(&bytes).unwrap()) - } - }; - Self { k: self.k, backend } - } -} - -impl std::fmt::Debug for KllSketch { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("KllSketch") - .field("k", &self.k) - .field("sketch_n", &self.count()) - .finish() - } -} - -// TODO: verify this -// Thread safety: The C++ library is not thread-safe by default, but since we're using it -// in a single-threaded context per accumulator instance and only sharing read-only operations, -// this should be safe. The actual sketch data is immutable once created. -unsafe impl Send for KllSketch {} -unsafe impl Sync for KllSketch {} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_kll_creation() { - let kll = KllSketch::new(200); - assert_eq!(kll.count(), 0); - assert_eq!(kll.k, 200); - } - - #[test] - fn test_kll_update() { - let mut kll = KllSketch::new(200); - kll.update(10.0); - kll.update(20.0); - kll.update(15.0); - assert_eq!(kll.count(), 3); - } - - #[test] - fn test_kll_quantile() { - let mut kll = KllSketch::new(200); - for i in 1..=10 { - kll.update(i as f64); - } - assert_eq!(kll.get_quantile(0.0), 1.0); - assert_eq!(kll.get_quantile(1.0), 10.0); - let median = kll.get_quantile(0.5); - assert!( - (5.0..=6.0).contains(&median), - "median should be between 5 and 6; got {median}" - ); - } - - #[test] - fn test_kll_merge() { - let mut kll1 = KllSketch::new(200); - let mut kll2 = KllSketch::new(200); - - for i in 1..=5 { - kll1.update(i as f64); - } - for i in 6..=10 { - kll2.update(i as f64); - } - - let merged = KllSketch::merge(vec![kll1, kll2]).unwrap(); - assert_eq!(merged.count(), 10); - assert_eq!(merged.get_quantile(0.0), 1.0); - assert_eq!(merged.get_quantile(1.0), 10.0); - } - - #[test] - fn test_msgpack_round_trip() { - let mut kll = KllSketch::new(200); - for i in 1..=5 { - kll.update(i as f64); - } - - let bytes = kll.serialize_msgpack(); - let deserialized = KllSketch::deserialize_msgpack(&bytes).unwrap(); - - assert_eq!(deserialized.k, 200); - assert_eq!(deserialized.count(), 5); - assert_eq!(deserialized.get_quantile(0.0), 1.0); - assert_eq!(deserialized.get_quantile(1.0), 5.0); - } - - #[test] - fn test_aggregate_kll() { - let values = [1.0, 2.0, 3.0, 4.0, 5.0]; - let bytes = KllSketch::aggregate_kll(200, &values).unwrap(); - let kll = KllSketch::deserialize_msgpack(&bytes).unwrap(); - assert_eq!(kll.count(), 5); - assert_eq!(kll.get_quantile(0.0), 1.0); - assert_eq!(kll.get_quantile(1.0), 5.0); - } - - #[test] - fn test_aggregate_kll_empty() { - assert!(KllSketch::aggregate_kll(200, &[]).is_none()); - } -} diff --git a/asap-common/sketch-core/src/kll_sketchlib.rs b/asap-common/sketch-core/src/kll_sketchlib.rs deleted file mode 100644 index d49d80b9..00000000 --- a/asap-common/sketch-core/src/kll_sketchlib.rs +++ /dev/null @@ -1,36 +0,0 @@ -use asap_sketchlib::{SketchInput, KLL}; - -/// Concrete KLL type from asap_sketchlib when sketchlib backend is enabled. -pub type SketchlibKll = KLL; - -/// Creates a fresh sketchlib KLL sketch with the requested accuracy parameter `k`. -pub fn new_sketchlib_kll(k: u16) -> SketchlibKll { - KLL::init_kll(k as i32) -} - -/// Updates a sketchlib KLL with one numeric observation. -pub fn sketchlib_kll_update(inner: &mut SketchlibKll, value: f64) { - // KLL accepts only numeric inputs. We intentionally ignore the error here because `value` - // is always numeric. - let _ = inner.update(&SketchInput::F64(value)); -} - -/// Queries a sketchlib KLL for the value at the requested quantile. -pub fn sketchlib_kll_quantile(inner: &SketchlibKll, q: f64) -> f64 { - inner.quantile(q) -} - -/// Merges `src` into `dst`. -pub fn sketchlib_kll_merge(dst: &mut SketchlibKll, src: &SketchlibKll) { - dst.merge(src); -} - -/// Serializes a sketchlib KLL into MessagePack bytes. -pub fn bytes_from_sketchlib_kll(inner: &SketchlibKll) -> Vec { - inner.serialize_to_bytes().unwrap() -} - -/// Deserializes a sketchlib KLL from MessagePack bytes. -pub fn sketchlib_kll_from_bytes(bytes: &[u8]) -> Result> { - Ok(KLL::deserialize_from_bytes(bytes)?) -} diff --git a/asap-common/sketch-core/src/lib.rs b/asap-common/sketch-core/src/lib.rs deleted file mode 100644 index 3cfda08d..00000000 --- a/asap-common/sketch-core/src/lib.rs +++ /dev/null @@ -1,19 +0,0 @@ -#[cfg(test)] -#[ctor::ctor] -fn init_sketch_legacy_for_tests() { - crate::config::force_legacy_mode_for_tests(); -} - -pub mod config; -pub mod count_min; -pub mod count_min_sketchlib; -pub mod count_min_with_heap; -pub mod count_min_with_heap_sketchlib; -pub mod count_sketch; -pub mod dd_sketch; -pub mod delta_set_aggregator; -pub mod hll_sketch; -pub mod hydra_kll; -pub mod kll; -pub mod kll_sketchlib; -pub mod set_aggregator; diff --git a/asap-common/sketch-core/src/set_aggregator.rs b/asap-common/sketch-core/src/set_aggregator.rs deleted file mode 100644 index c745f284..00000000 --- a/asap-common/sketch-core/src/set_aggregator.rs +++ /dev/null @@ -1,152 +0,0 @@ -// Adapted from QueryEngineRust/src/precompute_operators/set_aggregator_accumulator.rs -// Changes: -// - Renamed SetAggregatorAccumulator -> SetAggregator -// - values field is now HashSet instead of HashSet -// - add_key(&str) instead of add_key(KeyByLabelValues) -// - serialize_msgpack / deserialize_msgpack use StringSet { values: HashSet } -// wire format matching the Arroyo setaggregator_ UDF exactly (same as DeltaResult pattern) -// - merge_accumulators -> merge -// - Removed: deserialize_from_json, deserialize_from_bytes, deserialize_from_bytes_arroyo -// (QE-specific / buggy legacy formats stay in QE) -// - Removed: AggregateCore, SerializableToSink, MergeableAccumulator, MultipleSubpopulationAggregate impls -// - Removed: with_added (QE-specific constructor) - -use serde::{Deserialize, Serialize}; -use std::collections::HashSet; - -/// Set aggregator for tracking a set of unique string keys. -/// Wire format: StringSet { values: HashSet } in MessagePack — matches Arroyo setaggregator_ UDF. -#[derive(Debug, Clone)] -pub struct SetAggregator { - pub values: HashSet, -} - -impl SetAggregator { - pub fn new() -> Self { - Self { - values: HashSet::new(), - } - } - - pub fn insert(&mut self, key: &str) { - self.values.insert(key.to_string()); - } - - pub fn merge( - accumulators: Vec, - ) -> Result> { - if accumulators.is_empty() { - return Err("No accumulators to merge".into()); - } - - let mut merged = SetAggregator::new(); - for accumulator in accumulators { - merged.values.extend(accumulator.values); - } - - Ok(merged) - } - - /// Serialize to MessagePack — matches the Arroyo setaggregator_ UDF wire format exactly: - /// StringSet { values: HashSet } as a msgpack map. - pub fn serialize_msgpack(&self) -> Vec { - #[derive(Serialize)] - struct StringSet<'a> { - values: &'a HashSet, - } - let wrapper = StringSet { - values: &self.values, - }; - let mut buf = Vec::new(); - rmp_serde::encode::write(&mut buf, &wrapper).unwrap(); - buf - } - - /// Deserialize from MessagePack produced by the Arroyo setaggregator_ UDF. - pub fn deserialize_msgpack(buffer: &[u8]) -> Result> { - #[derive(Deserialize)] - struct StringSet { - values: HashSet, - } - let wrapper: StringSet = rmp_serde::from_slice(buffer) - .map_err(|e| format!("Failed to deserialize SetAggregator from MessagePack: {e}"))?; - Ok(Self { - values: wrapper.values, - }) - } -} - -impl Default for SetAggregator { - fn default() -> Self { - Self::new() - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_creation() { - let sa = SetAggregator::new(); - assert!(sa.values.is_empty()); - } - - #[test] - fn test_insert() { - let mut sa = SetAggregator::new(); - sa.insert("web"); - sa.insert("api"); - sa.insert("web"); // duplicate - assert_eq!(sa.values.len(), 2); - assert!(sa.values.contains("web")); - assert!(sa.values.contains("api")); - } - - #[test] - fn test_merge() { - let mut sa1 = SetAggregator::new(); - let mut sa2 = SetAggregator::new(); - - sa1.insert("web"); - sa1.insert("api"); - sa2.insert("api"); // duplicate - sa2.insert("db"); - - let merged = SetAggregator::merge(vec![sa1, sa2]).unwrap(); - assert_eq!(merged.values.len(), 3); - assert!(merged.values.contains("web")); - assert!(merged.values.contains("api")); - assert!(merged.values.contains("db")); - } - - #[test] - fn test_msgpack_round_trip() { - let mut sa = SetAggregator::new(); - sa.insert("web"); - sa.insert("api"); - - let bytes = sa.serialize_msgpack(); - let deserialized = SetAggregator::deserialize_msgpack(&bytes).unwrap(); - - assert_eq!(deserialized.values.len(), 2); - assert!(deserialized.values.contains("web")); - assert!(deserialized.values.contains("api")); - } - - #[test] - fn test_msgpack_matches_arroyo_format() { - // Verify wire format is StringSet { values: [...] } not a plain array. - // Arroyo's setaggregator_.rs serializes StringSet { values: HashSet }. - #[derive(Deserialize)] - struct StringSet { - values: HashSet, - } - let mut sa = SetAggregator::new(); - sa.insert("a"); - let bytes = sa.serialize_msgpack(); - let decoded: StringSet = - rmp_serde::from_slice(&bytes).expect("should decode as StringSet { values: ... }"); - assert!(decoded.values.contains("a")); - } -} diff --git a/asap-query-engine/Cargo.toml b/asap-query-engine/Cargo.toml index 2261aa4f..0bf130f9 100644 --- a/asap-query-engine/Cargo.toml +++ b/asap-query-engine/Cargo.toml @@ -5,7 +5,6 @@ edition.workspace = true [dependencies] # Internal crates (workspace) -sketch-core.workspace = true promql_utilities.workspace = true sql_utilities.workspace = true asap_types.workspace = true @@ -63,7 +62,7 @@ zstd = "0.13" reqwest = { version = "0.11", default-features = false, features = ["json", "rustls-tls"] } tracing-appender = "0.2" elastic_dsl_utilities.workspace = true -asap_sketchlib = { git = "https://github.com/ProjectASAP/asap_sketchlib" } +asap_sketchlib = { git = "https://github.com/ProjectASAP/asap_sketchlib", branch = "refactor/adopt-sketch-core-modules", features = ["asap-cli"] } # Persistence layer (SimpleMapStore parts / manifest / Tier-2 cache) moka = { version = "0.12", features = ["sync"] } memmap2 = "0.9" diff --git a/asap-query-engine/Dockerfile b/asap-query-engine/Dockerfile index f54cdce6..2401dc4e 100644 --- a/asap-query-engine/Dockerfile +++ b/asap-query-engine/Dockerfile @@ -16,9 +16,6 @@ RUN apt-get update && apt-get install -y --no-install-recommends \ # Copy the asap-common directory COPY asap-common ./asap-common -# Copy path dependencies of asap-query-engine -COPY asap-common/sketch-core ./asap-common/sketch-core - COPY Cargo.toml ./ COPY Cargo.lock ./ COPY asap-query-engine/Cargo.toml ./asap-query-engine/ diff --git a/asap-query-engine/src/drivers/ingest/otel.rs b/asap-query-engine/src/drivers/ingest/otel.rs index 72f7ba72..c27b5913 100644 --- a/asap-query-engine/src/drivers/ingest/otel.rs +++ b/asap-query-engine/src/drivers/ingest/otel.rs @@ -669,8 +669,7 @@ async fn route_modified_otlp_sketches_to_precompute( // one in-flight delta per (metric, labels) so // the next full snapshot replaces the current // cache entry cleanly. - let accumulator: Box = if dp.encoding - == ENCODING_PROTO_DELTA + let accumulator: Box = if dp.encoding == ENCODING_PROTO_DELTA || dp.encoding == ENCODING_MSGPACK_DELTA { let Some(base) = ingest_state @@ -714,11 +713,7 @@ async fn route_modified_otlp_sketches_to_precompute( .insert(series_key.clone(), merged.clone_boxed_core()); merged } else { - match decode_modified_otlp_sketch_bytes( - dp.kind, - dp.encoding, - &dp.sketch, - ) { + match decode_modified_otlp_sketch_bytes(dp.kind, dp.encoding, &dp.sketch) { Ok(acc) => { ingest_state .sketch_snapshots @@ -978,19 +973,16 @@ pub(crate) fn apply_modified_otlp_delta_bytes( DDSketch / HLL / CountSketch / CountMin are wired" ) .into()), - (ENCODING_MSGPACK_DELTA, _) => Err( - "MSGPACK_DELTA encoding is not yet wired; PR G covers PROTO_DELTA only" - .into(), - ), + (ENCODING_MSGPACK_DELTA, _) => { + Err("MSGPACK_DELTA encoding is not yet wired; PR G covers PROTO_DELTA only".into()) + } (ENCODING_PROTO, _) | (ENCODING_MSGPACK, _) => Err(format!( "encoding {encoding} is a full-state frame — route through \ `decode_modified_otlp_sketch_bytes` and replace the cached \ accumulator, not through the delta applier" ) .into()), - (other, _) => Err( - format!("unknown modified-OTLP sketch encoding {other}").into(), - ), + (other, _) => Err(format!("unknown modified-OTLP sketch encoding {other}").into()), } } @@ -1337,14 +1329,12 @@ mod dispatcher_tests { use super::*; use crate::data_model::AggregateCore; use crate::precompute_operators::{DDSketchAccumulator, HllSketchAccumulator}; - use sketch_core::dd_sketch::DdSketch; - use sketch_core::hll_sketch::HllVariant; + use asap_sketchlib::asap::dd_sketch::DdSketch; + use asap_sketchlib::asap::hll_sketch::HllVariant; #[test] fn apply_modified_otlp_delta_bytes_ddsketch_round_trip() { - use asap_otel_proto::sketchlib::v1::{ - DdSketchBucketDelta, DdSketchDelta as PbDelta, - }; + use asap_otel_proto::sketchlib::v1::{DdSketchBucketDelta, DdSketchDelta as PbDelta}; use prost::Message; // Base sketch represents the last full snapshot the agent sent. @@ -1354,8 +1344,14 @@ mod dispatcher_tests { let bytes = PbDelta { buckets: vec![ - DdSketchBucketDelta { index: 0, d_count: 10 }, - DdSketchBucketDelta { index: 2, d_count: 20 }, + DdSketchBucketDelta { + index: 0, + d_count: 10, + }, + DdSketchBucketDelta { + index: 2, + d_count: 20, + }, ], d_count: 30, d_sum: 70.0, @@ -1374,10 +1370,7 @@ mod dispatcher_tests { ) .expect("apply ok"); - let dd = acc - .as_any() - .downcast_ref::() - .unwrap(); + let dd = acc.as_any().downcast_ref::().unwrap(); assert_eq!(dd.inner.store_counts, vec![11, 2, 23]); assert_eq!(dd.inner.count, 36); assert_eq!(dd.inner.min, 0.5); @@ -1405,18 +1398,10 @@ mod dispatcher_tests { } .encode_to_vec(); - apply_modified_otlp_delta_bytes( - SketchKind::Hll, - ENCODING_PROTO_DELTA, - &mut acc, - &bytes, - ) - .expect("apply ok"); + apply_modified_otlp_delta_bytes(SketchKind::Hll, ENCODING_PROTO_DELTA, &mut acc, &bytes) + .expect("apply ok"); - let hll = acc - .as_any() - .downcast_ref::() - .unwrap(); + let hll = acc.as_any().downcast_ref::().unwrap(); assert_eq!(hll.inner.registers, vec![4, 5, 6, 7]); } @@ -1437,16 +1422,11 @@ mod dispatcher_tests { #[test] fn apply_rejects_full_state_encoding() { - let mut acc: Box = - Box::new(DDSketchAccumulator::new(0.01)); - let err = apply_modified_otlp_delta_bytes( - SketchKind::DdSketch, - ENCODING_PROTO, - &mut acc, - &[], - ) - .expect_err("expected full-state-rejection error") - .to_string(); + let mut acc: Box = Box::new(DDSketchAccumulator::new(0.01)); + let err = + apply_modified_otlp_delta_bytes(SketchKind::DdSketch, ENCODING_PROTO, &mut acc, &[]) + .expect_err("expected full-state-rejection error") + .to_string(); assert!(err.contains("full-state frame")); } diff --git a/asap-query-engine/src/lib.rs b/asap-query-engine/src/lib.rs index 37996c55..367b227d 100644 --- a/asap-query-engine/src/lib.rs +++ b/asap-query-engine/src/lib.rs @@ -2,13 +2,13 @@ #[ctor::ctor] fn init_sketch_backend_for_tests() { #[cfg(feature = "sketchlib-tests")] - let _ = sketch_core::config::configure( - sketch_core::config::ImplMode::Sketchlib, - sketch_core::config::ImplMode::Legacy, - sketch_core::config::ImplMode::Sketchlib, + let _ = asap_sketchlib::asap::config::configure( + asap_sketchlib::asap::config::ImplMode::Sketchlib, + asap_sketchlib::asap::config::ImplMode::Legacy, + asap_sketchlib::asap::config::ImplMode::Sketchlib, ); #[cfg(not(feature = "sketchlib-tests"))] - sketch_core::config::force_legacy_mode_for_tests(); + asap_sketchlib::asap::config::force_legacy_mode_for_tests(); } pub mod data_model; diff --git a/asap-query-engine/src/main.rs b/asap-query-engine/src/main.rs index 391a87e9..9e106f22 100644 --- a/asap-query-engine/src/main.rs +++ b/asap-query-engine/src/main.rs @@ -5,7 +5,7 @@ use std::sync::Arc; use tokio::signal; use tracing::{error, info, warn}; -use sketch_core::config::{self, ImplMode}; +use asap_sketchlib::asap::config::{self, ImplMode}; use query_engine_rust::data_model::enums::{InputFormat, LockStrategy, StreamingEngine}; use query_engine_rust::drivers::AdapterConfig; @@ -941,11 +941,8 @@ mod tests { #[test] fn no_cold_no_forward_yields_no_fallback() { - let cfg = AdapterConfig::from_prom_with_optional_cold( - "http://prom:9090".into(), - false, - None, - ); + let cfg = + AdapterConfig::from_prom_with_optional_cold("http://prom:9090".into(), false, None); assert!( cfg.fallback.is_none(), "without cold-store and without forward, no fallback should be installed", @@ -954,11 +951,8 @@ mod tests { #[test] fn no_cold_with_forward_yields_prom_fallback() { - let cfg = AdapterConfig::from_prom_with_optional_cold( - "http://prom:9090".into(), - true, - None, - ); + let cfg = + AdapterConfig::from_prom_with_optional_cold("http://prom:9090".into(), true, None); assert!( cfg.fallback.is_some(), "forward_unsupported=true must install Prom fallback", diff --git a/asap-query-engine/src/precompute_engine/ingest_handler.rs b/asap-query-engine/src/precompute_engine/ingest_handler.rs index 2ba5b576..a9ec8dff 100644 --- a/asap-query-engine/src/precompute_engine/ingest_handler.rs +++ b/asap-query-engine/src/precompute_engine/ingest_handler.rs @@ -551,15 +551,11 @@ mod tests { /// hand-merged sequence of full sketches. #[tokio::test] async fn delta_path_reconstitutes_cumulative_state() { - use crate::drivers::ingest::otel::{ - apply_modified_otlp_delta_bytes, SketchKind, - }; + use crate::drivers::ingest::otel::{apply_modified_otlp_delta_bytes, SketchKind}; use crate::precompute_operators::DDSketchAccumulator; - use asap_otel_proto::sketchlib::v1::{ - DdSketchBucketDelta, DdSketchDelta as PbDelta, - }; + use asap_otel_proto::sketchlib::v1::{DdSketchBucketDelta, DdSketchDelta as PbDelta}; + use asap_sketchlib::asap::dd_sketch::DdSketch; use prost::Message; - use sketch_core::dd_sketch::DdSketch; const ENCODING_PROTO_DELTA: i32 = 2; @@ -579,8 +575,14 @@ mod tests { // First delta adds to bucket 0 and bucket 2. let d1 = PbDelta { buckets: vec![ - DdSketchBucketDelta { index: 0, d_count: 10 }, - DdSketchBucketDelta { index: 2, d_count: 20 }, + DdSketchBucketDelta { + index: 0, + d_count: 10, + }, + DdSketchBucketDelta { + index: 2, + d_count: 20, + }, ], d_count: 30, d_sum: 70.0, @@ -594,13 +596,8 @@ mod tests { .get(series_key) .unwrap() .clone_boxed_core(); - apply_modified_otlp_delta_bytes( - SketchKind::DdSketch, - ENCODING_PROTO_DELTA, - &mut acc1, - &d1, - ) - .expect("apply first delta"); + apply_modified_otlp_delta_bytes(SketchKind::DdSketch, ENCODING_PROTO_DELTA, &mut acc1, &d1) + .expect("apply first delta"); state .sketch_snapshots .insert(series_key.to_string(), acc1.clone_boxed_core()); @@ -608,7 +605,10 @@ mod tests { // Second delta — picks up on top of the first, proving the // cache refresh is transitive. let d2 = PbDelta { - buckets: vec![DdSketchBucketDelta { index: 1, d_count: 5 }], + buckets: vec![DdSketchBucketDelta { + index: 1, + d_count: 5, + }], d_count: 5, d_sum: 10.0, new_max: 6.0, @@ -621,18 +621,10 @@ mod tests { .get(series_key) .unwrap() .clone_boxed_core(); - apply_modified_otlp_delta_bytes( - SketchKind::DdSketch, - ENCODING_PROTO_DELTA, - &mut acc2, - &d2, - ) - .expect("apply second delta"); + apply_modified_otlp_delta_bytes(SketchKind::DdSketch, ENCODING_PROTO_DELTA, &mut acc2, &d2) + .expect("apply second delta"); - let final_dd = acc2 - .as_any() - .downcast_ref::() - .unwrap(); + let final_dd = acc2.as_any().downcast_ref::().unwrap(); // Base [1,2,3] + d1 [+10 on 0, +20 on 2] = [11,2,23]; // + d2 [+5 on 1] = [11,7,23]. assert_eq!(final_dd.inner.store_counts, vec![11, 7, 23]); diff --git a/asap-query-engine/src/precompute_engine/worker.rs b/asap-query-engine/src/precompute_engine/worker.rs index 93a4b5d7..df2e8520 100644 --- a/asap-query-engine/src/precompute_engine/worker.rs +++ b/asap-query-engine/src/precompute_engine/worker.rs @@ -964,8 +964,8 @@ mod tests { use crate::precompute_operators::datasketches_kll_accumulator::DatasketchesKLLAccumulator; use crate::precompute_operators::multiple_sum_accumulator::MultipleSumAccumulator; use crate::precompute_operators::sum_accumulator::SumAccumulator; + use asap_sketchlib::asap::kll::KllSketch; use asap_types::enums::{AggregationType, WindowType}; - use sketch_core::kll::KllSketch; fn make_agg_config( id: u64, diff --git a/asap-query-engine/src/precompute_operators/count_min_sketch_accumulator.rs b/asap-query-engine/src/precompute_operators/count_min_sketch_accumulator.rs index 902b8d8b..b311dbc1 100644 --- a/asap-query-engine/src/precompute_operators/count_min_sketch_accumulator.rs +++ b/asap-query-engine/src/precompute_operators/count_min_sketch_accumulator.rs @@ -2,13 +2,13 @@ use crate::data_model::{ AggregateCore, AggregationType, KeyByLabelValues, MergeableAccumulator, MultipleSubpopulationAggregate, SerializableToSink, }; +use asap_sketchlib::asap::count_min::{CountMinDelta, CountMinSketch}; use serde_json::Value; -use sketch_core::count_min::{CountMinDelta, CountMinSketch}; use std::collections::HashMap; use promql_utilities::query_logics::enums::Statistic; -/// Count-Min Sketch accumulator — wraps sketch_core::CountMinSketch. +/// Count-Min Sketch accumulator — wraps asap_sketchlib::asap::CountMinSketch. /// Core struct, update/merge/serde logic live in sketch-core. /// This file retains QE-specific trait impls, legacy deserializers, and JSON output. #[derive(Debug, Clone)] @@ -186,12 +186,9 @@ impl CountMinSketchAccumulator { use asap_otel_proto::sketchlib::v1::CountMinDelta as PbDelta; use prost::Message; - let pb = PbDelta::decode(buffer) - .map_err(|e| format!("decode CountMinDelta: {e}"))?; + let pb = PbDelta::decode(buffer).map_err(|e| format!("decode CountMinDelta: {e}"))?; - if pb.cell_rows.len() != pb.cell_cols.len() - || pb.cell_rows.len() != pb.d_counts.len() - { + if pb.cell_rows.len() != pb.cell_cols.len() || pb.cell_rows.len() != pb.d_counts.len() { return Err(format!( "CountMinDelta packed-array length mismatch: \ cell_rows={}, cell_cols={}, d_counts={}", @@ -335,11 +332,8 @@ impl AggregateCore for CountMinSketchAccumulator { self } - fn as_any_mut(&mut self) -> &mut dyn std::any::Any { - self - } fn merge_with( @@ -412,7 +406,11 @@ impl AggregateCore for CountMinSketchAccumulator { } let row_totals = matrix.iter().map(|r| r.iter().sum::()); let min_total = row_totals.fold(f64::INFINITY, f64::min); - Ok(if min_total.is_finite() { min_total } else { 0.0 }) + Ok(if min_total.is_finite() { + min_total + } else { + 0.0 + }) } other => Err(format!( "CountMinSketchAccumulator: statistic {:?} not supported \ diff --git a/asap-query-engine/src/precompute_operators/count_min_sketch_with_heap_accumulator.rs b/asap-query-engine/src/precompute_operators/count_min_sketch_with_heap_accumulator.rs index 18076553..53c61917 100644 --- a/asap-query-engine/src/precompute_operators/count_min_sketch_with_heap_accumulator.rs +++ b/asap-query-engine/src/precompute_operators/count_min_sketch_with_heap_accumulator.rs @@ -2,13 +2,13 @@ use crate::data_model::{ AggregateCore, AggregationType, KeyByLabelValues, MergeableAccumulator, MultipleSubpopulationAggregate, SerializableToSink, }; +use asap_sketchlib::asap::count_min_with_heap::{CountMinSketchWithHeap, HeapItem}; use serde_json::Value; -use sketch_core::count_min_with_heap::{CountMinSketchWithHeap, HeapItem}; use std::collections::HashMap; use promql_utilities::query_logics::enums::Statistic; -/// Count-Min Sketch with Heap accumulator — wraps sketch_core::CountMinSketchWithHeap. +/// Count-Min Sketch with Heap accumulator — wraps asap_sketchlib::asap::CountMinSketchWithHeap. /// Core struct, update/merge/serde logic live in sketch-core. /// This file retains QE-specific trait impls, legacy deserializers, and JSON output. /// @@ -20,7 +20,7 @@ pub struct CountMinSketchWithHeapAccumulator { } // Re-export HeapItem so existing code using CountMinSketchWithHeapAccumulator::HeapItem still works. -pub use sketch_core::count_min_with_heap::HeapItem as HeapItemReexport; +pub use asap_sketchlib::asap::count_min_with_heap::HeapItem as HeapItemReexport; impl CountMinSketchWithHeapAccumulator { pub fn new(row_num: usize, col_num: usize, heap_size: usize) -> Self { @@ -150,11 +150,8 @@ impl AggregateCore for CountMinSketchWithHeapAccumulator { self } - fn as_any_mut(&mut self) -> &mut dyn std::any::Any { - self - } fn merge_with( diff --git a/asap-query-engine/src/precompute_operators/count_sketch_accumulator.rs b/asap-query-engine/src/precompute_operators/count_sketch_accumulator.rs index 8c7d35e2..76f09413 100644 --- a/asap-query-engine/src/precompute_operators/count_sketch_accumulator.rs +++ b/asap-query-engine/src/precompute_operators/count_sketch_accumulator.rs @@ -1,4 +1,4 @@ -//! Count Sketch accumulator — wraps `sketch_core::count_sketch::CountSketch`. +//! Count Sketch accumulator — wraps `asap_sketchlib::asap::count_sketch::CountSketch`. //! //! This is the concrete accumulator reached from the modified-OTLP //! `Metric.data = CountSketch{…}` hot path (PR C-CountSketch). Its @@ -20,8 +20,8 @@ //! round-trip works end-to-end without that richer query surface. use crate::data_model::{AggregateCore, AggregationType, KeyByLabelValues, SerializableToSink}; +use asap_sketchlib::asap::count_sketch::{CountSketch, CountSketchDelta}; use serde_json::Value; -use sketch_core::count_sketch::{CountSketch, CountSketchDelta}; use std::collections::HashMap; /// Count Sketch accumulator — inner matrix of signed counts. @@ -159,12 +159,9 @@ impl CountSketchAccumulator { use asap_otel_proto::sketchlib::v1::CountSketchDelta as PbDelta; use prost::Message; - let pb = PbDelta::decode(buffer) - .map_err(|e| format!("decode CountSketchDelta: {e}"))?; + let pb = PbDelta::decode(buffer).map_err(|e| format!("decode CountSketchDelta: {e}"))?; - if pb.cell_rows.len() != pb.cell_cols.len() - || pb.cell_rows.len() != pb.d_counts.len() - { + if pb.cell_rows.len() != pb.cell_cols.len() || pb.cell_rows.len() != pb.d_counts.len() { return Err(format!( "CountSketchDelta packed-array length mismatch: \ cell_rows={}, cell_cols={}, d_counts={}", @@ -221,11 +218,8 @@ impl AggregateCore for CountSketchAccumulator { self } - fn as_any_mut(&mut self) -> &mut dyn std::any::Any { - self - } fn merge_with( @@ -272,7 +266,7 @@ impl AggregateCore for CountSketchAccumulator { // but not a heavy-hitter answer). Hash compatibility note: // this relies on the agent and backend using the // sketchlib HashSpec; sketchlib-go's `portableHashSpec` - // is the canonical seed list, and `sketch_core::CountSketch` + // is the canonical seed list, and `asap_sketchlib::asap::CountSketch` // hashes against the same spec. match statistic { Statistic::Topk | Statistic::Count => { @@ -308,7 +302,7 @@ impl AggregateCore for CountSketchAccumulator { /// /// Hash compatibility with the agent is via the sketchlib hash /// spec; the agent's `sketchlib-go::CountSketch` and the -/// backend's `sketch_core::count_sketch::CountSketch` must use +/// backend's `asap_sketchlib::asap::count_sketch::CountSketch` must use /// the same seed list (sketchlib's `portableHashSpec` / /// `default_hash_spec`). fn count_sketch_query_key(matrix: &Vec>, key: &str) -> f64 { diff --git a/asap-query-engine/src/precompute_operators/datasketches_kll_accumulator.rs b/asap-query-engine/src/precompute_operators/datasketches_kll_accumulator.rs index bbc1bfd3..fb7ffbb9 100644 --- a/asap-query-engine/src/precompute_operators/datasketches_kll_accumulator.rs +++ b/asap-query-engine/src/precompute_operators/datasketches_kll_accumulator.rs @@ -2,9 +2,9 @@ use crate::data_model::{ AggregateCore, AggregationType, AuxStats, MergeableAccumulator, SerializableToSink, SingleSubpopulationAggregate, }; +use asap_sketchlib::asap::kll::KllSketch; use base64::{engine::general_purpose, Engine as _}; use serde_json::Value; -use sketch_core::kll::KllSketch; use std::collections::HashMap; #[cfg(feature = "extra_debugging")] use std::time::Instant; @@ -12,7 +12,7 @@ use tracing::debug; use promql_utilities::query_logics::enums::Statistic; -/// KLL sketch accumulator — wraps sketch_core::KllSketch. +/// KLL sketch accumulator — wraps asap_sketchlib::asap::KllSketch. /// Core struct, update/merge/serde logic live in sketch-core. /// This file retains QE-specific trait impls, legacy deserializers, and JSON output. pub struct DatasketchesKLLAccumulator { @@ -258,11 +258,8 @@ impl AggregateCore for DatasketchesKLLAccumulator { self } - fn as_any_mut(&mut self) -> &mut dyn std::any::Any { - self - } fn merge_with( diff --git a/asap-query-engine/src/precompute_operators/dd_sketch_accumulator.rs b/asap-query-engine/src/precompute_operators/dd_sketch_accumulator.rs index 239b3516..c0cee54b 100644 --- a/asap-query-engine/src/precompute_operators/dd_sketch_accumulator.rs +++ b/asap-query-engine/src/precompute_operators/dd_sketch_accumulator.rs @@ -1,4 +1,4 @@ -//! DDSketch accumulator — wraps `sketch_core::dd_sketch::DdSketch`. +//! DDSketch accumulator — wraps `asap_sketchlib::asap::dd_sketch::DdSketch`. //! //! Concrete accumulator reached from the modified-OTLP //! `Metric.data = DDSketch{…}` hot path (PR C-CountSketch follow-up). @@ -12,8 +12,8 @@ //! works end-to-end without that richer query surface. use crate::data_model::{AggregateCore, AggregationType, KeyByLabelValues, SerializableToSink}; +use asap_sketchlib::asap::dd_sketch::{DdSketch, DdSketchDelta}; use serde_json::Value; -use sketch_core::dd_sketch::{DdSketch, DdSketchDelta}; use std::collections::HashMap; /// DDSketch accumulator — inner log-bucketed sketch. @@ -107,8 +107,7 @@ impl DDSketchAccumulator { use asap_otel_proto::sketchlib::v1::DdSketchDelta as PbDelta; use prost::Message; - let pb = PbDelta::decode(buffer) - .map_err(|e| format!("decode DDSketchDelta: {e}"))?; + let pb = PbDelta::decode(buffer).map_err(|e| format!("decode DDSketchDelta: {e}"))?; let buckets = pb .buckets @@ -160,11 +159,8 @@ impl AggregateCore for DDSketchAccumulator { self } - fn as_any_mut(&mut self) -> &mut dyn std::any::Any { - self - } fn merge_with( @@ -375,9 +371,7 @@ mod tests { #[test] fn test_apply_proto_delta_bytes_round_trip() { - use asap_otel_proto::sketchlib::v1::{ - DdSketchBucketDelta, DdSketchDelta as PbDelta, - }; + use asap_otel_proto::sketchlib::v1::{DdSketchBucketDelta, DdSketchDelta as PbDelta}; use prost::Message; let mut acc = DDSketchAccumulator::new(0.01); diff --git a/asap-query-engine/src/precompute_operators/delta_set_aggregator_accumulator.rs b/asap-query-engine/src/precompute_operators/delta_set_aggregator_accumulator.rs index 9b00ad90..3825bfcb 100644 --- a/asap-query-engine/src/precompute_operators/delta_set_aggregator_accumulator.rs +++ b/asap-query-engine/src/precompute_operators/delta_set_aggregator_accumulator.rs @@ -2,8 +2,8 @@ use crate::data_model::{ AggregateCore, AggregationType, KeyByLabelValues, MergeableAccumulator, MultipleSubpopulationAggregate, SerializableToSink, }; +use asap_sketchlib::asap::delta_set_aggregator::{deserialize_msgpack, serialize_msgpack}; use serde_json::Value; -use sketch_core::delta_set_aggregator::{deserialize_msgpack, serialize_msgpack}; use std::collections::{HashMap, HashSet}; use promql_utilities::query_logics::enums::Statistic; @@ -219,11 +219,8 @@ impl AggregateCore for DeltaSetAggregatorAccumulator { self } - fn as_any_mut(&mut self) -> &mut dyn std::any::Any { - self - } fn merge_with( diff --git a/asap-query-engine/src/precompute_operators/hll_sketch_accumulator.rs b/asap-query-engine/src/precompute_operators/hll_sketch_accumulator.rs index 9c5bf651..0e46ec77 100644 --- a/asap-query-engine/src/precompute_operators/hll_sketch_accumulator.rs +++ b/asap-query-engine/src/precompute_operators/hll_sketch_accumulator.rs @@ -1,4 +1,4 @@ -//! HLL accumulator — wraps `sketch_core::hll_sketch::HllSketch`. +//! HLL accumulator — wraps `asap_sketchlib::asap::hll_sketch::HllSketch`. //! //! Concrete accumulator reached from the modified-OTLP //! `Metric.data = HLLSketch{…}` hot path (PR C-CountSketch follow-up). @@ -12,8 +12,8 @@ //! store round-trip works end-to-end without that richer query surface. use crate::data_model::{AggregateCore, AggregationType, KeyByLabelValues, SerializableToSink}; +use asap_sketchlib::asap::hll_sketch::{HllDelta, HllSketch, HllVariant}; use serde_json::Value; -use sketch_core::hll_sketch::{HllDelta, HllSketch, HllVariant}; use std::collections::HashMap; /// HLL accumulator — inner register array + variant metadata. @@ -95,7 +95,7 @@ impl HllSketchAccumulator { let variant = match proto_variant { ProtoVariant::Unspecified => HllVariant::Unspecified, ProtoVariant::Regular => HllVariant::Regular, - ProtoVariant::Datafusion => HllVariant::Datafusion, + ProtoVariant::ErtlMle => HllVariant::Datafusion, ProtoVariant::Hip => HllVariant::Hip, }; let inner = HllSketch::from_raw( @@ -124,8 +124,7 @@ impl HllSketchAccumulator { use asap_otel_proto::sketchlib::v1::HllDelta as PbDelta; use prost::Message; - let pb = PbDelta::decode(buffer) - .map_err(|e| format!("decode HLLDelta: {e}"))?; + let pb = PbDelta::decode(buffer).map_err(|e| format!("decode HLLDelta: {e}"))?; let updates = pb .updates @@ -170,11 +169,8 @@ impl AggregateCore for HllSketchAccumulator { self } - fn as_any_mut(&mut self) -> &mut dyn std::any::Any { - self - } fn merge_with( @@ -238,8 +234,8 @@ impl AggregateCore for HllSketchAccumulator { /// (linear-counting) and large-range (32-bit space) corrections /// from the original Flajolet et al. paper. /// -/// Inlined here rather than added as a method on `sketch_core::HllSketch` -/// because the existing `sketch_core` types only expose merge / +/// Inlined here rather than added as a method on `asap_sketchlib::asap::HllSketch` +/// because the existing `asap_sketchlib::asap` types only expose merge / /// serialize today; adding a query method there would force a /// cross-crate change. fn hll_cardinality_estimate(registers: &[u8]) -> f64 { diff --git a/asap-query-engine/src/precompute_operators/hydra_kll_accumulator.rs b/asap-query-engine/src/precompute_operators/hydra_kll_accumulator.rs index c8948a8a..f9ed22bc 100644 --- a/asap-query-engine/src/precompute_operators/hydra_kll_accumulator.rs +++ b/asap-query-engine/src/precompute_operators/hydra_kll_accumulator.rs @@ -5,13 +5,13 @@ use crate::{ }, KeyByLabelValues, }; +use asap_sketchlib::asap::hydra_kll::HydraKllSketch; use base64::{engine::general_purpose, Engine as _}; -use sketch_core::hydra_kll::HydraKllSketch; use std::collections::HashMap; use promql_utilities::query_logics::enums::Statistic; -/// HydraKLL sketch accumulator — wraps sketch_core::HydraKllSketch. +/// HydraKLL sketch accumulator — wraps asap_sketchlib::asap::HydraKllSketch. /// Core struct, update/merge/serde logic live in sketch-core. /// This file retains QE-specific trait impls and JSON output. #[derive(Debug, Clone)] @@ -88,11 +88,8 @@ impl AggregateCore for HydraKllSketchAccumulator { self } - fn as_any_mut(&mut self) -> &mut dyn std::any::Any { - self - } fn merge_with( diff --git a/asap-query-engine/src/precompute_operators/set_aggregator_accumulator.rs b/asap-query-engine/src/precompute_operators/set_aggregator_accumulator.rs index d234363e..4baa69d5 100644 --- a/asap-query-engine/src/precompute_operators/set_aggregator_accumulator.rs +++ b/asap-query-engine/src/precompute_operators/set_aggregator_accumulator.rs @@ -2,13 +2,13 @@ use crate::data_model::{ AggregateCore, AggregationType, KeyByLabelValues, MergeableAccumulator, MultipleSubpopulationAggregate, SerializableToSink, }; +use asap_sketchlib::asap::set_aggregator::SetAggregator; use serde_json::Value; -use sketch_core::set_aggregator::SetAggregator; use std::collections::{HashMap, HashSet}; use promql_utilities::query_logics::enums::Statistic; -/// Set aggregator accumulator — wraps sketch_core::SetAggregator. +/// Set aggregator accumulator — wraps asap_sketchlib::asap::SetAggregator. /// Core struct, merge/serde logic live in sketch-core. /// This file retains QE-specific trait impls, KeyByLabelValues conversion, /// and legacy deserializers. @@ -154,11 +154,8 @@ impl AggregateCore for SetAggregatorAccumulator { self } - fn as_any_mut(&mut self) -> &mut dyn std::any::Any { - self - } fn merge_with( diff --git a/asap-query-engine/src/stores/promsketch_store/query.rs b/asap-query-engine/src/stores/promsketch_store/query.rs index 0de029bb..2f610eac 100644 --- a/asap-query-engine/src/stores/promsketch_store/query.rs +++ b/asap-query-engine/src/stores/promsketch_store/query.rs @@ -1,4 +1,4 @@ -use asap_sketchlib::{EHSketchList, SketchInput, UniformSampling}; +use asap_sketchlib::{EHSketchList, DataInput, UniformSampling}; use super::series::PromSketchMemSeries; @@ -80,7 +80,7 @@ fn eval_kll_quantile( .ok_or("no volumes cover the requested time range for KLL")?; merged - .query(&SketchInput::F64(phi)) + .query(&DataInput::F64(phi)) .map_err(|e| -> Box { e.into() }) } @@ -174,7 +174,7 @@ mod tests { // Insert values 1..=100 at successive timestamps for i in 1..=100u64 { - let input = SketchInput::F64(i as f64); + let input = DataInput::F64(i as f64); if let Some(ref mut eh) = series.sketch_instances.eh_kll { eh.update(i, &input); } @@ -190,7 +190,7 @@ mod tests { .ensure_initialized(PromSketchType::USampling, &config); for i in 1..=1000u64 { - let input = SketchInput::F64(i as f64); + let input = DataInput::F64(i as f64); if let Some(ref mut eh) = series.sketch_instances.eh_sampling { eh.update(i, &input); } @@ -206,7 +206,7 @@ mod tests { .ensure_initialized(PromSketchType::EHUniv, &config); for i in 1..=100u64 { - let input = SketchInput::F64(i as f64); + let input = DataInput::F64(i as f64); if let Some(ref mut eh) = series.sketch_instances.eh_univ { eh.update(i, &input, 1); } diff --git a/asap-query-engine/src/stores/promsketch_store/series.rs b/asap-query-engine/src/stores/promsketch_store/series.rs index cdffe49f..c5c26f95 100644 --- a/asap-query-engine/src/stores/promsketch_store/series.rs +++ b/asap-query-engine/src/stores/promsketch_store/series.rs @@ -1,5 +1,5 @@ use asap_sketchlib::{ - EHSketchList, EHUnivOptimized, ExponentialHistogram, SketchInput, UniformSampling, KLL, + EHSketchList, EHUnivOptimized, ExponentialHistogram, DataInput, UniformSampling, KLL, }; use super::config::PromSketchConfig; @@ -68,7 +68,7 @@ impl PromSketchInstances { /// Insert a data point into all active sketches. pub fn insert(&mut self, time: u64, value: f64) { - let input = SketchInput::F64(value); + let input = DataInput::F64(value); if let Some(ref mut eh) = self.eh_univ { // EHUnivOptimized::update(time, key, frequency_count) diff --git a/asap-query-engine/src/tests/test_utilities/engine_factories.rs b/asap-query-engine/src/tests/test_utilities/engine_factories.rs index 80cc67b8..cbaaeae6 100644 --- a/asap-query-engine/src/tests/test_utilities/engine_factories.rs +++ b/asap-query-engine/src/tests/test_utilities/engine_factories.rs @@ -637,7 +637,7 @@ pub async fn assert_old_new_match(engine: &SimpleEngine, query: &str, query_time .build_query_execution_context_promql(query.to_string(), query_time_sec) .expect("Failed to build context"); - let old_results = engine + let (old_results, _) = engine .execute_query_pipeline(&context, false) .expect("Old pipeline failed"); diff --git a/asap-query-engine/tests/e2e_modified_otlp_sketch_path.rs b/asap-query-engine/tests/e2e_modified_otlp_sketch_path.rs index d91fe0dc..68a0ed9a 100644 --- a/asap-query-engine/tests/e2e_modified_otlp_sketch_path.rs +++ b/asap-query-engine/tests/e2e_modified_otlp_sketch_path.rs @@ -1191,7 +1191,8 @@ async fn e2e_count_min_sketch_msgpack_modified_otlp_path() { // Build a known sketch in sketch-core and serialize with msgpack — this // is what the Go producer (sketchlib-go) will emit once PR I's matching // Go-side work lands. - let mut cms = sketch_core::count_min::CountMinSketch::new(rows as usize, cols as usize); + let mut cms = + asap_sketchlib::asap::count_min::CountMinSketch::new(rows as usize, cols as usize); cms.update("user_a", 1.0); cms.update("user_b", 1.0); cms.update("user_a", 1.0); @@ -1207,7 +1208,7 @@ async fn e2e_count_min_sketch_msgpack_modified_otlp_path() { post_otlp_http(&client, otlp_http_port, req).await; // Watermark advance using an empty msgpack sketch. - let empty = sketch_core::count_min::CountMinSketch::new(rows as usize, cols as usize); + let empty = asap_sketchlib::asap::count_min::CountMinSketch::new(rows as usize, cols as usize); let watermark_req = build_count_min_msgpack_export_request( metric_name, service_label, diff --git a/asap-query-engine/tests/e2e_precompute_equivalence.rs b/asap-query-engine/tests/e2e_precompute_equivalence.rs index 6ddf4c77..aa894b86 100644 --- a/asap-query-engine/tests/e2e_precompute_equivalence.rs +++ b/asap-query-engine/tests/e2e_precompute_equivalence.rs @@ -7,12 +7,12 @@ //! 3. Advances the watermark past the window boundary to close it //! 4. Drains captured outputs and verifies equivalence with ArroYo-format accumulators +use asap_sketchlib::asap::kll::KllSketch; use asap_types::aggregation_config::AggregationConfig; use asap_types::enums::{AggregationType, WindowType}; use flate2::{write::GzEncoder, Compression}; use prost::Message; use serde_json::json; -use sketch_core::kll::KllSketch; use std::collections::HashMap; use std::io::Write; use std::sync::Arc; From 60e8f5b171ee50ef1fb08ea4b84b39976f473350 Mon Sep 17 00:00:00 2001 From: Zeying Zhu Date: Fri, 1 May 2026 15:32:44 -0400 Subject: [PATCH 2/2] refactor: switch consumer imports to asap_sketchlib::sketches::* MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Update PR #73 against the reorganized asap_sketchlib (PR #36): the runtime sketches no longer live under a dedicated `asap::` module — they were merged into the existing `src/sketches/` layout (single home per sketch concept, ASAP-runtime types appended to the file that already holds the high-throughput in-process variant). Mechanical path swaps in asap-query-engine: - `asap_sketchlib::asap::dd_sketch::*` → `::sketches::ddsketch::*` - `asap_sketchlib::asap::count_min::*` → `::sketches::countmin::*` - `asap_sketchlib::asap::count_sketch::*` → `::sketches::count::*` - `asap_sketchlib::asap::hll_sketch::*` → `::sketches::hll::*` - `asap_sketchlib::asap::kll::*` → `::sketches::kll::*` - `asap_sketchlib::asap::count_min_with_heap::*` → `::sketches::cms_heap::*` - `asap_sketchlib::asap::hydra_kll::*` → `::sketches::hydra_kll::*` - `asap_sketchlib::asap::set_aggregator::*` → `::sketches::set_aggregator::*` - `asap_sketchlib::asap::delta_set_aggregator::*`→ `::sketches::delta_set_aggregator::*` - `asap_sketchlib::asap::config::*` → `::asap_runtime::*` Naming-conflict renames carried through to the consumers: - `HllDelta` → `HllSketchDelta` (octo_delta::HllDelta still wins the short name) - `HeapItem` → `CmsHeapItem` (common::input::HeapItem still wins the short name) main.rs aliases `asap_sketchlib::asap_runtime as config` so the existing clap derive references (`config::DEFAULT_CMS_IMPL`, `config::configure(...)`) still work without touching the rest of the bin. Tests: - `cargo build --workspace` → clean - `cargo test -p query_engine_rust --lib precompute_operators` → 141 passed, 0 failed Depends on ProjectASAP/asap_sketchlib#36 (force-pushed `e473ccc`). Co-Authored-By: Claude Opus 4.7 (1M context) --- Cargo.lock | 175 +----------------- asap-query-engine/Cargo.toml | 5 +- asap-query-engine/src/drivers/ingest/otel.rs | 4 +- .../src/drivers/query/servers/http.rs | 18 +- .../src/engines/simple_engine.rs | 15 +- asap-query-engine/src/lib.rs | 13 -- asap-query-engine/src/main.rs | 22 --- .../precompute_engine/accumulator_factory.rs | 10 +- .../src/precompute_engine/ingest_handler.rs | 2 +- .../src/precompute_engine/worker.rs | 2 +- .../count_min_sketch_accumulator.rs | 47 ++--- .../count_min_sketch_with_heap_accumulator.rs | 63 +++---- .../count_sketch_accumulator.rs | 20 +- .../datasketches_kll_accumulator.rs | 55 ++---- .../dd_sketch_accumulator.rs | 8 +- .../delta_set_aggregator_accumulator.rs | 9 +- .../hll_sketch_accumulator.rs | 12 +- .../hydra_kll_accumulator.rs | 26 +-- .../increase_accumulator.rs | 3 - .../min_max_accumulator.rs | 3 - .../multiple_increase_accumulator.rs | 3 - .../multiple_min_max_accumulator.rs | 3 - .../multiple_sum_accumulator.rs | 3 - .../set_aggregator_accumulator.rs | 13 +- .../sketch_envelope_accumulator.rs | 3 - .../precompute_operators/sum_accumulator.rs | 3 - .../tests/e2e_modified_otlp_sketch_path.rs | 17 +- .../tests/e2e_precompute_equivalence.rs | 26 +-- asap-query-engine/tests/test_both_backends.rs | 30 --- 29 files changed, 170 insertions(+), 443 deletions(-) delete mode 100644 asap-query-engine/tests/test_both_backends.rs diff --git a/Cargo.lock b/Cargo.lock index 37af1546..febd8591 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -423,11 +423,9 @@ dependencies = [ [[package]] name = "asap_sketchlib" version = "0.1.0" -source = "git+https://github.com/ProjectASAP/asap_sketchlib?branch=refactor%2Fadopt-sketch-core-modules#63051a830e029679db9fe8da3cd4488773a18bbe" +source = "git+https://github.com/ProjectASAP/asap_sketchlib?branch=refactor%2Fadopt-sketch-core-modules#d84ff152c7ac7c90b97bf2fbe0d88f28c147d7a6" dependencies = [ "bytes", - "clap 4.6.0", - "dsrs", "prost", "prost-build", "protoc-bin-vendored", @@ -592,12 +590,6 @@ dependencies = [ "tracing", ] -[[package]] -name = "base64" -version = "0.13.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9e1b586273c5702936fe7b7d6896644d8be71e6314cfe09d3167c95f712589e8" - [[package]] name = "base64" version = "0.21.7" @@ -684,17 +676,6 @@ dependencies = [ "alloc-stdlib", ] -[[package]] -name = "bstr" -version = "0.2.17" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ba3569f383e8f1598449f1a423e72e99569137b47740b1da11ef19af3d5c3223" -dependencies = [ - "lazy_static", - "memchr", - "regex-automata 0.1.10", -] - [[package]] name = "bumpalo" version = "3.20.2" @@ -893,17 +874,6 @@ version = "1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c8d4a3bb8b1e0c1050499d1815f5ab16d04f0959b233085fb31653fbfc9d98f9" -[[package]] -name = "codespan-reporting" -version = "0.13.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "af491d569909a7e4dee0ad7db7f5341fef5c614d5b8ec8cf765732aba3cff681" -dependencies = [ - "serde", - "termcolor", - "unicode-width 0.2.2", -] - [[package]] name = "colorchoice" version = "1.0.5" @@ -1105,78 +1075,6 @@ dependencies = [ "memchr", ] -[[package]] -name = "ctor" -version = "0.2.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "32a2785755761f3ddc1492979ce1e48d2c00d09311c39e4466429188f3dd6501" -dependencies = [ - "quote", - "syn 2.0.117", -] - -[[package]] -name = "cxx" -version = "1.0.194" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "747d8437319e3a2f43d93b341c137927ca70c0f5dabeea7a005a73665e247c7e" -dependencies = [ - "cc", - "cxx-build", - "cxxbridge-cmd", - "cxxbridge-flags", - "cxxbridge-macro", - "foldhash 0.2.0", - "link-cplusplus", -] - -[[package]] -name = "cxx-build" -version = "1.0.194" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b0f4697d190a142477b16aef7da8a99bfdc41e7e8b1687583c0d23a79c7afc1e" -dependencies = [ - "cc", - "codespan-reporting", - "indexmap 2.13.1", - "proc-macro2", - "quote", - "scratch", - "syn 2.0.117", -] - -[[package]] -name = "cxxbridge-cmd" -version = "1.0.194" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d0956799fa8678d4c50eed028f2de1c0552ae183c76e976cf7ca8c4e36a7c328" -dependencies = [ - "clap 4.6.0", - "codespan-reporting", - "indexmap 2.13.1", - "proc-macro2", - "quote", - "syn 2.0.117", -] - -[[package]] -name = "cxxbridge-flags" -version = "1.0.194" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "23384a836ab4f0ad98ace7e3955ad2de39de42378ab487dc28d3990392cb283a" - -[[package]] -name = "cxxbridge-macro" -version = "1.0.194" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e6acc6b5822b9526adfb4fc377b67128fdd60aac757cc4a741a6278603f763cf" -dependencies = [ - "indexmap 2.13.1", - "proc-macro2", - "quote", - "syn 2.0.117", -] - [[package]] name = "dashmap" version = "5.5.3" @@ -1657,22 +1555,6 @@ dependencies = [ "syn 2.0.117", ] -[[package]] -name = "dsrs" -version = "0.6.1" -source = "git+https://github.com/ProjectASAP/datasketches-rs?rev=d748ec75c80fff21f7b24897244dd1c895df2e9a#d748ec75c80fff21f7b24897244dd1c895df2e9a" -dependencies = [ - "base64 0.13.1", - "bstr", - "cxx", - "cxx-build", - "memchr", - "rmp-serde", - "serde", - "structopt", - "thin-dst", -] - [[package]] name = "either" version = "1.15.0" @@ -1786,12 +1668,6 @@ version = "0.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d9c4f5dac5e15c24eb999c26181a6ca40b39fe946cbe4c263c7209467bc83af2" -[[package]] -name = "foldhash" -version = "0.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "77ce24cb58228fbb8aa041425bb1050850ac19177686ea6e0f41a70416f56fdb" - [[package]] name = "form_urlencoded" version = "1.2.2" @@ -2022,7 +1898,7 @@ version = "0.15.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9229cfe53dfd69f0609a49f65461bd93001ea1ef889cd5529dd176593f5338a1" dependencies = [ - "foldhash 0.1.5", + "foldhash", ] [[package]] @@ -2650,15 +2526,6 @@ dependencies = [ "vcpkg", ] -[[package]] -name = "link-cplusplus" -version = "1.0.12" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7f78c730aaa7d0b9336a299029ea49f9ee53b0ed06e9202e8cb7db9bae7b8c82" -dependencies = [ - "cc", -] - [[package]] name = "linux-raw-sys" version = "0.12.1" @@ -2766,7 +2633,7 @@ version = "0.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d1525a2a28c7f4fa0fc98bb91ae755d1e2d1505079e05539e35bc876b5d65ae9" dependencies = [ - "regex-automata 0.4.14", + "regex-automata", ] [[package]] @@ -2856,7 +2723,7 @@ version = "0.50.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" dependencies = [ - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -3531,11 +3398,9 @@ dependencies = [ "clap 4.6.0", "crc32fast", "criterion", - "ctor", "dashmap 5.5.3", "datafusion", "datafusion_summary_library", - "dsrs", "elastic_dsl_utilities", "flate2", "form_urlencoded", @@ -3750,16 +3615,10 @@ checksum = "e10754a14b9137dd7b1e3e5b0493cc9171fdd105e0ab477f51b72e7f3ac0e276" dependencies = [ "aho-corasick", "memchr", - "regex-automata 0.4.14", + "regex-automata", "regex-syntax", ] -[[package]] -name = "regex-automata" -version = "0.1.10" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6c230d73fb8d8c1b9c0b3135c5142a8acee3a0558fb8db5cf1cb65f8d7862132" - [[package]] name = "regex-automata" version = "0.4.14" @@ -3945,12 +3804,6 @@ version = "1.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" -[[package]] -name = "scratch" -version = "1.0.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d68f2ec51b097e4c1a75b681a8bec621909b5e91f15bb7b840c4f2f7b01148b2" - [[package]] name = "sct" version = "0.7.1" @@ -4237,6 +4090,7 @@ dependencies = [ "cfg-if", "libc", "psm", + "windows-sys 0.52.0", "windows-sys 0.59.0", ] @@ -4395,15 +4249,6 @@ dependencies = [ "windows-sys 0.52.0", ] -[[package]] -name = "termcolor" -version = "1.4.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "06794f8f6c5c898b3275aebefa6b8a1cb24cd2c6c79397ab15774837a0bc5755" -dependencies = [ - "winapi-util", -] - [[package]] name = "textwrap" version = "0.11.0" @@ -4413,12 +4258,6 @@ dependencies = [ "unicode-width 0.1.14", ] -[[package]] -name = "thin-dst" -version = "1.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "db3c46be180f1af9673ebb27bc1235396f61ef6965b3fe0dbb2e624deb604f0e" - [[package]] name = "thiserror" version = "1.0.69" @@ -4801,7 +4640,7 @@ dependencies = [ "matchers", "nu-ansi-term", "once_cell", - "regex-automata 0.4.14", + "regex-automata", "sharded-slab", "smallvec", "thread_local", diff --git a/asap-query-engine/Cargo.toml b/asap-query-engine/Cargo.toml index 0bf130f9..90c4c2ef 100644 --- a/asap-query-engine/Cargo.toml +++ b/asap-query-engine/Cargo.toml @@ -39,7 +39,6 @@ urlencoding = "2.1" flate2 = "1.0" async-trait = "0.1" xxhash-rust = { version = "0.8", features = ["xxh32", "xxh64"] } -dsrs = { git = "https://github.com/ProjectASAP/datasketches-rs", rev = "d748ec75c80fff21f7b24897244dd1c895df2e9a" } base64 = "0.21" hex = "0.4" sqlparser = "0.59.0" @@ -62,7 +61,7 @@ zstd = "0.13" reqwest = { version = "0.11", default-features = false, features = ["json", "rustls-tls"] } tracing-appender = "0.2" elastic_dsl_utilities.workspace = true -asap_sketchlib = { git = "https://github.com/ProjectASAP/asap_sketchlib", branch = "refactor/adopt-sketch-core-modules", features = ["asap-cli"] } +asap_sketchlib = { git = "https://github.com/ProjectASAP/asap_sketchlib", branch = "refactor/adopt-sketch-core-modules" } # Persistence layer (SimpleMapStore parts / manifest / Tier-2 cache) moka = { version = "0.12", features = ["sync"] } memmap2 = "0.9" @@ -85,7 +84,6 @@ name = "e2e_quickstart_resource_test" path = "src/bin/e2e_quickstart_resource_test.rs" [dev-dependencies] -ctor = "0.2" tempfile = "3.20.0" criterion = { version = "0.5", features = ["html_reports"] } @@ -100,4 +98,3 @@ default = [] lock_profiling = [] # Enable extra debugging output extra_debugging = [] -sketchlib-tests = [] diff --git a/asap-query-engine/src/drivers/ingest/otel.rs b/asap-query-engine/src/drivers/ingest/otel.rs index c27b5913..cbe34887 100644 --- a/asap-query-engine/src/drivers/ingest/otel.rs +++ b/asap-query-engine/src/drivers/ingest/otel.rs @@ -1329,8 +1329,8 @@ mod dispatcher_tests { use super::*; use crate::data_model::AggregateCore; use crate::precompute_operators::{DDSketchAccumulator, HllSketchAccumulator}; - use asap_sketchlib::asap::dd_sketch::DdSketch; - use asap_sketchlib::asap::hll_sketch::HllVariant; + use asap_sketchlib::sketches::ddsketch::DdSketch; + use asap_sketchlib::sketches::hll::HllVariant; #[test] fn apply_modified_otlp_delta_bytes_ddsketch_round_trip() { diff --git a/asap-query-engine/src/drivers/query/servers/http.rs b/asap-query-engine/src/drivers/query/servers/http.rs index 20ac2ac5..256d26ba 100644 --- a/asap-query-engine/src/drivers/query/servers/http.rs +++ b/asap-query-engine/src/drivers/query/servers/http.rs @@ -429,8 +429,7 @@ async fn handle_instant_query( } }; - let response = - process_query_request(&state, &parsed_request, start_time, HashMap::new()).await; + let response = process_query_request(&state, &parsed_request, start_time, HashMap::new()).await; srv_metrics::record_query_outcome( srv_metrics::QUERY_TYPE_INSTANT, query_status_label(&response), @@ -552,10 +551,7 @@ async fn handle_instant_query_post( total_duration.as_secs_f64() * 1000.0 ); - srv_metrics::record_query_outcome( - srv_metrics::QUERY_TYPE_INSTANT, - query_status_label(&result), - ); + srv_metrics::record_query_outcome(srv_metrics::QUERY_TYPE_INSTANT, query_status_label(&result)); result } @@ -730,10 +726,7 @@ async fn handle_range_query( }; let response = process_range_query_request(&state, &parsed_request, start_time).await; - srv_metrics::record_query_outcome( - srv_metrics::QUERY_TYPE_RANGE, - query_status_label(&response), - ); + srv_metrics::record_query_outcome(srv_metrics::QUERY_TYPE_RANGE, query_status_label(&response)); response } @@ -794,10 +787,7 @@ async fn handle_range_query_post(State(state): State, body: Bytes) -> }; let response = process_range_query_request(&state, &parsed_request, start_time).await; - srv_metrics::record_query_outcome( - srv_metrics::QUERY_TYPE_RANGE, - query_status_label(&response), - ); + srv_metrics::record_query_outcome(srv_metrics::QUERY_TYPE_RANGE, query_status_label(&response)); response } diff --git a/asap-query-engine/src/engines/simple_engine.rs b/asap-query-engine/src/engines/simple_engine.rs index 0f5b48e0..530917ef 100644 --- a/asap-query-engine/src/engines/simple_engine.rs +++ b/asap-query-engine/src/engines/simple_engine.rs @@ -876,7 +876,14 @@ impl SimpleEngine { plan: &StoreQueryPlan, do_merge: bool, agg_info: &AggregationIdInfo, - ) -> Result<(MergedOutputsMap, Option, Option<(u64, u64)>), String> { + ) -> Result< + ( + MergedOutputsMap, + Option, + Option<(u64, u64)>, + ), + String, + > { // Query and merge values let values_map = self.execute_store_query(&plan.values_query).map_err(|e| { warn!("Error querying store for values: {}", e); @@ -961,11 +968,7 @@ impl SimpleEngine { target.1, filtered.len() ); - self.merge_precomputed_outputs( - &filtered, - do_merge, - agg_info.aggregation_type_for_value, - ) + self.merge_precomputed_outputs(&filtered, do_merge, agg_info.aggregation_type_for_value) }; let merge_duration = merge_start_time.elapsed(); diff --git a/asap-query-engine/src/lib.rs b/asap-query-engine/src/lib.rs index 367b227d..659f99e5 100644 --- a/asap-query-engine/src/lib.rs +++ b/asap-query-engine/src/lib.rs @@ -1,16 +1,3 @@ -#[cfg(test)] -#[ctor::ctor] -fn init_sketch_backend_for_tests() { - #[cfg(feature = "sketchlib-tests")] - let _ = asap_sketchlib::asap::config::configure( - asap_sketchlib::asap::config::ImplMode::Sketchlib, - asap_sketchlib::asap::config::ImplMode::Legacy, - asap_sketchlib::asap::config::ImplMode::Sketchlib, - ); - #[cfg(not(feature = "sketchlib-tests"))] - asap_sketchlib::asap::config::force_legacy_mode_for_tests(); -} - pub mod data_model; pub mod drivers; pub mod engines; diff --git a/asap-query-engine/src/main.rs b/asap-query-engine/src/main.rs index 9e106f22..26f2d38e 100644 --- a/asap-query-engine/src/main.rs +++ b/asap-query-engine/src/main.rs @@ -5,8 +5,6 @@ use std::sync::Arc; use tokio::signal; use tracing::{error, info, warn}; -use asap_sketchlib::asap::config::{self, ImplMode}; - use query_engine_rust::data_model::enums::{InputFormat, LockStrategy, StreamingEngine}; use query_engine_rust::drivers::AdapterConfig; use query_engine_rust::precompute_engine::config::LateDataPolicy; @@ -150,18 +148,6 @@ struct Args { #[arg(long)] promsketch_config: Option, - /// Backend implementation for Count-Min Sketch (legacy | sketchlib) - #[arg(long, value_enum, default_value_t = config::DEFAULT_CMS_IMPL)] - sketch_cms_impl: ImplMode, - - /// Backend implementation for KLL Sketch (legacy | sketchlib) - #[arg(long, value_enum, default_value_t = config::DEFAULT_KLL_IMPL)] - sketch_kll_impl: ImplMode, - - /// Backend implementation for Count-Min-With-Heap (legacy | sketchlib) - #[arg(long, value_enum, default_value_t = config::DEFAULT_CMWH_IMPL)] - sketch_cmwh_impl: ImplMode, - /// Enable OTLP metrics ingest (gRPC + HTTP) #[arg(long)] enable_otel_ingest: bool, @@ -297,14 +283,6 @@ struct Args { async fn main() -> Result<()> { let args = Args::parse(); - // Configure sketch-core backends before any sketch operations. - config::configure( - args.sketch_cms_impl, - args.sketch_kll_impl, - args.sketch_cmwh_impl, - ) - .expect("sketch backend already initialised"); - // Create output directory fs::create_dir_all(&args.output_dir)?; diff --git a/asap-query-engine/src/precompute_engine/accumulator_factory.rs b/asap-query-engine/src/precompute_engine/accumulator_factory.rs index 8e572640..840bb088 100644 --- a/asap-query-engine/src/precompute_engine/accumulator_factory.rs +++ b/asap-query-engine/src/precompute_engine/accumulator_factory.rs @@ -295,11 +295,11 @@ impl DDSketchAccumulatorUpdater { impl AccumulatorUpdater for DDSketchAccumulatorUpdater { fn update_single(&mut self, value: f64, _timestamp_ms: i64) { // sketch-core's DdSketch (the inner of DDSketchAccumulator) - // exposes `insert(f64)` for single-value ingestion. The + // exposes `update(f64)` for single-value ingestion. The // worker calls this when a raw OTLP datapoint matches an // aggregation typed as DDSketch — the sketch-merge path // uses `merge_with` directly. - self.acc.inner.insert(value); + self.acc.inner.update(value); } fn update_keyed(&mut self, _key: &KeyByLabelValues, value: f64, timestamp_ms: i64) { @@ -734,9 +734,9 @@ pub fn create_accumulator_updater(config: &AggregationConfig) -> Box { - Box::new(DDSketchAccumulatorUpdater::new(ddsketch_alpha_param(config))) - } + AggregationType::DDSketch => Box::new(DDSketchAccumulatorUpdater::new( + ddsketch_alpha_param(config), + )), other => { tracing::warn!( "Unknown aggregation_type '{:?}', defaulting to SingleSubpopulation Sum", diff --git a/asap-query-engine/src/precompute_engine/ingest_handler.rs b/asap-query-engine/src/precompute_engine/ingest_handler.rs index a9ec8dff..c7bf3359 100644 --- a/asap-query-engine/src/precompute_engine/ingest_handler.rs +++ b/asap-query-engine/src/precompute_engine/ingest_handler.rs @@ -554,7 +554,7 @@ mod tests { use crate::drivers::ingest::otel::{apply_modified_otlp_delta_bytes, SketchKind}; use crate::precompute_operators::DDSketchAccumulator; use asap_otel_proto::sketchlib::v1::{DdSketchBucketDelta, DdSketchDelta as PbDelta}; - use asap_sketchlib::asap::dd_sketch::DdSketch; + use asap_sketchlib::sketches::ddsketch::DdSketch; use prost::Message; const ENCODING_PROTO_DELTA: i32 = 2; diff --git a/asap-query-engine/src/precompute_engine/worker.rs b/asap-query-engine/src/precompute_engine/worker.rs index df2e8520..b3f08393 100644 --- a/asap-query-engine/src/precompute_engine/worker.rs +++ b/asap-query-engine/src/precompute_engine/worker.rs @@ -964,7 +964,7 @@ mod tests { use crate::precompute_operators::datasketches_kll_accumulator::DatasketchesKLLAccumulator; use crate::precompute_operators::multiple_sum_accumulator::MultipleSumAccumulator; use crate::precompute_operators::sum_accumulator::SumAccumulator; - use asap_sketchlib::asap::kll::KllSketch; + use asap_sketchlib::sketches::kll::KllSketch; use asap_types::enums::{AggregationType, WindowType}; fn make_agg_config( diff --git a/asap-query-engine/src/precompute_operators/count_min_sketch_accumulator.rs b/asap-query-engine/src/precompute_operators/count_min_sketch_accumulator.rs index b311dbc1..ca6a06a2 100644 --- a/asap-query-engine/src/precompute_operators/count_min_sketch_accumulator.rs +++ b/asap-query-engine/src/precompute_operators/count_min_sketch_accumulator.rs @@ -2,14 +2,14 @@ use crate::data_model::{ AggregateCore, AggregationType, KeyByLabelValues, MergeableAccumulator, MultipleSubpopulationAggregate, SerializableToSink, }; -use asap_sketchlib::asap::count_min::{CountMinDelta, CountMinSketch}; +use asap_sketchlib::sketches::countmin::{CountMinSketch, CountMinSketchDelta}; use serde_json::Value; use std::collections::HashMap; use promql_utilities::query_logics::enums::Statistic; -/// Count-Min Sketch accumulator — wraps asap_sketchlib::asap::CountMinSketch. -/// Core struct, update/merge/serde logic live in sketch-core. +/// Count-Min Sketch accumulator — wraps asap_sketchlib::sketches::CountMinSketch. +/// Core struct, update/merge/serde logic live in `asap_sketchlib::sketches`. /// This file retains QE-specific trait impls, legacy deserializers, and JSON output. #[derive(Debug, Clone)] pub struct CountMinSketchAccumulator { @@ -29,7 +29,7 @@ impl CountMinSketchAccumulator { } pub fn query_key(&self, key: &KeyByLabelValues) -> f64 { - self.inner.query_key(&key.to_semicolon_str()) + self.inner.estimate(&key.to_semicolon_str()) } pub fn deserialize_from_json(data: &Value) -> Result> { @@ -64,7 +64,8 @@ impl CountMinSketchAccumulator { buffer: &[u8], ) -> Result> { Ok(Self { - inner: CountMinSketch::deserialize_msgpack(buffer)?, + inner: CountMinSketch::deserialize_msgpack(buffer) + .map_err(|e| -> Box { e.to_string().into() })?, }) } @@ -76,7 +77,8 @@ impl CountMinSketchAccumulator { /// uses — this method is the modified-OTLP entrypoint for PR I). pub fn from_msgpack_bytes(buffer: &[u8]) -> Result> { Ok(Self { - inner: CountMinSketch::deserialize_msgpack(buffer)?, + inner: CountMinSketch::deserialize_msgpack(buffer) + .map_err(|e| -> Box { e.to_string().into() })?, }) } @@ -205,7 +207,7 @@ impl CountMinSketchAccumulator { .zip(pb.d_counts.iter()) .map(|((r, c), dc)| (*r, *c, *dc)) .collect(); - let delta = CountMinDelta { + let delta = CountMinSketchDelta { rows: pb.rows, cols: pb.cols, cells, @@ -286,10 +288,10 @@ impl CountMinSketchAccumulator { } // Check dimensions are consistent - let row_num = cms_accumulators[0].inner.row_num; - let col_num = cms_accumulators[0].inner.col_num; + let rows = cms_accumulators[0].inner.rows(); + let cols = cms_accumulators[0].inner.cols(); for acc in &cms_accumulators { - if acc.inner.row_num != row_num || acc.inner.col_num != col_num { + if acc.inner.rows() != rows || acc.inner.cols() != cols { return Err( "Cannot merge CountMinSketch accumulators with different dimensions".into(), ); @@ -308,14 +310,14 @@ impl CountMinSketchAccumulator { impl SerializableToSink for CountMinSketchAccumulator { fn serialize_to_json(&self) -> Value { serde_json::json!({ - "row_num": self.inner.row_num, - "col_num": self.inner.col_num, + "row_num": self.inner.rows(), + "col_num": self.inner.cols(), "sketch": self.inner.sketch() }) } fn serialize_to_bytes(&self) -> Vec { - self.inner.serialize_msgpack() + self.inner.serialize_msgpack().unwrap_or_default() } } @@ -444,11 +446,12 @@ impl MergeableAccumulator for CountMinSketchAccumulat if accumulators.is_empty() { return Err("No accumulators to merge".into()); } - let inners: Vec = accumulators.into_iter().map(|acc| acc.inner).collect(); - let merged_inner = CountMinSketch::merge(inners)?; - Ok(Self { - inner: merged_inner, - }) + let mut iter = accumulators.into_iter(); + let mut merged = iter.next().unwrap(); + for acc in iter { + merged.inner.merge(&acc.inner)?; + } + Ok(merged) } } @@ -459,8 +462,8 @@ mod tests { #[test] fn test_count_min_sketch_creation() { let cms = CountMinSketchAccumulator::new(4, 1000); - assert_eq!(cms.inner.row_num, 4); - assert_eq!(cms.inner.col_num, 1000); + assert_eq!(cms.inner.rows(), 4); + assert_eq!(cms.inner.cols(), 1000); let sketch = cms.inner.sketch(); assert_eq!(sketch.len(), 4); assert_eq!(sketch[0].len(), 1000); @@ -539,8 +542,8 @@ mod tests { let deserialized = CountMinSketchAccumulator::deserialize_from_bytes_arroyo(&bytes).unwrap(); - assert_eq!(deserialized.inner.row_num, 2); - assert_eq!(deserialized.inner.col_num, 3); + assert_eq!(deserialized.inner.rows(), 2); + assert_eq!(deserialized.inner.cols(), 3); let deser_sketch = deserialized.inner.sketch(); assert_eq!(deser_sketch[0][1], 42.0); assert_eq!(deser_sketch[1][2], 100.0); diff --git a/asap-query-engine/src/precompute_operators/count_min_sketch_with_heap_accumulator.rs b/asap-query-engine/src/precompute_operators/count_min_sketch_with_heap_accumulator.rs index 53c61917..4e1ef6d9 100644 --- a/asap-query-engine/src/precompute_operators/count_min_sketch_with_heap_accumulator.rs +++ b/asap-query-engine/src/precompute_operators/count_min_sketch_with_heap_accumulator.rs @@ -2,25 +2,22 @@ use crate::data_model::{ AggregateCore, AggregationType, KeyByLabelValues, MergeableAccumulator, MultipleSubpopulationAggregate, SerializableToSink, }; -use asap_sketchlib::asap::count_min_with_heap::{CountMinSketchWithHeap, HeapItem}; +use asap_sketchlib::sketches::cms_heap::{CmsHeapItem, CountMinSketchWithHeap}; use serde_json::Value; use std::collections::HashMap; use promql_utilities::query_logics::enums::Statistic; -/// Count-Min Sketch with Heap accumulator — wraps asap_sketchlib::asap::CountMinSketchWithHeap. -/// Core struct, update/merge/serde logic live in sketch-core. +/// Count-Min Sketch with Heap accumulator — wraps `asap_sketchlib::sketches::CountMinSketchWithHeap`. +/// Core struct, update/merge/serde logic live in `asap_sketchlib::sketches::cms_heap`. /// This file retains QE-specific trait impls, legacy deserializers, and JSON output. -/// -/// NOTE (bug, do not fix): QueryEngineRust uses xxhash-rust::xxh32; the Arroyo template uses -/// twox-hash::XxHash32. Bucket assignments differ. Tracked separately. #[derive(Debug, Clone)] pub struct CountMinSketchWithHeapAccumulator { pub inner: CountMinSketchWithHeap, } // Re-export HeapItem so existing code using CountMinSketchWithHeapAccumulator::HeapItem still works. -pub use asap_sketchlib::asap::count_min_with_heap::HeapItem as HeapItemReexport; +pub use asap_sketchlib::sketches::cms_heap::CmsHeapItem as HeapItemReexport; impl CountMinSketchWithHeapAccumulator { pub fn new(row_num: usize, col_num: usize, heap_size: usize) -> Self { @@ -31,7 +28,7 @@ impl CountMinSketchWithHeapAccumulator { pub fn query_key(&self, key: &KeyByLabelValues) -> f64 { let key_string = key.labels.join(";"); - self.inner.query_key(&key_string) + self.inner.estimate(&key_string) } /// This function seems will never be used anymore. Keep it for possible future use. @@ -74,7 +71,7 @@ impl CountMinSketchWithHeapAccumulator { let value = item["value"] .as_f64() .ok_or("Missing or invalid 'value' in heap item")?; - topk_heap.push(HeapItem { key, value }); + topk_heap.push(CmsHeapItem { key, value }); } Ok(Self { @@ -88,7 +85,8 @@ impl CountMinSketchWithHeapAccumulator { buffer: &[u8], ) -> Result> { Ok(Self { - inner: CountMinSketchWithHeap::deserialize_msgpack(buffer)?, + inner: CountMinSketchWithHeap::deserialize_msgpack(buffer) + .map_err(|e| -> Box { e.to_string().into() })?, }) } @@ -124,8 +122,8 @@ impl SerializableToSink for CountMinSketchWithHeapAccumulator { .collect(); serde_json::json!({ - "row_num": self.inner.row_num, - "col_num": self.inner.col_num, + "row_num": self.inner.rows(), + "col_num": self.inner.cols(), "heap_size": self.inner.heap_size, "sketch": self.inner.sketch_matrix(), "topk_heap": heap_items @@ -133,7 +131,7 @@ impl SerializableToSink for CountMinSketchWithHeapAccumulator { } fn serialize_to_bytes(&self) -> Vec { - self.inner.serialize_msgpack() + self.inner.serialize_msgpack().unwrap_or_default() } } @@ -219,12 +217,12 @@ impl MergeableAccumulator for CountMinSketchW if accumulators.is_empty() { return Err("No accumulators to merge".into()); } - let inners: Vec = - accumulators.into_iter().map(|acc| acc.inner).collect(); - let merged_inner = CountMinSketchWithHeap::merge(inners)?; - Ok(Self { - inner: merged_inner, - }) + let mut iter = accumulators.into_iter(); + let mut merged = iter.next().unwrap(); + for acc in iter { + merged.inner.merge(&acc.inner)?; + } + Ok(merged) } } @@ -235,8 +233,8 @@ mod tests { #[test] fn test_count_min_sketch_with_heap_creation() { let cms = CountMinSketchWithHeapAccumulator::new(4, 1000, 20); - assert_eq!(cms.inner.row_num, 4); - assert_eq!(cms.inner.col_num, 1000); + assert_eq!(cms.inner.rows(), 4); + assert_eq!(cms.inner.cols(), 1000); assert_eq!(cms.inner.heap_size, 20); assert_eq!(cms.inner.topk_heap_items().len(), 0); } @@ -259,11 +257,11 @@ mod tests { vec![0.0, 20.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], ]; let heap1 = vec![ - HeapItem { + CmsHeapItem { key: "key1".to_string(), value: 100.0, }, - HeapItem { + CmsHeapItem { key: "key2".to_string(), value: 50.0, }, @@ -273,11 +271,11 @@ mod tests { vec![0.0, 15.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], ]; let heap2 = vec![ - HeapItem { + CmsHeapItem { key: "key3".to_string(), value: 75.0, }, - HeapItem { + CmsHeapItem { key: "key1".to_string(), value: 80.0, }, @@ -305,8 +303,8 @@ mod tests { let result = CountMinSketchWithHeapAccumulator::merge_accumulators(vec![cms.clone()]); assert!(result.is_ok()); let merged = result.unwrap(); - assert_eq!(merged.inner.row_num, cms.inner.row_num); - assert_eq!(merged.inner.col_num, cms.inner.col_num); + assert_eq!(merged.inner.rows(), cms.inner.rows()); + assert_eq!(merged.inner.cols(), cms.inner.cols()); assert_eq!(merged.inner.heap_size, cms.inner.heap_size); } @@ -316,17 +314,14 @@ mod tests { let cms2 = CountMinSketchWithHeapAccumulator::new(3, 10, 5); let result = CountMinSketchWithHeapAccumulator::merge_accumulators(vec![cms1, cms2]); assert!(result.is_err()); - assert!(result - .unwrap_err() - .to_string() - .contains("different dimensions")); + assert!(result.unwrap_err().to_string().contains("dimension")); } #[test] fn test_count_min_sketch_with_heap_serialization() { // Use from_legacy_matrix for a controlled state that round-trips correctly with both backends. let sketch = vec![vec![0.0, 42.0, 0.0], vec![0.0, 0.0, 100.0]]; - let topk_heap = vec![HeapItem { + let topk_heap = vec![CmsHeapItem { key: "test_key".to_string(), value: 99.0, }]; @@ -338,8 +333,8 @@ mod tests { let deserialized = CountMinSketchWithHeapAccumulator::deserialize_from_bytes_arroyo(&bytes).unwrap(); - assert_eq!(deserialized.inner.row_num, 2); - assert_eq!(deserialized.inner.col_num, 3); + assert_eq!(deserialized.inner.rows(), 2); + assert_eq!(deserialized.inner.cols(), 3); assert_eq!(deserialized.inner.heap_size, 5); assert_eq!(deserialized.inner.sketch_matrix()[0][1], 42.0); // [1][2] may be 100 (legacy, no hash collision) or 199 (100+99 when test_key hashes there) diff --git a/asap-query-engine/src/precompute_operators/count_sketch_accumulator.rs b/asap-query-engine/src/precompute_operators/count_sketch_accumulator.rs index 76f09413..00238eb6 100644 --- a/asap-query-engine/src/precompute_operators/count_sketch_accumulator.rs +++ b/asap-query-engine/src/precompute_operators/count_sketch_accumulator.rs @@ -1,4 +1,4 @@ -//! Count Sketch accumulator — wraps `asap_sketchlib::asap::count_sketch::CountSketch`. +//! Count Sketch accumulator — wraps `asap_sketchlib::sketches::count::CountSketch`. //! //! This is the concrete accumulator reached from the modified-OTLP //! `Metric.data = CountSketch{…}` hot path (PR C-CountSketch). Its @@ -20,7 +20,7 @@ //! round-trip works end-to-end without that richer query surface. use crate::data_model::{AggregateCore, AggregationType, KeyByLabelValues, SerializableToSink}; -use asap_sketchlib::asap::count_sketch::{CountSketch, CountSketchDelta}; +use asap_sketchlib::sketches::count::{CountSketch, CountSketchDelta}; use serde_json::Value; use std::collections::HashMap; @@ -194,14 +194,14 @@ impl CountSketchAccumulator { impl SerializableToSink for CountSketchAccumulator { fn serialize_to_json(&self) -> Value { serde_json::json!({ - "row_num": self.inner.row_num, - "col_num": self.inner.col_num, + "row_num": self.inner.rows, + "col_num": self.inner.cols, "sketch": self.inner.sketch(), }) } fn serialize_to_bytes(&self) -> Vec { - self.inner.serialize_msgpack() + self.inner.serialize_msgpack().unwrap_or_default() } } @@ -266,7 +266,7 @@ impl AggregateCore for CountSketchAccumulator { // but not a heavy-hitter answer). Hash compatibility note: // this relies on the agent and backend using the // sketchlib HashSpec; sketchlib-go's `portableHashSpec` - // is the canonical seed list, and `asap_sketchlib::asap::CountSketch` + // is the canonical seed list, and `asap_sketchlib::sketches::CountSketch` // hashes against the same spec. match statistic { Statistic::Topk | Statistic::Count => { @@ -302,7 +302,7 @@ impl AggregateCore for CountSketchAccumulator { /// /// Hash compatibility with the agent is via the sketchlib hash /// spec; the agent's `sketchlib-go::CountSketch` and the -/// backend's `asap_sketchlib::asap::count_sketch::CountSketch` must use +/// backend's `asap_sketchlib::sketches::count::CountSketch` must use /// the same seed list (sketchlib's `portableHashSpec` / /// `default_hash_spec`). fn count_sketch_query_key(matrix: &Vec>, key: &str) -> f64 { @@ -507,10 +507,10 @@ mod tests { 2, 3, ); - let bytes = original.serialize_msgpack(); + let bytes = original.serialize_msgpack().unwrap(); let acc = CountSketchAccumulator::from_msgpack_bytes(&bytes).expect("decode ok"); - assert_eq!(acc.inner.row_num, 2); - assert_eq!(acc.inner.col_num, 3); + assert_eq!(acc.inner.rows, 2); + assert_eq!(acc.inner.cols, 3); assert_eq!(acc.inner.sketch(), original.sketch()); } diff --git a/asap-query-engine/src/precompute_operators/datasketches_kll_accumulator.rs b/asap-query-engine/src/precompute_operators/datasketches_kll_accumulator.rs index fb7ffbb9..d4abf23a 100644 --- a/asap-query-engine/src/precompute_operators/datasketches_kll_accumulator.rs +++ b/asap-query-engine/src/precompute_operators/datasketches_kll_accumulator.rs @@ -2,7 +2,7 @@ use crate::data_model::{ AggregateCore, AggregationType, AuxStats, MergeableAccumulator, SerializableToSink, SingleSubpopulationAggregate, }; -use asap_sketchlib::asap::kll::KllSketch; +use asap_sketchlib::sketches::kll::KllSketch; use base64::{engine::general_purpose, Engine as _}; use serde_json::Value; use std::collections::HashMap; @@ -12,9 +12,9 @@ use tracing::debug; use promql_utilities::query_logics::enums::Statistic; -/// KLL sketch accumulator — wraps asap_sketchlib::asap::KllSketch. -/// Core struct, update/merge/serde logic live in sketch-core. -/// This file retains QE-specific trait impls, legacy deserializers, and JSON output. +/// KLL sketch accumulator — wraps asap_sketchlib::sketches::KllSketch. +/// Core struct, update/merge/serde logic live in `asap_sketchlib::sketches`. +/// This file retains QE-specific trait impls and JSON output. pub struct DatasketchesKLLAccumulator { pub inner: KllSketch, } @@ -31,42 +31,19 @@ impl DatasketchesKLLAccumulator { } pub fn get_quantile(&self, quantile: f64) -> f64 { - self.inner.get_quantile(quantile) - } - - pub fn deserialize_from_json(data: &Value) -> Result> { - // Mirror Python implementation: expects {"sketch": base64_encoded_string} - let sketch_b64 = data["sketch"] - .as_str() - .ok_or("Missing or invalid 'sketch' field")?; - - let sketch_bytes = general_purpose::STANDARD - .decode(sketch_b64) - .map_err(|e| format!("Failed to decode base64 sketch data: {e}"))?; - - // TODO: remove this hardcoding once FlinkSketch serializes k in its output - Ok(Self { - inner: KllSketch::from_dsrs_bytes(&sketch_bytes, 200)?, - }) - } - - pub fn deserialize_from_bytes(buffer: &[u8]) -> Result> { - // Mirror Python implementation: deserialize sketch directly from bytes - // TODO: remove this hardcoding once FlinkSketch serializes k in its output - Ok(Self { - inner: KllSketch::from_dsrs_bytes(buffer, 200)?, - }) + self.inner.quantile(quantile) } pub fn deserialize_from_bytes_arroyo( buffer: &[u8], ) -> Result> { debug!( - "Deserializing DatasketchesKLLAccumulator from Arroyo MessagePack buffer of size {}", + "Deserializing DatasketchesKLLAccumulator from MessagePack buffer of size {}", buffer.len() ); Ok(Self { - inner: KllSketch::deserialize_msgpack(buffer)?, + inner: KllSketch::deserialize_msgpack(buffer) + .map_err(|e| -> Box { e.to_string().into() })?, }) } @@ -81,7 +58,8 @@ impl DatasketchesKLLAccumulator { /// serializes its full internal state to msgpack. pub fn from_msgpack_bytes(buffer: &[u8]) -> Result> { Ok(Self { - inner: KllSketch::deserialize_msgpack(buffer)?, + inner: KllSketch::deserialize_msgpack(buffer) + .map_err(|e| -> Box { e.to_string().into() })?, }) } @@ -241,7 +219,7 @@ impl SerializableToSink for DatasketchesKLLAccumulator { } fn serialize_to_bytes(&self) -> Vec { - self.inner.serialize_msgpack() + self.inner.serialize_msgpack().unwrap_or_default() } } @@ -385,11 +363,12 @@ impl MergeableAccumulator for DatasketchesKLLAccumul if accumulators.is_empty() { return Err("No accumulators to merge".into()); } - let inners: Vec = accumulators.into_iter().map(|acc| acc.inner).collect(); - let merged_inner = KllSketch::merge(inners)?; - Ok(Self { - inner: merged_inner, - }) + let mut iter = accumulators.into_iter(); + let mut merged = iter.next().unwrap(); + for acc in iter { + merged.inner.merge(&acc.inner)?; + } + Ok(merged) } } diff --git a/asap-query-engine/src/precompute_operators/dd_sketch_accumulator.rs b/asap-query-engine/src/precompute_operators/dd_sketch_accumulator.rs index c0cee54b..f0203bc4 100644 --- a/asap-query-engine/src/precompute_operators/dd_sketch_accumulator.rs +++ b/asap-query-engine/src/precompute_operators/dd_sketch_accumulator.rs @@ -1,4 +1,4 @@ -//! DDSketch accumulator — wraps `asap_sketchlib::asap::dd_sketch::DdSketch`. +//! DDSketch accumulator — wraps `asap_sketchlib::sketches::ddsketch::DdSketch`. //! //! Concrete accumulator reached from the modified-OTLP //! `Metric.data = DDSketch{…}` hot path (PR C-CountSketch follow-up). @@ -12,7 +12,7 @@ //! works end-to-end without that richer query surface. use crate::data_model::{AggregateCore, AggregationType, KeyByLabelValues, SerializableToSink}; -use asap_sketchlib::asap::dd_sketch::{DdSketch, DdSketchDelta}; +use asap_sketchlib::sketches::ddsketch::{DdSketch, DdSketchDelta}; use serde_json::Value; use std::collections::HashMap; @@ -142,7 +142,7 @@ impl SerializableToSink for DDSketchAccumulator { } fn serialize_to_bytes(&self) -> Vec { - self.inner.serialize_msgpack() + self.inner.serialize_msgpack().unwrap_or_default() } } @@ -354,7 +354,7 @@ mod tests { #[test] fn test_from_msgpack_bytes_round_trip() { let original = DdSketch::from_raw(0.01, vec![5, 10, 15, 20], -2, 50, 150.0, 0.25, 8.0); - let bytes = original.serialize_msgpack(); + let bytes = original.serialize_msgpack().unwrap(); let acc = DDSketchAccumulator::from_msgpack_bytes(&bytes).expect("decode ok"); assert_eq!(acc.inner.alpha, 0.01); assert_eq!(acc.inner.store_counts, vec![5, 10, 15, 20]); diff --git a/asap-query-engine/src/precompute_operators/delta_set_aggregator_accumulator.rs b/asap-query-engine/src/precompute_operators/delta_set_aggregator_accumulator.rs index 3825bfcb..18360be4 100644 --- a/asap-query-engine/src/precompute_operators/delta_set_aggregator_accumulator.rs +++ b/asap-query-engine/src/precompute_operators/delta_set_aggregator_accumulator.rs @@ -2,7 +2,7 @@ use crate::data_model::{ AggregateCore, AggregationType, KeyByLabelValues, MergeableAccumulator, MultipleSubpopulationAggregate, SerializableToSink, }; -use asap_sketchlib::asap::delta_set_aggregator::{deserialize_msgpack, serialize_msgpack}; +use asap_sketchlib::sketches::delta_set_aggregator::{deserialize_msgpack, serialize_msgpack}; use serde_json::Value; use std::collections::{HashMap, HashSet}; @@ -10,7 +10,7 @@ use promql_utilities::query_logics::enums::Statistic; /// Accumulator that tracks sets of added and removed keys. /// Used for delta aggregation to track changes in cardinality. -/// Wire format (DeltaResult) and msgpack serde live in sketch-core. +/// Wire format (DeltaResult) and msgpack serde live in `asap_sketchlib::sketches`. #[derive(Debug, Clone)] pub struct DeltaSetAggregatorAccumulator { pub added: HashSet, @@ -153,7 +153,8 @@ impl DeltaSetAggregatorAccumulator { buffer: &[u8], ) -> Result> { // Delegate to sketch-core canonical DeltaResult msgpack format - let delta = deserialize_msgpack(buffer)?; + let delta = deserialize_msgpack(buffer) + .map_err(|e| -> Box { e.to_string().into() })?; let mut added = HashSet::new(); for item in &delta.added { @@ -202,7 +203,7 @@ impl SerializableToSink for DeltaSetAggregatorAccumulator { .iter() .map(|key| key.to_semicolon_str()) .collect(); - serialize_msgpack(&added, &removed) + serialize_msgpack(&added, &removed).unwrap_or_default() } } diff --git a/asap-query-engine/src/precompute_operators/hll_sketch_accumulator.rs b/asap-query-engine/src/precompute_operators/hll_sketch_accumulator.rs index 0e46ec77..aa5cf48a 100644 --- a/asap-query-engine/src/precompute_operators/hll_sketch_accumulator.rs +++ b/asap-query-engine/src/precompute_operators/hll_sketch_accumulator.rs @@ -1,4 +1,4 @@ -//! HLL accumulator — wraps `asap_sketchlib::asap::hll_sketch::HllSketch`. +//! HLL accumulator — wraps `asap_sketchlib::sketches::hll::HllSketch`. //! //! Concrete accumulator reached from the modified-OTLP //! `Metric.data = HLLSketch{…}` hot path (PR C-CountSketch follow-up). @@ -12,7 +12,7 @@ //! store round-trip works end-to-end without that richer query surface. use crate::data_model::{AggregateCore, AggregationType, KeyByLabelValues, SerializableToSink}; -use asap_sketchlib::asap::hll_sketch::{HllDelta, HllSketch, HllVariant}; +use asap_sketchlib::sketches::hll::{HllSketch, HllSketchDelta, HllVariant}; use serde_json::Value; use std::collections::HashMap; @@ -131,7 +131,7 @@ impl HllSketchAccumulator { .into_iter() .map(|u| (u.index, u.value as u8)) .collect(); - let delta = HllDelta { updates }; + let delta = HllSketchDelta { updates }; self.inner .apply_delta(&delta) .map_err(|e| format!("apply HLLDelta: {e}"))?; @@ -152,7 +152,7 @@ impl SerializableToSink for HllSketchAccumulator { } fn serialize_to_bytes(&self) -> Vec { - self.inner.serialize_msgpack() + self.inner.serialize_msgpack().unwrap_or_default() } } @@ -234,7 +234,7 @@ impl AggregateCore for HllSketchAccumulator { /// (linear-counting) and large-range (32-bit space) corrections /// from the original Flajolet et al. paper. /// -/// Inlined here rather than added as a method on `asap_sketchlib::asap::HllSketch` +/// Inlined here rather than added as a method on `asap_sketchlib::sketches::HllSketch` /// because the existing `asap_sketchlib::asap` types only expose merge / /// serialize today; adding a query method there would force a /// cross-crate change. @@ -441,7 +441,7 @@ mod tests { 2.5, 42.0, ); - let bytes = original.serialize_msgpack(); + let bytes = original.serialize_msgpack().unwrap(); let acc = HllSketchAccumulator::from_msgpack_bytes(&bytes).expect("decode ok"); assert_eq!(acc.inner.variant, HllVariant::Hip); assert_eq!(acc.inner.precision, 3); diff --git a/asap-query-engine/src/precompute_operators/hydra_kll_accumulator.rs b/asap-query-engine/src/precompute_operators/hydra_kll_accumulator.rs index f9ed22bc..38b827ad 100644 --- a/asap-query-engine/src/precompute_operators/hydra_kll_accumulator.rs +++ b/asap-query-engine/src/precompute_operators/hydra_kll_accumulator.rs @@ -5,14 +5,14 @@ use crate::{ }, KeyByLabelValues, }; -use asap_sketchlib::asap::hydra_kll::HydraKllSketch; +use asap_sketchlib::sketches::hydra_kll::HydraKllSketch; use base64::{engine::general_purpose, Engine as _}; use std::collections::HashMap; use promql_utilities::query_logics::enums::Statistic; -/// HydraKLL sketch accumulator — wraps asap_sketchlib::asap::HydraKllSketch. -/// Core struct, update/merge/serde logic live in sketch-core. +/// HydraKLL sketch accumulator — wraps asap_sketchlib::sketches::HydraKllSketch. +/// Core struct, update/merge/serde logic live in `asap_sketchlib::sketches`. /// This file retains QE-specific trait impls and JSON output. #[derive(Debug, Clone)] pub struct HydraKllSketchAccumulator { @@ -38,25 +38,26 @@ impl HydraKllSketchAccumulator { buffer: &[u8], ) -> Result> { Ok(Self { - inner: HydraKllSketch::deserialize_msgpack(buffer)?, + inner: HydraKllSketch::deserialize_msgpack(buffer) + .map_err(|e| -> Box { e.to_string().into() })?, }) } pub fn query_key(&self, key: &KeyByLabelValues, quantile: f64) -> f64 { - self.inner.query(&key.to_semicolon_str(), quantile) + self.inner.quantile(&key.to_semicolon_str(), quantile) } } impl SerializableToSink for HydraKllSketchAccumulator { fn serialize_to_json(&self) -> serde_json::Value { // Mirror Python implementation: {"sketch": base64_encoded_string} - let sketch_bytes = self.inner.serialize_msgpack(); + let sketch_bytes = self.inner.serialize_msgpack().unwrap_or_default(); let sketch_b64 = general_purpose::STANDARD.encode(&sketch_bytes); serde_json::json!({ "sketch": sketch_b64 }) } fn serialize_to_bytes(&self) -> Vec { - self.inner.serialize_msgpack() + self.inner.serialize_msgpack().unwrap_or_default() } } @@ -67,11 +68,12 @@ impl MergeableAccumulator for HydraKllSketchAccumulat if accumulators.is_empty() { return Err("No accumulators to merge".into()); } - let inners: Vec = accumulators.into_iter().map(|acc| acc.inner).collect(); - let merged_inner = HydraKllSketch::merge(inners)?; - Ok(Self { - inner: merged_inner, - }) + let mut iter = accumulators.into_iter(); + let mut merged = iter.next().unwrap(); + for acc in iter { + merged.inner.merge(&acc.inner)?; + } + Ok(merged) } } diff --git a/asap-query-engine/src/precompute_operators/increase_accumulator.rs b/asap-query-engine/src/precompute_operators/increase_accumulator.rs index 0a912563..7d804e8f 100644 --- a/asap-query-engine/src/precompute_operators/increase_accumulator.rs +++ b/asap-query-engine/src/precompute_operators/increase_accumulator.rs @@ -216,11 +216,8 @@ impl AggregateCore for IncreaseAccumulator { self } - fn as_any_mut(&mut self) -> &mut dyn std::any::Any { - self - } fn merge_with( diff --git a/asap-query-engine/src/precompute_operators/min_max_accumulator.rs b/asap-query-engine/src/precompute_operators/min_max_accumulator.rs index 744ea1e4..957e9c80 100644 --- a/asap-query-engine/src/precompute_operators/min_max_accumulator.rs +++ b/asap-query-engine/src/precompute_operators/min_max_accumulator.rs @@ -155,11 +155,8 @@ impl AggregateCore for MinMaxAccumulator { self } - fn as_any_mut(&mut self) -> &mut dyn std::any::Any { - self - } fn merge_with( diff --git a/asap-query-engine/src/precompute_operators/multiple_increase_accumulator.rs b/asap-query-engine/src/precompute_operators/multiple_increase_accumulator.rs index 87dace44..0cb8abd0 100644 --- a/asap-query-engine/src/precompute_operators/multiple_increase_accumulator.rs +++ b/asap-query-engine/src/precompute_operators/multiple_increase_accumulator.rs @@ -238,11 +238,8 @@ impl AggregateCore for MultipleIncreaseAccumulator { self } - fn as_any_mut(&mut self) -> &mut dyn std::any::Any { - self - } fn merge_with( diff --git a/asap-query-engine/src/precompute_operators/multiple_min_max_accumulator.rs b/asap-query-engine/src/precompute_operators/multiple_min_max_accumulator.rs index 0687930f..90f0391e 100644 --- a/asap-query-engine/src/precompute_operators/multiple_min_max_accumulator.rs +++ b/asap-query-engine/src/precompute_operators/multiple_min_max_accumulator.rs @@ -211,11 +211,8 @@ impl AggregateCore for MultipleMinMaxAccumulator { self } - fn as_any_mut(&mut self) -> &mut dyn std::any::Any { - self - } fn merge_with( diff --git a/asap-query-engine/src/precompute_operators/multiple_sum_accumulator.rs b/asap-query-engine/src/precompute_operators/multiple_sum_accumulator.rs index 5c83f17d..5f565356 100644 --- a/asap-query-engine/src/precompute_operators/multiple_sum_accumulator.rs +++ b/asap-query-engine/src/precompute_operators/multiple_sum_accumulator.rs @@ -201,11 +201,8 @@ impl AggregateCore for MultipleSumAccumulator { self } - fn as_any_mut(&mut self) -> &mut dyn std::any::Any { - self - } fn merge_with( diff --git a/asap-query-engine/src/precompute_operators/set_aggregator_accumulator.rs b/asap-query-engine/src/precompute_operators/set_aggregator_accumulator.rs index 4baa69d5..73745220 100644 --- a/asap-query-engine/src/precompute_operators/set_aggregator_accumulator.rs +++ b/asap-query-engine/src/precompute_operators/set_aggregator_accumulator.rs @@ -2,14 +2,14 @@ use crate::data_model::{ AggregateCore, AggregationType, KeyByLabelValues, MergeableAccumulator, MultipleSubpopulationAggregate, SerializableToSink, }; -use asap_sketchlib::asap::set_aggregator::SetAggregator; +use asap_sketchlib::sketches::set_aggregator::SetAggregator; use serde_json::Value; use std::collections::{HashMap, HashSet}; use promql_utilities::query_logics::enums::Statistic; -/// Set aggregator accumulator — wraps asap_sketchlib::asap::SetAggregator. -/// Core struct, merge/serde logic live in sketch-core. +/// Set aggregator accumulator — wraps asap_sketchlib::sketches::SetAggregator. +/// Core struct, merge/serde logic live in `asap_sketchlib::sketches`. /// This file retains QE-specific trait impls, KeyByLabelValues conversion, /// and legacy deserializers. #[derive(Debug, Clone)] @@ -92,7 +92,8 @@ impl SetAggregatorAccumulator { pub fn deserialize_from_bytes_arroyo( buffer: &[u8], ) -> Result> { - let sa = SetAggregator::deserialize_msgpack(buffer)?; + let sa = SetAggregator::deserialize_msgpack(buffer) + .map_err(|e| -> Box { e.to_string().into() })?; let added = sa .values .into_iter() @@ -106,9 +107,9 @@ impl SetAggregatorAccumulator { pub fn serialize_to_bytes_arroyo(&self) -> Vec { let mut sa = SetAggregator::new(); for key in &self.added { - sa.insert(&key.to_semicolon_str()); + sa.update(&key.to_semicolon_str()); } - sa.serialize_msgpack() + sa.serialize_msgpack().unwrap_or_default() } } diff --git a/asap-query-engine/src/precompute_operators/sketch_envelope_accumulator.rs b/asap-query-engine/src/precompute_operators/sketch_envelope_accumulator.rs index 4b7f846c..07fbf575 100644 --- a/asap-query-engine/src/precompute_operators/sketch_envelope_accumulator.rs +++ b/asap-query-engine/src/precompute_operators/sketch_envelope_accumulator.rs @@ -93,11 +93,8 @@ impl AggregateCore for SketchEnvelopeAccumulator { self } - fn as_any_mut(&mut self) -> &mut dyn std::any::Any { - self - } fn merge_with( diff --git a/asap-query-engine/src/precompute_operators/sum_accumulator.rs b/asap-query-engine/src/precompute_operators/sum_accumulator.rs index f715e45b..696b66d9 100644 --- a/asap-query-engine/src/precompute_operators/sum_accumulator.rs +++ b/asap-query-engine/src/precompute_operators/sum_accumulator.rs @@ -89,11 +89,8 @@ impl AggregateCore for SumAccumulator { self } - fn as_any_mut(&mut self) -> &mut dyn std::any::Any { - self - } fn merge_with( diff --git a/asap-query-engine/tests/e2e_modified_otlp_sketch_path.rs b/asap-query-engine/tests/e2e_modified_otlp_sketch_path.rs index 68a0ed9a..338145d9 100644 --- a/asap-query-engine/tests/e2e_modified_otlp_sketch_path.rs +++ b/asap-query-engine/tests/e2e_modified_otlp_sketch_path.rs @@ -1192,11 +1192,11 @@ async fn e2e_count_min_sketch_msgpack_modified_otlp_path() { // is what the Go producer (sketchlib-go) will emit once PR I's matching // Go-side work lands. let mut cms = - asap_sketchlib::asap::count_min::CountMinSketch::new(rows as usize, cols as usize); + asap_sketchlib::sketches::countmin::CountMinSketch::new(rows as usize, cols as usize); cms.update("user_a", 1.0); cms.update("user_b", 1.0); cms.update("user_a", 1.0); - let sketch_bytes = cms.serialize_msgpack(); + let sketch_bytes = cms.serialize_msgpack().expect("serialize CMS msgpack"); let client = reqwest::Client::new(); let req = build_count_min_msgpack_export_request( @@ -1208,12 +1208,15 @@ async fn e2e_count_min_sketch_msgpack_modified_otlp_path() { post_otlp_http(&client, otlp_http_port, req).await; // Watermark advance using an empty msgpack sketch. - let empty = asap_sketchlib::asap::count_min::CountMinSketch::new(rows as usize, cols as usize); + let empty = + asap_sketchlib::sketches::countmin::CountMinSketch::new(rows as usize, cols as usize); let watermark_req = build_count_min_msgpack_export_request( metric_name, service_label, 2_000_000_000, - empty.serialize_msgpack(), + empty + .serialize_msgpack() + .expect("serialize empty CMS msgpack"), ); post_otlp_http(&client, otlp_http_port, watermark_req).await; @@ -1234,10 +1237,10 @@ async fn e2e_count_min_sketch_msgpack_modified_otlp_path() { .downcast_ref::() .expect("captured accumulator should be CountMinSketchAccumulator"); - // user_a was updated twice → query_key("user_a") should estimate ≥ 2. + // user_a was updated twice → estimate("user_a") should be ≥ 2. assert!( - cms_acc.inner.query_key("user_a") >= 2.0, + cms_acc.inner.estimate("user_a") >= 2.0, "CountMinSketch msgpack round-trip should preserve user_a count (got {})", - cms_acc.inner.query_key("user_a") + cms_acc.inner.estimate("user_a") ); } diff --git a/asap-query-engine/tests/e2e_precompute_equivalence.rs b/asap-query-engine/tests/e2e_precompute_equivalence.rs index aa894b86..55439e89 100644 --- a/asap-query-engine/tests/e2e_precompute_equivalence.rs +++ b/asap-query-engine/tests/e2e_precompute_equivalence.rs @@ -1,13 +1,13 @@ //! End-to-end integration tests: precompute engine output equivalence -//! with ArroYo sketch format. +//! with the wire-format sketch encoding. //! //! Each test: //! 1. Starts a PrecomputeEngine backed by a CapturingOutputSink //! 2. Sends Prometheus remote write samples via HTTP (Snappy-compressed protobuf) //! 3. Advances the watermark past the window boundary to close it -//! 4. Drains captured outputs and verifies equivalence with ArroYo-format accumulators +//! 4. Drains captured outputs and verifies equivalence with wire-format accumulators -use asap_sketchlib::asap::kll::KllSketch; +use asap_sketchlib::sketches::kll::KllSketch; use asap_types::aggregation_config::AggregationConfig; use asap_types::enums::{AggregationType, WindowType}; use flate2::{write::GzEncoder, Compression}; @@ -162,10 +162,10 @@ fn gzip_hex(bytes: &[u8]) -> String { hex::encode(encoder.finish().unwrap()) } -// ─── test 1: DatasketchesKLL output matches ArroYo KLL ────────────────────── +// ─── test 1: DatasketchesKLL output matches wire-format KLL ───────────────── /// Full e2e: send KLL samples through the HTTP ingest → PrecomputeEngine stack, -/// then verify the emitted DatasketchesKLLAccumulator matches what ArroYo's +/// then verify the emitted DatasketchesKLLAccumulator matches what the wire-format /// KllSketch::aggregate_kll would produce for the same values. #[tokio::test] async fn e2e_kll_output_matches_arroyo() { @@ -242,7 +242,7 @@ async fn e2e_kll_output_matches_arroyo() { .downcast_ref::() .expect("captured accumulator should be DatasketchesKLLAccumulator"); - // Build the ArroYo-format equivalent and deserialize it + // Build the wire-format equivalent and deserialize it let arroyo_bytes = KllSketch::aggregate_kll(k, &values).expect("KllSketch::aggregate_kll failed"); let arroyo_json = json!({ @@ -254,11 +254,11 @@ async fn e2e_kll_output_matches_arroyo() { let streaming_config_for_deser = StreamingConfig::new(agg_map); let (_arroyo_output, arroyo_acc_box) = PrecomputedOutput::deserialize_from_json_arroyo(&arroyo_json, &streaming_config_for_deser) - .expect("ArroYo KLL deserialization failed"); + .expect("wire-format KLL deserialization failed"); let arroyo_acc = arroyo_acc_box .as_any() .downcast_ref::() - .expect("ArroYo payload should deserialize to DatasketchesKLLAccumulator"); + .expect("wire-format payload should deserialize to DatasketchesKLLAccumulator"); // Window metadata assert_eq!(handcrafted_output.aggregation_id, agg_id); @@ -284,11 +284,11 @@ async fn e2e_kll_output_matches_arroyo() { } } -// ─── test 2: MultipleSum output matches ArroYo MultipleSum ────────────────── +// ─── test 2: MultipleSum output matches wire-format MultipleSum ───────────── /// Full e2e: send MultipleSum samples (grouped by "host") through the HTTP /// ingest → PrecomputeEngine stack, then verify the emitted -/// MultipleSumAccumulator matches the ArroYo MessagePack-encoded sums map. +/// MultipleSumAccumulator matches the wire-format MessagePack-encoded sums map. #[tokio::test] async fn e2e_multiple_sum_output_matches_arroyo() { let port = 19401u16; @@ -356,7 +356,7 @@ async fn e2e_multiple_sum_output_matches_arroyo() { .downcast_ref::() .expect("captured accumulator should be MultipleSumAccumulator"); - // Build the ArroYo-format equivalent and deserialize it + // Build the wire-format equivalent and deserialize it let mut expected_sums: HashMap = HashMap::new(); expected_sums.insert("A".to_string(), 6.0); let arroyo_bytes = rmp_serde::to_vec(&expected_sums).expect("msgpack encoding failed"); @@ -369,11 +369,11 @@ async fn e2e_multiple_sum_output_matches_arroyo() { let streaming_config_for_deser = StreamingConfig::new(agg_map); let (_arroyo_output, arroyo_acc_box) = PrecomputedOutput::deserialize_from_json_arroyo(&arroyo_json, &streaming_config_for_deser) - .expect("ArroYo MultipleSum deserialization failed"); + .expect("wire-format MultipleSum deserialization failed"); let arroyo_acc = arroyo_acc_box .as_any() .downcast_ref::() - .expect("ArroYo payload should deserialize to MultipleSumAccumulator"); + .expect("wire-format payload should deserialize to MultipleSumAccumulator"); // Window metadata assert_eq!(handcrafted_output.aggregation_id, agg_id); diff --git a/asap-query-engine/tests/test_both_backends.rs b/asap-query-engine/tests/test_both_backends.rs deleted file mode 100644 index 56437562..00000000 --- a/asap-query-engine/tests/test_both_backends.rs +++ /dev/null @@ -1,30 +0,0 @@ -//! Integration test that runs the library test suite with the sketchlib backend. -//! -//! When you run `cargo test -p query_engine_rust` (without --features sketchlib-tests), -//! the lib tests run with the legacy backend. This test spawns a second run with the -//! sketchlib backend so both modes are exercised in one `cargo test` invocation. -//! -//! This test is only compiled when sketchlib-tests is NOT enabled, to avoid recursion. - -#[cfg(not(feature = "sketchlib-tests"))] -#[test] -fn test_sketchlib_backend() { - use std::process::Command; - - let status = Command::new(env!("CARGO")) - .args([ - "test", - "-p", - "query_engine_rust", - "--lib", - "--features", - "sketchlib-tests", - ]) - .status() - .expect("failed to spawn cargo test"); - - assert!( - status.success(), - "sketchlib backend tests failed (run `cargo test -p query_engine_rust --lib --features sketchlib-tests` for details)" - ); -}