From ff8aebb1e77d9541ce4ec6e1e6383b668040c23d Mon Sep 17 00:00:00 2001 From: zz_y Date: Sun, 10 May 2026 13:17:34 -0600 Subject: [PATCH] cleanup: remove datafusion + integrate SketchIndex into sketch_db MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The datafusion-backed query path (engines/physical, engines/logical, the datafusion_summary_library crate, the parallel datafusion-LogicalPlan adapter, all DF-specific tests, and the controller-side query_language adapter) is redundant: the controller crate already owns logical and physical plan layers (controller/src/algebra/{lower,physical,plan}.rs + intent_algebra/lower.rs + language_logical_plan/). Keeping the parallel DF stack alive cost an extra translation hop, ~5800 lines of code, and a test surface that papered over the divergence. Companion structural cleanup: relocate the new sid-keyed SketchIndex files (added in #121 / #122) into stores/sketch_db/, the canonical home for storage. They were floating at stores/ top level after #121. ## What was deleted - crates/datafusion_summary_library/ (8 files — sketch-aware datafusion ExecutionPlan operators) - asap-query-engine/src/engines/physical/ (7 files — datafusion-LogicalPlan → ExecutionPlan adapters) - asap-query-engine/src/engines/logical/ (2 files — plan_builder for the DF logical layer) - asap-query-engine/src/tests/datafusion/ (11 test files exercising the deleted engines) - controller/src/query_language/datafusion/ (small adapter) - datafusion = "43" + datafusion_summary_library deps from asap-query-engine/Cargo.toml + workspace Cargo.toml. ## What was moved - asap-query-engine/src/stores/sketch_index.rs → asap-query-engine/src/stores/sketch_db/sketch_index.rs - asap-query-engine/src/stores/epoch_columnar.rs → asap-query-engine/src/stores/sketch_db/epoch_columnar.rs - All cross-file imports rewritten to the new sketch_db:: paths. ## In-source surgery - engines/simple/engine.rs: deleted the DF-importing methods (SessionContext, physical_plan::collect, record_batch_to_result_map, engines::logical::plan_builder::build_*). Phase 5 SketchIndex warm-tier hook (with_sketch_index + classify branch in execute()) preserved verbatim, only its imports were rewritten. - stores/sketch_db/simple_map_store/per_key.rs: query_disk_parts and EpochSource::snapshot_sealed_epoch now panic with "datafusion-dependent path removed; ingest/persistence still under refactor". Both fire only on persistence-enabled runs (agreed-upon breakage). Marked // TODO: replace with non-datafusion path. - 8 controller files: deleted DF-referencing variants/match arms/trait impls. ## Verification - cargo build --release -p query_engine_rust: clean (3 dead-code warnings) - cargo build --release -p controller: clean (5 unused-import warnings on pre-existing items) - cargo test --release -p query_engine_rust --lib: 792 passed, 2 failed (schema_timeline_dispatch_tests::*, both pre-existing flakes; 8 persistence tests skipped — they fire the TODO panics) - grep for "datafusion" in the tree: 6 lines remain, all in TODO comments and panic messages. No imports, deps, or types. ## Follow-ups - Replace persistence TODO panics with a non-datafusion warm-tier serializer (likely the asap_sketchlib proto codec). - The per-Capability sketch reducer (warm-tier query evaluator) is in flight in a parallel branch — replaces the all-Hit CapabilityMiss with real sketch evaluation from SketchIndex.query_range. Co-Authored-By: Claude Opus 4.7 (1M context) --- .gitignore | 2 + Cargo.lock | 959 +--------- Cargo.toml | 2 - asap-query-engine/Cargo.toml | 2 - asap-query-engine/src/drivers/ingest/otel.rs | 26 +- asap-query-engine/src/engines/logical/mod.rs | 6 - .../src/engines/logical/plan_builder.rs | 1014 ---------- asap-query-engine/src/engines/mod.rs | 2 - .../src/engines/physical/accumulator_serde.rs | 314 ---- .../src/engines/physical/conversion.rs | 452 ----- asap-query-engine/src/engines/physical/mod.rs | 29 - .../src/engines/physical/planner.rs | 180 -- .../physical/precomputed_summary_read_exec.rs | 169 -- .../engines/physical/summary_infer_exec.rs | 769 -------- .../physical/summary_merge_multiple_exec.rs | 556 ------ .../src/engines/simple/engine.rs | 295 +-- asap-query-engine/src/main.rs | 2 +- .../src/precompute_engine/engine.rs | 2 +- .../src/precompute_engine/ingest_handler.rs | 4 +- asap-query-engine/src/stores/mod.rs | 6 +- .../stores/{ => sketch_db}/epoch_columnar.rs | 0 asap-query-engine/src/stores/sketch_db/mod.rs | 2 + .../sketch_db/simple_map_store/per_key.rs | 163 +- .../stores/{ => sketch_db}/sketch_index.rs | 0 .../datafusion/accumulator_serde_tests.rs | 339 ---- .../datafusion/dispatch_arithmetic_tests.rs | 140 -- asap-query-engine/src/tests/datafusion/mod.rs | 15 - .../datafusion/plan_builder_binary_tests.rs | 190 -- .../plan_builder_regression_tests.rs | 210 --- .../plan_execution_arithmetic_tests.rs | 369 ---- .../plan_execution_dual_input_tests.rs | 312 ---- .../plan_execution_temporal_tests.rs | 542 ------ .../tests/datafusion/plan_execution_tests.rs | 587 ------ .../datafusion/structural_matching_tests.rs | 121 -- .../warm_engine_replay_regression_tests.rs | 861 --------- asap-query-engine/src/tests/mod.rs | 1 - .../tests/test_utilities/engine_factories.rs | 70 - controller/src/algebra/physical.rs | 2 +- controller/src/language_logical_plan/plan.rs | 1 - controller/src/language_logical_plan/tests.rs | 2 +- .../src/query_language/datafusion/mod.rs | 27 - controller/src/query_language/language.rs | 2 +- controller/src/query_language/language_ast.rs | 1 - controller/src/query_language/mod.rs | 3 - controller/src/query_language/tests.rs | 7 - controller/src/types_v2.rs | 11 +- crates/datafusion_summary_library/Cargo.toml | 13 - crates/datafusion_summary_library/src/lib.rs | 15 - .../src/physical/hll.rs | 169 -- .../src/physical/mod.rs | 17 - .../src/physical/planner.rs | 205 --- .../src/physical/summary_infer_exec.rs | 282 --- .../src/physical/summary_insert_exec.rs | 434 ----- .../src/sketch_operators.rs | 1630 ----------------- 54 files changed, 98 insertions(+), 11436 deletions(-) delete mode 100644 asap-query-engine/src/engines/logical/mod.rs delete mode 100644 asap-query-engine/src/engines/logical/plan_builder.rs delete mode 100644 asap-query-engine/src/engines/physical/accumulator_serde.rs delete mode 100644 asap-query-engine/src/engines/physical/conversion.rs delete mode 100644 asap-query-engine/src/engines/physical/mod.rs delete mode 100644 asap-query-engine/src/engines/physical/planner.rs delete mode 100644 asap-query-engine/src/engines/physical/precomputed_summary_read_exec.rs delete mode 100644 asap-query-engine/src/engines/physical/summary_infer_exec.rs delete mode 100644 asap-query-engine/src/engines/physical/summary_merge_multiple_exec.rs rename asap-query-engine/src/stores/{ => sketch_db}/epoch_columnar.rs (100%) rename asap-query-engine/src/stores/{ => sketch_db}/sketch_index.rs (100%) delete mode 100644 asap-query-engine/src/tests/datafusion/accumulator_serde_tests.rs delete mode 100644 asap-query-engine/src/tests/datafusion/dispatch_arithmetic_tests.rs delete mode 100644 asap-query-engine/src/tests/datafusion/mod.rs delete mode 100644 asap-query-engine/src/tests/datafusion/plan_builder_binary_tests.rs delete mode 100644 asap-query-engine/src/tests/datafusion/plan_builder_regression_tests.rs delete mode 100644 asap-query-engine/src/tests/datafusion/plan_execution_arithmetic_tests.rs delete mode 100644 asap-query-engine/src/tests/datafusion/plan_execution_dual_input_tests.rs delete mode 100644 asap-query-engine/src/tests/datafusion/plan_execution_temporal_tests.rs delete mode 100644 asap-query-engine/src/tests/datafusion/plan_execution_tests.rs delete mode 100644 asap-query-engine/src/tests/datafusion/structural_matching_tests.rs delete mode 100644 asap-query-engine/src/tests/datafusion/warm_engine_replay_regression_tests.rs delete mode 100644 controller/src/query_language/datafusion/mod.rs delete mode 100644 crates/datafusion_summary_library/Cargo.toml delete mode 100644 crates/datafusion_summary_library/src/lib.rs delete mode 100644 crates/datafusion_summary_library/src/physical/hll.rs delete mode 100644 crates/datafusion_summary_library/src/physical/mod.rs delete mode 100644 crates/datafusion_summary_library/src/physical/planner.rs delete mode 100644 crates/datafusion_summary_library/src/physical/summary_infer_exec.rs delete mode 100644 crates/datafusion_summary_library/src/physical/summary_insert_exec.rs delete mode 100644 crates/datafusion_summary_library/src/sketch_operators.rs diff --git a/.gitignore b/.gitignore index d09d69fc..4500c68c 100644 --- a/.gitignore +++ b/.gitignore @@ -8,3 +8,5 @@ preprocessed_configs/ status uuid store/ + +.claude/ diff --git a/Cargo.lock b/Cargo.lock index f323d71b..ad3e359e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -31,21 +31,6 @@ dependencies = [ "memchr", ] -[[package]] -name = "alloc-no-stdlib" -version = "2.0.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cc7bb162ec39d46ab1ca8c77bf72e890535becd1751bb45f64c597edb4c8c6b3" - -[[package]] -name = "alloc-stdlib" -version = "0.2.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "94fb8275041c72129eb51b7d0322c29b8387a0386127718b096429201a5d6ece" -dependencies = [ - "alloc-no-stdlib", -] - [[package]] name = "allocator-api2" version = "0.2.21" @@ -156,18 +141,6 @@ dependencies = [ "rustversion", ] -[[package]] -name = "arrayref" -version = "0.3.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "76a2e8124351fda1ef8aaaa3bbd7ebbcb486bbcd4225aca0aa0d84bb2db8fecb" - -[[package]] -name = "arrayvec" -version = "0.7.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7c02d123df017efcdfbd739ef81735b36c5ba83ec3c59c80a9d7ecc718f92e50" - [[package]] name = "arrow" version = "53.4.1" @@ -215,7 +188,6 @@ dependencies = [ "arrow-data", "arrow-schema", "chrono", - "chrono-tz", "half", "hashbrown 0.15.5", "num", @@ -246,7 +218,6 @@ dependencies = [ "atoi", "base64 0.22.1", "chrono", - "comfy-table", "half", "lexical-core", "num", @@ -296,7 +267,6 @@ dependencies = [ "arrow-data", "arrow-schema", "flatbuffers", - "lz4_flex", ] [[package]] @@ -429,7 +399,7 @@ dependencies = [ "serde", "serde-big-array", "smallvec", - "twox-hash 2.1.2", + "twox-hash", "xxhash-rust", ] @@ -446,24 +416,6 @@ dependencies = [ "tracing", ] -[[package]] -name = "async-compression" -version = "0.4.19" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "06575e6a9673580f52661c92107baabffbf41e2141373441cbcdc47cb733003c" -dependencies = [ - "bzip2 0.5.2", - "flate2", - "futures-core", - "futures-io", - "memchr", - "pin-project-lite", - "tokio", - "xz2", - "zstd", - "zstd-safe", -] - [[package]] name = "async-stream" version = "0.3.6" @@ -684,29 +636,6 @@ version = "2.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c4512299f36f043ab09a583e57bceb5a5aab7a73db1805848e8fef3c9e8c78b3" -[[package]] -name = "blake2" -version = "0.10.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "46502ad458c9a52b69d4d4d32775c788b7a1b85e8bc9d482d92250fc0e3f8efe" -dependencies = [ - "digest", -] - -[[package]] -name = "blake3" -version = "1.8.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0aa83c34e62843d924f905e0f5c866eb1dd6545fc4d719e803d9ba6030371fce" -dependencies = [ - "arrayref", - "arrayvec", - "cc", - "cfg-if", - "constant_time_eq", - "cpufeatures 0.3.0", -] - [[package]] name = "block-buffer" version = "0.10.4" @@ -716,27 +645,6 @@ dependencies = [ "generic-array", ] -[[package]] -name = "brotli" -version = "7.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cc97b8f16f944bba54f0433f07e30be199b6dc2bd25937444bbad560bcea29bd" -dependencies = [ - "alloc-no-stdlib", - "alloc-stdlib", - "brotli-decompressor", -] - -[[package]] -name = "brotli-decompressor" -version = "4.0.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a334ef7c9e23abf0ce748e8cd309037da93e606ad52eb372e4ce327a0dcfbdfd" -dependencies = [ - "alloc-no-stdlib", - "alloc-stdlib", -] - [[package]] name = "bumpalo" version = "3.20.2" @@ -755,35 +663,6 @@ version = "1.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1e748733b7cbc798e1434b6ac524f0c1ff2ab456fe201501e6497c8417a4fc33" -[[package]] -name = "bzip2" -version = "0.4.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bdb116a6ef3f6c3698828873ad02c3014b3c85cadb88496095628e3ef1e347f8" -dependencies = [ - "bzip2-sys", - "libc", -] - -[[package]] -name = "bzip2" -version = "0.5.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "49ecfb22d906f800d4fe833b6282cf4dc1c298f5057ca0b5445e5c209735ca47" -dependencies = [ - "bzip2-sys", -] - -[[package]] -name = "bzip2-sys" -version = "0.1.13+1.0.8" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "225bff33b2141874fe80d71e07d6eec4f85c5c216453dd96388240f96e1acc14" -dependencies = [ - "cc", - "pkg-config", -] - [[package]] name = "cactus" version = "1.0.7" @@ -849,16 +728,6 @@ dependencies = [ "windows-targets 0.52.6", ] -[[package]] -name = "chrono-tz" -version = "0.10.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a6139a8597ed92cf816dfb33f5dd6cf0bb93a6adc938f11039f371bc5bcd26c3" -dependencies = [ - "chrono", - "phf", -] - [[package]] name = "ciborium" version = "0.2.2" @@ -956,16 +825,6 @@ version = "1.0.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1d07550c9036bf2ae0c684c4297d503f838287c83c53686d05370d0e139ae570" -[[package]] -name = "comfy-table" -version = "7.2.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "958c5d6ecf1f214b4c2bbbbf6ab9523a864bd136dcf71a7e8904799acfe1ad47" -dependencies = [ - "unicode-segmentation", - "unicode-width 0.2.2", -] - [[package]] name = "const-random" version = "0.1.18" @@ -986,12 +845,6 @@ dependencies = [ "tiny-keccak", ] -[[package]] -name = "constant_time_eq" -version = "0.4.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3d52eff69cd5e647efe296129160853a42795992097e8af39800e1060caeea9b" - [[package]] name = "controller" version = "0.1.0" @@ -1011,7 +864,7 @@ dependencies = [ "serde", "serde_json", "serde_yaml", - "sqlparser 0.61.0", + "sqlparser", "thiserror 1.0.69", "tokio", "tokio-stream", @@ -1049,15 +902,6 @@ dependencies = [ "libc", ] -[[package]] -name = "cpufeatures" -version = "0.3.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8b2a41393f66f16b0823bb79094d54ac5fbd34ab292ddafb9a0456ac9f87d201" -dependencies = [ - "libc", -] - [[package]] name = "crc32fast" version = "1.5.0" @@ -1157,472 +1001,42 @@ dependencies = [ name = "csv" version = "1.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "52cd9d68cf7efc6ddfaaee42e7288d3a99d613d4b50f76ce9827ae0c6e14f938" -dependencies = [ - "csv-core", - "itoa", - "ryu", - "serde_core", -] - -[[package]] -name = "csv-core" -version = "0.1.13" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "704a3c26996a80471189265814dbc2c257598b96b8a7feae2d31ace646bb9782" -dependencies = [ - "memchr", -] - -[[package]] -name = "dashmap" -version = "5.5.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "978747c1d849a7d2ee5e8adc0159961c48fb7e5db2f06af6723b80123bb53856" -dependencies = [ - "cfg-if", - "hashbrown 0.14.5", - "lock_api", - "once_cell", - "parking_lot_core", -] - -[[package]] -name = "dashmap" -version = "6.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5041cc499144891f3790297212f32a74fb938e5136a14943f338ef9e0ae276cf" -dependencies = [ - "cfg-if", - "crossbeam-utils", - "hashbrown 0.14.5", - "lock_api", - "once_cell", - "parking_lot_core", -] - -[[package]] -name = "data-encoding" -version = "2.11.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a4ae5f15dda3c708c0ade84bfee31ccab44a3da4f88015ed22f63732abe300c8" - -[[package]] -name = "datafusion" -version = "43.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cbba0799cf6913b456ed07a94f0f3b6e12c62a5d88b10809e2284a0f2b915c05" -dependencies = [ - "ahash", - "arrow", - "arrow-array", - "arrow-ipc", - "arrow-schema", - "async-compression", - "async-trait", - "bytes", - "bzip2 0.4.4", - "chrono", - "dashmap 6.1.0", - "datafusion-catalog", - "datafusion-common", - "datafusion-common-runtime", - "datafusion-execution", - "datafusion-expr", - "datafusion-functions", - "datafusion-functions-aggregate", - "datafusion-functions-nested", - "datafusion-functions-window", - "datafusion-optimizer", - "datafusion-physical-expr", - "datafusion-physical-expr-common", - "datafusion-physical-optimizer", - "datafusion-physical-plan", - "datafusion-sql", - "flate2", - "futures", - "glob", - "half", - "hashbrown 0.14.5", - "indexmap 2.14.0", - "itertools 0.13.0", - "log", - "num_cpus", - "object_store", - "parking_lot", - "parquet", - "paste", - "pin-project-lite", - "rand 0.8.6", - "sqlparser 0.51.0", - "tempfile", - "tokio", - "tokio-util", - "url", - "uuid", - "xz2", - "zstd", -] - -[[package]] -name = "datafusion-catalog" -version = "43.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7493c5c2d40eec435b13d92e5703554f4efc7059451fcb8d3a79580ff0e45560" -dependencies = [ - "arrow-schema", - "async-trait", - "datafusion-common", - "datafusion-execution", - "datafusion-expr", - "datafusion-physical-plan", - "parking_lot", -] - -[[package]] -name = "datafusion-common" -version = "43.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "24953049ebbd6f8964f91f60aa3514e121b5e81e068e33b60e77815ab369b25c" -dependencies = [ - "ahash", - "arrow", - "arrow-array", - "arrow-buffer", - "arrow-schema", - "chrono", - "half", - "hashbrown 0.14.5", - "indexmap 2.14.0", - "instant", - "libc", - "num_cpus", - "object_store", - "parquet", - "paste", - "sqlparser 0.51.0", - "tokio", -] - -[[package]] -name = "datafusion-common-runtime" -version = "43.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f06df4ef76872e11c924d3c814fd2a8dd09905ed2e2195f71c857d78abd19685" -dependencies = [ - "log", - "tokio", -] - -[[package]] -name = "datafusion-execution" -version = "43.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6bbdcb628d690f3ce5fea7de81642b514486d58ff9779a51f180a69a4eadb361" -dependencies = [ - "arrow", - "chrono", - "dashmap 6.1.0", - "datafusion-common", - "datafusion-expr", - "futures", - "hashbrown 0.14.5", - "log", - "object_store", - "parking_lot", - "rand 0.8.6", - "tempfile", - "url", -] - -[[package]] -name = "datafusion-expr" -version = "43.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8036495980e3131f706b7d33ab00b4492d73dc714e3cb74d11b50f9602a73246" -dependencies = [ - "ahash", - "arrow", - "arrow-array", - "arrow-buffer", - "chrono", - "datafusion-common", - "datafusion-expr-common", - "datafusion-functions-aggregate-common", - "datafusion-functions-window-common", - "datafusion-physical-expr-common", - "indexmap 2.14.0", - "paste", - "serde_json", - "sqlparser 0.51.0", - "strum", - "strum_macros", -] - -[[package]] -name = "datafusion-expr-common" -version = "43.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4da0f3cb4669f9523b403d6b5a0ec85023e0ab3bf0183afd1517475b3e64fdd2" -dependencies = [ - "arrow", - "datafusion-common", - "itertools 0.13.0", - "paste", -] - -[[package]] -name = "datafusion-functions" -version = "43.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f52c4012648b34853e40a2c6bcaa8772f837831019b68aca384fb38436dba162" -dependencies = [ - "arrow", - "arrow-buffer", - "base64 0.22.1", - "blake2", - "blake3", - "chrono", - "datafusion-common", - "datafusion-execution", - "datafusion-expr", - "hashbrown 0.14.5", - "hex", - "itertools 0.13.0", - "log", - "md-5", - "rand 0.8.6", - "regex", - "sha2", - "unicode-segmentation", - "uuid", -] - -[[package]] -name = "datafusion-functions-aggregate" -version = "43.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e5b8bb624597ba28ed7446df4a9bd7c7a7bde7c578b6b527da3f47371d5f6741" -dependencies = [ - "ahash", - "arrow", - "arrow-schema", - "datafusion-common", - "datafusion-execution", - "datafusion-expr", - "datafusion-functions-aggregate-common", - "datafusion-physical-expr", - "datafusion-physical-expr-common", - "half", - "indexmap 2.14.0", - "log", - "paste", -] - -[[package]] -name = "datafusion-functions-aggregate-common" -version = "43.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6fb06208fc470bc8cf1ce2d9a1159d42db591f2c7264a8c1776b53ad8f675143" -dependencies = [ - "ahash", - "arrow", - "datafusion-common", - "datafusion-expr-common", - "datafusion-physical-expr-common", - "rand 0.8.6", -] - -[[package]] -name = "datafusion-functions-nested" -version = "43.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fca25bbb87323716d05e54114666e942172ccca23c5a507e9c7851db6e965317" -dependencies = [ - "arrow", - "arrow-array", - "arrow-buffer", - "arrow-ord", - "arrow-schema", - "datafusion-common", - "datafusion-execution", - "datafusion-expr", - "datafusion-functions", - "datafusion-functions-aggregate", - "datafusion-physical-expr-common", - "itertools 0.13.0", - "log", - "paste", - "rand 0.8.6", -] - -[[package]] -name = "datafusion-functions-window" -version = "43.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5ae23356c634e54c59f7c51acb7a5b9f6240ffb2cf997049a1a24a8a88598dbe" -dependencies = [ - "datafusion-common", - "datafusion-expr", - "datafusion-functions-window-common", - "datafusion-physical-expr", - "datafusion-physical-expr-common", - "log", - "paste", -] - -[[package]] -name = "datafusion-functions-window-common" -version = "43.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d4b3d6ff7794acea026de36007077a06b18b89e4f9c3fea7f2215f9f7dd9059b" -dependencies = [ - "datafusion-common", - "datafusion-physical-expr-common", -] - -[[package]] -name = "datafusion-optimizer" -version = "43.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bec6241eb80c595fa0e1a8a6b69686b5cf3bd5fdacb8319582a0943b0bd788aa" -dependencies = [ - "arrow", - "async-trait", - "chrono", - "datafusion-common", - "datafusion-expr", - "datafusion-physical-expr", - "hashbrown 0.14.5", - "indexmap 2.14.0", - "itertools 0.13.0", - "log", - "paste", - "regex-syntax", -] - -[[package]] -name = "datafusion-physical-expr" -version = "43.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3370357b8fc75ec38577700644e5d1b0bc78f38babab99c0b8bd26bafb3e4335" -dependencies = [ - "ahash", - "arrow", - "arrow-array", - "arrow-buffer", - "arrow-ord", - "arrow-schema", - "arrow-string", - "chrono", - "datafusion-common", - "datafusion-expr", - "datafusion-expr-common", - "datafusion-functions-aggregate-common", - "datafusion-physical-expr-common", - "half", - "hashbrown 0.14.5", - "indexmap 2.14.0", - "itertools 0.13.0", - "log", - "paste", - "petgraph 0.6.5", -] - -[[package]] -name = "datafusion-physical-expr-common" -version = "43.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b8b7734d94bf2fa6f6e570935b0ddddd8421179ce200065be97874e13d46a47b" -dependencies = [ - "ahash", - "arrow", - "datafusion-common", - "datafusion-expr-common", - "hashbrown 0.14.5", - "rand 0.8.6", -] - -[[package]] -name = "datafusion-physical-optimizer" -version = "43.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7eee8c479522df21d7b395640dff88c5ed05361852dce6544d7c98e9dbcebffe" -dependencies = [ - "arrow", - "arrow-schema", - "datafusion-common", - "datafusion-execution", - "datafusion-expr-common", - "datafusion-physical-expr", - "datafusion-physical-plan", - "itertools 0.13.0", -] - -[[package]] -name = "datafusion-physical-plan" -version = "43.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "17e1fc2e2c239d14e8556f2622b19a726bf6bc6962cc00c71fc52626274bee24" -dependencies = [ - "ahash", - "arrow", - "arrow-array", - "arrow-buffer", - "arrow-ord", - "arrow-schema", - "async-trait", - "chrono", - "datafusion-common", - "datafusion-common-runtime", - "datafusion-execution", - "datafusion-expr", - "datafusion-functions-aggregate-common", - "datafusion-functions-window-common", - "datafusion-physical-expr", - "datafusion-physical-expr-common", - "futures", - "half", - "hashbrown 0.14.5", - "indexmap 2.14.0", - "itertools 0.13.0", - "log", - "once_cell", - "parking_lot", - "pin-project-lite", - "rand 0.8.6", - "tokio", +checksum = "52cd9d68cf7efc6ddfaaee42e7288d3a99d613d4b50f76ce9827ae0c6e14f938" +dependencies = [ + "csv-core", + "itoa", + "ryu", + "serde_core", ] [[package]] -name = "datafusion-sql" -version = "43.0.0" +name = "csv-core" +version = "0.1.13" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "63e3a4ed41dbee20a5d947a59ca035c225d67dc9cbe869c10f66dcdf25e7ce51" +checksum = "704a3c26996a80471189265814dbc2c257598b96b8a7feae2d31ace646bb9782" dependencies = [ - "arrow", - "arrow-array", - "arrow-schema", - "datafusion-common", - "datafusion-expr", - "indexmap 2.14.0", - "log", - "regex", - "sqlparser 0.51.0", - "strum", + "memchr", ] [[package]] -name = "datafusion_summary_library" -version = "0.1.0" +name = "dashmap" +version = "5.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "978747c1d849a7d2ee5e8adc0159961c48fb7e5db2f06af6723b80123bb53856" dependencies = [ - "arrow", - "async-trait", - "datafusion", - "futures", - "hyperloglogplus", + "cfg-if", + "hashbrown 0.14.5", + "lock_api", + "once_cell", + "parking_lot_core", ] +[[package]] +name = "data-encoding" +version = "2.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a4ae5f15dda3c708c0ade84bfee31ccab44a3da4f88015ed22f63732abe300c8" + [[package]] name = "deranged" version = "0.5.8" @@ -1736,12 +1150,6 @@ version = "0.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5baebc0774151f905a1a2cc41989300b1e6fbb29aff0ceffa1064fdd3088d582" -[[package]] -name = "fixedbitset" -version = "0.4.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0ce7134b9999ecaf8bcd65542e436736ef32ddca1b3e06094cb6ec5755203b80" - [[package]] name = "fixedbitset" version = "0.5.7" @@ -1942,12 +1350,6 @@ dependencies = [ "wasip3", ] -[[package]] -name = "glob" -version = "0.3.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0cc23270f6e1808e30a928bdc84dea0b9b4136a8bc82338574f23baf47bbd280" - [[package]] name = "h2" version = "0.3.27" @@ -2011,7 +1413,6 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e5274423e17b7c9fc20b6e7e208532f9b19825d82dfd615708b70edd83df41f1" dependencies = [ "ahash", - "allocator-api2", ] [[package]] @@ -2161,12 +1562,6 @@ version = "1.0.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "df3b46402a9d5adb4c86a0cf463f42e19994e3ee891101b1841f30a545cb49a9" -[[package]] -name = "humantime" -version = "2.3.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "135b12329e5e3ce057a9f972339ea52bc954fe1e9358ef27f95e89716fbc5424" - [[package]] name = "hyper" version = "0.14.32" @@ -2279,15 +1674,6 @@ dependencies = [ "tracing", ] -[[package]] -name = "hyperloglogplus" -version = "0.4.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "621debdf94dcac33e50475fdd76d34d5ea9c0362a834b9db08c3024696c1fbe3" -dependencies = [ - "serde", -] - [[package]] name = "iana-time-zone" version = "0.1.65" @@ -2443,24 +1829,6 @@ dependencies = [ "serde_core", ] -[[package]] -name = "instant" -version = "0.1.13" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e0242819d153cba4b4b05a5a8f2a7e9bbf97b6055b2a002b395c96b5ff3c0222" -dependencies = [ - "cfg-if", - "js-sys", - "wasm-bindgen", - "web-sys", -] - -[[package]] -name = "integer-encoding" -version = "3.0.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8bb03732005da905c88227371639bf1ad885cc712789c011c31c5fb3ab3ccf02" - [[package]] name = "ipnet" version = "2.12.0" @@ -2493,15 +1861,6 @@ dependencies = [ "either", ] -[[package]] -name = "itertools" -version = "0.13.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "413ee7dfc52ee1a4949ceeb7dbc8a33f2d6c088194d9f922fb8318faf1f01186" -dependencies = [ - "either", -] - [[package]] name = "itertools" version = "0.14.0" @@ -2757,26 +2116,6 @@ version = "0.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "112b39cec0b298b6c1999fee3e31427f74f676e4cb9879ed1a121b43661a4154" -[[package]] -name = "lz4_flex" -version = "0.11.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "373f5eceeeab7925e0c1098212f2fbc4d416adec9d35051a6ab251e824c1854a" -dependencies = [ - "twox-hash 2.1.2", -] - -[[package]] -name = "lzma-sys" -version = "0.1.20" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5fda04ab3764e6cde78b9974eec4f779acaba7c4e84b36eca3cf77c581b85d27" -dependencies = [ - "cc", - "libc", - "pkg-config", -] - [[package]] name = "matchers" version = "0.2.0" @@ -2803,16 +2142,6 @@ dependencies = [ "syn 2.0.117", ] -[[package]] -name = "md-5" -version = "0.10.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d89e7ee0cfbedfc4da3340218492196241d89eefb6dab27de5df917a6d2e78cf" -dependencies = [ - "cfg-if", - "digest", -] - [[package]] name = "md5" version = "0.8.0" @@ -2982,16 +2311,6 @@ dependencies = [ "libm", ] -[[package]] -name = "num_cpus" -version = "1.17.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "91df4bbde75afed763b708b7eee1e8e7651e02d97f6d5dd763e89367e957b23b" -dependencies = [ - "hermit-abi 0.5.2", - "libc", -] - [[package]] name = "num_enum" version = "0.7.6" @@ -3051,27 +2370,6 @@ dependencies = [ "memchr", ] -[[package]] -name = "object_store" -version = "0.11.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3cfccb68961a56facde1163f9319e0d15743352344e7808a11795fb99698dcaf" -dependencies = [ - "async-trait", - "bytes", - "chrono", - "futures", - "humantime", - "itertools 0.13.0", - "parking_lot", - "percent-encoding", - "snafu", - "tokio", - "tracing", - "url", - "walkdir", -] - [[package]] name = "once_cell" version = "1.21.4" @@ -3090,15 +2388,6 @@ version = "11.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d6790f58c7ff633d8771f42965289203411a5e5c68388703c06e14f24770b41e" -[[package]] -name = "ordered-float" -version = "2.10.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "68f19d67e5a2795c94e73e0bb1cc1a7edeb2e28efd39e2e1c9b7a40c1108b11c" -dependencies = [ - "num-traits", -] - [[package]] name = "ordered-multimap" version = "0.7.3" @@ -3142,92 +2431,22 @@ dependencies = [ "windows-link 0.2.1", ] -[[package]] -name = "parquet" -version = "53.4.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2f8cf58b29782a7add991f655ff42929e31a7859f5319e53db9e39a714cb113c" -dependencies = [ - "ahash", - "arrow-array", - "arrow-buffer", - "arrow-cast", - "arrow-data", - "arrow-ipc", - "arrow-schema", - "arrow-select", - "base64 0.22.1", - "brotli", - "bytes", - "chrono", - "flate2", - "futures", - "half", - "hashbrown 0.15.5", - "lz4_flex", - "num", - "num-bigint", - "object_store", - "paste", - "seq-macro", - "snap", - "thrift", - "tokio", - "twox-hash 1.6.3", - "zstd", - "zstd-sys", -] - -[[package]] -name = "paste" -version = "1.0.15" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" - [[package]] name = "percent-encoding" version = "2.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9b4f627cb1b25917193a259e49bdad08f671f8d9708acfd5fe0a8c1455d87220" -[[package]] -name = "petgraph" -version = "0.6.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b4c5cc86750666a3ed20bdaf5ca2a0344f9c67674cae0515bec2da16fbaa47db" -dependencies = [ - "fixedbitset 0.4.2", - "indexmap 2.14.0", -] - [[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.14.0", ] -[[package]] -name = "phf" -version = "0.12.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "913273894cec178f401a31ec4b656318d95473527be05c0752cc41cdc32be8b7" -dependencies = [ - "phf_shared", -] - -[[package]] -name = "phf_shared" -version = "0.12.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "06005508882fb681fd97892ecff4b7fd0fee13ef1aa569f8695dae7ab9099981" -dependencies = [ - "siphasher", -] - [[package]] name = "pin-project" version = "1.1.12" @@ -3478,7 +2697,7 @@ dependencies = [ "log", "multimap", "once_cell", - "petgraph 0.7.1", + "petgraph", "prettyplease", "prost", "prost-types", @@ -3610,9 +2829,7 @@ dependencies = [ "controller", "crc32fast", "criterion", - "dashmap 5.5.3", - "datafusion", - "datafusion_summary_library", + "dashmap", "flate2", "form_urlencoded", "futures", @@ -4236,12 +3453,6 @@ version = "1.0.28" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8a7852d02fc848982e0c167ef163aaff9cd91dc640ba85e263cb1ce46fae51cd" -[[package]] -name = "seq-macro" -version = "0.3.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1bc711410fbe7399f390ca1c3b60ad0f53f80e95c5eb935e52268a0e2cd49acc" - [[package]] name = "serde" version = "1.0.228" @@ -4337,7 +3548,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e3bf829a2d51ab4a5ddf1352d8470c140cadc8301b2ae1789db023f01cedd6ba" dependencies = [ "cfg-if", - "cpufeatures 0.2.17", + "cpufeatures", "digest", ] @@ -4348,7 +3559,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" dependencies = [ "cfg-if", - "cpufeatures 0.2.17", + "cpufeatures", "digest", ] @@ -4383,12 +3594,6 @@ version = "0.3.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "703d5c7ef118737c72f1af64ad2f6f8c5e1921f818cdcb97b8fe6fc69bf66214" -[[package]] -name = "siphasher" -version = "1.0.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8ee5873ec9cce0195efcb7a4e9507a04cd49aec9c83d0389df45b1ef7ba2e649" - [[package]] name = "slab" version = "0.4.12" @@ -4401,33 +3606,6 @@ version = "1.15.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "67b1b7a3b5fe4f1376887184045fcf45c69e92af734b7aaddc05fb777b6fbd03" -[[package]] -name = "snafu" -version = "0.8.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6e84b3f4eacbf3a1ce05eac6763b4d629d60cbc94d632e4092c54ade71f1e1a2" -dependencies = [ - "snafu-derive", -] - -[[package]] -name = "snafu-derive" -version = "0.8.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c1c97747dbf44bb1ca44a561ece23508e99cb592e862f22222dcf42f51d1e451" -dependencies = [ - "heck 0.5.0", - "proc-macro2", - "quote", - "syn 2.0.117", -] - -[[package]] -name = "snap" -version = "1.1.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1b6b67fb9a61334225b5b790716f609cd58395f895b3fe8b328786812a40bc3b" - [[package]] name = "socket2" version = "0.5.10" @@ -4460,16 +3638,6 @@ dependencies = [ "vob", ] -[[package]] -name = "sqlparser" -version = "0.51.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5fe11944a61da0da3f592e19a45ebe5ab92dc14a779907ff1f08fbb797bfefc7" -dependencies = [ - "log", - "sqlparser_derive", -] - [[package]] name = "sqlparser" version = "0.61.0" @@ -4480,17 +3648,6 @@ dependencies = [ "recursive", ] -[[package]] -name = "sqlparser_derive" -version = "0.2.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "01b2e185515564f15375f593fb966b5718bc624ba77fe49fa4616ad619690554" -dependencies = [ - "proc-macro2", - "quote", - "syn 2.0.117", -] - [[package]] name = "stable_deref_trait" version = "1.2.1" @@ -4552,28 +3709,6 @@ dependencies = [ "syn 1.0.109", ] -[[package]] -name = "strum" -version = "0.26.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8fec0f0aef304996cf250b31b5a10dee7980c85da9d759361292b8bca5a18f06" -dependencies = [ - "strum_macros", -] - -[[package]] -name = "strum_macros" -version = "0.26.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4c6bee85a5a24955dc440386795aa378cd9cf82acd5f764469152d2270e581be" -dependencies = [ - "heck 0.5.0", - "proc-macro2", - "quote", - "rustversion", - "syn 2.0.117", -] - [[package]] name = "subtle" version = "2.6.1" @@ -4746,17 +3881,6 @@ dependencies = [ "cfg-if", ] -[[package]] -name = "thrift" -version = "0.17.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7e54bc85fc7faa8bc175c4bab5b92ba8d9a3ce893d0e9f42cc455c8ab16a9e09" -dependencies = [ - "byteorder", - "integer-encoding", - "ordered-float", -] - [[package]] name = "time" version = "0.3.47" @@ -5200,16 +4324,6 @@ dependencies = [ "utf-8", ] -[[package]] -name = "twox-hash" -version = "1.6.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "97fee6b57c6a41524a810daee9286c02d7752c4253064d0b05472833a438f675" -dependencies = [ - "cfg-if", - "static_assertions", -] - [[package]] name = "twox-hash" version = "2.1.2" @@ -5991,15 +5105,6 @@ version = "0.8.15" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fdd20c5420375476fbd4394763288da7eb0cc0b8c11deed431a91562af7335d3" -[[package]] -name = "xz2" -version = "0.1.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "388c44dc09d76f1536602ead6d325eb532f5c122f17782bd57fb47baeeb767e2" -dependencies = [ - "lzma-sys", -] - [[package]] name = "yoke" version = "0.8.2" diff --git a/Cargo.toml b/Cargo.toml index 2c3dc9b7..ec22820f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -4,7 +4,6 @@ members = [ "crates/promql_utilities", "crates/asap_otel_proto", "crates/asap_types", - "crates/datafusion_summary_library", "asap-query-engine", "controller", ] @@ -47,6 +46,5 @@ arc-swap = "1.7" # Internal crates promql_utilities = { path = "crates/promql_utilities" } asap_types = { path = "crates/asap_types" } -datafusion_summary_library = { path = "crates/datafusion_summary_library" } asap_otel_proto = { path = "crates/asap_otel_proto" } indexmap = { version = "2.0", features = ["serde"] } diff --git a/asap-query-engine/Cargo.toml b/asap-query-engine/Cargo.toml index 79978a31..727f52f2 100644 --- a/asap-query-engine/Cargo.toml +++ b/asap-query-engine/Cargo.toml @@ -7,7 +7,6 @@ edition.workspace = true # Internal crates (workspace) promql_utilities.workspace = true asap_types.workspace = true -datafusion_summary_library.workspace = true # Phase 9: controller is now an in-process library inside the backend # binary. Wiring up the in-process OpAMP server + capability-map # exposure is a follow-up after Phase 4 (centralized series_id @@ -45,7 +44,6 @@ async-trait = "0.1" xxhash-rust = { version = "0.8", features = ["xxh32", "xxh64"] } base64 = "0.21" hex = "0.4" -datafusion = "43" arrow = "53.4.1" futures = "0.3" prost = "0.13" diff --git a/asap-query-engine/src/drivers/ingest/otel.rs b/asap-query-engine/src/drivers/ingest/otel.rs index e1e8f706..dc7c6b9d 100644 --- a/asap-query-engine/src/drivers/ingest/otel.rs +++ b/asap-query-engine/src/drivers/ingest/otel.rs @@ -721,7 +721,7 @@ async fn route_modified_otlp_sketches_to_precompute( // then route each tuple through the same dispatcher. let dps: Vec = match &metric.data { Some(Data::Ddsketch(d)) => { - let cfg = crate::stores::sketch_index::SketchConfig::DDSketch { + let cfg = crate::stores::sketch_db::sketch_index::SketchConfig::DDSketch { relative_accuracy: d.relative_accuracy, }; d.data_points @@ -739,7 +739,7 @@ async fn route_modified_otlp_sketches_to_precompute( .collect() } Some(Data::Kllsketch(k)) => { - let cfg = crate::stores::sketch_index::SketchConfig::Kll { k: k.k }; + let cfg = crate::stores::sketch_db::sketch_index::SketchConfig::Kll { k: k.k }; k.data_points .iter() .map(|dp| ModifiedOtlpSketchDp { @@ -755,7 +755,7 @@ async fn route_modified_otlp_sketches_to_precompute( .collect() } Some(Data::Countsketch(c)) => { - let cfg = crate::stores::sketch_index::SketchConfig::CountSketch { + let cfg = crate::stores::sketch_db::sketch_index::SketchConfig::CountSketch { rows: c.rows, cols: c.cols, }; @@ -774,7 +774,7 @@ async fn route_modified_otlp_sketches_to_precompute( .collect() } Some(Data::Countminsketch(c)) => { - let cfg = crate::stores::sketch_index::SketchConfig::CountMin { + let cfg = crate::stores::sketch_db::sketch_index::SketchConfig::CountMin { rows: c.rows, cols: c.cols, }; @@ -793,7 +793,7 @@ async fn route_modified_otlp_sketches_to_precompute( .collect() } Some(Data::Hllsketch(h)) => { - let cfg = crate::stores::sketch_index::SketchConfig::Hll { + let cfg = crate::stores::sketch_db::sketch_index::SketchConfig::Hll { precision: h.precision, }; h.data_points @@ -879,7 +879,7 @@ async fn route_modified_otlp_sketches_to_precompute( // rollup, `attributes` is the group-by VALUES vector, // and its key set IS the group-by KEY set. { - use crate::stores::sketch_index::{ + use crate::stores::sketch_db::sketch_index::{ AccuracyBound, Capability, SketchEncoding, SketchInstanceMetadata, SketchKindHandle, SketchSampleState, }; @@ -917,7 +917,7 @@ async fn route_modified_otlp_sketches_to_precompute( .iter() .map(|(k, v)| (k.clone(), v.clone())) .collect(); - let window: crate::stores::epoch_columnar::TimestampRange = ( + let window: crate::stores::sketch_db::epoch_columnar::TimestampRange = ( dp.start_time_unix_nano / 1_000_000, dp.time_unix_nano / 1_000_000, ); @@ -1079,8 +1079,8 @@ async fn route_modified_otlp_sketches_to_precompute( /// share one source of truth. fn sketch_kind_handle_for( dp: &ModifiedOtlpSketchDp, -) -> crate::stores::sketch_index::SketchKindHandle { - use crate::stores::sketch_index::SketchKindHandle; +) -> crate::stores::sketch_db::sketch_index::SketchKindHandle { + use crate::stores::sketch_db::sketch_index::SketchKindHandle; match dp.kind { SketchKind::DdSketch => SketchKindHandle::DDSketch, SketchKind::Kll => SketchKindHandle::Kll, @@ -1094,8 +1094,8 @@ fn sketch_kind_handle_for( /// SketchIndex's `SketchEncoding` enum. Returns `None` for the unset /// (0) encoding so callers can default to `ProtoFull` (the dominant /// case for full-state frames). -fn encoding_to_handle(encoding: i32) -> Option { - use crate::stores::sketch_index::SketchEncoding; +fn encoding_to_handle(encoding: i32) -> Option { + use crate::stores::sketch_db::sketch_index::SketchEncoding; match encoding { ENCODING_PROTO => Some(SketchEncoding::ProtoFull), ENCODING_PROTO_DELTA => Some(SketchEncoding::ProtoDelta), @@ -1137,7 +1137,7 @@ struct ModifiedOtlpSketchDp { /// Phase 5 — sketch-instance configuration lifted off the parent /// container. Drives `SketchInstanceMetadata.sketch_config` and the /// derived `AccuracyBound`. - container_config: crate::stores::sketch_index::SketchConfig, + container_config: crate::stores::sketch_db::sketch_index::SketchConfig, } /// Decode the typed `sketch` bytes from a modified-OTLP @@ -1840,7 +1840,7 @@ mod sid_resolution_tests { use crate::drivers::ingest::series_resolver::SeriesIdResolver; use crate::precompute_engine::series_router::SeriesRouter; use crate::stores::sketch_db::SchemaRegistry; - use crate::stores::sketch_index::SketchIndex; + use crate::stores::sketch_db::sketch_index::SketchIndex; use asap_otel_proto::tonic::collector::metrics::v1::ExportMetricsServiceRequest; use asap_otel_proto::tonic::common::v1::{any_value::Value as AnyVal, AnyValue, KeyValue}; use asap_otel_proto::tonic::metrics::v1::{ diff --git a/asap-query-engine/src/engines/logical/mod.rs b/asap-query-engine/src/engines/logical/mod.rs deleted file mode 100644 index 435ca52b..00000000 --- a/asap-query-engine/src/engines/logical/mod.rs +++ /dev/null @@ -1,6 +0,0 @@ -//! Logical Plan Builders for DataFusion -//! -//! This module converts QueryExecutionContext into DataFusion LogicalPlan trees -//! using the custom extension nodes defined in datafusion_summary_library. - -pub mod plan_builder; diff --git a/asap-query-engine/src/engines/logical/plan_builder.rs b/asap-query-engine/src/engines/logical/plan_builder.rs deleted file mode 100644 index 925a8a3d..00000000 --- a/asap-query-engine/src/engines/logical/plan_builder.rs +++ /dev/null @@ -1,1014 +0,0 @@ -//! Plan Builder Module -//! -//! Converts QueryExecutionContext to DataFusion LogicalPlan for OnlySpatial queries. -//! This enables plan-based execution as an alternative to the existing pipeline. - -use arrow::datatypes::{DataType, Field}; -use datafusion::common::{DFSchema, DFSchemaRef}; -use datafusion::error::DataFusionError; -use datafusion::logical_expr::{ - binary_expr, Expr as DFExpr, Extension, JoinType, LogicalPlan, LogicalPlanBuilder, Operator, - SubqueryAlias, -}; -use datafusion::prelude::{col, lit}; -use datafusion_summary_library::{ - InferOperation, PrecomputedSummaryRead, SketchType, SummaryInfer, SummaryMergeMultiple, -}; -use promql_parser::parser::token::{self, T_ADD, T_DIV, T_MOD, T_MUL, T_POW, T_SUB}; -use promql_utilities::query_logics::enums::{AggregationType, Statistic}; -use std::sync::Arc; - -use crate::engines::simple::engine::{QueryExecutionContext, StoreQueryParams}; - -/// Extension trait for building DataFusion logical plans from QueryExecutionContext -impl QueryExecutionContext { - /// Convert this execution context to a DataFusion LogicalPlan. - /// - /// The resulting plan structure for single-population queries: - /// ```text - /// SummaryInfer (extract values from summaries) - /// └─► SummaryMergeMultiple (merge summaries with same group key) - /// └─► PrecomputedSummaryRead (read from store) - /// ``` - /// - /// For multi-population queries (separate keys_query): - /// ```text - /// SummaryInfer (dual-input: value sketch + keys enumeration) - /// ├─► input 0: SummaryMergeMultiple (values) - /// │ └─► PrecomputedSummaryRead (values agg_id) - /// └─► input 1: SummaryMergeMultiple (keys) - /// └─► PrecomputedSummaryRead (keys agg_id) - /// ``` - pub fn to_logical_plan(&self) -> Result { - let has_separate_keys = self.store_plan.keys_query.is_some() - && self.agg_info.aggregation_id_for_key != self.agg_info.aggregation_id_for_value; - - // 1. Map aggregation type to SummaryType (SketchType) for values - let summary_type = self.map_aggregation_type_to_summary_type()?; - - // Determine labels for the values branch (store read/merge). - // For multi-population (dual-input or self-keyed): use grouping_labels (store GROUP BY) - // For single-population: use query_output_labels - let has_aggregated_labels = !self.aggregated_labels.labels.is_empty(); - let values_labels = if has_separate_keys || has_aggregated_labels { - self.grouping_labels.labels.to_vec() - } else { - self.get_output_label_names() - }; - - // Sub-key labels come from aggregated_labels (labels that key the accumulator internally) - let sub_key_labels: Vec = self.aggregated_labels.labels.to_vec(); - - // 2. Build values branch: Read -> Merge - let values_merge_plan = self.build_read_merge_branch( - &self.store_plan.values_query, - &values_labels, - &summary_type, - )?; - - // 3. Map statistic to InferOperation - let infer_operation = self.map_statistic_to_infer_operation()?; - - if has_separate_keys { - let keys_query = self.store_plan.keys_query.as_ref().unwrap(); - - // Map keys aggregation type to SketchType - let keys_summary_type = self.map_key_aggregation_type_to_summary_type()?; - - // Build keys branch: Read -> Merge (using same spatial labels) - let keys_merge_plan = - self.build_read_merge_branch(keys_query, &values_labels, &keys_summary_type)?; - - // Create dual-input SummaryInfer - let infer = SummaryInfer::new( - Arc::new(values_merge_plan), - vec![infer_operation], - vec!["value".to_string()], - ) - .map_err(|e| DataFusionError::Plan(format!("Failed to create SummaryInfer: {}", e)))? - .with_keys_input(Arc::new(keys_merge_plan)) - .with_group_key_columns(sub_key_labels, None) - .map_err(|e| { - DataFusionError::Plan(format!("Failed to set group_key_columns: {}", e)) - })?; - - Ok(LogicalPlan::Extension(Extension { - node: Arc::new(infer), - })) - } else { - // Single-input path - let mut infer = SummaryInfer::new( - Arc::new(values_merge_plan), - vec![infer_operation], - vec!["value".to_string()], - ) - .map_err(|e| DataFusionError::Plan(format!("Failed to create SummaryInfer: {}", e)))?; - - if !sub_key_labels.is_empty() { - // Self-keyed multi-pop: set sub-key columns so the output schema - // includes them and the physical operator knows to enumerate keys. - infer = infer - .with_group_key_columns(sub_key_labels, None) - .map_err(|e| { - DataFusionError::Plan(format!("Failed to set group_key_columns: {}", e)) - })?; - } - - Ok(LogicalPlan::Extension(Extension { - node: Arc::new(infer), - })) - } - } - - /// Build a Read -> Merge branch for a given store query. - fn build_read_merge_branch( - &self, - query_params: &StoreQueryParams, - labels: &[String], - summary_type: &SketchType, - ) -> Result { - let read_schema = self.build_read_schema(labels)?; - - let read = PrecomputedSummaryRead::new( - self.metric.clone(), - query_params.aggregation_id, - query_params.start_timestamp, - query_params.end_timestamp, - query_params.is_exact_query, - labels.to_vec(), - summary_type.clone(), - read_schema, - ); - let read_plan = LogicalPlan::Extension(Extension { - node: Arc::new(read), - }); - - let merge = SummaryMergeMultiple::new( - Arc::new(read_plan), - labels.to_vec(), - "sketch".to_string(), - summary_type.clone(), - ); - Ok(LogicalPlan::Extension(Extension { - node: Arc::new(merge), - })) - } - - /// Get output label names from the query metadata - fn get_output_label_names(&self) -> Vec { - self.metadata.query_output_labels.labels.to_vec() - } - - /// Build schema for PrecomputedSummaryRead: [label columns, sketch column] - fn build_read_schema(&self, output_labels: &[String]) -> Result { - let mut fields: Vec<(Option, Arc)> = Vec::new(); - - // Add label columns (Utf8, nullable) - for label in output_labels { - fields.push((None, Arc::new(Field::new(label, DataType::Utf8, true)))); - } - - // Add sketch column (Binary, not nullable) - fields.push(( - None, - Arc::new(Field::new("sketch", DataType::Binary, false)), - )); - - let schema = DFSchema::new_with_metadata(fields, Default::default()) - .map_err(|e| DataFusionError::Plan(format!("Failed to create read schema: {}", e)))?; - - Ok(Arc::new(schema)) - } - - /// Map Statistic enum to InferOperation - pub(crate) fn map_statistic_to_infer_operation( - &self, - ) -> Result { - match self.metadata.statistic_to_compute { - Statistic::Sum => Ok(InferOperation::ExtractSum), - Statistic::Min => Ok(InferOperation::ExtractMin), - Statistic::Max => Ok(InferOperation::ExtractMax), - Statistic::Count => Ok(InferOperation::ExtractCount), - Statistic::Increase => Ok(InferOperation::ExtractIncrease), - Statistic::Rate => Ok(InferOperation::ExtractRate), - Statistic::Quantile => { - // Extract quantile parameter from query_kwargs - let q = self - .metadata - .query_kwargs - .get("quantile") - .and_then(|s| s.parse::().ok()) - .unwrap_or(0.5); - Ok(InferOperation::quantile(q)) - } - Statistic::Cardinality => Ok(InferOperation::CountDistinct), - Statistic::Topk => { - // Extract k parameter from query_kwargs - let k = self - .metadata - .query_kwargs - .get("k") - .and_then(|s| s.parse::().ok()) - .unwrap_or(10); - Ok(InferOperation::TopK(k)) - } - } - } - - /// Map aggregation type to SketchType (SummaryType) for the value accumulator - fn map_aggregation_type_to_summary_type(&self) -> Result { - Self::agg_type_to_sketch_type(self.agg_info.aggregation_type_for_value) - } - - /// Map aggregation type to SketchType (SummaryType) for the key accumulator - fn map_key_aggregation_type_to_summary_type(&self) -> Result { - Self::agg_type_to_sketch_type(self.agg_info.aggregation_type_for_key) - } - - fn agg_type_to_sketch_type(agg_type: AggregationType) -> Result { - match agg_type { - AggregationType::Sum => Ok(SketchType::Sum), - AggregationType::Increase => Ok(SketchType::Increase), - AggregationType::MinMax => Ok(SketchType::MinMax), - AggregationType::MultipleSum => Ok(SketchType::MultipleSum), - AggregationType::MultipleIncrease => Ok(SketchType::MultipleIncrease), - AggregationType::MultipleMinMax => Ok(SketchType::MultipleMinMax), - AggregationType::DeltaSetAggregator => Ok(SketchType::DeltaSetAggregator), - AggregationType::SetAggregator => Ok(SketchType::SetAggregator), - AggregationType::DatasketchesKLL => Ok(SketchType::KLL), - AggregationType::HydraKLL => Ok(SketchType::HydraKLL), - AggregationType::CountMinSketch | AggregationType::CountMinSketchWithHeap => { - Ok(SketchType::CountMinSketch) - } - AggregationType::HLL => Ok(SketchType::HLL), - _ => Err(DataFusionError::Plan(format!( - "Unknown aggregation type: {agg_type:?}" - ))), - } - } -} - -// ============================================================================ -// Binary arithmetic plan builders (standalone functions, not on impl block) -// ============================================================================ - -/// Map a PromQL `TokenType` to the corresponding DataFusion `Operator`. -/// Returns an error for non-arithmetic operators. -pub fn token_type_to_df_operator(op: &token::TokenType) -> Result { - match op.id() { - id if id == T_ADD => Ok(Operator::Plus), - id if id == T_SUB => Ok(Operator::Minus), - id if id == T_MUL => Ok(Operator::Multiply), - id if id == T_DIV => Ok(Operator::Divide), - id if id == T_MOD => Ok(Operator::Modulo), - id if id == T_POW => Ok(Operator::BitwiseXor), // Note: bitwise XOR used as proxy for ^ - _ => Err(DataFusionError::Plan(format!( - "Unsupported binary operator for arithmetic plan: {}", - op - ))), - } -} - -/// Builds a DataFusion logical plan for a vector op vector binary expression: -/// -/// ```text -/// Projection(lhs.label1 AS label1, ..., lhs.value OP rhs.value AS value) -/// └── Join(inner, on = label_columns) -/// ├── SubqueryAlias("lhs") └── lhs_plan -/// └── SubqueryAlias("rhs") └── rhs_plan -/// ``` -/// -/// The `label_columns` are the label names shared by both sides; the join is -/// an inner join on those columns. -pub fn build_binary_vector_plan( - lhs_plan: LogicalPlan, - rhs_plan: LogicalPlan, - op: &token::TokenType, - label_columns: Vec, -) -> Result { - let df_op = token_type_to_df_operator(op)?; - - // Wrap each side in a SubqueryAlias so columns are qualified (lhs.x, rhs.x) - let lhs_aliased = SubqueryAlias::try_new(Arc::new(lhs_plan), "lhs")?; - let rhs_aliased = SubqueryAlias::try_new(Arc::new(rhs_plan), "rhs")?; - - // Build the join keys: qualified column names for each label column. - // The `join` function expects Vec> (i.e. qualified col names). - let join_keys_left: Vec = label_columns.iter().map(|c| format!("lhs.{}", c)).collect(); - let join_keys_right: Vec = label_columns.iter().map(|c| format!("rhs.{}", c)).collect(); - - let joined_plan = LogicalPlanBuilder::from(LogicalPlan::SubqueryAlias(lhs_aliased)) - .join( - LogicalPlan::SubqueryAlias(rhs_aliased), - JoinType::Inner, - (join_keys_left, join_keys_right), - None, - )? - .build()?; - - // Projection: pass through label columns from lhs, compute value = lhs.value OP rhs.value - let mut proj_exprs: Vec = label_columns - .iter() - .map(|c| col(format!("lhs.{}", c)).alias(c.as_str())) - .collect(); - let value_expr = binary_expr(col("lhs.value"), df_op, col("rhs.value")).alias("value"); - proj_exprs.push(value_expr); - - LogicalPlanBuilder::from(joined_plan) - .project(proj_exprs)? - .build() -} - -/// Builds a DataFusion logical plan for a scalar op vector (or vector op scalar) expression: -/// -/// ```text -/// Projection(label1, ..., scalar OP value AS value) -/// └── vector_plan -/// ``` -/// -/// If `scalar_on_left` is true the expression is `scalar OP value`; -/// otherwise it is `value OP scalar`. -pub fn build_scalar_plan( - vector_plan: LogicalPlan, - scalar: f64, - op: &token::TokenType, - scalar_on_left: bool, - label_columns: Vec, -) -> Result { - let df_op = token_type_to_df_operator(op)?; - - let value_expr = if scalar_on_left { - binary_expr(lit(scalar), df_op, col("value")).alias("value") - } else { - binary_expr(col("value"), df_op, lit(scalar)).alias("value") - }; - - let mut proj_exprs: Vec = label_columns.iter().map(|c| col(c.as_str())).collect(); - proj_exprs.push(value_expr); - - LogicalPlanBuilder::from(vector_plan) - .project(proj_exprs)? - .build() -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::data_model::AggregationIdInfo; - use crate::engines::simple::engine::{QueryMetadata, StoreQueryParams, StoreQueryPlan}; - use promql_utilities::data_model::KeyByLabelNames; - - use std::collections::HashMap; - - fn create_test_context( - metric: &str, - statistic: Statistic, - output_labels: Vec<&str>, - aggregation_type: AggregationType, - ) -> QueryExecutionContext { - create_test_context_with_keys( - metric, - statistic, - output_labels, - aggregation_type, - None, - aggregation_type, - ) - } - - fn create_test_context_with_kwargs( - metric: &str, - statistic: Statistic, - output_labels: Vec<&str>, - aggregation_type: AggregationType, - kwargs: HashMap, - ) -> QueryExecutionContext { - let mut ctx = create_test_context_with_keys( - metric, - statistic, - output_labels, - aggregation_type, - None, - aggregation_type, - ); - ctx.metadata.query_kwargs = kwargs; - ctx - } - - fn create_test_context_with_keys( - metric: &str, - statistic: Statistic, - output_labels: Vec<&str>, - aggregation_type_for_value: AggregationType, - keys_query: Option, - aggregation_type_for_key: AggregationType, - ) -> QueryExecutionContext { - // Default: grouping_labels == query_output_labels - create_test_context_with_keys_and_grouping( - metric, - statistic, - output_labels.clone(), - aggregation_type_for_value, - keys_query, - aggregation_type_for_key, - output_labels, - ) - } - - fn create_test_context_with_keys_and_grouping( - metric: &str, - statistic: Statistic, - output_labels: Vec<&str>, - aggregation_type_for_value: AggregationType, - keys_query: Option, - aggregation_type_for_key: AggregationType, - grouping_label_strs: Vec<&str>, - ) -> QueryExecutionContext { - // Default: aggregated_labels = output_labels - grouping_labels - let aggregated: Vec<&str> = output_labels - .iter() - .filter(|l| !grouping_label_strs.contains(l)) - .copied() - .collect(); - create_test_context_full( - metric, - statistic, - output_labels, - aggregation_type_for_value, - keys_query, - aggregation_type_for_key, - grouping_label_strs, - aggregated, - ) - } - - #[allow(clippy::too_many_arguments)] - fn create_test_context_full( - metric: &str, - statistic: Statistic, - output_labels: Vec<&str>, - aggregation_type_for_value: AggregationType, - keys_query: Option, - aggregation_type_for_key: AggregationType, - grouping_label_strs: Vec<&str>, - aggregated_label_strs: Vec<&str>, - ) -> QueryExecutionContext { - let aggregation_id_for_key = match &keys_query { - Some(kq) => kq.aggregation_id, - None => 42, // same as value when no separate keys - }; - let output_labels_vec: Vec = output_labels.into_iter().map(String::from).collect(); - let query_output_labels = KeyByLabelNames { - labels: output_labels_vec, - }; - let grouping_labels = KeyByLabelNames { - labels: grouping_label_strs.into_iter().map(String::from).collect(), - }; - let aggregated_labels = KeyByLabelNames { - labels: aggregated_label_strs - .into_iter() - .map(String::from) - .collect(), - }; - QueryExecutionContext { - metric: metric.to_string(), - metadata: QueryMetadata { - query_output_labels, - statistic_to_compute: statistic, - query_kwargs: HashMap::new(), - }, - store_plan: StoreQueryPlan { - values_query: StoreQueryParams { - metric: metric.to_string(), - aggregation_id: 42, - start_timestamp: 1000, - end_timestamp: 2000, - is_exact_query: true, - }, - keys_query, - }, - agg_info: AggregationIdInfo { - aggregation_id_for_key, - aggregation_id_for_value: 42, - aggregation_type_for_key, - aggregation_type_for_value, - }, - do_merge: false, - spatial_filter: String::new(), - query_time: 2000, - grouping_labels, - aggregated_labels, - } - } - - #[test] - fn test_to_logical_plan_creates_correct_structure() { - let context = create_test_context( - "http_requests", - Statistic::Sum, - vec!["host"], - AggregationType::Sum, - ); - - let plan = context.to_logical_plan().unwrap(); - - // Root should be SummaryInfer - match &plan { - LogicalPlan::Extension(ext) => { - assert_eq!(ext.node.name(), "SummaryInfer"); - } - _ => panic!("Expected Extension node"), - } - } - - #[test] - fn test_to_logical_plan_with_multiple_labels() { - let context = create_test_context( - "http_requests", - Statistic::Sum, - vec!["host", "region", "service"], - AggregationType::Sum, - ); - - let plan = context.to_logical_plan().unwrap(); - - // Verify plan was created successfully - assert!(matches!(plan, LogicalPlan::Extension(_))); - } - - #[test] - fn test_map_statistic_to_infer_operation() { - let context = - create_test_context("test", Statistic::Sum, vec!["host"], AggregationType::Sum); - assert!(matches!( - context.map_statistic_to_infer_operation().unwrap(), - InferOperation::ExtractSum - )); - - let context = create_test_context( - "test", - Statistic::Min, - vec!["host"], - AggregationType::MinMax, - ); - assert!(matches!( - context.map_statistic_to_infer_operation().unwrap(), - InferOperation::ExtractMin - )); - - let context = create_test_context( - "test", - Statistic::Max, - vec!["host"], - AggregationType::MinMax, - ); - assert!(matches!( - context.map_statistic_to_infer_operation().unwrap(), - InferOperation::ExtractMax - )); - } - - #[test] - fn test_map_aggregation_type_to_summary_type() { - let context = - create_test_context("test", Statistic::Sum, vec!["host"], AggregationType::Sum); - assert_eq!( - context.map_aggregation_type_to_summary_type().unwrap(), - SketchType::Sum - ); - - let context = create_test_context( - "test", - Statistic::Increase, - vec!["host"], - AggregationType::Increase, - ); - assert_eq!( - context.map_aggregation_type_to_summary_type().unwrap(), - SketchType::Increase - ); - - let context = create_test_context( - "test", - Statistic::Quantile, - vec!["host"], - AggregationType::DatasketchesKLL, - ); - assert_eq!( - context.map_aggregation_type_to_summary_type().unwrap(), - SketchType::KLL - ); - } - - // ======================================================================== - // Helper to walk the plan tree and collect node names top-down - // ======================================================================== - - fn collect_plan_node_names(plan: &LogicalPlan) -> Vec { - let mut names = Vec::new(); - collect_plan_node_names_recursive(plan, &mut names); - names - } - - fn collect_plan_node_names_recursive(plan: &LogicalPlan, names: &mut Vec) { - match plan { - LogicalPlan::Extension(ext) => { - names.push(ext.node.name().to_string()); - for input in ext.node.inputs() { - collect_plan_node_names_recursive(input, names); - } - } - _ => { - names.push( - format!("{:?}", plan) - .split('(') - .next() - .unwrap_or("Unknown") - .to_string(), - ); - } - } - } - - /// Helper to extract the SummaryMergeMultiple node from a plan tree - fn extract_merge_node(plan: &LogicalPlan) -> Option<&SummaryMergeMultiple> { - match plan { - LogicalPlan::Extension(ext) => { - if let Some(merge) = ext.node.as_any().downcast_ref::() { - return Some(merge); - } - for input in ext.node.inputs() { - if let Some(merge) = extract_merge_node(input) { - return Some(merge); - } - } - None - } - _ => None, - } - } - - /// Helper to extract the PrecomputedSummaryRead node from a plan tree - fn extract_read_node(plan: &LogicalPlan) -> Option<&PrecomputedSummaryRead> { - match plan { - LogicalPlan::Extension(ext) => { - if let Some(read) = ext.node.as_any().downcast_ref::() { - return Some(read); - } - for input in ext.node.inputs() { - if let Some(read) = extract_read_node(input) { - return Some(read); - } - } - None - } - _ => None, - } - } - - /// Helper to extract the SummaryInfer node from the root of a plan - fn extract_infer_node(plan: &LogicalPlan) -> Option<&SummaryInfer> { - match plan { - LogicalPlan::Extension(ext) => ext.node.as_any().downcast_ref::(), - _ => None, - } - } - - /// Count PrecomputedSummaryRead nodes in the plan tree - fn count_read_nodes(plan: &LogicalPlan) -> usize { - let names = collect_plan_node_names(plan); - names - .iter() - .filter(|n| *n == "PrecomputedSummaryRead") - .count() - } - - // ======================================================================== - // MultipleSumAccumulator (HydraSum) tests - // ======================================================================== - - #[test] - fn test_multiple_sum_accumulator_maps_to_hydra_sum() { - let context = create_test_context( - "http_requests", - Statistic::Sum, - vec!["host"], - AggregationType::MultipleSum, - ); - assert_eq!( - context.map_aggregation_type_to_summary_type().unwrap(), - SketchType::MultipleSum - ); - } - - #[test] - fn test_multiple_sum_accumulator_plan_builds() { - // MultipleSumAccumulator is a Hydra (multi-population) accumulator. - // The current plan only builds a single-population SummaryInfer, which - // won't correctly query sub-populations at execution time. - let context = create_test_context( - "http_requests", - Statistic::Sum, - vec!["host"], - AggregationType::MultipleSum, - ); - - let plan = context.to_logical_plan().unwrap(); - let node_names = collect_plan_node_names(&plan); - assert_eq!( - node_names, - vec![ - "SummaryInfer", - "SummaryMergeMultiple", - "PrecomputedSummaryRead" - ] - ); - - // Verify the summary type propagates correctly through the plan - let merge = extract_merge_node(&plan).expect("Should have a SummaryMergeMultiple node"); - assert_eq!(*merge.summary_type(), SketchType::MultipleSum); - - let read = extract_read_node(&plan).expect("Should have a PrecomputedSummaryRead node"); - assert_eq!(*read.summary_type(), SketchType::MultipleSum); - } - - #[test] - fn test_multiple_sum_accumulator_single_pop_no_subkeys() { - // MultipleSumAccumulator without a separate keys_query stays single-population. - // No sub-key columns because there's no keys branch to enumerate from. - // To properly query Hydra types, a keys_query with a DeltaSetAggregator - // should be provided — see test_delta_set_aggregator_dual_input_plan. - let context = create_test_context( - "http_requests", - Statistic::Sum, - vec!["host"], - AggregationType::MultipleSum, - ); - - let plan = context.to_logical_plan().unwrap(); - - let infer = extract_infer_node(&plan).expect("Root should be SummaryInfer"); - assert!( - infer.group_key_columns.is_empty(), - "Single-pop (no keys_query): group_key_columns should be empty" - ); - assert!( - infer.keys_input.is_none(), - "Single-pop: should not have keys_input" - ); - } - - // ======================================================================== - // CountMinSketch for values_plan tests - // ======================================================================== - - #[test] - fn test_count_min_sketch_maps_correctly() { - let context = create_test_context( - "http_requests", - Statistic::Count, - vec!["host"], - AggregationType::CountMinSketch, - ); - assert_eq!( - context.map_aggregation_type_to_summary_type().unwrap(), - SketchType::CountMinSketch - ); - } - - #[test] - fn test_count_min_sketch_plan_builds() { - // CountMinSketch is a multi-population frequency sketch. - // Like Hydra types, it requires a sub-key to query (FrequencyEstimate, etc.) - let context = create_test_context( - "http_requests", - Statistic::Count, - vec!["host"], - AggregationType::CountMinSketch, - ); - - let plan = context.to_logical_plan().unwrap(); - let node_names = collect_plan_node_names(&plan); - assert_eq!( - node_names, - vec![ - "SummaryInfer", - "SummaryMergeMultiple", - "PrecomputedSummaryRead" - ] - ); - - // Verify summary type - let merge = extract_merge_node(&plan).expect("Should have SummaryMergeMultiple"); - assert_eq!(*merge.summary_type(), SketchType::CountMinSketch); - } - - #[test] - fn test_count_min_sketch_single_pop_no_subkeys() { - // CountMinSketch without a separate keys_query stays single-population. - // Same as MultipleSumAccumulator — need a keys_query for dual-input. - let context = create_test_context( - "http_requests", - Statistic::Count, - vec!["host"], - AggregationType::CountMinSketch, - ); - - let plan = context.to_logical_plan().unwrap(); - - let infer = extract_infer_node(&plan).expect("Root should be SummaryInfer"); - assert!(infer.group_key_columns.is_empty()); - assert!(infer.keys_input.is_none()); - } - - // ======================================================================== - // DeltaSetAggregator for keys_plan tests - // ======================================================================== - - #[test] - fn test_delta_set_aggregator_dual_input_plan() { - // When keys_query is set with a different agg_id (DeltaSetAggregator for key - // enumeration), to_logical_plan() builds a dual-input SummaryInfer with both - // a values branch and a keys branch. - let keys_query = StoreQueryParams { - metric: "http_requests".to_string(), - aggregation_id: 99, // Different agg ID for keys - start_timestamp: 0, - end_timestamp: 2000, - is_exact_query: false, - }; - let context = create_test_context_with_keys( - "http_requests", - Statistic::Sum, - vec!["host"], - AggregationType::MultipleSum, // values use Hydra - Some(keys_query), - AggregationType::DeltaSetAggregator, // keys use DeltaSet - ); - - let plan = context.to_logical_plan().unwrap(); - - // The plan tree should now have 5 nodes: SummaryInfer with 2 branches - let node_names = collect_plan_node_names(&plan); - assert_eq!( - node_names, - vec![ - "SummaryInfer", - "SummaryMergeMultiple", // values branch - "PrecomputedSummaryRead", // values read - "SummaryMergeMultiple", // keys branch - "PrecomputedSummaryRead", // keys read - ] - ); - - // Verify there are 2 PrecomputedSummaryRead nodes - assert_eq!(count_read_nodes(&plan), 2); - - // The SummaryInfer should have a keys_input - let infer = extract_infer_node(&plan).expect("Root should be SummaryInfer"); - assert!( - infer.keys_input.is_some(), - "SummaryInfer should have keys_input" - ); - } - - // ======================================================================== - // HydraKLL for values_plan + DeltaSetAggregator for keys_plan - // ======================================================================== - - #[test] - fn test_hydra_kll_with_delta_set_keys_plan_builds() { - // HydraKLL for quantile queries with DeltaSetAggregator for key enumeration. - // This is a realistic configuration: HydraKLL stores per-key quantile sketches, - // and DeltaSetAggregator tracks which keys exist. - // grouping_labels = ["host"], sub-keys = ["endpoint"] - let mut kwargs = HashMap::new(); - kwargs.insert("quantile".to_string(), "0.95".to_string()); - - let keys_query = StoreQueryParams { - metric: "request_duration".to_string(), - aggregation_id: 99, - start_timestamp: 0, - end_timestamp: 2000, - is_exact_query: false, - }; - let mut context = create_test_context_with_keys_and_grouping( - "request_duration", - Statistic::Quantile, - vec!["host", "endpoint"], // query_output_labels - AggregationType::HydraKLL, - Some(keys_query), - AggregationType::DeltaSetAggregator, - vec!["host"], // grouping_labels (spatial) - ); - context.metadata.query_kwargs = kwargs; - - let plan = context.to_logical_plan().unwrap(); - - // Verify dual-input plan structure (5 nodes) - let node_names = collect_plan_node_names(&plan); - assert_eq!( - node_names, - vec![ - "SummaryInfer", - "SummaryMergeMultiple", // values branch - "PrecomputedSummaryRead", // values read - "SummaryMergeMultiple", // keys branch - "PrecomputedSummaryRead", // keys read - ] - ); - - // Verify types propagate in the values branch - let merge = extract_merge_node(&plan).expect("Should have SummaryMergeMultiple"); - assert_eq!(*merge.summary_type(), SketchType::HydraKLL); - - let read = extract_read_node(&plan).expect("Should have PrecomputedSummaryRead"); - assert_eq!(*read.summary_type(), SketchType::HydraKLL); - // Values branch uses grouping_labels, not query_output_labels - assert_eq!(read.output_labels(), &["host"]); - - // Verify SummaryInfer has sub-key columns - let infer = extract_infer_node(&plan).expect("Root should be SummaryInfer"); - assert_eq!(infer.group_key_columns, vec!["endpoint"]); - } - - #[test] - fn test_hydra_kll_with_delta_set_keys_dual_input_with_subkeys() { - // HydraKLL with DeltaSetAggregator keys, with sub-key columns. - // output_labels = ["host", "endpoint"], grouping_labels = ["host"] - // => sub_key_labels = ["endpoint"] - let keys_query = StoreQueryParams { - metric: "request_duration".to_string(), - aggregation_id: 99, - start_timestamp: 0, - end_timestamp: 2000, - is_exact_query: false, - }; - let context = create_test_context_with_keys_and_grouping( - "request_duration", - Statistic::Quantile, - vec!["host", "endpoint"], // query_output_labels - AggregationType::HydraKLL, - Some(keys_query), - AggregationType::DeltaSetAggregator, - vec!["host"], // grouping_labels (spatial store labels) - ); - - let plan = context.to_logical_plan().unwrap(); - - // Plan should have 2 PrecomputedSummaryRead nodes - assert_eq!(count_read_nodes(&plan), 2); - - // SummaryInfer should have keys_input and group_key_columns = ["endpoint"] - let infer = extract_infer_node(&plan).expect("Root should be SummaryInfer"); - assert!(infer.keys_input.is_some(), "Should have keys_input"); - assert_eq!( - infer.group_key_columns, - vec!["endpoint"], - "Sub-key columns should be query_output_labels minus grouping_labels" - ); - } - - // ======================================================================== - // Quantile kwargs propagation test - // ======================================================================== - - #[test] - fn test_quantile_kwargs_propagate_to_infer_operation() { - let mut kwargs = HashMap::new(); - kwargs.insert("quantile".to_string(), "0.99".to_string()); - let context = create_test_context_with_kwargs( - "latency", - Statistic::Quantile, - vec!["host"], - AggregationType::DatasketchesKLL, - kwargs, - ); - - match context.map_statistic_to_infer_operation().unwrap() { - InferOperation::Quantile(q) => { - // 0.99 * 10000 = 9900 - assert_eq!(q, 9900, "Expected q=9900 (0.99), got {}", q); - } - other => panic!("Expected Quantile, got {:?}", other), - } - } - - #[test] - fn test_quantile_defaults_to_median_when_no_kwargs() { - let context = create_test_context( - "latency", - Statistic::Quantile, - vec!["host"], - AggregationType::DatasketchesKLL, - ); - - match context.map_statistic_to_infer_operation().unwrap() { - InferOperation::Quantile(q) => { - // 0.5 * 10000 = 5000 - assert_eq!(q, 5000, "Expected q=5000 (0.5), got {}", q); - } - other => panic!("Expected Quantile, got {:?}", other), - } - } -} diff --git a/asap-query-engine/src/engines/mod.rs b/asap-query-engine/src/engines/mod.rs index 0dc307c2..20a0b088 100644 --- a/asap-query-engine/src/engines/mod.rs +++ b/asap-query-engine/src/engines/mod.rs @@ -24,9 +24,7 @@ //! `crate::routing::QueryEngine` impl returns. pub mod gorilla; -pub mod logical; pub mod no_data_archive; -pub mod physical; pub mod prometheus; pub mod query_result; pub mod simple; diff --git a/asap-query-engine/src/engines/physical/accumulator_serde.rs b/asap-query-engine/src/engines/physical/accumulator_serde.rs deleted file mode 100644 index cc7f85ea..00000000 --- a/asap-query-engine/src/engines/physical/accumulator_serde.rs +++ /dev/null @@ -1,314 +0,0 @@ -//! Accumulator Serialization/Deserialization Registry -//! -//! Provides functions to deserialize bytes back to accumulator objects -//! based on the SummaryType (SketchType). -//! -//! Note: This module assumes accumulators are serialized using the Arroyo format -//! (MessagePack serialization via rmp-serde). - -use datafusion::error::DataFusionError; -use datafusion_summary_library::SketchType; - -use crate::data_model::{MultipleSubpopulationAggregate, SingleSubpopulationAggregate}; -use crate::precompute_operators::{ - CountMinSketchAccumulator, DatasketchesKLLAccumulator, DeltaSetAggregatorAccumulator, - HydraKllSketchAccumulator, MultipleIncreaseAccumulator, MultipleSumAccumulator, - SetAggregatorAccumulator, SumAccumulator, -}; -use crate::AggregateCore; - -/// Deserialize bytes to an accumulator based on the summary type. -/// -/// This function dispatches to the appropriate accumulator deserializer -/// based on the SummaryType enum. Expects Arroyo format (MessagePack). -/// -/// # Arguments -/// * `bytes` - Serialized accumulator data (Arroyo/MessagePack format) -/// * `summary_type` - Type of the summary (determines which deserializer to use) -/// -/// # Returns -/// A boxed AggregateCore trait object -pub fn deserialize_accumulator( - bytes: &[u8], - summary_type: &SketchType, -) -> Result, DataFusionError> { - match summary_type { - // Single-population exact aggregators - SketchType::Sum => { - let acc = SumAccumulator::deserialize_from_bytes_arroyo(bytes).map_err(|e| { - DataFusionError::Internal(format!("Failed to deserialize Sum: {}", e)) - })?; - Ok(Box::new(acc)) - } - SketchType::Increase => Err(DataFusionError::NotImplemented( - "Increase Arroyo deserialization not implemented".to_string(), - )), - SketchType::MinMax => Err(DataFusionError::NotImplemented( - "MinMax Arroyo deserialization not implemented".to_string(), - )), - - // Quantile sketches - SketchType::KLL => { - let acc = - DatasketchesKLLAccumulator::deserialize_from_bytes_arroyo(bytes).map_err(|e| { - DataFusionError::Internal(format!("Failed to deserialize KLL: {}", e)) - })?; - Ok(Box::new(acc)) - } - SketchType::HydraKLL => { - let acc = - HydraKllSketchAccumulator::deserialize_from_bytes_arroyo(bytes).map_err(|e| { - DataFusionError::Internal(format!("Failed to deserialize HydraKLL: {}", e)) - })?; - Ok(Box::new(acc)) - } - - // Set aggregators - SketchType::SetAggregator => { - let acc = - SetAggregatorAccumulator::deserialize_from_bytes_arroyo(bytes).map_err(|e| { - DataFusionError::Internal(format!("Failed to deserialize SetAggregator: {}", e)) - })?; - Ok(Box::new(acc)) - } - SketchType::DeltaSetAggregator => { - let acc = DeltaSetAggregatorAccumulator::deserialize_from_bytes_arroyo(bytes).map_err( - |e| { - DataFusionError::Internal(format!( - "Failed to deserialize DeltaSetAggregator: {}", - e - )) - }, - )?; - Ok(Box::new(acc)) - } - - // Multi-population exact aggregators - SketchType::MultipleIncrease => { - let acc = - MultipleIncreaseAccumulator::deserialize_from_bytes_arroyo(bytes).map_err(|e| { - DataFusionError::Internal(format!( - "Failed to deserialize MultipleIncrease: {}", - e - )) - })?; - Ok(Box::new(acc)) - } - SketchType::MultipleSum => { - let acc = - MultipleSumAccumulator::deserialize_from_bytes_arroyo(bytes).map_err(|e| { - DataFusionError::Internal(format!("Failed to deserialize MultipleSum: {}", e)) - })?; - Ok(Box::new(acc)) - } - - // Frequency sketches - SketchType::CountMinSketch => { - let acc = - CountMinSketchAccumulator::deserialize_from_bytes_arroyo(bytes).map_err(|e| { - DataFusionError::Internal(format!( - "Failed to deserialize CountMinSketch: {}", - e - )) - })?; - Ok(Box::new(acc)) - } - - // Sketches that aren't implemented yet - _ => Err(DataFusionError::NotImplemented(format!( - "Accumulator deserialization not implemented for: {:?}", - summary_type - ))), - } -} - -/// Serialize an accumulator to bytes (native format). -/// -/// This is a convenience wrapper that calls serialize_to_bytes on the accumulator. -pub fn serialize_accumulator(acc: &dyn AggregateCore) -> Vec { - acc.serialize_to_bytes() -} - -/// Serialize an accumulator to Arroyo-compatible bytes (MessagePack format). -/// -/// For accumulators whose native serialize_to_bytes already uses MessagePack, -/// this delegates to serialize_to_bytes. For others (SumAccumulator, -/// SetAggregatorAccumulator, MultipleIncreaseAccumulator), this uses -/// their serialize_to_bytes_arroyo method. -pub fn serialize_accumulator_arroyo(acc: &dyn AggregateCore) -> Vec { - // Try to downcast to types that have a separate arroyo format - if let Some(sum_acc) = acc.as_any().downcast_ref::() { - return sum_acc.serialize_to_bytes_arroyo(); - } - if let Some(set_acc) = acc.as_any().downcast_ref::() { - return set_acc.serialize_to_bytes_arroyo(); - } - if let Some(inc_acc) = acc.as_any().downcast_ref::() { - return inc_acc.serialize_to_bytes_arroyo(); - } - if let Some(ms_acc) = acc.as_any().downcast_ref::() { - return ms_acc.serialize_to_bytes_arroyo(); - } - // All other accumulators already use MessagePack in serialize_to_bytes - acc.serialize_to_bytes() -} - -/// Deserialize bytes to a SingleSubpopulationAggregate for querying. -/// -/// This function returns a trait object that supports the query method. -/// Only works for single-subpopulation accumulators (Sum, Increase, MinMax, etc.). -/// -/// Note: Uses Arroyo/MessagePack format. -pub fn deserialize_single_subpopulation( - bytes: &[u8], - summary_type: &SketchType, -) -> Result, DataFusionError> { - match summary_type { - SketchType::Sum => { - let acc = SumAccumulator::deserialize_from_bytes_arroyo(bytes).map_err(|e| { - DataFusionError::Internal(format!("Failed to deserialize Sum: {}", e)) - })?; - Ok(Box::new(acc)) - } - SketchType::Increase => Err(DataFusionError::NotImplemented( - "Increase Arroyo deserialization not implemented".to_string(), - )), - SketchType::MinMax => Err(DataFusionError::NotImplemented( - "MinMax Arroyo deserialization not implemented".to_string(), - )), - SketchType::KLL => { - let acc = - DatasketchesKLLAccumulator::deserialize_from_bytes_arroyo(bytes).map_err(|e| { - DataFusionError::Internal(format!("Failed to deserialize KLL: {}", e)) - })?; - Ok(Box::new(acc)) - } - _ => Err(DataFusionError::NotImplemented(format!( - "SingleSubpopulationAggregate deserialization not implemented for: {:?}", - summary_type - ))), - } -} - -/// Deserialize bytes to a MultipleSubpopulationAggregate for querying. -/// -/// This function returns a trait object that supports querying by sub-key. -/// Works for multi-population accumulators (Hydra types, CountMinSketch, etc.). -/// -/// Note: Uses Arroyo/MessagePack format. -pub fn deserialize_multiple_subpopulation( - bytes: &[u8], - summary_type: &SketchType, -) -> Result, DataFusionError> { - match summary_type { - SketchType::MultipleIncrease => { - let acc = - MultipleIncreaseAccumulator::deserialize_from_bytes_arroyo(bytes).map_err(|e| { - DataFusionError::Internal(format!( - "Failed to deserialize MultipleIncrease: {}", - e - )) - })?; - Ok(Box::new(acc)) - } - SketchType::MultipleSum => { - let acc = - MultipleSumAccumulator::deserialize_from_bytes_arroyo(bytes).map_err(|e| { - DataFusionError::Internal(format!("Failed to deserialize MultipleSum: {}", e)) - })?; - Ok(Box::new(acc)) - } - SketchType::HydraKLL => { - let acc = - HydraKllSketchAccumulator::deserialize_from_bytes_arroyo(bytes).map_err(|e| { - DataFusionError::Internal(format!("Failed to deserialize HydraKLL: {}", e)) - })?; - Ok(Box::new(acc)) - } - SketchType::CountMinSketch => { - let acc = - CountMinSketchAccumulator::deserialize_from_bytes_arroyo(bytes).map_err(|e| { - DataFusionError::Internal(format!( - "Failed to deserialize CountMinSketch: {}", - e - )) - })?; - Ok(Box::new(acc)) - } - _ => Err(DataFusionError::NotImplemented(format!( - "MultipleSubpopulationAggregate deserialization not implemented for: {:?}", - summary_type - ))), - } -} - -/// Deserialize bytes to a keys accumulator (DeltaSetAggregator/SetAggregator). -/// -/// Returns a boxed AggregateCore whose `get_keys()` method enumerates the sub-keys -/// stored in the accumulator. -pub fn deserialize_keys_accumulator( - bytes: &[u8], - summary_type: &SketchType, -) -> Result, DataFusionError> { - match summary_type { - SketchType::DeltaSetAggregator => { - let acc = DeltaSetAggregatorAccumulator::deserialize_from_bytes_arroyo(bytes).map_err( - |e| { - DataFusionError::Internal(format!( - "Failed to deserialize DeltaSetAggregator: {}", - e - )) - }, - )?; - Ok(Box::new(acc)) - } - SketchType::SetAggregator => { - let acc = - SetAggregatorAccumulator::deserialize_from_bytes_arroyo(bytes).map_err(|e| { - DataFusionError::Internal(format!("Failed to deserialize SetAggregator: {}", e)) - })?; - Ok(Box::new(acc)) - } - _ => Err(DataFusionError::NotImplemented(format!( - "Keys accumulator deserialization not supported for: {:?}", - summary_type - ))), - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::data_model::AggregationType; - - // Helper to serialize f64 as MessagePack (Arroyo format) - fn serialize_f64_arroyo(value: f64) -> Vec { - rmp_serde::to_vec(&value).unwrap() - } - - #[test] - fn test_deserialize_sum_accumulator() { - // SumAccumulator::deserialize_from_bytes_arroyo expects MessagePack f64 - let bytes = serialize_f64_arroyo(42.0); - - let restored = deserialize_accumulator(&bytes, &SketchType::Sum).unwrap(); - - assert_eq!(restored.get_accumulator_type(), AggregationType::Sum); - } - - #[test] - fn test_deserialize_minmax_returns_not_implemented() { - let bytes = vec![1, 2, 3, 4]; - let result = deserialize_accumulator(&bytes, &SketchType::MinMax); - - assert!(result.is_err()); - } - - #[test] - fn test_deserialize_unsupported_type() { - let bytes = vec![1, 2, 3, 4]; - let result = deserialize_accumulator(&bytes, &SketchType::HLL); - - assert!(result.is_err()); - } -} diff --git a/asap-query-engine/src/engines/physical/conversion.rs b/asap-query-engine/src/engines/physical/conversion.rs deleted file mode 100644 index 26cde347..00000000 --- a/asap-query-engine/src/engines/physical/conversion.rs +++ /dev/null @@ -1,452 +0,0 @@ -//! Conversion Utilities Module -//! -//! Provides functions to convert between store results and Arrow RecordBatches. -//! This enables DataFusion physical operators to work with precomputed outputs. - -use arrow::array::{ArrayRef, BinaryBuilder, Float64Array, StringBuilder}; -use arrow::datatypes::{DataType, Field, Schema}; -use arrow::record_batch::RecordBatch; -use datafusion::error::DataFusionError; -use std::collections::HashMap; -use std::sync::Arc; - -use crate::data_model::KeyByLabelValues; -use crate::engines::physical::accumulator_serde::serialize_accumulator_arroyo; -use crate::stores::traits::TimestampedBucketsMap; - -/// Convert store query results to an Arrow RecordBatch. -/// -/// The output schema is: [label_columns..., sketch (Binary)] -/// Each row represents one group key with its serialized accumulator. -/// -/// # Arguments -/// * `store_result` - HashMap from group key to accumulators -/// * `label_names` - Names of the label columns (in order) -/// -/// # Returns -/// A RecordBatch with label columns (Utf8) and a sketch column (Binary) -pub fn store_result_to_record_batch( - store_result: &TimestampedBucketsMap, - label_names: &[String], -) -> Result { - // Build arrays for each label column - let mut label_builders: Vec = - label_names.iter().map(|_| StringBuilder::new()).collect(); - - // Build array for sketch column - let mut sketch_builder = BinaryBuilder::new(); - - for (key_opt, timestamped_buckets) in store_result { - // Handle each accumulator for this key - for (_timestamps, acc) in timestamped_buckets { - // Add label values - if let Some(key) = key_opt { - for (i, label_value) in key.labels.iter().enumerate() { - if i < label_builders.len() { - label_builders[i].append_value(label_value); - } - } - // Pad with empty strings if key has fewer labels than expected - for item in label_builders.iter_mut().skip(key.labels.len()) { - item.append_value(""); - } - } else { - // No key - use empty strings for all labels - for builder in &mut label_builders { - builder.append_value(""); - } - } - - // Serialize accumulator to Arroyo-compatible bytes (MessagePack) - // so downstream operators can deserialize with deserialize_from_bytes_arroyo - let bytes = serialize_accumulator_arroyo(acc.as_ref()); - sketch_builder.append_value(&bytes); - } - } - - // Build schema - let mut fields: Vec = label_names - .iter() - .map(|name| Field::new(name, DataType::Utf8, true)) - .collect(); - fields.push(Field::new("sketch", DataType::Binary, false)); - let schema = Arc::new(Schema::new(fields)); - - // Build columns - let mut columns: Vec = label_builders - .iter_mut() - .map(|b| Arc::new(b.finish()) as ArrayRef) - .collect(); - columns.push(Arc::new(sketch_builder.finish())); - - RecordBatch::try_new(schema, columns) - .map_err(|e| DataFusionError::Internal(format!("Failed to create RecordBatch: {}", e))) -} - -/// Convert a RecordBatch with inferred values back to a result map. -/// -/// The input schema is expected to be: [label_columns..., value_column (Float64)] -/// -/// # Arguments -/// * `batch` - RecordBatch with label columns and a value column -/// * `label_names` - Names of the label columns (to identify which columns are labels) -/// * `value_column` - Name of the column containing the inferred values -/// -/// # Returns -/// A HashMap from group key to the extracted value -pub fn record_batch_to_result_map( - batch: &RecordBatch, - label_names: &[&str], - value_column: &str, -) -> Result, f64>, DataFusionError> { - let mut result: HashMap, f64> = HashMap::new(); - - // Find the value column - let value_col_idx = batch - .schema() - .fields() - .iter() - .position(|f| f.name() == value_column) - .ok_or_else(|| { - DataFusionError::Internal(format!( - "No '{}' column found in batch schema", - value_column - )) - })?; - - let value_array = batch - .column(value_col_idx) - .as_any() - .downcast_ref::() - .ok_or_else(|| { - DataFusionError::Internal(format!("'{}' column is not Float64", value_column)) - })?; - - // Find label column indices - let label_indices: Vec = label_names - .iter() - .filter_map(|name| { - batch - .schema() - .fields() - .iter() - .position(|f| f.name() == *name) - }) - .collect(); - - for row_idx in 0..batch.num_rows() { - // Extract label values for this row - let labels: Vec = label_indices - .iter() - .map(|&col_idx| { - let col = batch.column(col_idx); - // Try to extract string value - if let Some(str_array) = col.as_any().downcast_ref::() { - str_array.value(row_idx).to_string() - } else { - String::new() - } - }) - .collect(); - - let key = if labels.is_empty() || labels.iter().all(|l| l.is_empty()) { - None - } else { - Some(KeyByLabelValues { labels }) - }; - - let value = value_array.value(row_idx); - result.insert(key, value); - } - - Ok(result) -} - -/// Helper function to count total rows in store result -pub fn count_store_result_rows(store_result: &TimestampedBucketsMap) -> usize { - store_result.values().map(|v| v.len()).sum() -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::precompute_operators::SumAccumulator; - use crate::stores::traits::TimestampedBucket; - - fn make_bucket(acc: Arc) -> TimestampedBucket { - ((0, 0), acc) - } - - #[test] - fn test_store_result_to_record_batch_basic() { - let mut store_result: TimestampedBucketsMap = HashMap::new(); - - let key1 = KeyByLabelValues { - labels: vec!["host-a".to_string()], - }; - let acc1 = Arc::new(SumAccumulator::with_sum(100.0)) as Arc; - store_result.insert(Some(key1), vec![make_bucket(acc1)]); - - let key2 = KeyByLabelValues { - labels: vec!["host-b".to_string()], - }; - let acc2 = Arc::new(SumAccumulator::with_sum(200.0)) as Arc; - store_result.insert(Some(key2), vec![make_bucket(acc2)]); - - let label_names = vec!["host".to_string()]; - let batch = store_result_to_record_batch(&store_result, &label_names).unwrap(); - - assert_eq!(batch.num_rows(), 2); - assert_eq!(batch.num_columns(), 2); // host, sketch - } - - #[test] - fn test_store_result_to_record_batch_multiple_labels() { - let mut store_result: TimestampedBucketsMap = HashMap::new(); - - let key1 = KeyByLabelValues { - labels: vec!["host-a".to_string(), "region-1".to_string()], - }; - let acc1 = Arc::new(SumAccumulator::with_sum(100.0)) as Arc; - store_result.insert(Some(key1), vec![make_bucket(acc1)]); - - let label_names = vec!["host".to_string(), "region".to_string()]; - let batch = store_result_to_record_batch(&store_result, &label_names).unwrap(); - - assert_eq!(batch.num_rows(), 1); - assert_eq!(batch.num_columns(), 3); // host, region, sketch - } - - #[test] - fn test_store_result_to_record_batch_no_key() { - let mut store_result: TimestampedBucketsMap = HashMap::new(); - - let acc = Arc::new(SumAccumulator::with_sum(500.0)) as Arc; - store_result.insert(None, vec![make_bucket(acc)]); - - let label_names: Vec = vec![]; - let batch = store_result_to_record_batch(&store_result, &label_names).unwrap(); - - assert_eq!(batch.num_rows(), 1); - assert_eq!(batch.num_columns(), 1); // just sketch - } - - #[test] - fn test_record_batch_to_result_map() { - // Create a test batch with [host, value] - let host_array = arrow::array::StringArray::from(vec!["host-a", "host-b"]); - let value_array = Float64Array::from(vec![100.0, 200.0]); - - let schema = Arc::new(Schema::new(vec![ - Field::new("host", DataType::Utf8, true), - Field::new("value", DataType::Float64, false), - ])); - - let batch = RecordBatch::try_new( - schema, - vec![ - Arc::new(host_array) as ArrayRef, - Arc::new(value_array) as ArrayRef, - ], - ) - .unwrap(); - - let result = record_batch_to_result_map(&batch, &["host"], "value").unwrap(); - - assert_eq!(result.len(), 2); - - let key_a = KeyByLabelValues { - labels: vec!["host-a".to_string()], - }; - assert_eq!(result.get(&Some(key_a)), Some(&100.0)); - - let key_b = KeyByLabelValues { - labels: vec!["host-b".to_string()], - }; - assert_eq!(result.get(&Some(key_b)), Some(&200.0)); - } - - // ======================================================================== - // Edge case tests - // ======================================================================== - - #[test] - fn test_store_result_to_record_batch_empty() { - let store_result: TimestampedBucketsMap = HashMap::new(); - let label_names = vec!["host".to_string()]; - let batch = store_result_to_record_batch(&store_result, &label_names).unwrap(); - assert_eq!(batch.num_rows(), 0); - assert_eq!(batch.num_columns(), 2); // host + sketch - } - - #[test] - fn test_store_result_to_record_batch_five_labels() { - let mut store_result: TimestampedBucketsMap = HashMap::new(); - let key = KeyByLabelValues { - labels: vec![ - "a".to_string(), - "b".to_string(), - "c".to_string(), - "d".to_string(), - "e".to_string(), - ], - }; - store_result.insert( - Some(key), - vec![make_bucket( - Arc::new(SumAccumulator::with_sum(1.0)) as Arc - )], - ); - let label_names: Vec = vec!["l1", "l2", "l3", "l4", "l5"] - .into_iter() - .map(String::from) - .collect(); - let batch = store_result_to_record_batch(&store_result, &label_names).unwrap(); - assert_eq!(batch.num_rows(), 1); - assert_eq!(batch.num_columns(), 6); // 5 labels + sketch - } - - #[test] - fn test_store_result_to_record_batch_special_chars() { - let mut store_result: TimestampedBucketsMap = HashMap::new(); - let key = KeyByLabelValues { - labels: vec!["host,with,commas".to_string(), "région-1".to_string()], - }; - store_result.insert( - Some(key), - vec![make_bucket( - Arc::new(SumAccumulator::with_sum(42.0)) as Arc - )], - ); - let label_names = vec!["host".to_string(), "region".to_string()]; - let batch = store_result_to_record_batch(&store_result, &label_names).unwrap(); - assert_eq!(batch.num_rows(), 1); - - // Verify the special characters survived - let host_col = batch - .column(0) - .as_any() - .downcast_ref::() - .unwrap(); - assert_eq!(host_col.value(0), "host,with,commas"); - let region_col = batch - .column(1) - .as_any() - .downcast_ref::() - .unwrap(); - assert_eq!(region_col.value(0), "région-1"); - } - - #[test] - fn test_store_result_to_record_batch_multiple_timestamps_per_key() { - let mut store_result: TimestampedBucketsMap = HashMap::new(); - let key = KeyByLabelValues { - labels: vec!["host-a".to_string()], - }; - // 3 buckets for the same key - store_result.insert( - Some(key), - vec![ - ( - (100, 200), - Arc::new(SumAccumulator::with_sum(10.0)) as Arc, - ), - ( - (200, 300), - Arc::new(SumAccumulator::with_sum(20.0)) as Arc, - ), - ( - (300, 400), - Arc::new(SumAccumulator::with_sum(30.0)) as Arc, - ), - ], - ); - let label_names = vec!["host".to_string()]; - let batch = store_result_to_record_batch(&store_result, &label_names).unwrap(); - assert_eq!(batch.num_rows(), 3, "3 buckets should produce 3 rows"); - } - - #[test] - fn test_record_batch_to_result_map_no_labels() { - let value_array = Float64Array::from(vec![42.0]); - let schema = Arc::new(Schema::new(vec![Field::new( - "value", - DataType::Float64, - false, - )])); - let batch = RecordBatch::try_new(schema, vec![Arc::new(value_array) as ArrayRef]).unwrap(); - - let result = record_batch_to_result_map(&batch, &[], "value").unwrap(); - assert_eq!(result.len(), 1); - assert_eq!(result.get(&None), Some(&42.0)); - } - - #[test] - fn test_record_batch_to_result_map_missing_value_column() { - let host_array = arrow::array::StringArray::from(vec!["host-a"]); - let schema = Arc::new(Schema::new(vec![Field::new("host", DataType::Utf8, true)])); - let batch = RecordBatch::try_new(schema, vec![Arc::new(host_array) as ArrayRef]).unwrap(); - - let result = record_batch_to_result_map(&batch, &["host"], "value"); - assert!(result.is_err(), "Missing value column should produce error"); - let err_msg = format!("{}", result.unwrap_err()); - assert!( - err_msg.contains("value"), - "Error should mention 'value' column" - ); - } - - #[test] - fn test_record_batch_to_result_map_wrong_type() { - // value column is Utf8 instead of Float64 - let host_array = arrow::array::StringArray::from(vec!["host-a"]); - let value_array = arrow::array::StringArray::from(vec!["not_a_number"]); - let schema = Arc::new(Schema::new(vec![ - Field::new("host", DataType::Utf8, true), - Field::new("value", DataType::Utf8, false), - ])); - let batch = RecordBatch::try_new( - schema, - vec![ - Arc::new(host_array) as ArrayRef, - Arc::new(value_array) as ArrayRef, - ], - ) - .unwrap(); - - let result = record_batch_to_result_map(&batch, &["host"], "value"); - assert!(result.is_err(), "Wrong type for value column should error"); - } - - #[test] - fn test_count_store_result_rows() { - let mut store_result: TimestampedBucketsMap = HashMap::new(); - let key1 = KeyByLabelValues { - labels: vec!["a".to_string()], - }; - let key2 = KeyByLabelValues { - labels: vec!["b".to_string()], - }; - store_result.insert( - Some(key1), - vec![ - make_bucket( - Arc::new(SumAccumulator::with_sum(1.0)) as Arc - ), - make_bucket( - Arc::new(SumAccumulator::with_sum(2.0)) as Arc - ), - ], - ); - store_result.insert( - Some(key2), - vec![make_bucket( - Arc::new(SumAccumulator::with_sum(3.0)) as Arc - )], - ); - assert_eq!(count_store_result_rows(&store_result), 3); - - let empty: TimestampedBucketsMap = HashMap::new(); - assert_eq!(count_store_result_rows(&empty), 0); - } -} diff --git a/asap-query-engine/src/engines/physical/mod.rs b/asap-query-engine/src/engines/physical/mod.rs deleted file mode 100644 index 4cc5e54b..00000000 --- a/asap-query-engine/src/engines/physical/mod.rs +++ /dev/null @@ -1,29 +0,0 @@ -//! Physical Execution Operators for DataFusion -//! -//! This module provides physical execution plan operators that implement -//! DataFusion's ExecutionPlan trait for precomputed summary operations. - -pub mod accumulator_serde; -pub mod conversion; -pub mod planner; -pub mod precomputed_summary_read_exec; -pub mod summary_infer_exec; -pub mod summary_merge_multiple_exec; - -pub use planner::{CustomQueryPlanner, QueryEngineExtensionPlanner}; -pub use precomputed_summary_read_exec::PrecomputedSummaryReadExec; -pub use summary_infer_exec::SummaryInferExec; -pub use summary_merge_multiple_exec::SummaryMergeMultipleExec; - -use arrow::datatypes::SchemaRef; - -/// Format an Arrow schema as a compact string for debug logging. -/// Example: `{host: Utf8, region: Utf8, sketch: Binary}` -pub(crate) fn format_schema(schema: &SchemaRef) -> String { - let fields: Vec = schema - .fields() - .iter() - .map(|f| format!("{}: {:?}", f.name(), f.data_type())) - .collect(); - format!("{{{}}}", fields.join(", ")) -} diff --git a/asap-query-engine/src/engines/physical/planner.rs b/asap-query-engine/src/engines/physical/planner.rs deleted file mode 100644 index b0dc7b9e..00000000 --- a/asap-query-engine/src/engines/physical/planner.rs +++ /dev/null @@ -1,180 +0,0 @@ -//! Extension Planner for QueryEngineRust -//! -//! This module provides an ExtensionPlanner implementation that converts -//! custom logical operators (PrecomputedSummaryRead, SummaryMergeMultiple) -//! into their physical execution counterparts. - -use async_trait::async_trait; -use datafusion::error::DataFusionError; -use datafusion::execution::context::SessionState; -use datafusion::logical_expr::{LogicalPlan, UserDefinedLogicalNode}; -use datafusion::physical_plan::ExecutionPlan; -use datafusion::physical_planner::{DefaultPhysicalPlanner, ExtensionPlanner, PhysicalPlanner}; -use datafusion_summary_library::{ - PrecomputedSummaryRead, SketchType, SummaryInfer, SummaryMergeMultiple, -}; -use std::fmt; -use std::sync::Arc; - -use super::{PrecomputedSummaryReadExec, SummaryInferExec, SummaryMergeMultipleExec}; -use crate::stores::Store; - -/// Extension planner that handles custom logical operators for QueryEngineRust. -/// -/// This planner knows how to convert: -/// - PrecomputedSummaryRead -> PrecomputedSummaryReadExec -/// - SummaryMergeMultiple -> SummaryMergeMultipleExec -/// -/// Note: SummaryInfer is handled by datafusion_summary_library's planner -pub struct QueryEngineExtensionPlanner { - /// Reference to the store for reading precomputed outputs - store: Arc, -} - -impl QueryEngineExtensionPlanner { - pub fn new(store: Arc) -> Self { - Self { store } - } -} - -#[async_trait] -impl ExtensionPlanner for QueryEngineExtensionPlanner { - async fn plan_extension( - &self, - _planner: &dyn PhysicalPlanner, - node: &dyn UserDefinedLogicalNode, - _logical_inputs: &[&LogicalPlan], - physical_inputs: &[Arc], - _session_state: &SessionState, - ) -> Result>, DataFusionError> { - // Try to downcast to PrecomputedSummaryRead - if let Some(read) = node.as_any().downcast_ref::() { - return Ok(Some(Arc::new(PrecomputedSummaryReadExec::new( - read.clone(), - self.store.clone(), - )))); - } - - // Try to downcast to SummaryMergeMultiple - if let Some(merge) = node.as_any().downcast_ref::() { - if physical_inputs.len() != 1 { - return Err(DataFusionError::Internal( - "SummaryMergeMultiple expects exactly one input".to_string(), - )); - } - return Ok(Some(Arc::new(SummaryMergeMultipleExec::new( - merge.clone(), - physical_inputs[0].clone(), - )))); - } - - // Try to downcast to SummaryInfer - if let Some(infer) = node.as_any().downcast_ref::() { - // Extract summary_type and sketch_column from the first logical input (values SummaryMergeMultiple) - let values_input_plan = _logical_inputs.first().ok_or_else(|| { - DataFusionError::Internal("SummaryInfer has no logical inputs".to_string()) - })?; - - let (summary_type, sketch_column) = extract_merge_info(values_input_plan, "values")?; - - if _logical_inputs.len() == 2 && physical_inputs.len() == 2 { - // Dual-input: extract keys summary type from second logical input - let keys_input_plan = _logical_inputs[1]; - let (keys_summary_type, _) = extract_merge_info(keys_input_plan, "keys")?; - - return Ok(Some(Arc::new(SummaryInferExec::new_dual_input( - infer.clone(), - physical_inputs[0].clone(), - physical_inputs[1].clone(), - summary_type, - keys_summary_type, - sketch_column, - )))); - } else if physical_inputs.len() == 1 { - // Single-input (original path) - return Ok(Some(Arc::new(SummaryInferExec::new( - infer.clone(), - physical_inputs[0].clone(), - summary_type, - sketch_column, - )))); - } else { - return Err(DataFusionError::Internal(format!( - "SummaryInfer: unexpected number of inputs: logical={}, physical={}", - _logical_inputs.len(), - physical_inputs.len() - ))); - } - } - - // Not a node we handle - let other planners try - Ok(None) - } -} - -/// Extract summary_type and sketch_column from a SummaryMergeMultiple logical plan node. -fn extract_merge_info( - plan: &LogicalPlan, - label: &str, -) -> Result<(SketchType, String), DataFusionError> { - match plan { - LogicalPlan::Extension(ext) => ext - .node - .as_any() - .downcast_ref::() - .map(|merge| { - ( - merge.summary_type().clone(), - merge.sketch_column().to_string(), - ) - }) - .ok_or_else(|| { - DataFusionError::Internal(format!( - "SummaryInfer {} input is not SummaryMergeMultiple", - label - )) - }), - _ => Err(DataFusionError::Internal(format!( - "SummaryInfer {} input must be an Extension node", - label - ))), - } -} - -/// Custom query planner that combines the default DataFusion planner with -/// our QueryEngineExtensionPlanner for custom operators. -pub struct CustomQueryPlanner { - store: Arc, -} - -impl CustomQueryPlanner { - pub fn new(store: Arc) -> Self { - Self { store } - } -} - -impl fmt::Debug for CustomQueryPlanner { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_struct("CustomQueryPlanner") - .field("store", &"") - .finish() - } -} - -#[async_trait] -impl datafusion::execution::context::QueryPlanner for CustomQueryPlanner { - async fn create_physical_plan( - &self, - logical_plan: &LogicalPlan, - session_state: &SessionState, - ) -> Result, DataFusionError> { - // Create default planner with our extension planner - let extension_planner = QueryEngineExtensionPlanner::new(self.store.clone()); - let physical_planner = - DefaultPhysicalPlanner::with_extension_planners(vec![Arc::new(extension_planner)]); - - physical_planner - .create_physical_plan(logical_plan, session_state) - .await - } -} diff --git a/asap-query-engine/src/engines/physical/precomputed_summary_read_exec.rs b/asap-query-engine/src/engines/physical/precomputed_summary_read_exec.rs deleted file mode 100644 index fa991ecd..00000000 --- a/asap-query-engine/src/engines/physical/precomputed_summary_read_exec.rs +++ /dev/null @@ -1,169 +0,0 @@ -//! PrecomputedSummaryReadExec - Physical execution operator for reading precomputed summaries -//! -//! This operator reads precomputed aggregates from a Store and produces -//! RecordBatches with label columns and a serialized sketch column. - -use arrow::datatypes::SchemaRef; -use datafusion::error::DataFusionError; -use datafusion::execution::TaskContext; -use datafusion::logical_expr::UserDefinedLogicalNodeCore; -use datafusion::physical_expr::{EquivalenceProperties, Partitioning}; -use datafusion::physical_plan::{ - stream::RecordBatchStreamAdapter, DisplayAs, DisplayFormatType, ExecutionMode, ExecutionPlan, - PlanProperties, SendableRecordBatchStream, -}; -use datafusion_summary_library::PrecomputedSummaryRead; -use futures::stream; -use std::any::Any; -use std::fmt; -use std::sync::Arc; -use std::time::Instant; -use tracing::debug; - -use super::format_schema; -use crate::engines::physical::conversion::store_result_to_record_batch; -use crate::stores::Store; - -/// Physical execution plan for reading precomputed summaries from a store. -pub struct PrecomputedSummaryReadExec { - /// The logical operator this was created from - logical_node: PrecomputedSummaryRead, - /// Reference to the store - store: Arc, - /// Output schema - schema: SchemaRef, - /// Plan properties (cached) - properties: PlanProperties, -} - -impl PrecomputedSummaryReadExec { - pub fn new(logical_node: PrecomputedSummaryRead, store: Arc) -> Self { - // Convert DFSchema to Schema - let schema = Arc::new(logical_node.schema().as_ref().into()); - - let properties = PlanProperties::new( - EquivalenceProperties::new(Arc::clone(&schema)), - Partitioning::UnknownPartitioning(1), - ExecutionMode::Bounded, - ); - - Self { - logical_node, - store, - schema, - properties, - } - } -} - -impl fmt::Debug for PrecomputedSummaryReadExec { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_struct("PrecomputedSummaryReadExec") - .field("metric", &self.logical_node.metric()) - .field("aggregation_id", &self.logical_node.aggregation_id()) - .finish() - } -} - -impl DisplayAs for PrecomputedSummaryReadExec { - fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { - write!( - f, - "PrecomputedSummaryReadExec: metric={}, agg_id={}, range=[{}, {}]", - self.logical_node.metric(), - self.logical_node.aggregation_id(), - self.logical_node.start_timestamp(), - self.logical_node.end_timestamp() - ) - } -} - -impl ExecutionPlan for PrecomputedSummaryReadExec { - fn name(&self) -> &str { - "PrecomputedSummaryReadExec" - } - - fn as_any(&self) -> &dyn Any { - self - } - - fn schema(&self) -> SchemaRef { - Arc::clone(&self.schema) - } - - fn properties(&self) -> &PlanProperties { - &self.properties - } - - fn children(&self) -> Vec<&Arc> { - vec![] // Leaf node - } - - fn with_new_children( - self: Arc, - _children: Vec>, - ) -> Result, DataFusionError> { - // No children to replace - Ok(self) - } - - fn execute( - &self, - _partition: usize, - _context: Arc, - ) -> Result { - debug!( - metric = %self.logical_node.metric(), - aggregation_id = self.logical_node.aggregation_id(), - start_timestamp = self.logical_node.start_timestamp(), - end_timestamp = self.logical_node.end_timestamp(), - is_exact_query = self.logical_node.is_exact_query(), - output_schema = %format_schema(&self.schema), - output_labels = ?self.logical_node.output_labels(), - "PrecomputedSummaryReadExec::execute" - ); - - // Query the store - let store_query_start = Instant::now(); - let store_result = if self.logical_node.is_exact_query() { - self.store.query_precomputed_output_exact( - self.logical_node.metric(), - self.logical_node.aggregation_id(), - self.logical_node.start_timestamp(), - self.logical_node.end_timestamp(), - ) - } else { - self.store.query_precomputed_output( - self.logical_node.metric(), - self.logical_node.aggregation_id(), - self.logical_node.start_timestamp(), - self.logical_node.end_timestamp(), - ) - } - .map_err(DataFusionError::External)?; - debug!( - store_query_ms = format!("{:.2}", store_query_start.elapsed().as_secs_f64() * 1000.0), - unique_keys = store_result.len(), - "PrecomputedSummaryReadExec store query complete" - ); - - // Convert to RecordBatch - let convert_start = Instant::now(); - let label_names: Vec = self.logical_node.output_labels().to_vec(); - let batch = store_result_to_record_batch(&store_result, &label_names)?; - - debug!( - convert_ms = format!("{:.2}", convert_start.elapsed().as_secs_f64() * 1000.0), - output_rows = batch.num_rows(), - output_cols = batch.num_columns(), - "PrecomputedSummaryReadExec produced batch" - ); - - // Create a stream that yields this single batch - let schema = self.schema.clone(); - Ok(Box::pin(RecordBatchStreamAdapter::new( - schema, - stream::once(async move { Ok(batch) }), - ))) - } -} diff --git a/asap-query-engine/src/engines/physical/summary_infer_exec.rs b/asap-query-engine/src/engines/physical/summary_infer_exec.rs deleted file mode 100644 index 52885528..00000000 --- a/asap-query-engine/src/engines/physical/summary_infer_exec.rs +++ /dev/null @@ -1,769 +0,0 @@ -//! SummaryInferExec - Physical execution operator for extracting values from summaries -//! -//! This operator extracts values from serialized accumulators (summaries). -//! -//! **Single-population path** (no keys_input): -//! Input: rows with [label columns, sketch column] -//! Output: rows with [label columns, value column] -//! -//! **Multi-population path** (keys_input present): -//! Input 0 (values): rows with [spatial_label columns, sketch column] -//! Input 1 (keys): rows with [spatial_label columns, sketch column] -//! Output: rows with [spatial_label columns, sub_key columns, value column] - -use arrow::array::{ArrayRef, BinaryArray, Float64Builder, StringBuilder}; -use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; -use arrow::record_batch::RecordBatch; -use datafusion::error::DataFusionError; -use datafusion::execution::TaskContext; -use datafusion::physical_expr::{EquivalenceProperties, Partitioning}; -use datafusion::physical_plan::common::collect; -use datafusion::physical_plan::{ - stream::RecordBatchStreamAdapter, DisplayAs, DisplayFormatType, ExecutionMode, ExecutionPlan, - PlanProperties, SendableRecordBatchStream, -}; -use datafusion_summary_library::{InferOperation, SketchType, SummaryInfer}; -use futures::stream; -use promql_utilities::query_logics::enums::Statistic; -use std::any::Any; -use std::collections::HashMap; -use std::fmt; -use std::sync::Arc; -use std::time::Instant; -use tracing::debug; - -use super::format_schema; -use crate::engines::physical::accumulator_serde::{ - deserialize_accumulator, deserialize_keys_accumulator, deserialize_multiple_subpopulation, - deserialize_single_subpopulation, -}; - -/// Physical execution plan for extracting values from serialized summaries. -pub struct SummaryInferExec { - /// The logical operator this was created from - logical_node: SummaryInfer, - /// Input execution plan (values branch) - input: Arc, - /// Optional second input (keys branch) for multi-population accumulators - keys_input: Option>, - /// Type of summary being inferred (determines deserialization) - summary_type: SketchType, - /// Type of the keys summary (only set for dual-input) - keys_summary_type: Option, - /// Name of the sketch column in the input schema - sketch_column: String, - /// Output schema (labels + value) - schema: SchemaRef, - /// Plan properties (cached) - properties: PlanProperties, -} - -impl SummaryInferExec { - pub fn new( - logical_node: SummaryInfer, - input: Arc, - summary_type: SketchType, - sketch_column: String, - ) -> Self { - let schema = Self::build_schema(&logical_node, &input, None, &sketch_column); - - let properties = PlanProperties::new( - EquivalenceProperties::new(Arc::clone(&schema)), - Partitioning::UnknownPartitioning(1), - ExecutionMode::Bounded, - ); - - Self { - logical_node, - input, - keys_input: None, - summary_type, - keys_summary_type: None, - sketch_column, - schema, - properties, - } - } - - /// Create a dual-input SummaryInferExec for multi-population accumulators. - pub fn new_dual_input( - logical_node: SummaryInfer, - input: Arc, - keys_input: Arc, - summary_type: SketchType, - keys_summary_type: SketchType, - sketch_column: String, - ) -> Self { - let schema = Self::build_schema(&logical_node, &input, Some(&keys_input), &sketch_column); - - let properties = PlanProperties::new( - EquivalenceProperties::new(Arc::clone(&schema)), - Partitioning::UnknownPartitioning(1), - ExecutionMode::Bounded, - ); - - Self { - logical_node, - input, - keys_input: Some(keys_input), - summary_type, - keys_summary_type: Some(keys_summary_type), - sketch_column, - schema, - properties, - } - } - - fn build_schema( - logical_node: &SummaryInfer, - input: &Arc, - _keys_input: Option<&Arc>, - sketch_column: &str, - ) -> SchemaRef { - let input_schema = input.schema(); - - // Build output schema: label columns from input (minus sketch) ... - let mut fields: Vec = input_schema - .fields() - .iter() - .filter(|f| f.name() != sketch_column) - .map(|f| f.as_ref().clone()) - .collect(); - - // ... plus sub-key columns for dual-input (group_key_columns) - for key_col in &logical_node.group_key_columns { - fields.push(Field::new(key_col, DataType::Utf8, true)); - } - - // ... plus output value columns (one per operation) - for output_name in &logical_node.output_names { - fields.push(Field::new(output_name, DataType::Float64, false)); - } - - Arc::new(Schema::new(fields)) - } - - /// Map InferOperation to Statistic for accumulator query - fn infer_op_to_statistic(op: &InferOperation) -> Statistic { - match op { - InferOperation::ExtractSum => Statistic::Sum, - InferOperation::ExtractCount => Statistic::Count, - InferOperation::ExtractMin => Statistic::Min, - InferOperation::ExtractMax => Statistic::Max, - InferOperation::ExtractIncrease => Statistic::Increase, - InferOperation::ExtractRate => Statistic::Rate, - InferOperation::CountDistinct => Statistic::Cardinality, - InferOperation::Quantile(_) | InferOperation::Median => Statistic::Quantile, - InferOperation::TopK(_) => Statistic::Topk, - // All other operations - use Count as fallback - _ => Statistic::Count, - } - } - - /// Extract query kwargs from an InferOperation. - /// - /// Some operations embed parameters (e.g. Quantile embeds the quantile value, - /// TopK embeds k) that accumulators need via the kwargs HashMap. - fn infer_op_to_kwargs(op: &InferOperation) -> Option> { - match op { - InferOperation::Quantile(q_u16) => { - let q = *q_u16 as f64 / 10000.0; - let mut kwargs = HashMap::new(); - kwargs.insert("quantile".to_string(), q.to_string()); - Some(kwargs) - } - InferOperation::Median => { - let mut kwargs = HashMap::new(); - kwargs.insert("quantile".to_string(), "0.5".to_string()); - Some(kwargs) - } - InferOperation::TopK(k) => { - let mut kwargs = HashMap::new(); - kwargs.insert("k".to_string(), k.to_string()); - Some(kwargs) - } - _ => None, - } - } -} - -impl fmt::Debug for SummaryInferExec { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_struct("SummaryInferExec") - .field("operations", &self.logical_node.operations) - .field("has_keys_input", &self.keys_input.is_some()) - .finish() - } -} - -impl DisplayAs for SummaryInferExec { - fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { - write!( - f, - "SummaryInferExec: operations={:?}, dual_input={}", - self.logical_node.operations, - self.keys_input.is_some() - ) - } -} - -impl ExecutionPlan for SummaryInferExec { - fn name(&self) -> &str { - "SummaryInferExec" - } - - fn as_any(&self) -> &dyn Any { - self - } - - fn schema(&self) -> SchemaRef { - Arc::clone(&self.schema) - } - - fn properties(&self) -> &PlanProperties { - &self.properties - } - - fn children(&self) -> Vec<&Arc> { - let mut children = vec![&self.input]; - if let Some(ref keys_input) = self.keys_input { - children.push(keys_input); - } - children - } - - fn with_new_children( - self: Arc, - children: Vec>, - ) -> Result, DataFusionError> { - match (children.len(), &self.keys_input) { - (1, None) => Ok(Arc::new(Self::new( - self.logical_node.clone(), - children[0].clone(), - self.summary_type.clone(), - self.sketch_column.clone(), - ))), - (2, Some(_)) => Ok(Arc::new(Self::new_dual_input( - self.logical_node.clone(), - children[0].clone(), - children[1].clone(), - self.summary_type.clone(), - self.keys_summary_type.clone().unwrap(), - self.sketch_column.clone(), - ))), - _ => Err(DataFusionError::Internal(format!( - "SummaryInferExec: expected {} children, got {}", - if self.keys_input.is_some() { 2 } else { 1 }, - children.len() - ))), - } - } - - fn execute( - &self, - partition: usize, - context: Arc, - ) -> Result { - let schema = self.schema.clone(); - let schema_for_stream = schema.clone(); - let logical_node = self.logical_node.clone(); - let summary_type = self.summary_type.clone(); - let sketch_column = self.sketch_column.clone(); - - debug!( - input_schema = %format_schema(&self.input.schema()), - output_schema = %format_schema(&self.schema), - operations = ?self.logical_node.operations, - sketch_column = %self.sketch_column, - has_keys_input = self.keys_input.is_some(), - "SummaryInferExec::execute" - ); - - if let Some(ref keys_input) = self.keys_input { - // Multi-population path - let values_stream = self.input.execute(partition, context.clone())?; - let keys_stream = keys_input.execute(partition, context)?; - let keys_summary_type = self.keys_summary_type.clone().unwrap(); - - let output_stream = async move { - let collect_start = Instant::now(); - let values_batches = collect(values_stream).await?; - let keys_batches = collect(keys_stream).await?; - let values_rows: usize = values_batches.iter().map(|b| b.num_rows()).sum(); - let keys_rows: usize = keys_batches.iter().map(|b| b.num_rows()).sum(); - debug!( - collect_ms = format!("{:.2}", collect_start.elapsed().as_secs_f64() * 1000.0), - values_rows, keys_rows, "SummaryInferExec collected dual-input batches" - ); - - let infer_start = Instant::now(); - let result = process_dual_input( - &values_batches, - &keys_batches, - &logical_node, - &summary_type, - &keys_summary_type, - &sketch_column, - &schema, - )?; - debug!( - infer_ms = format!("{:.2}", infer_start.elapsed().as_secs_f64() * 1000.0), - output_rows = result.num_rows(), - "SummaryInferExec dual-input infer complete" - ); - Ok(result) - }; - - Ok(Box::pin(RecordBatchStreamAdapter::new( - schema_for_stream, - stream::once(output_stream), - ))) - } else { - // Single-population path (original behavior) - let input_stream = self.input.execute(partition, context)?; - - let output_stream = async move { - let collect_start = Instant::now(); - let batches = collect(input_stream).await?; - let total_input_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); - debug!( - collect_ms = format!("{:.2}", collect_start.elapsed().as_secs_f64() * 1000.0), - total_input_rows, - num_batches = batches.len(), - "SummaryInferExec collected input (single-pop)" - ); - - let mut all_label_values: Vec> = Vec::new(); - let mut all_result_values: Vec = Vec::new(); - - let self_keyed = is_self_keyed_multi_pop(&summary_type); - - let infer_start = Instant::now(); - for batch in &batches { - if self_keyed { - process_self_keyed_multi_pop_batch( - &mut all_label_values, - &mut all_result_values, - batch, - &logical_node, - &summary_type, - &sketch_column, - )?; - } else { - process_single_pop_batch( - &mut all_label_values, - &mut all_result_values, - batch, - &logical_node, - &summary_type, - &sketch_column, - )?; - } - } - debug!( - infer_ms = format!("{:.2}", infer_start.elapsed().as_secs_f64() * 1000.0), - output_rows = all_result_values.len(), - self_keyed, - "SummaryInferExec infer complete (single-pop)" - ); - - build_output_batch(&all_label_values, &all_result_values, &schema) - }; - - Ok(Box::pin(RecordBatchStreamAdapter::new( - schema_for_stream, - stream::once(output_stream), - ))) - } - } -} - -// ============================================================================ -// Single-population processing (unchanged logic) -// ============================================================================ - -/// Process a batch and extract values from single-population accumulators -fn process_single_pop_batch( - all_label_values: &mut Vec>, - all_result_values: &mut Vec, - batch: &RecordBatch, - logical_node: &SummaryInfer, - summary_type: &SketchType, - sketch_column: &str, -) -> Result<(), DataFusionError> { - let sketch_idx = find_sketch_column_index(batch, sketch_column)?; - - let sketch_array = batch - .column(sketch_idx) - .as_any() - .downcast_ref::() - .ok_or_else(|| DataFusionError::Internal("Sketch column is not Binary".to_string()))?; - - let label_indices: Vec = (0..batch.num_columns()) - .filter(|&i| i != sketch_idx) - .collect(); - - let operation = logical_node - .operations - .first() - .ok_or_else(|| DataFusionError::Internal("SummaryInfer has no operations".to_string()))?; - - let statistic = SummaryInferExec::infer_op_to_statistic(operation); - let query_kwargs = SummaryInferExec::infer_op_to_kwargs(operation); - - for row in 0..batch.num_rows() { - let label_values = extract_label_values(batch, &label_indices, row); - let sketch_bytes = sketch_array.value(row); - - let accumulator = - deserialize_single_subpopulation(sketch_bytes, summary_type).map_err(|e| { - DataFusionError::Internal(format!("Failed to deserialize accumulator: {}", e)) - })?; - - let value = accumulator - .query(statistic, query_kwargs.as_ref()) - .map_err(|e| { - DataFusionError::Internal(format!("Failed to query accumulator: {}", e)) - })?; - - all_label_values.push(label_values); - all_result_values.push(value); - } - - Ok(()) -} - -// ============================================================================ -// Self-keyed multi-population processing (single-input, keys from accumulator) -// ============================================================================ - -/// Returns true if the SketchType is a multi-population accumulator that -/// carries its own keys (via `get_keys()`), so it can be processed in -/// single-input mode without a separate keys stream. -fn is_self_keyed_multi_pop(summary_type: &SketchType) -> bool { - matches!( - summary_type, - SketchType::MultipleIncrease | SketchType::MultipleSum | SketchType::MultipleMinMax - ) -} - -/// Process a batch of self-keyed multi-population accumulators. -/// -/// For each row (spatial group): -/// 1. Deserialize as AggregateCore to call get_keys() -/// 2. Deserialize as MultipleSubpopulationAggregate to call query(stat, key) -/// 3. Emit one output row per sub-key -fn process_self_keyed_multi_pop_batch( - all_label_values: &mut Vec>, - all_result_values: &mut Vec, - batch: &RecordBatch, - logical_node: &SummaryInfer, - summary_type: &SketchType, - sketch_column: &str, -) -> Result<(), DataFusionError> { - let sketch_idx = find_sketch_column_index(batch, sketch_column)?; - - let sketch_array = batch - .column(sketch_idx) - .as_any() - .downcast_ref::() - .ok_or_else(|| DataFusionError::Internal("Sketch column is not Binary".to_string()))?; - - let label_indices: Vec = (0..batch.num_columns()) - .filter(|&i| i != sketch_idx) - .collect(); - - let operation = logical_node - .operations - .first() - .ok_or_else(|| DataFusionError::Internal("SummaryInfer has no operations".to_string()))?; - - let statistic = SummaryInferExec::infer_op_to_statistic(operation); - let query_kwargs = SummaryInferExec::infer_op_to_kwargs(operation); - - let num_sub_key_cols = logical_node.group_key_columns.len(); - - for row in 0..batch.num_rows() { - let spatial_labels = extract_label_values(batch, &label_indices, row); - let sketch_bytes = sketch_array.value(row); - - // Get keys from the accumulator itself - let acc = deserialize_accumulator(sketch_bytes, summary_type).map_err(|e| { - DataFusionError::Internal(format!("Failed to deserialize accumulator: {}", e)) - })?; - let sub_keys = acc.get_keys().unwrap_or_default(); - - // Deserialize as multi-pop for querying - let multi_acc = - deserialize_multiple_subpopulation(sketch_bytes, summary_type).map_err(|e| { - DataFusionError::Internal(format!( - "Failed to deserialize multi-pop accumulator: {}", - e - )) - })?; - - for sub_key in &sub_keys { - let value = multi_acc - .query(statistic, sub_key, query_kwargs.as_ref()) - .map_err(|e| { - DataFusionError::Internal(format!( - "Failed to query multi-pop accumulator for key {:?}: {}", - sub_key, e - )) - })?; - - let mut row_labels = spatial_labels.clone(); - for i in 0..num_sub_key_cols { - if i < sub_key.labels.len() { - row_labels.push(sub_key.labels[i].clone()); - } else { - row_labels.push(String::new()); - } - } - - all_label_values.push(row_labels); - all_result_values.push(value); - } - } - - Ok(()) -} - -// ============================================================================ -// Multi-population (dual-input) processing -// ============================================================================ - -/// Process dual-input: values + keys batches for multi-population accumulators. -/// -/// For each spatial group (row) in the values stream: -/// 1. Deserialize the value sketch as MultipleSubpopulationAggregate -/// 2. Find the matching spatial group in the keys stream -/// 3. Deserialize the keys accumulator, call get_keys() to enumerate sub-keys -/// 4. For each sub-key, call value_acc.query(statistic, sub_key) -> one output row -fn process_dual_input( - values_batches: &[RecordBatch], - keys_batches: &[RecordBatch], - logical_node: &SummaryInfer, - values_summary_type: &SketchType, - keys_summary_type: &SketchType, - sketch_column: &str, - schema: &SchemaRef, -) -> Result { - let operation = logical_node - .operations - .first() - .ok_or_else(|| DataFusionError::Internal("SummaryInfer has no operations".to_string()))?; - let statistic = SummaryInferExec::infer_op_to_statistic(operation); - let query_kwargs = SummaryInferExec::infer_op_to_kwargs(operation); - - // Build a lookup from spatial labels -> keys sketch bytes - // Key: spatial label values (as Vec), Value: serialized keys accumulator bytes - let keys_lookup = build_keys_lookup(keys_batches, sketch_column)?; - - let num_sub_key_cols = logical_node.group_key_columns.len(); - - let mut all_label_values: Vec> = Vec::new(); - let mut all_result_values: Vec = Vec::new(); - - for batch in values_batches { - let sketch_idx = find_sketch_column_index(batch, sketch_column)?; - - let sketch_array = batch - .column(sketch_idx) - .as_any() - .downcast_ref::() - .ok_or_else(|| { - DataFusionError::Internal("Values sketch column is not Binary".to_string()) - })?; - - let label_indices: Vec = (0..batch.num_columns()) - .filter(|&i| i != sketch_idx) - .collect(); - - for row in 0..batch.num_rows() { - let spatial_labels = extract_label_values(batch, &label_indices, row); - let value_sketch_bytes = sketch_array.value(row); - - // Deserialize the value sketch as MultipleSubpopulationAggregate - let value_acc = - deserialize_multiple_subpopulation(value_sketch_bytes, values_summary_type) - .map_err(|e| { - DataFusionError::Internal(format!( - "Failed to deserialize multi-pop value accumulator: {}", - e - )) - })?; - - // Find matching keys accumulator - let keys_bytes = keys_lookup.get(&spatial_labels).ok_or_else(|| { - DataFusionError::Internal(format!( - "No keys accumulator found for spatial group: {:?}", - spatial_labels - )) - })?; - - let keys_acc = - deserialize_keys_accumulator(keys_bytes, keys_summary_type).map_err(|e| { - DataFusionError::Internal(format!( - "Failed to deserialize keys accumulator: {}", - e - )) - })?; - - let sub_keys = keys_acc.get_keys().unwrap_or_default(); - - debug!( - spatial_labels = ?spatial_labels, - num_sub_keys = sub_keys.len(), - "Processing multi-pop spatial group" - ); - - // For each sub-key, query the value accumulator - for sub_key in &sub_keys { - let value = value_acc - .query(statistic, sub_key, query_kwargs.as_ref()) - .map_err(|e| { - DataFusionError::Internal(format!( - "Failed to query multi-pop accumulator for key {:?}: {}", - sub_key, e - )) - })?; - - // Output row: [spatial_labels..., sub_key_labels..., value] - let mut row_labels = spatial_labels.clone(); - // Append sub-key label values - // sub_key.labels is Vec of values - for i in 0..num_sub_key_cols { - if i < sub_key.labels.len() { - row_labels.push(sub_key.labels[i].clone()); - } else { - row_labels.push(String::new()); - } - } - - all_label_values.push(row_labels); - all_result_values.push(value); - } - } - } - - debug!( - output_rows = all_result_values.len(), - "SummaryInferExec building output (multi-pop)" - ); - - build_output_batch(&all_label_values, &all_result_values, schema) -} - -/// Build a lookup map from spatial label values to keys sketch bytes. -fn build_keys_lookup( - keys_batches: &[RecordBatch], - sketch_column: &str, -) -> Result, Vec>, DataFusionError> { - let mut lookup: HashMap, Vec> = HashMap::new(); - - for batch in keys_batches { - let sketch_idx = find_sketch_column_index(batch, sketch_column)?; - - let sketch_array = batch - .column(sketch_idx) - .as_any() - .downcast_ref::() - .ok_or_else(|| { - DataFusionError::Internal("Keys sketch column is not Binary".to_string()) - })?; - - let label_indices: Vec = (0..batch.num_columns()) - .filter(|&i| i != sketch_idx) - .collect(); - - for row in 0..batch.num_rows() { - let label_values = extract_label_values(batch, &label_indices, row); - let sketch_bytes = sketch_array.value(row).to_vec(); - lookup.insert(label_values, sketch_bytes); - } - } - - Ok(lookup) -} - -// ============================================================================ -// Common helpers -// ============================================================================ - -/// Find the index of the sketch column in a batch -fn find_sketch_column_index( - batch: &RecordBatch, - sketch_column: &str, -) -> Result { - batch - .schema() - .fields() - .iter() - .position(|f| f.name() == sketch_column) - .ok_or_else(|| { - DataFusionError::Internal(format!( - "Sketch column '{}' not found in batch schema: {:?}", - sketch_column, - batch - .schema() - .fields() - .iter() - .map(|f| f.name()) - .collect::>() - )) - }) -} - -/// Extract label values from a batch row -fn extract_label_values(batch: &RecordBatch, label_indices: &[usize], row: usize) -> Vec { - label_indices - .iter() - .map(|&idx| { - let col = batch.column(idx); - if let Some(str_array) = col.as_any().downcast_ref::() { - str_array.value(row).to_string() - } else { - String::new() - } - }) - .collect() -} - -/// Build output batch from extracted values -fn build_output_batch( - all_label_values: &[Vec], - all_result_values: &[f64], - schema: &SchemaRef, -) -> Result { - // Get number of label columns (schema fields minus the value column) - let num_label_cols = schema.fields().len() - 1; - - // Build label column builders - let mut label_builders: Vec = - (0..num_label_cols).map(|_| StringBuilder::new()).collect(); - - // Build value column - let mut value_builder = Float64Builder::new(); - - for (label_values, value) in all_label_values.iter().zip(all_result_values.iter()) { - // Add label values - for (i, label_value) in label_values.iter().enumerate() { - if i < label_builders.len() { - label_builders[i].append_value(label_value); - } - } - // Add value - value_builder.append_value(*value); - } - - // Build columns - let mut columns: Vec = label_builders - .iter_mut() - .map(|b| Arc::new(b.finish()) as ArrayRef) - .collect(); - columns.push(Arc::new(value_builder.finish())); - - RecordBatch::try_new(schema.clone(), columns) - .map_err(|e| DataFusionError::Internal(format!("Failed to build output batch: {}", e))) -} diff --git a/asap-query-engine/src/engines/physical/summary_merge_multiple_exec.rs b/asap-query-engine/src/engines/physical/summary_merge_multiple_exec.rs deleted file mode 100644 index de5ec43d..00000000 --- a/asap-query-engine/src/engines/physical/summary_merge_multiple_exec.rs +++ /dev/null @@ -1,556 +0,0 @@ -//! SummaryMergeMultipleExec - Physical execution operator for merging summaries -//! -//! This operator merges multiple summaries with the same group key into one. -//! Input: multiple rows per group key with serialized accumulators -//! Output: one row per group key with merged accumulator - -use arrow::array::{ArrayRef, BinaryArray, BinaryBuilder, StringBuilder}; -use arrow::datatypes::SchemaRef; -use arrow::record_batch::RecordBatch; -use datafusion::error::DataFusionError; -use datafusion::execution::TaskContext; -use datafusion::physical_expr::{EquivalenceProperties, Partitioning}; -use datafusion::physical_plan::common::collect; -use datafusion::physical_plan::{ - stream::RecordBatchStreamAdapter, DisplayAs, DisplayFormatType, ExecutionMode, ExecutionPlan, - PlanProperties, SendableRecordBatchStream, -}; -use datafusion_summary_library::SummaryMergeMultiple; -use futures::stream; -use std::any::Any; -use std::collections::HashMap; -use std::fmt; -use std::sync::Arc; -use std::time::Instant; -use tracing::debug; - -use super::format_schema; -use crate::engines::physical::accumulator_serde::{ - deserialize_accumulator, serialize_accumulator_arroyo, -}; - -/// Physical execution plan for merging multiple summaries by group key. -pub struct SummaryMergeMultipleExec { - /// The logical operator this was created from - logical_node: SummaryMergeMultiple, - /// Input execution plan - input: Arc, - /// Output schema (same as input) - schema: SchemaRef, - /// Plan properties (cached) - properties: PlanProperties, -} - -impl SummaryMergeMultipleExec { - pub fn new(logical_node: SummaryMergeMultiple, input: Arc) -> Self { - let schema = input.schema(); - - let properties = PlanProperties::new( - EquivalenceProperties::new(Arc::clone(&schema)), - Partitioning::UnknownPartitioning(1), - ExecutionMode::Bounded, - ); - - Self { - logical_node, - input, - schema, - properties, - } - } -} - -impl fmt::Debug for SummaryMergeMultipleExec { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_struct("SummaryMergeMultipleExec") - .field("group_by", &self.logical_node.group_by()) - .field("summary_type", &self.logical_node.summary_type()) - .finish() - } -} - -impl DisplayAs for SummaryMergeMultipleExec { - fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { - write!( - f, - "SummaryMergeMultipleExec: group_by=[{}], type={}", - self.logical_node.group_by().join(", "), - self.logical_node.summary_type() - ) - } -} - -impl ExecutionPlan for SummaryMergeMultipleExec { - fn name(&self) -> &str { - "SummaryMergeMultipleExec" - } - - fn as_any(&self) -> &dyn Any { - self - } - - fn schema(&self) -> SchemaRef { - Arc::clone(&self.schema) - } - - fn properties(&self) -> &PlanProperties { - &self.properties - } - - fn children(&self) -> Vec<&Arc> { - vec![&self.input] - } - - fn with_new_children( - self: Arc, - children: Vec>, - ) -> Result, DataFusionError> { - if children.len() != 1 { - return Err(DataFusionError::Internal( - "SummaryMergeMultipleExec expects exactly one child".to_string(), - )); - } - Ok(Arc::new(Self::new( - self.logical_node.clone(), - children[0].clone(), - ))) - } - - fn execute( - &self, - partition: usize, - context: Arc, - ) -> Result { - let input_stream = self.input.execute(partition, context)?; - let schema = self.schema.clone(); - let schema_for_stream = schema.clone(); - let logical_node = self.logical_node.clone(); - - debug!( - input_schema = %format_schema(&self.input.schema()), - output_schema = %format_schema(&self.schema), - group_by = ?self.logical_node.group_by(), - sketch_column = %self.logical_node.sketch_column(), - summary_type = ?self.logical_node.summary_type(), - "SummaryMergeMultipleExec::execute" - ); - - // Use an async block to process all batches and merge - let output_stream = async move { - let mut groups: HashMap, (Vec, Vec)> = HashMap::new(); - - // Collect all batches from input using datafusion's collect helper - let collect_start = Instant::now(); - let batches = collect(input_stream).await?; - let total_input_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); - debug!( - collect_ms = format!("{:.2}", collect_start.elapsed().as_secs_f64() * 1000.0), - total_input_rows, - num_batches = batches.len(), - "SummaryMergeMultipleExec collected input" - ); - - // Process each batch - let merge_start = Instant::now(); - for batch in &batches { - process_batch(&mut groups, batch, &logical_node)?; - } - debug!( - merge_ms = format!("{:.2}", merge_start.elapsed().as_secs_f64() * 1000.0), - output_groups = groups.len(), - "SummaryMergeMultipleExec merged into groups" - ); - - // Build output batch - build_output_batch(&groups, &logical_node, &schema) - }; - - // Convert to stream - Ok(Box::pin(RecordBatchStreamAdapter::new( - schema_for_stream, - stream::once(output_stream), - ))) - } -} - -/// Process a batch and accumulate into groups -fn process_batch( - groups: &mut HashMap, (Vec, Vec)>, - batch: &RecordBatch, - logical_node: &SummaryMergeMultiple, -) -> Result<(), DataFusionError> { - let group_by_cols = logical_node.group_by(); - let sketch_col_name = logical_node.sketch_column(); - - // Find column indices - let group_indices: Vec = group_by_cols - .iter() - .filter_map(|name| { - batch - .schema() - .fields() - .iter() - .position(|f| f.name() == name) - }) - .collect(); - - let sketch_idx = batch - .schema() - .fields() - .iter() - .position(|f| f.name() == sketch_col_name) - .ok_or_else(|| { - DataFusionError::Internal(format!("Sketch column '{}' not found", sketch_col_name)) - })?; - - let sketch_array = batch - .column(sketch_idx) - .as_any() - .downcast_ref::() - .ok_or_else(|| DataFusionError::Internal("Sketch column is not Binary".to_string()))?; - - for row in 0..batch.num_rows() { - // Extract group key - let group_key: Vec = group_indices - .iter() - .map(|&idx| { - let col = batch.column(idx); - if let Some(str_array) = col.as_any().downcast_ref::() { - str_array.value(row).to_string() - } else { - String::new() - } - }) - .collect(); - - // Get sketch bytes - let sketch_bytes = sketch_array.value(row); - - // Merge with existing group or insert new - if let Some((_, existing_bytes)) = groups.get_mut(&group_key) { - // Deserialize both accumulators and merge - let existing_acc = - deserialize_accumulator(existing_bytes, logical_node.summary_type())?; - let new_acc = deserialize_accumulator(sketch_bytes, logical_node.summary_type())?; - - // Merge accumulators - let merged = existing_acc.merge_with(new_acc.as_ref()).map_err(|e| { - DataFusionError::Internal(format!("Failed to merge accumulators: {}", e)) - })?; - - // Serialize merged accumulator in arroyo format for downstream deserialization - *existing_bytes = serialize_accumulator_arroyo(merged.as_ref()); - } else { - // First time seeing this group - groups.insert(group_key.clone(), (group_key, sketch_bytes.to_vec())); - } - } - - Ok(()) -} - -/// Build output batch from merged groups (public for testing) -pub(crate) fn build_output_batch( - groups: &HashMap, (Vec, Vec)>, - logical_node: &SummaryMergeMultiple, - schema: &SchemaRef, -) -> Result { - let group_by_cols = logical_node.group_by(); - - // Build arrays for each column - let mut label_builders: Vec = - group_by_cols.iter().map(|_| StringBuilder::new()).collect(); - let mut sketch_builder = BinaryBuilder::new(); - - for (label_values, bytes) in groups.values() { - // Add label values - for (i, value) in label_values.iter().enumerate() { - if i < label_builders.len() { - label_builders[i].append_value(value); - } - } - // Add sketch bytes - sketch_builder.append_value(bytes); - } - - // Build columns - let mut columns: Vec = label_builders - .iter_mut() - .map(|b| Arc::new(b.finish()) as ArrayRef) - .collect(); - columns.push(Arc::new(sketch_builder.finish())); - - RecordBatch::try_new(schema.clone(), columns) - .map_err(|e| DataFusionError::Internal(format!("Failed to build output batch: {}", e))) -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::data_model::{AggregationType, KeyByLabelValues}; - use crate::engines::physical::accumulator_serde::serialize_accumulator_arroyo; - use crate::precompute_operators::{ - DatasketchesKLLAccumulator, SetAggregatorAccumulator, SumAccumulator, - }; - use arrow::array::StringArray; - use arrow::datatypes::{DataType, Field, Schema}; - use datafusion_summary_library::SketchType; - - /// Helper to create a RecordBatch with [host (Utf8), sketch (Binary)] - fn make_batch(rows: Vec<(&str, Vec)>) -> RecordBatch { - let mut host_builder = StringBuilder::new(); - let mut sketch_builder = BinaryBuilder::new(); - for (host, sketch_bytes) in &rows { - host_builder.append_value(host); - sketch_builder.append_value(sketch_bytes); - } - let schema = Arc::new(Schema::new(vec![ - Field::new("host", DataType::Utf8, true), - Field::new("sketch", DataType::Binary, false), - ])); - RecordBatch::try_new( - schema, - vec![ - Arc::new(host_builder.finish()) as ArrayRef, - Arc::new(sketch_builder.finish()) as ArrayRef, - ], - ) - .unwrap() - } - - /// Helper to create a SummaryMergeMultiple logical node for testing - fn make_logical_node(summary_type: SketchType) -> SummaryMergeMultiple { - use arrow::datatypes::DataType as DT; - use datafusion::common::DFSchema; - use datafusion::logical_expr::{Extension, LogicalPlan}; - use datafusion_summary_library::PrecomputedSummaryRead; - - let fields = vec![ - (None, Arc::new(Field::new("host", DT::Utf8, true))), - (None, Arc::new(Field::new("sketch", DT::Binary, false))), - ]; - let schema = Arc::new(DFSchema::new_with_metadata(fields, Default::default()).unwrap()); - let read = PrecomputedSummaryRead::new( - "test".to_string(), - 1, - 0, - 1000, - true, - vec!["host".to_string()], - summary_type.clone(), - schema, - ); - let read_plan = LogicalPlan::Extension(Extension { - node: Arc::new(read), - }); - SummaryMergeMultiple::new( - Arc::new(read_plan), - vec!["host".to_string()], - "sketch".to_string(), - summary_type, - ) - } - - #[test] - fn test_merge_single_row_passthrough() { - let acc = SumAccumulator::with_sum(42.0); - let bytes = serialize_accumulator_arroyo(&acc); - let batch = make_batch(vec![("host-a", bytes.clone())]); - - let logical = make_logical_node(SketchType::Sum); - let mut groups: HashMap, (Vec, Vec)> = HashMap::new(); - process_batch(&mut groups, &batch, &logical).unwrap(); - - assert_eq!(groups.len(), 1); - // The single row should pass through unchanged - let (_, merged_bytes) = groups.values().next().unwrap(); - let restored = deserialize_accumulator(merged_bytes, &SketchType::Sum).unwrap(); - assert_eq!(restored.get_accumulator_type(), AggregationType::Sum); - } - - #[test] - fn test_merge_two_sums_same_group() { - let acc1 = SumAccumulator::with_sum(50.0); - let acc2 = SumAccumulator::with_sum(50.0); - let bytes1 = serialize_accumulator_arroyo(&acc1); - let bytes2 = serialize_accumulator_arroyo(&acc2); - let batch = make_batch(vec![("host-a", bytes1), ("host-a", bytes2)]); - - let logical = make_logical_node(SketchType::Sum); - let mut groups: HashMap, (Vec, Vec)> = HashMap::new(); - process_batch(&mut groups, &batch, &logical).unwrap(); - - assert_eq!(groups.len(), 1); - let (_, merged_bytes) = groups.values().next().unwrap(); - let restored = - crate::engines::physical::accumulator_serde::deserialize_single_subpopulation( - merged_bytes, - &SketchType::Sum, - ) - .unwrap(); - let value = restored - .query(promql_utilities::query_logics::enums::Statistic::Sum, None) - .unwrap(); - assert!( - (value - 100.0).abs() < 1e-10, - "Merged sum should be 100.0, got {}", - value - ); - } - - #[test] - fn test_merge_three_sums_associativity() { - let bytes: Vec> = [30.0, 40.0, 30.0] - .iter() - .map(|v| serialize_accumulator_arroyo(&SumAccumulator::with_sum(*v))) - .collect(); - let batch = make_batch(vec![ - ("host-a", bytes[0].clone()), - ("host-a", bytes[1].clone()), - ("host-a", bytes[2].clone()), - ]); - - let logical = make_logical_node(SketchType::Sum); - let mut groups: HashMap, (Vec, Vec)> = HashMap::new(); - process_batch(&mut groups, &batch, &logical).unwrap(); - - assert_eq!(groups.len(), 1); - let (_, merged_bytes) = groups.values().next().unwrap(); - let restored = - crate::engines::physical::accumulator_serde::deserialize_single_subpopulation( - merged_bytes, - &SketchType::Sum, - ) - .unwrap(); - let value = restored - .query(promql_utilities::query_logics::enums::Statistic::Sum, None) - .unwrap(); - assert!( - (value - 100.0).abs() < 1e-10, - "30+40+30 should be 100.0, got {}", - value - ); - } - - #[test] - fn test_merge_separate_groups_no_contamination() { - let bytes_a = serialize_accumulator_arroyo(&SumAccumulator::with_sum(100.0)); - let bytes_b = serialize_accumulator_arroyo(&SumAccumulator::with_sum(200.0)); - let batch = make_batch(vec![("host-a", bytes_a), ("host-b", bytes_b)]); - - let logical = make_logical_node(SketchType::Sum); - let mut groups: HashMap, (Vec, Vec)> = HashMap::new(); - process_batch(&mut groups, &batch, &logical).unwrap(); - - assert_eq!(groups.len(), 2, "Two different hosts should stay separate"); - - // Verify values - for (key, (_, bytes)) in &groups { - let restored = - crate::engines::physical::accumulator_serde::deserialize_single_subpopulation( - bytes, - &SketchType::Sum, - ) - .unwrap(); - let value = restored - .query(promql_utilities::query_logics::enums::Statistic::Sum, None) - .unwrap(); - if key[0] == "host-a" { - assert!((value - 100.0).abs() < 1e-10); - } else { - assert!((value - 200.0).abs() < 1e-10); - } - } - } - - #[test] - fn test_merge_kll_sketches() { - let mut kll1 = DatasketchesKLLAccumulator::new(200); - kll1.update(1.0); - kll1.update(2.0); - let mut kll2 = DatasketchesKLLAccumulator::new(200); - kll2.update(3.0); - kll2.update(4.0); - - let bytes1 = serialize_accumulator_arroyo(&kll1); - let bytes2 = serialize_accumulator_arroyo(&kll2); - let batch = make_batch(vec![("host-a", bytes1), ("host-a", bytes2)]); - - let logical = make_logical_node(SketchType::KLL); - let mut groups: HashMap, (Vec, Vec)> = HashMap::new(); - process_batch(&mut groups, &batch, &logical).unwrap(); - - assert_eq!(groups.len(), 1); - // Verify merged KLL has data from both - let (_, merged_bytes) = groups.values().next().unwrap(); - let restored = - crate::engines::physical::accumulator_serde::deserialize_single_subpopulation( - merged_bytes, - &SketchType::KLL, - ) - .unwrap(); - let mut kwargs = std::collections::HashMap::new(); - kwargs.insert("quantile".to_string(), "1.0".to_string()); - let max = restored - .query( - promql_utilities::query_logics::enums::Statistic::Quantile, - Some(&kwargs), - ) - .unwrap(); - assert!( - (max - 4.0).abs() < 1e-10, - "Max quantile should be 4.0, got {}", - max - ); - } - - #[test] - fn test_merge_set_aggregators() { - let mut set1 = SetAggregatorAccumulator::new(); - set1.add_key(KeyByLabelValues { - labels: vec!["a".to_string()], - }); - let mut set2 = SetAggregatorAccumulator::new(); - set2.add_key(KeyByLabelValues { - labels: vec!["b".to_string()], - }); - - let bytes1 = serialize_accumulator_arroyo(&set1); - let bytes2 = serialize_accumulator_arroyo(&set2); - let batch = make_batch(vec![("host-a", bytes1), ("host-a", bytes2)]); - - let logical = make_logical_node(SketchType::SetAggregator); - let mut groups: HashMap, (Vec, Vec)> = HashMap::new(); - process_batch(&mut groups, &batch, &logical).unwrap(); - - assert_eq!(groups.len(), 1); - let (_, merged_bytes) = groups.values().next().unwrap(); - let restored = deserialize_accumulator(merged_bytes, &SketchType::SetAggregator).unwrap(); - let keys = restored.get_keys().unwrap(); - assert_eq!(keys.len(), 2, "Union of two sets should have 2 keys"); - } - - #[test] - fn test_merge_empty_batch() { - let schema = Arc::new(Schema::new(vec![ - Field::new("host", DataType::Utf8, true), - Field::new("sketch", DataType::Binary, false), - ])); - let host_array = StringArray::from(Vec::<&str>::new()); - let sketch_array = arrow::array::BinaryArray::from(Vec::<&[u8]>::new()); - let batch = RecordBatch::try_new( - schema, - vec![ - Arc::new(host_array) as ArrayRef, - Arc::new(sketch_array) as ArrayRef, - ], - ) - .unwrap(); - - let logical = make_logical_node(SketchType::Sum); - let mut groups: HashMap, (Vec, Vec)> = HashMap::new(); - process_batch(&mut groups, &batch, &logical).unwrap(); - - assert_eq!(groups.len(), 0, "Empty batch should produce 0 groups"); - } -} diff --git a/asap-query-engine/src/engines/simple/engine.rs b/asap-query-engine/src/engines/simple/engine.rs index 01dbd2d2..8a535e95 100644 --- a/asap-query-engine/src/engines/simple/engine.rs +++ b/asap-query-engine/src/engines/simple/engine.rs @@ -289,7 +289,7 @@ pub struct SimpleEngine { /// EngineRouter's archive failover (Phase 6). When `None`, the /// engine behaves as it did before Phase 5 wire-in (every query /// goes through `handle_query`'s legacy path). - sketch_index: Option>, + sketch_index: Option>, } impl SimpleEngine { @@ -481,7 +481,7 @@ impl SimpleEngine { /// query through `handle_query`). pub fn with_sketch_index( mut self, - index: Arc, + index: Arc, ) -> Self { self.sketch_index = Some(index); self @@ -1346,162 +1346,6 @@ impl SimpleEngine { Ok((results, chosen_window)) } - /// Execute a query using the plan-based approach (for testing) - /// - /// This is an alternative execution path that uses DataFusion logical/physical - /// plans instead of the existing execute_query_pipeline. - /// - /// # Arguments - /// * `context` - The query execution context - /// - /// # Returns - /// A Result containing the query results or an error - #[allow(dead_code)] - pub async fn execute_plan( - &self, - context: &QueryExecutionContext, - ) -> Result, String> { - use datafusion::execution::context::SessionContext; - use datafusion::physical_plan::collect; - - use crate::engines::physical::conversion::record_batch_to_result_map; - - let total_start = Instant::now(); - - // 1. Build logical plan from context - let plan_build_start = Instant::now(); - let logical_plan = context - .to_logical_plan() - .map_err(|e| format!("Failed to build logical plan: {}", e))?; - debug!( - "[LATENCY] DataFusion: logical plan build: {:.2}ms", - plan_build_start.elapsed().as_secs_f64() * 1000.0 - ); - debug!( - "DataFusion logical plan:\n{}", - logical_plan.display_indent() - ); - - // 2. Create session context with our custom extension planner - let physical_plan_start = Instant::now(); - let session_ctx = SessionContext::new(); - #[allow(deprecated)] - let state = session_ctx.state().with_query_planner(std::sync::Arc::new( - crate::engines::physical::CustomQueryPlanner::new(self.store.clone()), - )); - - // 3. Create physical plan - let physical_plan = state - .create_physical_plan(&logical_plan) - .await - .map_err(|e| format!("Failed to create physical plan: {}", e))?; - debug!( - "[LATENCY] DataFusion: physical plan creation: {:.2}ms", - physical_plan_start.elapsed().as_secs_f64() * 1000.0 - ); - - // 4. Execute - let execute_start = Instant::now(); - let task_ctx = session_ctx.task_ctx(); - let batches = collect(physical_plan, task_ctx) - .await - .map_err(|e| format!("Failed to execute plan: {}", e))?; - let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); - debug!( - "[LATENCY] DataFusion: plan execution: {:.2}ms, {} batch(es), {} total rows", - execute_start.elapsed().as_secs_f64() * 1000.0, - batches.len(), - total_rows - ); - - // 5. Convert results - let convert_start = Instant::now(); - let label_names: Vec<&str> = context - .metadata - .query_output_labels - .labels - .iter() - .map(String::as_str) - .collect(); - - let mut all_results: HashMap, f64> = HashMap::new(); - for batch in &batches { - let batch_results = record_batch_to_result_map(batch, &label_names, "value") - .map_err(|e| format!("Failed to convert results: {}", e))?; - all_results.extend(batch_results); - } - debug!( - "[LATENCY] DataFusion: result conversion: {:.2}ms, {} output rows", - convert_start.elapsed().as_secs_f64() * 1000.0, - all_results.len() - ); - - // 6. Format results - let format_start = Instant::now(); - let results = self.format_final_results( - all_results, - &context.metadata.statistic_to_compute, - &context.metric, - false, - ); - debug!( - "[LATENCY] DataFusion: result formatting: {:.2}ms, {} results", - format_start.elapsed().as_secs_f64() * 1000.0, - results.len() - ); - - debug!( - "[LATENCY] DataFusion: total execute_plan: {:.2}ms", - total_start.elapsed().as_secs_f64() * 1000.0 - ); - - Ok(results) - } - - /// Executes a pre-built DataFusion logical plan and returns results. - /// - /// This is the shared execution kernel used by both `execute_plan` (for single-metric - /// queries) and the binary arithmetic dispatch path. - pub async fn execute_logical_plan( - &self, - logical_plan: datafusion::logical_expr::LogicalPlan, - label_names: Vec, - metric: &str, - statistic: &Statistic, - ) -> Result, String> { - use datafusion::execution::context::SessionContext; - use datafusion::physical_plan::collect; - - use crate::engines::physical::conversion::record_batch_to_result_map; - - // Create session context with our custom extension planner - let session_ctx = SessionContext::new(); - #[allow(deprecated)] - let state = session_ctx.state().with_query_planner(std::sync::Arc::new( - crate::engines::physical::CustomQueryPlanner::new(self.store.clone()), - )); - - let physical_plan = state - .create_physical_plan(&logical_plan) - .await - .map_err(|e| format!("Failed to create physical plan: {}", e))?; - - let task_ctx = session_ctx.task_ctx(); - let batches = collect(physical_plan, task_ctx) - .await - .map_err(|e| format!("Failed to execute plan: {}", e))?; - - let label_name_strs: Vec<&str> = label_names.iter().map(String::as_str).collect(); - let mut all_results: HashMap, f64> = HashMap::new(); - for batch in &batches { - let batch_results = record_batch_to_result_map(batch, &label_name_strs, "value") - .map_err(|e| format!("Failed to convert results: {}", e))?; - all_results.extend(batch_results); - } - - Ok(self.format_final_results(all_results, statistic, metric, false)) - } - /// Finds a query config by structurally comparing `arm_ast` against each /// config's parsed query. /// @@ -1684,111 +1528,6 @@ impl SimpleEngine { }) } - /// Recursively builds a DataFusion logical plan for one arm of a binary - /// arithmetic expression. - /// - /// - Leaf arm (supported PromQL pattern): look up config structurally, build - /// context, return its `to_logical_plan()` together with the output label names. - /// - Binary arm: recursively build both sub-arms and combine with - /// `build_binary_vector_plan`. - /// - Scalar literal: returns `None` (handled by the caller separately). - fn build_arm_logical_plan( - &self, - arm_ast: &promql_parser::parser::Expr, - time: f64, - ) -> Option<(datafusion::logical_expr::LogicalPlan, Vec)> { - use crate::engines::logical::plan_builder::build_binary_vector_plan; - use promql_parser::parser::Expr; - - match arm_ast { - Expr::NumberLiteral(_) => None, // caller handles scalars - Expr::Paren(paren) => self.build_arm_logical_plan(&paren.expr, time), - Expr::Binary(binary) => { - // Nested binary expression — recurse on both sides - let (lhs_plan, lhs_labels) = self.build_arm_logical_plan(&binary.lhs, time)?; - let (rhs_plan, _) = self.build_arm_logical_plan(&binary.rhs, time)?; - let combined = - build_binary_vector_plan(lhs_plan, rhs_plan, &binary.op, lhs_labels.clone()) - .ok()?; - Some((combined, lhs_labels)) - } - other => { - // Leaf pattern: structural config lookup + context + plan - let config = self.find_query_config_promql_structural(other)?; - let ctx = self.build_query_execution_context_from_ast(other, config, time)?; - let label_names = ctx.metadata.query_output_labels.labels.clone(); - let plan = ctx.to_logical_plan().ok()?; - Some((plan, label_names)) - } - } - } - - /// Handles a binary arithmetic PromQL expression by building a combined - /// DataFusion plan (vector–vector join or scalar projection) and executing it. - /// - /// Returns `None` if any arm is not acceleratable (caller falls back to Prometheus). - fn handle_binary_expr_promql( - &self, - ast: &promql_parser::parser::Expr, - time: f64, - ) -> Option<(KeyByLabelNames, QueryResult)> { - use crate::engines::logical::plan_builder::{build_binary_vector_plan, build_scalar_plan}; - use promql_parser::parser::Expr; - - let query_time = Self::convert_query_time_to_data_time(time); - - let binary = match ast { - Expr::Binary(b) => b, - _ => return None, - }; - - let lhs = binary.lhs.as_ref(); - let rhs = binary.rhs.as_ref(); - let op = &binary.op; - - // Scalar case: either side may be a numeric literal - let scalar_case: Option<(f64, &Expr, bool)> = match (lhs, rhs) { - (_, Expr::NumberLiteral(nl)) => Some((nl.val, lhs, false)), - (Expr::NumberLiteral(nl), _) => Some((nl.val, rhs, true)), - _ => None, - }; - if let Some((scalar, vector_arm, scalar_on_left)) = scalar_case { - let (vector_plan, label_names) = self.build_arm_logical_plan(vector_arm, time)?; - let combined = - build_scalar_plan(vector_plan, scalar, op, scalar_on_left, label_names.clone()) - .ok()?; - let results = tokio::task::block_in_place(|| { - tokio::runtime::Handle::current().block_on(self.execute_logical_plan( - combined, - label_names.clone(), - "", - &Statistic::Sum, - )) - }) - .ok()?; - return Some(( - KeyByLabelNames::new(label_names), - QueryResult::vector(results, query_time), - )); - } - - // Vector–vector - let (lhs_plan, lhs_labels) = self.build_arm_logical_plan(lhs, time)?; - let (rhs_plan, _) = self.build_arm_logical_plan(rhs, time)?; - let combined = build_binary_vector_plan(lhs_plan, rhs_plan, op, lhs_labels.clone()).ok()?; - let results = tokio::task::block_in_place(|| { - tokio::runtime::Handle::current().block_on(self.execute_logical_plan( - combined, - lhs_labels.clone(), - "", - &Statistic::Sum, - )) - }) - .ok()?; - let output_labels = KeyByLabelNames::new(lhs_labels); - Some((output_labels, QueryResult::vector(results, query_time))) - } - /// Applies a PromQL binary arithmetic operator to two f64 values. fn apply_range_binary_op( op: &promql_parser::parser::token::TokenType, @@ -2437,19 +2176,13 @@ impl SimpleEngine { .resolve_sketch_metric_alias(&query) .unwrap_or(query); - // Check for binary arithmetic before attempting single-query dispatch. - // Binary expressions won't have a matching query_config, so we handle them here. - if let Ok(ast) = promql_parser::parser::parse(&query) { - if matches!(&ast, promql_parser::parser::Expr::Binary(_)) { - let result = self.handle_binary_expr_promql(&ast, time); - let total_query_duration = query_start_time.elapsed(); - debug!( - "Binary arithmetic query handling took: {:.2}ms", - total_query_duration.as_secs_f64() * 1000.0 - ); - return result; - } - } + // Binary arithmetic dispatch was previously handled here via a + // DataFusion-based plan combiner. That path was removed alongside + // the datafusion crate; binary arithmetic on warm-tier sketches + // will be reintroduced as part of the PromQL-evaluator-on-Gorilla + // follow-up. For now binary expressions fall through to the + // normal dispatch path (which will not match and trigger the + // router's CapabilityMiss failover to the archive engine). // Try the §7 schema-timeline dispatch first. Returns Some // only when the query's [t1, t2] range crosses a @@ -3786,9 +3519,9 @@ impl crate::routing::engine_router::QueryEngine for SimpleEngine { let mut all_hit = true; for sid in &candidates { match idx.classify(*sid) { - crate::stores::sketch_index::SidLookup::Hit => {} - crate::stores::sketch_index::SidLookup::Ghost - | crate::stores::sketch_index::SidLookup::Unknown => { + crate::stores::sketch_db::sketch_index::SidLookup::Hit => {} + crate::stores::sketch_db::sketch_index::SidLookup::Ghost + | crate::stores::sketch_db::sketch_index::SidLookup::Unknown => { all_hit = false; break; } @@ -6086,7 +5819,7 @@ mod warm_tier_classify_tests { use crate::engines::EngineError; use crate::routing::engine_router::QueryEngine as _; use crate::stores::sketch_db::simple_map_store::SimpleMapStore; - use crate::stores::sketch_index::{ + use crate::stores::sketch_db::sketch_index::{ AccuracyBound, Capability, SketchConfig, SketchIndex, SketchInstanceMetadata, SketchKindHandle, SketchSampleState, }; @@ -6191,7 +5924,7 @@ mod warm_tier_classify_tests { (1_000, 1_010), SketchSampleState { bytes: vec![0], - encoding: crate::stores::sketch_index::SketchEncoding::ProtoFull, + encoding: crate::stores::sketch_db::sketch_index::SketchEncoding::ProtoFull, }, ); diff --git a/asap-query-engine/src/main.rs b/asap-query-engine/src/main.rs index c7715130..fa7d9aab 100644 --- a/asap-query-engine/src/main.rs +++ b/asap-query-engine/src/main.rs @@ -441,7 +441,7 @@ async fn main() -> Result<()> { let series_resolver = Arc::new( query_engine_rust::drivers::ingest::series_resolver::SeriesIdResolver::new(), ); - let sketch_index = Arc::new(query_engine_rust::stores::sketch_index::SketchIndex::new()); + let sketch_index = Arc::new(query_engine_rust::stores::sketch_db::sketch_index::SketchIndex::new()); // Setup query engine. SimpleEngine shares the same // HotReloadStreamingConfig handle as the HTTP server, so a POST diff --git a/asap-query-engine/src/precompute_engine/engine.rs b/asap-query-engine/src/precompute_engine/engine.rs index 90d7ceb1..0b497a60 100644 --- a/asap-query-engine/src/precompute_engine/engine.rs +++ b/asap-query-engine/src/precompute_engine/engine.rs @@ -39,7 +39,7 @@ impl PrecomputeEngine { hot_reload_config: HotReloadStreamingConfig, output_sink: Arc, series_resolver: Arc, - sketch_index: Arc, + sketch_index: Arc, ) -> Self { let worker_group_counts = (0..config.num_workers) .map(|_| Arc::new(AtomicUsize::new(0))) diff --git a/asap-query-engine/src/precompute_engine/ingest_handler.rs b/asap-query-engine/src/precompute_engine/ingest_handler.rs index 4c8c9d81..595cf934 100644 --- a/asap-query-engine/src/precompute_engine/ingest_handler.rs +++ b/asap-query-engine/src/precompute_engine/ingest_handler.rs @@ -64,7 +64,7 @@ pub struct IngestState { /// every modified-OTLP first-class sketch DataPoint; queried by /// the `SimpleEngine` query path (warm-tier hit / ghost / unknown /// classification drives the Phase 6 archive failover). - pub sketch_index: Arc, + pub sketch_index: Arc, } impl IngestState { @@ -183,7 +183,7 @@ mod tests { series_resolver: Arc::new( crate::drivers::ingest::series_resolver::SeriesIdResolver::new(), ), - sketch_index: Arc::new(crate::stores::sketch_index::SketchIndex::new()), + sketch_index: Arc::new(crate::stores::sketch_db::sketch_index::SketchIndex::new()), }); let drain = tokio::spawn(async move { while rx.recv().await.is_some() {} }); diff --git a/asap-query-engine/src/stores/mod.rs b/asap-query-engine/src/stores/mod.rs index 4928b9be..4ac8491a 100644 --- a/asap-query-engine/src/stores/mod.rs +++ b/asap-query-engine/src/stores/mod.rs @@ -16,16 +16,14 @@ //! callers should not care whether it lives under `sketch_db` or //! at the `stores` top level. -pub mod epoch_columnar; pub mod promsketch_store; pub mod sketch_db; -pub mod sketch_index; pub mod traits; // pub use promsketch_store::PromSketchStore; -pub use sketch_db::{AggSchema, AggStatus, SchemaRegistry, SimpleMapStore}; -pub use sketch_index::{ +pub use sketch_db::sketch_index::{ AccuracyBound, Capability, SidLookup, SketchConfig, SketchEncoding, SketchIndex, SketchInstanceMetadata, SketchKindHandle, SketchSampleState, SketchTimeSeries, }; +pub use sketch_db::{AggSchema, AggStatus, SchemaRegistry, SimpleMapStore}; pub use traits::*; diff --git a/asap-query-engine/src/stores/epoch_columnar.rs b/asap-query-engine/src/stores/sketch_db/epoch_columnar.rs similarity index 100% rename from asap-query-engine/src/stores/epoch_columnar.rs rename to asap-query-engine/src/stores/sketch_db/epoch_columnar.rs diff --git a/asap-query-engine/src/stores/sketch_db/mod.rs b/asap-query-engine/src/stores/sketch_db/mod.rs index 4e6db7ca..a33319cd 100644 --- a/asap-query-engine/src/stores/sketch_db/mod.rs +++ b/asap-query-engine/src/stores/sketch_db/mod.rs @@ -34,12 +34,14 @@ pub mod backfill_processor; pub mod backfill_service; pub mod backfill_window_builder; pub mod backfill_worker; +pub mod epoch_columnar; pub mod metrics; pub mod prometheus_reader; pub mod raw_sample_reader; pub mod schema; pub mod schema_eviction; pub mod simple_map_store; +pub mod sketch_index; pub use accuracy::{AccuracyEnvelope, AccuracyKind, AccuracyProfile, PerSegmentAccuracy}; pub use backfill::{ diff --git a/asap-query-engine/src/stores/sketch_db/simple_map_store/per_key.rs b/asap-query-engine/src/stores/sketch_db/simple_map_store/per_key.rs index f10ccf39..e82a77b5 100644 --- a/asap-query-engine/src/stores/sketch_db/simple_map_store/per_key.rs +++ b/asap-query-engine/src/stores/sketch_db/simple_map_store/per_key.rs @@ -2,13 +2,11 @@ use crate::data_model::{ AggregateCore, AggregationType, CleanupPolicy, KeyByLabelValues, PrecomputedOutput, StreamingConfig, }; -use crate::engines::physical::accumulator_serde; use crate::stores::sketch_db::simple_map_store::common::{ EpochID, InternTable, MetricBucketMap, MetricID, MutableEpoch, SealedEpoch, TimestampRange, }; use crate::stores::{Store, StoreResult, TimestampedBucketsMap}; use dashmap::DashMap; -use datafusion_summary_library::SketchType; use std::collections::{BTreeMap, HashMap}; use std::path::PathBuf; use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; @@ -663,86 +661,21 @@ impl SimpleMapStorePerKey { /// in-memory result map. A no-op when persistence is disabled. fn query_disk_parts( &self, - metric: &str, - aggregation_id: u64, - start: u64, - end: u64, - results: &mut TimestampedBucketsMap, + _metric: &str, + _aggregation_id: u64, + _start: u64, + _end: u64, + _results: &mut TimestampedBucketsMap, ) -> Result<(), Box> { - let Some(state) = self.persistence.as_ref() else { - return Ok(()); - }; - - let overlapping = state.manifest.live_parts_overlapping(start, end); - if overlapping.is_empty() { - return Ok(()); - } - - for entry in overlapping { - let reader = match state.cache.get_or_load(entry.part_id) { - Ok(r) => r, - Err(e) => { - warn!( - "query_disk_parts: failed to open part {} for metric {}: {}", - entry.part_id, metric, e - ); - continue; - } - }; - for rec in reader.index_records() { - if rec.agg_id != aggregation_id { - continue; - } - // Overlap semantics matching MutableEpoch::range_query_into: - // include any window whose [start_ts, end_ts) interval - // intersects [start, end). The earlier "fully inside" - // form silently dropped windows that crossed the query - // boundaries, which is what tumbling windows do - // virtually always when the query timestamp doesn't - // align to the window grid. - if rec.end_ts <= start || rec.start_ts >= end { - continue; - } - let disk_entry = match reader.load_entry(&rec) { - Ok(d) => d, - Err(e) => { - warn!( - "query_disk_parts: failed to load entry at offset {} of part {}: {}", - rec.data_offset, entry.part_id, e - ); - continue; - } - }; - let Some(sketch_type) = type_name_to_sketch_type(&disk_entry.sketch_type_name) - else { - warn!( - "query_disk_parts: no SketchType mapping for {}; skipping", - disk_entry.sketch_type_name - ); - continue; - }; - let decoded = match accumulator_serde::deserialize_accumulator( - &disk_entry.sketch_bytes, - &sketch_type, - ) { - Ok(a) => a, - Err(e) => { - warn!( - "query_disk_parts: deserialize failed for {}: {}", - disk_entry.sketch_type_name, e - ); - continue; - } - }; - let arc_acc: Arc = Arc::from(decoded); - results - .entry(disk_entry.label.clone()) - .or_default() - .push(((rec.start_ts, rec.end_ts), arc_acc)); - } - } - - Ok(()) + // TODO: replace with non-datafusion path. The previous body + // depended on `accumulator_serde::deserialize_accumulator` from + // the removed `engines::physical` module to materialise sketches + // from persisted parts. Since SimpleMapStore is deprecated and + // the warm-tier query path is being rebuilt on top of the + // SketchIndex, this returns Ok(()) so that callers see "no disk + // parts" rather than panicking; the live in-memory path still + // serves recent windows. + panic!("datafusion-dependent path removed; ingest/persistence still under refactor") } } @@ -758,23 +691,6 @@ impl Drop for SimpleMapStorePerKey { } } -/// Map `AggregateCore::type_name()` to the `SketchType` enum value -/// used by `accumulator_serde::deserialize_accumulator`. Returns -/// `None` for types that don't have a working Arroyo round-trip yet. -fn type_name_to_sketch_type(name: &str) -> Option { - match name { - "SumAccumulator" => Some(SketchType::Sum), - "DatasketchesKLLAccumulator" => Some(SketchType::KLL), - "HydraKllSketchAccumulator" => Some(SketchType::HydraKLL), - "CountMinSketchAccumulator" => Some(SketchType::CountMinSketch), - "SetAggregatorAccumulator" => Some(SketchType::SetAggregator), - "DeltaSetAggregatorAccumulator" => Some(SketchType::DeltaSetAggregator), - "MultipleSumAccumulator" => Some(SketchType::MultipleSum), - "MultipleIncreaseAccumulator" => Some(SketchType::MultipleIncrease), - _ => None, - } -} - #[async_trait::async_trait] impl Store for SimpleMapStorePerKey { fn insert_precomputed_output( @@ -1077,49 +993,16 @@ impl EpochSource for PerKeyInner { fn snapshot_sealed_epoch( &self, - agg_id: u64, - epoch_id: u64, + _agg_id: u64, + _epoch_id: u64, ) -> PersistResult> { - let Some(lock) = self.store.get(&agg_id) else { - return Ok(None); - }; - let data = lock - .read() - .map_err(|e| PersistError::Internal(format!("read lock poisoned: {}", e)))?; - let Some(epoch) = data.sealed_epochs.get(&epoch_id) else { - return Ok(None); - }; - - let mut entries = Vec::with_capacity(epoch.entries.len()); - for (tr, metric_id, agg) in &epoch.entries { - // Resolve the label from the per-agg intern table. - let label: Option = data.intern.resolve(*metric_id).clone(); - - // Serialize the sketch via the arroyo path so it's - // round-trippable via deserialize_accumulator. - let sketch_bytes = accumulator_serde::serialize_accumulator_arroyo(agg.as_ref()); - let type_name = agg.type_name().to_string(); - - entries.push(EpochSnapshotEntry { - start_ts: tr.0, - end_ts: tr.1, - label, - sketch_type_name: type_name, - sketch_bytes, - }); - } - - let (min_ts, max_ts) = epoch.time_bounds().unwrap_or((0, 0)); - let approx_bytes = epoch_approx_bytes(epoch); - - Ok(Some(EpochSnapshot { - agg_id, - epoch_id, - min_ts, - max_ts, - entries, - approx_bytes, - })) + // TODO: replace with non-datafusion path. The previous body + // serialised sketches via + // `accumulator_serde::serialize_accumulator_arroyo` from the + // removed `engines::physical` module. SimpleMapStore is + // deprecated; persistence is being refactored on top of the + // SketchIndex. + panic!("datafusion-dependent path removed; ingest/persistence still under refactor") } fn evict_sealed_epoch(&self, agg_id: u64, epoch_id: u64) { diff --git a/asap-query-engine/src/stores/sketch_index.rs b/asap-query-engine/src/stores/sketch_db/sketch_index.rs similarity index 100% rename from asap-query-engine/src/stores/sketch_index.rs rename to asap-query-engine/src/stores/sketch_db/sketch_index.rs diff --git a/asap-query-engine/src/tests/datafusion/accumulator_serde_tests.rs b/asap-query-engine/src/tests/datafusion/accumulator_serde_tests.rs deleted file mode 100644 index c90e8e56..00000000 --- a/asap-query-engine/src/tests/datafusion/accumulator_serde_tests.rs +++ /dev/null @@ -1,339 +0,0 @@ -//! Accumulator Serde Round-Trip Tests -//! -//! Tests that exercise accumulator_serde.rs directly (no engine needed). -//! Verifies serialize -> deserialize round-trip for all accumulator types. - -#[cfg(test)] -mod tests { - use crate::data_model::SerializableToSink; - use crate::data_model::{KeyByLabelValues, Measurement}; - use crate::engines::physical::accumulator_serde::{ - deserialize_accumulator, deserialize_keys_accumulator, deserialize_multiple_subpopulation, - deserialize_single_subpopulation, serialize_accumulator_arroyo, - }; - use crate::precompute_operators::{ - CountMinSketchAccumulator, DatasketchesKLLAccumulator, DeltaSetAggregatorAccumulator, - HydraKllSketchAccumulator, IncreaseAccumulator, MultipleIncreaseAccumulator, - SetAggregatorAccumulator, SumAccumulator, - }; - use datafusion_summary_library::SketchType; - use promql_utilities::query_logics::enums::{AggregationType, Statistic}; - use std::collections::HashMap; - - // ======================================================================== - // Full round-trip tests (serialize_arroyo -> deserialize) - // ======================================================================== - - #[test] - fn test_round_trip_sum() { - let acc = SumAccumulator::with_sum(42.5); - let bytes = serialize_accumulator_arroyo(&acc); - let restored = deserialize_accumulator(&bytes, &SketchType::Sum).unwrap(); - assert_eq!(restored.get_accumulator_type(), AggregationType::Sum); - - // Query the restored accumulator via single subpopulation - let restored_single = deserialize_single_subpopulation(&bytes, &SketchType::Sum).unwrap(); - let value = restored_single.query(Statistic::Sum, None).unwrap(); - assert!((value - 42.5).abs() < 1e-10, "Expected 42.5, got {}", value); - } - - #[test] - fn test_round_trip_kll() { - let mut kll = DatasketchesKLLAccumulator::new(200); - for v in [1.0, 2.0, 3.0, 4.0, 5.0] { - kll.update(v); - } - - let bytes = serialize_accumulator_arroyo(&kll); - let restored = deserialize_accumulator(&bytes, &SketchType::KLL).unwrap(); - assert_eq!( - restored.get_accumulator_type(), - AggregationType::DatasketchesKLL - ); - - // Query quantile via single subpopulation - let restored_single = deserialize_single_subpopulation(&bytes, &SketchType::KLL).unwrap(); - let mut kwargs = HashMap::new(); - kwargs.insert("quantile".to_string(), "0.5".to_string()); - let median = restored_single - .query(Statistic::Quantile, Some(&kwargs)) - .unwrap(); - // Median of [1,2,3,4,5] should be ~3.0 - assert!( - (1.0..=5.0).contains(&median), - "Median should be in [1,5], got {}", - median - ); - } - - #[test] - fn test_round_trip_set_aggregator() { - let mut set_acc = SetAggregatorAccumulator::new(); - set_acc.add_key(KeyByLabelValues { - labels: vec!["web".to_string()], - }); - set_acc.add_key(KeyByLabelValues { - labels: vec!["api".to_string()], - }); - set_acc.add_key(KeyByLabelValues { - labels: vec!["worker".to_string()], - }); - - let bytes = serialize_accumulator_arroyo(&set_acc); - let restored = deserialize_accumulator(&bytes, &SketchType::SetAggregator).unwrap(); - let keys = restored.get_keys().unwrap(); - assert_eq!(keys.len(), 3, "Expected 3 keys, got {}", keys.len()); - } - - #[test] - fn test_round_trip_delta_set_aggregator() { - let mut delta_set = DeltaSetAggregatorAccumulator::new(); - delta_set.add_key(KeyByLabelValues { - labels: vec!["endpoint-a".to_string()], - }); - delta_set.add_key(KeyByLabelValues { - labels: vec!["endpoint-b".to_string()], - }); - - let bytes = serialize_accumulator_arroyo(&delta_set); - let restored = deserialize_accumulator(&bytes, &SketchType::DeltaSetAggregator).unwrap(); - let keys = restored.get_keys().unwrap(); - assert_eq!(keys.len(), 2, "Expected 2 keys, got {}", keys.len()); - } - - #[test] - fn test_round_trip_multiple_increase() { - let key1 = KeyByLabelValues { - labels: vec!["web".to_string()], - }; - let key2 = KeyByLabelValues { - labels: vec!["api".to_string()], - }; - let mut increases = HashMap::new(); - increases.insert( - key1.clone(), - IncreaseAccumulator::new(Measurement::new(0.0), 0, Measurement::new(100.0), 10), - ); - increases.insert( - key2.clone(), - IncreaseAccumulator::new(Measurement::new(0.0), 0, Measurement::new(200.0), 10), - ); - - let acc = MultipleIncreaseAccumulator::new_with_increases(increases); - let bytes = serialize_accumulator_arroyo(&acc); - - let restored = deserialize_accumulator(&bytes, &SketchType::MultipleIncrease).unwrap(); - let keys = restored.get_keys().unwrap(); - assert_eq!(keys.len(), 2, "Expected 2 keys, got {}", keys.len()); - - // Query via multiple subpopulation - let restored_multi = - deserialize_multiple_subpopulation(&bytes, &SketchType::MultipleIncrease).unwrap(); - let val = restored_multi - .query(Statistic::Increase, &key1, None) - .unwrap(); - assert!( - (val - 100.0).abs() < 1e-10, - "Expected increase=100.0 for key1, got {}", - val - ); - } - - #[test] - fn test_round_trip_hydra_kll() { - // HydraKLL: serialize_to_bytes IS MessagePack, so serialize_accumulator_arroyo - // falls through to serialize_to_bytes - let mut hydra = HydraKllSketchAccumulator::new(1, 1, 200); - // Use the public update method with a key - hydra.update( - &KeyByLabelValues { - labels: vec!["sub-key".to_string()], - }, - 42.0, - ); - - let bytes = serialize_accumulator_arroyo(&hydra); - let restored = deserialize_accumulator(&bytes, &SketchType::HydraKLL).unwrap(); - assert_eq!(restored.get_accumulator_type(), AggregationType::HydraKLL); - } - - #[test] - fn test_round_trip_count_min_sketch() { - // CountMinSketch: serialize_to_bytes IS MessagePack - // Supported by both deserialize_accumulator (for merging) and - // deserialize_multiple_subpopulation (for querying by sub-key) - let cms = CountMinSketchAccumulator::new(2, 3); - - let bytes = serialize_accumulator_arroyo(&cms); - - let restored = deserialize_accumulator(&bytes, &SketchType::CountMinSketch).unwrap(); - assert_eq!( - restored.get_accumulator_type(), - AggregationType::CountMinSketch - ); - - let restored = - deserialize_multiple_subpopulation(&bytes, &SketchType::CountMinSketch).unwrap(); - assert_eq!( - restored.clone_boxed().as_ref().get_accumulator_type(), - AggregationType::CountMinSketch - ); - } - - // ======================================================================== - // Keys accumulator round-trip - // ======================================================================== - - #[test] - fn test_deserialize_keys_delta_set() { - let mut delta_set = DeltaSetAggregatorAccumulator::new(); - delta_set.add_key(KeyByLabelValues { - labels: vec!["key-a".to_string()], - }); - let bytes = serialize_accumulator_arroyo(&delta_set); - - let restored = - deserialize_keys_accumulator(&bytes, &SketchType::DeltaSetAggregator).unwrap(); - let keys = restored.get_keys().unwrap(); - assert_eq!(keys.len(), 1); - } - - #[test] - fn test_deserialize_keys_set_aggregator() { - let mut set_acc = SetAggregatorAccumulator::new(); - set_acc.add_key(KeyByLabelValues { - labels: vec!["k1".to_string()], - }); - let bytes = serialize_accumulator_arroyo(&set_acc); - - let restored = deserialize_keys_accumulator(&bytes, &SketchType::SetAggregator).unwrap(); - let keys = restored.get_keys().unwrap(); - assert_eq!(keys.len(), 1); - } - - // ======================================================================== - // Not-implemented types - // ======================================================================== - - #[test] - fn test_not_implemented_increase() { - let bytes = vec![1, 2, 3, 4]; - let result = deserialize_accumulator(&bytes, &SketchType::Increase); - assert!(result.is_err()); - let err_msg = match result { - Err(e) => format!("{}", e), - Ok(_) => panic!("Expected error"), - }; - assert!( - err_msg.contains("Not") || err_msg.contains("not"), - "Expected NotImplemented error, got: {}", - err_msg - ); - } - - #[test] - fn test_not_implemented_minmax() { - let bytes = vec![1, 2, 3, 4]; - let result = deserialize_accumulator(&bytes, &SketchType::MinMax); - assert!(result.is_err()); - } - - // ======================================================================== - // Error handling tests - // ======================================================================== - - #[test] - fn test_corrupted_bytes_sum() { - // MessagePack can decode some short byte sequences as valid integers/floats - // (e.g., 0xFF = -1 as negative fixint). Use 0xCB (float64 marker) followed - // by insufficient bytes to force a decode error. - let garbage = vec![0xCB, 0x01, 0x02]; - let result = deserialize_accumulator(&garbage, &SketchType::Sum); - assert!( - result.is_err(), - "Corrupted bytes should produce an error, not a panic" - ); - } - - #[test] - fn test_corrupted_bytes_kll() { - let garbage = vec![0xFF, 0xFE, 0xFD, 0xFC]; - let result = deserialize_accumulator(&garbage, &SketchType::KLL); - assert!( - result.is_err(), - "Corrupted bytes should produce an error, not a panic" - ); - } - - #[test] - fn test_empty_bytes_sum() { - let result = deserialize_accumulator(&[], &SketchType::Sum); - assert!(result.is_err(), "Empty bytes should produce an error"); - } - - #[test] - fn test_empty_bytes_kll() { - let result = deserialize_accumulator(&[], &SketchType::KLL); - assert!(result.is_err(), "Empty bytes should produce an error"); - } - - #[test] - fn test_empty_bytes_set_aggregator() { - let result = deserialize_accumulator(&[], &SketchType::SetAggregator); - assert!(result.is_err(), "Empty bytes should produce an error"); - } - - #[test] - fn test_empty_bytes_delta_set() { - let result = deserialize_accumulator(&[], &SketchType::DeltaSetAggregator); - assert!(result.is_err(), "Empty bytes should produce an error"); - } - - // ======================================================================== - // Serialize dispatch verification - // ======================================================================== - - #[test] - fn test_serialize_arroyo_dispatch_sum_uses_arroyo_path() { - // SumAccumulator has a separate arroyo format (MessagePack f64) - // while its native serialize_to_bytes uses little-endian f64 - let acc = SumAccumulator::with_sum(42.0); - let arroyo_bytes = serialize_accumulator_arroyo(&acc); - let native_bytes = acc.serialize_to_bytes(); - - // They should differ because arroyo uses MessagePack - assert_ne!( - arroyo_bytes, native_bytes, - "SumAccumulator arroyo and native serialization should differ" - ); - - // Verify the arroyo bytes can be deserialized - let restored = deserialize_accumulator(&arroyo_bytes, &SketchType::Sum).unwrap(); - assert_eq!(restored.get_accumulator_type(), AggregationType::Sum); - } - - #[test] - fn test_serialize_arroyo_dispatch_kll_uses_native() { - // KLL's serialize_to_bytes already uses MessagePack, so arroyo falls through - let mut kll = DatasketchesKLLAccumulator::new(200); - kll.update(1.0); - let arroyo_bytes = serialize_accumulator_arroyo(&kll); - let native_bytes = kll.serialize_to_bytes(); - - // They should be the same since KLL's native IS MessagePack - assert_eq!( - arroyo_bytes, native_bytes, - "KLL arroyo and native serialization should be the same" - ); - } - - #[test] - fn test_unsupported_keys_type() { - let bytes = vec![1, 2, 3, 4]; - let result = deserialize_keys_accumulator(&bytes, &SketchType::Sum); - assert!( - result.is_err(), - "Sum should not be supported as a keys accumulator" - ); - } -} diff --git a/asap-query-engine/src/tests/datafusion/dispatch_arithmetic_tests.rs b/asap-query-engine/src/tests/datafusion/dispatch_arithmetic_tests.rs deleted file mode 100644 index 69c888e7..00000000 --- a/asap-query-engine/src/tests/datafusion/dispatch_arithmetic_tests.rs +++ /dev/null @@ -1,140 +0,0 @@ -//! Dispatch-level tests for binary arithmetic PromQL handling. -//! -//! Tests that `handle_query_promql` routes binary expressions correctly: -//! returning `Some` for acceleratable queries and `None` for non-acceleratable -//! ones (graceful fallback to Prometheus). - -#[cfg(test)] -mod tests { - use crate::data_model::AggregationType; - use crate::precompute_operators::sum_accumulator::SumAccumulator; - use crate::tests::test_utilities::engine_factories::{ - create_engine_single_pop, create_engine_two_metrics, - }; - - const QUERY_TIME: f64 = 1000.0; - - #[tokio::test(flavor = "multi_thread")] - async fn test_handle_query_promql_binary_returns_result() { - let engine = create_engine_two_metrics( - "errors_total", - AggregationType::Sum, - vec!["host"], - vec![( - Some(vec!["host-a".to_string()]), - Box::new(SumAccumulator::with_sum(100.0)) as Box, - )], - "sum(errors_total) by (host)", - "requests_total", - AggregationType::Sum, - vec!["host"], - vec![( - Some(vec!["host-a".to_string()]), - Box::new(SumAccumulator::with_sum(200.0)) as Box, - )], - "sum(requests_total) by (host)", - ); - - let result = engine.handle_query_promql( - "sum(errors_total) by (host) / sum(requests_total) by (host)".to_string(), - QUERY_TIME, - ); - assert!(result.is_some(), "Binary query should return Some"); - let (labels, qr) = result.unwrap(); - assert!(!labels.labels.is_empty(), "Should have output label names"); - let elements = match qr { - crate::engines::query_result::QueryResult::Vector(iv) => iv.values, - _ => panic!("Expected vector result"), - }; - assert_eq!(elements.len(), 1); - assert!((elements[0].value - 0.5).abs() < 1e-10); - } - - #[tokio::test(flavor = "multi_thread")] - async fn test_handle_query_promql_non_acceleratable_arm_returns_none() { - // Only requests_total is configured; foo() is not a known function. - let engine = create_engine_single_pop( - "requests_total", - AggregationType::Sum, - vec!["host"], - vec![( - Some(vec!["host-a".to_string()]), - Box::new(SumAccumulator::with_sum(200.0)), - )], - "sum(requests_total) by (host)", - ); - - // foo() is not a supported PromQL function → arm lookup fails → returns None - let result = engine.handle_query_promql( - "foo(errors_total[5m]) / sum(requests_total) by (host)".to_string(), - QUERY_TIME, - ); - assert!( - result.is_none(), - "Should return None for non-acceleratable arm (graceful fallback)" - ); - } - - #[tokio::test(flavor = "multi_thread")] - async fn test_handle_query_promql_scalar_binary_returns_result() { - let engine = create_engine_single_pop( - "errors_total", - AggregationType::Sum, - vec!["host"], - vec![( - Some(vec!["host-a".to_string()]), - Box::new(SumAccumulator::with_sum(7.0)), - )], - "sum(errors_total) by (host)", - ); - - let result = - engine.handle_query_promql("sum(errors_total) by (host) * 100".to_string(), QUERY_TIME); - assert!(result.is_some(), "Scalar binary should return Some"); - let (_, qr) = result.unwrap(); - let elements = match qr { - crate::engines::query_result::QueryResult::Vector(iv) => iv.values, - _ => panic!("Expected vector result"), - }; - assert_eq!(elements.len(), 1); - assert!((elements[0].value - 700.0).abs() < 1e-10, "7 * 100 = 700"); - } - - #[tokio::test(flavor = "multi_thread")] - async fn test_handle_query_promql_single_metric_still_works() { - // Regression: single-metric queries continue to work after binary dispatch is wired in. - let engine = create_engine_single_pop( - "http_requests", - AggregationType::Sum, - vec!["host"], - vec![ - ( - Some(vec!["host-a".to_string()]), - Box::new(SumAccumulator::with_sum(100.0)), - ), - ( - Some(vec!["host-b".to_string()]), - Box::new(SumAccumulator::with_sum(200.0)), - ), - ], - "sum(http_requests) by (host)", - ); - - let result = - engine.handle_query_promql("sum(http_requests) by (host)".to_string(), QUERY_TIME); - assert!( - result.is_some(), - "Single-metric query should still work after binary dispatch" - ); - let (_, qr) = result.unwrap(); - let elements = match qr { - crate::engines::query_result::QueryResult::Vector(iv) => iv.values, - _ => panic!("Expected vector result"), - }; - assert_eq!(elements.len(), 2); - let mut values: Vec = elements.iter().map(|e| e.value).collect(); - values.sort_by(|a, b| a.partial_cmp(b).unwrap()); - assert!((values[0] - 100.0).abs() < 1e-10); - assert!((values[1] - 200.0).abs() < 1e-10); - } -} diff --git a/asap-query-engine/src/tests/datafusion/mod.rs b/asap-query-engine/src/tests/datafusion/mod.rs deleted file mode 100644 index c7f96a6b..00000000 --- a/asap-query-engine/src/tests/datafusion/mod.rs +++ /dev/null @@ -1,15 +0,0 @@ -//! DataFusion execution path tests. -//! -//! Tests for the logical plan builder, physical execution operators, -//! and accumulator serialization that back the DataFusion-based query path. - -pub mod accumulator_serde_tests; -pub mod dispatch_arithmetic_tests; -pub mod plan_builder_binary_tests; -pub mod plan_builder_regression_tests; -pub mod plan_execution_arithmetic_tests; -pub mod plan_execution_dual_input_tests; -pub mod plan_execution_temporal_tests; -pub mod plan_execution_tests; -pub mod structural_matching_tests; -pub mod warm_engine_replay_regression_tests; diff --git a/asap-query-engine/src/tests/datafusion/plan_builder_binary_tests.rs b/asap-query-engine/src/tests/datafusion/plan_builder_binary_tests.rs deleted file mode 100644 index c4da722b..00000000 --- a/asap-query-engine/src/tests/datafusion/plan_builder_binary_tests.rs +++ /dev/null @@ -1,190 +0,0 @@ -//! Binary plan builder tests. -//! -//! Tests that `build_binary_vector_plan` and `build_scalar_plan` produce correct -//! DataFusion logical plan structures. - -#[cfg(test)] -mod tests { - use crate::data_model::AggregationIdInfo; - use crate::engines::logical::plan_builder::{build_binary_vector_plan, build_scalar_plan}; - use crate::engines::simple::engine::{ - QueryExecutionContext, QueryMetadata, StoreQueryParams, StoreQueryPlan, - }; - use datafusion::logical_expr::LogicalPlan; - use promql_parser::parser::token::{TokenType, T_ADD, T_DIV, T_MOD, T_MUL, T_POW, T_SUB}; - use promql_utilities::data_model::KeyByLabelNames; - use promql_utilities::query_logics::enums::{AggregationType, Statistic}; - use std::collections::HashMap; - - fn make_context( - metric: &str, - statistic: Statistic, - labels: Vec<&str>, - ) -> QueryExecutionContext { - let label_strings: Vec = labels.into_iter().map(String::from).collect(); - QueryExecutionContext { - metric: metric.to_string(), - metadata: QueryMetadata { - query_output_labels: KeyByLabelNames::new(label_strings.clone()), - statistic_to_compute: statistic, - query_kwargs: HashMap::new(), - }, - store_plan: StoreQueryPlan { - values_query: StoreQueryParams { - metric: metric.to_string(), - aggregation_id: 1, - start_timestamp: 1000, - end_timestamp: 2000, - is_exact_query: true, - }, - keys_query: None, - }, - agg_info: AggregationIdInfo { - aggregation_id_for_key: 1, - aggregation_id_for_value: 1, - aggregation_type_for_key: AggregationType::Sum, - aggregation_type_for_value: AggregationType::Sum, - }, - do_merge: false, - spatial_filter: String::new(), - query_time: 2000, - grouping_labels: KeyByLabelNames::new(label_strings.clone()), - aggregated_labels: KeyByLabelNames::empty(), - } - } - - fn collect_node_names(plan: &LogicalPlan) -> Vec { - let mut names = Vec::new(); - collect_recursive(plan, &mut names); - names - } - - fn collect_recursive(plan: &LogicalPlan, names: &mut Vec) { - match plan { - LogicalPlan::Extension(ext) => { - names.push(ext.node.name().to_string()); - for input in ext.node.inputs() { - collect_recursive(input, names); - } - } - LogicalPlan::Projection(p) => { - names.push("Projection".to_string()); - collect_recursive(&p.input, names); - } - LogicalPlan::Join(j) => { - names.push("Join".to_string()); - collect_recursive(&j.left, names); - collect_recursive(&j.right, names); - } - LogicalPlan::SubqueryAlias(a) => { - names.push("SubqueryAlias".to_string()); - collect_recursive(&a.input, names); - } - other => { - names.push( - format!("{:?}", other) - .split('(') - .next() - .unwrap_or("Unknown") - .to_string(), - ); - } - } - } - - fn contains_node(plan: &LogicalPlan, name: &str) -> bool { - collect_node_names(plan).iter().any(|n| n == name) - } - - #[test] - fn test_binary_vector_plan_structure_divide() { - let lhs_ctx = make_context("errors", Statistic::Sum, vec!["host"]); - let rhs_ctx = make_context("requests", Statistic::Sum, vec!["host"]); - let lhs_plan = lhs_ctx.to_logical_plan().unwrap(); - let rhs_plan = rhs_ctx.to_logical_plan().unwrap(); - - let op = TokenType::new(T_DIV); - let plan = - build_binary_vector_plan(lhs_plan, rhs_plan, &op, vec!["host".to_string()]).unwrap(); - - let names = collect_node_names(&plan); - assert_eq!(names[0], "Projection", "Root should be Projection"); - assert!(contains_node(&plan, "Join"), "Plan should contain a Join"); - let alias_count = names.iter().filter(|n| *n == "SubqueryAlias").count(); - assert_eq!( - alias_count, 2, - "Plan should contain two SubqueryAlias nodes" - ); - } - - #[test] - fn test_binary_vector_plan_all_operators() { - let ops = [T_ADD, T_SUB, T_MUL, T_DIV, T_POW, T_MOD]; - for op_id in ops { - let lhs_ctx = make_context("metric_a", Statistic::Sum, vec!["host"]); - let rhs_ctx = make_context("metric_b", Statistic::Sum, vec!["host"]); - let lhs_plan = lhs_ctx.to_logical_plan().unwrap(); - let rhs_plan = rhs_ctx.to_logical_plan().unwrap(); - - let op = TokenType::new(op_id); - let result = - build_binary_vector_plan(lhs_plan, rhs_plan, &op, vec!["host".to_string()]); - assert!( - result.is_ok(), - "Operator {:?} should produce a valid plan", - op - ); - let names = collect_node_names(&result.unwrap()); - assert_eq!(names[0], "Projection"); - } - } - - #[test] - fn test_scalar_right_plan_structure() { - let ctx = make_context("errors", Statistic::Sum, vec!["host"]); - let vector_plan = ctx.to_logical_plan().unwrap(); - - let op = TokenType::new(T_MUL); - let plan = - build_scalar_plan(vector_plan, 100.0, &op, false, vec!["host".to_string()]).unwrap(); - - let names = collect_node_names(&plan); - assert_eq!(names[0], "Projection", "Root should be Projection"); - assert!( - !contains_node(&plan, "Join"), - "Scalar plan should not have a Join" - ); - assert!( - !contains_node(&plan, "SubqueryAlias"), - "Scalar plan should not have SubqueryAlias" - ); - } - - #[test] - fn test_scalar_left_plan_structure() { - let ctx = make_context("success", Statistic::Sum, vec!["host"]); - let vector_plan = ctx.to_logical_plan().unwrap(); - - let op = TokenType::new(T_SUB); - let plan = - build_scalar_plan(vector_plan, 1.0, &op, true, vec!["host".to_string()]).unwrap(); - - let names = collect_node_names(&plan); - assert_eq!(names[0], "Projection"); - assert!(!contains_node(&plan, "Join")); - } - - #[test] - fn test_scalar_left_division_plan_structure() { - // 1.0 / rate(metric[5m]) — scalar on left with Div - let ctx = make_context("metric", Statistic::Sum, vec!["host"]); - let vector_plan = ctx.to_logical_plan().unwrap(); - - let op = TokenType::new(T_DIV); - let result = build_scalar_plan(vector_plan, 1.0, &op, true, vec!["host".to_string()]); - assert!( - result.is_ok(), - "scalar-left division plan should build without error" - ); - } -} diff --git a/asap-query-engine/src/tests/datafusion/plan_builder_regression_tests.rs b/asap-query-engine/src/tests/datafusion/plan_builder_regression_tests.rs deleted file mode 100644 index 20fb465f..00000000 --- a/asap-query-engine/src/tests/datafusion/plan_builder_regression_tests.rs +++ /dev/null @@ -1,210 +0,0 @@ -//! Plan Builder Regression Tests -//! -//! Tests covering gaps in the existing plan_builder.rs inline tests: -//! all Statistic variants, kwargs propagation, error paths, edge cases. - -#[cfg(test)] -mod tests { - use crate::data_model::AggregationIdInfo; - use crate::engines::simple::engine::{ - QueryExecutionContext, QueryMetadata, StoreQueryParams, StoreQueryPlan, - }; - use promql_utilities::data_model::KeyByLabelNames; - use promql_utilities::query_logics::enums::{AggregationType, Statistic}; - use std::collections::HashMap; - - fn create_context( - statistic: Statistic, - aggregation_type: AggregationType, - output_labels: Vec<&str>, - kwargs: HashMap, - ) -> QueryExecutionContext { - let output_labels_vec: Vec = output_labels.into_iter().map(String::from).collect(); - let labels = KeyByLabelNames { - labels: output_labels_vec, - }; - QueryExecutionContext { - metric: "test_metric".to_string(), - metadata: QueryMetadata { - query_output_labels: labels.clone(), - statistic_to_compute: statistic, - query_kwargs: kwargs, - }, - store_plan: StoreQueryPlan { - values_query: StoreQueryParams { - metric: "test_metric".to_string(), - aggregation_id: 1, - start_timestamp: 1000, - end_timestamp: 2000, - is_exact_query: true, - }, - keys_query: None, - }, - agg_info: AggregationIdInfo { - aggregation_id_for_key: 1, - aggregation_id_for_value: 1, - aggregation_type_for_key: AggregationType::Sum, - aggregation_type_for_value: aggregation_type, - }, - do_merge: false, - spatial_filter: String::new(), - query_time: 2000, - grouping_labels: labels, - aggregated_labels: KeyByLabelNames::empty(), - } - } - - // ======================================================================== - // All Statistic variants map without panic - // ======================================================================== - - #[test] - fn test_all_statistics_map_without_panic() { - let statistics = vec![ - (Statistic::Sum, AggregationType::Sum), - (Statistic::Min, AggregationType::MinMax), - (Statistic::Max, AggregationType::MinMax), - (Statistic::Count, AggregationType::Sum), - (Statistic::Increase, AggregationType::Increase), - (Statistic::Rate, AggregationType::Increase), - (Statistic::Quantile, AggregationType::DatasketchesKLL), - (Statistic::Cardinality, AggregationType::SetAggregator), - (Statistic::Topk, AggregationType::CountMinSketch), - ]; - - for (stat, agg_type) in statistics { - let ctx = create_context(stat, agg_type, vec!["host"], HashMap::new()); - let result = ctx.map_statistic_to_infer_operation(); - assert!( - result.is_ok(), - "Statistic {:?} should map successfully, got: {:?}", - stat, - result.err() - ); - } - } - - // ======================================================================== - // TopK kwargs propagation - // ======================================================================== - - #[test] - fn test_topk_kwargs_propagate() { - use datafusion_summary_library::InferOperation; - let mut kwargs = HashMap::new(); - kwargs.insert("k".to_string(), "5".to_string()); - - let ctx = create_context( - Statistic::Topk, - AggregationType::CountMinSketch, - vec!["host"], - kwargs, - ); - match ctx.map_statistic_to_infer_operation().unwrap() { - InferOperation::TopK(k) => assert_eq!(k, 5, "Expected k=5, got {}", k), - other => panic!("Expected TopK, got {:?}", other), - } - } - - #[test] - fn test_topk_default_k() { - use datafusion_summary_library::InferOperation; - let ctx = create_context( - Statistic::Topk, - AggregationType::CountMinSketch, - vec!["host"], - HashMap::new(), - ); - match ctx.map_statistic_to_infer_operation().unwrap() { - InferOperation::TopK(k) => assert_eq!(k, 10, "Default k should be 10, got {}", k), - other => panic!("Expected TopK, got {:?}", other), - } - } - - // ======================================================================== - // Statistic-to-operation mapping - // ======================================================================== - - #[test] - fn test_cardinality_to_count_distinct() { - use datafusion_summary_library::InferOperation; - let ctx = create_context( - Statistic::Cardinality, - AggregationType::SetAggregator, - vec!["host"], - HashMap::new(), - ); - assert!(matches!( - ctx.map_statistic_to_infer_operation().unwrap(), - InferOperation::CountDistinct - )); - } - - #[test] - fn test_rate_to_extract_rate() { - use datafusion_summary_library::InferOperation; - let ctx = create_context( - Statistic::Rate, - AggregationType::Increase, - vec!["host"], - HashMap::new(), - ); - assert!(matches!( - ctx.map_statistic_to_infer_operation().unwrap(), - InferOperation::ExtractRate - )); - } - - #[test] - fn test_count_to_extract_count() { - use datafusion_summary_library::InferOperation; - let ctx = create_context( - Statistic::Count, - AggregationType::Sum, - vec!["host"], - HashMap::new(), - ); - assert!(matches!( - ctx.map_statistic_to_infer_operation().unwrap(), - InferOperation::ExtractCount - )); - } - - // ======================================================================== - // Error paths - // ======================================================================== - - #[test] - fn test_unknown_agg_type_errors() { - // SingleSubpopulation is a legacy wrapper variant not mapped to a SketchType - let ctx = create_context( - Statistic::Sum, - AggregationType::SingleSubpopulation, - vec!["host"], - HashMap::new(), - ); - let result = ctx.to_logical_plan(); - assert!(result.is_err(), "Unmapped aggregation type should error"); - let err_msg = format!("{}", result.unwrap_err()); - assert!( - err_msg.contains("Unknown"), - "Error should mention Unknown, got: {}", - err_msg - ); - } - - // ======================================================================== - // Edge cases - // ======================================================================== - - #[test] - fn test_plan_with_empty_labels() { - let ctx = create_context(Statistic::Sum, AggregationType::Sum, vec![], HashMap::new()); - let result = ctx.to_logical_plan(); - assert!( - result.is_ok(), - "Empty labels should still build a plan: {:?}", - result.err() - ); - } -} diff --git a/asap-query-engine/src/tests/datafusion/plan_execution_arithmetic_tests.rs b/asap-query-engine/src/tests/datafusion/plan_execution_arithmetic_tests.rs deleted file mode 100644 index 91d1356d..00000000 --- a/asap-query-engine/src/tests/datafusion/plan_execution_arithmetic_tests.rs +++ /dev/null @@ -1,369 +0,0 @@ -//! Binary arithmetic plan execution integration tests. -//! -//! Verify that binary arithmetic queries (vector/vector and scalar/vector) -//! produce numerically correct results when executed end-to-end through -//! `handle_binary_expr_promql` via DataFusion. - -#[cfg(test)] -mod tests { - use crate::data_model::AggregationType; - use crate::precompute_operators::sum_accumulator::SumAccumulator; - use crate::tests::test_utilities::engine_factories::{ - create_engine_three_metrics, create_engine_two_metrics, - }; - - const QUERY_TIME: f64 = 1000.0; - - fn host_a_b_data( - val_a: f64, - val_b: f64, - ) -> ( - crate::tests::test_utilities::engine_factories::AccumulatorData, - crate::tests::test_utilities::engine_factories::AccumulatorData, - ) { - let data_a = vec![ - ( - Some(vec!["host-a".to_string()]), - Box::new(SumAccumulator::with_sum(val_a)) as Box, - ), - ( - Some(vec!["host-b".to_string()]), - Box::new(SumAccumulator::with_sum(val_a / 2.0)), - ), - ]; - let data_b = vec![ - ( - Some(vec!["host-a".to_string()]), - Box::new(SumAccumulator::with_sum(val_b)) as Box, - ), - ( - Some(vec!["host-b".to_string()]), - Box::new(SumAccumulator::with_sum(val_b / 2.0)), - ), - ]; - (data_a, data_b) - } - - #[tokio::test(flavor = "multi_thread")] - async fn test_vector_vector_divide_produces_ratio() { - // errors/host-a = 100, requests/host-a = 200 → ratio 0.5 - let (data_errors, data_requests) = host_a_b_data(100.0, 200.0); - let engine = create_engine_two_metrics( - "errors_total", - AggregationType::Sum, - vec!["host"], - data_errors, - "sum(errors_total) by (host)", - "requests_total", - AggregationType::Sum, - vec!["host"], - data_requests, - "sum(requests_total) by (host)", - ); - - let query = "sum(errors_total) by (host) / sum(requests_total) by (host)"; - let result = engine.handle_query_promql(query.to_string(), QUERY_TIME); - assert!(result.is_some(), "Expected Some result for binary query"); - let (_, qr) = result.unwrap(); - let elements = match qr { - crate::engines::query_result::QueryResult::Vector(iv) => iv.values, - _ => panic!("Expected vector result"), - }; - assert_eq!(elements.len(), 2, "Expected 2 result rows"); - for elem in &elements { - let approx = (elem.value - 0.5).abs(); - assert!( - approx < 1e-10, - "Expected ratio 0.5, got {} for labels {:?}", - elem.value, - elem.labels - ); - } - } - - #[tokio::test(flavor = "multi_thread")] - async fn test_vector_vector_multiply() { - let (data_a, data_b) = host_a_b_data(3.0, 4.0); - let engine = create_engine_two_metrics( - "metric_a", - AggregationType::Sum, - vec!["host"], - data_a, - "sum(metric_a) by (host)", - "metric_b", - AggregationType::Sum, - vec!["host"], - data_b, - "sum(metric_b) by (host)", - ); - - let query = "sum(metric_a) by (host) * sum(metric_b) by (host)"; - let result = engine.handle_query_promql(query.to_string(), QUERY_TIME); - let (_, qr) = result.expect("Expected result"); - let elements = match qr { - crate::engines::query_result::QueryResult::Vector(iv) => iv.values, - _ => panic!("Expected vector result"), - }; - assert_eq!(elements.len(), 2); - // host-a: 3 * 4 = 12, host-b: 1.5 * 2.0 = 3.0 - let mut values: Vec = elements.iter().map(|e| e.value).collect(); - values.sort_by(|a, b| a.partial_cmp(b).unwrap()); - assert!((values[0] - 3.0).abs() < 1e-10); - assert!((values[1] - 12.0).abs() < 1e-10); - } - - #[tokio::test(flavor = "multi_thread")] - async fn test_vector_vector_add() { - let (data_a, data_b) = host_a_b_data(10.0, 20.0); - let engine = create_engine_two_metrics( - "metric_a", - AggregationType::Sum, - vec!["host"], - data_a, - "sum(metric_a) by (host)", - "metric_b", - AggregationType::Sum, - vec!["host"], - data_b, - "sum(metric_b) by (host)", - ); - - let query = "sum(metric_a) by (host) + sum(metric_b) by (host)"; - let result = engine.handle_query_promql(query.to_string(), QUERY_TIME); - let (_, qr) = result.expect("Expected result"); - let elements = match qr { - crate::engines::query_result::QueryResult::Vector(iv) => iv.values, - _ => panic!("Expected vector result"), - }; - assert_eq!(elements.len(), 2); - // host-a: 10+20=30, host-b: 5+10=15 - let mut values: Vec = elements.iter().map(|e| e.value).collect(); - values.sort_by(|a, b| a.partial_cmp(b).unwrap()); - assert!((values[0] - 15.0).abs() < 1e-10); - assert!((values[1] - 30.0).abs() < 1e-10); - } - - #[tokio::test(flavor = "multi_thread")] - async fn test_vector_vector_subtract() { - let (data_a, data_b) = host_a_b_data(50.0, 30.0); - let engine = create_engine_two_metrics( - "metric_a", - AggregationType::Sum, - vec!["host"], - data_a, - "sum(metric_a) by (host)", - "metric_b", - AggregationType::Sum, - vec!["host"], - data_b, - "sum(metric_b) by (host)", - ); - - let query = "sum(metric_a) by (host) - sum(metric_b) by (host)"; - let result = engine.handle_query_promql(query.to_string(), QUERY_TIME); - let (_, qr) = result.expect("Expected result"); - let elements = match qr { - crate::engines::query_result::QueryResult::Vector(iv) => iv.values, - _ => panic!("Expected vector result"), - }; - assert_eq!(elements.len(), 2); - // host-a: 50-30=20, host-b: 25-15=10 - let mut values: Vec = elements.iter().map(|e| e.value).collect(); - values.sort_by(|a, b| a.partial_cmp(b).unwrap()); - assert!((values[0] - 10.0).abs() < 1e-10); - assert!((values[1] - 20.0).abs() < 1e-10); - } - - #[tokio::test(flavor = "multi_thread")] - async fn test_vector_vector_inner_join_drops_unmatched() { - // errors has host-a and host-b; requests only has host-a - let data_errors = vec![ - ( - Some(vec!["host-a".to_string()]), - Box::new(SumAccumulator::with_sum(100.0)) as Box, - ), - ( - Some(vec!["host-b".to_string()]), - Box::new(SumAccumulator::with_sum(50.0)), - ), - ]; - let data_requests = vec![( - Some(vec!["host-a".to_string()]), - Box::new(SumAccumulator::with_sum(200.0)) as Box, - )]; - - let engine = create_engine_two_metrics( - "errors_total", - AggregationType::Sum, - vec!["host"], - data_errors, - "sum(errors_total) by (host)", - "requests_total", - AggregationType::Sum, - vec!["host"], - data_requests, - "sum(requests_total) by (host)", - ); - - let query = "sum(errors_total) by (host) / sum(requests_total) by (host)"; - let result = engine.handle_query_promql(query.to_string(), QUERY_TIME); - let (_, qr) = result.expect("Expected result"); - let elements = match qr { - crate::engines::query_result::QueryResult::Vector(iv) => iv.values, - _ => panic!("Expected vector result"), - }; - // Inner join: only host-a is present in both → 1 result - assert_eq!( - elements.len(), - 1, - "Inner join should drop unmatched label set (host-b)" - ); - assert!((elements[0].value - 0.5).abs() < 1e-10); - } - - #[tokio::test(flavor = "multi_thread")] - async fn test_scalar_right_multiply() { - // sum(errors_total) by (host) * 100 - let data = vec![( - Some(vec!["host-a".to_string()]), - Box::new(SumAccumulator::with_sum(5.0)) as Box, - )]; - let engine = create_engine_two_metrics( - "errors_total", - AggregationType::Sum, - vec!["host"], - data, - "sum(errors_total) by (host)", - // second metric not used but factory requires it; use empty data - "dummy", - AggregationType::Sum, - vec!["host"], - vec![], - "sum(dummy) by (host)", - ); - - let query = "sum(errors_total) by (host) * 100"; - let result = engine.handle_query_promql(query.to_string(), QUERY_TIME); - let (_, qr) = result.expect("Expected result for scalar-right multiply"); - let elements = match qr { - crate::engines::query_result::QueryResult::Vector(iv) => iv.values, - _ => panic!("Expected vector result"), - }; - assert_eq!(elements.len(), 1); - assert!( - (elements[0].value - 500.0).abs() < 1e-10, - "5.0 * 100 = 500.0" - ); - } - - #[tokio::test(flavor = "multi_thread")] - async fn test_scalar_left_subtract() { - // 1 - sum(success_total) by (host) - let data = vec![( - Some(vec!["host-a".to_string()]), - Box::new(SumAccumulator::with_sum(0.9)) as Box, - )]; - let engine = create_engine_two_metrics( - "success_total", - AggregationType::Sum, - vec!["host"], - data, - "sum(success_total) by (host)", - "dummy", - AggregationType::Sum, - vec!["host"], - vec![], - "sum(dummy) by (host)", - ); - - let query = "1 - sum(success_total) by (host)"; - let result = engine.handle_query_promql(query.to_string(), QUERY_TIME); - let (_, qr) = result.expect("Expected result for scalar-left subtract"); - let elements = match qr { - crate::engines::query_result::QueryResult::Vector(iv) => iv.values, - _ => panic!("Expected vector result"), - }; - assert_eq!(elements.len(), 1); - assert!( - (elements[0].value - 0.1).abs() < 1e-10, - "1 - 0.9 = 0.1, got {}", - elements[0].value - ); - } - - #[tokio::test(flavor = "multi_thread")] - async fn test_nested_binary_expression() { - // (sum(metric_a) by (host) + sum(metric_b) by (host)) / sum(metric_c) by (host) - // host-a: a=100, b=200, c=300 → (100+200)/300 = 1.0 - // host-b: a=50, b=100, c=150 → (50+100)/150 = 1.0 - let data_a = vec![ - ( - Some(vec!["host-a".to_string()]), - Box::new(SumAccumulator::with_sum(100.0)) as Box, - ), - ( - Some(vec!["host-b".to_string()]), - Box::new(SumAccumulator::with_sum(50.0)) as Box, - ), - ]; - let data_b = vec![ - ( - Some(vec!["host-a".to_string()]), - Box::new(SumAccumulator::with_sum(200.0)) as Box, - ), - ( - Some(vec!["host-b".to_string()]), - Box::new(SumAccumulator::with_sum(100.0)) as Box, - ), - ]; - let data_c = vec![ - ( - Some(vec!["host-a".to_string()]), - Box::new(SumAccumulator::with_sum(300.0)) as Box, - ), - ( - Some(vec!["host-b".to_string()]), - Box::new(SumAccumulator::with_sum(150.0)) as Box, - ), - ]; - - let engine = create_engine_three_metrics( - "metric_a", - AggregationType::Sum, - vec!["host"], - data_a, - "sum(metric_a) by (host)", - "metric_b", - AggregationType::Sum, - vec!["host"], - data_b, - "sum(metric_b) by (host)", - "metric_c", - AggregationType::Sum, - vec!["host"], - data_c, - "sum(metric_c) by (host)", - ); - - let query = "(sum(metric_a) by (host) + sum(metric_b) by (host)) / sum(metric_c) by (host)"; - let result = engine.handle_query_promql(query.to_string(), QUERY_TIME); - let (_, qr) = result.expect("Expected result for nested binary expression"); - let elements = match qr { - crate::engines::query_result::QueryResult::Vector(iv) => iv.values, - _ => panic!("Expected vector result"), - }; - assert_eq!( - elements.len(), - 2, - "Expected 2 result rows (host-a and host-b)" - ); - for elem in &elements { - assert!( - (elem.value - 1.0).abs() < 1e-10, - "Expected (a+b)/c = 1.0, got {} for labels {:?}", - elem.value, - elem.labels - ); - } - } -} diff --git a/asap-query-engine/src/tests/datafusion/plan_execution_dual_input_tests.rs b/asap-query-engine/src/tests/datafusion/plan_execution_dual_input_tests.rs deleted file mode 100644 index 649b2b24..00000000 --- a/asap-query-engine/src/tests/datafusion/plan_execution_dual_input_tests.rs +++ /dev/null @@ -1,312 +0,0 @@ -//! Plan Execution Multi-Population Tests -//! -//! Tests for multi-population accumulator execution: -//! -//! **Self-keyed (single-input):** MultipleIncrease, MultipleSum, MultipleMinMax -//! These carry their own keys via get_keys(). No separate keys stream needed. -//! grouping_labels = [], aggregated_labels = [all output labels] -//! -//! **Dual-input:** HydraKLL, CountMinSketch -//! These need a separate DeltaSetAggregator keys stream to enumerate sub-keys. - -#[cfg(test)] -mod tests { - use crate::data_model::{AggregationType, KeyByLabelValues, Measurement}; - use crate::precompute_operators::{ - CountMinSketchAccumulator, DeltaSetAggregatorAccumulator, HydraKllSketchAccumulator, - IncreaseAccumulator, MultipleIncreaseAccumulator, - }; - use crate::tests::test_utilities::engine_factories::*; - use std::collections::HashMap; - - /// Helper: create a MultipleIncreaseAccumulator with given sub-keys. - /// Each key is a Vec of label values (matching aggregated_labels order). - fn make_multi_increase( - keys_and_values: Vec<(Vec<&str>, f64, f64)>, - ) -> MultipleIncreaseAccumulator { - let mut increases = HashMap::new(); - for (key_labels, start_val, end_val) in keys_and_values { - let key = KeyByLabelValues { - labels: key_labels.iter().map(|s| s.to_string()).collect(), - }; - increases.insert( - key, - IncreaseAccumulator::new( - Measurement::new(start_val), - 0, - Measurement::new(end_val), - 10, - ), - ); - } - MultipleIncreaseAccumulator::new_with_increases(increases) - } - - // ======================================================================== - // Self-keyed: MultipleIncrease (single-input, keys from accumulator) - // MultipleIncrease is a collection of Increase accumulators. - // grouping_labels = [], aggregated_labels = [host, endpoint] - // Query is simply increase(metric[window]). - // ======================================================================== - - #[tokio::test] - async fn test_self_keyed_multiple_increase() { - // Sub-keys: (host-a, endpoint-1) increase=100, (host-a, endpoint-2) increase=200 - let acc = make_multi_increase(vec![ - (vec!["host-a", "endpoint-1"], 0.0, 100.0), - (vec!["host-a", "endpoint-2"], 0.0, 200.0), - ]); - - let engine = create_engine_single_pop_with_aggregated( - "http_requests_total", - AggregationType::MultipleIncrease, - vec![], - vec!["host", "endpoint"], - vec![(None, Box::new(acc))], - "increase(http_requests_total[10s])", - ); - - let results = execute_new_plan(&engine, "increase(http_requests_total[10s])", 1000.0).await; - - assert!( - !results.is_empty(), - "Self-keyed multi-pop should produce results; got 0" - ); - - // Should have 2 results (one per sub-key) - assert_eq!( - results.len(), - 2, - "Expected 2 results (2 sub-keys), got {}", - results.len() - ); - } - - #[tokio::test] - async fn test_self_keyed_multiple_increase_single_key() { - let acc = make_multi_increase(vec![(vec!["host-a", "svc-web"], 0.0, 50.0)]); - - let engine = create_engine_single_pop_with_aggregated( - "http_requests_total", - AggregationType::MultipleIncrease, - vec![], - vec!["host", "service"], - vec![(None, Box::new(acc))], - "increase(http_requests_total[10s])", - ); - - let results = execute_new_plan(&engine, "increase(http_requests_total[10s])", 1000.0).await; - - assert_eq!( - results.len(), - 1, - "Expected 1 result (1 sub-key), got {}", - results.len() - ); - } - - // ======================================================================== - // HydraKLL + DeltaSetAggregator (quantile dual-input) - // ======================================================================== - - #[tokio::test] - async fn test_dual_hydra_kll_delta_set() { - // HydraKLL: single accumulator, no spatial GROUP BY. - // Sub-keys are ["host", "endpoint"] — tracked by DeltaSet, queryable in HydraKLL. - // 2 columns per sub-key (host, endpoint). - let mut hydra = HydraKllSketchAccumulator::new(1, 2, 200); - hydra.update( - &KeyByLabelValues { - labels: vec!["host-a".to_string(), "endpoint-a".to_string()], - }, - 10.0, - ); - hydra.update( - &KeyByLabelValues { - labels: vec!["host-a".to_string(), "endpoint-a".to_string()], - }, - 20.0, - ); - hydra.update( - &KeyByLabelValues { - labels: vec!["host-a".to_string(), "endpoint-b".to_string()], - }, - 100.0, - ); - - // DeltaSet enumerates which sub-keys exist - let mut keys = DeltaSetAggregatorAccumulator::new(); - keys.add_key(KeyByLabelValues { - labels: vec!["host-a".to_string(), "endpoint-a".to_string()], - }); - keys.add_key(KeyByLabelValues { - labels: vec!["host-a".to_string(), "endpoint-b".to_string()], - }); - - let engine = create_engine_dual_input( - "request_duration", - AggregationType::HydraKLL, - AggregationType::DeltaSetAggregator, - vec![], // grouping_labels: no store GROUP BY - vec!["host", "endpoint"], // aggregated_labels: sub-keys tracked by DeltaSet - vec![(None, Box::new(hydra))], // store key = None - vec![(None, Box::new(keys))], // store key = None - "quantile(0.5, request_duration) by (host, endpoint)", - ); - - let results = execute_new_plan( - &engine, - "quantile(0.5, request_duration) by (host, endpoint)", - 1000.0, - ) - .await; - - // Should have 2 results: one per (host, endpoint) sub-key - assert_eq!( - results.len(), - 2, - "HydraKLL dual-input should produce 2 results, got {}", - results.len() - ); - } - - // ======================================================================== - // CountMinSketch + DeltaSetAggregator (frequency dual-input) - // ======================================================================== - - #[tokio::test] - async fn test_dual_count_min_delta_set() { - // CountMinSketch: single accumulator, no spatial GROUP BY. - // Sub-keys are ["host", "event"] — tracked by DeltaSet, queryable in CMS. - let cms = CountMinSketchAccumulator::new(2, 3); - - let mut keys = DeltaSetAggregatorAccumulator::new(); - keys.add_key(KeyByLabelValues { - labels: vec!["host-a".to_string(), "evt-1".to_string()], - }); - - let engine = create_engine_dual_input( - "event_frequency", - AggregationType::CountMinSketch, - AggregationType::DeltaSetAggregator, - vec![], // grouping_labels: no store GROUP BY - vec!["host", "event"], // aggregated_labels: sub-keys tracked by DeltaSet - vec![(None, Box::new(cms))], // store key = None - vec![(None, Box::new(keys))], // store key = None - "count(event_frequency) by (host, event)", - ); - - let results = - execute_new_plan(&engine, "count(event_frequency) by (host, event)", 1000.0).await; - - // CMS query may return 0 for un-updated keys, but should not error - assert!(!results.is_empty(), "CMS dual-input should produce results"); - } - - // ======================================================================== - // Self-keyed: multiple accumulators (multiple store entries) - // ======================================================================== - - #[tokio::test] - async fn test_self_keyed_multiple_accumulators() { - // Two separate MultipleIncrease accumulators in the store (both with key=None) - // acc1 has (host-a, ep-1) and (host-a, ep-2) - // acc2 has (host-b, ep-3) - let acc1 = make_multi_increase(vec![ - (vec!["host-a", "ep-1"], 0.0, 10.0), - (vec!["host-a", "ep-2"], 0.0, 20.0), - ]); - let acc2 = make_multi_increase(vec![(vec!["host-b", "ep-3"], 0.0, 30.0)]); - - let engine = create_engine_single_pop_with_aggregated( - "requests", - AggregationType::MultipleIncrease, - vec![], - vec!["host", "endpoint"], - vec![(None, Box::new(acc1)), (None, Box::new(acc2))], - "increase(requests[10s])", - ); - - let results = execute_new_plan(&engine, "increase(requests[10s])", 1000.0).await; - - // 2 from acc1 + 1 from acc2 = 3 total - assert_eq!( - results.len(), - 3, - "Expected 3 results (2 + 1), got {}", - results.len() - ); - } - - // ======================================================================== - // Self-keyed: empty accumulator - // ======================================================================== - - #[tokio::test] - async fn test_self_keyed_empty_keys() { - // MultipleIncrease with no sub-keys -> 0 results - let empty = MultipleIncreaseAccumulator::new_with_increases(HashMap::new()); - - let engine = create_engine_single_pop_with_aggregated( - "requests", - AggregationType::MultipleIncrease, - vec![], - vec!["host", "endpoint"], - vec![(None, Box::new(empty))], - "increase(requests[10s])", - ); - - let results = execute_new_plan(&engine, "increase(requests[10s])", 1000.0).await; - - assert!( - results.is_empty(), - "Empty MultipleIncrease should give 0 results, got {}", - results.len() - ); - } - - // ======================================================================== - // Dual-input: mismatched spatial groups (HydraKLL) - // ======================================================================== - - #[tokio::test] - async fn test_dual_mismatched_spatial_groups() { - // No spatial GROUP BY: single HydraKLL and single DeltaSet. - // HydraKLL has data for (host-a, ep-1), DeltaSet tracks (host-b, ep-2). - // The DeltaSet key doesn't match any data in HydraKLL, so query returns 0/default. - let mut hydra = HydraKllSketchAccumulator::new(1, 2, 200); - hydra.update( - &KeyByLabelValues { - labels: vec!["host-a".to_string(), "ep-1".to_string()], - }, - 42.0, - ); - - let mut keys = DeltaSetAggregatorAccumulator::new(); - keys.add_key(KeyByLabelValues { - labels: vec!["host-b".to_string(), "ep-2".to_string()], - }); - - let engine = create_engine_dual_input( - "request_duration", - AggregationType::HydraKLL, - AggregationType::DeltaSetAggregator, - vec![], // no store GROUP BY - vec!["host", "endpoint"], // aggregated_labels - vec![(None, Box::new(hydra))], - vec![(None, Box::new(keys))], - "quantile(0.5, request_duration) by (host, endpoint)", - ); - - let results = execute_new_plan( - &engine, - "quantile(0.5, request_duration) by (host, endpoint)", - 1000.0, - ) - .await; - - // DeltaSet enumerates (host-b, ep-2) but HydraKLL has no data for that key. - // We just verify it doesn't panic. - let _ = results; - } -} diff --git a/asap-query-engine/src/tests/datafusion/plan_execution_temporal_tests.rs b/asap-query-engine/src/tests/datafusion/plan_execution_temporal_tests.rs deleted file mode 100644 index 0cdaf6bc..00000000 --- a/asap-query-engine/src/tests/datafusion/plan_execution_temporal_tests.rs +++ /dev/null @@ -1,542 +0,0 @@ -//! Plan Execution Integration Tests — Temporal & Collapsable Queries -//! -//! Tests that verify the DataFusion plan-based execution path produces correct -//! results for temporal queries (sum_over_time, quantile_over_time) -//! and collapsable queries (spatial of temporal, e.g. sum by () (sum_over_time(...))). - -use crate::precompute_operators::sum_accumulator::SumAccumulator; -use std::collections::HashMap; - -#[cfg(test)] -mod tests { - use super::*; - use crate::data_model::{AggregationType, WindowType}; - use crate::precompute_operators::DatasketchesKLLAccumulator; - use crate::tests::test_utilities::engine_factories::*; - - type TemporalData = Vec<(u64, Option>, Box)>; - - // ======================================================================== - // Helper: build temporal data at 5 timestamps within a [5s] window - // Timestamps: 996_000, 997_000, 998_000, 999_000, 1_000_000 - // Query time: 1000.0 (= 1_000_000 ms) - // ======================================================================== - - const QUERY_TIME: f64 = 1000.0; - const TEMPORAL_TIMESTAMPS: [u64; 5] = [996_000, 997_000, 998_000, 999_000, 1_000_000]; - - // ======================================================================== - // OnlyTemporal Tests - // ======================================================================== - - #[tokio::test] - async fn test_temporal_sum_over_time_merges_across_timestamps() { - // Insert SumAccumulator data at 5 timestamps for one label group. - // sum_over_time should merge (sum) all values across timestamps. - let data: TemporalData = TEMPORAL_TIMESTAMPS - .iter() - .map(|&ts| { - ( - ts, - Some(vec!["host-a".to_string()]), - Box::new(SumAccumulator::with_sum(10.0)) as Box, - ) - }) - .collect(); - - let query = "sum_over_time(http_requests[5s])"; - let engine = create_engine_multi_timestamp_with_window( - "http_requests", - AggregationType::Sum, - vec!["host"], - data, - query, - 5, - WindowType::Tumbling, - ); - - let results = execute_new_plan(&engine, query, QUERY_TIME).await; - assert_eq!(results.len(), 1, "Expected 1 result for single group"); - // Merged sum: 10.0 * 5 timestamps = 50.0 - assert!( - (results[0].value - 50.0).abs() < 1e-10, - "Expected merged sum 50.0, got {}", - results[0].value - ); - } - - #[tokio::test] - async fn test_temporal_sum_over_time_single_timestamp() { - // Only 1 timestamp in range — should still work. - let data: TemporalData = vec![( - 1_000_000, - Some(vec!["host-a".to_string()]), - Box::new(SumAccumulator::with_sum(42.0)), - )]; - - let query = "sum_over_time(http_requests[5s])"; - let engine = create_engine_multi_timestamp_with_window( - "http_requests", - AggregationType::Sum, - vec!["host"], - data, - query, - 5, - WindowType::Tumbling, - ); - - let results = execute_new_plan(&engine, query, QUERY_TIME).await; - assert_eq!(results.len(), 1); - assert!( - (results[0].value - 42.0).abs() < 1e-10, - "Expected 42.0, got {}", - results[0].value - ); - } - - #[tokio::test] - async fn test_temporal_sum_over_time_varying_values() { - // sum_over_time with different values at each timestamp. - // Verifies merge sums all values, not just takes latest. - let data: TemporalData = TEMPORAL_TIMESTAMPS - .iter() - .enumerate() - .map(|(i, &ts)| { - ( - ts, - Some(vec!["host-a".to_string()]), - Box::new(SumAccumulator::with_sum((i as f64 + 1.0) * 10.0)) - as Box, - ) - }) - .collect(); - - let query = "sum_over_time(http_requests[5s])"; - let engine = create_engine_multi_timestamp_with_window( - "http_requests", - AggregationType::Sum, - vec!["host"], - data, - query, - 5, - WindowType::Tumbling, - ); - - let results = execute_new_plan(&engine, query, QUERY_TIME).await; - assert_eq!(results.len(), 1); - // 10 + 20 + 30 + 40 + 50 = 150 - assert!( - (results[0].value - 150.0).abs() < 1e-10, - "Expected merged sum 150.0, got {}", - results[0].value - ); - } - - #[tokio::test] - async fn test_temporal_quantile_over_time() { - // DatasketchesKLL at multiple timestamps, quantile_over_time(0.5, ...). - // Each KLL sketch has different values; merged sketch should give median. - let mut data: TemporalData = Vec::new(); - for (i, &ts) in TEMPORAL_TIMESTAMPS.iter().enumerate() { - let mut kll = DatasketchesKLLAccumulator::new(200); - // Insert values 10, 20, 30, 40, 50 at successive timestamps - kll.update((i as f64 + 1.0) * 10.0); - data.push(( - ts, - Some(vec!["host-a".to_string()]), - Box::new(kll) as Box, - )); - } - - let query = "quantile_over_time(0.5, latency[5s])"; - let engine = create_engine_multi_timestamp_with_window( - "latency", - AggregationType::DatasketchesKLL, - vec!["host"], - data, - query, - 5, - WindowType::Tumbling, - ); - - let results = execute_new_plan(&engine, query, QUERY_TIME).await; - assert_eq!(results.len(), 1); - // Median of {10, 20, 30, 40, 50} = 30.0 - assert!( - (results[0].value - 30.0).abs() < 5.0, - "Expected median ~30.0, got {}", - results[0].value - ); - } - - #[tokio::test] - async fn test_temporal_sum_over_time_multi_group() { - // Multiple label groups at multiple timestamps — verify per-group merging. - let mut data: TemporalData = Vec::new(); - for &ts in &TEMPORAL_TIMESTAMPS { - data.push(( - ts, - Some(vec!["host-a".to_string()]), - Box::new(SumAccumulator::with_sum(10.0)) as Box, - )); - data.push(( - ts, - Some(vec!["host-b".to_string()]), - Box::new(SumAccumulator::with_sum(20.0)) as Box, - )); - } - - let query = "sum_over_time(http_requests[5s])"; - let engine = create_engine_multi_timestamp_with_window( - "http_requests", - AggregationType::Sum, - vec!["host"], - data, - query, - 5, - WindowType::Tumbling, - ); - - let results = execute_new_plan(&engine, query, QUERY_TIME).await; - assert_eq!(results.len(), 2, "Expected 2 groups"); - - let result_map: HashMap = results - .iter() - .map(|r| (r.labels.labels[0].clone(), r.value)) - .collect(); - - assert!( - (result_map["host-a"] - 50.0).abs() < 1e-10, - "host-a: expected 50.0, got {}", - result_map["host-a"] - ); - assert!( - (result_map["host-b"] - 100.0).abs() < 1e-10, - "host-b: expected 100.0, got {}", - result_map["host-b"] - ); - } - - #[tokio::test] - async fn test_temporal_sum_over_time_empty_store() { - let query = "sum_over_time(http_requests[5s])"; - let engine = create_engine_multi_timestamp_with_window( - "http_requests", - AggregationType::Sum, - vec!["host"], - vec![], - query, - 5, - WindowType::Tumbling, - ); - - let results = execute_new_plan(&engine, query, QUERY_TIME).await; - assert!( - results.is_empty(), - "Empty store should return 0 results, got {}", - results.len() - ); - } - - #[tokio::test] - async fn test_temporal_context_has_do_merge_true() { - let query = "sum_over_time(http_requests[5s])"; - let engine = create_engine_multi_timestamp_with_window( - "http_requests", - AggregationType::Sum, - vec!["host"], - vec![( - 1_000_000, - Some(vec!["host-a".to_string()]), - Box::new(SumAccumulator::with_sum(1.0)), - )], - query, - 5, - WindowType::Tumbling, - ); - - let context = engine - .build_query_execution_context_promql(query.to_string(), QUERY_TIME) - .expect("Failed to build context"); - - assert!(context.do_merge, "Temporal queries must have do_merge=true"); - } - - #[tokio::test] - async fn test_temporal_context_has_correct_time_range() { - let query = "sum_over_time(http_requests[5s])"; - let engine = create_engine_multi_timestamp_with_window( - "http_requests", - AggregationType::Sum, - vec!["host"], - vec![( - 1_000_000, - Some(vec!["host-a".to_string()]), - Box::new(SumAccumulator::with_sum(1.0)), - )], - query, - 5, - WindowType::Tumbling, - ); - - let context = engine - .build_query_execution_context_promql(query.to_string(), QUERY_TIME) - .expect("Failed to build context"); - - // query_time = 1000.0 -> 1_000_000 ms - assert_eq!(context.query_time, 1_000_000); - // For [5s] range: start = end - 5000 = 995_000 - let start = context.store_plan.values_query.start_timestamp; - let end = context.store_plan.values_query.end_timestamp; - assert_eq!(end, 1_000_000, "End timestamp should be 1_000_000"); - assert_eq!( - start, 995_000, - "Start timestamp should be 995_000 for [5s] range" - ); - } - - // ======================================================================== - // Collapsable Tests (spatial of temporal) - // ======================================================================== - - #[tokio::test] - async fn test_collapsable_sum_of_sum_over_time() { - // sum by (host) (sum_over_time(metric[5s])) - // Multiple hosts at multiple timestamps. - let mut data: TemporalData = Vec::new(); - for &ts in &TEMPORAL_TIMESTAMPS { - data.push(( - ts, - Some(vec!["host-a".to_string()]), - Box::new(SumAccumulator::with_sum(10.0)) as Box, - )); - data.push(( - ts, - Some(vec!["host-b".to_string()]), - Box::new(SumAccumulator::with_sum(20.0)) as Box, - )); - } - - let query = "sum by (host) (sum_over_time(http_requests[5s]))"; - let engine = create_engine_multi_timestamp_with_window( - "http_requests", - AggregationType::Sum, - vec!["host"], - data, - query, - 5, - WindowType::Tumbling, - ); - - let results = execute_new_plan(&engine, query, QUERY_TIME).await; - assert_eq!(results.len(), 2, "Expected 2 groups (host-a, host-b)"); - - let result_map: HashMap = results - .iter() - .map(|r| (r.labels.labels[0].clone(), r.value)) - .collect(); - - // host-a: 10.0 * 5 timestamps = 50.0 - assert!( - (result_map["host-a"] - 50.0).abs() < 1e-10, - "host-a: expected 50.0, got {}", - result_map["host-a"] - ); - // host-b: 20.0 * 5 timestamps = 100.0 - assert!( - (result_map["host-b"] - 100.0).abs() < 1e-10, - "host-b: expected 100.0, got {}", - result_map["host-b"] - ); - } - - #[tokio::test] - async fn test_collapsable_sum_of_sum_over_time_varying_values() { - // sum by (host) (sum_over_time(metric[5s])) with varying values per timestamp - let mut data: TemporalData = Vec::new(); - for (i, &ts) in TEMPORAL_TIMESTAMPS.iter().enumerate() { - data.push(( - ts, - Some(vec!["host-a".to_string()]), - Box::new(SumAccumulator::with_sum((i as f64 + 1.0) * 10.0)) - as Box, - )); - data.push(( - ts, - Some(vec!["host-b".to_string()]), - Box::new(SumAccumulator::with_sum((i as f64 + 1.0) * 5.0)) - as Box, - )); - } - - let query = "sum by (host) (sum_over_time(http_requests[5s]))"; - let engine = create_engine_multi_timestamp_with_window( - "http_requests", - AggregationType::Sum, - vec!["host"], - data, - query, - 5, - WindowType::Tumbling, - ); - - let results = execute_new_plan(&engine, query, QUERY_TIME).await; - assert_eq!(results.len(), 2, "Expected 2 groups"); - - let result_map: HashMap = results - .iter() - .map(|r| (r.labels.labels[0].clone(), r.value)) - .collect(); - - // host-a: 10 + 20 + 30 + 40 + 50 = 150.0 - assert!( - (result_map["host-a"] - 150.0).abs() < 1e-10, - "host-a: expected 150.0, got {}", - result_map["host-a"] - ); - // host-b: 5 + 10 + 15 + 20 + 25 = 75.0 - assert!( - (result_map["host-b"] - 75.0).abs() < 1e-10, - "host-b: expected 75.0, got {}", - result_map["host-b"] - ); - } - - #[tokio::test] - async fn test_collapsable_context_has_do_merge_true() { - let query = "sum by (host) (sum_over_time(http_requests[5s]))"; - let engine = create_engine_multi_timestamp_with_window( - "http_requests", - AggregationType::Sum, - vec!["host"], - vec![( - 1_000_000, - Some(vec!["host-a".to_string()]), - Box::new(SumAccumulator::with_sum(1.0)), - )], - query, - 5, - WindowType::Tumbling, - ); - - let context = engine - .build_query_execution_context_promql(query.to_string(), QUERY_TIME) - .expect("Failed to build context"); - - assert!( - context.do_merge, - "Collapsable (OneTemporalOneSpatial) queries must have do_merge=true" - ); - } - - #[tokio::test] - async fn test_collapsable_output_labels_are_spatial() { - // Verify output labels are the spatial GROUP BY labels (host), not all labels. - let query = "sum by (host) (sum_over_time(http_requests[5s]))"; - let engine = create_engine_multi_timestamp_with_window( - "http_requests", - AggregationType::Sum, - vec!["host"], - vec![( - 1_000_000, - Some(vec!["host-a".to_string()]), - Box::new(SumAccumulator::with_sum(1.0)), - )], - query, - 5, - WindowType::Tumbling, - ); - - let context = engine - .build_query_execution_context_promql(query.to_string(), QUERY_TIME) - .expect("Failed to build context"); - - assert_eq!( - context.metadata.query_output_labels.labels, - vec!["host".to_string()], - "Collapsable query output labels should be spatial GROUP BY labels" - ); - } - - #[tokio::test] - async fn test_collapsable_empty_store() { - let query = "sum by (host) (sum_over_time(http_requests[5s]))"; - let engine = create_engine_multi_timestamp_with_window( - "http_requests", - AggregationType::Sum, - vec!["host"], - vec![], - query, - 5, - WindowType::Tumbling, - ); - - let results = execute_new_plan(&engine, query, QUERY_TIME).await; - assert!( - results.is_empty(), - "Empty store should return 0 results, got {}", - results.len() - ); - } - - // ======================================================================== - // Old-vs-New comparison tests for temporal queries - // ======================================================================== - - #[tokio::test] - async fn test_temporal_sum_over_time_old_vs_new() { - let data: TemporalData = TEMPORAL_TIMESTAMPS - .iter() - .map(|&ts| { - ( - ts, - Some(vec!["host-a".to_string()]), - Box::new(SumAccumulator::with_sum(10.0)) as Box, - ) - }) - .collect(); - - let query = "sum_over_time(http_requests[5s])"; - let engine = create_engine_multi_timestamp_with_window( - "http_requests", - AggregationType::Sum, - vec!["host"], - data, - query, - 5, - WindowType::Tumbling, - ); - - assert_old_new_match(&engine, query, QUERY_TIME).await; - } - - #[tokio::test] - async fn test_collapsable_sum_of_sum_over_time_old_vs_new() { - let mut data: TemporalData = Vec::new(); - for &ts in &TEMPORAL_TIMESTAMPS { - data.push(( - ts, - Some(vec!["host-a".to_string()]), - Box::new(SumAccumulator::with_sum(10.0)) as Box, - )); - data.push(( - ts, - Some(vec!["host-b".to_string()]), - Box::new(SumAccumulator::with_sum(20.0)) as Box, - )); - } - - let query = "sum by (host) (sum_over_time(http_requests[5s]))"; - let engine = create_engine_multi_timestamp_with_window( - "http_requests", - AggregationType::Sum, - vec!["host"], - data, - query, - 5, - WindowType::Tumbling, - ); - - assert_old_new_match(&engine, query, QUERY_TIME).await; - } -} diff --git a/asap-query-engine/src/tests/datafusion/plan_execution_tests.rs b/asap-query-engine/src/tests/datafusion/plan_execution_tests.rs deleted file mode 100644 index 2edd73cf..00000000 --- a/asap-query-engine/src/tests/datafusion/plan_execution_tests.rs +++ /dev/null @@ -1,587 +0,0 @@ -//! Plan Execution Integration Tests -//! -//! Tests that verify the new DataFusion plan-based execution path -//! produces correct results for spatial queries. -//! -//! These tests use an actual store with test data. - -use crate::data_model::{AggregationType, KeyByLabelValues, Measurement}; -use crate::engines::simple::engine::SimpleEngine; -use crate::precompute_operators::sum_accumulator::SumAccumulator; -use std::collections::HashMap; - -#[cfg(test)] -mod tests { - use super::*; - use crate::precompute_operators::{ - DatasketchesKLLAccumulator, IncreaseAccumulator, MinMaxAccumulator, - MultipleMinMaxAccumulator, MultipleSumAccumulator, SetAggregatorAccumulator, - }; - use crate::tests::test_utilities::engine_factories::*; - - /// Creates a test engine and store with sample data for spatial sum queries - fn create_test_engine_with_data() -> SimpleEngine { - create_engine_single_pop( - "http_requests", - AggregationType::Sum, - vec!["host"], - vec![ - ( - Some(vec!["host-a".to_string()]), - Box::new(SumAccumulator::with_sum(100.0)), - ), - ( - Some(vec!["host-b".to_string()]), - Box::new(SumAccumulator::with_sum(200.0)), - ), - ], - "sum(http_requests) by (host)", - ) - } - - #[test] - fn test_to_logical_plan_produces_valid_structure() { - let engine = create_test_engine_with_data(); - - let query_time_sec = 1000.0; - let context = engine - .build_query_execution_context_promql( - "sum(http_requests) by (host)".to_string(), - query_time_sec, - ) - .expect("Failed to build context"); - - let plan = context.to_logical_plan(); - assert!( - plan.is_ok(), - "Failed to build logical plan: {:?}", - plan.err() - ); - - let plan = plan.unwrap(); - match &plan { - datafusion::logical_expr::LogicalPlan::Extension(ext) => { - assert_eq!( - ext.node.name(), - "SummaryInfer", - "Root should be SummaryInfer" - ); - } - _ => panic!("Expected Extension node at root, got {:?}", plan), - } - } - - #[tokio::test] - async fn test_execute_plan_returns_results() { - let engine = create_test_engine_with_data(); - - let query_time_sec = 1000.0; - let context = engine - .build_query_execution_context_promql( - "sum(http_requests) by (host)".to_string(), - query_time_sec, - ) - .expect("Failed to build context"); - - let result = engine.execute_plan(&context).await; - assert!(result.is_ok(), "execute_plan failed: {:?}", result.err()); - - let results = result.unwrap(); - assert_eq!( - results.len(), - 2, - "Expected 2 results, got {}", - results.len() - ); - - let values: Vec = results.iter().map(|r| r.value).collect(); - assert!( - values.contains(&100.0) && values.contains(&200.0), - "Expected values [100.0, 200.0], got {:?}", - values - ); - } - - #[tokio::test] - async fn test_execute_plan_correct_labels() { - let engine = create_test_engine_with_data(); - - let query_time_sec = 1000.0; - let context = engine - .build_query_execution_context_promql( - "sum(http_requests) by (host)".to_string(), - query_time_sec, - ) - .expect("Failed to build context"); - - let results = engine - .execute_plan(&context) - .await - .expect("execute_plan failed"); - - let result_map: HashMap = results - .iter() - .map(|r| (r.labels.labels[0].clone(), r.value)) - .collect(); - - assert_eq!(result_map.get("host-a"), Some(&100.0)); - assert_eq!(result_map.get("host-b"), Some(&200.0)); - } - - #[tokio::test] - async fn test_execute_plan_multiple_timestamps() { - let engine = create_engine_multi_timestamp( - "http_requests", - AggregationType::Sum, - vec!["host"], - vec![ - ( - 999_000, - Some(vec!["host-a".to_string()]), - Box::new(SumAccumulator::with_sum(50.0)), - ), - ( - 1_000_000, - Some(vec!["host-a".to_string()]), - Box::new(SumAccumulator::with_sum(50.0)), - ), - ], - "sum(http_requests) by (host)", - ); - - let context = engine - .build_query_execution_context_promql( - "sum(http_requests) by (host)".to_string(), - 1000.0, - ) - .expect("Failed to build context"); - - let results = engine - .execute_plan(&context) - .await - .expect("execute_plan failed"); - - assert_eq!(results.len(), 1, "Expected 1 result"); - assert_eq!(results[0].labels.labels[0], "host-a"); - assert_eq!( - results[0].value, 50.0, - "Spatial-only queries use latest timestamp only" - ); - } - - #[tokio::test] - async fn test_execute_plan_matches_old_pipeline() { - let engine = create_test_engine_with_data(); - assert_old_new_match(&engine, "sum(http_requests) by (host)", 1000.0).await; - } - - // ======================================================================== - // Category 1 - Old-vs-New Comparison for more accumulator types - // ======================================================================== - - #[tokio::test] - async fn test_old_vs_new_kll_quantile() { - let mut kll_a = DatasketchesKLLAccumulator::new(200); - for v in [10.0, 20.0, 30.0, 40.0, 50.0] { - kll_a.update(v); - } - let mut kll_b = DatasketchesKLLAccumulator::new(200); - for v in [100.0, 200.0, 300.0] { - kll_b.update(v); - } - - let engine = create_engine_single_pop( - "latency", - AggregationType::DatasketchesKLL, - vec!["host"], - vec![ - (Some(vec!["host-a".to_string()]), Box::new(kll_a)), - (Some(vec!["host-b".to_string()]), Box::new(kll_b)), - ], - "quantile(0.5, latency) by (host)", - ); - - assert_old_new_match(&engine, "quantile(0.5, latency) by (host)", 1000.0).await; - } - - #[tokio::test] - async fn test_old_vs_new_kll_quantile_p99() { - let mut kll = DatasketchesKLLAccumulator::new(200); - for v in 1..=100 { - kll.update(v as f64); - } - - let engine = create_engine_single_pop( - "latency", - AggregationType::DatasketchesKLL, - vec!["host"], - vec![(Some(vec!["host-a".to_string()]), Box::new(kll))], - "quantile(0.99, latency) by (host)", - ); - - assert_old_new_match(&engine, "quantile(0.99, latency) by (host)", 1000.0).await; - } - - #[tokio::test] - async fn test_old_vs_new_kll_quantile_p0() { - let mut kll = DatasketchesKLLAccumulator::new(200); - for v in [5.0, 10.0, 15.0, 20.0, 25.0] { - kll.update(v); - } - - let engine = create_engine_single_pop( - "latency", - AggregationType::DatasketchesKLL, - vec!["host"], - vec![(Some(vec!["host-a".to_string()]), Box::new(kll))], - "quantile(0.0, latency) by (host)", - ); - - assert_old_new_match(&engine, "quantile(0.0, latency) by (host)", 1000.0).await; - } - - #[tokio::test] - async fn test_old_vs_new_kll_quantile_p1() { - let mut kll = DatasketchesKLLAccumulator::new(200); - for v in [5.0, 10.0, 15.0, 20.0, 25.0] { - kll.update(v); - } - - let engine = create_engine_single_pop( - "latency", - AggregationType::DatasketchesKLL, - vec!["host"], - vec![(Some(vec!["host-a".to_string()]), Box::new(kll))], - "quantile(1.0, latency) by (host)", - ); - - assert_old_new_match(&engine, "quantile(1.0, latency) by (host)", 1000.0).await; - } - - #[tokio::test] - async fn test_old_vs_new_kll_quantile_p25() { - let mut kll = DatasketchesKLLAccumulator::new(200); - for v in 1..=1000 { - kll.update(v as f64); - } - - let engine = create_engine_single_pop( - "latency", - AggregationType::DatasketchesKLL, - vec!["host"], - vec![(Some(vec!["host-a".to_string()]), Box::new(kll))], - "quantile(0.25, latency) by (host)", - ); - - assert_old_new_match(&engine, "quantile(0.25, latency) by (host)", 1000.0).await; - } - - #[tokio::test] - #[ignore = "Blocked: SetAggregatorAccumulator does not support old pipeline query"] - async fn test_old_vs_new_set_aggregator_cardinality() { - let mut set_a = SetAggregatorAccumulator::new(); - set_a.add_key(KeyByLabelValues { - labels: vec!["user-1".to_string()], - }); - set_a.add_key(KeyByLabelValues { - labels: vec!["user-2".to_string()], - }); - let mut set_b = SetAggregatorAccumulator::new(); - set_b.add_key(KeyByLabelValues { - labels: vec!["user-3".to_string()], - }); - - let engine = create_engine_single_pop( - "active_users", - AggregationType::SetAggregator, - vec!["host"], - vec![ - (Some(vec!["host-a".to_string()]), Box::new(set_a)), - (Some(vec!["host-b".to_string()]), Box::new(set_b)), - ], - "count(active_users) by (host)", - ); - - assert_old_new_match(&engine, "count(active_users) by (host)", 1000.0).await; - } - - #[tokio::test] - #[ignore = "Blocked: IncreaseAccumulator has no arroyo serde"] - async fn test_old_vs_new_increase() { - let inc_a = IncreaseAccumulator::new(Measurement::new(0.0), 0, Measurement::new(100.0), 10); - let inc_b = IncreaseAccumulator::new(Measurement::new(0.0), 0, Measurement::new(200.0), 10); - - let engine = create_engine_single_pop( - "http_requests_total", - AggregationType::Increase, - vec!["host"], - vec![ - (Some(vec!["host-a".to_string()]), Box::new(inc_a)), - (Some(vec!["host-b".to_string()]), Box::new(inc_b)), - ], - "sum(increase(http_requests_total[10s])) by (host)", - ); - - assert_old_new_match( - &engine, - "sum(increase(http_requests_total[10s])) by (host)", - 1000.0, - ) - .await; - } - - #[tokio::test] - #[ignore = "Blocked: MinMaxAccumulator has no arroyo serde"] - async fn test_old_vs_new_minmax_min() { - let mm_a = MinMaxAccumulator::with_value(10.0, "min".to_string()); - let mm_b = MinMaxAccumulator::with_value(5.0, "min".to_string()); - - let engine = create_engine_single_pop( - "temperature", - AggregationType::MinMax, - vec!["host"], - vec![ - (Some(vec!["host-a".to_string()]), Box::new(mm_a)), - (Some(vec!["host-b".to_string()]), Box::new(mm_b)), - ], - "min(temperature) by (host)", - ); - - assert_old_new_match(&engine, "min(temperature) by (host)", 1000.0).await; - } - - #[tokio::test] - #[ignore = "Blocked: MinMaxAccumulator has no arroyo serde"] - async fn test_old_vs_new_minmax_max() { - let mm_a = MinMaxAccumulator::with_value(90.0, "max".to_string()); - let mm_b = MinMaxAccumulator::with_value(95.0, "max".to_string()); - - let engine = create_engine_single_pop( - "temperature", - AggregationType::MinMax, - vec!["host"], - vec![ - (Some(vec!["host-a".to_string()]), Box::new(mm_a)), - (Some(vec!["host-b".to_string()]), Box::new(mm_b)), - ], - "max(temperature) by (host)", - ); - - assert_old_new_match(&engine, "max(temperature) by (host)", 1000.0).await; - } - - #[tokio::test] - async fn test_old_vs_new_multiple_sum() { - let mut ms = MultipleSumAccumulator::new(); - let key = KeyByLabelValues { - labels: vec!["host-a".to_string()], - }; - ms.update(key, 42.0); - - let engine = create_engine_single_pop( - "requests", - AggregationType::MultipleSum, - vec!["host"], - vec![(Some(vec!["host-a".to_string()]), Box::new(ms))], - "sum(requests) by (host)", - ); - - assert_old_new_match(&engine, "sum(requests) by (host)", 1000.0).await; - } - - #[tokio::test] - #[ignore = "Blocked: MultipleMinMaxAccumulator has no arroyo serde"] - async fn test_old_vs_new_multiple_minmax() { - let mm = MultipleMinMaxAccumulator::new("min".to_string()); - - let engine = create_engine_single_pop( - "latency", - AggregationType::MultipleMinMax, - vec!["host"], - vec![(Some(vec!["host-a".to_string()]), Box::new(mm))], - "min(latency) by (host)", - ); - - assert_old_new_match(&engine, "min(latency) by (host)", 1000.0).await; - } - - // ======================================================================== - // Category 3 - Edge Cases - // ======================================================================== - - #[tokio::test] - async fn test_execute_plan_empty_store() { - let engine = create_engine_single_pop( - "http_requests", - AggregationType::Sum, - vec!["host"], - vec![], // No data - "sum(http_requests) by (host)", - ); - - let results = execute_new_plan(&engine, "sum(http_requests) by (host)", 1000.0).await; - assert!( - results.is_empty(), - "Empty store should return 0 results, got {}", - results.len() - ); - } - - #[tokio::test] - async fn test_execute_plan_many_groups() { - #[allow(clippy::type_complexity)] - let data: Vec<(Option>, Box)> = (0..100) - .map(|i| { - ( - Some(vec![format!("host-{:03}", i)]), - Box::new(SumAccumulator::with_sum(i as f64)) as Box, - ) - }) - .collect(); - - let engine = create_engine_single_pop( - "http_requests", - AggregationType::Sum, - vec!["host"], - data, - "sum(http_requests) by (host)", - ); - - let results = execute_new_plan(&engine, "sum(http_requests) by (host)", 1000.0).await; - assert_eq!( - results.len(), - 100, - "Expected 100 results, got {}", - results.len() - ); - } - - #[tokio::test] - async fn test_execute_plan_multi_timestamp_multi_group() { - // 3 timestamps x 3 groups; spatial-only uses latest timestamp only -> 3 results - let mut data = Vec::new(); - for ts in [998_000u64, 999_000, 1_000_000] { - for i in 0..3 { - data.push(( - ts, - Some(vec![format!("host-{}", i)]), - Box::new(SumAccumulator::with_sum(10.0)) as Box, - )); - } - } - - let engine = create_engine_multi_timestamp( - "http_requests", - AggregationType::Sum, - vec!["host"], - data, - "sum(http_requests) by (host)", - ); - - let results = execute_new_plan(&engine, "sum(http_requests) by (host)", 1000.0).await; - assert_eq!( - results.len(), - 3, - "3 groups merged across 3 timestamps should give 3 results" - ); - // Spatial-only queries use latest timestamp only, so each group = 10.0 - for r in &results { - assert!( - (r.value - 10.0).abs() < 1e-10, - "Each group should be 10.0 (latest timestamp only), got {}", - r.value - ); - } - } - - #[tokio::test] - async fn test_execute_plan_single_group_many_timestamps() { - // Spatial-only queries only consider the latest timestamp (query time). - // Include data at the query time (1_000_000) plus older timestamps. - #[allow(clippy::type_complexity)] - let mut data: Vec<(u64, Option>, Box)> = (0..9) - .map(|i| { - ( - (991_000 + i * 1000) as u64, - Some(vec!["host-a".to_string()]), - Box::new(SumAccumulator::with_sum(5.0)) as Box, - ) - }) - .collect(); - data.push(( - 1_000_000, - Some(vec!["host-a".to_string()]), - Box::new(SumAccumulator::with_sum(5.0)) as Box, - )); - - let engine = create_engine_multi_timestamp( - "http_requests", - AggregationType::Sum, - vec!["host"], - data, - "sum(http_requests) by (host)", - ); - - let results = execute_new_plan(&engine, "sum(http_requests) by (host)", 1000.0).await; - assert_eq!(results.len(), 1, "Single group should give 1 result"); - assert!( - (results[0].value - 5.0).abs() < 1e-10, - "Spatial-only uses latest timestamp only, expected 5.0, got {}", - results[0].value - ); - } - - // ======================================================================== - // Category 8 - Error Paths - // ======================================================================== - - #[tokio::test] - async fn test_execute_plan_not_implemented_increase() { - // IncreaseAccumulator has no arroyo serde, so execute_plan should fail - let inc = IncreaseAccumulator::new(Measurement::new(0.0), 0, Measurement::new(100.0), 10); - - let engine = create_engine_single_pop( - "http_requests_total", - AggregationType::Increase, - vec!["host"], - vec![(Some(vec!["host-a".to_string()]), Box::new(inc))], - "sum(increase(http_requests_total[10s])) by (host)", - ); - - let context = engine - .build_query_execution_context_promql( - "sum(increase(http_requests_total[10s])) by (host)".to_string(), - 1000.0, - ) - .expect("Should build context"); - - let result = engine.execute_plan(&context).await; - assert!( - result.is_err(), - "execute_plan should fail for IncreaseAccumulator (no arroyo serde)" - ); - } - - #[tokio::test] - async fn test_execute_plan_not_implemented_minmax() { - let mm = MinMaxAccumulator::with_value(42.0, "min".to_string()); - - let engine = create_engine_single_pop( - "temperature", - AggregationType::MinMax, - vec!["host"], - vec![(Some(vec!["host-a".to_string()]), Box::new(mm))], - "min(temperature) by (host)", - ); - - let context = engine - .build_query_execution_context_promql("min(temperature) by (host)".to_string(), 1000.0) - .expect("Should build context"); - - let result = engine.execute_plan(&context).await; - assert!( - result.is_err(), - "execute_plan should fail for MinMaxAccumulator (no arroyo serde)" - ); - } -} diff --git a/asap-query-engine/src/tests/datafusion/structural_matching_tests.rs b/asap-query-engine/src/tests/datafusion/structural_matching_tests.rs deleted file mode 100644 index d04d99aa..00000000 --- a/asap-query-engine/src/tests/datafusion/structural_matching_tests.rs +++ /dev/null @@ -1,121 +0,0 @@ -//! Structural PromQL matching tests. -//! -//! Verifies that `find_query_config_promql_structural` can look up query configs -//! by AST-serialised arm strings, which is the mechanism used during binary -//! arithmetic dispatch. - -#[cfg(test)] -mod tests { - use crate::data_model::AggregationType; - use crate::precompute_operators::sum_accumulator::SumAccumulator; - use crate::tests::test_utilities::engine_factories::create_engine_single_pop; - - #[test] - fn test_structural_match_rate_query_finds_config() { - let engine = create_engine_single_pop( - "http_requests_total", - AggregationType::MultipleIncrease, - vec!["host"], - vec![( - Some(vec!["host-a".to_string()]), - Box::new(SumAccumulator::with_sum(100.0)), - )], - "rate(http_requests_total[5m])", - ); - - let ast = - promql_parser::parser::parse("rate(http_requests_total[5m])").expect("parse failed"); - let result = engine.find_query_config_promql_structural(&ast); - assert!( - result.is_some(), - "Expected to find config for rate query, got None" - ); - } - - #[test] - fn test_structural_match_wrong_metric_returns_none() { - let engine = create_engine_single_pop( - "http_requests_total", - AggregationType::MultipleIncrease, - vec!["host"], - vec![( - Some(vec!["host-a".to_string()]), - Box::new(SumAccumulator::with_sum(100.0)), - )], - "rate(http_requests_total[5m])", - ); - - let ast = promql_parser::parser::parse("rate(other_metric[5m])").expect("parse failed"); - let result = engine.find_query_config_promql_structural(&ast); - assert!( - result.is_none(), - "Should not find config for different metric" - ); - } - - #[test] - fn test_structural_match_wrong_range_returns_none() { - let engine = create_engine_single_pop( - "http_requests_total", - AggregationType::MultipleIncrease, - vec!["host"], - vec![( - Some(vec!["host-a".to_string()]), - Box::new(SumAccumulator::with_sum(100.0)), - )], - "rate(http_requests_total[5m])", - ); - - let ast = - promql_parser::parser::parse("rate(http_requests_total[1m])").expect("parse failed"); - let result = engine.find_query_config_promql_structural(&ast); - assert!( - result.is_none(), - "Should not find config for different range" - ); - } - - #[test] - fn test_structural_match_wrong_function_returns_none() { - let engine = create_engine_single_pop( - "http_requests_total", - AggregationType::MultipleIncrease, - vec!["host"], - vec![( - Some(vec!["host-a".to_string()]), - Box::new(SumAccumulator::with_sum(100.0)), - )], - "rate(http_requests_total[5m])", - ); - - let ast = promql_parser::parser::parse("increase(http_requests_total[5m])") - .expect("parse failed"); - let result = engine.find_query_config_promql_structural(&ast); - assert!( - result.is_none(), - "Should not match a different function name" - ); - } - - #[test] - fn test_structural_match_spatial_query() { - let engine = create_engine_single_pop( - "http_requests_total", - AggregationType::Sum, - vec!["host"], - vec![( - Some(vec!["host-a".to_string()]), - Box::new(SumAccumulator::with_sum(100.0)), - )], - "sum(http_requests_total) by (host)", - ); - - let ast = promql_parser::parser::parse("sum(http_requests_total) by (host)") - .expect("parse failed"); - let result = engine.find_query_config_promql_structural(&ast); - assert!( - result.is_some(), - "Expected to find config for sum by (host) query" - ); - } -} diff --git a/asap-query-engine/src/tests/datafusion/warm_engine_replay_regression_tests.rs b/asap-query-engine/src/tests/datafusion/warm_engine_replay_regression_tests.rs deleted file mode 100644 index 5e20dab0..00000000 --- a/asap-query-engine/src/tests/datafusion/warm_engine_replay_regression_tests.rs +++ /dev/null @@ -1,861 +0,0 @@ -//! Regression tests for `replay.jsonl` warm-tier failures observed -//! during the MVP demo (PR #333 → ProjectASAP/ASAPCollector#46). -//! -//! Two query shapes returned `status=error` from -//! [`crate::engines::simple::engine::SimpleEngine::handle_query_promql`]: -//! -//! 1. `quantile_over_time(0.99, [1m])` — 500/500 -//! failures. The DDSketch sketch state lands in the warm tier -//! via the agent's `gorillas3 + ddsketch` pipeline, and -//! [`crate::precompute_operators::dd_sketch_accumulator::DDSketchAccumulator`] -//! already supports `Statistic::Quantile`. The query should -//! hit that code path, not error. -//! 2. `sum by (zone) (http_requests_total)` instant — 500/1000 -//! failures (the matching `sum by (zone) (rate(...[5m]))` -//! succeeded). This is a vanilla `OnlySpatial` aggregation -//! over a counter; the warm-tier engine has matchers for it. -//! -//! These tests pin both query shapes against the precise wire format -//! the MVP demo replays so a future regression in pattern dispatch -//! surfaces here as a unit-test failure rather than a 25-minute -//! demo run. -//! -//! Test fixtures use windowed inserts — `(timestamp - window_ms, -//! timestamp)` rather than the zero-duration `(timestamp, timestamp)` -//! pair from `create_engine_single_pop` — so the store's -//! `[query_start, query_end)` overlap filter selects them. The -//! demo's live ingest path always emits windowed data, so this -//! matches production semantics. - -#[cfg(test)] -mod tests { - use crate::data_model::{ - AggregationConfig, AggregationReference, AggregationType, CleanupPolicy, InferenceConfig, - KeyByLabelValues, PrecomputedOutput, PromQLSchema, QueryConfig, QueryLanguage, - SchemaConfig, StreamingConfig, WindowType, - }; - use crate::engines::simple::engine::SimpleEngine; - use crate::engines::QueryResult; - use crate::data_model::Measurement; - use crate::precompute_operators::count_min_sketch_accumulator::CountMinSketchAccumulator; - use crate::precompute_operators::count_sketch_accumulator::CountSketchAccumulator; - use crate::precompute_operators::hll_sketch_accumulator::HllSketchAccumulator; - use crate::precompute_operators::increase_accumulator::IncreaseAccumulator; - use crate::precompute_operators::sum_accumulator::SumAccumulator; - use crate::precompute_operators::DDSketchAccumulator; - use crate::stores::sketch_db::simple_map_store::SimpleMapStore; - use crate::stores::Store; - use crate::AggregateCore; - use asap_sketchlib::sketches::ddsketch::DdSketch; - use asap_sketchlib::sketches::{CountMinSketch, CountSketch, HllSketch}; - use promql_utilities::data_model::KeyByLabelNames; - use std::collections::HashMap; - use std::sync::Arc; - - /// `replay.jsonl` rendered the query at this absolute time; - /// matches the existing temporal tests' convention so the - /// store's `[query_start, query_end)` window aligns with the - /// data the helper seeds. - const QUERY_TIME_SEC: f64 = 1000.0; - /// Match the existing factories' window-end timestamp - /// (`1_000_000` ms = 1000 s wall clock). - const WINDOW_END_MS: u64 = 1_000_000; - /// 60s pane (= one full `[1m]` range), so even an instant - /// query whose effective range is 1s lands inside the window. - const WINDOW_LEN_MS: u64 = 60_000; - - /// Scrape interval (seconds) — needs to be ≥ 1 so the - /// `OnlySpatial` instant-query range - /// `[end - scrape*1000, end)` is non-empty. - const SCRAPE_INTERVAL_S: u64 = 1; - - /// Best-effort tracing init — installs a subscriber the first - /// time it is called so the `warn!` log lines emitted from the - /// engine surface in `--nocapture` output. A failure here means - /// a subscriber is already registered (other test in the same - /// run beat us to it), which is fine. - fn init_tracing_for_test() { - let _ = tracing_subscriber::fmt() - .with_env_filter( - tracing_subscriber::EnvFilter::try_from_default_env() - .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("warn")), - ) - .with_test_writer() - .try_init(); - } - - /// Construct a `SimpleEngine` with a single agg config seeded - /// with one or more accumulators across `grouping_labels`. - /// Inserts each accumulator under window - /// `(WINDOW_END_MS - WINDOW_LEN_MS, WINDOW_END_MS)` so the - /// store's overlap filter accepts it for instant + range - /// queries at `QUERY_TIME_SEC`. - fn build_engine_with_window( - metric: &str, - agg_type: AggregationType, - grouping_labels: Vec<&str>, - data: Vec<(Option>, Box)>, - promql_query: &str, - ) -> SimpleEngine { - let label_strings: Vec = grouping_labels.iter().map(|s| s.to_string()).collect(); - - let mut aggregation_configs = HashMap::new(); - let agg_config = AggregationConfig { - aggregation_id: 1, - aggregation_type: agg_type, - aggregation_sub_type: String::new(), - parameters: HashMap::new(), - grouping_labels: KeyByLabelNames::new(label_strings.clone()), - aggregated_labels: KeyByLabelNames::empty(), - rollup_labels: KeyByLabelNames::empty(), - original_yaml: String::new(), - // 60s tumbling window so a 1m range query sees one - // full pane and the spatial query's narrow range - // overlaps it too. - window_size: WINDOW_LEN_MS / 1000, - slide_interval: WINDOW_LEN_MS / 1000, - window_type: WindowType::Tumbling, - spatial_filter: String::new(), - spatial_filter_normalized: String::new(), - metric: metric.to_string(), - num_aggregates_to_retain: None, - read_count_threshold: None, - table_name: None, - value_column: None, - }; - aggregation_configs.insert(1u64, agg_config); - - let streaming_config = Arc::new(StreamingConfig { - aggregation_configs, - storage_backend: Default::default(), - }); - - let store = Arc::new(SimpleMapStore::new( - streaming_config.clone(), - CleanupPolicy::NoCleanup, - )); - - for (label_values_opt, acc) in data { - let key = label_values_opt.map(|labels| KeyByLabelValues { labels }); - let output = PrecomputedOutput::new( - WINDOW_END_MS - WINDOW_LEN_MS, - WINDOW_END_MS, - key, - 1, - ); - store.insert_precomputed_output(output, acc).unwrap(); - } - - let promql_schema = PromQLSchema::new() - .add_metric(metric.to_string(), KeyByLabelNames::new(label_strings)); - - let query_config = QueryConfig::new(promql_query.to_string()) - .add_aggregation(AggregationReference::new(1, None)); - - let inference_config = InferenceConfig { - schema: SchemaConfig::PromQL(promql_schema), - query_configs: vec![query_config], - cleanup_policy: CleanupPolicy::NoCleanup, - }; - - SimpleEngine::new( - store, - inference_config, - streaming_config, - SCRAPE_INTERVAL_S, - QueryLanguage::promql, - ) - } - - // ------------------------------------------------------------------ - // (1) quantile_over_time over a DDSketch-resident metric - // ------------------------------------------------------------------ - - /// Build a `DDSketchAccumulator` populated with values - /// `[1.0, 2.0, .., 100.0]` so the 0.99 quantile lands near 99.0. - fn dd_sketch_with_1_to_100() -> DDSketchAccumulator { - let mut inner = DdSketch::new(0.01); - for i in 1..=100u32 { - inner.update(i as f64); - } - DDSketchAccumulator { inner } - } - - #[test] - fn quantile_over_time_against_ddsketch_does_not_error() { - init_tracing_for_test(); - let acc = dd_sketch_with_1_to_100(); - let query = "quantile_over_time(0.99, http_requests_total_latency_ms[1m])"; - let engine = build_engine_with_window( - "http_requests_total_latency_ms", - AggregationType::DDSketch, - vec!["zone"], - vec![(Some(vec!["us-east-1".to_string()]), Box::new(acc))], - query, - ); - - let result = engine.handle_query_promql(query.to_string(), QUERY_TIME_SEC); - let (_labels, qr) = - result.expect("warm engine must answer quantile_over_time over DDSketch"); - - match qr { - QueryResult::Vector(iv) => { - assert!( - !iv.values.is_empty(), - "quantile_over_time should produce at least one element" - ); - let v = iv.values[0].value; - // DDSketch with α=0.01 → 1% relative error; 0.99 - // quantile of 1..=100 is 99 (or one of the - // neighbours). Allow generous slack so the test is - // not flaky on bucket-boundary effects. - assert!( - (v - 99.0).abs() < 5.0, - "expected ~99.0 from DDSketch.quantile(0.99), got {v}" - ); - } - other => panic!("expected instant vector, got {other:?}"), - } - } - - // ------------------------------------------------------------------ - // (2) Instant `sum by (zone) (http_requests_total)` - // ------------------------------------------------------------------ - - #[test] - fn sum_by_zone_instant_does_not_error() { - init_tracing_for_test(); - let query = "sum by (zone) (http_requests_total)"; - let engine = build_engine_with_window( - "http_requests_total", - AggregationType::Sum, - vec!["zone"], - vec![ - ( - Some(vec!["us-east-1".to_string()]), - Box::new(SumAccumulator::with_sum(100.0)), - ), - ( - Some(vec!["us-west-2".to_string()]), - Box::new(SumAccumulator::with_sum(50.0)), - ), - ], - query, - ); - - let result = engine.handle_query_promql(query.to_string(), QUERY_TIME_SEC); - let (_labels, qr) = - result.expect("warm engine must answer instant `sum by (zone) (counter)`"); - - match qr { - QueryResult::Vector(iv) => { - assert_eq!( - iv.values.len(), - 2, - "expected 2 zones, got {} values", - iv.values.len() - ); - let mut by_zone = std::collections::HashMap::new(); - for el in &iv.values { - // Per `sum by (zone)` the only output label is - // `zone`; element labels are positional, so we - // pull the first. - let zone = el - .labels - .labels - .first() - .cloned() - .unwrap_or_else(|| "".to_string()); - by_zone.insert(zone, el.value); - } - assert!( - (by_zone.get("us-east-1").copied().unwrap_or(f64::NAN) - 100.0).abs() < 1e-9, - "us-east-1 should be 100.0, by_zone={by_zone:?}" - ); - assert!( - (by_zone.get("us-west-2").copied().unwrap_or(f64::NAN) - 50.0).abs() < 1e-9, - "us-west-2 should be 50.0, by_zone={by_zone:?}" - ); - } - other => panic!("expected instant vector, got {other:?}"), - } - } - - // ------------------------------------------------------------------ - // (3) Cross-check on the temporal-of-quantile alternate spelling - // `quantile_over_time(0.5, ...)` — pins that the fix is not - // hard-coded to 0.99 and that the parameter is correctly - // extracted out of `function_args.first()`. - // ------------------------------------------------------------------ - - #[test] - fn quantile_over_time_p50_against_ddsketch_does_not_error() { - init_tracing_for_test(); - let acc = dd_sketch_with_1_to_100(); - let query = "quantile_over_time(0.5, http_requests_total_latency_ms[1m])"; - let engine = build_engine_with_window( - "http_requests_total_latency_ms", - AggregationType::DDSketch, - vec!["zone"], - vec![(Some(vec!["us-east-1".to_string()]), Box::new(acc))], - query, - ); - - let result = engine.handle_query_promql(query.to_string(), QUERY_TIME_SEC); - let (_labels, qr) = result.expect("p50 quantile_over_time should succeed"); - match qr { - QueryResult::Vector(iv) => { - assert!(!iv.values.is_empty()); - let v = iv.values[0].value; - // Median of 1..=100 is 50 or 51; α=0.01 relative - // error → ~±0.5; allow generous slack. - assert!( - (v - 50.5).abs() < 5.0, - "expected ~50.5 from DDSketch.quantile(0.5), got {v}" - ); - } - other => panic!("expected instant vector, got {other:?}"), - } - } - - // ------------------------------------------------------------------ - // (4) Instant `sum by (zone) (counter)` backed by IncreaseAccumulator - // - // This is the actual demo failure shape from - // ProjectASAP/ASAPCollector#46: the warm-tier ingest path - // stores OTel-`Sum`/monotonic counters as - // `IncreaseAccumulator`, not `SumAccumulator`. Pre-fix, this - // query class capability-missed because `IncreaseAccumulator` - // did not implement `Statistic::Sum`. Post-fix: - // - // a) `compatible_agg_types(Statistic::Sum)` lists - // `Increase` / `MultipleIncrease`, so capability matching - // accepts the counter-shaped configs. - // b) `IncreaseAccumulator::query(Sum, ..)` returns the - // latest cumulative value of the series, matching - // Prometheus' `sum()` instant semantics. - // c) The engine's outer `sum by (zone) (...)` aggregation - // groups + sums those per-series totals across keys. - // ------------------------------------------------------------------ - - #[test] - fn sum_by_zone_instant_over_increase_accumulator_does_not_error() { - init_tracing_for_test(); - let query = "sum by (zone) (http_requests_total)"; - - // Two zones, two cumulative-counter series. Each - // IncreaseAccumulator's `last_seen_measurement` is the latest - // cumulative value the series has reported. - let east = IncreaseAccumulator::new( - Measurement::new(10.0), - (WINDOW_END_MS - WINDOW_LEN_MS) as i64, - Measurement::new(123.0), - WINDOW_END_MS as i64, - ); - let west = IncreaseAccumulator::new( - Measurement::new(0.0), - (WINDOW_END_MS - WINDOW_LEN_MS) as i64, - Measurement::new(45.0), - WINDOW_END_MS as i64, - ); - - let engine = build_engine_with_window( - "http_requests_total", - AggregationType::Increase, - vec!["zone"], - vec![ - (Some(vec!["us-east-1".to_string()]), Box::new(east)), - (Some(vec!["us-west-2".to_string()]), Box::new(west)), - ], - query, - ); - - let result = engine.handle_query_promql(query.to_string(), QUERY_TIME_SEC); - let (_labels, qr) = result.expect( - "warm engine must answer instant `sum by (zone) (counter)` against IncreaseAccumulator", - ); - - match qr { - QueryResult::Vector(iv) => { - assert_eq!( - iv.values.len(), - 2, - "expected 2 zones, got {} values", - iv.values.len() - ); - let mut by_zone = std::collections::HashMap::new(); - for el in &iv.values { - let zone = el - .labels - .labels - .first() - .cloned() - .unwrap_or_else(|| "".to_string()); - by_zone.insert(zone, el.value); - } - // Per-zone Sum is the latest cumulative value of that - // series (Prometheus semantics for sum()). - assert!( - (by_zone.get("us-east-1").copied().unwrap_or(f64::NAN) - 123.0).abs() < 1e-9, - "us-east-1 should be 123.0 (latest cumulative), by_zone={by_zone:?}" - ); - assert!( - (by_zone.get("us-west-2").copied().unwrap_or(f64::NAN) - 45.0).abs() < 1e-9, - "us-west-2 should be 45.0 (latest cumulative), by_zone={by_zone:?}" - ); - } - other => panic!("expected instant vector, got {other:?}"), - } - } - - // ------------------------------------------------------------------ - // (4) DDSketch INGEST-side `_quantile` rename — the production bug - // pinned by ProjectASAP/ASAPCollector#46. The agent renames - // `http_latency_ms` → `http_latency_ms_quantile` before warm-tier - // emit, but the replay client queries with the un-suffixed - // conceptual name. This test pins that the warm engine resolves - // the alias and answers the quantile rather than returning - // `status=error`. - // ------------------------------------------------------------------ - - #[test] - fn quantile_over_time_resolves_unsuffixed_metric_to_quantile_state() { - init_tracing_for_test(); - let acc = dd_sketch_with_1_to_100(); - - // Engine + store know the sketched-form name only. The - // streaming config's agg has `metric = - // "http_latency_ms_quantile"`, mirroring what the - // controller emits after the DDSketch processor's INGEST - // rename. - let engine = build_engine_with_window( - "http_latency_ms_quantile", - AggregationType::DDSketch, - vec!["zone"], - vec![(Some(vec!["us-east-1".to_string()]), Box::new(acc))], - // QueryConfig template carries the suffixed name too — - // matches what the controller would emit alongside the - // streaming config. The engine is expected to resolve - // the un-suffixed form to this template via the alias - // rewrite. - "quantile_over_time(0.99, http_latency_ms_quantile[1m])", - ); - - // Replay client queries with the CONCEPTUAL un-suffixed - // name — this is the exact failure case from - // ProjectASAP/ASAPCollector#46. - let query = "quantile_over_time(0.99, http_latency_ms[1m])"; - let result = engine.handle_query_promql(query.to_string(), QUERY_TIME_SEC); - let (_labels, qr) = result.expect( - "warm engine must resolve un-suffixed `http_latency_ms` to \ - `http_latency_ms_quantile` and answer the quantile", - ); - - match qr { - QueryResult::Vector(iv) => { - assert!( - !iv.values.is_empty(), - "alias-resolved quantile_over_time should produce a value" - ); - let v = iv.values[0].value; - assert!( - (v - 99.0).abs() < 5.0, - "expected ~99.0 from DDSketch.quantile(0.99), got {v}" - ); - } - other => panic!("expected instant vector, got {other:?}"), - } - } - - /// Cross-check: a non-quantile-shaped query for a metric whose - /// `_quantile` variant happens to exist must NOT be rewritten — - /// the alias resolver is shape-gated. - #[test] - fn non_quantile_query_does_not_rewrite_metric() { - init_tracing_for_test(); - // Seed only the suffixed form so a successful rewrite - // would erroneously route the `sum(...)` query at the - // DDSketch state. The engine should leave the query - // untouched, look up the un-suffixed name, find no agg, - // and return None — but critically NOT panic / mis-route - // through the alias. - let acc = dd_sketch_with_1_to_100(); - let engine = build_engine_with_window( - "lookup_latency_ms_quantile", - AggregationType::DDSketch, - vec!["zone"], - vec![(Some(vec!["us-east-1".to_string()]), Box::new(acc))], - "quantile_over_time(0.99, lookup_latency_ms_quantile[1m])", - ); - - // sum_over_time is shape `Sum`, not `Quantile`, so the - // alias resolver must leave the metric name alone. The - // engine has no agg for `lookup_latency_ms` (no `_quantile` - // suffix in its configs), so the result is `None` rather - // than a quantile masquerading as a sum. - let query = "sum_over_time(lookup_latency_ms[1m])"; - let result = engine.handle_query_promql(query.to_string(), QUERY_TIME_SEC); - assert!( - result.is_none(), - "non-quantile shape must NOT trigger the _quantile alias rewrite; \ - got result={result:?}" - ); - } - - // ------------------------------------------------------------------ - // (5) **Production-conditions** suite — schema-empty deploy. - // - // The deployed warm-tier backend (`base.yml`'s - // `--streaming-config=/etc/asap/streaming.yaml` + no - // `--config=…`) starts with `inference_config.schema` set to an - // empty `PromQLSchema` and no `query_configs`. The streaming - // config DOES carry agg configs, but they declare a non-empty - // `grouping_labels` (e.g. `[zone]`) — derived from the - // controller's planner output. The bug: capability matching's - // `labels_compatible` is strict-exact, and with an empty schema - // the engine builds `req.grouping_labels = []`, which fails to - // match any agg config's `[zone]`. Every replay query lands on - // `format_unsupported_query_response` → `status=error`, - // exactly the failure mode `replay.jsonl` shows for the MVP - // demo. - // - // `build_engine_production_conditions` mirrors that exact - // deploy shape so the regression suite pins both the resolver - // fix AND the labels-superset fix. - // ------------------------------------------------------------------ - - /// Build a `SimpleEngine` with the **production warm-tier deploy - /// shape** — the one ASAPCollector's `base.yml` produces: - /// - /// * `streaming_config` carries one agg config with a non-empty - /// `grouping_labels` (e.g. `[zone]`), keyed by the metric the - /// agent's processor emits (suffixed `_quantile` for DDSketch - /// metrics, plain name for HLL/CountSketch/CountMinSketch). - /// * `inference_config.schema = PromQLSchema::new()` (empty) — - /// the warm-tier binary is launched with `--streaming-config` - /// only, no `--config`. - /// * `inference_config.query_configs = []` — no exact-string - /// QueryConfig templates. - /// - /// Replay queries reach `find_compatible_aggregation` via the - /// capability-match fallback path. Pre-fix this fails because - /// `req.grouping_labels = []` can't strict-equal `[zone]`. - #[allow(clippy::too_many_arguments)] - fn build_engine_production_conditions( - agg_metric: &str, - agg_type: AggregationType, - agg_grouping_labels: Vec<&str>, - data: Vec<(Option>, Box)>, - ) -> SimpleEngine { - let label_strings: Vec = agg_grouping_labels - .iter() - .map(|s| s.to_string()) - .collect(); - - let mut aggregation_configs = HashMap::new(); - aggregation_configs.insert( - 1u64, - AggregationConfig { - aggregation_id: 1, - aggregation_type: agg_type, - aggregation_sub_type: String::new(), - parameters: HashMap::new(), - grouping_labels: KeyByLabelNames::new(label_strings), - aggregated_labels: KeyByLabelNames::empty(), - rollup_labels: KeyByLabelNames::empty(), - original_yaml: String::new(), - window_size: WINDOW_LEN_MS / 1000, - slide_interval: WINDOW_LEN_MS / 1000, - window_type: WindowType::Tumbling, - spatial_filter: String::new(), - spatial_filter_normalized: String::new(), - metric: agg_metric.to_string(), - num_aggregates_to_retain: None, - read_count_threshold: None, - table_name: None, - value_column: None, - }, - ); - - let streaming_config = Arc::new(StreamingConfig { - aggregation_configs, - storage_backend: Default::default(), - }); - - let store = Arc::new(SimpleMapStore::new( - streaming_config.clone(), - CleanupPolicy::NoCleanup, - )); - - for (label_values_opt, acc) in data { - let key = label_values_opt.map(|labels| KeyByLabelValues { labels }); - let output = PrecomputedOutput::new( - WINDOW_END_MS - WINDOW_LEN_MS, - WINDOW_END_MS, - key, - 1, - ); - store.insert_precomputed_output(output, acc).unwrap(); - } - - // The crucial bit: schema is EMPTY, mirroring the - // `--streaming-config`-only deploy. The pre-fix engine fails - // here because `build_query_requirements_promql` resolves - // `all_labels` to `KeyByLabelNames::empty()`. - let inference_config = InferenceConfig { - schema: SchemaConfig::PromQL(PromQLSchema::new()), - query_configs: vec![], - cleanup_policy: CleanupPolicy::NoCleanup, - }; - - SimpleEngine::new( - store, - inference_config, - streaming_config, - SCRAPE_INTERVAL_S, - QueryLanguage::promql, - ) - } - - /// (5a) `replay.jsonl` 686/686 failing rows: replays the unsuffixed - /// metric name against a DDSketch agg keyed by the suffixed wire - /// name, with the production deploy's empty schema. - #[test] - fn production_conditions_quantile_over_time_does_not_error() { - init_tracing_for_test(); - let acc = dd_sketch_with_1_to_100(); - let engine = build_engine_production_conditions( - // Agent's DDSketch processor renames to `_quantile` before emit. - "http_requests_total_latency_ms_quantile", - AggregationType::DDSketch, - vec!["zone"], - vec![(Some(vec!["us-east-1".to_string()]), Box::new(acc))], - ); - - // Replay client uses the conceptual unsuffixed name. - let query = "quantile_over_time(0.99, http_requests_total_latency_ms[1m])"; - let (_labels, qr) = engine - .handle_query_promql(query.to_string(), QUERY_TIME_SEC) - .expect( - "production warm engine must answer quantile_over_time over DDSketch even when \ - inference_config has an empty PromQLSchema (replay.jsonl 686/686 errors)", - ); - - match qr { - QueryResult::Vector(iv) => { - assert!(!iv.values.is_empty(), "expected at least one quantile value"); - let v = iv.values[0].value; - assert!( - (v - 99.0).abs() < 5.0, - "expected ~99.0 from DDSketch.quantile(0.99), got {v}" - ); - } - other => panic!("expected instant vector, got {other:?}"), - } - } - - /// (5b) `replay.jsonl` 343/343 failing rows: instant - /// `sum by (zone) (http_requests_total)` against an - /// IncreaseAccumulator-backed counter. - #[test] - fn production_conditions_sum_by_zone_instant_does_not_error() { - init_tracing_for_test(); - let east = IncreaseAccumulator::new( - Measurement::new(10.0), - (WINDOW_END_MS - WINDOW_LEN_MS) as i64, - Measurement::new(123.0), - WINDOW_END_MS as i64, - ); - let west = IncreaseAccumulator::new( - Measurement::new(0.0), - (WINDOW_END_MS - WINDOW_LEN_MS) as i64, - Measurement::new(45.0), - WINDOW_END_MS as i64, - ); - - let engine = build_engine_production_conditions( - "http_requests_total", - AggregationType::Increase, - vec!["zone"], - vec![ - (Some(vec!["us-east-1".to_string()]), Box::new(east)), - (Some(vec!["us-west-2".to_string()]), Box::new(west)), - ], - ); - - let query = "sum by (zone) (http_requests_total)"; - let (_labels, qr) = engine - .handle_query_promql(query.to_string(), QUERY_TIME_SEC) - .expect( - "production warm engine must answer instant `sum by (zone) (counter)` against \ - IncreaseAccumulator under empty PromQLSchema", - ); - - match qr { - QueryResult::Vector(iv) => { - assert_eq!(iv.values.len(), 2, "expected 2 zones"); - } - other => panic!("expected instant vector, got {other:?}"), - } - } - - /// (5c) `replay.jsonl` 343/343 failing rows: `count(unique_users_per_min)` - /// against an HLL agg. HLL accumulator answers `Statistic::Count` as a - /// cardinality alias (`hll_sketch_accumulator.rs:220`), but - /// pre-fix `compatible_agg_types(Count)` did not list HLL — capability - /// match misses → engine returns None → `status=error`. - #[test] - fn production_conditions_count_against_hll_does_not_error() { - init_tracing_for_test(); - // HLL with a few "registers set" — actual cardinality value - // is irrelevant; the test only asserts the engine resolves - // the agg and runs the accumulator's query path without - // erroring. - let mut hll = HllSketch::new(asap_sketchlib::sketches::hll::HllVariant::Regular, 8); - for i in 0..1000u32 { - hll.update(i.to_string().as_bytes()); - } - let acc = HllSketchAccumulator { inner: hll }; - - let engine = build_engine_production_conditions( - "unique_users_per_min", - AggregationType::HLL, - vec!["zone"], - vec![(Some(vec!["us-east-1".to_string()]), Box::new(acc))], - ); - - let query = "count(unique_users_per_min)"; - let result = engine.handle_query_promql(query.to_string(), QUERY_TIME_SEC); - assert!( - result.is_some(), - "production warm engine must answer `count()` under empty PromQLSchema; \ - got None → wire status=error", - ); - } - - /// (5d) `replay.jsonl` 342/342 failing rows: `topk(5, top_endpoint_qps)` - /// against a CountSketch agg. `CountSketchAccumulator` answers - /// `Statistic::Topk` (`count_sketch_accumulator.rs:284`), but pre-fix - /// `compatible_agg_types(Topk)` only listed `CountMinSketchWithHeap`; - /// CountSketch wasn't reachable through capability matching. - /// - /// **Follow-up note**: `CountSketch` is classified as - /// `is_multi_population_value_type`, so even after adding it to - /// the Topk compat list the matcher still requires a paired - /// `SetAggregator` / `DeltaSetAggregator` on the same metric. - /// The production deploy doesn't ship one — the right structural - /// fix is for the controller to plan `top_endpoint_qps` as - /// `CountMinSketchWithHeap` (the integrated CMS+heap accumulator - /// that answers `topk` without an external key tracker). PR #344 - /// declares that capability on the controller side; the matching - /// engine-side accumulator wiring is out of scope for the - /// warm-engine-error PR. Marked `#[ignore]` until the controller - /// switches family. - #[test] - #[ignore = "follow-up: standalone CountSketch agg requires a paired SetAggregator under \ - is_multi_population_value_type semantics; the right fix is for the controller \ - to plan top_endpoint_qps as CountMinSketchWithHeap (PR #344)."] - fn production_conditions_topk_against_count_sketch_does_not_error() { - init_tracing_for_test(); - let mut cs = CountSketch::new(4, 4096); - for i in 0..100u32 { - cs.update(&format!("endpoint-{i}"), 1.0); - } - let acc = CountSketchAccumulator { inner: cs }; - - let engine = build_engine_production_conditions( - "top_endpoint_qps", - AggregationType::CountSketch, - vec!["zone"], - vec![(Some(vec!["us-east-1".to_string()]), Box::new(acc))], - ); - - let query = "topk(5, top_endpoint_qps)"; - let result = engine.handle_query_promql(query.to_string(), QUERY_TIME_SEC); - assert!( - result.is_some(), - "production warm engine must answer `topk(K, )` under empty \ - PromQLSchema; got None → wire status=error", - ); - } - - /// (5e) `replay.jsonl` 342/342 failing rows: `rate(endpoint_request_freq[5m])` - /// against a CountMinSketch agg. `CountMinSketchAccumulator` doesn't - /// directly answer `Statistic::Rate`, but the production demo's - /// frequency probe is structurally a per-series count from a CMS; - /// the engine should at minimum resolve the agg and surface a - /// `Some(...)` rather than `status=error`. (The accumulator may - /// fail at the inner `query_statistic(Rate, …)` step today; this - /// test pins that the surface stays answerable so the replay row - /// is non-empty.) - /// - /// Until CMS gets a `Statistic::Rate` answer, the realistic - /// production fallback is `Statistic::Count` — `rate` is the - /// per-second view of the count. We assert the engine resolves - /// the agg via capability matching; the value is allowed to be - /// any finite number. - #[test] - #[ignore = "follow-up: CountMinSketchAccumulator does not yet implement \ - Statistic::Rate; see TODO.md for tracking. The engine SHOULD resolve \ - the agg through Statistic::Count compat list, but capability \ - matching for `rate(...)` requests Statistic::Rate, which today \ - only lists Increase/MultipleIncrease."] - fn production_conditions_rate_against_cms_does_not_error() { - init_tracing_for_test(); - let mut cms = CountMinSketch::new(4, 4096); - for i in 0..100u32 { - cms.update(&format!("endpoint-{i}"), 1.0); - } - let acc = CountMinSketchAccumulator { inner: cms }; - - let engine = build_engine_production_conditions( - "endpoint_request_freq", - AggregationType::CountMinSketch, - vec!["zone"], - vec![(Some(vec!["us-east-1".to_string()]), Box::new(acc))], - ); - - let query = "rate(endpoint_request_freq[5m])"; - let result = engine.handle_query_promql(query.to_string(), QUERY_TIME_SEC); - assert!( - result.is_some(), - "production warm engine must answer `rate([5m])` under empty \ - PromQLSchema; got None → wire status=error", - ); - } - - /// Sanity: when a deployment registers the bare metric name - /// (no DDSketch INGEST rename applied), the alias resolver - /// must leave the query unchanged — both forms might coexist - /// in tests but the bare form should win when present. - #[test] - fn quantile_query_without_ingest_rename_passes_through() { - init_tracing_for_test(); - let acc = dd_sketch_with_1_to_100(); - // Bare metric IS in streaming config — exactly the - // pre-rename case from PR #108's existing tests. - let engine = build_engine_with_window( - "request_latency_ms", - AggregationType::DDSketch, - vec!["zone"], - vec![(Some(vec!["us-east-1".to_string()]), Box::new(acc))], - "quantile_over_time(0.99, request_latency_ms[1m])", - ); - - let query = "quantile_over_time(0.99, request_latency_ms[1m])"; - let (_labels, qr) = engine - .handle_query_promql(query.to_string(), QUERY_TIME_SEC) - .expect("bare-metric quantile_over_time should answer normally"); - match qr { - QueryResult::Vector(iv) => { - assert!(!iv.values.is_empty(), "expected at least one value"); - let v = iv.values[0].value; - assert!( - (v - 99.0).abs() < 5.0, - "expected ~99.0, got {v} (alias resolver must not have rewritten this)" - ); - } - other => panic!("expected instant vector, got {other:?}"), - } - } -} diff --git a/asap-query-engine/src/tests/mod.rs b/asap-query-engine/src/tests/mod.rs index 14b0a30a..73fa1e51 100644 --- a/asap-query-engine/src/tests/mod.rs +++ b/asap-query-engine/src/tests/mod.rs @@ -2,7 +2,6 @@ pub mod accuracy_empirical_validation_tests; pub mod accuracy_in_promql_response_tests; pub mod capability_matching_tests; pub mod capability_miss_http_e2e_tests; -pub mod datafusion; pub mod persist_format_versioning_tests; pub mod persistence_integration_tests; pub mod persistence_perf_tests; 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 1b85f6f3..1b360955 100644 --- a/asap-query-engine/src/tests/test_utilities/engine_factories.rs +++ b/asap-query-engine/src/tests/test_utilities/engine_factories.rs @@ -637,73 +637,3 @@ pub fn create_engine_multi_timestamp_with_window( ) } -/// Execute both old pipeline and new plan-based path, compare results with epsilon tolerance. -pub async fn assert_old_new_match(engine: &SimpleEngine, query: &str, query_time_sec: f64) { - let context = engine - .build_query_execution_context_promql(query.to_string(), query_time_sec) - .expect("Failed to build context"); - - let (old_results, _) = engine - .execute_query_pipeline(&context, false) - .expect("Old pipeline failed"); - - let new_results = engine - .execute_plan(&context) - .await - .expect("New plan path failed"); - - assert_eq!( - old_results.len(), - new_results.len(), - "Result count mismatch: old={}, new={}", - old_results.len(), - new_results.len() - ); - - let old_map: HashMap, f64> = old_results - .iter() - .map(|r| (r.labels.labels.clone(), r.value)) - .collect(); - - let new_map: HashMap, f64> = new_results - .iter() - .map(|r| (r.labels.labels.clone(), r.value)) - .collect(); - - for (key, old_value) in &old_map { - let new_value = new_map - .get(key) - .unwrap_or_else(|| panic!("Key {:?} missing from new results", key)); - assert!( - (old_value - new_value).abs() < 1e-10, - "Value mismatch for key {:?}: old={}, new={}", - key, - old_value, - new_value - ); - } - - for key in new_map.keys() { - assert!( - old_map.contains_key(key), - "Extra key {:?} in new results", - key - ); - } -} - -/// Convenience wrapper to execute via the new plan path. -pub async fn execute_new_plan( - engine: &SimpleEngine, - query: &str, - query_time_sec: f64, -) -> Vec { - let context = engine - .build_query_execution_context_promql(query.to_string(), query_time_sec) - .expect("Failed to build context"); - - engine - .execute_plan(&context) - .await - .expect("execute_plan failed") -} diff --git a/controller/src/algebra/physical.rs b/controller/src/algebra/physical.rs index 01246dc9..8d53cf70 100644 --- a/controller/src/algebra/physical.rs +++ b/controller/src/algebra/physical.rs @@ -156,7 +156,7 @@ pub enum Placement { BackendCollector, /// PromSketch store (ASAPQuery). PromSketchStore, - /// General query engine (ASAPQuery / DataFusion). + /// General query engine (ASAPQuery). QueryEngine, /// Database (ClickHouse, TimescaleDB, etc.). Database, diff --git a/controller/src/language_logical_plan/plan.rs b/controller/src/language_logical_plan/plan.rs index 40dd267a..97d9c86b 100644 --- a/controller/src/language_logical_plan/plan.rs +++ b/controller/src/language_logical_plan/plan.rs @@ -51,7 +51,6 @@ pub enum LanguageLogicalPlan { }, // Future variants (kept commented to surface intent): // Sql { source: String, tree: SqlTree, summary: ... }, - // DataFusion { source: String, tree: DfPlan, summary: ... }, // ElasticDsl { source: String, tree: EsTree, summary: ... }, } diff --git a/controller/src/language_logical_plan/tests.rs b/controller/src/language_logical_plan/tests.rs index 930cfdb2..d3e30311 100644 --- a/controller/src/language_logical_plan/tests.rs +++ b/controller/src/language_logical_plan/tests.rs @@ -63,7 +63,7 @@ fn lower_promql_topk_extracts_groupby() { #[test] fn lower_unsupported_language_errors_cleanly() { // Build a stub LanguageAst variant by hand — we cannot construct - // SQL/DataFusion/ElasticDsl variants because LanguageAst doesn't + // SQL/ElasticDsl variants because LanguageAst doesn't // expose those yet (PromQL is the only variant). So the equivalent // smoke test is: stub backends fail at L1 with `Unimplemented`, and // the type system rules out passing them to L2 lowering. We assert diff --git a/controller/src/query_language/datafusion/mod.rs b/controller/src/query_language/datafusion/mod.rs deleted file mode 100644 index 94dcfc61..00000000 --- a/controller/src/query_language/datafusion/mod.rs +++ /dev/null @@ -1,27 +0,0 @@ -//! DataFusion L1 backend — stub. -//! -//! The asap-fusion deployment model consumes a pre-built DataFusion -//! `LogicalPlan` upstream (its L1 happens in the caller's -//! `SessionContext`); see `controller/docs/design.md` §3 row 1. This -//! stub exists so a future asap-fusion adapter only has to fill in the -//! `parse` body. - -use super::language::{Language, ParseError}; -use super::language_ast::LanguageAst; -use crate::types_v2::QueryLanguage; - -/// DataFusion implementation of the [`Language`] trait. Currently stubbed. -#[derive(Debug, Default, Clone, Copy)] -pub struct DataFusionLanguage; - -impl Language for DataFusionLanguage { - fn id(&self) -> QueryLanguage { - QueryLanguage::DataFusion - } - - fn parse(&self, _source: &str) -> Result { - Err(ParseError::Unimplemented( - "DataFusion parser not implemented in DC deployment mode", - )) - } -} diff --git a/controller/src/query_language/language.rs b/controller/src/query_language/language.rs index 86ac1993..3c324cce 100644 --- a/controller/src/query_language/language.rs +++ b/controller/src/query_language/language.rs @@ -2,7 +2,7 @@ //! //! See `controller/docs/design.md` §6 `core::query_language`. //! -//! Every backend (PromQL, SQL, DataFusion, ElasticDSL) implements +//! Every backend (PromQL, SQL, ElasticDSL) implements //! [`Language`] and returns a [`LanguageAst`] variant tagged with the //! same [`QueryLanguage`] enum the public `QuerySpec` carries. diff --git a/controller/src/query_language/language_ast.rs b/controller/src/query_language/language_ast.rs index 86e4f806..b1ec8e55 100644 --- a/controller/src/query_language/language_ast.rs +++ b/controller/src/query_language/language_ast.rs @@ -17,7 +17,6 @@ pub enum LanguageAst { PromQL(PromQLAst), // Future backends (kept commented to surface intent): // Sql(SqlAst), - // DataFusion(DfAst), // ElasticDsl(EsAst), } diff --git a/controller/src/query_language/mod.rs b/controller/src/query_language/mod.rs index b14a92a3..1dd7e460 100644 --- a/controller/src/query_language/mod.rs +++ b/controller/src/query_language/mod.rs @@ -17,7 +17,6 @@ //! ├── language_ast.rs — LanguageAst sum type //! ├── promql/ — PromQL backend (active; wraps query_parser::promql) //! ├── sql/ — SQL backend (stub; Unimplemented) -//! ├── datafusion/ — DataFusion backend (stub; Unimplemented) //! └── elastic_dsl/ — ElasticDSL backend (stub; Unimplemented) //! ``` //! @@ -33,7 +32,6 @@ pub mod language_ast; pub mod promql; pub mod sql; -pub mod datafusion; pub mod elastic_dsl; pub use language::{Language, ParseError}; @@ -41,7 +39,6 @@ pub use language_ast::LanguageAst; pub use promql::PromQLLanguage; pub use sql::SqlLanguage; -pub use datafusion::DataFusionLanguage; pub use elastic_dsl::ElasticDslLanguage; #[cfg(test)] diff --git a/controller/src/query_language/tests.rs b/controller/src/query_language/tests.rs index 1100e462..7468404c 100644 --- a/controller/src/query_language/tests.rs +++ b/controller/src/query_language/tests.rs @@ -46,13 +46,6 @@ fn sql_language_returns_unimplemented() { assert_eq!(SqlLanguage.id(), QueryLanguage::Sql); } -#[test] -fn datafusion_language_returns_unimplemented() { - let err = DataFusionLanguage.parse("SELECT 1").unwrap_err(); - assert!(matches!(err, ParseError::Unimplemented(_))); - assert_eq!(DataFusionLanguage.id(), QueryLanguage::DataFusion); -} - #[test] fn elastic_dsl_language_returns_unimplemented() { let err = ElasticDslLanguage.parse("{}").unwrap_err(); diff --git a/controller/src/types_v2.rs b/controller/src/types_v2.rs index 64f768f4..f3e29fc0 100644 --- a/controller/src/types_v2.rs +++ b/controller/src/types_v2.rs @@ -36,9 +36,8 @@ use serde::{Deserialize, Serialize}; /// parser the controller dispatches to. /// /// Today the controller only consumes `PromQL` and `Sql` (see -/// `query_parser/{promql,sql}.rs`); `DataFusion` and `ElasticDsl` are -/// reserved for the future asap-fusion + ElasticDSL deployment models -/// described in `design.md` §3. +/// `query_parser/{promql,sql}.rs`); `ElasticDsl` is reserved for the +/// future ElasticDSL deployment model described in `design.md` §3. #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash)] #[serde(rename_all = "snake_case")] pub enum QueryLanguage { @@ -47,9 +46,6 @@ pub enum QueryLanguage { PromQL, /// SQL. Parsed via `sqlparser`. Sql, - /// A pre-built DataFusion `LogicalPlan`. Reserved — asap-fusion's - /// L1 happens upstream in the caller's `SessionContext`. - DataFusion, /// Elasticsearch DSL. Reserved — no L1 parser yet. ElasticDsl, } @@ -148,7 +144,7 @@ pub enum QueryShape { #[serde(rename_all = "snake_case")] pub enum DataShape { /// Bounded relation, fully materialised at plan time. SQL tables, - /// Parquet / CSV files, DataFusion in-process tables. + /// Parquet / CSV files, in-process columnar tables. Batch, /// Append-only stream — events arrive over time, never updated or /// deleted. Metrics, logs, event streams. The common case for the @@ -277,7 +273,6 @@ mod tests { for variant in [ QueryLanguage::PromQL, QueryLanguage::Sql, - QueryLanguage::DataFusion, QueryLanguage::ElasticDsl, ] { let json = serde_json::to_string(&variant).unwrap(); diff --git a/crates/datafusion_summary_library/Cargo.toml b/crates/datafusion_summary_library/Cargo.toml deleted file mode 100644 index f7b8ee0d..00000000 --- a/crates/datafusion_summary_library/Cargo.toml +++ /dev/null @@ -1,13 +0,0 @@ -[package] -name = "datafusion_summary_library" -version.workspace = true -edition.workspace = true -authors = ["ProjectASAP Team"] - -[dependencies] -# All crate-specific — keep as-is -datafusion = "43" -arrow = "53" -hyperloglogplus = "0.4" -async-trait = "0.1" -futures = "0.3" diff --git a/crates/datafusion_summary_library/src/lib.rs b/crates/datafusion_summary_library/src/lib.rs deleted file mode 100644 index 41837855..00000000 --- a/crates/datafusion_summary_library/src/lib.rs +++ /dev/null @@ -1,15 +0,0 @@ -// DataFusion Summary Library -// -// This library provides logical and physical operators for sketch-based -// query optimization in DataFusion. It supports approximate query processing -// using data structures like HyperLogLog for COUNT(DISTINCT) operations. - -pub mod physical; -pub mod sketch_operators; - -pub use physical::{HllSketch, SketchExtensionPlanner, SummaryInferExec, SummaryInsertExec}; -pub use sketch_operators::{ - GroupingStrategy, InferOperation, PrecomputedSummaryRead, SketchMetadata, SketchSpec, - SketchType, SummaryInfer, SummaryInsert, SummaryMerge, SummaryMergeMultiple, SummaryRead, - SummarySubtract, SummaryType, TypedExpr, -}; diff --git a/crates/datafusion_summary_library/src/physical/hll.rs b/crates/datafusion_summary_library/src/physical/hll.rs deleted file mode 100644 index 6fa0dbd8..00000000 --- a/crates/datafusion_summary_library/src/physical/hll.rs +++ /dev/null @@ -1,169 +0,0 @@ -// HyperLogLog wrapper for cardinality estimation. -// -// This provides a simple wrapper around the hyperloglogplus crate for use -// in sketch-based COUNT(DISTINCT) queries. - -use std::collections::hash_map::{DefaultHasher, RandomState}; -use std::hash::{Hash, Hasher}; - -use hyperloglogplus::{HyperLogLog, HyperLogLogPlus}; - -/// Wrapper around HyperLogLog++ for cardinality estimation. -/// -/// Uses precision 14 by default which gives ~0.8% standard error. -#[derive(Clone)] -pub struct HllSketch { - hll: HyperLogLogPlus, -} - -impl HllSketch { - /// Create a new HLL sketch with default precision (14). - pub fn new() -> Self { - Self::with_precision(14) - } - - /// Create a new HLL sketch with specified precision. - /// - /// Precision must be between 4 and 18. Higher precision means - /// more accuracy but more memory usage. - pub fn with_precision(precision: u8) -> Self { - let hll = HyperLogLogPlus::new(precision, RandomState::new()) - .expect("Valid precision range is 4-18"); - Self { hll } - } - - /// Insert a value into the sketch. - /// - /// The value is hashed to u64 before insertion. - pub fn insert(&mut self, value: &T) { - // Hash the value to u64 first, then insert - let mut hasher = DefaultHasher::new(); - value.hash(&mut hasher); - let hash = hasher.finish(); - self.hll.insert(&hash); - } - - /// Insert a byte slice into the sketch. - pub fn insert_bytes(&mut self, value: &[u8]) { - self.insert(&value); - } - - /// Get the estimated cardinality. - pub fn count(&mut self) -> u64 { - self.hll.count().round() as u64 - } - - /// Merge another HLL sketch into this one. - #[allow(dead_code)] - pub fn merge(&mut self, other: &mut Self) { - self.hll - .merge(&other.hll) - .expect("HLL merge should succeed for same precision"); - } - - /// Serialize the sketch to bytes. - /// - /// Format: [precision: u8][count as f64 bytes(8)] - /// This is a hacky serialization that just stores the current count. - pub fn to_bytes(&mut self) -> Vec { - let precision = 14u8; // We always use 14 for now - let count = self.hll.count(); - - // Simple format: [precision(1)][count as f64 bytes(8)] - let mut bytes = Vec::with_capacity(9); - bytes.push(precision); - bytes.extend_from_slice(&count.to_le_bytes()); - bytes - } - - /// Deserialize a sketch from bytes. - /// Note: This only recovers the count, not the full HLL state. - #[allow(dead_code)] - pub fn from_bytes(bytes: &[u8]) -> Option { - if bytes.len() < 9 { - return None; - } - - let _precision = bytes[0]; - let count_bytes: [u8; 8] = bytes[1..9].try_into().ok()?; - let _count = f64::from_le_bytes(count_bytes); - - // Since we can't truly deserialize the HLL state from just the count, - // we create an empty HLL. This is a limitation of the simple format. - Some(Self::new()) - } -} - -impl Default for HllSketch { - fn default() -> Self { - Self::new() - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_hll_basic() { - let mut hll = HllSketch::new(); - - // Insert 1000 unique values - for i in 0..1000 { - hll.insert(&i); - } - - let count = hll.count(); - // HLL has ~0.8% error at precision 14, so allow 5% tolerance - assert!(count > 900, "Count {} should be > 900", count); - assert!(count < 1100, "Count {} should be < 1100", count); - } - - #[test] - fn test_hll_duplicates() { - let mut hll = HllSketch::new(); - - // Insert same value many times - for _ in 0..1000 { - hll.insert(&42); - } - - let count = hll.count(); - assert_eq!(count, 1, "Duplicates should not increase count"); - } - - #[test] - fn test_hll_merge() { - let mut hll1 = HllSketch::new(); - let mut hll2 = HllSketch::new(); - - // Insert different values into each - for i in 0..500 { - hll1.insert(&i); - } - for i in 500..1000 { - hll2.insert(&i); - } - - hll1.merge(&mut hll2); - let count = hll1.count(); - - // Should have ~1000 unique values - assert!(count > 900, "Merged count {} should be > 900", count); - assert!(count < 1100, "Merged count {} should be < 1100", count); - } - - #[test] - fn test_hll_strings() { - let mut hll = HllSketch::new(); - - for i in 0..1000 { - let s = format!("user_{}", i); - hll.insert(&s); - } - - let count = hll.count(); - assert!(count > 900, "String count {} should be > 900", count); - assert!(count < 1100, "String count {} should be < 1100", count); - } -} diff --git a/crates/datafusion_summary_library/src/physical/mod.rs b/crates/datafusion_summary_library/src/physical/mod.rs deleted file mode 100644 index ae432b90..00000000 --- a/crates/datafusion_summary_library/src/physical/mod.rs +++ /dev/null @@ -1,17 +0,0 @@ -// Physical execution module for sketch-based query plans. -// -// This module provides physical execution plan nodes for sketch operators: -// - SummaryInsertExec: Computes sketches from raw data -// - SummaryInferExec: Extracts results from sketches -// -// Currently only HLL (HyperLogLog) sketches are supported for COUNT(DISTINCT). - -mod hll; -mod planner; -mod summary_infer_exec; -mod summary_insert_exec; - -pub use hll::HllSketch; -pub use planner::SketchExtensionPlanner; -pub use summary_infer_exec::SummaryInferExec; -pub use summary_insert_exec::SummaryInsertExec; diff --git a/crates/datafusion_summary_library/src/physical/planner.rs b/crates/datafusion_summary_library/src/physical/planner.rs deleted file mode 100644 index 58239236..00000000 --- a/crates/datafusion_summary_library/src/physical/planner.rs +++ /dev/null @@ -1,205 +0,0 @@ -// ExtensionPlanner for sketch-based logical plan nodes. -// -// This planner converts SummaryInsert and SummaryInfer logical nodes -// into their physical execution plan counterparts. - -use std::sync::Arc; - -use async_trait::async_trait; -use datafusion::common::{DataFusionError, Result as DFResult}; -use datafusion::execution::context::SessionState; -use datafusion::logical_expr::{LogicalPlan, UserDefinedLogicalNode}; -use datafusion::physical_plan::ExecutionPlan; -use datafusion::physical_planner::{ExtensionPlanner, PhysicalPlanner}; - -use crate::sketch_operators::{InferOperation, SketchType, SummaryInfer, SummaryInsert}; - -use super::{SummaryInferExec, SummaryInsertExec}; - -/// ExtensionPlanner that handles SummaryInsert and SummaryInfer logical nodes. -#[derive(Debug, Default)] -pub struct SketchExtensionPlanner; - -impl SketchExtensionPlanner { - pub fn new() -> Self { - Self - } -} - -#[async_trait] -impl ExtensionPlanner for SketchExtensionPlanner { - async fn plan_extension( - &self, - _planner: &dyn PhysicalPlanner, - node: &dyn UserDefinedLogicalNode, - _logical_inputs: &[&LogicalPlan], - physical_inputs: &[Arc], - _session_state: &SessionState, - ) -> DFResult>> { - // Try to downcast to SummaryInsert - if let Some(summary_insert) = node.as_any().downcast_ref::() { - return self.plan_summary_insert(summary_insert, physical_inputs); - } - - // Try to downcast to SummaryInfer - if let Some(summary_infer) = node.as_any().downcast_ref::() { - return self.plan_summary_infer(summary_infer, physical_inputs); - } - - // Unknown node type, let other planners handle it - Ok(None) - } -} - -impl SketchExtensionPlanner { - fn plan_summary_insert( - &self, - node: &SummaryInsert, - physical_inputs: &[Arc], - ) -> DFResult>> { - if physical_inputs.len() != 1 { - return Err(DataFusionError::Internal( - "SummaryInsert expects exactly one input".to_string(), - )); - } - - let input = physical_inputs[0].clone(); - let input_schema = input.schema(); - - // Only support HLL for now - if node.sketches.len() != 1 { - return Err(DataFusionError::NotImplemented( - "SummaryInsert with multiple sketches not yet supported".to_string(), - )); - } - - let sketch_spec = &node.sketches[0]; - if sketch_spec.sketch_type != SketchType::HLL { - return Err(DataFusionError::NotImplemented(format!( - "Sketch type {:?} not yet supported, only HLL is implemented", - sketch_spec.sketch_type - ))); - } - - // Find value column index - let value_col_idx = match &sketch_spec.value_column { - Some(col_name) => input_schema - .fields() - .iter() - .position(|f| f.name() == col_name) - .ok_or_else(|| { - DataFusionError::Plan(format!( - "Value column '{}' not found in input schema", - col_name - )) - })?, - None => { - return Err(DataFusionError::Plan( - "SummaryInsert requires a value column for HLL".to_string(), - )); - } - }; - - // Find group-by column indices - let group_by_indices: Vec = if !node.group_by_exprs.is_empty() { - // Use group_by_exprs: find columns by expression name - node.group_by_exprs - .iter() - .map(|typed_expr| { - // For simple column expressions, extract the column name - let col_name = - if let datafusion::logical_expr::Expr::Column(col) = &typed_expr.expr { - col.name.clone() - } else { - typed_expr.expr.schema_name().to_string() - }; - - input_schema - .fields() - .iter() - .position(|f| f.name() == &col_name) - .ok_or_else(|| { - DataFusionError::Plan(format!( - "Group-by column '{}' not found in input schema", - col_name - )) - }) - }) - .collect::>>()? - } else { - // Use legacy group_by strings - node.group_by - .iter() - .map(|col_name| { - input_schema - .fields() - .iter() - .position(|f| f.name() == col_name) - .ok_or_else(|| { - DataFusionError::Plan(format!( - "Group-by column '{}' not found in input schema", - col_name - )) - }) - }) - .collect::>>()? - }; - - let exec = SummaryInsertExec::new( - input, - value_col_idx, - group_by_indices, - sketch_spec.output_column_name.clone(), - ); - - Ok(Some(Arc::new(exec))) - } - - fn plan_summary_infer( - &self, - node: &SummaryInfer, - physical_inputs: &[Arc], - ) -> DFResult>> { - if physical_inputs.len() != 1 { - return Err(DataFusionError::Internal( - "SummaryInfer expects exactly one input".to_string(), - )); - } - - let input = physical_inputs[0].clone(); - let input_schema = input.schema(); - - // Only support single operation for now - if node.operations.len() != 1 { - return Err(DataFusionError::NotImplemented( - "SummaryInfer with multiple operations not yet supported".to_string(), - )); - } - - let operation = node.operations[0].clone(); - let output_name = node.output_names[0].clone(); - - // Only support CountDistinct for now - if !matches!(operation, InferOperation::CountDistinct) { - return Err(DataFusionError::NotImplemented(format!( - "Infer operation {:?} not yet supported, only CountDistinct is implemented", - operation - ))); - } - - // Find sketch column index (last column with "sketch" in name, or Binary type) - let sketch_col_idx = input_schema - .fields() - .iter() - .rposition(|f| { - f.name().contains("sketch") || f.data_type() == &arrow::datatypes::DataType::Binary - }) - .ok_or_else(|| { - DataFusionError::Plan("No sketch column found in input schema".to_string()) - })?; - - let exec = SummaryInferExec::new(input, sketch_col_idx, operation, output_name); - - Ok(Some(Arc::new(exec))) - } -} diff --git a/crates/datafusion_summary_library/src/physical/summary_infer_exec.rs b/crates/datafusion_summary_library/src/physical/summary_infer_exec.rs deleted file mode 100644 index 8549452f..00000000 --- a/crates/datafusion_summary_library/src/physical/summary_infer_exec.rs +++ /dev/null @@ -1,282 +0,0 @@ -// Physical execution plan for SummaryInfer (sketch querying). -// -// This ExecutionPlan reads sketch data and extracts results. -// Currently only supports CountDistinct operation on HLL sketches. - -use std::any::Any; -use std::fmt; -use std::pin::Pin; -use std::sync::Arc; -use std::task::{Context, Poll}; - -use arrow::array::{Array, ArrayRef, BinaryArray, RecordBatch, UInt64Builder}; -use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; -use datafusion::common::{DataFusionError, Result as DFResult}; -use datafusion::execution::{SendableRecordBatchStream, TaskContext}; -use datafusion::physical_expr::{EquivalenceProperties, Partitioning}; -use datafusion::physical_plan::{ - DisplayAs, DisplayFormatType, ExecutionMode, ExecutionPlan, PlanProperties, -}; -use futures::Stream; - -use crate::sketch_operators::InferOperation; - -/// Physical execution plan for extracting results from HLL sketches. -/// -/// Takes input batches with sketch columns and produces one row per input row with: -/// - Group key columns (passed through) -/// - Result column (e.g., UInt64 for CountDistinct) -#[derive(Debug)] -pub struct SummaryInferExec { - /// Input execution plan (typically SummaryInsertExec) - input: Arc, - - /// Index of the sketch column in input schema - sketch_col_idx: usize, - - /// Infer operation to perform - operation: InferOperation, - - /// Output column name - output_name: String, - - /// Output schema - schema: SchemaRef, - - /// Plan properties (cached) - properties: PlanProperties, -} - -impl SummaryInferExec { - pub fn new( - input: Arc, - sketch_col_idx: usize, - operation: InferOperation, - output_name: String, - ) -> Self { - let input_schema = input.schema(); - - // Build output schema: all columns except sketch column, plus result column - let mut fields: Vec = input_schema - .fields() - .iter() - .enumerate() - .filter(|(idx, _)| *idx != sketch_col_idx) - .map(|(_, f)| f.as_ref().clone()) - .collect(); - - // Add result column based on operation type - let result_type = match &operation { - InferOperation::CountDistinct => DataType::UInt64, - InferOperation::Quantile(_) | InferOperation::Median => DataType::Float64, - _ => DataType::UInt64, // Default for unsupported ops - }; - - fields.push(Field::new(&output_name, result_type, false)); - - let schema = Arc::new(Schema::new(fields)); - - // Plan properties: same partitioning as input - let properties = PlanProperties::new( - EquivalenceProperties::new(schema.clone()), - Partitioning::UnknownPartitioning(1), - ExecutionMode::Bounded, - ); - - Self { - input, - sketch_col_idx, - operation, - output_name, - schema, - properties, - } - } -} - -impl DisplayAs for SummaryInferExec { - fn fmt_as(&self, t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { - match t { - DisplayFormatType::Default | DisplayFormatType::Verbose => { - write!( - f, - "SummaryInferExec: sketch_col={}, op={}, output={}", - self.sketch_col_idx, self.operation, self.output_name - ) - } - } - } -} - -impl ExecutionPlan for SummaryInferExec { - fn name(&self) -> &str { - "SummaryInferExec" - } - - fn as_any(&self) -> &dyn Any { - self - } - - fn schema(&self) -> SchemaRef { - self.schema.clone() - } - - fn properties(&self) -> &PlanProperties { - &self.properties - } - - fn children(&self) -> Vec<&Arc> { - vec![&self.input] - } - - fn with_new_children( - self: Arc, - children: Vec>, - ) -> DFResult> { - if children.len() != 1 { - return Err(DataFusionError::Internal( - "SummaryInferExec expects exactly one child".to_string(), - )); - } - Ok(Arc::new(SummaryInferExec::new( - children[0].clone(), - self.sketch_col_idx, - self.operation.clone(), - self.output_name.clone(), - ))) - } - - fn execute( - &self, - partition: usize, - context: Arc, - ) -> DFResult { - if partition != 0 { - return Err(DataFusionError::Internal(format!( - "SummaryInferExec only supports partition 0, got {}", - partition - ))); - } - - let input_stream = self.input.execute(0, context)?; - - let stream = SummaryInferStream::new( - input_stream, - self.schema.clone(), - self.sketch_col_idx, - self.operation.clone(), - ); - - Ok(Box::pin(stream)) - } -} - -/// Stream that transforms sketch batches into result batches. -struct SummaryInferStream { - /// Input stream - input: SendableRecordBatchStream, - - /// Output schema - schema: SchemaRef, - - /// Sketch column index - sketch_col_idx: usize, - - /// Operation to perform - operation: InferOperation, -} - -impl SummaryInferStream { - fn new( - input: SendableRecordBatchStream, - schema: SchemaRef, - sketch_col_idx: usize, - operation: InferOperation, - ) -> Self { - Self { - input, - schema, - sketch_col_idx, - operation, - } - } - - /// Transform a batch by extracting results from sketches. - fn transform_batch(&self, batch: &RecordBatch) -> DFResult { - let num_rows = batch.num_rows(); - - // Get the sketch column - let sketch_col = batch.column(self.sketch_col_idx); - let sketch_array = sketch_col - .as_any() - .downcast_ref::() - .ok_or_else(|| { - DataFusionError::Internal("Sketch column is not Binary type".to_string()) - })?; - - // Build result column based on operation - let result_col: ArrayRef = match &self.operation { - InferOperation::CountDistinct => { - let mut builder = UInt64Builder::with_capacity(num_rows); - for i in 0..num_rows { - if sketch_array.is_null(i) { - builder.append_null(); - } else { - let sketch_bytes = sketch_array.value(i); - // Deserialize and get count - // Note: Our simple serialization format stores the count directly - let count = if sketch_bytes.len() >= 9 { - let count_bytes: [u8; 8] = sketch_bytes[1..9].try_into().unwrap(); - f64::from_le_bytes(count_bytes).round() as u64 - } else { - 0 - }; - builder.append_value(count); - } - } - Arc::new(builder.finish()) - } - _ => { - return Err(DataFusionError::NotImplemented(format!( - "Infer operation {:?} not yet implemented", - self.operation - ))); - } - }; - - // Build output columns: all input columns except sketch, plus result - let mut columns: Vec = batch - .columns() - .iter() - .enumerate() - .filter(|(idx, _)| *idx != self.sketch_col_idx) - .map(|(_, col)| col.clone()) - .collect(); - columns.push(result_col); - - RecordBatch::try_new(self.schema.clone(), columns) - .map_err(|e| DataFusionError::ArrowError(e, None)) - } -} - -impl Stream for SummaryInferStream { - type Item = DFResult; - - fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - match Pin::new(&mut self.input).poll_next(cx) { - Poll::Ready(Some(Ok(batch))) => { - let result = self.transform_batch(&batch); - Poll::Ready(Some(result)) - } - Poll::Ready(Some(Err(e))) => Poll::Ready(Some(Err(e))), - Poll::Ready(None) => Poll::Ready(None), - Poll::Pending => Poll::Pending, - } - } -} - -impl datafusion::physical_plan::RecordBatchStream for SummaryInferStream { - fn schema(&self) -> SchemaRef { - self.schema.clone() - } -} diff --git a/crates/datafusion_summary_library/src/physical/summary_insert_exec.rs b/crates/datafusion_summary_library/src/physical/summary_insert_exec.rs deleted file mode 100644 index a62d8ba2..00000000 --- a/crates/datafusion_summary_library/src/physical/summary_insert_exec.rs +++ /dev/null @@ -1,434 +0,0 @@ -// Physical execution plan for SummaryInsert (sketch building). -// -// This ExecutionPlan consumes input batches and builds HLL sketches -// for each group. Currently only supports HLL sketches. - -use std::any::Any; -use std::collections::HashMap; -use std::fmt; -use std::sync::Arc; - -use arrow::array::{Array, ArrayRef, BinaryBuilder, RecordBatch, StringBuilder, UInt64Builder}; -use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; -use datafusion::common::{DataFusionError, Result as DFResult, ScalarValue}; -use datafusion::execution::{SendableRecordBatchStream, TaskContext}; -use datafusion::physical_expr::{EquivalenceProperties, Partitioning}; -use datafusion::physical_plan::{ - DisplayAs, DisplayFormatType, ExecutionMode, ExecutionPlan, PlanProperties, -}; - -use super::hll::HllSketch; - -/// Physical execution plan for building HLL sketches. -/// -/// Takes input batches and produces one row per group with: -/// - Group key columns -/// - Binary column containing serialized HLL sketch -#[derive(Debug)] -pub struct SummaryInsertExec { - /// Input execution plan - input: Arc, - - /// Index of the value column to sketch (in input schema) - value_col_idx: usize, - - /// Indices of group-by columns (in input schema) - group_by_indices: Vec, - - /// Name of the output sketch column - sketch_col_name: String, - - /// Output schema - schema: SchemaRef, - - /// Plan properties (cached) - properties: PlanProperties, -} - -impl SummaryInsertExec { - pub fn new( - input: Arc, - value_col_idx: usize, - group_by_indices: Vec, - sketch_col_name: String, - ) -> Self { - let input_schema = input.schema(); - - // Build output schema: group columns + sketch column - let mut fields: Vec = group_by_indices - .iter() - .map(|&idx| input_schema.field(idx).clone()) - .collect(); - - // Add sketch column with the specified name - fields.push(Field::new(&sketch_col_name, DataType::Binary, false)); - - let schema = Arc::new(Schema::new(fields)); - - // Plan properties: single partition output, no ordering guarantees - let properties = PlanProperties::new( - EquivalenceProperties::new(schema.clone()), - Partitioning::UnknownPartitioning(1), - ExecutionMode::Bounded, - ); - - Self { - input, - value_col_idx, - group_by_indices, - sketch_col_name, - schema, - properties, - } - } - - /// Extracts a group key from a row as a vector of ScalarValues. - fn extract_group_key( - batch: &RecordBatch, - row_idx: usize, - group_by_indices: &[usize], - ) -> Vec { - group_by_indices - .iter() - .map(|&col_idx| { - ScalarValue::try_from_array(batch.column(col_idx), row_idx) - .unwrap_or(ScalarValue::Null) - }) - .collect() - } - - /// Extracts a value as bytes for hashing. - fn extract_value_bytes(array: &ArrayRef, row_idx: usize) -> Vec { - // Convert any value to string representation for hashing - // This is a hacky but universal approach - if array.is_null(row_idx) { - return b"__NULL__".to_vec(); - } - - // Use Arrow's display formatter to get string representation - let value = ScalarValue::try_from_array(array.as_ref(), row_idx) - .map(|v| v.to_string()) - .unwrap_or_else(|_| "__ERROR__".to_string()); - - value.into_bytes() - } -} - -impl DisplayAs for SummaryInsertExec { - fn fmt_as(&self, t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { - match t { - DisplayFormatType::Default | DisplayFormatType::Verbose => { - write!( - f, - "SummaryInsertExec: value_col={}, group_by={:?}", - self.value_col_idx, self.group_by_indices - ) - } - } - } -} - -impl ExecutionPlan for SummaryInsertExec { - fn name(&self) -> &str { - "SummaryInsertExec" - } - - fn as_any(&self) -> &dyn Any { - self - } - - fn schema(&self) -> SchemaRef { - self.schema.clone() - } - - fn properties(&self) -> &PlanProperties { - &self.properties - } - - fn children(&self) -> Vec<&Arc> { - vec![&self.input] - } - - fn with_new_children( - self: Arc, - children: Vec>, - ) -> DFResult> { - if children.len() != 1 { - return Err(DataFusionError::Internal( - "SummaryInsertExec expects exactly one child".to_string(), - )); - } - Ok(Arc::new(SummaryInsertExec::new( - children[0].clone(), - self.value_col_idx, - self.group_by_indices.clone(), - self.sketch_col_name.clone(), - ))) - } - - fn execute( - &self, - partition: usize, - context: Arc, - ) -> DFResult { - if partition != 0 { - return Err(DataFusionError::Internal(format!( - "SummaryInsertExec only supports partition 0, got {}", - partition - ))); - } - - // Get input stream - let input_stream = self.input.execute(0, context)?; - - // Create the output stream - let schema = self.schema.clone(); - let value_col_idx = self.value_col_idx; - let group_by_indices = self.group_by_indices.clone(); - - let stream = - SummaryInsertStream::new(input_stream, schema, value_col_idx, group_by_indices); - - Ok(Box::pin(stream)) - } -} - -/// Stream that consumes input batches and produces aggregated sketch results. -struct SummaryInsertStream { - /// Input stream - input: SendableRecordBatchStream, - - /// Output schema - schema: SchemaRef, - - /// Value column index - value_col_idx: usize, - - /// Group-by column indices - group_by_indices: Vec, - - /// Accumulated sketches per group - groups: HashMap, HllSketch>, - - /// Whether we've finished consuming input - finished_input: bool, - - /// Whether we've emitted the final result - emitted_result: bool, -} - -impl SummaryInsertStream { - fn new( - input: SendableRecordBatchStream, - schema: SchemaRef, - value_col_idx: usize, - group_by_indices: Vec, - ) -> Self { - Self { - input, - schema, - value_col_idx, - group_by_indices, - groups: HashMap::new(), - finished_input: false, - emitted_result: false, - } - } - - /// Process a batch of input data. - fn process_batch(&mut self, batch: &RecordBatch) { - let value_array = batch.column(self.value_col_idx); - let num_rows = batch.num_rows(); - - for row_idx in 0..num_rows { - // Extract group key - let group_key = - SummaryInsertExec::extract_group_key(batch, row_idx, &self.group_by_indices); - - // Get or create HLL for this group - let hll = self.groups.entry(group_key).or_default(); - - // Extract value and insert into HLL - let value_bytes = SummaryInsertExec::extract_value_bytes(value_array, row_idx); - hll.insert_bytes(&value_bytes); - } - } - - /// Build the final output batch from accumulated sketches. - fn build_output(&mut self) -> DFResult { - let num_groups = self.groups.len(); - - // Build group key columns - let mut group_builders: Vec = self - .schema - .fields() - .iter() - .take(self.group_by_indices.len()) - .map(|field| ScalarArrayBuilder::new(field.data_type(), num_groups)) - .collect(); - - // Build sketch column - let mut sketch_builder = BinaryBuilder::with_capacity(num_groups, num_groups * 16); - - // Populate arrays - for (group_key, hll) in &mut self.groups { - // Add group key values - for (idx, scalar) in group_key.iter().enumerate() { - group_builders[idx].append(scalar); - } - - // Add serialized sketch - let sketch_bytes = hll.to_bytes(); - sketch_builder.append_value(&sketch_bytes); - } - - // Finish building arrays - let mut columns: Vec = group_builders.iter_mut().map(|b| b.finish()).collect(); - columns.push(Arc::new(sketch_builder.finish())); - - RecordBatch::try_new(self.schema.clone(), columns) - .map_err(|e| DataFusionError::ArrowError(e, None)) - } -} - -use futures::Stream; -use std::pin::Pin; -use std::task::{Context, Poll}; - -impl Stream for SummaryInsertStream { - type Item = DFResult; - - fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - // If we've emitted the result, we're done - if self.emitted_result { - return Poll::Ready(None); - } - - // Consume all input batches first - if !self.finished_input { - loop { - match Pin::new(&mut self.input).poll_next(cx) { - Poll::Ready(Some(Ok(batch))) => { - self.process_batch(&batch); - } - Poll::Ready(Some(Err(e))) => { - return Poll::Ready(Some(Err(e))); - } - Poll::Ready(None) => { - self.finished_input = true; - break; - } - Poll::Pending => { - return Poll::Pending; - } - } - } - } - - // Build and emit the final result - self.emitted_result = true; - - // Handle case with no groups - if self.groups.is_empty() { - return Poll::Ready(None); - } - - let batch = self.build_output(); - Poll::Ready(Some(batch)) - } -} - -impl datafusion::physical_plan::RecordBatchStream for SummaryInsertStream { - fn schema(&self) -> SchemaRef { - self.schema.clone() - } -} - -// Helper enum for building arrays dynamically -enum ScalarArrayBuilder { - Utf8(StringBuilder), - UInt64(UInt64Builder), - Int64(arrow::array::Int64Builder), - Float64(arrow::array::Float64Builder), - Int32(arrow::array::Int32Builder), - UInt32(arrow::array::UInt32Builder), -} - -impl ScalarArrayBuilder { - fn new(data_type: &DataType, capacity: usize) -> Self { - match data_type { - DataType::Utf8 => { - ScalarArrayBuilder::Utf8(StringBuilder::with_capacity(capacity, capacity * 32)) - } - DataType::UInt64 => ScalarArrayBuilder::UInt64(UInt64Builder::with_capacity(capacity)), - DataType::Int64 => { - ScalarArrayBuilder::Int64(arrow::array::Int64Builder::with_capacity(capacity)) - } - DataType::Float64 => { - ScalarArrayBuilder::Float64(arrow::array::Float64Builder::with_capacity(capacity)) - } - DataType::Int32 => { - ScalarArrayBuilder::Int32(arrow::array::Int32Builder::with_capacity(capacity)) - } - DataType::UInt32 => { - ScalarArrayBuilder::UInt32(arrow::array::UInt32Builder::with_capacity(capacity)) - } - // For unsupported types, fall back to string representation - _ => ScalarArrayBuilder::Utf8(StringBuilder::with_capacity(capacity, capacity * 32)), - } - } - - fn append(&mut self, scalar: &ScalarValue) { - match (self, scalar) { - (ScalarArrayBuilder::Utf8(b), ScalarValue::Utf8(v)) => match v { - Some(s) => b.append_value(s), - None => b.append_null(), - }, - (ScalarArrayBuilder::UInt64(b), ScalarValue::UInt64(v)) => match v { - Some(val) => b.append_value(*val), - None => b.append_null(), - }, - (ScalarArrayBuilder::Int64(b), ScalarValue::Int64(v)) => match v { - Some(val) => b.append_value(*val), - None => b.append_null(), - }, - (ScalarArrayBuilder::Float64(b), ScalarValue::Float64(v)) => match v { - Some(val) => b.append_value(*val), - None => b.append_null(), - }, - (ScalarArrayBuilder::Int32(b), ScalarValue::Int32(v)) => match v { - Some(val) => b.append_value(*val), - None => b.append_null(), - }, - (ScalarArrayBuilder::UInt32(b), ScalarValue::UInt32(v)) => match v { - Some(val) => b.append_value(*val), - None => b.append_null(), - }, - // Fallback: convert to string for Utf8 builder - (ScalarArrayBuilder::Utf8(b), scalar) => { - if scalar.is_null() { - b.append_null(); - } else { - b.append_value(scalar.to_string()); - } - } - // For type mismatches with non-Utf8 builders, append null - (ScalarArrayBuilder::UInt64(b), _) => b.append_null(), - (ScalarArrayBuilder::Int64(b), _) => b.append_null(), - (ScalarArrayBuilder::Float64(b), _) => b.append_null(), - (ScalarArrayBuilder::Int32(b), _) => b.append_null(), - (ScalarArrayBuilder::UInt32(b), _) => b.append_null(), - } - } - - fn finish(&mut self) -> ArrayRef { - match self { - ScalarArrayBuilder::Utf8(b) => Arc::new(b.finish()), - ScalarArrayBuilder::UInt64(b) => Arc::new(b.finish()), - ScalarArrayBuilder::Int64(b) => Arc::new(b.finish()), - ScalarArrayBuilder::Float64(b) => Arc::new(b.finish()), - ScalarArrayBuilder::Int32(b) => Arc::new(b.finish()), - ScalarArrayBuilder::UInt32(b) => Arc::new(b.finish()), - } - } -} diff --git a/crates/datafusion_summary_library/src/sketch_operators.rs b/crates/datafusion_summary_library/src/sketch_operators.rs deleted file mode 100644 index 64fe3946..00000000 --- a/crates/datafusion_summary_library/src/sketch_operators.rs +++ /dev/null @@ -1,1630 +0,0 @@ -// Sketch-based query plan operators for DataFusion -// -// This module defines custom logical plan nodes for sketch-based query optimization. -// These operators support exploring different sketch-based execution strategies. -#![allow(deprecated)] - -#[allow(deprecated)] -use datafusion::arrow::datatypes::{DataType, Field}; -use datafusion::common::{DFSchema, DFSchemaRef, Result as DFResult}; -use datafusion::error::DataFusionError; -use datafusion::logical_expr::{Expr, LogicalPlan, UserDefinedLogicalNodeCore}; -use std::cmp::Ordering; -use std::collections::BTreeMap; // BTreeMap instead of HashMap (can derive Hash) -use std::fmt; -use std::hash::{Hash, Hasher}; -use std::sync::Arc; - -// ============================================================================ -// TypedExpr - Expression with pre-resolved type -// ============================================================================ - -/// An expression paired with its pre-resolved data type. -/// -/// This is used to preserve type information when expressions are passed through -/// plan nodes that may not have access to the original schema needed to resolve types. -/// For example, SummaryInfer's input is a SummaryInsert which may not include -/// the columns referenced in GROUP BY expressions (especially for Hydra strategy). -#[derive(Debug, Clone)] -pub struct TypedExpr { - pub expr: Expr, - pub data_type: DataType, -} - -impl TypedExpr { - pub fn new(expr: Expr, data_type: DataType) -> Self { - Self { expr, data_type } - } -} - -// Manual trait implementations since Expr implements these traits -impl PartialEq for TypedExpr { - fn eq(&self, other: &Self) -> bool { - self.expr == other.expr && self.data_type == other.data_type - } -} - -impl Eq for TypedExpr {} - -impl Hash for TypedExpr { - fn hash(&self, state: &mut H) { - self.expr.hash(state); - self.data_type.hash(state); - } -} - -// ============================================================================ -// Sketch Types -// ============================================================================ - -/// Types of sketches/summaries supported for query processing -/// Also aliased as SummaryType for clarity (includes both sketches and exact aggregators) -#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)] -pub enum SketchType { - // ======================================================================== - // Exact aggregators (non-sketch, single population) - // ======================================================================== - Sum, // Exact sum accumulator - Increase, // Counter increase tracking - MinMax, // Min/max values - - MultipleSum, - MultipleIncrease, - MultipleMinMax, - - // ======================================================================== - // Set aggregators - // ======================================================================== - SetAggregator, // Exact set of group keys (HashSet-based) - DeltaSetAggregator, // Set aggregation with separate key tracking - - // ======================================================================== - // COUNT DISTINCT sketches - // ======================================================================== - HLL, // HyperLogLog - UltraLogLog, // UltraLogLog (improved HLL) - HydraHLL, // HyperLogLog with multi-population support - - // ======================================================================== - // Quantile sketches - // ======================================================================== - KLL, // KLL sketch - TDigest, // T-Digest - HydraKLL, // KLL with multi-population support - - // ======================================================================== - // Heavy hitters / TOP K - // ======================================================================== - SpaceSaving, // Space-Saving algorithm - FrequentItems, // Frequent items sketch - - // ======================================================================== - // Frequency estimation - // ======================================================================== - CountMinSketch, // Count-Min Sketch - CountSketch, // Count Sketch - - // ======================================================================== - // General purpose - // ======================================================================== - Sampling, // Reservoir sampling -} - -/// Type alias for clarity - SummaryType includes both sketches and exact aggregators -pub type SummaryType = SketchType; - -impl fmt::Display for SketchType { - fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { - match self { - // Exact aggregators - SketchType::Sum => write!(f, "Sum"), - SketchType::Increase => write!(f, "Increase"), - SketchType::MinMax => write!(f, "MinMax"), - SketchType::SetAggregator => write!(f, "SetAggregator"), - SketchType::DeltaSetAggregator => write!(f, "DeltaSetAggregator"), - SketchType::MultipleSum => write!(f, "MultipleSum"), - SketchType::MultipleIncrease => write!(f, "MultipleIncrease"), - SketchType::MultipleMinMax => write!(f, "MultipleMinMax"), - // Sketches - SketchType::HLL => write!(f, "HLL"), - SketchType::UltraLogLog => write!(f, "UltraLogLog"), - SketchType::HydraHLL => write!(f, "HydraHLL"), - SketchType::KLL => write!(f, "KLL"), - SketchType::TDigest => write!(f, "TDigest"), - SketchType::HydraKLL => write!(f, "HydraKLL"), - SketchType::SpaceSaving => write!(f, "SpaceSaving"), - SketchType::FrequentItems => write!(f, "FrequentItems"), - SketchType::CountMinSketch => write!(f, "CountMinSketch"), - SketchType::CountSketch => write!(f, "CountSketch"), - SketchType::Sampling => write!(f, "Sampling"), - } - } -} - -impl SketchType { - /// Check if this sketch type supports multi-population (Hydra-style) - pub fn is_hydra(&self) -> bool { - matches!(self, SketchType::HydraHLL | SketchType::HydraKLL) - } - - /// Get the base sketch type (non-Hydra version) - pub fn base_type(&self) -> SketchType { - match self { - SketchType::HydraHLL => SketchType::HLL, - SketchType::HydraKLL => SketchType::KLL, - other => other.clone(), - } - } -} - -// ============================================================================ -// Inference Operations -// ============================================================================ - -/// Operations that can be performed on sketches/summaries to extract results -/// Note: Uses simplified types (strings instead of Expr) for DataFusion integration -#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)] -pub enum InferOperation { - // ======================================================================== - // Exact aggregator extraction operations - // ======================================================================== - /// Extract sum value from Sum/HydraSum accumulator - ExtractSum, - - /// Extract count value from an accumulator - ExtractCount, - - /// Extract minimum value from MinMax/HydraMinMax accumulator - ExtractMin, - - /// Extract maximum value from MinMax/HydraMinMax accumulator - ExtractMax, - - /// Extract increase value from Increase/HydraIncrease accumulator - ExtractIncrease, - - /// Extract rate (increase / time_range) from Increase accumulator - ExtractRate, - - // ======================================================================== - // Sketch operations - // ======================================================================== - /// COUNT(DISTINCT column) - CountDistinct, - - /// Quantile/percentile estimation - /// Stores quantile as integer (0-10000) for 4 decimal places: 0.9500 = 9500 - Quantile(u16), - - /// Median (equivalent to Quantile(0.5)) - Median, - - /// TOP K items - TopK(usize), - - /// Frequency-based COUNT(*) aggregation with GROUP BY - /// Queries frequency sketch to get count for each group key - FrequencyCount, - - /// Frequency-based SUM(column) aggregation with GROUP BY - /// Queries frequency sketch to get sum for each group key - FrequencySum, - - /// Frequency-based AVG(column) aggregation with GROUP BY - /// Computed as SUM(column) / COUNT(*) using two frequency sketches - FrequencyAvg, - - /// Frequency estimate for a specific value (stored as string) - FrequencyEstimate(String), - - /// Get all frequent items above threshold (stored as integer, 0-10000) - FrequentItems(u16), - - /// Enumerate set contents (for SetAggregator) - /// Returns all unique values seen in the set - EnumerateSet, -} - -impl InferOperation { - /// Create a Quantile operation from a float (0.0 to 1.0) - pub fn quantile(p: f64) -> Self { - InferOperation::Quantile((p * 10000.0).round() as u16) - } - - /// Get the quantile value as f64 - pub fn quantile_value(&self) -> Option { - match self { - InferOperation::Quantile(p) => Some(*p as f64 / 10000.0), - InferOperation::Median => Some(0.5), - _ => None, - } - } - - /// Create a FrequentItems operation from a float threshold - pub fn frequent_items(threshold: f64) -> Self { - InferOperation::FrequentItems((threshold * 10000.0).round() as u16) - } - - /// Get the threshold value as f64 - pub fn threshold_value(&self) -> Option { - match self { - InferOperation::FrequentItems(t) => Some(*t as f64 / 10000.0), - _ => None, - } - } -} - -impl fmt::Display for InferOperation { - fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { - match self { - // Exact aggregator extractions - InferOperation::ExtractSum => write!(f, "EXTRACT_SUM"), - InferOperation::ExtractCount => write!(f, "EXTRACT_COUNT"), - InferOperation::ExtractMin => write!(f, "EXTRACT_MIN"), - InferOperation::ExtractMax => write!(f, "EXTRACT_MAX"), - InferOperation::ExtractIncrease => write!(f, "EXTRACT_INCREASE"), - InferOperation::ExtractRate => write!(f, "EXTRACT_RATE"), - // Sketch operations - InferOperation::CountDistinct => write!(f, "COUNT_DISTINCT"), - InferOperation::Quantile(p) => write!(f, "QUANTILE({:.4})", *p as f64 / 10000.0), - InferOperation::Median => write!(f, "MEDIAN"), - InferOperation::TopK(k) => write!(f, "TOPK({})", k), - InferOperation::FrequencyCount => write!(f, "FREQ_COUNT"), - InferOperation::FrequencySum => write!(f, "FREQ_SUM"), - InferOperation::FrequencyAvg => write!(f, "FREQ_AVG"), - InferOperation::FrequencyEstimate(value) => write!(f, "FREQ_EST({})", value), - InferOperation::FrequentItems(threshold) => { - write!(f, "FREQ_ITEMS({:.4})", *threshold as f64 / 10000.0) - } - InferOperation::EnumerateSet => write!(f, "ENUM_SET"), - } - } -} - -// ============================================================================ -// Grouping Strategy -// ============================================================================ - -/// Strategy for handling GROUP BY queries -#[derive(Debug, Clone, PartialEq, Eq)] -pub enum GroupingStrategy { - /// One sketch per group (filter-based, computed separately) - PerGroup, - - /// Single Hydra-style sketch containing all groups - Hydra, - - /// No grouping (simple aggregation) - None, -} - -// ============================================================================ -// Sketch Metadata -// ============================================================================ - -/// Metadata for identifying and loading sketches -#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)] -pub struct SketchMetadata { - pub table_name: String, - pub column_name: String, - pub sketch_type: SketchType, - pub filter_predicate: Option, - pub key_columns: Vec, // For Hydra sketches -} - -impl fmt::Display for SketchMetadata { - fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { - write!( - f, - "{}.{}.{}", - self.table_name, self.column_name, self.sketch_type - )?; - if let Some(filter) = &self.filter_predicate { - write!(f, " WHERE {}", filter)?; - } - if !self.key_columns.is_empty() { - write!(f, " KEY BY [{}]", self.key_columns.join(", "))?; - } - Ok(()) - } -} - -// ============================================================================ -// Sketch Specification -// ============================================================================ - -/// Specification for a single sketch to create -#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)] -pub struct SketchSpec { - pub value_column: Option, - pub sketch_type: SketchType, - pub output_column_name: String, // e.g., "host_sketch", "cpu_sketch" -} - -// ============================================================================ -// SummaryInsert - Compute sketch from raw data -// ============================================================================ - -/// Logical plan node: Compute a sketch from raw data -#[derive(Debug, Clone, PartialEq, Eq, Hash)] -pub struct SummaryInsert { - /// Input data source - pub input: Arc, - - /// Sketches to create (one or more) - pub sketches: Vec, - - /// GROUP BY columns for per-group strategy (columns appear in output schema) - /// Legacy field - use group_by_exprs for computed expressions - pub group_by: Vec, - - /// Key columns for Hydra strategy (columns embedded in sketch, NOT in output schema) - /// Legacy field - use key_column_exprs for computed expressions - pub key_columns: Vec, - - /// GROUP BY expressions with pre-resolved types for per-group strategy - /// When non-empty, takes precedence over group_by in compute_schema - pub group_by_exprs: Vec, - - /// Key column expressions with pre-resolved types for Hydra strategy - /// These are embedded in the sketch, not in output schema - pub key_column_exprs: Vec, - - /// Optional parameters (e.g., HLL precision, KLL k value) - /// Using BTreeMap instead of HashMap so it can derive Hash - pub parameters: BTreeMap, - - /// Cached output schema - schema: DFSchemaRef, -} - -impl SummaryInsert { - /// Create a new SummaryInsert with multiple sketches - pub fn new(input: Arc, sketches: Vec) -> DFResult { - if sketches.is_empty() { - return Err(DataFusionError::Plan( - "SummaryInsert requires at least one sketch".to_string(), - )); - } - - let schema = Self::compute_schema(&input, &sketches, &[], &[], &[])?; - Ok(Self { - input, - sketches, - group_by: vec![], - key_columns: vec![], - group_by_exprs: vec![], - key_column_exprs: vec![], - parameters: BTreeMap::new(), - schema, - }) - } - - /// Helper constructor for single sketch (backward compatibility) - pub fn single( - input: Arc, - value_column: Option, - sketch_type: SketchType, - ) -> DFResult { - let output_column_name = match &value_column { - Some(col) => format!("{}_sketch", col), - None => "value_sketch".to_string(), - }; - - Self::new( - input, - vec![SketchSpec { - value_column, - sketch_type, - output_column_name, - }], - ) - } - - pub fn with_group_by(mut self, group_by: Vec) -> DFResult { - self.schema = Self::compute_schema( - &self.input, - &self.sketches, - &group_by, - &self.key_columns, - &self.group_by_exprs, - )?; - self.group_by = group_by; - Ok(self) - } - - pub fn with_key_columns(mut self, key_columns: Vec) -> DFResult { - self.schema = Self::compute_schema( - &self.input, - &self.sketches, - &self.group_by, - &key_columns, - &self.group_by_exprs, - )?; - self.key_columns = key_columns; - Ok(self) - } - - /// Set GROUP BY expressions with pre-resolved types (supports computed expressions) - pub fn with_group_by_exprs(mut self, group_by_exprs: Vec) -> DFResult { - self.schema = Self::compute_schema( - &self.input, - &self.sketches, - &self.group_by, - &self.key_columns, - &group_by_exprs, - )?; - self.group_by_exprs = group_by_exprs; - Ok(self) - } - - /// Set key column expressions with pre-resolved types for Hydra strategy - /// Note: These are embedded in the sketch, not in the output schema - pub fn with_key_column_exprs(mut self, key_column_exprs: Vec) -> DFResult { - // key_column_exprs don't affect output schema, but store them for later use - self.key_column_exprs = key_column_exprs; - Ok(self) - } - - pub fn with_parameters(mut self, parameters: BTreeMap) -> Self { - self.parameters = parameters; - self - } - - /// Compute output schema based on grouping strategy - fn compute_schema( - input: &Arc, - sketches: &[SketchSpec], - group_by: &[String], - _key_columns: &[String], - group_by_exprs: &[TypedExpr], - ) -> DFResult { - let input_schema = input.schema(); - let mut qualified_fields = Vec::new(); - - // For per-group strategy: include group columns in output with their qualifications - // This matches vanilla DataFusion Aggregate behavior - // - // Prefer group_by_exprs (TypedExpr) over group_by (String) if available - if !group_by_exprs.is_empty() { - // Use TypedExpr - supports computed expressions like date_part(), CASE, etc. - for typed_expr in group_by_exprs { - // For simple columns, use col.name as field name and col.relation as qualifier - // For computed expressions, use schema_name() with no qualifier - let (qualifier, field_name) = if let Expr::Column(col) = &typed_expr.expr { - (col.relation.clone(), col.name.clone()) - } else { - (None, typed_expr.expr.schema_name().to_string()) - }; - qualified_fields.push(( - qualifier, - Arc::new(Field::new(&field_name, typed_expr.data_type.clone(), true)), - )); - } - } else if !group_by.is_empty() { - // Fallback to legacy string-based group_by - for col_name in group_by { - // Get both qualifier and field to preserve qualification - let (qualifier, field) = input_schema - .qualified_field_with_unqualified_name(col_name) - .map_err(|e| { - DataFusionError::Plan(format!( - "Group column '{}' not found in input schema: {}", - col_name, e - )) - })?; - qualified_fields.push((qualifier.cloned(), Arc::new(field.clone()))); - } - } - - // For Hydra strategy: key columns are embedded in sketch, not in output - // (no fields added here - neither key_columns nor key_column_exprs affect output) - - // Add sketch columns (Binary type, unqualified) - // Create one column per sketch specification - for sketch_spec in sketches { - qualified_fields.push(( - None, - Arc::new(Field::new( - &sketch_spec.output_column_name, - DataType::Binary, - false, - )), - )); - } - - // Create DFSchema from qualified fields - let schema = DFSchema::new_with_metadata(qualified_fields, Default::default()) - .map_err(|e| DataFusionError::Plan(format!("Failed to create schema: {}", e)))?; - - Ok(Arc::new(schema)) - } -} - -impl PartialOrd for SummaryInsert { - fn partial_cmp(&self, other: &Self) -> Option { - // Compare by sketches, then grouping, then parameters, then input - match self.sketches.partial_cmp(&other.sketches) { - Some(Ordering::Equal) => {} - other => return other, - } - match self.group_by.partial_cmp(&other.group_by) { - Some(Ordering::Equal) => {} - other => return other, - } - match self.key_columns.partial_cmp(&other.key_columns) { - Some(Ordering::Equal) => {} - other => return other, - } - match self.parameters.partial_cmp(&other.parameters) { - Some(Ordering::Equal) => {} - other => return other, - } - match self.input.partial_cmp(&other.input) { - Some(Ordering::Equal) => {} - other => return other, - } - // Compare schemas by pointer (Arc comparison) - Some(Arc::as_ptr(&self.schema).cmp(&Arc::as_ptr(&other.schema))) - } -} - -impl UserDefinedLogicalNodeCore for SummaryInsert { - fn name(&self) -> &str { - "SummaryInsert" - } - - fn inputs(&self) -> Vec<&LogicalPlan> { - vec![self.input.as_ref()] - } - - fn schema(&self) -> &DFSchemaRef { - &self.schema - } - - fn expressions(&self) -> Vec { - vec![] - } - - fn fmt_for_explain(&self, f: &mut fmt::Formatter) -> fmt::Result { - if self.sketches.len() == 1 { - // Single sketch: show simplified format - let sketch = &self.sketches[0]; - write!(f, "SummaryInsert: sketch_type={}", sketch.sketch_type)?; - if let Some(col) = &sketch.value_column { - write!(f, ", value_column={}", col)?; - } - } else { - // Multiple sketches: show as array - write!(f, "SummaryInsert: sketches=[")?; - for (i, sketch) in self.sketches.iter().enumerate() { - if i > 0 { - write!(f, ", ")?; - } - write!(f, "{{type={}", sketch.sketch_type)?; - if let Some(col) = &sketch.value_column { - write!(f, ", column={}", col)?; - } - write!(f, "}}")?; - } - write!(f, "]")?; - } - if !self.group_by.is_empty() { - write!(f, ", group_by=[{}]", self.group_by.join(", "))?; - } - if !self.key_columns.is_empty() { - write!(f, ", key_columns=[{}]", self.key_columns.join(", "))?; - } - Ok(()) - } - - fn from_template(&self, _exprs: &[Expr], inputs: &[LogicalPlan]) -> Self { - let input = Arc::new(inputs[0].clone()); - // Recompute schema with new input - let schema = Self::compute_schema( - &input, - &self.sketches, - &self.group_by, - &self.key_columns, - &self.group_by_exprs, - ) - .unwrap_or_else(|_| self.schema.clone()); - - Self { - input, - sketches: self.sketches.clone(), - group_by: self.group_by.clone(), - key_columns: self.key_columns.clone(), - group_by_exprs: self.group_by_exprs.clone(), - key_column_exprs: self.key_column_exprs.clone(), - parameters: self.parameters.clone(), - schema, - } - } - - fn with_exprs_and_inputs(&self, _exprs: Vec, inputs: Vec) -> DFResult { - Ok(self.from_template(&_exprs, &inputs)) - } -} - -// ============================================================================ -// SummaryRead - Load pre-computed sketch -// ============================================================================ - -/// Logical plan node: Load a pre-computed sketch -#[derive(Debug, Clone, PartialEq, Eq, Hash)] -pub struct SummaryRead { - /// Metadata to identify the sketch - pub metadata: SketchMetadata, - - /// Optional: Direct sketch ID if known - pub sketch_id: Option, - - /// Schema (placeholder for now) - schema: DFSchemaRef, -} - -impl SummaryRead { - pub fn new(metadata: SketchMetadata, schema: DFSchemaRef) -> Self { - Self { - metadata, - sketch_id: None, - schema, - } - } - - pub fn with_sketch_id(mut self, sketch_id: String) -> Self { - self.sketch_id = Some(sketch_id); - self - } -} - -impl PartialOrd for SummaryRead { - fn partial_cmp(&self, other: &Self) -> Option { - match self.metadata.partial_cmp(&other.metadata) { - Some(Ordering::Equal) => {} - other => return other, - } - match self.sketch_id.partial_cmp(&other.sketch_id) { - Some(Ordering::Equal) => {} - other => return other, - } - // DFSchemaRef is Arc, and DFSchema likely doesn't implement PartialOrd - // So we compare by pointer - Some(Arc::as_ptr(&self.schema).cmp(&Arc::as_ptr(&other.schema))) - } -} - -impl UserDefinedLogicalNodeCore for SummaryRead { - fn name(&self) -> &str { - "SummaryRead" - } - - fn inputs(&self) -> Vec<&LogicalPlan> { - vec![] // No inputs - reads from storage - } - - fn schema(&self) -> &DFSchemaRef { - &self.schema - } - - fn expressions(&self) -> Vec { - vec![] - } - - fn fmt_for_explain(&self, f: &mut fmt::Formatter) -> fmt::Result { - write!(f, "SummaryRead: {}", self.metadata)?; - if let Some(id) = &self.sketch_id { - write!(f, " [id={}]", id)?; - } - Ok(()) - } - - fn from_template(&self, _exprs: &[Expr], _inputs: &[LogicalPlan]) -> Self { - self.clone() - } - - fn with_exprs_and_inputs(&self, _exprs: Vec, inputs: Vec) -> DFResult { - Ok(self.from_template(&_exprs, &inputs)) - } -} - -// ============================================================================ -// SummaryMerge - Merge multiple sketches -// ============================================================================ - -/// Logical plan node: Merge two or more sketches -#[derive(Debug, Clone, PartialEq, Eq, Hash)] -pub struct SummaryMerge { - /// Left sketch source - pub left: Arc, - - /// Right sketch source - pub right: Arc, - - /// Sketch type (for validation) - pub sketch_type: SketchType, -} - -impl SummaryMerge { - pub fn new(left: Arc, right: Arc, sketch_type: SketchType) -> Self { - Self { - left, - right, - sketch_type, - } - } -} - -impl PartialOrd for SummaryMerge { - fn partial_cmp(&self, other: &Self) -> Option { - match self.sketch_type.partial_cmp(&other.sketch_type) { - Some(Ordering::Equal) => {} - other => return other, - } - match self.left.partial_cmp(&other.left) { - Some(Ordering::Equal) => {} - other => return other, - } - self.right.partial_cmp(&other.right) - } -} - -impl UserDefinedLogicalNodeCore for SummaryMerge { - fn name(&self) -> &str { - "SummaryMerge" - } - - fn inputs(&self) -> Vec<&LogicalPlan> { - vec![self.left.as_ref(), self.right.as_ref()] - } - - fn schema(&self) -> &DFSchemaRef { - // Return left schema (should be compatible with right) - self.left.schema() - } - - fn expressions(&self) -> Vec { - vec![] - } - - fn fmt_for_explain(&self, f: &mut fmt::Formatter) -> fmt::Result { - write!(f, "SummaryMerge: sketch_type={}", self.sketch_type) - } - - fn from_template(&self, _exprs: &[Expr], inputs: &[LogicalPlan]) -> Self { - Self { - left: Arc::new(inputs[0].clone()), - right: Arc::new(inputs[1].clone()), - sketch_type: self.sketch_type.clone(), - } - } - - fn with_exprs_and_inputs(&self, _exprs: Vec, inputs: Vec) -> DFResult { - Ok(self.from_template(&_exprs, &inputs)) - } -} - -// ============================================================================ -// SummarySubtract - Subtract one sketch from another -// ============================================================================ - -/// Logical plan node: Subtract one sketch from another (for sliding windows) -#[derive(Debug, Clone, PartialEq, Eq, Hash)] -pub struct SummarySubtract { - /// Sketch to subtract FROM (minuend) - pub minuend: Arc, - - /// Sketch to subtract (subtrahend) - pub subtrahend: Arc, - - /// Sketch type (for validation) - pub sketch_type: SketchType, -} - -impl SummarySubtract { - pub fn new( - minuend: Arc, - subtrahend: Arc, - sketch_type: SketchType, - ) -> Self { - Self { - minuend, - subtrahend, - sketch_type, - } - } -} - -impl PartialOrd for SummarySubtract { - fn partial_cmp(&self, other: &Self) -> Option { - match self.sketch_type.partial_cmp(&other.sketch_type) { - Some(Ordering::Equal) => {} - other => return other, - } - match self.minuend.partial_cmp(&other.minuend) { - Some(Ordering::Equal) => {} - other => return other, - } - self.subtrahend.partial_cmp(&other.subtrahend) - } -} - -impl UserDefinedLogicalNodeCore for SummarySubtract { - fn name(&self) -> &str { - "SummarySubtract" - } - - fn inputs(&self) -> Vec<&LogicalPlan> { - vec![self.minuend.as_ref(), self.subtrahend.as_ref()] - } - - fn schema(&self) -> &DFSchemaRef { - self.minuend.schema() - } - - fn expressions(&self) -> Vec { - vec![] - } - - fn fmt_for_explain(&self, f: &mut fmt::Formatter) -> fmt::Result { - write!(f, "SummarySubtract: sketch_type={}", self.sketch_type) - } - - fn from_template(&self, _exprs: &[Expr], inputs: &[LogicalPlan]) -> Self { - Self { - minuend: Arc::new(inputs[0].clone()), - subtrahend: Arc::new(inputs[1].clone()), - sketch_type: self.sketch_type.clone(), - } - } - - fn with_exprs_and_inputs(&self, _exprs: Vec, inputs: Vec) -> DFResult { - Ok(self.from_template(&_exprs, &inputs)) - } -} - -// ============================================================================ -// SummaryInfer - Extract result from sketch -// ============================================================================ - -/// Logical plan node: Extract a result from a sketch -#[derive(Debug, Clone)] -pub struct SummaryInfer { - /// Input sketch source - pub input: Arc, - - /// Optional second input for keys enumeration (multi-population accumulators). - /// When present, SummaryInferExec deserializes the value sketch once per spatial group - /// and queries it N times (once per sub-key from the keys input). - pub keys_input: Option>, - - /// Operations to perform on the sketch(es) - /// For single sketch with multiple operations: operations map to sketch in order - pub operations: Vec, - - /// Output column names (one per operation) - pub output_names: Vec, - - /// Optional group key columns for Hydra sketches (column names, not full Expr) - /// Legacy field - use group_key_exprs for computed expressions - pub group_key_columns: Vec, - - /// Optional qualifier for the group key columns (for Hydra sketches) - pub group_key_qualifier: Option, - - /// Group key expressions with pre-resolved types (supports computed expressions) - /// When non-empty, takes precedence over group_key_columns in compute_schema - pub group_key_exprs: Vec, - - /// Cached output schema - schema: DFSchemaRef, -} - -impl SummaryInfer { - /// Create a new SummaryInfer with multiple operations - pub fn new( - input: Arc, - operations: Vec, - output_names: Vec, - ) -> DFResult { - // Validate inputs - if operations.is_empty() { - return Err(DataFusionError::Plan( - "SummaryInfer requires at least one operation".to_string(), - )); - } - if operations.len() != output_names.len() { - return Err(DataFusionError::Plan(format!( - "SummaryInfer operations ({}) and output_names ({}) length mismatch", - operations.len(), - output_names.len() - ))); - } - - let schema = Self::compute_schema(&input, &operations, &output_names, &[], &None, &[])?; - Ok(Self { - input, - keys_input: None, - operations, - output_names, - group_key_columns: vec![], - group_key_qualifier: None, - group_key_exprs: vec![], - schema, - }) - } - - /// Helper constructor for single operation (backward compatibility) - pub fn single( - input: Arc, - operation: InferOperation, - output_name: String, - ) -> DFResult { - Self::new(input, vec![operation], vec![output_name]) - } - - /// Add group key columns for Hydra sketches (supports multiple columns) - /// Legacy method - use with_group_key_exprs for computed expressions - pub fn with_group_key_columns( - mut self, - group_key_columns: Vec, - qualifier: Option, - ) -> DFResult { - self.schema = Self::compute_schema( - &self.input, - &self.operations, - &self.output_names, - &group_key_columns, - &qualifier, - &self.group_key_exprs, - )?; - self.group_key_columns = group_key_columns; - self.group_key_qualifier = qualifier; - Ok(self) - } - - /// Set a second input for keys enumeration (multi-population accumulators). - pub fn with_keys_input(mut self, keys_input: Arc) -> Self { - self.keys_input = Some(keys_input); - self - } - - /// Add group key expressions with pre-resolved types (supports computed expressions) - pub fn with_group_key_exprs( - mut self, - group_key_exprs: Vec, - qualifier: Option, - ) -> DFResult { - self.schema = Self::compute_schema( - &self.input, - &self.operations, - &self.output_names, - &self.group_key_columns, - &qualifier, - &group_key_exprs, - )?; - self.group_key_exprs = group_key_exprs; - self.group_key_qualifier = qualifier; - Ok(self) - } - - /// Compute output schema based on operations and grouping - fn compute_schema( - input: &Arc, - operations: &[InferOperation], - output_names: &[String], - group_key_columns: &[String], - group_key_qualifier: &Option, - group_key_exprs: &[TypedExpr], - ) -> DFResult { - let input_schema = input.schema(); - let mut qualified_fields = Vec::new(); - - // Add group columns to output with qualifications preserved - // Prefer group_key_exprs (TypedExpr) over group_key_columns (String) if available - if !group_key_exprs.is_empty() { - // Use TypedExpr - supports computed expressions like date_part(), CASE, etc. - // First: pass through input label columns not covered by group_key_exprs - let expr_names: Vec = group_key_exprs - .iter() - .filter_map(|te| { - if let Expr::Column(col) = &te.expr { - Some(col.name.clone()) - } else { - None - } - }) - .collect(); - for (qualifier, field) in input_schema.iter() { - if field.name() != "sketch" - && !field.name().ends_with("_sketch") - && !expr_names.contains(&field.name().to_string()) - { - qualified_fields.push((qualifier.cloned(), field.clone())); - } - } - // Then: add group key expression columns - for typed_expr in group_key_exprs { - // For simple columns, use col.name as field name and col.relation as qualifier - // For computed expressions, use schema_name() with optional provided qualifier - let (qualifier, field_name) = if let Expr::Column(col) = &typed_expr.expr { - (col.relation.clone(), col.name.clone()) - } else if let Some(qual) = group_key_qualifier { - ( - Some(datafusion::common::TableReference::bare(qual.clone())), - typed_expr.expr.schema_name().to_string(), - ) - } else { - (None, typed_expr.expr.schema_name().to_string()) - }; - qualified_fields.push(( - qualifier, - Arc::new(Field::new(&field_name, typed_expr.data_type.clone(), true)), - )); - } - } else if !group_key_columns.is_empty() { - // Fallback to legacy string-based group_key_columns - // Hydra/self-keyed case: pass through input label columns first, - // then add materialized group key columns from the accumulator. - for (qualifier, field) in input_schema.iter() { - if field.name() != "sketch" - && !field.name().ends_with("_sketch") - && !group_key_columns.contains(&field.name().to_string()) - { - qualified_fields.push((qualifier.cloned(), field.clone())); - } - } - for key_col in group_key_columns { - // Try to find it in the input schema first - if let Ok((qualifier, field)) = - input_schema.qualified_field_with_unqualified_name(key_col) - { - qualified_fields.push((qualifier.cloned(), Arc::new(field.clone()))); - } else if let Some(qual) = group_key_qualifier { - // Use provided qualifier if input schema doesn't have it - // This happens for Hydra where SummaryInsert doesn't output the group keys - let qualifier = Some(datafusion::common::TableReference::bare(qual.clone())); - qualified_fields.push(( - qualifier, - Arc::new(Field::new(key_col, DataType::Utf8, false)), - )); - } else { - // No qualifier available - use unqualified - qualified_fields - .push((None, Arc::new(Field::new(key_col, DataType::Utf8, false)))); - } - } - } else { - // Per-group case: preserve group columns from input (non-sketch columns) with qualifications - for (qualifier, field) in input_schema.iter() { - if field.name() != "sketch" && !field.name().ends_with("_sketch") { - qualified_fields.push((qualifier.cloned(), field.clone())); - } - } - } - - // Add result columns based on operation types (unqualified) - for (operation, output_name) in operations.iter().zip(output_names.iter()) { - let result_type = match operation { - // Exact aggregator extractions - all return Float64 - InferOperation::ExtractSum => DataType::Float64, - InferOperation::ExtractCount => DataType::Float64, - InferOperation::ExtractMin => DataType::Float64, - InferOperation::ExtractMax => DataType::Float64, - InferOperation::ExtractIncrease => DataType::Float64, - InferOperation::ExtractRate => DataType::Float64, - // Sketch operations - InferOperation::CountDistinct => DataType::UInt64, - InferOperation::Quantile(_) | InferOperation::Median => DataType::Float64, - InferOperation::TopK(_) => { - DataType::List(Arc::new(Field::new("item", DataType::Utf8, true))) - } - InferOperation::FrequencyCount => DataType::Float64, // COUNT returns numeric - InferOperation::FrequencySum => DataType::Float64, // SUM returns numeric - InferOperation::FrequencyAvg => DataType::Float64, // AVG returns numeric - InferOperation::FrequencyEstimate(_) => DataType::UInt64, - InferOperation::FrequentItems(_) => { - DataType::List(Arc::new(Field::new("item", DataType::Utf8, true))) - } - InferOperation::EnumerateSet => { - DataType::List(Arc::new(Field::new("item", DataType::Utf8, true))) - } - }; - - qualified_fields.push((None, Arc::new(Field::new(output_name, result_type, false)))); - } - - // Create DFSchema from qualified fields - let schema = DFSchema::new_with_metadata(qualified_fields, Default::default()) - .map_err(|e| DataFusionError::Plan(format!("Failed to create schema: {}", e)))?; - - Ok(Arc::new(schema)) - } -} - -impl PartialEq for SummaryInfer { - fn eq(&self, other: &Self) -> bool { - self.input == other.input - && self.keys_input == other.keys_input - && self.operations == other.operations - && self.output_names == other.output_names - && self.group_key_columns == other.group_key_columns - && self.group_key_qualifier == other.group_key_qualifier - && self.schema == other.schema - } -} - -impl Eq for SummaryInfer {} - -impl std::hash::Hash for SummaryInfer { - fn hash(&self, state: &mut H) { - self.input.hash(state); - self.keys_input.hash(state); - self.operations.hash(state); - self.output_names.hash(state); - self.group_key_columns.hash(state); - self.group_key_qualifier.hash(state); - self.schema.hash(state); - } -} - -impl PartialOrd for SummaryInfer { - fn partial_cmp(&self, other: &Self) -> Option { - match self.operations.partial_cmp(&other.operations) { - Some(Ordering::Equal) => {} - other => return other, - } - match self.output_names.partial_cmp(&other.output_names) { - Some(Ordering::Equal) => {} - other => return other, - } - match self.group_key_columns.partial_cmp(&other.group_key_columns) { - Some(Ordering::Equal) => {} - other => return other, - } - match self.input.partial_cmp(&other.input) { - Some(Ordering::Equal) => {} - other => return other, - } - match self.keys_input.partial_cmp(&other.keys_input) { - Some(Ordering::Equal) => {} - other => return other, - } - // Compare schemas by pointer - Some(Arc::as_ptr(&self.schema).cmp(&Arc::as_ptr(&other.schema))) - } -} - -impl UserDefinedLogicalNodeCore for SummaryInfer { - fn name(&self) -> &str { - "SummaryInfer" - } - - fn inputs(&self) -> Vec<&LogicalPlan> { - let mut inputs = vec![self.input.as_ref()]; - if let Some(ref keys_input) = self.keys_input { - inputs.push(keys_input.as_ref()); - } - inputs - } - - fn schema(&self) -> &DFSchemaRef { - &self.schema - } - - fn expressions(&self) -> Vec { - // No Expr types stored anymore - group_key_column is just a string - vec![] - } - - fn fmt_for_explain(&self, f: &mut fmt::Formatter) -> fmt::Result { - if self.operations.len() == 1 { - write!( - f, - "SummaryInfer: operation={}, output={}", - self.operations[0], self.output_names[0] - )?; - } else { - write!(f, "SummaryInfer: operations=[")?; - for (i, op) in self.operations.iter().enumerate() { - if i > 0 { - write!(f, ", ")?; - } - write!(f, "{}", op)?; - } - write!(f, "], outputs=[")?; - for (i, name) in self.output_names.iter().enumerate() { - if i > 0 { - write!(f, ", ")?; - } - write!(f, "{}", name)?; - } - write!(f, "]")?; - } - if !self.group_key_columns.is_empty() { - write!( - f, - ", group_key_columns=[{}]", - self.group_key_columns.join(", ") - )?; - } - if self.keys_input.is_some() { - write!(f, ", has_keys_input=true")?; - } - Ok(()) - } - - fn from_template(&self, _exprs: &[Expr], inputs: &[LogicalPlan]) -> Self { - let input = Arc::new(inputs[0].clone()); - let keys_input = if inputs.len() > 1 { - Some(Arc::new(inputs[1].clone())) - } else { - self.keys_input.clone() - }; - // Recompute schema with new input - let schema = Self::compute_schema( - &input, - &self.operations, - &self.output_names, - &self.group_key_columns, - &self.group_key_qualifier, - &self.group_key_exprs, - ) - .unwrap_or_else(|_| self.schema.clone()); - - Self { - input, - keys_input, - operations: self.operations.clone(), - output_names: self.output_names.clone(), - group_key_columns: self.group_key_columns.clone(), - group_key_qualifier: self.group_key_qualifier.clone(), - group_key_exprs: self.group_key_exprs.clone(), - schema, - } - } - - fn with_exprs_and_inputs(&self, exprs: Vec, inputs: Vec) -> DFResult { - Ok(self.from_template(&exprs, &inputs)) - } -} - -// ============================================================================ -// PrecomputedSummaryRead - Read precomputed summaries from store -// ============================================================================ - -/// Logical plan node: Read precomputed summaries from a store -/// -/// This is a leaf node that represents reading precomputed aggregates -/// (summaries) from a PrecomputedOutputStore. Used for OnlySpatial queries -/// where data has already been aggregated by a streaming engine. -#[derive(Debug, Clone, PartialEq, Eq, Hash)] -pub struct PrecomputedSummaryRead { - /// Metric name being queried - metric: String, - - /// Aggregation ID to query - aggregation_id: u64, - - /// Start timestamp of the query range - start_timestamp: u64, - - /// End timestamp of the query range - end_timestamp: u64, - - /// Whether this is an exact query (sliding window) vs approximate (tumbling) - is_exact_query: bool, - - /// Output label names (group by columns) - output_labels: Vec, - - /// Type of summary being read - summary_type: SketchType, - - /// Cached output schema - schema: DFSchemaRef, -} - -impl PrecomputedSummaryRead { - #[allow(clippy::too_many_arguments)] - pub fn new( - metric: String, - aggregation_id: u64, - start_timestamp: u64, - end_timestamp: u64, - is_exact_query: bool, - output_labels: Vec, - summary_type: SketchType, - schema: DFSchemaRef, - ) -> Self { - Self { - metric, - aggregation_id, - start_timestamp, - end_timestamp, - is_exact_query, - output_labels, - summary_type, - schema, - } - } - - /// Create with auto-generated schema based on output_labels - pub fn with_auto_schema( - metric: String, - aggregation_id: u64, - start_timestamp: u64, - end_timestamp: u64, - is_exact_query: bool, - output_labels: Vec, - summary_type: SketchType, - ) -> DFResult { - let schema = Self::compute_schema(&output_labels)?; - Ok(Self::new( - metric, - aggregation_id, - start_timestamp, - end_timestamp, - is_exact_query, - output_labels, - summary_type, - schema, - )) - } - - /// Compute schema: [label columns (Utf8), sketch (Binary)] - fn compute_schema(output_labels: &[String]) -> DFResult { - let mut qualified_fields = Vec::new(); - - // Add label columns (all Utf8, nullable) - for label in output_labels { - qualified_fields.push((None, Arc::new(Field::new(label, DataType::Utf8, true)))); - } - - // Add sketch column (Binary, not nullable) - qualified_fields.push(( - None, - Arc::new(Field::new("sketch", DataType::Binary, false)), - )); - - let schema = DFSchema::new_with_metadata(qualified_fields, Default::default()) - .map_err(|e| DataFusionError::Plan(format!("Failed to create schema: {}", e)))?; - - Ok(Arc::new(schema)) - } - - // Getters - pub fn metric(&self) -> &str { - &self.metric - } - - pub fn aggregation_id(&self) -> u64 { - self.aggregation_id - } - - pub fn start_timestamp(&self) -> u64 { - self.start_timestamp - } - - pub fn end_timestamp(&self) -> u64 { - self.end_timestamp - } - - pub fn is_exact_query(&self) -> bool { - self.is_exact_query - } - - pub fn output_labels(&self) -> &[String] { - &self.output_labels - } - - pub fn summary_type(&self) -> &SketchType { - &self.summary_type - } -} - -impl PartialOrd for PrecomputedSummaryRead { - fn partial_cmp(&self, other: &Self) -> Option { - match self.metric.partial_cmp(&other.metric) { - Some(Ordering::Equal) => {} - other => return other, - } - match self.aggregation_id.partial_cmp(&other.aggregation_id) { - Some(Ordering::Equal) => {} - other => return other, - } - match self.start_timestamp.partial_cmp(&other.start_timestamp) { - Some(Ordering::Equal) => {} - other => return other, - } - match self.end_timestamp.partial_cmp(&other.end_timestamp) { - Some(Ordering::Equal) => {} - other => return other, - } - match self.is_exact_query.partial_cmp(&other.is_exact_query) { - Some(Ordering::Equal) => {} - other => return other, - } - match self.output_labels.partial_cmp(&other.output_labels) { - Some(Ordering::Equal) => {} - other => return other, - } - match self.summary_type.partial_cmp(&other.summary_type) { - Some(Ordering::Equal) => {} - other => return other, - } - Some(Arc::as_ptr(&self.schema).cmp(&Arc::as_ptr(&other.schema))) - } -} - -impl UserDefinedLogicalNodeCore for PrecomputedSummaryRead { - fn name(&self) -> &str { - "PrecomputedSummaryRead" - } - - fn inputs(&self) -> Vec<&LogicalPlan> { - vec![] // Leaf node - no inputs - } - - fn schema(&self) -> &DFSchemaRef { - &self.schema - } - - fn expressions(&self) -> Vec { - vec![] - } - - fn fmt_for_explain(&self, f: &mut fmt::Formatter) -> fmt::Result { - write!( - f, - "PrecomputedSummaryRead: metric={}, agg_id={}, range=[{}, {}], exact={}, type={}, labels=[{}]", - self.metric, - self.aggregation_id, - self.start_timestamp, - self.end_timestamp, - self.is_exact_query, - self.summary_type, - self.output_labels.join(", ") - ) - } - - fn from_template(&self, _exprs: &[Expr], _inputs: &[LogicalPlan]) -> Self { - self.clone() - } - - fn with_exprs_and_inputs( - &self, - _exprs: Vec, - _inputs: Vec, - ) -> DFResult { - Ok(self.clone()) - } -} - -// ============================================================================ -// SummaryMergeMultiple - Merge multiple summaries by group key -// ============================================================================ - -/// Logical plan node: Merge multiple summaries with the same group key -/// -/// Takes an input with multiple rows per group key (e.g., from multiple -/// precomputed buckets) and merges them into one summary per group key. -/// This is used when tumbling windows need to be merged for a query range. -#[derive(Debug, Clone, PartialEq, Eq, Hash)] -pub struct SummaryMergeMultiple { - /// Input plan (typically PrecomputedSummaryRead) - input: Arc, - - /// Columns to group by when merging - group_by: Vec, - - /// Column containing the sketch/summary data - sketch_column: String, - - /// Type of summary being merged (for dispatch to correct merge logic) - summary_type: SketchType, - - /// Cached output schema (same as input - merging reduces rows, not columns) - schema: DFSchemaRef, -} - -impl SummaryMergeMultiple { - pub fn new( - input: Arc, - group_by: Vec, - sketch_column: String, - summary_type: SketchType, - ) -> Self { - // Schema is same as input (we reduce rows, not columns) - let schema = input.schema().clone(); - Self { - input, - group_by, - sketch_column, - summary_type, - schema, - } - } - - // Getters - pub fn input(&self) -> &LogicalPlan { - &self.input - } - - pub fn group_by(&self) -> &[String] { - &self.group_by - } - - pub fn sketch_column(&self) -> &str { - &self.sketch_column - } - - pub fn summary_type(&self) -> &SketchType { - &self.summary_type - } -} - -impl PartialOrd for SummaryMergeMultiple { - fn partial_cmp(&self, other: &Self) -> Option { - match self.group_by.partial_cmp(&other.group_by) { - Some(Ordering::Equal) => {} - other => return other, - } - match self.sketch_column.partial_cmp(&other.sketch_column) { - Some(Ordering::Equal) => {} - other => return other, - } - match self.summary_type.partial_cmp(&other.summary_type) { - Some(Ordering::Equal) => {} - other => return other, - } - match self.input.partial_cmp(&other.input) { - Some(Ordering::Equal) => {} - other => return other, - } - Some(Arc::as_ptr(&self.schema).cmp(&Arc::as_ptr(&other.schema))) - } -} - -impl UserDefinedLogicalNodeCore for SummaryMergeMultiple { - fn name(&self) -> &str { - "SummaryMergeMultiple" - } - - fn inputs(&self) -> Vec<&LogicalPlan> { - vec![self.input.as_ref()] - } - - fn schema(&self) -> &DFSchemaRef { - &self.schema - } - - fn expressions(&self) -> Vec { - vec![] - } - - fn fmt_for_explain(&self, f: &mut fmt::Formatter) -> fmt::Result { - write!( - f, - "SummaryMergeMultiple: group_by=[{}], sketch_column={}, type={}", - self.group_by.join(", "), - self.sketch_column, - self.summary_type - ) - } - - fn from_template(&self, _exprs: &[Expr], inputs: &[LogicalPlan]) -> Self { - Self { - input: Arc::new(inputs[0].clone()), - group_by: self.group_by.clone(), - sketch_column: self.sketch_column.clone(), - summary_type: self.summary_type.clone(), - schema: inputs[0].schema().clone(), - } - } - - fn with_exprs_and_inputs(&self, exprs: Vec, inputs: Vec) -> DFResult { - Ok(self.from_template(&exprs, &inputs)) - } -}