risingwave_expr/aggregate/
def.rs

1// Copyright 2023 RisingWave Labs
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
15//! Aggregation function definitions.
16
17use std::fmt::Display;
18use std::iter::Peekable;
19use std::str::FromStr;
20
21use anyhow::Context;
22use enum_as_inner::EnumAsInner;
23use itertools::Itertools;
24use risingwave_common::bail;
25use risingwave_common::types::{DataType, Datum};
26use risingwave_common::util::sort_util::{ColumnOrder, OrderType};
27use risingwave_common::util::value_encoding::DatumFromProtoExt;
28pub use risingwave_pb::expr::agg_call::PbKind as PbAggKind;
29use risingwave_pb::expr::{
30    PbAggCall, PbAggType, PbExprNode, PbInputRef, PbUserDefinedFunctionMetadata,
31};
32
33use crate::Result;
34use crate::expr::{BoxedExpression, ExpectExt, LiteralExpression, Token, build_from_prost};
35
36/// Represents an aggregation function.
37// TODO(runji):
38//  remove this struct from the expression module.
39//  this module only cares about aggregate functions themselves.
40//  advanced features like order by, filter, distinct, etc. should be handled by the upper layer.
41#[derive(Debug, Clone)]
42pub struct AggCall {
43    /// Aggregation type for constructing agg state.
44    pub agg_type: AggType,
45
46    /// Arguments of aggregation function input.
47    pub args: AggArgs,
48
49    /// The return type of aggregation function.
50    pub return_type: DataType,
51
52    /// Order requirements specified in order by clause of agg call
53    pub column_orders: Vec<ColumnOrder>,
54
55    /// Filter of aggregation.
56    pub filter: Option<BoxedExpression>,
57
58    /// Should deduplicate the input before aggregation.
59    pub distinct: bool,
60
61    /// Constant arguments.
62    pub direct_args: Vec<LiteralExpression>,
63}
64
65impl AggCall {
66    pub fn from_protobuf(agg_call: &PbAggCall) -> Result<Self> {
67        let agg_type = AggType::from_protobuf_flatten(
68            agg_call.get_kind()?,
69            agg_call.udf.as_ref(),
70            agg_call.scalar.as_ref(),
71        )?;
72        let args = AggArgs::from_protobuf(agg_call.get_args())?;
73        let column_orders = agg_call
74            .get_order_by()
75            .iter()
76            .map(|col_order| {
77                let col_idx = col_order.get_column_index() as usize;
78                let order_type = OrderType::from_protobuf(col_order.get_order_type().unwrap());
79                ColumnOrder::new(col_idx, order_type)
80            })
81            .collect();
82        let filter = match agg_call.filter {
83            Some(ref pb_filter) => Some(build_from_prost(pb_filter)?), /* TODO: non-strict filter in streaming */
84            None => None,
85        };
86        let direct_args = agg_call
87            .direct_args
88            .iter()
89            .map(|arg| {
90                let data_type = DataType::from(arg.get_type().unwrap());
91                LiteralExpression::new(
92                    data_type.clone(),
93                    Datum::from_protobuf(arg.get_datum().unwrap(), &data_type).unwrap(),
94                )
95            })
96            .collect_vec();
97        Ok(AggCall {
98            agg_type,
99            args,
100            return_type: DataType::from(agg_call.get_return_type()?),
101            column_orders,
102            filter,
103            distinct: agg_call.distinct,
104            direct_args,
105        })
106    }
107
108    /// Build an `AggCall` from a string.
109    ///
110    /// # Syntax
111    ///
112    /// ```text
113    /// (<name>:<type> [<index>:<type>]* [distinct] [orderby [<index>:<asc|desc>]*])
114    /// ```
115    pub fn from_pretty(s: impl AsRef<str>) -> Self {
116        let tokens = crate::expr::lexer(s.as_ref());
117        Parser::new(tokens.into_iter()).parse_aggregation()
118    }
119
120    pub fn with_filter(mut self, filter: BoxedExpression) -> Self {
121        self.filter = Some(filter);
122        self
123    }
124}
125
126struct Parser<Iter: Iterator> {
127    tokens: Peekable<Iter>,
128}
129
130impl<Iter: Iterator<Item = Token>> Parser<Iter> {
131    fn new(tokens: Iter) -> Self {
132        Self {
133            tokens: tokens.peekable(),
134        }
135    }
136
137    fn parse_aggregation(&mut self) -> AggCall {
138        assert_eq!(self.tokens.next(), Some(Token::LParen), "Expected a (");
139        let func = self.parse_function();
140        assert_eq!(self.tokens.next(), Some(Token::Colon), "Expected a Colon");
141        let ty = self.parse_type();
142
143        let mut distinct = false;
144        let mut children = Vec::new();
145        let mut column_orders = Vec::new();
146        while matches!(self.tokens.peek(), Some(Token::Index(_))) {
147            children.push(self.parse_arg());
148        }
149        if matches!(self.tokens.peek(), Some(Token::Literal(s)) if s == "distinct") {
150            distinct = true;
151            self.tokens.next(); // Consume
152        }
153        if matches!(self.tokens.peek(), Some(Token::Literal(s)) if s == "orderby") {
154            self.tokens.next(); // Consume
155            while matches!(self.tokens.peek(), Some(Token::Index(_))) {
156                column_orders.push(self.parse_orderkey());
157            }
158        }
159        self.tokens.next(); // Consume the RParen
160
161        AggCall {
162            agg_type: AggType::from_protobuf_flatten(func, None, None).unwrap(),
163            args: AggArgs {
164                data_types: children.iter().map(|(_, ty)| ty.clone()).collect(),
165                val_indices: children.iter().map(|(idx, _)| *idx).collect(),
166            },
167            return_type: ty,
168            column_orders,
169            filter: None,
170            distinct,
171            direct_args: Vec::new(),
172        }
173    }
174
175    fn parse_type(&mut self) -> DataType {
176        match self.tokens.next().expect("Unexpected end of input") {
177            Token::Literal(name) => name.parse::<DataType>().expect_str("type", &name),
178            t => panic!("Expected a Literal, got {t:?}"),
179        }
180    }
181
182    fn parse_arg(&mut self) -> (usize, DataType) {
183        let idx = match self.tokens.next().expect("Unexpected end of input") {
184            Token::Index(idx) => idx,
185            t => panic!("Expected an Index, got {t:?}"),
186        };
187        assert_eq!(self.tokens.next(), Some(Token::Colon), "Expected a Colon");
188        let ty = self.parse_type();
189        (idx, ty)
190    }
191
192    fn parse_function(&mut self) -> PbAggKind {
193        match self.tokens.next().expect("Unexpected end of input") {
194            Token::Literal(name) => {
195                PbAggKind::from_str_name(&name.to_uppercase()).expect_str("function", &name)
196            }
197            t => panic!("Expected a Literal, got {t:?}"),
198        }
199    }
200
201    fn parse_orderkey(&mut self) -> ColumnOrder {
202        let idx = match self.tokens.next().expect("Unexpected end of input") {
203            Token::Index(idx) => idx,
204            t => panic!("Expected an Index, got {t:?}"),
205        };
206        assert_eq!(self.tokens.next(), Some(Token::Colon), "Expected a Colon");
207        let order = match self.tokens.next().expect("Unexpected end of input") {
208            Token::Literal(s) if s == "asc" => OrderType::ascending(),
209            Token::Literal(s) if s == "desc" => OrderType::descending(),
210            t => panic!("Expected asc or desc, got {t:?}"),
211        };
212        ColumnOrder::new(idx, order)
213    }
214}
215
216/// Aggregate function kind.
217#[derive(Debug, Clone, PartialEq, Eq, Hash, EnumAsInner)]
218pub enum AggType {
219    /// Built-in aggregate function.
220    ///
221    /// The associated value should not be `UserDefined` or `WrapScalar`.
222    Builtin(PbAggKind),
223
224    /// User defined aggregate function.
225    UserDefined(PbUserDefinedFunctionMetadata),
226
227    /// Wrap a scalar function that takes a list as input as an aggregation function.
228    WrapScalar(PbExprNode),
229}
230
231impl Display for AggType {
232    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
233        match self {
234            Self::Builtin(kind) => write!(f, "{}", kind.as_str_name().to_lowercase()),
235            Self::UserDefined(_) => write!(f, "udaf"),
236            Self::WrapScalar(_) => write!(f, "wrap_scalar"),
237        }
238    }
239}
240
241/// `FromStr` for builtin aggregate functions.
242impl FromStr for AggType {
243    type Err = ();
244
245    fn from_str(s: &str) -> Result<Self, Self::Err> {
246        let kind = PbAggKind::from_str(s)?;
247        Ok(AggType::Builtin(kind))
248    }
249}
250
251impl From<PbAggKind> for AggType {
252    fn from(pb: PbAggKind) -> Self {
253        assert!(!matches!(
254            pb,
255            PbAggKind::Unspecified | PbAggKind::UserDefined | PbAggKind::WrapScalar
256        ));
257        AggType::Builtin(pb)
258    }
259}
260
261impl AggType {
262    pub fn from_protobuf_flatten(
263        pb_kind: PbAggKind,
264        user_defined: Option<&PbUserDefinedFunctionMetadata>,
265        scalar: Option<&PbExprNode>,
266    ) -> Result<Self> {
267        match pb_kind {
268            PbAggKind::UserDefined => {
269                let user_defined = user_defined.context("expect user defined")?;
270                Ok(AggType::UserDefined(user_defined.clone()))
271            }
272            PbAggKind::WrapScalar => {
273                let scalar = scalar.context("expect scalar")?;
274                Ok(AggType::WrapScalar(scalar.clone()))
275            }
276            PbAggKind::Unspecified => bail!("Unrecognized agg."),
277            _ => Ok(AggType::Builtin(pb_kind)),
278        }
279    }
280
281    pub fn to_protobuf_simple(&self) -> PbAggKind {
282        match self {
283            Self::Builtin(pb) => *pb,
284            Self::UserDefined(_) => PbAggKind::UserDefined,
285            Self::WrapScalar(_) => PbAggKind::WrapScalar,
286        }
287    }
288
289    pub fn from_protobuf(pb_type: &PbAggType) -> Result<Self> {
290        match PbAggKind::try_from(pb_type.kind).context("no such aggregate function type")? {
291            PbAggKind::Unspecified => bail!("Unrecognized agg."),
292            PbAggKind::UserDefined => Ok(AggType::UserDefined(pb_type.get_udf_meta()?.clone())),
293            PbAggKind::WrapScalar => Ok(AggType::WrapScalar(pb_type.get_scalar_expr()?.clone())),
294            kind => Ok(AggType::Builtin(kind)),
295        }
296    }
297
298    pub fn to_protobuf(&self) -> PbAggType {
299        match self {
300            Self::Builtin(kind) => PbAggType {
301                kind: *kind as _,
302                udf_meta: None,
303                scalar_expr: None,
304            },
305            Self::UserDefined(udf_meta) => PbAggType {
306                kind: PbAggKind::UserDefined as _,
307                udf_meta: Some(udf_meta.clone()),
308                scalar_expr: None,
309            },
310            Self::WrapScalar(scalar_expr) => PbAggType {
311                kind: PbAggKind::WrapScalar as _,
312                udf_meta: None,
313                scalar_expr: Some(scalar_expr.clone()),
314            },
315        }
316    }
317}
318
319/// Macros to generate match arms for `AggType`.
320/// IMPORTANT: These macros must be carefully maintained especially when adding new
321/// `AggType`/`PbAggKind` variants.
322pub mod agg_types {
323    /// [`AggType`](super::AggType)s that should've been rewritten to other kinds. These kinds
324    /// should not appear when generating physical plan nodes.
325    #[macro_export]
326    macro_rules! rewritten {
327        () => {
328            AggType::Builtin(
329                PbAggKind::Avg
330                    | PbAggKind::StddevPop
331                    | PbAggKind::StddevSamp
332                    | PbAggKind::VarPop
333                    | PbAggKind::VarSamp
334                    | PbAggKind::Grouping
335                    // ApproxPercentile always uses custom agg executors,
336                    // rather than an aggregation operator
337                    | PbAggKind::ApproxPercentile
338                    | PbAggKind::ArgMin
339                    | PbAggKind::ArgMax
340            )
341        };
342    }
343    pub use rewritten;
344
345    /// [`AggType`](super::AggType)s of which the aggregate results are not affected by the
346    /// user given ORDER BY clause.
347    #[macro_export]
348    macro_rules! result_unaffected_by_order_by {
349        () => {
350            AggType::Builtin(PbAggKind::BitAnd
351                | PbAggKind::BitOr
352                | PbAggKind::BitXor // XOR is commutative and associative
353                | PbAggKind::BoolAnd
354                | PbAggKind::BoolOr
355                | PbAggKind::Min
356                | PbAggKind::Max
357                | PbAggKind::Sum
358                | PbAggKind::Sum0
359                | PbAggKind::Count
360                | PbAggKind::Avg
361                | PbAggKind::ApproxCountDistinct
362                | PbAggKind::VarPop
363                | PbAggKind::VarSamp
364                | PbAggKind::StddevPop
365                | PbAggKind::StddevSamp)
366        };
367    }
368    pub use result_unaffected_by_order_by;
369
370    /// [`AggType`](super::AggType)s that must be called with ORDER BY clause. These are
371    /// slightly different from variants not in [`result_unaffected_by_order_by`], in that
372    /// variants returned by this macro should be banned while the others should just be warned.
373    #[macro_export]
374    macro_rules! must_have_order_by {
375        () => {
376            AggType::Builtin(
377                PbAggKind::FirstValue
378                    | PbAggKind::LastValue
379                    | PbAggKind::PercentileCont
380                    | PbAggKind::PercentileDisc
381                    | PbAggKind::Mode,
382            )
383        };
384    }
385    pub use must_have_order_by;
386
387    /// [`AggType`](super::AggType)s of which the aggregate results are not affected by the
388    /// user given DISTINCT keyword.
389    #[macro_export]
390    macro_rules! result_unaffected_by_distinct {
391        () => {
392            AggType::Builtin(
393                PbAggKind::BitAnd
394                    | PbAggKind::BitOr
395                    | PbAggKind::BoolAnd
396                    | PbAggKind::BoolOr
397                    | PbAggKind::Min
398                    | PbAggKind::Max
399                    | PbAggKind::ApproxCountDistinct,
400            )
401        };
402    }
403    pub use result_unaffected_by_distinct;
404
405    /// [`AggType`](crate::aggregate::AggType)s that are simply cannot 2-phased.
406    #[macro_export]
407    macro_rules! simply_cannot_two_phase {
408        () => {
409            AggType::Builtin(
410                PbAggKind::StringAgg
411                    | PbAggKind::ApproxCountDistinct
412                    | PbAggKind::ArrayAgg
413                    | PbAggKind::JsonbAgg
414                    | PbAggKind::JsonbObjectAgg
415                    | PbAggKind::FirstValue
416                    | PbAggKind::LastValue
417                    | PbAggKind::PercentileCont
418                    | PbAggKind::PercentileDisc
419                    | PbAggKind::Mode
420                    // FIXME(wrj): move `BoolAnd` and `BoolOr` out
421                    //  after we support general merge in stateless_simple_agg
422                    | PbAggKind::BoolAnd
423                    | PbAggKind::BoolOr
424                    | PbAggKind::BitAnd
425                    | PbAggKind::BitOr
426            )
427            | AggType::UserDefined(_)
428            | AggType::WrapScalar(_)
429        };
430    }
431    pub use simply_cannot_two_phase;
432
433    /// [`AggType`](super::AggType)s that are implemented with a single value state (so-called
434    /// stateless).
435    #[macro_export]
436    macro_rules! single_value_state {
437        () => {
438            AggType::Builtin(
439                PbAggKind::Sum
440                    | PbAggKind::Sum0
441                    | PbAggKind::Count
442                    | PbAggKind::BitAnd
443                    | PbAggKind::BitOr
444                    | PbAggKind::BitXor
445                    | PbAggKind::BoolAnd
446                    | PbAggKind::BoolOr
447                    | PbAggKind::ApproxCountDistinct
448                    | PbAggKind::InternalLastSeenValue
449                    | PbAggKind::ApproxPercentile,
450            ) | AggType::UserDefined(_)
451        };
452    }
453    pub use single_value_state;
454
455    /// [`AggType`](super::AggType)s that are implemented with a single value state (so-called
456    /// stateless) iff the input is append-only.
457    #[macro_export]
458    macro_rules! single_value_state_iff_in_append_only {
459        () => {
460            AggType::Builtin(PbAggKind::Max | PbAggKind::Min)
461        };
462    }
463    pub use single_value_state_iff_in_append_only;
464
465    /// [`AggType`](super::AggType)s that are implemented with a materialized input state.
466    #[macro_export]
467    macro_rules! materialized_input_state {
468        () => {
469            AggType::Builtin(
470                PbAggKind::Min
471                    | PbAggKind::Max
472                    | PbAggKind::FirstValue
473                    | PbAggKind::LastValue
474                    | PbAggKind::StringAgg
475                    | PbAggKind::ArrayAgg
476                    | PbAggKind::JsonbAgg
477                    | PbAggKind::JsonbObjectAgg
478                    | PbAggKind::PercentileCont
479                    | PbAggKind::PercentileDisc
480                    | PbAggKind::Mode,
481            ) | AggType::WrapScalar(_)
482        };
483    }
484    pub use materialized_input_state;
485
486    /// Ordered-set aggregate functions.
487    #[macro_export]
488    macro_rules! ordered_set {
489        () => {
490            AggType::Builtin(
491                PbAggKind::PercentileCont
492                    | PbAggKind::PercentileDisc
493                    | PbAggKind::Mode
494                    | PbAggKind::ApproxPercentile,
495            )
496        };
497    }
498    pub use ordered_set;
499}
500
501impl AggType {
502    /// Get the total phase agg kind from the partial phase agg kind.
503    pub fn partial_to_total(&self) -> Option<Self> {
504        match self {
505            AggType::Builtin(
506                PbAggKind::BitXor
507                | PbAggKind::Min
508                | PbAggKind::Max
509                | PbAggKind::Sum
510                | PbAggKind::InternalLastSeenValue,
511            ) => Some(self.clone()),
512            AggType::Builtin(PbAggKind::Sum0 | PbAggKind::Count) => {
513                Some(Self::Builtin(PbAggKind::Sum0))
514            }
515            agg_types::simply_cannot_two_phase!() => None,
516            agg_types::rewritten!() => None,
517            // invalid variants
518            AggType::Builtin(
519                PbAggKind::Unspecified | PbAggKind::UserDefined | PbAggKind::WrapScalar,
520            ) => None,
521        }
522    }
523}
524
525/// An aggregation function may accept 0, 1 or 2 arguments.
526#[derive(Clone, Debug, Default)]
527pub struct AggArgs {
528    data_types: Box<[DataType]>,
529    val_indices: Box<[usize]>,
530}
531
532impl AggArgs {
533    pub fn from_protobuf(args: &[PbInputRef]) -> Result<Self> {
534        Ok(AggArgs {
535            data_types: args
536                .iter()
537                .map(|arg| DataType::from(arg.get_type().unwrap()))
538                .collect(),
539            val_indices: args.iter().map(|arg| arg.get_index() as usize).collect(),
540        })
541    }
542
543    /// return the types of arguments.
544    pub fn arg_types(&self) -> &[DataType] {
545        &self.data_types
546    }
547
548    /// return the indices of the arguments in [`risingwave_common::array::StreamChunk`].
549    pub fn val_indices(&self) -> &[usize] {
550        &self.val_indices
551    }
552}
553
554impl FromIterator<(DataType, usize)> for AggArgs {
555    fn from_iter<T: IntoIterator<Item = (DataType, usize)>>(iter: T) -> Self {
556        let (data_types, val_indices): (Vec<_>, Vec<_>) = iter.into_iter().unzip();
557        AggArgs {
558            data_types: data_types.into(),
559            val_indices: val_indices.into(),
560        }
561    }
562}
risingwave_expr/aggregate/def.rs

risingwave_expr/aggregate/
def.rs
