query/optimizer/
infer_aggregate_input_bounds.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::sync::Arc;
16
17use arrow_schema::DataType;
18use datafusion::functions_aggregate::min_max::{max_udaf, min_udaf};
19use datafusion::functions_aggregate::sum::sum_udaf;
20use datafusion_common::tree_node::{Transformed, TreeNode};
21use datafusion_common::{Column, DFSchema, Result, ScalarValue};
22use datafusion_expr::expr::AggregateFunction;
23use datafusion_expr::utils::split_conjunction;
24use datafusion_expr::{Aggregate, Expr, ExprSchemable, Filter, LogicalPlan, Operator};
25use datafusion_optimizer::optimizer::ApplyOrder;
26use datafusion_optimizer::{OptimizerConfig, OptimizerRule};
27
28use crate::dist_plan::MergeScanLogicalPlan;
29
30/// Derives an input bound from an existential HAVING condition and MIN/MAX.
31///
32/// For example, `MIN(x), SUM(CASE WHEN p AND x < u THEN 1 ELSE 0 END)` with
33/// HAVING `SUM(...) > 0` can discard rows outside `x < u`: they contribute zero
34/// to SUM, and every surviving group has a witness below u, preserving its MIN.
35/// MAX permits the symmetric lower bound. The original HAVING must remain.
36///
37/// This deliberately requires exactly these two aggregates and ordinary column
38/// grouping. Extra aggregates, grouping sets, floating-point extrema and arbitrary
39/// expressions need separate preservation proofs. It must run before MergeScan
40/// hides remote inputs, as well as in the non-distributed logical optimizer.
41#[derive(Debug)]
42pub(crate) struct InferAggregateInputBoundsRule;
43
44impl OptimizerRule for InferAggregateInputBoundsRule {
45    fn name(&self) -> &str {
46        "infer_aggregate_input_bounds"
47    }
48
49    fn apply_order(&self) -> Option<ApplyOrder> {
50        Some(ApplyOrder::TopDown)
51    }
52
53    fn rewrite(
54        &self,
55        plan: LogicalPlan,
56        _config: &dyn OptimizerConfig,
57    ) -> Result<Transformed<LogicalPlan>> {
58        let LogicalPlan::Filter(mut having) = plan else {
59            return Ok(Transformed::no(plan));
60        };
61        let LogicalPlan::Aggregate(aggregate) = having.input.as_ref() else {
62            return Ok(Transformed::no(LogicalPlan::Filter(having)));
63        };
64        // The pre-MergeScan pass already handled remote inputs. Their filters
65        // are hidden from this pass, so re-inference can add useless local work.
66        if aggregate.input.exists(|plan| {
67            Ok(matches!(plan, LogicalPlan::Extension(extension)
68                if extension.node.as_any().is::<MergeScanLogicalPlan>()))
69        })? {
70            return Ok(Transformed::no(LogicalPlan::Filter(having)));
71        }
72        let Some(bound) = infer_bound(aggregate, &having.predicate)? else {
73            return Ok(Transformed::no(LogicalPlan::Filter(having)));
74        };
75        if already_filtered(&aggregate.input, &bound) {
76            return Ok(Transformed::no(LogicalPlan::Filter(having)));
77        }
78
79        let mut aggregate = aggregate.clone();
80        aggregate.input = Arc::new(LogicalPlan::Filter(Filter::try_new(
81            bound,
82            aggregate.input,
83        )?));
84        having.input = Arc::new(LogicalPlan::Aggregate(aggregate));
85        Ok(Transformed::yes(LogicalPlan::Filter(having)))
86    }
87}
88
89fn infer_bound(aggregate: &Aggregate, having: &Expr) -> Result<Option<Expr>> {
90    if aggregate.aggr_expr.len() != 2
91        || aggregate.group_expr.is_empty()
92        || !aggregate
93            .group_expr
94            .iter()
95            .all(|e| matches!(e, Expr::Column(_)))
96        || !safe_predicate(having, aggregate.schema.as_ref())?
97    {
98        return Ok(None);
99    }
100
101    for conjunct in split_conjunction(having) {
102        let Some(sum_index) = positive_sum_index(conjunct, aggregate) else {
103            continue;
104        };
105        let Some(sum) = plain_aggregate(&aggregate.aggr_expr[sum_index]) else {
106            continue;
107        };
108        let Some(extremum) = plain_aggregate(&aggregate.aggr_expr[1 - sum_index]) else {
109            continue;
110        };
111        if let Some(bound) = bound_from_aggregates(sum, extremum, &aggregate.input)? {
112            return Ok(Some(bound));
113        }
114    }
115    Ok(None)
116}
117
118fn positive_sum_index(predicate: &Expr, aggregate: &Aggregate) -> Option<usize> {
119    let Expr::BinaryExpr(comparison) = predicate else {
120        return None;
121    };
122    let (output, threshold, operator) = match (comparison.left.as_ref(), comparison.right.as_ref())
123    {
124        (Expr::Column(column), literal) => (column, literal, comparison.op),
125        (literal, Expr::Column(column)) => (column, literal, comparison.op.swap()?),
126        _ => return None,
127    };
128    if !((operator == Operator::Gt && integer_literal(threshold) == Some(0))
129        || (operator == Operator::GtEq && integer_literal(threshold) == Some(1)))
130    {
131        return None;
132    }
133    let index = aggregate.schema.index_of_column(output).ok()?;
134    index
135        .checked_sub(aggregate.group_expr.len())
136        .filter(|&i| i < 2)
137}
138
139fn plain_aggregate(expr: &Expr) -> Option<&AggregateFunction> {
140    let expr = match expr {
141        Expr::Alias(alias) => alias.expr.as_ref(),
142        _ => expr,
143    };
144    let Expr::AggregateFunction(function) = expr else {
145        return None;
146    };
147    let params = &function.params;
148    (params.args.len() == 1
149        && !params.distinct
150        && params.filter.is_none()
151        && params.order_by.is_empty()
152        && params.null_treatment.is_none())
153    .then_some(function)
154}
155
156fn bound_from_aggregates(
157    sum: &AggregateFunction,
158    extremum: &AggregateFunction,
159    input: &LogicalPlan,
160) -> Result<Option<Expr>> {
161    let schema = input.schema();
162    // Compare function identities, not names: a user-defined "sum" need not
163    // have the built-in aggregate's zero-contribution semantics.
164    if sum.func != sum_udaf() || (extremum.func != min_udaf() && extremum.func != max_udaf()) {
165        return Ok(None);
166    }
167    let argument = &extremum.params.args[0];
168    let argument_type = argument.get_type(schema)?;
169    if !(argument_type.is_integer() || matches!(argument_type, DataType::Timestamp(..)))
170        || !safe_operand(argument, schema)?
171    {
172        return Ok(None);
173    }
174    let Some(condition) = indicator_condition(&sum.params.args[0]) else {
175        return Ok(None);
176    };
177    if !safe_predicate(condition, schema)?
178        || !columns_are_from_scan(argument, input)?
179        || !columns_are_from_scan(condition, input)?
180    {
181        return Ok(None);
182    }
183    Ok(split_conjunction(condition)
184        .into_iter()
185        .find_map(|predicate| extremum_bound(predicate, argument, extremum, &argument_type)))
186}
187
188fn indicator_condition(expr: &Expr) -> Option<&Expr> {
189    let Expr::Case(case) = expr else {
190        return None;
191    };
192    let [(condition, contribution)] = case.when_then_expr.as_slice() else {
193        return None;
194    };
195    (case.expr.is_none()
196        && integer_literal(contribution) == Some(1)
197        && case.else_expr.as_deref().and_then(integer_literal) == Some(0))
198    .then_some(condition)
199}
200
201fn extremum_bound(
202    predicate: &Expr,
203    argument: &Expr,
204    extremum: &AggregateFunction,
205    argument_type: &DataType,
206) -> Option<Expr> {
207    let Expr::BinaryExpr(comparison) = predicate else {
208        return None;
209    };
210    let (literal, operator) = if comparison.left.as_ref() == argument {
211        (comparison.right.as_ref(), comparison.op)
212    } else if comparison.right.as_ref() == argument {
213        (comparison.left.as_ref(), comparison.op.swap()?)
214    } else {
215        return None;
216    };
217    let Expr::Literal(bound, _) = literal else {
218        return None;
219    };
220    if bound.is_null() || &bound.data_type() != argument_type {
221        return None;
222    }
223    let preserves_extremum = if extremum.func == min_udaf() {
224        matches!(operator, Operator::Lt | Operator::LtEq)
225    } else {
226        matches!(operator, Operator::Gt | Operator::GtEq)
227    };
228    preserves_extremum.then(|| predicate.clone())
229}
230
231fn integer_literal(expr: &Expr) -> Option<i64> {
232    match expr {
233        Expr::Literal(ScalarValue::Int64(integer), _) => *integer,
234        Expr::Literal(ScalarValue::Int32(integer), _) => integer.map(i64::from),
235        _ => None,
236    }
237}
238
239/// Rejects explicitly volatile, fallible or subquery expressions. Column
240/// references additionally require input-plan provenance checks.
241fn safe_predicate(expr: &Expr, schema: &DFSchema) -> Result<bool> {
242    match expr {
243        Expr::BinaryExpr(binary) => Ok(matches!(
244            binary.op,
245            Operator::And
246                | Operator::Or
247                | Operator::Eq
248                | Operator::NotEq
249                | Operator::Lt
250                | Operator::LtEq
251                | Operator::Gt
252                | Operator::GtEq
253        ) && safe_predicate(&binary.left, schema)?
254            && safe_predicate(&binary.right, schema)?),
255        Expr::IsNull(expr) | Expr::IsNotNull(expr) => safe_operand(expr, schema),
256        _ => safe_operand(expr, schema),
257    }
258}
259
260fn safe_operand(expr: &Expr, schema: &DFSchema) -> Result<bool> {
261    match expr {
262        Expr::Column(_) | Expr::Literal(..) => Ok(true),
263        // Timestamp-to-Int64 preserves its native unit and cannot fail. Keep
264        // the cast intact rather than assuming timestamps are nanoseconds.
265        Expr::Cast(cast) => Ok(cast.field.data_type() == &DataType::Int64
266            && matches!(cast.expr.as_ref(), Expr::Column(_))
267            && matches!(cast.expr.get_type(schema)?, DataType::Timestamp(..))),
268        _ => Ok(false),
269    }
270}
271
272/// A projected column can hide division, a fallible cast or a volatile UDF.
273/// PushDownFilter may inline it below the WHERE that made evaluation safe.
274fn columns_are_from_scan(expr: &Expr, input: &LogicalPlan) -> Result<bool> {
275    for column in expr.column_refs() {
276        if !column_is_from_scan(column, input)? {
277            return Ok(false);
278        }
279    }
280    Ok(true)
281}
282
283fn column_is_from_scan(column: &Column, input: &LogicalPlan) -> Result<bool> {
284    let index = input.schema().index_of_column(column)?;
285    match input {
286        LogicalPlan::TableScan(_) => Ok(true),
287        LogicalPlan::Filter(filter) => column_is_from_scan(column, &filter.input),
288        LogicalPlan::SubqueryAlias(alias) => column_at_is_from_scan(index, &alias.input),
289        LogicalPlan::Projection(projection) => {
290            let expr = match &projection.expr[index] {
291                Expr::Alias(alias) => alias.expr.as_ref(),
292                expr => expr,
293            };
294            Ok(safe_operand(expr, projection.input.schema())?
295                && columns_are_from_scan(expr, &projection.input)?)
296        }
297        LogicalPlan::Join(join) => {
298            match (
299                join.left.schema().has_column(column),
300                join.right.schema().has_column(column),
301            ) {
302                (true, false) => column_is_from_scan(column, &join.left),
303                (false, true) => column_is_from_scan(column, &join.right),
304                _ => Ok(false),
305            }
306        }
307        LogicalPlan::Union(union) => {
308            for input in &union.inputs {
309                if !column_at_is_from_scan(index, input)? {
310                    return Ok(false);
311                }
312            }
313            Ok(true)
314        }
315        // Unknown plan nodes may synthesize columns or hide their origin.
316        _ => Ok(false),
317    }
318}
319
320fn column_at_is_from_scan(index: usize, input: &LogicalPlan) -> Result<bool> {
321    let (qualifier, field) = input.schema().qualified_field(index);
322    column_is_from_scan(&Column::new(qualifier.cloned(), field.name()), input)
323}
324
325fn already_filtered(input: &LogicalPlan, bound: &Expr) -> bool {
326    match input {
327        LogicalPlan::Filter(filter) => {
328            split_conjunction(&filter.predicate).contains(&bound)
329                || already_filtered(&filter.input, bound)
330        }
331        LogicalPlan::TableScan(scan) => scan
332            .filters
333            .iter()
334            .any(|predicate| split_conjunction(predicate).contains(&bound)),
335        _ => false,
336    }
337}
338
339#[cfg(test)]
340mod tests {
341    use arrow::array::{
342        BooleanArray, DictionaryArray, Float64Array, Int64Array, StringArray,
343        TimestampMillisecondArray, TimestampNanosecondArray, UInt32Array,
344    };
345    use arrow::datatypes::UInt32Type;
346    use arrow::record_batch::RecordBatch;
347    use arrow::util::pretty::pretty_format_batches;
348    use datafusion::prelude::SessionContext;
349    use datafusion_optimizer::OptimizerContext;
350    use datafusion_optimizer::optimizer::Optimizer;
351    use datafusion_optimizer::push_down_filter::PushDownFilter;
352    use datafusion_optimizer::simplify_expressions::SimplifyExpressions;
353
354    use super::*;
355
356    fn context() -> SessionContext {
357        // Covers an earlier nonmatching row, both interval endpoints, NULL
358        // values/predicates/group keys, and groups with no matching rows.
359        let keys = vec![
360            Some("a"),
361            Some("a"),
362            Some("a"),
363            Some("a"),
364            Some("a"),
365            Some("b"),
366            Some("c"),
367            Some("c"),
368            None,
369            None,
370            None,
371        ];
372        let numbers = vec![
373            Some(5),
374            Some(10),
375            Some(20),
376            Some(25),
377            None,
378            Some(25),
379            Some(15),
380            Some(4),
381            Some(2),
382            Some(12),
383            Some(22),
384        ];
385        let flags = vec![
386            Some(false),
387            Some(true),
388            Some(true),
389            Some(true),
390            Some(true),
391            Some(true),
392            None,
393            Some(false),
394            Some(false),
395            Some(true),
396            Some(false),
397        ];
398        let services = DictionaryArray::<UInt32Type>::try_new(
399            UInt32Array::from(
400                flags
401                    .iter()
402                    .map(|flag| flag.map(|f| u32::from(!f)))
403                    .collect::<Vec<_>>(),
404            ),
405            Arc::new(StringArray::from(vec!["cart", "other"])),
406        )
407        .unwrap();
408        let batch = RecordBatch::try_from_iter(vec![
409            ("k", Arc::new(StringArray::from(keys)) as _),
410            ("service", Arc::new(services) as _),
411            ("v", Arc::new(Int64Array::from(numbers.clone())) as _),
412            (
413                "f",
414                Arc::new(Float64Array::from(vec![1.0; numbers.len()])) as _,
415            ),
416            (
417                "ts",
418                Arc::new(TimestampNanosecondArray::from(numbers.clone())) as _,
419            ),
420            (
421                "ts_ms",
422                Arc::new(TimestampMillisecondArray::from(numbers)) as _,
423            ),
424            ("flag", Arc::new(BooleanArray::from(flags)) as _),
425        ])
426        .unwrap();
427        let ctx = SessionContext::new();
428        ctx.register_batch("t", batch).unwrap();
429        ctx
430    }
431
432    async fn analyzed(ctx: &SessionContext, sql: &str) -> LogicalPlan {
433        let state = ctx.state();
434        let plan = state.create_logical_plan(sql).await.unwrap();
435        let plan = state
436            .analyzer()
437            .execute_and_check(plan, state.config_options(), |_, _| {})
438            .unwrap();
439        Optimizer::with_rules(vec![Arc::new(SimplifyExpressions::new())])
440            .optimize(plan, &OptimizerContext::new(), |_, _| {})
441            .unwrap()
442    }
443
444    fn optimize(plan: LogicalPlan) -> LogicalPlan {
445        Optimizer::with_rules(vec![Arc::new(InferAggregateInputBoundsRule)])
446            .optimize(plan, &OptimizerContext::new(), |_, _| {})
447            .unwrap()
448    }
449
450    fn aggregate_input_bound(plan: &LogicalPlan) -> Option<&Expr> {
451        if let LogicalPlan::Aggregate(aggregate) = plan
452            && let LogicalPlan::Filter(filter) = aggregate.input.as_ref()
453        {
454            return Some(&filter.predicate);
455        }
456        plan.inputs().into_iter().find_map(aggregate_input_bound)
457    }
458
459    async fn rows(ctx: &SessionContext, plan: LogicalPlan) -> String {
460        let batches = ctx
461            .execute_logical_plan(plan)
462            .await
463            .unwrap()
464            .collect()
465            .await
466            .unwrap();
467        pretty_format_batches(&batches).unwrap().to_string()
468    }
469
470    #[tokio::test]
471    async fn existential_having_preserves_extrema_and_conditional_sum() {
472        let ctx = context();
473        for (extremum, argument, condition, having, expected_bound) in [
474            (
475                "MIN",
476                "v",
477                "service = 'cart' AND v >= 10 AND v < 20",
478                "n > 0",
479                "t.v < Int64(20)",
480            ),
481            ("MIN", "v", "flag AND v <= 20", "n >= 1", "t.v <= Int64(20)"),
482            (
483                "MAX",
484                "v",
485                "flag AND v >= 10 AND v < 20",
486                "0 < n",
487                "t.v >= Int64(10)",
488            ),
489            ("MAX", "v", "flag AND 10 < v", "1 <= n", "t.v > Int64(10)"),
490            (
491                "MIN",
492                "CAST(ts AS BIGINT)",
493                "flag AND CAST(ts AS BIGINT) < 20",
494                "n > 0",
495                "CAST(t.ts AS Int64) < Int64(20)",
496            ),
497            (
498                "MIN",
499                "CAST(ts_ms AS BIGINT)",
500                "flag AND CAST(ts_ms AS BIGINT) < 20",
501                "n > 0",
502                "CAST(t.ts_ms AS Int64) < Int64(20)",
503            ),
504            (
505                "MIN",
506                "ts",
507                "flag AND ts < TIMESTAMP '1970-01-01 00:00:00.000000020'",
508                "n > 0",
509                "t.ts < TimestampNanosecond(20, None)",
510            ),
511        ] {
512            let sql = format!(
513                "SELECT k, {extremum}({argument}) AS endpoint,
514                SUM(CASE WHEN {condition} THEN 1 ELSE 0 END) AS n
515                FROM t GROUP BY k HAVING {having} ORDER BY k NULLS FIRST"
516            );
517            let original = analyzed(&ctx, &sql).await;
518            let optimized = optimize(original.clone());
519            assert_eq!(
520                aggregate_input_bound(&optimized).unwrap().to_string(),
521                expected_bound,
522                "{sql}"
523            );
524            assert_eq!(
525                optimized,
526                optimize(optimized.clone()),
527                "rule must be idempotent: {sql}"
528            );
529            assert_eq!(
530                rows(&ctx, original).await,
531                rows(&ctx, optimized).await,
532                "{sql}"
533            );
534        }
535    }
536
537    #[tokio::test]
538    async fn guarded_derived_operands_are_not_pruned() {
539        let ctx = SessionContext::new();
540        let batch = RecordBatch::try_from_iter(vec![
541            ("k", Arc::new(StringArray::from(vec!["a", "a", "b"])) as _),
542            ("v", Arc::new(Int64Array::from(vec![12, 5, 25])) as _),
543            ("d", Arc::new(Int64Array::from(vec![0, 1, 1])) as _),
544            (
545                "raw",
546                Arc::new(StringArray::from(vec!["bad", "5", "25"])) as _,
547            ),
548        ])
549        .unwrap();
550        ctx.register_batch("guarded", batch).unwrap();
551        ctx.sql("CREATE VIEW guarded_view AS SELECT k, v / d AS r FROM guarded WHERE abs(d) > 0")
552            .await
553            .unwrap()
554            .collect()
555            .await
556            .unwrap();
557        for source in [
558            "SELECT k, v / d AS r FROM guarded WHERE abs(d) > 0",
559            "SELECT k, quotient AS r FROM (SELECT k, v / d AS quotient FROM guarded WHERE abs(d) > 0) q",
560            "SELECT k, CAST(raw AS BIGINT) AS r FROM guarded WHERE abs(d) > 0",
561            "SELECT k, r FROM guarded_view",
562            "SELECT k, v AS r FROM guarded UNION ALL SELECT k, v / d AS r FROM guarded WHERE abs(d) > 0",
563            "SELECT l.k, l.r FROM guarded_view l JOIN guarded q ON l.k = q.k",
564        ] {
565            let sql = format!(
566                "SELECT k, MIN(r) FROM ({source}) s GROUP BY k
567                HAVING SUM(CASE WHEN r < 20 THEN 1 ELSE 0 END) > 0 ORDER BY k"
568            );
569            let original = analyzed(&ctx, &sql).await;
570            let optimized = optimize(original.clone());
571            assert_eq!(
572                rows(&ctx, original.clone()).await,
573                rows(&ctx, optimized.clone()).await,
574                "{sql}"
575            );
576            assert_eq!(original, optimized, "{sql}");
577        }
578    }
579
580    #[tokio::test]
581    async fn scan_column_aliases_still_allow_bounds() {
582        let ctx = context();
583        for source in [
584            "SELECT k, v AS r FROM t",
585            "SELECT k, renamed AS r FROM (SELECT k, v AS renamed FROM t) q",
586            "SELECT k, CAST(ts AS BIGINT) AS r FROM t",
587            "SELECT k, v AS r FROM t UNION ALL SELECT k, v AS r FROM t",
588            "SELECT l.k, l.v AS r FROM t l JOIN t r ON l.k = r.k",
589        ] {
590            let sql = format!(
591                "SELECT k, MIN(r) FROM ({source}) s GROUP BY k
592                HAVING SUM(CASE WHEN r < 20 THEN 1 ELSE 0 END) > 0 ORDER BY k NULLS FIRST"
593            );
594            let original = analyzed(&ctx, &sql).await;
595            let optimized = optimize(original.clone());
596            assert!(aggregate_input_bound(&optimized).is_some(), "{sql}");
597            assert_eq!(
598                rows(&ctx, original).await,
599                rows(&ctx, optimized).await,
600                "{sql}"
601            );
602        }
603    }
604
605    #[tokio::test]
606    async fn unsafe_or_unsupported_aggregates_are_not_pruned() {
607        let ctx = context();
608        let indicator = "SUM(CASE WHEN flag AND v < 20 THEN 1 ELSE 0 END)";
609        for sql in [
610            format!("SELECT k, MIN(v), COUNT(*) FROM t GROUP BY k HAVING {indicator}>0"),
611            format!("SELECT k, MIN(v), MAX(v) FROM t GROUP BY k HAVING {indicator}>0"),
612            format!("SELECT k, MIN(f) FROM t GROUP BY k HAVING {indicator}>0"),
613            format!("SELECT k, MAX(v) FROM t GROUP BY k HAVING {indicator}>0"),
614            format!("SELECT MIN(v) FROM t HAVING {indicator}>0"),
615            format!("SELECT k, MIN(v) FROM t GROUP BY ROLLUP(k) HAVING {indicator}>0"),
616            format!("SELECT k, MIN(v) FROM t GROUP BY k HAVING {indicator}=0"),
617            format!("SELECT k, MIN(v) FROM t GROUP BY k HAVING {indicator}>0 OR MIN(v)>0"),
618            format!("SELECT k, MIN(v) FILTER (WHERE flag) FROM t GROUP BY k HAVING {indicator}>0"),
619            "SELECT k, MIN(v) FROM t GROUP BY k HAVING SUM(CASE WHEN v<20 THEN 1 ELSE 2 END)>0".into(),
620            "SELECT k, MIN(v) FROM t GROUP BY k HAVING SUM(CASE WHEN v<20 THEN 1 END)>0".into(),
621            "SELECT k, MIN(v) FROM t GROUP BY k HAVING SUM(DISTINCT CASE WHEN v<20 THEN 1 ELSE 0 END)>0".into(),
622            "SELECT k, MIN(v) FROM t GROUP BY k HAVING SUM(CASE WHEN v<20 OR flag THEN 1 ELSE 0 END)>0".into(),
623            "SELECT k, MIN(v) FROM t GROUP BY k HAVING SUM(CASE WHEN random()>0.5 AND v<20 THEN 1 ELSE 0 END)>0".into(),
624            "SELECT k, MIN(CAST(k AS BIGINT)) FROM t GROUP BY k HAVING SUM(CASE WHEN CAST(k AS BIGINT)<20 THEN 1 ELSE 0 END)>0".into(),
625            "SELECT k, MIN(v) FROM t GROUP BY k HAVING SUM(CASE WHEN v<20 THEN 1 WHEN flag THEN 1 ELSE 0 END)>0".into(),
626            "SELECT k, MIN(f) FROM t GROUP BY k HAVING SUM(CASE WHEN f<20.0 THEN 1 ELSE 0 END)>0".into(),
627            "SELECT k, MIN(v) FROM t GROUP BY k HAVING SUM(CASE WHEN v<20 THEN 1 ELSE 0 END)>0 AND SUM(CASE WHEN v<100 THEN 1 ELSE 0 END)>0".into(),
628        ] {
629            let original = analyzed(&ctx, &sql).await;
630            assert_eq!(original, optimize(original.clone()), "{sql}");
631        }
632    }
633
634    #[tokio::test]
635    async fn nonpositive_thresholds_preserve_zero_match_groups() {
636        let ctx = context();
637        let expected = [
638            "+---+----------+",
639            "| k | endpoint |",
640            "+---+----------+",
641            "|   | 2        |",
642            "| a | 5        |",
643            "| b | 25       |",
644            "| c | 4        |",
645            "+---+----------+",
646        ]
647        .join("\n");
648        for threshold in [">= 0", "> -1"] {
649            let sql = format!(
650                "SELECT k, MIN(v) AS endpoint FROM t GROUP BY k
651                HAVING SUM(CASE WHEN flag AND v < 20 THEN 1 ELSE 0 END) {threshold}
652                ORDER BY k NULLS FIRST"
653            );
654            let original = analyzed(&ctx, &sql).await;
655            let optimized = optimize(original.clone());
656            assert_eq!(original, optimized, "{sql}");
657            assert_eq!(expected, rows(&ctx, optimized).await, "{sql}");
658        }
659    }
660
661    #[tokio::test]
662    async fn bound_survives_filter_pushdown_without_being_added_again() {
663        let ctx = context();
664        let original = analyzed(
665            &ctx,
666            "SELECT k, MIN(v) FROM t GROUP BY k
667            HAVING SUM(CASE WHEN flag AND v<20 THEN 1 ELSE 0 END)>0",
668        )
669        .await;
670        let optimizer = Optimizer::with_rules(vec![
671            Arc::new(InferAggregateInputBoundsRule),
672            Arc::new(PushDownFilter::new()),
673        ]);
674        let optimized = optimizer
675            .optimize(original, &OptimizerContext::new(), |_, _| {})
676            .unwrap();
677        let repeated = optimizer
678            .optimize(optimized.clone(), &OptimizerContext::new(), |_, _| {})
679            .unwrap();
680        assert_eq!(optimized, repeated);
681        assert!(aggregate_input_bound(&optimized).is_some());
682    }
683}
query/optimizer/infer_aggregate_input_bounds.rs

query/optimizer/
infer_aggregate_input_bounds.rs
