query/dist_plan/
commutativity.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::{
19    StateMergeHelper, get_aggr_func, is_all_aggr_exprs_steppable,
20};
21use common_telemetry::debug;
22use datafusion::error::Result as DfResult;
23use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion};
24use datafusion_expr::expr::NullTreatment;
25use datafusion_expr::{Expr, LogicalPlan, UserDefinedLogicalNode};
26use promql::extension_plan::{
27    EmptyMetric, InstantManipulate, RangeManipulate, SeriesDivide, SeriesNormalize,
28};
29use store_api::metric_engine_consts::DATA_SCHEMA_TSID_COLUMN_NAME;
30
31use crate::dist_plan::MergeScanLogicalPlan;
32use crate::dist_plan::analyzer::AliasMapping;
33use crate::dist_plan::merge_sort::{MergeSortLogicalPlan, merge_sort_transformer};
34
35pub struct StepTransformAction {
36    extra_parent_plans: Vec<LogicalPlan>,
37    new_child_plan: Option<LogicalPlan>,
38}
39
40/// generate the upper aggregation plan that will execute on the frontend.
41/// Basically a logical plan resembling the following:
42/// Projection:
43///     Aggregate:
44///
45/// from Aggregate
46///
47/// The upper Projection exists sole to make sure parent plan can recognize the output
48/// of the upper aggregation plan.
49pub fn step_aggr_to_upper_aggr(
50    aggr_plan: &LogicalPlan,
51) -> datafusion_common::Result<StepTransformAction> {
52    let LogicalPlan::Aggregate(input_aggr) = aggr_plan else {
53        return Err(datafusion_common::DataFusionError::Plan(
54            "step_aggr_to_upper_aggr only accepts Aggregate plan".to_string(),
55        ));
56    };
57    if !is_all_aggr_exprs_steppable(&input_aggr.aggr_expr) {
58        return Err(datafusion_common::DataFusionError::NotImplemented(format!(
59            "Some aggregate expressions are not steppable in [{}]",
60            input_aggr
61                .aggr_expr
62                .iter()
63                .map(|e| e.to_string())
64                .collect::<Vec<_>>()
65                .join(", ")
66        )));
67    }
68
69    let step_aggr_plan = StateMergeHelper::split_aggr_node(input_aggr.clone())?;
70
71    // TODO(discord9): remove duplication
72    let ret = StepTransformAction {
73        extra_parent_plans: vec![step_aggr_plan.upper_merge.clone()],
74        new_child_plan: Some(step_aggr_plan.lower_state.clone()),
75    };
76    Ok(ret)
77}
78
79#[allow(dead_code)]
80pub enum Commutativity {
81    Commutative,
82    PartialCommutative,
83    ConditionalCommutative(Option<Transformer>),
84    TransformedCommutative {
85        /// Return plans from parent to child order
86        transformer: Option<StageTransformer>,
87    },
88    NonCommutative,
89    Unimplemented,
90    /// For unrelated plans like DDL
91    Unsupported,
92}
93
94pub struct Categorizer {}
95
96impl Categorizer {
97    pub fn check_plan(
98        plan: &LogicalPlan,
99        partition_cols: Option<AliasMapping>,
100    ) -> DfResult<Commutativity> {
101        // Subquery is treated separately in `inspect_plan_with_subquery`. To avoid rewrite the
102        // "maybe rewritten" plan, stop the check here.
103        if has_subquery(plan)? {
104            return Ok(Commutativity::Unimplemented);
105        }
106
107        let partition_cols = partition_cols.unwrap_or_default();
108
109        let comm = match plan {
110            LogicalPlan::Projection(proj) => {
111                for expr in &proj.expr {
112                    let commutativity = Self::check_expr(expr);
113                    if !matches!(commutativity, Commutativity::Commutative) {
114                        return Ok(commutativity);
115                    }
116                }
117                Commutativity::Commutative
118            }
119            LogicalPlan::Filter(filter) => Self::check_expr(&filter.predicate),
120            LogicalPlan::Window(_) => Commutativity::Unimplemented,
121            LogicalPlan::Aggregate(aggr) => {
122                // The pinned Substrait codec drops IGNORE NULLS, so keep these aggregates on the frontend.
123                if aggr.aggr_expr.iter().any(|expr| {
124                    get_aggr_func(expr).is_some_and(|func| {
125                        func.params.null_treatment == Some(NullTreatment::IgnoreNulls)
126                    })
127                }) {
128                    return Ok(Commutativity::NonCommutative);
129                }
130                // The state/merge split maps each group expression to one output column,
131                // which doesn't hold for grouping sets.
132                let has_grouping_set = aggr
133                    .group_expr
134                    .iter()
135                    .any(|expr| matches!(expr, Expr::GroupingSet(_)));
136                let is_all_steppable =
137                    !has_grouping_set && is_all_aggr_exprs_steppable(&aggr.aggr_expr);
138                let matches_partition = Self::check_partition(&aggr.group_expr, &partition_cols);
139                if !matches_partition && is_all_steppable {
140                    debug!("Plan is steppable: {plan}");
141                    return Ok(Commutativity::TransformedCommutative {
142                        transformer: Some(Arc::new(|plan: &LogicalPlan| {
143                            debug!("Before Step optimize: {plan}");
144                            let ret = step_aggr_to_upper_aggr(plan);
145                            ret.inspect_err(|err| {
146                                common_telemetry::error!("Failed to step aggregate plan: {err:?}");
147                            })
148                            .map(|s| TransformerAction {
149                                extra_parent_plans: s.extra_parent_plans,
150                                new_child_plan: s.new_child_plan,
151                            })
152                        })),
153                    });
154                }
155                if !matches_partition {
156                    return Ok(Commutativity::NonCommutative);
157                }
158                for expr in &aggr.aggr_expr {
159                    let commutativity = Self::check_expr(expr);
160                    if !matches!(commutativity, Commutativity::Commutative) {
161                        return Ok(commutativity);
162                    }
163                }
164                // all group by expressions are partition columns can push down, unless
165                // another push down(including `Limit` or `Sort`) is already in progress(which will then prevent next cond commutative node from being push down).
166                // TODO(discord9): This is a temporary solution(that works), a better description of
167                // commutativity is needed under this situation.
168                Commutativity::ConditionalCommutative(None)
169            }
170            LogicalPlan::Sort(_sort) => {
171                if partition_cols.is_empty() {
172                    return Ok(Commutativity::Commutative);
173                }
174
175                // sort plan needs to consider column priority
176                // Change Sort to MergeSort which assumes the input streams are already sorted hence can be more efficient.
177                Commutativity::ConditionalCommutative(Some(Arc::new(merge_sort_transformer)))
178            }
179            LogicalPlan::Join(_) => Commutativity::NonCommutative,
180            LogicalPlan::Repartition(_) => {
181                // unsupported? or non-commutative
182                Commutativity::Unimplemented
183            }
184            LogicalPlan::Union(_) => Commutativity::Unimplemented,
185            LogicalPlan::TableScan(_) => Commutativity::Commutative,
186            LogicalPlan::EmptyRelation(_) => Commutativity::NonCommutative,
187            LogicalPlan::Subquery(_) => Commutativity::Unimplemented,
188            LogicalPlan::SubqueryAlias(_) => Commutativity::Commutative,
189            LogicalPlan::Limit(limit) => {
190                // Only execute `fetch` on remote nodes.
191                // wait for https://github.com/apache/arrow-datafusion/pull/7669
192                if partition_cols.is_empty() && limit.fetch.is_some() {
193                    Commutativity::Commutative
194                } else if limit.skip.is_none() && limit.fetch.is_some() {
195                    Commutativity::PartialCommutative
196                } else {
197                    Commutativity::Unimplemented
198                }
199            }
200            LogicalPlan::Extension(extension) => {
201                Self::check_extension_plan(extension.node.as_ref() as _, &partition_cols)
202            }
203            LogicalPlan::Distinct(_) => {
204                if partition_cols.is_empty() {
205                    Commutativity::Commutative
206                } else {
207                    Commutativity::PartialCommutative
208                }
209            }
210            LogicalPlan::Unnest(_) => Commutativity::Commutative,
211            LogicalPlan::Statement(_) => Commutativity::Unsupported,
212            LogicalPlan::Values(_) => Commutativity::Unsupported,
213            LogicalPlan::Explain(_) => Commutativity::Unsupported,
214            LogicalPlan::Analyze(_) => Commutativity::Unsupported,
215            LogicalPlan::DescribeTable(_) => Commutativity::Unsupported,
216            LogicalPlan::Dml(_) => Commutativity::Unsupported,
217            LogicalPlan::Ddl(_) => Commutativity::Unsupported,
218            LogicalPlan::Copy(_) => Commutativity::Unsupported,
219            LogicalPlan::RecursiveQuery(_) => Commutativity::Unsupported,
220        };
221
222        Ok(comm)
223    }
224
225    pub fn check_extension_plan(
226        plan: &dyn UserDefinedLogicalNode,
227        partition_cols: &AliasMapping,
228    ) -> Commutativity {
229        match plan.name() {
230            name if name == SeriesDivide::name() => {
231                let series_divide = plan.as_any().downcast_ref::<SeriesDivide>().unwrap();
232                // Metric engine `__tsid` uniquely identifies a time-series. Treat a series divide
233                // that keys by `__tsid` as commutative across regions so it can be pushed down.
234                if series_divide
235                    .tags()
236                    .iter()
237                    .any(|tag| tag == DATA_SCHEMA_TSID_COLUMN_NAME)
238                {
239                    return Commutativity::Commutative;
240                }
241
242                let tags = series_divide.tags().iter().collect::<HashSet<_>>();
243
244                for all_alias in partition_cols.values() {
245                    let all_alias = all_alias.iter().map(|c| &c.name).collect::<HashSet<_>>();
246                    if tags.intersection(&all_alias).count() == 0 {
247                        return Commutativity::NonCommutative;
248                    }
249                }
250
251                Commutativity::Commutative
252            }
253            name if name == SeriesNormalize::name()
254                || name == InstantManipulate::name()
255                || name == RangeManipulate::name() =>
256            {
257                // They should always follows Series Divide.
258                // Either all commutative or all non-commutative (which will be blocked by SeriesDivide).
259                Commutativity::Commutative
260            }
261            name if name == EmptyMetric::name()
262                || name == MergeScanLogicalPlan::name()
263                || name == MergeSortLogicalPlan::name() =>
264            {
265                Commutativity::Unimplemented
266            }
267            _ => Commutativity::Unsupported,
268        }
269    }
270
271    pub fn check_expr(expr: &Expr) -> Commutativity {
272        #[allow(deprecated)]
273        match expr {
274            Expr::Column(_)
275            | Expr::ScalarVariable(_, _)
276            | Expr::Literal(_, _)
277            | Expr::BinaryExpr(_)
278            | Expr::Not(_)
279            | Expr::IsNotNull(_)
280            | Expr::IsNull(_)
281            | Expr::IsTrue(_)
282            | Expr::IsFalse(_)
283            | Expr::IsNotTrue(_)
284            | Expr::IsNotFalse(_)
285            | Expr::Negative(_)
286            | Expr::Between(_)
287            | Expr::Exists(_)
288            | Expr::InList(_)
289            | Expr::Like(_)
290            | Expr::Case(_) => Commutativity::Commutative,
291            // Annotation collection must not affect distribution. Keep scalar UDFs pushdownable;
292            // TODO: Preserve datanode PromQL annotations through plan decoding, return them in
293            // Flight stream metadata, and merge them on the frontend.
294            Expr::ScalarFunction(_) => Commutativity::Commutative,
295            Expr::AggregateFunction(_udaf) => Commutativity::Commutative,
296
297            Expr::SimilarTo(_)
298            | Expr::IsUnknown(_)
299            | Expr::IsNotUnknown(_)
300            | Expr::WindowFunction(_)
301            | Expr::InSubquery(_)
302            | Expr::ScalarSubquery(_)
303            | Expr::HigherOrderFunction(_)
304            | Expr::Lambda(_)
305            | Expr::LambdaVariable(_)
306            | Expr::Wildcard { .. } => Commutativity::Unimplemented,
307
308            Expr::Alias(alias) => Self::check_expr(&alias.expr),
309            Expr::Cast(cast) => Self::check_expr(&cast.expr),
310            Expr::TryCast(try_cast) => Self::check_expr(&try_cast.expr),
311
312            Expr::Unnest(_)
313            | Expr::GroupingSet(_)
314            | Expr::Placeholder(_)
315            | Expr::OuterReferenceColumn(_, _)
316            | Expr::SetComparison(_) => Commutativity::Unimplemented,
317        }
318    }
319
320    /// Return true if the given expr and partition cols satisfied the rule.
321    /// In this case the plan can be treated as fully commutative.
322    ///
323    /// So only if every partition column is itself one of `exprs`, return true.
324    /// Otherwise return false.
325    ///
326    /// An expression that only references a partition column, like `substr(host, 3, 1)`
327    /// or `k % 2`, doesn't count: it can put rows from different partitions into the same
328    /// group.
329    fn check_partition(exprs: &[Expr], partition_cols: &AliasMapping) -> bool {
330        let group_cols = exprs
331            .iter()
332            .filter_map(|expr| {
333                let mut expr = expr;
334                while let Expr::Alias(alias) = expr {
335                    expr = &alias.expr;
336                }
337                match expr {
338                    Expr::Column(column) => Some(column.name.clone()),
339                    _ => None,
340                }
341            })
342            .collect::<HashSet<_>>();
343        for all_alias in partition_cols.values() {
344            let all_alias = all_alias
345                .iter()
346                .map(|c| c.name.clone())
347                .collect::<HashSet<_>>();
348            // check if group columns intersect with all alias of partition columns
349            // is empty, if it's empty, not all partition columns show up in `exprs`
350            if group_cols.intersection(&all_alias).count() == 0 {
351                return false;
352            }
353        }
354
355        true
356    }
357}
358
359#[cfg(test)]
360mod tests {
361    use std::collections::{BTreeMap, BTreeSet};
362
363    use datafusion_common::Column;
364    use datafusion_expr::LogicalPlanBuilder;
365    use datafusion_expr::expr::ScalarFunction;
366    use datafusion_functions::core::coalesce;
367    use promql::functions::{
368        NativeHistogramDrop, NativeHistogramFraction, NativeHistogramQuantile,
369    };
370
371    use super::*;
372
373    #[test]
374    fn series_divide_by_tsid_is_commutative() {
375        let input = LogicalPlanBuilder::empty(false).build().unwrap();
376        let series_divide = SeriesDivide::new(
377            vec![DATA_SCHEMA_TSID_COLUMN_NAME.to_string()],
378            "ts".to_string(),
379            input,
380        );
381
382        let partition_cols: AliasMapping = BTreeMap::from([(
383            "some_partition_col".to_string(),
384            BTreeSet::from([Column::from_name("some_partition_col")]),
385        )]);
386
387        let commutativity = Categorizer::check_extension_plan(&series_divide, &partition_cols);
388        assert!(matches!(commutativity, Commutativity::Commutative));
389    }
390
391    #[test]
392    fn annotated_histogram_helpers_do_not_block_pushdown() {
393        for udf in [
394            NativeHistogramQuantile::scalar_udf(),
395            NativeHistogramFraction::scalar_udf(),
396            NativeHistogramDrop::warning_bool_false_udf(String::new(), None),
397        ] {
398            let helper = Expr::ScalarFunction(ScalarFunction::new_udf(Arc::new(udf), vec![]));
399            let expr = Expr::ScalarFunction(ScalarFunction::new_udf(coalesce(), vec![helper]));
400            assert!(matches!(
401                Categorizer::check_expr(&expr),
402                Commutativity::Commutative
403            ));
404        }
405    }
406}
407
408pub type Transformer = Arc<dyn Fn(&LogicalPlan) -> Option<LogicalPlan>>;
409
410/// Returns transformer action that need to be applied
411pub type StageTransformer = Arc<dyn Fn(&LogicalPlan) -> DfResult<TransformerAction>>;
412
413/// The Action that a transformer should take on the plan.
414pub struct TransformerAction {
415    /// list of plans that need to be applied to parent plans, in the order of parent to child.
416    /// i.e. if this returns `[Projection, Aggregate]`, then the parent plan should be transformed to
417    /// ```ignore
418    /// Original Parent Plan:
419    ///     Projection:
420    ///         Aggregate:
421    ///             MergeScan: ...
422    /// ```
423    pub extra_parent_plans: Vec<LogicalPlan>,
424    /// new child plan, if None, use the original plan.
425    pub new_child_plan: Option<LogicalPlan>,
426}
427
428pub fn partial_commutative_transformer(plan: &LogicalPlan) -> Option<LogicalPlan> {
429    Some(plan.clone())
430}
431
432fn has_subquery(plan: &LogicalPlan) -> DfResult<bool> {
433    let mut found = false;
434    plan.apply_expressions(|e| {
435        e.apply(|x| {
436            if matches!(
437                x,
438                Expr::Exists(_) | Expr::InSubquery(_) | Expr::ScalarSubquery(_)
439            ) {
440                found = true;
441                Ok(TreeNodeRecursion::Stop)
442            } else {
443                Ok(TreeNodeRecursion::Continue)
444            }
445        })
446    })?;
447    Ok(found)
448}
449
450#[cfg(test)]
451mod test {
452    use datafusion_expr::{LogicalPlanBuilder, Sort};
453
454    use super::*;
455
456    #[test]
457    fn sort_on_empty_partition() {
458        let plan = LogicalPlan::Sort(Sort {
459            expr: vec![],
460            input: Arc::new(LogicalPlanBuilder::empty(false).build().unwrap()),
461            fetch: None,
462        });
463        assert!(matches!(
464            Categorizer::check_plan(&plan, Some(Default::default())).unwrap(),
465            Commutativity::Commutative
466        ));
467    }
468}
query/dist_plan/commutativity.rs

query/dist_plan/
commutativity.rs
