From fd2306ffaa3ec2d760f05df3fcf3791efb1ff991 Mon Sep 17 00:00:00 2001 From: Huaijin Date: Wed, 12 Aug 2026 15:23:50 +0800 Subject: [PATCH] fix(search): preserve indexed top-n aggregation results (#13779) - close: https://github.com/openobserve/openobserve/issues/13363 ## Summary - Prevent `AggregateTopkExec` from truncating file-local indexed aggregation rows before the final aggregate merges duplicate group keys. - Keep the Tantivy `SimpleTopN` optimization enabled while removing only the conflicting downstream top-k node. - Add a regression test that constructs the conflicting physical plan, verifies `AggregateTopkExec` is present before leader index optimization, and verifies it is removed afterward while `SimpleTopN` remains selected. ## Root cause `SimpleTopN` emits partial count rows per `(file, group key)`. `AggregateTopkExec` previously ranked and truncated those individual rows before the final aggregate could combine identical keys across files. With enough files, this silently undercounted results and removed entire groups while `is_partial` remained false. The rewrite removes the conflicting node bottom-up. This is important because `AggregateTopkExec` is inside the follower plan already serialized by `RemoteScanExec`; rebuilding the parent updates those serialized plan bytes before attaching `SimpleTopN`. ## Verification Using 121 Parquet files and 121 Tantivy index files containing 93,601 records across 12 groups: | Result | Before | After | |---|---:|---:| | Groups returned | 9 | 12 | | Sum of grouped counts | 83,201 | 93,601 | | Baseline `COUNT(*)` | 93,601 | 93,601 | | `is_partial` | false | false | Tests run: - `cargo fmt --all -- --check` - `cargo test -p search test_leader_rule_topn_removes_conflicting_aggregate_topk -- --nocapture` - `cargo test -p search datafusion::optimizer::physical_optimizer::index_optimizer::tests -- --nocapture` - `cargo test -p search aggregate_topk -- --nocapture` - `cargo check -p search` --- .../physical_optimizer/index_optimizer/mod.rs | 42 +++++++++++++++++-- 1 file changed, 38 insertions(+), 4 deletions(-) diff --git a/src/search/src/datafusion/optimizer/physical_optimizer/index_optimizer/mod.rs b/src/search/src/datafusion/optimizer/physical_optimizer/index_optimizer/mod.rs index c1640b316d..48e8ec64d5 100644 --- a/src/search/src/datafusion/optimizer/physical_optimizer/index_optimizer/mod.rs +++ b/src/search/src/datafusion/optimizer/physical_optimizer/index_optimizer/mod.rs @@ -40,7 +40,10 @@ mod topn; mod utils; use crate::datafusion::{ - distributed_plan::{empty_exec::NewEmptyExec, remote_scan_exec::RemoteScanExec}, + distributed_plan::{ + aggregate_topk_exec::AggregateTopkExec, empty_exec::NewEmptyExec, + remote_scan_exec::RemoteScanExec, + }, optimizer::physical_optimizer::index_optimizer::{ count::is_simple_count, distinct::is_simple_distinct, @@ -276,12 +279,19 @@ impl TreeNodeRewriter for LeaderIndexOptimizer { #[derive(Debug)] struct IndexOptimizerRewrite { index_optimizer_mode: IndexOptimizeMode, + remove_aggregate_topk: bool, } impl IndexOptimizerRewrite { fn new(index_optimizer_mode: IndexOptimizeMode) -> Self { + // SimpleTopN produces one partial aggregate row per (file, group key). + // AggregateTopkExec ranks those individual rows before the final aggregate can merge + // duplicate group keys, so the two optimizations cannot safely be combined. + let remove_aggregate_topk = + matches!(index_optimizer_mode, IndexOptimizeMode::SimpleTopN(..)); IndexOptimizerRewrite { index_optimizer_mode, + remove_aggregate_topk, } } } @@ -290,6 +300,10 @@ impl TreeNodeRewriter for IndexOptimizerRewrite { type Node = Arc; fn f_up(&mut self, node: Arc) -> Result> { + if self.remove_aggregate_topk && node.downcast_ref::().is_some() { + return Ok(Transformed::yes(node.children()[0].clone())); + } + if let Some(remote) = node.downcast_ref::() { let remote = Arc::new( remote @@ -342,7 +356,8 @@ mod tests { use crate::datafusion::{ distributed_plan::node::RemoteScanNode, optimizer::physical_optimizer::{ - index_optimizer::utils::tests::get_remote_scan, remote_scan::RemoteScanRule, + aggregate_topk::AggregateTopkRule, index_optimizer::utils::tests::get_remote_scan, + remote_scan::RemoteScanRule, }, table_provider::empty_table::NewEmptyTable, }; @@ -463,7 +478,7 @@ mod tests { } #[tokio::test] - async fn test_leader_rule_topn() { + async fn test_leader_rule_topn_removes_conflicting_aggregate_topk() { let schema = Arc::new(Schema::new(vec![ Field::new("_timestamp", DataType::Int64, false), Field::new("name", DataType::Utf8, false), @@ -483,7 +498,7 @@ mod tests { .with_config(SessionConfig::new().with_target_partitions(12)) .with_runtime_env(Arc::new(RuntimeEnvBuilder::new().build().unwrap())) .with_physical_optimizer_rule(Arc::new(remote_scan_rule)) - .with_physical_optimizer_rule(Arc::new(LeaderIndexOptimizerRule::new(index_fields))) + .with_physical_optimizer_rule(Arc::new(AggregateTopkRule::new(10))) .with_default_features() .build(); let ctx = SessionContext::new_with_state(state); @@ -494,6 +509,25 @@ mod tests { let plan = ctx.state().create_logical_plan(sql).await.unwrap(); let plan = ctx.state().create_physical_plan(&plan).await.unwrap(); + let plan_before_index_optimizer = datafusion::physical_plan::displayable(plan.as_ref()) + .indent(true) + .to_string(); + assert!( + plan_before_index_optimizer.contains("AggregateTopkExec"), + "the regression setup must contain AggregateTopkExec before SimpleTopN is selected:\n{plan_before_index_optimizer}" + ); + + let plan = LeaderIndexOptimizerRule::new(index_fields) + .optimize(plan, &ConfigOptions::new()) + .unwrap(); + let plan_after_index_optimizer = datafusion::physical_plan::displayable(plan.as_ref()) + .indent(true) + .to_string(); + assert!( + !plan_after_index_optimizer.contains("AggregateTopkExec"), + "SimpleTopN must not truncate file-local rows before the final aggregate:\n{plan_after_index_optimizer}" + ); + let remote_scan = get_remote_scan(plan); assert_eq!( remote_scan[0].index_optimize_mode(),