mz_transform/canonicalization/
topk_elision.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10//! Remove TopK operators with both an offset of zero and no limit.
11
12use mz_expr::MirRelationExpr;
13use mz_expr::visit::Visit;
14
15use crate::TransformCtx;
16
17/// Remove TopK operators with both an offset of zero and no limit.
18#[derive(Debug)]
19pub struct TopKElision;
20
21impl crate::Transform for TopKElision {
22    fn name(&self) -> &'static str {
23        "TopKElision"
24    }
25
26    #[mz_ore::instrument(
27        target = "optimizer",
28        level = "debug",
29        fields(path.segment = "topk_elision")
30    )]
31    fn actually_perform_transform(
32        &self,
33        relation: &mut MirRelationExpr,
34        _: &mut TransformCtx,
35    ) -> Result<(), crate::TransformError> {
36        relation.visit_mut_post(&mut Self::action)?;
37        mz_repr::explain::trace_plan(&*relation);
38        Ok(())
39    }
40}
41
42impl TopKElision {
43    /// Remove TopK operators with both an offset of zero and no limit.
44    pub fn action(relation: &mut MirRelationExpr) {
45        if let MirRelationExpr::TopK {
46            input,
47            group_key: _,
48            order_key: _,
49            limit,
50            offset,
51            monotonic: _,
52            expected_group_size: _,
53        } = relation
54        {
55            // The limit is not set if it either `None` or literal `Null`.
56            if limit.as_ref().map_or(true, |l| l.is_literal_null()) && *offset == 0 {
57                *relation = input.take_dangerous();
58            } else if limit.as_ref().and_then(|l| l.as_literal_int64()) == Some(0) {
59                relation.take_safely(None);
60            }
61        }
62    }
63}