mz_sql/plan/
plan_utils.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//! Helper code used throughout the planner.
11
12use std::fmt;
13
14use mz_repr::RelationDesc;
15
16use crate::ast::Ident;
17use crate::normalize;
18use crate::plan::PlanError;
19use crate::plan::query::SelectOptionExtracted;
20
21/// Renames the columns in `desc` with the names in `column_names` if
22/// `column_names` is non-empty.
23///
24/// Returns an error if the length of `column_names` is greater than the arity
25/// of `desc`.
26pub fn maybe_rename_columns(
27    context: impl fmt::Display,
28    desc: &mut RelationDesc,
29    column_names: &[Ident],
30) -> Result<(), PlanError> {
31    if column_names.len() > desc.typ().column_types.len() {
32        sql_bail!(
33            "{0} definition names {1} column{2}, but {0} has {3} column{4}",
34            context,
35            column_names.len(),
36            if column_names.len() == 1 { "" } else { "s" },
37            desc.typ().column_types.len(),
38            if desc.typ().column_types.len() == 1 {
39                ""
40            } else {
41                "s"
42            },
43        )
44    }
45
46    for (i, name) in column_names.iter().enumerate() {
47        *desc.get_name_mut(i) = normalize::column_name(name.clone());
48    }
49
50    Ok(())
51}
52
53/// Specifies the side of a join.
54///
55/// Intended for use in error messages.
56#[derive(Debug, Clone, Copy)]
57pub enum JoinSide {
58    /// The left side.
59    Left,
60    /// The right side.
61    Right,
62}
63
64impl fmt::Display for JoinSide {
65    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
66        match self {
67            JoinSide::Left => f.write_str("left"),
68            JoinSide::Right => f.write_str("right"),
69        }
70    }
71}
72
73/// Specifies a bundle of group size query hints.
74///
75/// This struct bridges from old to new syntax for group size query hints,
76/// making it easier to pass these hints along and make use of a group size
77/// hint configuration.
78#[derive(Debug, Default, Clone, Copy)]
79pub struct GroupSizeHints {
80    pub aggregate_input_group_size: Option<u64>,
81    pub distinct_on_input_group_size: Option<u64>,
82    pub limit_input_group_size: Option<u64>,
83}
84
85impl TryFrom<SelectOptionExtracted> for GroupSizeHints {
86    type Error = PlanError;
87
88    /// Creates group size hints from extracted `SELECT` `OPTIONS` validating that
89    /// either the old `EXPECTED GROUP SIZE` syntax was used or alternatively the
90    /// new syntax with `AGGREGATE INPUT GROUP SIZE`, `DISTINCT ON INPUT GROUP SIZE`,
91    /// and `LIMIT INPUT GROUP SIZE`. If the two syntax versions are mixed in the
92    /// same `OPTIONS` clause, an error is returned.[^1]
93    /// [^1] <https://github.com/MaterializeInc/materialize/blob/main/doc/developer/design/20230829_topk_size_hint.md>
94    fn try_from(select_option_extracted: SelectOptionExtracted) -> Result<Self, Self::Error> {
95        let SelectOptionExtracted {
96            expected_group_size,
97            aggregate_input_group_size,
98            distinct_on_input_group_size,
99            limit_input_group_size,
100            ..
101        } = select_option_extracted;
102        if expected_group_size.is_some()
103            && (aggregate_input_group_size.is_some()
104                || distinct_on_input_group_size.is_some()
105                || limit_input_group_size.is_some())
106        {
107            Err(PlanError::InvalidGroupSizeHints)
108        } else {
109            let aggregate_input_group_size = aggregate_input_group_size.or(expected_group_size);
110            let distinct_on_input_group_size = distinct_on_input_group_size.or(expected_group_size);
111            let limit_input_group_size = limit_input_group_size.or(expected_group_size);
112            Ok(GroupSizeHints {
113                aggregate_input_group_size,
114                distinct_on_input_group_size,
115                limit_input_group_size,
116            })
117        }
118    }
119}