query/optimizer/
count_nest_aggr.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::collections::HashSet;
16use std::sync::Arc;
17
18use common_function::aggrs::aggr_wrapper::get_aggr_func;
19use common_function::aggrs::extremum::Extremum;
20use datafusion::config::ConfigOptions;
21use datafusion::functions_aggregate::count::count_udaf;
22use datafusion::logical_expr::{Extension, LogicalPlan, LogicalPlanBuilder, Sort};
23use datafusion_common::Result;
24use datafusion_common::tree_node::{Transformed, TreeNode};
25use datafusion_expr::{Expr, UserDefinedLogicalNodeCore, lit};
26use promql::extension_plan::{InstantManipulate, SeriesDivide, SeriesNormalize};
27use store_api::metric_engine_consts::DATA_SCHEMA_TSID_COLUMN_NAME;
28
29use crate::QueryEngineContext;
30use crate::optimizer::ExtensionAnalyzerRule;
31
32/// Rewrites `count(<presence-preserving-agg>(<vector_selector>) by (...))` into a presence-based
33/// group count.
34///
35/// This stays intentionally narrow:
36/// - the outer aggregate must be plain `count`
37/// - the inner aggregate must be a plain aggregate whose result existence is equivalent to input
38///   group existence
39/// - the inner input must be the direct instant-vector-selector plan
40/// - the outer count must only group by the evaluation timestamp
41#[derive(Debug)]
42pub struct CountNestAggrRule;
43
44impl ExtensionAnalyzerRule for CountNestAggrRule {
45    fn analyze(
46        &self,
47        plan: LogicalPlan,
48        _ctx: &QueryEngineContext,
49        _config: &ConfigOptions,
50    ) -> Result<LogicalPlan> {
51        plan.transform_down(&Self::rewrite_plan).map(|x| x.data)
52    }
53}
54
55impl CountNestAggrRule {
56    fn rewrite_plan(plan: LogicalPlan) -> Result<Transformed<LogicalPlan>> {
57        let LogicalPlan::Sort(sort) = plan else {
58            return Ok(Transformed::no(plan));
59        };
60
61        if let Some(rewritten) = Self::try_rewrite_sort(&sort)? {
62            Ok(Transformed::yes(rewritten))
63        } else {
64            Ok(Transformed::no(LogicalPlan::Sort(sort)))
65        }
66    }
67
68    fn try_rewrite_sort(sort: &Sort) -> Result<Option<LogicalPlan>> {
69        if sort.fetch.is_some() {
70            return Ok(None);
71        }
72
73        let LogicalPlan::Aggregate(outer_agg) = sort.input.as_ref() else {
74            return Ok(None);
75        };
76        if outer_agg.group_expr.len() != 1 || outer_agg.aggr_expr.len() != 1 {
77            return Ok(None);
78        }
79        let outer_time_expr = outer_agg.group_expr[0].clone();
80        let outer_count_arg =
81            match Self::aggregate_if(&outer_agg.aggr_expr[0], |name| name == "count") {
82                Some((_, arg)) => arg,
83                None => return Ok(None),
84            };
85
86        let LogicalPlan::Sort(inner_sort) = outer_agg.input.as_ref() else {
87            return Ok(None);
88        };
89        if inner_sort.fetch.is_some() {
90            return Ok(None);
91        }
92
93        let LogicalPlan::Aggregate(inner_agg) = inner_sort.input.as_ref() else {
94            return Ok(None);
95        };
96        if inner_agg.aggr_expr.len() != 1 || inner_agg.group_expr.is_empty() {
97            return Ok(None);
98        }
99        let (inner_is_count, inner_value_expr) =
100            match Self::aggregate_if(&inner_agg.aggr_expr[0], |name| {
101                Self::is_supported_inner_aggregate(name)
102            }) {
103                Some((name, arg)) => (name == "count", arg),
104                None => return Ok(None),
105            };
106        let Expr::Column(_) = inner_value_expr else {
107            return Ok(None);
108        };
109
110        let Expr::Column(outer_count_column) = outer_count_arg else {
111            return Ok(None);
112        };
113        let inner_output_field = inner_agg.schema.field(inner_agg.group_expr.len());
114        if outer_count_column.name != *inner_output_field.name() {
115            return Ok(None);
116        }
117
118        if !Self::is_projection_chain_to_instant(inner_agg.input.as_ref()) {
119            return Ok(None);
120        }
121
122        if !inner_agg
123            .group_expr
124            .iter()
125            .all(|expr| matches!(expr, Expr::Column(_)))
126        {
127            return Ok(None);
128        }
129
130        let Some(time_expr_pos) = inner_agg
131            .group_expr
132            .iter()
133            .position(|expr| expr == &outer_time_expr)
134        else {
135            return Ok(None);
136        };
137
138        let mut presence_group_exprs = Vec::with_capacity(inner_agg.group_expr.len());
139        presence_group_exprs.push(outer_time_expr.clone());
140        presence_group_exprs.extend(
141            inner_agg
142                .group_expr
143                .iter()
144                .enumerate()
145                .filter(|(idx, _)| *idx != time_expr_pos)
146                .map(|(_, expr)| expr.clone()),
147        );
148
149        let mut required_input_columns =
150            Self::collect_required_input_columns(&presence_group_exprs, inner_value_expr);
151        required_input_columns.extend(Self::collect_required_instant_columns(
152            inner_agg.input.as_ref(),
153        ));
154        let presence_source = Self::rebuild_projection_chain_to_instant(
155            inner_agg.input.as_ref(),
156            &required_input_columns,
157        )?;
158
159        let outer_value_name = outer_agg
160            .schema
161            .field(outer_agg.group_expr.len())
162            .name()
163            .clone();
164        let mut presence_input = LogicalPlanBuilder::from(presence_source);
165        if !inner_is_count {
166            presence_input = presence_input.filter(inner_value_expr.clone().is_not_null())?;
167        }
168        let presence_input = presence_input
169            .project(presence_group_exprs.clone())?
170            .distinct()?
171            .build()?;
172
173        let rewritten = LogicalPlanBuilder::from(presence_input)
174            .aggregate(
175                outer_agg.group_expr.clone(),
176                vec![count_udaf().call(vec![lit(1_i64)]).alias(outer_value_name)],
177            )?
178            .sort(sort.expr.clone())?
179            .build()?;
180
181        Ok(Some(rewritten))
182    }
183
184    fn collect_required_input_columns(group_exprs: &[Expr], value_expr: &Expr) -> HashSet<String> {
185        let mut required = HashSet::new();
186
187        for expr in group_exprs {
188            if let Expr::Column(column) = expr {
189                required.insert(column.name.clone());
190            }
191        }
192        if let Expr::Column(column) = value_expr {
193            // Keep the value column in the pruned instant input so `InstantManipulate`
194            // can still perform stale-NaN filtering before we project down to keys.
195            required.insert(column.name.clone());
196        }
197
198        required
199    }
200
201    fn collect_required_instant_columns(plan: &LogicalPlan) -> HashSet<String> {
202        let mut required = HashSet::new();
203        Self::collect_required_instant_columns_into(plan, &mut required);
204        required
205    }
206
207    fn collect_required_instant_columns_into(plan: &LogicalPlan, required: &mut HashSet<String>) {
208        match plan {
209            LogicalPlan::Projection(projection) => {
210                Self::collect_required_instant_columns_into(projection.input.as_ref(), required);
211            }
212            LogicalPlan::Extension(extension) => {
213                for expr in extension.node.expressions() {
214                    if let Expr::Column(column) = expr {
215                        required.insert(column.name);
216                    }
217                }
218
219                if extension.node.as_any().is::<SeriesDivide>()
220                    && extension.node.inputs()[0]
221                        .schema()
222                        .fields()
223                        .iter()
224                        .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME)
225                {
226                    required.insert(DATA_SCHEMA_TSID_COLUMN_NAME.to_string());
227                }
228
229                if let Some(input) = extension.node.inputs().into_iter().next() {
230                    Self::collect_required_instant_columns_into(input, required);
231                }
232            }
233            _ => {}
234        }
235    }
236
237    fn aggregate_if<F>(expr: &Expr, accept_name: F) -> Option<(&str, &Expr)>
238    where
239        F: FnOnce(&str) -> bool,
240    {
241        // The PromQL planner aliases `prom_min`/`prom_max` to the builtin names.
242        let func = get_aggr_func(expr)?;
243        let name = func.func.name();
244        if !accept_name(name)
245            || func.params.filter.is_some()
246            || func.params.distinct
247            || !func.params.order_by.is_empty()
248            || func.params.args.len() != 1
249        {
250            return None;
251        }
252
253        Some((name, &func.params.args[0]))
254    }
255
256    fn is_supported_inner_aggregate(name: &str) -> bool {
257        // PromQL `min`/`max` plan to `prom_min`/`prom_max`, which return NULL for the same
258        // groups as the builtin ones.
259        matches!(
260            name,
261            "count"
262                | "sum"
263                | "avg"
264                | "min"
265                | "max"
266                | Extremum::MIN_NAME
267                | Extremum::MAX_NAME
268                | "stddev_pop"
269                | "var_pop"
270        )
271    }
272
273    fn is_projection_chain_to_instant(plan: &LogicalPlan) -> bool {
274        let mut current = plan;
275        loop {
276            match current {
277                LogicalPlan::Projection(projection) => current = projection.input.as_ref(),
278                LogicalPlan::Extension(ext) => {
279                    return ext.node.as_any().is::<InstantManipulate>();
280                }
281                _ => return false,
282            }
283        }
284    }
285
286    fn rebuild_projection_chain_to_instant(
287        plan: &LogicalPlan,
288        required_columns: &HashSet<String>,
289    ) -> Result<LogicalPlan> {
290        match plan {
291            LogicalPlan::Projection(projection) => {
292                let input = Self::rebuild_projection_chain_to_instant(
293                    projection.input.as_ref(),
294                    required_columns,
295                )?;
296                LogicalPlanBuilder::from(input)
297                    .project(projection.expr.clone())?
298                    .build()
299            }
300            LogicalPlan::Extension(extension) => {
301                if let Some(instant) = extension.node.as_any().downcast_ref::<InstantManipulate>() {
302                    let input =
303                        Self::prune_instant_input(extension.node.inputs()[0], required_columns)?;
304                    return Ok(LogicalPlan::Extension(Extension {
305                        node: Arc::new(instant.with_exprs_and_inputs(vec![], vec![input])?),
306                    }));
307                }
308
309                Ok(plan.clone())
310            }
311            _ => Ok(plan.clone()),
312        }
313    }
314
315    fn prune_instant_input(
316        plan: &LogicalPlan,
317        required_columns: &HashSet<String>,
318    ) -> Result<LogicalPlan> {
319        match plan {
320            LogicalPlan::Extension(extension) => {
321                if let Some(normalize) = extension.node.as_any().downcast_ref::<SeriesNormalize>() {
322                    let input =
323                        Self::prune_instant_input(extension.node.inputs()[0], required_columns)?;
324                    return Ok(LogicalPlan::Extension(Extension {
325                        node: Arc::new(normalize.with_exprs_and_inputs(vec![], vec![input])?),
326                    }));
327                }
328
329                if let Some(divide) = extension.node.as_any().downcast_ref::<SeriesDivide>() {
330                    let divide_input = extension.node.inputs()[0].clone();
331
332                    let projection_exprs = divide_input
333                        .schema()
334                        .fields()
335                        .iter()
336                        .filter(|field| required_columns.contains(field.name()))
337                        .map(|field| {
338                            Expr::Column(datafusion_common::Column::from_name(field.name().clone()))
339                        })
340                        .collect::<Vec<_>>();
341                    let projected_input = LogicalPlanBuilder::from(divide_input)
342                        .project(projection_exprs)?
343                        .build()?;
344
345                    return Ok(LogicalPlan::Extension(Extension {
346                        node: Arc::new(
347                            divide.with_exprs_and_inputs(vec![], vec![projected_input])?,
348                        ),
349                    }));
350                }
351
352                Ok(plan.clone())
353            }
354            _ => Ok(plan.clone()),
355        }
356    }
357}
query/optimizer/count_nest_aggr.rs

query/optimizer/
count_nest_aggr.rs
