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`
This commit is contained in:
Huaijin
2026-08-12 07:23:50 +00:00
committed by GitHub
parent 3d06ed0423
commit fd2306ffaa
@@ -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<dyn ExecutionPlan>;
fn f_up(&mut self, node: Arc<dyn ExecutionPlan>) -> Result<Transformed<Self::Node>> {
if self.remove_aggregate_topk && node.downcast_ref::<AggregateTopkExec>().is_some() {
return Ok(Transformed::yes(node.children()[0].clone()));
}
if let Some(remote) = node.downcast_ref::<RemoteScanExec>() {
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(),