Skip to main content

mz_expr/relation/
func.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#![allow(missing_docs)]
11
12use std::cmp::{max, min};
13use std::iter::Sum;
14use std::ops::Deref;
15use std::str::FromStr;
16use std::{fmt, iter};
17
18use chrono::{DateTime, NaiveDateTime, NaiveTime, Utc};
19use dec::OrderedDecimal;
20use itertools::{Either, Itertools};
21use mz_ore::cast::CastFrom;
22
23use mz_ore::str::separated;
24use mz_ore::{soft_assert_eq_no_log, soft_assert_or_log};
25use mz_repr::adt::array::ArrayDimension;
26use mz_repr::adt::date::Date;
27use mz_repr::adt::interval::Interval;
28use mz_repr::adt::numeric::{self, Numeric, NumericMaxScale};
29use mz_repr::adt::regex::{Regex as ReprRegex, RegexCompilationError};
30use mz_repr::adt::timestamp::{CheckedTimestamp, TimestampLike};
31use mz_repr::{
32    ColumnName, Datum, Diff, ReprColumnType, ReprRelationType, Row, RowArena, RowPacker, SharedRow,
33    SqlColumnType, SqlRelationType, SqlScalarType, datum_size,
34};
35use num::{CheckedAdd, Integer, Signed, ToPrimitive};
36use ordered_float::OrderedFloat;
37use regex::Regex;
38use serde::{Deserialize, Serialize};
39use smallvec::SmallVec;
40
41use crate::EvalError;
42use crate::WindowFrameBound::{
43    CurrentRow, OffsetFollowing, OffsetPreceding, UnboundedFollowing, UnboundedPreceding,
44};
45use crate::WindowFrameUnits::{Groups, Range, Rows};
46use crate::explain::{HumanizedExpr, HumanizerMode};
47use crate::relation::{
48    ColumnOrder, WindowFrame, WindowFrameBound, WindowFrameUnits, compare_columns,
49};
50use crate::scalar::func::{add_timestamp_months, jsonb_stringify};
51
52// TODO(jamii) be careful about overflow in sum/avg
53// see https://timely.zulipchat.com/#narrow/stream/186635-engineering/topic/additional.20work/near/163507435
54
55fn max_string<'a, I>(datums: I) -> Datum<'a>
56where
57    I: IntoIterator<Item = Datum<'a>>,
58{
59    match datums
60        .into_iter()
61        .filter(|d| !d.is_null())
62        .max_by(|a, b| a.unwrap_str().cmp(b.unwrap_str()))
63    {
64        Some(datum) => datum,
65        None => Datum::Null,
66    }
67}
68
69fn max_datum<'a, I, DatumType>(datums: I) -> Datum<'a>
70where
71    I: IntoIterator<Item = Datum<'a>>,
72    DatumType: TryFrom<Datum<'a>> + Ord,
73    <DatumType as TryFrom<Datum<'a>>>::Error: std::fmt::Debug,
74    Datum<'a>: From<Option<DatumType>>,
75{
76    let x: Option<DatumType> = datums
77        .into_iter()
78        .filter(|d| !d.is_null())
79        .map(|d| DatumType::try_from(d).expect("unexpected type"))
80        .max();
81
82    x.into()
83}
84
85fn min_datum<'a, I, DatumType>(datums: I) -> Datum<'a>
86where
87    I: IntoIterator<Item = Datum<'a>>,
88    DatumType: TryFrom<Datum<'a>> + Ord,
89    <DatumType as TryFrom<Datum<'a>>>::Error: std::fmt::Debug,
90    Datum<'a>: From<Option<DatumType>>,
91{
92    let x: Option<DatumType> = datums
93        .into_iter()
94        .filter(|d| !d.is_null())
95        .map(|d| DatumType::try_from(d).expect("unexpected type"))
96        .min();
97
98    x.into()
99}
100
101fn min_string<'a, I>(datums: I) -> Datum<'a>
102where
103    I: IntoIterator<Item = Datum<'a>>,
104{
105    match datums
106        .into_iter()
107        .filter(|d| !d.is_null())
108        .min_by(|a, b| a.unwrap_str().cmp(b.unwrap_str()))
109    {
110        Some(datum) => datum,
111        None => Datum::Null,
112    }
113}
114
115fn sum_datum<'a, I, DatumType, ResultType>(datums: I) -> Datum<'a>
116where
117    I: IntoIterator<Item = Datum<'a>>,
118    DatumType: TryFrom<Datum<'a>>,
119    <DatumType as TryFrom<Datum<'a>>>::Error: std::fmt::Debug,
120    ResultType: From<DatumType> + Sum + Into<Datum<'a>>,
121{
122    let mut datums = datums.into_iter().filter(|d| !d.is_null()).peekable();
123    if datums.peek().is_none() {
124        Datum::Null
125    } else {
126        let x = datums
127            .map(|d| ResultType::from(DatumType::try_from(d).expect("unexpected type")))
128            .sum::<ResultType>();
129        x.into()
130    }
131}
132
133/// Count-aware signed-integer sum. Accumulates `Σ value·diff` in `i128`, which
134/// matches the width of the dataflow's `Accum::SimpleNumber` accumulator (see
135/// `build_accumulable` and `finalize_accum` in `mz_compute::render::reduce`);
136/// `narrow` then reproduces that variant's `finalize_accum` arm. Unlike
137/// `expand_counts`, this consumes the multiplicity directly, so it is linear in
138/// the number of distinct values and correct for negative diffs (retractions),
139/// which `expand_counts` would silently drop.
140///
141/// Returns `Datum::Null` when no non-null value was accumulated, matching
142/// `finalize_accum`'s null handling: its `is_zero` check on `SimpleNumber`
143/// requires both a zero running sum and a zero non-null count.
144fn sum_signed_int_counted<'a, I, N>(datums: I, narrow: N) -> Datum<'a>
145where
146    I: IntoIterator<Item = (Datum<'a>, Diff)>,
147    N: FnOnce(i128) -> Datum<'a>,
148{
149    let mut accum: i128 = 0;
150    let mut non_nulls = Diff::ZERO;
151    for (datum, diff) in datums {
152        if datum.is_null() {
153            continue;
154        }
155        let value = match datum {
156            Datum::Int16(i) => i128::from(i),
157            Datum::Int32(i) => i128::from(i),
158            Datum::Int64(i) => i128::from(i),
159            other => panic!("unexpected non-integer datum in signed sum: {other:?}"),
160        };
161        // The dataflow accumulates `value * diff` in an `Overflowing<i128>`; we
162        // mirror that. Genuine i128 overflow would require summands far beyond
163        // any realistic input, so wrapping matches the dataflow's production
164        // behavior.
165        accum = accum.wrapping_add(value.wrapping_mul(i128::from(diff.into_inner())));
166        non_nulls += diff;
167    }
168    if accum == 0 && non_nulls.is_zero() {
169        Datum::Null
170    } else {
171        narrow(accum)
172    }
173}
174
175fn sum_numeric<'a, I>(datums: I) -> Datum<'a>
176where
177    I: IntoIterator<Item = Datum<'a>>,
178{
179    let mut cx = numeric::cx_datum();
180    let mut sum = Numeric::zero();
181    let mut empty = true;
182    for d in datums {
183        if !d.is_null() {
184            empty = false;
185            cx.add(&mut sum, &d.unwrap_numeric().0);
186        }
187    }
188    match empty {
189        true => Datum::Null,
190        false => Datum::from(sum),
191    }
192}
193
194fn count<'a, I>(datums: I) -> Datum<'a>
195where
196    I: IntoIterator<Item = (Datum<'a>, Diff)>,
197{
198    // Count is accumulable: rather than expand each `(datum, diff)` into `diff`
199    // copies and count them, we sum the diffs directly. A net-negative count is
200    // possible (the surface does not define behavior in that case) and surfaces
201    // here as a negative result.
202    // TODO(jkosh44) This should error when the count can't fit inside of an `i64` instead of returning a negative result.
203    let mut count = Diff::ZERO;
204    for (datum, diff) in datums {
205        if !datum.is_null() {
206            count += diff;
207        }
208    }
209    Datum::from(count.into_inner())
210}
211
212fn any<'a, I>(datums: I) -> Datum<'a>
213where
214    I: IntoIterator<Item = Datum<'a>>,
215{
216    datums
217        .into_iter()
218        .fold(Datum::False, |state, next| match (state, next) {
219            (Datum::True, _) | (_, Datum::True) => Datum::True,
220            (Datum::Null, _) | (_, Datum::Null) => Datum::Null,
221            _ => Datum::False,
222        })
223}
224
225fn all<'a, I>(datums: I) -> Datum<'a>
226where
227    I: IntoIterator<Item = Datum<'a>>,
228{
229    datums
230        .into_iter()
231        .fold(Datum::True, |state, next| match (state, next) {
232            (Datum::False, _) | (_, Datum::False) => Datum::False,
233            (Datum::Null, _) | (_, Datum::Null) => Datum::Null,
234            _ => Datum::True,
235        })
236}
237
238fn string_agg<'a, I>(datums: I, temp_storage: &'a RowArena, order_by: &[ColumnOrder]) -> Datum<'a>
239where
240    I: IntoIterator<Item = Datum<'a>>,
241{
242    const EMPTY_SEP: &str = "";
243
244    let datums = order_aggregate_datums(datums, order_by);
245    let mut sep_value_pairs = datums.into_iter().filter_map(|d| {
246        if d.is_null() {
247            return None;
248        }
249        let mut value_sep = d.unwrap_list().iter();
250        match (value_sep.next().unwrap(), value_sep.next().unwrap()) {
251            (Datum::Null, _) => None,
252            (Datum::String(val), Datum::Null) => Some((EMPTY_SEP, val)),
253            (Datum::String(val), Datum::String(sep)) => Some((sep, val)),
254            _ => unreachable!(),
255        }
256    });
257
258    let mut s = String::default();
259    match sep_value_pairs.next() {
260        // First value not prefixed by its separator
261        Some((_, value)) => s.push_str(value),
262        // If no non-null values sent, return NULL.
263        None => return Datum::Null,
264    }
265
266    for (sep, value) in sep_value_pairs {
267        s.push_str(sep);
268        s.push_str(value);
269    }
270
271    Datum::String(temp_storage.push_string(s))
272}
273
274fn jsonb_agg<'a, I>(datums: I, temp_storage: &'a RowArena, order_by: &[ColumnOrder]) -> Datum<'a>
275where
276    I: IntoIterator<Item = Datum<'a>>,
277{
278    let datums = order_aggregate_datums(datums, order_by);
279    temp_storage.make_datum(|packer| {
280        packer.push_list(datums.into_iter().filter(|d| !d.is_null()));
281    })
282}
283
284fn dict_agg<'a, I>(datums: I, temp_storage: &'a RowArena, order_by: &[ColumnOrder]) -> Datum<'a>
285where
286    I: IntoIterator<Item = Datum<'a>>,
287{
288    let datums = order_aggregate_datums(datums, order_by);
289    temp_storage.make_datum(|packer| {
290        let mut datums: Vec<_> = datums
291            .into_iter()
292            .filter_map(|d| {
293                if d.is_null() {
294                    return None;
295                }
296                let mut list = d.unwrap_list().iter();
297                let key = list.next().unwrap();
298                let val = list.next().unwrap();
299                if key.is_null() {
300                    // TODO(benesch): this should produce an error, but
301                    // aggregate functions cannot presently produce errors.
302                    None
303                } else {
304                    Some((key.unwrap_str(), val))
305                }
306            })
307            .collect();
308        // datums are ordered by any ORDER BY clause now, and we want to preserve
309        // the last entry for each key, but we also need to present unique and sorted
310        // keys to push_dict. Use sort_by here, which is stable, and so will preserve
311        // the ORDER BY order. Then reverse and dedup to retain the last of each
312        // key. Reverse again so we're back in push_dict order.
313        datums.sort_by_key(|(k, _v)| *k);
314        datums.reverse();
315        datums.dedup_by_key(|(k, _v)| *k);
316        datums.reverse();
317        packer.push_dict(datums);
318    })
319}
320
321/// Assuming datums is a List, sort them by the 2nd through Nth elements
322/// corresponding to order_by, then return the 1st element.
323///
324/// Near the usages of this function, we sometimes want to produce Datums with a shorter lifetime
325/// than 'a. We have to actually perform the shortening of the lifetime here, inside this function,
326/// because if we were to simply return `impl Iterator<Item = Datum<'a>>`, that wouldn't be
327/// covariant in the item type, because opaque types are always invariant. (Contrast this with how
328/// we perform the shortening _inside_ this function: the input of the `map` is known to
329/// specifically be `std::vec::IntoIter`, which is known to be covariant.)
330pub fn order_aggregate_datums<'a: 'b, 'b, I>(
331    datums: I,
332    order_by: &[ColumnOrder],
333) -> impl Iterator<Item = Datum<'b>>
334where
335    I: IntoIterator<Item = Datum<'a>>,
336{
337    order_aggregate_datums_with_rank_inner(datums, order_by)
338        .into_iter()
339        // (`payload` is coerced here to `Datum<'b>` in the argument of the closure)
340        .map(|(payload, _order_datums)| payload)
341}
342
343/// Assuming datums is a List, sort them by the 2nd through Nth elements
344/// corresponding to order_by, then return the 1st element and computed order by expression.
345fn order_aggregate_datums_with_rank<'a, I>(
346    datums: I,
347    order_by: &[ColumnOrder],
348) -> impl Iterator<Item = (Datum<'a>, Row)>
349where
350    I: IntoIterator<Item = Datum<'a>>,
351{
352    order_aggregate_datums_with_rank_inner(datums, order_by)
353        .into_iter()
354        .map(|(payload, order_by_datums)| (payload, Row::pack(order_by_datums)))
355}
356
357fn order_aggregate_datums_with_rank_inner<'a, I>(
358    datums: I,
359    order_by: &[ColumnOrder],
360) -> Vec<(Datum<'a>, Vec<Datum<'a>>)>
361where
362    I: IntoIterator<Item = Datum<'a>>,
363{
364    let mut decoded: Vec<(Datum, Vec<Datum>)> = datums
365        .into_iter()
366        .map(|d| {
367            let list = d.unwrap_list();
368            let mut list_it = list.iter();
369            let payload = list_it.next().unwrap();
370
371            // We decode the order_by Datums here instead of the comparison function, because the
372            // comparison function is expected to be called `O(log n)` times on each input row.
373            // The only downside is that the decoded data might be bigger, but I think that's fine,
374            // because:
375            // - if we have a window partition so big that this would create a memory problem, then
376            //   the non-incrementalness of window functions will create a serious CPU problem
377            //   anyway,
378            // - and anyhow various other parts of the window function code already do decoding
379            //   upfront.
380            let mut order_by_datums = Vec::with_capacity(order_by.len());
381            for _ in 0..order_by.len() {
382                order_by_datums.push(
383                    list_it
384                        .next()
385                        .expect("must have exactly the same number of Datums as `order_by`"),
386                );
387            }
388
389            (payload, order_by_datums)
390        })
391        .collect();
392
393    let mut sort_by =
394        |(payload_left, left_order_by_datums): &(Datum, Vec<Datum>),
395         (payload_right, right_order_by_datums): &(Datum, Vec<Datum>)| {
396            compare_columns(
397                order_by,
398                left_order_by_datums,
399                right_order_by_datums,
400                || payload_left.cmp(payload_right),
401            )
402        };
403    // `sort_unstable_by` can be faster and uses less memory than `sort_by`. An unstable sort is
404    // enough here, because if two elements are equal in our `compare` function, then the elements
405    // are actually binary-equal (because of the `tiebreaker` given to `compare_columns`), so it
406    // doesn't matter what order they end up in.
407    decoded.sort_unstable_by(&mut sort_by);
408    decoded
409}
410
411fn array_concat<'a, I>(datums: I, temp_storage: &'a RowArena, order_by: &[ColumnOrder]) -> Datum<'a>
412where
413    I: IntoIterator<Item = Datum<'a>>,
414{
415    let datums = order_aggregate_datums(datums, order_by);
416    let datums: Vec<_> = datums
417        .into_iter()
418        .map(|d| d.unwrap_array().elements().iter())
419        .flatten()
420        .collect();
421    let dims = ArrayDimension {
422        lower_bound: 1,
423        length: datums.len(),
424    };
425    temp_storage.make_datum(|packer| {
426        packer.try_push_array(&[dims], datums).unwrap();
427    })
428}
429
430fn list_concat<'a, I>(datums: I, temp_storage: &'a RowArena, order_by: &[ColumnOrder]) -> Datum<'a>
431where
432    I: IntoIterator<Item = Datum<'a>>,
433{
434    let datums = order_aggregate_datums(datums, order_by);
435    temp_storage.make_datum(|packer| {
436        packer.push_list(datums.into_iter().map(|d| d.unwrap_list().iter()).flatten());
437    })
438}
439
440/// The expected input is in the format of `[((OriginalRow, [EncodedArgs]), OrderByExprs...)]`
441/// The output is in the format of `[result_value, original_row]`.
442/// See an example at `lag_lead`, where the input-output formats are similar.
443fn row_number<'a, I>(
444    datums: I,
445    callers_temp_storage: &'a RowArena,
446    order_by: &[ColumnOrder],
447) -> Datum<'a>
448where
449    I: IntoIterator<Item = Datum<'a>>,
450{
451    // We want to use our own temp_storage here, to avoid flooding `callers_temp_storage` with a
452    // large number of new datums. This is because we don't want to make an assumption about
453    // whether the caller creates a new temp_storage between window partitions.
454    let temp_storage = RowArena::new();
455    let datums = row_number_no_list(datums, &temp_storage, order_by);
456
457    callers_temp_storage.make_datum(|packer| {
458        packer.push_list(datums);
459    })
460}
461
462/// Like `row_number`, but doesn't perform the final wrapping in a list, returning an Iterator
463/// instead.
464fn row_number_no_list<'a: 'b, 'b, I>(
465    datums: I,
466    callers_temp_storage: &'b RowArena,
467    order_by: &[ColumnOrder],
468) -> impl Iterator<Item = Datum<'b>>
469where
470    I: IntoIterator<Item = Datum<'a>>,
471{
472    let datums = order_aggregate_datums(datums, order_by);
473
474    callers_temp_storage.reserve(datums.size_hint().0);
475    #[allow(clippy::disallowed_methods)]
476    datums
477        .into_iter()
478        .map(|d| d.unwrap_list().iter())
479        .flatten()
480        .zip(1i64..)
481        .map(|(d, i)| {
482            callers_temp_storage.make_datum(|packer| {
483                packer.push_list_with(|packer| {
484                    packer.push(Datum::Int64(i));
485                    packer.push(d);
486                });
487            })
488        })
489}
490
491/// The expected input is in the format of `[((OriginalRow, [EncodedArgs]), OrderByExprs...)]`
492/// The output is in the format of `[result_value, original_row]`.
493/// See an example at `lag_lead`, where the input-output formats are similar.
494fn rank<'a, I>(datums: I, callers_temp_storage: &'a RowArena, order_by: &[ColumnOrder]) -> Datum<'a>
495where
496    I: IntoIterator<Item = Datum<'a>>,
497{
498    let temp_storage = RowArena::new();
499    let datums = rank_no_list(datums, &temp_storage, order_by);
500
501    callers_temp_storage.make_datum(|packer| {
502        packer.push_list(datums);
503    })
504}
505
506/// Like `rank`, but doesn't perform the final wrapping in a list, returning an Iterator
507/// instead.
508fn rank_no_list<'a: 'b, 'b, I>(
509    datums: I,
510    callers_temp_storage: &'b RowArena,
511    order_by: &[ColumnOrder],
512) -> impl Iterator<Item = Datum<'b>>
513where
514    I: IntoIterator<Item = Datum<'a>>,
515{
516    // Keep the row used for ordering around, as it is used to determine the rank
517    let datums = order_aggregate_datums_with_rank(datums, order_by);
518
519    let mut datums = datums
520        .into_iter()
521        .map(|(d0, order_row)| {
522            d0.unwrap_list()
523                .iter()
524                .map(move |d1| (d1, order_row.clone()))
525        })
526        .flatten();
527
528    callers_temp_storage.reserve(datums.size_hint().0);
529    datums
530        .next()
531        .map_or(vec![], |(first_datum, first_order_row)| {
532            // Folding with (last order_by row, last assigned rank,
533            // row number, output vec)
534            datums.fold(
535                (first_order_row, 1, 1, vec![(first_datum, 1)]),
536                |mut acc, (next_datum, next_order_row)| {
537                let (ref mut acc_row, ref mut acc_rank, ref mut acc_row_num, ref mut output) = acc;
538                *acc_row_num += 1;
539                // Identity is based on the order_by expression
540                if *acc_row != next_order_row {
541                    *acc_rank = *acc_row_num;
542                    *acc_row = next_order_row;
543                }
544
545                (*output).push((next_datum, *acc_rank));
546                acc
547            })
548        }.3).into_iter().map(|(d, i)| {
549        callers_temp_storage.make_datum(|packer| {
550            packer.push_list_with(|packer| {
551                packer.push(Datum::Int64(i));
552                packer.push(d);
553            });
554        })
555    })
556}
557
558/// The expected input is in the format of `[((OriginalRow, [EncodedArgs]), OrderByExprs...)]`
559/// The output is in the format of `[result_value, original_row]`.
560/// See an example at `lag_lead`, where the input-output formats are similar.
561fn dense_rank<'a, I>(
562    datums: I,
563    callers_temp_storage: &'a RowArena,
564    order_by: &[ColumnOrder],
565) -> Datum<'a>
566where
567    I: IntoIterator<Item = Datum<'a>>,
568{
569    let temp_storage = RowArena::new();
570    let datums = dense_rank_no_list(datums, &temp_storage, order_by);
571
572    callers_temp_storage.make_datum(|packer| {
573        packer.push_list(datums);
574    })
575}
576
577/// Like `dense_rank`, but doesn't perform the final wrapping in a list, returning an Iterator
578/// instead.
579fn dense_rank_no_list<'a: 'b, 'b, I>(
580    datums: I,
581    callers_temp_storage: &'b RowArena,
582    order_by: &[ColumnOrder],
583) -> impl Iterator<Item = Datum<'b>>
584where
585    I: IntoIterator<Item = Datum<'a>>,
586{
587    // Keep the row used for ordering around, as it is used to determine the rank
588    let datums = order_aggregate_datums_with_rank(datums, order_by);
589
590    let mut datums = datums
591        .into_iter()
592        .map(|(d0, order_row)| {
593            d0.unwrap_list()
594                .iter()
595                .map(move |d1| (d1, order_row.clone()))
596        })
597        .flatten();
598
599    callers_temp_storage.reserve(datums.size_hint().0);
600    datums
601        .next()
602        .map_or(vec![], |(first_datum, first_order_row)| {
603            // Folding with (last order_by row, last assigned rank,
604            // output vec)
605            datums.fold(
606                (first_order_row, 1, vec![(first_datum, 1)]),
607                |mut acc, (next_datum, next_order_row)| {
608                let (ref mut acc_row, ref mut acc_rank, ref mut output) = acc;
609                // Identity is based on the order_by expression
610                if *acc_row != next_order_row {
611                    *acc_rank += 1;
612                    *acc_row = next_order_row;
613                }
614
615                (*output).push((next_datum, *acc_rank));
616                acc
617            })
618        }.2).into_iter().map(|(d, i)| {
619        callers_temp_storage.make_datum(|packer| {
620            packer.push_list_with(|packer| {
621                packer.push(Datum::Int64(i));
622                packer.push(d);
623            });
624        })
625    })
626}
627
628/// The expected input is in the format of `[((OriginalRow, EncodedArgs), OrderByExprs...)]`
629/// For example,
630///
631/// lag(x*y, 1, null) over (partition by x+y order by x-y, x/y)
632///
633/// list of:
634/// row(
635///   row(
636///     row(#0, #1),
637///     row((#0 * #1), 1, null)
638///   ),
639///   (#0 - #1),
640///   (#0 / #1)
641/// )
642///
643/// The output is in the format of `[result_value, original_row]`, e.g.
644/// list of:
645/// row(
646///   42,
647///   row(7, 8)
648/// )
649fn lag_lead<'a, I>(
650    datums: I,
651    callers_temp_storage: &'a RowArena,
652    order_by: &[ColumnOrder],
653    lag_lead_type: &LagLeadType,
654    ignore_nulls: &bool,
655) -> Datum<'a>
656where
657    I: IntoIterator<Item = Datum<'a>>,
658{
659    let temp_storage = RowArena::new();
660    let iter = lag_lead_no_list(datums, &temp_storage, order_by, lag_lead_type, ignore_nulls);
661    callers_temp_storage.make_datum(|packer| {
662        packer.push_list(iter);
663    })
664}
665
666/// Like `lag_lead`, but doesn't perform the final wrapping in a list, returning an Iterator
667/// instead.
668fn lag_lead_no_list<'a: 'b, 'b, I>(
669    datums: I,
670    callers_temp_storage: &'b RowArena,
671    order_by: &[ColumnOrder],
672    lag_lead_type: &LagLeadType,
673    ignore_nulls: &bool,
674) -> impl Iterator<Item = Datum<'b>>
675where
676    I: IntoIterator<Item = Datum<'a>>,
677{
678    // Sort the datums according to the ORDER BY expressions and return the (OriginalRow, EncodedArgs) record
679    let datums = order_aggregate_datums(datums, order_by);
680
681    // Take the (OriginalRow, EncodedArgs) records and unwrap them into separate datums.
682    // EncodedArgs = (InputValue, Offset, DefaultValue) for Lag/Lead
683    // (`OriginalRow` is kept in a record form, as we don't need to look inside that.)
684    let (orig_rows, unwrapped_args): (Vec<_>, Vec<_>) = datums
685        .into_iter()
686        .map(|d| {
687            let mut iter = d.unwrap_list().iter();
688            let original_row = iter.next().unwrap();
689            let (input_value, offset, default_value) =
690                unwrap_lag_lead_encoded_args(iter.next().unwrap());
691            (original_row, (input_value, offset, default_value))
692        })
693        .unzip();
694
695    let result = lag_lead_inner(unwrapped_args, lag_lead_type, ignore_nulls);
696
697    callers_temp_storage.reserve(result.len());
698    result
699        .into_iter()
700        .zip_eq(orig_rows)
701        .map(|(result_value, original_row)| {
702            callers_temp_storage.make_datum(|packer| {
703                packer.push_list_with(|packer| {
704                    packer.push(result_value);
705                    packer.push(original_row);
706                });
707            })
708        })
709}
710
711/// lag/lead's arguments are in a record. This function unwraps this record.
712fn unwrap_lag_lead_encoded_args(encoded_args: Datum) -> (Datum, Datum, Datum) {
713    let mut encoded_args_iter = encoded_args.unwrap_list().iter();
714    let (input_value, offset, default_value) = (
715        encoded_args_iter.next().unwrap(),
716        encoded_args_iter.next().unwrap(),
717        encoded_args_iter.next().unwrap(),
718    );
719    (input_value, offset, default_value)
720}
721
722/// Each element of `args` has the 3 arguments evaluated for a single input row.
723/// Returns the results for each input row.
724fn lag_lead_inner<'a>(
725    args: Vec<(Datum<'a>, Datum<'a>, Datum<'a>)>,
726    lag_lead_type: &LagLeadType,
727    ignore_nulls: &bool,
728) -> Vec<Datum<'a>> {
729    if *ignore_nulls {
730        lag_lead_inner_ignore_nulls(args, lag_lead_type)
731    } else {
732        lag_lead_inner_respect_nulls(args, lag_lead_type)
733    }
734}
735
736fn lag_lead_inner_respect_nulls<'a>(
737    args: Vec<(Datum<'a>, Datum<'a>, Datum<'a>)>,
738    lag_lead_type: &LagLeadType,
739) -> Vec<Datum<'a>> {
740    let mut result: Vec<Datum> = Vec::with_capacity(args.len());
741    for (idx, (_, offset, default_value)) in args.iter().enumerate() {
742        // Null offsets are acceptable, and always return null
743        if offset.is_null() {
744            result.push(Datum::Null);
745            continue;
746        }
747
748        let idx = i64::try_from(idx).expect("Array index does not fit in i64");
749        let offset = i64::from(offset.unwrap_int32());
750        let offset = match lag_lead_type {
751            LagLeadType::Lag => -offset,
752            LagLeadType::Lead => offset,
753        };
754
755        // Get a Datum from `datums`. Return None if index is out of range.
756        let datums_get = |i: i64| -> Option<Datum> {
757            match u64::try_from(i) {
758                Ok(i) => args
759                    .get(usize::cast_from(i))
760                    .map(|d| Some(d.0)) // succeeded in getting a Datum from the vec
761                    .unwrap_or(None), // overindexing
762                Err(_) => None, // underindexing (negative index)
763            }
764        };
765
766        let lagged_value = datums_get(idx + offset).unwrap_or(*default_value);
767
768        result.push(lagged_value);
769    }
770
771    result
772}
773
774// `i64` indexes get involved in this function because it's convenient to allow negative indexes and
775// have `datums_get` fail on them, and thus handle the beginning and end of the input vector
776// uniformly, rather than checking underflow separately during index manipulations.
777#[allow(clippy::as_conversions)]
778fn lag_lead_inner_ignore_nulls<'a>(
779    args: Vec<(Datum<'a>, Datum<'a>, Datum<'a>)>,
780    lag_lead_type: &LagLeadType,
781) -> Vec<Datum<'a>> {
782    // We check here once that even the largest index fits in `i64`, and then do silent `as`
783    // conversions from `usize` indexes to `i64` indexes throughout this function.
784    if i64::try_from(args.len()).is_err() {
785        panic!("window partition way too big")
786    }
787    // Preparation: Make sure we can jump over a run of nulls in constant time, i.e., regardless of
788    // how many nulls the run has. The following skip tables will point to the next non-null index.
789    let mut skip_nulls_backward = vec![None; args.len()];
790    let mut last_non_null: i64 = -1;
791    let pairs = args
792        .iter()
793        .enumerate()
794        .zip_eq(skip_nulls_backward.iter_mut());
795    for ((i, (d, _, _)), slot) in pairs {
796        if d.is_null() {
797            *slot = Some(last_non_null);
798        } else {
799            last_non_null = i as i64;
800        }
801    }
802    let mut skip_nulls_forward = vec![None; args.len()];
803    let mut last_non_null: i64 = args.len() as i64;
804    let pairs = args
805        .iter()
806        .enumerate()
807        .rev()
808        .zip_eq(skip_nulls_forward.iter_mut().rev());
809    for ((i, (d, _, _)), slot) in pairs {
810        if d.is_null() {
811            *slot = Some(last_non_null);
812        } else {
813            last_non_null = i as i64;
814        }
815    }
816
817    // The actual computation.
818    let mut result: Vec<Datum> = Vec::with_capacity(args.len());
819    for (idx, (_, offset, default_value)) in args.iter().enumerate() {
820        // Null offsets are acceptable, and always return null
821        if offset.is_null() {
822            result.push(Datum::Null);
823            continue;
824        }
825
826        let idx = idx as i64; // checked at the beginning of the function that len() fits
827        let offset = i64::cast_from(offset.unwrap_int32());
828        let offset = match lag_lead_type {
829            LagLeadType::Lag => -offset,
830            LagLeadType::Lead => offset,
831        };
832        let increment = offset.signum();
833
834        // Get a Datum from `datums`. Return None if index is out of range.
835        let datums_get = |i: i64| -> Option<Datum> {
836            match u64::try_from(i) {
837                Ok(i) => args
838                    .get(usize::cast_from(i))
839                    .map(|d| Some(d.0)) // succeeded in getting a Datum from the vec
840                    .unwrap_or(None), // overindexing
841                Err(_) => None, // underindexing (negative index)
842            }
843        };
844
845        let lagged_value = if increment != 0 {
846            // We start j from idx, and step j until we have seen an abs(offset) number of non-null
847            // values or reach the beginning or end of the partition.
848            //
849            // If offset is big, then this is slow: Considering the entire function, it's
850            // `O(partition_size * offset)`.
851            // However, a common use case is an offset of 1, for which this doesn't matter.
852            // TODO: For larger offsets, we could have a completely different implementation
853            // that starts the inner loop from the index where we found the previous result:
854            // https://github.com/MaterializeInc/materialize/pull/29287#discussion_r1738695174
855            let mut j = idx;
856            for _ in 0..num::abs(offset) {
857                j += increment;
858                // Jump over a run of nulls
859                if datums_get(j).is_some_and(|d| d.is_null()) {
860                    let ju = j as usize; // `j >= 0` because of the above `is_some_and`
861                    if increment > 0 {
862                        j = skip_nulls_forward[ju].expect("checked above that it's null");
863                    } else {
864                        j = skip_nulls_backward[ju].expect("checked above that it's null");
865                    }
866                }
867                if datums_get(j).is_none() {
868                    break;
869                }
870            }
871            match datums_get(j) {
872                Some(datum) => datum,
873                None => *default_value,
874            }
875        } else {
876            assert_eq!(offset, 0);
877            let datum = datums_get(idx).expect("known to exist");
878            if !datum.is_null() {
879                datum
880            } else {
881                // Not clear what should the semantics be here. See
882                // https://github.com/MaterializeInc/database-issues/issues/8497
883                // (We used to run into an infinite loop in this case, so panicking is
884                // better.)
885                panic!("0 offset in lag/lead IGNORE NULLS");
886            }
887        };
888
889        result.push(lagged_value);
890    }
891
892    result
893}
894
895/// The expected input is in the format of [((OriginalRow, InputValue), OrderByExprs...)]
896fn first_value<'a, I>(
897    datums: I,
898    callers_temp_storage: &'a RowArena,
899    order_by: &[ColumnOrder],
900    window_frame: &WindowFrame,
901) -> Datum<'a>
902where
903    I: IntoIterator<Item = Datum<'a>>,
904{
905    let temp_storage = RowArena::new();
906    let iter = first_value_no_list(datums, &temp_storage, order_by, window_frame);
907    callers_temp_storage.make_datum(|packer| {
908        packer.push_list(iter);
909    })
910}
911
912/// Like `first_value`, but doesn't perform the final wrapping in a list, returning an Iterator
913/// instead.
914fn first_value_no_list<'a: 'b, 'b, I>(
915    datums: I,
916    callers_temp_storage: &'b RowArena,
917    order_by: &[ColumnOrder],
918    window_frame: &WindowFrame,
919) -> impl Iterator<Item = Datum<'b>>
920where
921    I: IntoIterator<Item = Datum<'a>>,
922{
923    // Sort the datums according to the ORDER BY expressions and return the (OriginalRow, InputValue) record
924    let datums = order_aggregate_datums(datums, order_by);
925
926    // Decode the input (OriginalRow, InputValue) into separate datums
927    let (orig_rows, args): (Vec<_>, Vec<_>) = datums
928        .into_iter()
929        .map(|d| {
930            let mut iter = d.unwrap_list().iter();
931            let original_row = iter.next().unwrap();
932            let arg = iter.next().unwrap();
933
934            (original_row, arg)
935        })
936        .unzip();
937
938    let results = first_value_inner(args, window_frame);
939
940    callers_temp_storage.reserve(results.len());
941    results
942        .into_iter()
943        .zip_eq(orig_rows)
944        .map(|(result_value, original_row)| {
945            callers_temp_storage.make_datum(|packer| {
946                packer.push_list_with(|packer| {
947                    packer.push(result_value);
948                    packer.push(original_row);
949                });
950            })
951        })
952}
953
954fn first_value_inner<'a>(datums: Vec<Datum<'a>>, window_frame: &WindowFrame) -> Vec<Datum<'a>> {
955    let length = datums.len();
956    let mut result: Vec<Datum> = Vec::with_capacity(length);
957    for (idx, current_datum) in datums.iter().enumerate() {
958        let first_value = match &window_frame.start_bound {
959            // Always return the current value
960            WindowFrameBound::CurrentRow => *current_datum,
961            WindowFrameBound::UnboundedPreceding => {
962                if let WindowFrameBound::OffsetPreceding(end_offset) = &window_frame.end_bound {
963                    let end_offset = usize::cast_from(*end_offset);
964
965                    // If the frame ends before the first row, return null
966                    if idx < end_offset {
967                        Datum::Null
968                    } else {
969                        datums[0]
970                    }
971                } else {
972                    datums[0]
973                }
974            }
975            WindowFrameBound::OffsetPreceding(offset) => {
976                let start_offset = usize::cast_from(*offset);
977                let start_idx = idx.saturating_sub(start_offset);
978                if let WindowFrameBound::OffsetPreceding(end_offset) = &window_frame.end_bound {
979                    let end_offset = usize::cast_from(*end_offset);
980
981                    // If the frame is empty or ends before the first row, return null
982                    if start_offset < end_offset || idx < end_offset {
983                        Datum::Null
984                    } else {
985                        datums[start_idx]
986                    }
987                } else {
988                    datums[start_idx]
989                }
990            }
991            WindowFrameBound::OffsetFollowing(offset) => {
992                let start_offset = usize::cast_from(*offset);
993                let start_idx = idx.saturating_add(start_offset);
994                if let WindowFrameBound::OffsetFollowing(end_offset) = &window_frame.end_bound {
995                    // If the frame is empty or starts after the last row, return null
996                    if offset > end_offset || start_idx >= length {
997                        Datum::Null
998                    } else {
999                        datums[start_idx]
1000                    }
1001                } else {
1002                    datums
1003                        .get(start_idx)
1004                        .map(|d| d.clone())
1005                        .unwrap_or(Datum::Null)
1006                }
1007            }
1008            // Forbidden during planning
1009            WindowFrameBound::UnboundedFollowing => unreachable!(),
1010        };
1011        result.push(first_value);
1012    }
1013    result
1014}
1015
1016/// The expected input is in the format of [((OriginalRow, InputValue), OrderByExprs...)]
1017fn last_value<'a, I>(
1018    datums: I,
1019    callers_temp_storage: &'a RowArena,
1020    order_by: &[ColumnOrder],
1021    window_frame: &WindowFrame,
1022) -> Datum<'a>
1023where
1024    I: IntoIterator<Item = Datum<'a>>,
1025{
1026    let temp_storage = RowArena::new();
1027    let iter = last_value_no_list(datums, &temp_storage, order_by, window_frame);
1028    callers_temp_storage.make_datum(|packer| {
1029        packer.push_list(iter);
1030    })
1031}
1032
1033/// Like `last_value`, but doesn't perform the final wrapping in a list, returning an Iterator
1034/// instead.
1035fn last_value_no_list<'a: 'b, 'b, I>(
1036    datums: I,
1037    callers_temp_storage: &'b RowArena,
1038    order_by: &[ColumnOrder],
1039    window_frame: &WindowFrame,
1040) -> impl Iterator<Item = Datum<'b>>
1041where
1042    I: IntoIterator<Item = Datum<'a>>,
1043{
1044    // Sort the datums according to the ORDER BY expressions and return the ((OriginalRow, InputValue), OrderByRow) record
1045    // The OrderByRow is kept around because it is required to compute the peer groups in RANGE mode
1046    let datums = order_aggregate_datums_with_rank(datums, order_by);
1047
1048    // Decode the input (OriginalRow, InputValue) into separate datums, while keeping the OrderByRow
1049    let size_hint = datums.size_hint().0;
1050    let mut args = Vec::with_capacity(size_hint);
1051    let mut original_rows = Vec::with_capacity(size_hint);
1052    let mut order_by_rows = Vec::with_capacity(size_hint);
1053    for (d, order_by_row) in datums.into_iter() {
1054        let mut iter = d.unwrap_list().iter();
1055        let original_row = iter.next().unwrap();
1056        let arg = iter.next().unwrap();
1057        order_by_rows.push(order_by_row);
1058        original_rows.push(original_row);
1059        args.push(arg);
1060    }
1061
1062    let results = last_value_inner(args, &order_by_rows, window_frame);
1063
1064    callers_temp_storage.reserve(results.len());
1065    results
1066        .into_iter()
1067        .zip_eq(original_rows)
1068        .map(|(result_value, original_row)| {
1069            callers_temp_storage.make_datum(|packer| {
1070                packer.push_list_with(|packer| {
1071                    packer.push(result_value);
1072                    packer.push(original_row);
1073                });
1074            })
1075        })
1076}
1077
1078fn last_value_inner<'a>(
1079    args: Vec<Datum<'a>>,
1080    order_by_rows: &Vec<Row>,
1081    window_frame: &WindowFrame,
1082) -> Vec<Datum<'a>> {
1083    let length = args.len();
1084    let mut results: Vec<Datum> = Vec::with_capacity(length);
1085    for (idx, (current_datum, order_by_row)) in args.iter().zip_eq(order_by_rows).enumerate() {
1086        let last_value = match &window_frame.end_bound {
1087            WindowFrameBound::CurrentRow => match &window_frame.units {
1088                // Always return the current value when in ROWS mode
1089                WindowFrameUnits::Rows => *current_datum,
1090                WindowFrameUnits::Range => {
1091                    // When in RANGE mode, return the last value of the peer group
1092                    // The peer group is the group of rows with the same ORDER BY value
1093                    // Note: Range is only supported for the default window frame (RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW),
1094                    // which is why it does not appear in the other branches
1095                    let target_idx = order_by_rows[idx..]
1096                        .iter()
1097                        .enumerate()
1098                        .take_while(|(_, row)| *row == order_by_row)
1099                        .last()
1100                        .unwrap()
1101                        .0
1102                        + idx;
1103                    args[target_idx]
1104                }
1105                // GROUPS is not supported, and forbidden during planning
1106                WindowFrameUnits::Groups => unreachable!(),
1107            },
1108            WindowFrameBound::UnboundedFollowing => {
1109                if let WindowFrameBound::OffsetFollowing(start_offset) = &window_frame.start_bound {
1110                    let start_offset = usize::cast_from(*start_offset);
1111
1112                    // If the frame starts after the last row of the window, return null
1113                    if idx + start_offset > length - 1 {
1114                        Datum::Null
1115                    } else {
1116                        args[length - 1]
1117                    }
1118                } else {
1119                    args[length - 1]
1120                }
1121            }
1122            WindowFrameBound::OffsetFollowing(offset) => {
1123                let end_offset = usize::cast_from(*offset);
1124                let end_idx = idx.saturating_add(end_offset);
1125                if let WindowFrameBound::OffsetFollowing(start_offset) = &window_frame.start_bound {
1126                    let start_offset = usize::cast_from(*start_offset);
1127                    let start_idx = idx.saturating_add(start_offset);
1128
1129                    // If the frame is empty or starts after the last row of the window, return null
1130                    if end_offset < start_offset || start_idx >= length {
1131                        Datum::Null
1132                    } else {
1133                        // Return the last valid element in the window
1134                        args.get(end_idx).unwrap_or(&args[length - 1]).clone()
1135                    }
1136                } else {
1137                    args.get(end_idx).unwrap_or(&args[length - 1]).clone()
1138                }
1139            }
1140            WindowFrameBound::OffsetPreceding(offset) => {
1141                let end_offset = usize::cast_from(*offset);
1142                let end_idx = idx.saturating_sub(end_offset);
1143                if idx < end_offset {
1144                    // If the frame ends before the first row, return null
1145                    Datum::Null
1146                } else if let WindowFrameBound::OffsetPreceding(start_offset) =
1147                    &window_frame.start_bound
1148                {
1149                    // If the frame is empty, return null
1150                    if offset > start_offset {
1151                        Datum::Null
1152                    } else {
1153                        args[end_idx]
1154                    }
1155                } else {
1156                    args[end_idx]
1157                }
1158            }
1159            // Forbidden during planning
1160            WindowFrameBound::UnboundedPreceding => unreachable!(),
1161        };
1162        results.push(last_value);
1163    }
1164    results
1165}
1166
1167/// Executes `FusedValueWindowFunc` on a reduction group.
1168/// The expected input is in the format of `[((OriginalRow, (Args1, Args2, ...)), OrderByExprs...)]`
1169/// where `Args1`, `Args2`, are the arguments of each of the fused functions. For functions that
1170/// have only a single argument (first_value/last_value), these are simple values. For functions
1171/// that have multiple arguments (lag/lead), these are also records.
1172fn fused_value_window_func<'a, I>(
1173    input_datums: I,
1174    callers_temp_storage: &'a RowArena,
1175    funcs: &Vec<AggregateFunc>,
1176    order_by: &Vec<ColumnOrder>,
1177) -> Datum<'a>
1178where
1179    I: IntoIterator<Item = Datum<'a>>,
1180{
1181    let temp_storage = RowArena::new();
1182    let iter = fused_value_window_func_no_list(input_datums, &temp_storage, funcs, order_by);
1183    callers_temp_storage.make_datum(|packer| {
1184        packer.push_list(iter);
1185    })
1186}
1187
1188/// Like `fused_value_window_func`, but doesn't perform the final wrapping in a list, returning an
1189/// Iterator instead.
1190fn fused_value_window_func_no_list<'a: 'b, 'b, I>(
1191    input_datums: I,
1192    callers_temp_storage: &'b RowArena,
1193    funcs: &Vec<AggregateFunc>,
1194    order_by: &Vec<ColumnOrder>,
1195) -> impl Iterator<Item = Datum<'b>>
1196where
1197    I: IntoIterator<Item = Datum<'a>>,
1198{
1199    let has_last_value = funcs
1200        .iter()
1201        .any(|f| matches!(f, AggregateFunc::LastValue { .. }));
1202
1203    let input_datums_with_ranks = order_aggregate_datums_with_rank(input_datums, order_by);
1204
1205    let size_hint = input_datums_with_ranks.size_hint().0;
1206    let mut encoded_argsss = vec![Vec::with_capacity(size_hint); funcs.len()];
1207    let mut original_rows = Vec::with_capacity(size_hint);
1208    let mut order_by_rows = Vec::with_capacity(size_hint);
1209    for (d, order_by_row) in input_datums_with_ranks {
1210        let mut iter = d.unwrap_list().iter();
1211        let original_row = iter.next().unwrap();
1212        original_rows.push(original_row);
1213        let mut argss_iter = iter.next().unwrap().unwrap_list().iter();
1214        for i in 0..funcs.len() {
1215            let encoded_args = argss_iter.next().unwrap();
1216            encoded_argsss[i].push(encoded_args);
1217        }
1218        if has_last_value {
1219            order_by_rows.push(order_by_row);
1220        }
1221    }
1222
1223    let mut results_per_row = vec![Vec::with_capacity(funcs.len()); original_rows.len()];
1224    for (func, encoded_argss) in funcs.iter().zip_eq(encoded_argsss) {
1225        let results = match func {
1226            AggregateFunc::LagLead {
1227                order_by: inner_order_by,
1228                lag_lead,
1229                ignore_nulls,
1230            } => {
1231                assert_eq!(order_by, inner_order_by);
1232                let unwrapped_argss = encoded_argss
1233                    .into_iter()
1234                    .map(|encoded_args| unwrap_lag_lead_encoded_args(encoded_args))
1235                    .collect();
1236                lag_lead_inner(unwrapped_argss, lag_lead, ignore_nulls)
1237            }
1238            AggregateFunc::FirstValue {
1239                order_by: inner_order_by,
1240                window_frame,
1241            } => {
1242                assert_eq!(order_by, inner_order_by);
1243                // (No unwrapping to do on the args here, because there is only 1 arg, so it's not
1244                // wrapped into a record.)
1245                first_value_inner(encoded_argss, window_frame)
1246            }
1247            AggregateFunc::LastValue {
1248                order_by: inner_order_by,
1249                window_frame,
1250            } => {
1251                assert_eq!(order_by, inner_order_by);
1252                // (No unwrapping to do on the args here, because there is only 1 arg, so it's not
1253                // wrapped into a record.)
1254                last_value_inner(encoded_argss, &order_by_rows, window_frame)
1255            }
1256            _ => panic!("unknown window function in FusedValueWindowFunc"),
1257        };
1258        for (results, result) in results_per_row.iter_mut().zip_eq(results) {
1259            results.push(result);
1260        }
1261    }
1262
1263    callers_temp_storage.reserve(2 * original_rows.len());
1264    results_per_row
1265        .into_iter()
1266        .enumerate()
1267        .map(move |(i, results)| {
1268            callers_temp_storage.make_datum(|packer| {
1269                packer.push_list_with(|packer| {
1270                    packer
1271                        .push(callers_temp_storage.make_datum(|packer| packer.push_list(results)));
1272                    packer.push(original_rows[i]);
1273                });
1274            })
1275        })
1276}
1277
1278/// `input_datums` is an entire window partition.
1279/// The expected input is in the format of `[((OriginalRow, InputValue), OrderByExprs...)]`
1280/// See also in the comment in `window_func_applied_to`.
1281///
1282/// `wrapped_aggregate`: e.g., for `sum(...) OVER (...)`, this is the `sum(...)`.
1283///
1284/// Note that this `order_by` doesn't have expressions, only `ColumnOrder`s. For an explanation,
1285/// see the comment on `WindowExprType`.
1286fn window_aggr<'a, I, A>(
1287    input_datums: I,
1288    callers_temp_storage: &'a RowArena,
1289    wrapped_aggregate: &AggregateFunc,
1290    order_by: &[ColumnOrder],
1291    window_frame: &WindowFrame,
1292) -> Datum<'a>
1293where
1294    I: IntoIterator<Item = Datum<'a>>,
1295    A: OneByOneAggr,
1296{
1297    let temp_storage = RowArena::new();
1298    let iter = window_aggr_no_list::<I, A>(
1299        input_datums,
1300        &temp_storage,
1301        wrapped_aggregate,
1302        order_by,
1303        window_frame,
1304    );
1305    callers_temp_storage.make_datum(|packer| {
1306        packer.push_list(iter);
1307    })
1308}
1309
1310/// Like `window_aggr`, but doesn't perform the final wrapping in a list, returning an Iterator
1311/// instead.
1312fn window_aggr_no_list<'a: 'b, 'b, I, A>(
1313    input_datums: I,
1314    callers_temp_storage: &'b RowArena,
1315    wrapped_aggregate: &AggregateFunc,
1316    order_by: &[ColumnOrder],
1317    window_frame: &WindowFrame,
1318) -> impl Iterator<Item = Datum<'b>>
1319where
1320    I: IntoIterator<Item = Datum<'a>>,
1321    A: OneByOneAggr,
1322{
1323    // Sort the datums according to the ORDER BY expressions and return the ((OriginalRow, InputValue), OrderByRow) record
1324    // The OrderByRow is kept around because it is required to compute the peer groups in RANGE mode
1325    let datums = order_aggregate_datums_with_rank(input_datums, order_by);
1326
1327    // Decode the input (OriginalRow, InputValue) into separate datums, while keeping the OrderByRow
1328    let size_hint = datums.size_hint().0;
1329    let mut args: Vec<Datum> = Vec::with_capacity(size_hint);
1330    let mut original_rows: Vec<Datum> = Vec::with_capacity(size_hint);
1331    let mut order_by_rows = Vec::with_capacity(size_hint);
1332    for (d, order_by_row) in datums.into_iter() {
1333        let mut iter = d.unwrap_list().iter();
1334        let original_row = iter.next().unwrap();
1335        let arg = iter.next().unwrap();
1336        order_by_rows.push(order_by_row);
1337        original_rows.push(original_row);
1338        args.push(arg);
1339    }
1340
1341    let results = window_aggr_inner::<A>(
1342        args,
1343        &order_by_rows,
1344        wrapped_aggregate,
1345        order_by,
1346        window_frame,
1347        callers_temp_storage,
1348    );
1349
1350    callers_temp_storage.reserve(results.len());
1351    results
1352        .into_iter()
1353        .zip_eq(original_rows)
1354        .map(|(result_value, original_row)| {
1355            callers_temp_storage.make_datum(|packer| {
1356                packer.push_list_with(|packer| {
1357                    packer.push(result_value);
1358                    packer.push(original_row);
1359                });
1360            })
1361        })
1362}
1363
1364fn window_aggr_inner<'a, A>(
1365    mut args: Vec<Datum<'a>>,
1366    order_by_rows: &Vec<Row>,
1367    wrapped_aggregate: &AggregateFunc,
1368    order_by: &[ColumnOrder],
1369    window_frame: &WindowFrame,
1370    temp_storage: &'a RowArena,
1371) -> Vec<Datum<'a>>
1372where
1373    A: OneByOneAggr,
1374{
1375    let length = args.len();
1376    let mut result: Vec<Datum> = Vec::with_capacity(length);
1377
1378    // In this degenerate case, all results would be `wrapped_aggregate.default()` (usually null).
1379    // However, this currently can't happen, because
1380    // - Groups frame mode is currently not supported;
1381    // - Range frame mode is currently supported only for the default frame, which includes the
1382    //   current row.
1383    soft_assert_or_log!(
1384        !((matches!(window_frame.units, WindowFrameUnits::Groups)
1385            || matches!(window_frame.units, WindowFrameUnits::Range))
1386            && !window_frame.includes_current_row()),
1387        "window frame without current row"
1388    );
1389
1390    if (matches!(
1391        window_frame.start_bound,
1392        WindowFrameBound::UnboundedPreceding
1393    ) && matches!(window_frame.end_bound, WindowFrameBound::UnboundedFollowing))
1394        || (order_by.is_empty()
1395            && (matches!(window_frame.units, WindowFrameUnits::Groups)
1396                || matches!(window_frame.units, WindowFrameUnits::Range))
1397            && window_frame.includes_current_row())
1398    {
1399        // Either
1400        //  - UNBOUNDED frame in both directions, or
1401        //  - There is no ORDER BY and the frame is such that the current peer group is included.
1402        //    (The current peer group will be the whole partition if there is no ORDER BY.)
1403        // We simply need to compute the aggregate once, on the entire partition, and each input
1404        // row will get this one aggregate value as result.
1405        let result_value =
1406            wrapped_aggregate.eval(args.into_iter().map(|d| (d, Diff::ONE)), temp_storage);
1407        // Every row will get the above aggregate as result.
1408        for _ in 0..length {
1409            result.push(result_value);
1410        }
1411    } else {
1412        fn rows_between_unbounded_preceding_and_current_row<'a, A>(
1413            args: Vec<Datum<'a>>,
1414            result: &mut Vec<Datum<'a>>,
1415            mut one_by_one_aggr: A,
1416            temp_storage: &'a RowArena,
1417        ) where
1418            A: OneByOneAggr,
1419        {
1420            for current_arg in args.into_iter() {
1421                one_by_one_aggr.give(&current_arg);
1422                let result_value = one_by_one_aggr.get_current_aggregate(temp_storage);
1423                result.push(result_value);
1424            }
1425        }
1426
1427        fn groups_between_unbounded_preceding_and_current_row<'a, A>(
1428            args: Vec<Datum<'a>>,
1429            order_by_rows: &Vec<Row>,
1430            result: &mut Vec<Datum<'a>>,
1431            mut one_by_one_aggr: A,
1432            temp_storage: &'a RowArena,
1433        ) where
1434            A: OneByOneAggr,
1435        {
1436            let mut peer_group_start = 0;
1437            while peer_group_start < args.len() {
1438                // Find the boundaries of the current peer group.
1439                // peer_group_start will point to the first element of the peer group,
1440                // peer_group_end will point to _just after_ the last element of the peer group.
1441                let mut peer_group_end = peer_group_start + 1;
1442                while peer_group_end < args.len()
1443                    && order_by_rows[peer_group_start] == order_by_rows[peer_group_end]
1444                {
1445                    // The peer group goes on while the OrderByRows not differ.
1446                    peer_group_end += 1;
1447                }
1448                // Let's compute the aggregate (which will be the same for all records in this
1449                // peer group).
1450                for current_arg in args[peer_group_start..peer_group_end].iter() {
1451                    one_by_one_aggr.give(current_arg);
1452                }
1453                let agg_for_peer_group = one_by_one_aggr.get_current_aggregate(temp_storage);
1454                // Put the above aggregate into each record in the peer group.
1455                for _ in args[peer_group_start..peer_group_end].iter() {
1456                    result.push(agg_for_peer_group);
1457                }
1458                // Point to the start of the next peer group.
1459                peer_group_start = peer_group_end;
1460            }
1461        }
1462
1463        fn rows_between_offset_and_offset<'a>(
1464            args: Vec<Datum<'a>>,
1465            result: &mut Vec<Datum<'a>>,
1466            wrapped_aggregate: &AggregateFunc,
1467            temp_storage: &'a RowArena,
1468            offset_start: i64,
1469            offset_end: i64,
1470        ) {
1471            let len = args
1472                .len()
1473                .to_i64()
1474                .expect("window partition's len should fit into i64");
1475            for i in 0..len {
1476                let i = i.to_i64().expect("window partition shouldn't be super big");
1477                // Trim the start of the frame to make it not reach over the start of the window
1478                // partition.
1479                let frame_start = max(i + offset_start, 0)
1480                    .to_usize()
1481                    .expect("The max made sure it's not negative");
1482                // Trim the end of the frame to make it not reach over the end of the window
1483                // partition.
1484                let frame_end = min(i + offset_end, len - 1).to_usize();
1485                match frame_end {
1486                    Some(frame_end) => {
1487                        if frame_start <= frame_end {
1488                            // Compute the aggregate on the frame.
1489                            // TODO:
1490                            // This implementation is quite slow if the frame is large: we do an
1491                            // inner loop over the entire frame, and compute the aggregate from
1492                            // scratch. We could do better:
1493                            //  - For invertible aggregations we could do a rolling aggregation.
1494                            //  - There are various tricks for min/max as well, making use of either
1495                            //    the fixed size of the window, or that we are not retracting
1496                            //    arbitrary elements but doing queue operations. E.g., see
1497                            //    http://codercareer.blogspot.com/2012/02/no-33-maximums-in-sliding-windows.html
1498                            let frame_values = args[frame_start..=frame_end]
1499                                .iter()
1500                                .map(|d| (*d, Diff::ONE));
1501                            let result_value = wrapped_aggregate.eval(frame_values, temp_storage);
1502                            result.push(result_value);
1503                        } else {
1504                            // frame_start > frame_end, so this is an empty frame.
1505                            let result_value = wrapped_aggregate.default();
1506                            result.push(result_value);
1507                        }
1508                    }
1509                    None => {
1510                        // frame_end would be negative, so this is an empty frame.
1511                        let result_value = wrapped_aggregate.default();
1512                        result.push(result_value);
1513                    }
1514                }
1515            }
1516        }
1517
1518        match (
1519            &window_frame.units,
1520            &window_frame.start_bound,
1521            &window_frame.end_bound,
1522        ) {
1523            // Cases where one edge of the frame is CurrentRow.
1524            // Note that these cases could be merged into the more general cases below where one
1525            // edge is some offset (with offset = 0), but the CurrentRow cases probably cover 95%
1526            // of user queries, so let's make this simple and fast.
1527            (Rows, UnboundedPreceding, CurrentRow) => {
1528                rows_between_unbounded_preceding_and_current_row::<A>(
1529                    args,
1530                    &mut result,
1531                    A::new(wrapped_aggregate, false),
1532                    temp_storage,
1533                );
1534            }
1535            (Rows, CurrentRow, UnboundedFollowing) => {
1536                // Same as above, but reverse.
1537                args.reverse();
1538                rows_between_unbounded_preceding_and_current_row::<A>(
1539                    args,
1540                    &mut result,
1541                    A::new(wrapped_aggregate, true),
1542                    temp_storage,
1543                );
1544                result.reverse();
1545            }
1546            (Range, UnboundedPreceding, CurrentRow) => {
1547                // Note that for the default frame, the RANGE frame mode is identical to the GROUPS
1548                // frame mode.
1549                groups_between_unbounded_preceding_and_current_row::<A>(
1550                    args,
1551                    order_by_rows,
1552                    &mut result,
1553                    A::new(wrapped_aggregate, false),
1554                    temp_storage,
1555                );
1556            }
1557            // The next several cases all call `rows_between_offset_and_offset`. Note that the
1558            // offset passed to `rows_between_offset_and_offset` should be negated when it's
1559            // PRECEDING.
1560            (Rows, OffsetPreceding(start_prec), OffsetPreceding(end_prec)) => {
1561                let start_prec = start_prec.to_i64().expect(
1562                    "window frame start OFFSET shouldn't be super big (the planning ensured this)",
1563                );
1564                let end_prec = end_prec.to_i64().expect(
1565                    "window frame end OFFSET shouldn't be super big (the planning ensured this)",
1566                );
1567                rows_between_offset_and_offset(
1568                    args,
1569                    &mut result,
1570                    wrapped_aggregate,
1571                    temp_storage,
1572                    -start_prec,
1573                    -end_prec,
1574                );
1575            }
1576            (Rows, OffsetPreceding(start_prec), OffsetFollowing(end_fol)) => {
1577                let start_prec = start_prec.to_i64().expect(
1578                    "window frame start OFFSET shouldn't be super big (the planning ensured this)",
1579                );
1580                let end_fol = end_fol.to_i64().expect(
1581                    "window frame end OFFSET shouldn't be super big (the planning ensured this)",
1582                );
1583                rows_between_offset_and_offset(
1584                    args,
1585                    &mut result,
1586                    wrapped_aggregate,
1587                    temp_storage,
1588                    -start_prec,
1589                    end_fol,
1590                );
1591            }
1592            (Rows, OffsetFollowing(start_fol), OffsetFollowing(end_fol)) => {
1593                let start_fol = start_fol.to_i64().expect(
1594                    "window frame start OFFSET shouldn't be super big (the planning ensured this)",
1595                );
1596                let end_fol = end_fol.to_i64().expect(
1597                    "window frame end OFFSET shouldn't be super big (the planning ensured this)",
1598                );
1599                rows_between_offset_and_offset(
1600                    args,
1601                    &mut result,
1602                    wrapped_aggregate,
1603                    temp_storage,
1604                    start_fol,
1605                    end_fol,
1606                );
1607            }
1608            (Rows, OffsetFollowing(_), OffsetPreceding(_)) => {
1609                unreachable!() // The planning ensured that this nonsensical case can't happen
1610            }
1611            (Rows, OffsetPreceding(start_prec), CurrentRow) => {
1612                let start_prec = start_prec.to_i64().expect(
1613                    "window frame start OFFSET shouldn't be super big (the planning ensured this)",
1614                );
1615                let end_fol = 0;
1616                rows_between_offset_and_offset(
1617                    args,
1618                    &mut result,
1619                    wrapped_aggregate,
1620                    temp_storage,
1621                    -start_prec,
1622                    end_fol,
1623                );
1624            }
1625            (Rows, CurrentRow, OffsetFollowing(end_fol)) => {
1626                let start_fol = 0;
1627                let end_fol = end_fol.to_i64().expect(
1628                    "window frame end OFFSET shouldn't be super big (the planning ensured this)",
1629                );
1630                rows_between_offset_and_offset(
1631                    args,
1632                    &mut result,
1633                    wrapped_aggregate,
1634                    temp_storage,
1635                    start_fol,
1636                    end_fol,
1637                );
1638            }
1639            (Rows, CurrentRow, CurrentRow) => {
1640                // We could have a more efficient implementation for this, but this is probably
1641                // super rare. (Might be more common with RANGE or GROUPS frame mode, though!)
1642                let start_fol = 0;
1643                let end_fol = 0;
1644                rows_between_offset_and_offset(
1645                    args,
1646                    &mut result,
1647                    wrapped_aggregate,
1648                    temp_storage,
1649                    start_fol,
1650                    end_fol,
1651                );
1652            }
1653            (Rows, CurrentRow, OffsetPreceding(_))
1654            | (Rows, UnboundedFollowing, _)
1655            | (Rows, _, UnboundedPreceding)
1656            | (Rows, OffsetFollowing(..), CurrentRow) => {
1657                unreachable!() // The planning ensured that these nonsensical cases can't happen
1658            }
1659            (Rows, UnboundedPreceding, UnboundedFollowing) => {
1660                // This is handled by the complicated if condition near the beginning of this
1661                // function.
1662                unreachable!()
1663            }
1664            (Rows, UnboundedPreceding, OffsetPreceding(_))
1665            | (Rows, UnboundedPreceding, OffsetFollowing(_))
1666            | (Rows, OffsetPreceding(..), UnboundedFollowing)
1667            | (Rows, OffsetFollowing(..), UnboundedFollowing) => {
1668                // Unsupported. Bail in the planner.
1669                // https://github.com/MaterializeInc/database-issues/issues/6720
1670                unreachable!()
1671            }
1672            (Range, _, _) => {
1673                // Unsupported.
1674                // The planner doesn't allow Range frame mode for now (except for the default
1675                // frame), see https://github.com/MaterializeInc/database-issues/issues/6585
1676                // Note that it would be easy to handle (Range, CurrentRow, UnboundedFollowing):
1677                // it would be similar to (Rows, CurrentRow, UnboundedFollowing), but would call
1678                // groups_between_unbounded_preceding_current_row.
1679                unreachable!()
1680            }
1681            (Groups, _, _) => {
1682                // Unsupported.
1683                // The planner doesn't allow Groups frame mode for now, see
1684                // https://github.com/MaterializeInc/database-issues/issues/6588
1685                unreachable!()
1686            }
1687        }
1688    }
1689
1690    result
1691}
1692
1693/// Computes a bundle of fused window aggregations.
1694/// The input is similar to `window_aggr`, but `InputValue` is not just a single value, but a record
1695/// where each component is the input to one of the aggregations.
1696fn fused_window_aggr<'a, I, A>(
1697    input_datums: I,
1698    callers_temp_storage: &'a RowArena,
1699    wrapped_aggregates: &Vec<AggregateFunc>,
1700    order_by: &Vec<ColumnOrder>,
1701    window_frame: &WindowFrame,
1702) -> Datum<'a>
1703where
1704    I: IntoIterator<Item = Datum<'a>>,
1705    A: OneByOneAggr,
1706{
1707    let temp_storage = RowArena::new();
1708    let iter = fused_window_aggr_no_list::<_, A>(
1709        input_datums,
1710        &temp_storage,
1711        wrapped_aggregates,
1712        order_by,
1713        window_frame,
1714    );
1715    callers_temp_storage.make_datum(|packer| {
1716        packer.push_list(iter);
1717    })
1718}
1719
1720/// Like `fused_window_aggr`, but doesn't perform the final wrapping in a list, returning an
1721/// Iterator instead.
1722fn fused_window_aggr_no_list<'a: 'b, 'b, I, A>(
1723    input_datums: I,
1724    callers_temp_storage: &'b RowArena,
1725    wrapped_aggregates: &Vec<AggregateFunc>,
1726    order_by: &Vec<ColumnOrder>,
1727    window_frame: &WindowFrame,
1728) -> impl Iterator<Item = Datum<'b>>
1729where
1730    I: IntoIterator<Item = Datum<'a>>,
1731    A: OneByOneAggr,
1732{
1733    // Sort the datums according to the ORDER BY expressions and return the ((OriginalRow, InputValue), OrderByRow) record
1734    // The OrderByRow is kept around because it is required to compute the peer groups in RANGE mode
1735    let datums = order_aggregate_datums_with_rank(input_datums, order_by);
1736
1737    let size_hint = datums.size_hint().0;
1738    let mut argss = vec![Vec::with_capacity(size_hint); wrapped_aggregates.len()];
1739    let mut original_rows = Vec::with_capacity(size_hint);
1740    let mut order_by_rows = Vec::with_capacity(size_hint);
1741    for (d, order_by_row) in datums {
1742        let mut iter = d.unwrap_list().iter();
1743        let original_row = iter.next().unwrap();
1744        original_rows.push(original_row);
1745        let args_iter = iter.next().unwrap().unwrap_list().iter();
1746        // Push each argument into the respective list
1747        for (args, arg) in argss.iter_mut().zip_eq(args_iter) {
1748            args.push(arg);
1749        }
1750        order_by_rows.push(order_by_row);
1751    }
1752
1753    let mut results_per_row =
1754        vec![Vec::with_capacity(wrapped_aggregates.len()); original_rows.len()];
1755    for (wrapped_aggr, args) in wrapped_aggregates.iter().zip_eq(argss) {
1756        let results = window_aggr_inner::<A>(
1757            args,
1758            &order_by_rows,
1759            wrapped_aggr,
1760            order_by,
1761            window_frame,
1762            callers_temp_storage,
1763        );
1764        for (results, result) in results_per_row.iter_mut().zip_eq(results) {
1765            results.push(result);
1766        }
1767    }
1768
1769    callers_temp_storage.reserve(2 * original_rows.len());
1770    results_per_row
1771        .into_iter()
1772        .enumerate()
1773        .map(move |(i, results)| {
1774            callers_temp_storage.make_datum(|packer| {
1775                packer.push_list_with(|packer| {
1776                    packer
1777                        .push(callers_temp_storage.make_datum(|packer| packer.push_list(results)));
1778                    packer.push(original_rows[i]);
1779                });
1780            })
1781        })
1782}
1783
1784/// An implementation of an aggregation where we can send in the input elements one-by-one, and
1785/// can also ask the current aggregate at any moment. (This just delegates to other aggregation
1786/// evaluation approaches.)
1787pub trait OneByOneAggr {
1788    /// The `reverse` parameter makes the aggregations process input elements in reverse order.
1789    /// This has an effect only for non-commutative aggregations, e.g. `list_agg`. These are
1790    /// currently only some of the Basic aggregations. (Basic aggregations are handled by
1791    /// `NaiveOneByOneAggr`).
1792    fn new(agg: &AggregateFunc, reverse: bool) -> Self;
1793    /// Pushes one input element into the aggregation.
1794    fn give(&mut self, d: &Datum);
1795    /// Returns the value of the aggregate computed on the given values so far.
1796    fn get_current_aggregate<'a>(&self, temp_storage: &'a RowArena) -> Datum<'a>;
1797}
1798
1799/// Naive implementation of [OneByOneAggr], suitable for stuff like const folding, but too slow for
1800/// rendering. This relies only on infrastructure available in `mz-expr`. It simply saves all the
1801/// given input, and calls the given [AggregateFunc]'s `eval` method when asked about the current
1802/// aggregate. (For Accumulable and Hierarchical aggregations, the rendering has more efficient
1803/// implementations, but for Basic aggregations even the rendering uses this naive implementation.)
1804#[derive(Debug)]
1805pub struct NaiveOneByOneAggr {
1806    agg: AggregateFunc,
1807    input: Vec<Row>,
1808    reverse: bool,
1809}
1810
1811impl OneByOneAggr for NaiveOneByOneAggr {
1812    fn new(agg: &AggregateFunc, reverse: bool) -> Self {
1813        NaiveOneByOneAggr {
1814            agg: agg.clone(),
1815            input: Vec::new(),
1816            reverse,
1817        }
1818    }
1819
1820    fn give(&mut self, d: &Datum) {
1821        let mut row = Row::default();
1822        row.packer().push(d);
1823        self.input.push(row);
1824    }
1825
1826    fn get_current_aggregate<'a>(&self, temp_storage: &'a RowArena) -> Datum<'a> {
1827        temp_storage.make_datum(|packer| {
1828            packer.push(if !self.reverse {
1829                self.agg.eval(
1830                    self.input.iter().map(|r| (r.unpack_first(), Diff::ONE)),
1831                    temp_storage,
1832                )
1833            } else {
1834                self.agg.eval(
1835                    self.input
1836                        .iter()
1837                        .rev()
1838                        .map(|r| (r.unpack_first(), Diff::ONE)),
1839                    temp_storage,
1840                )
1841            });
1842        })
1843    }
1844}
1845
1846/// Identify whether the given aggregate function is Lag or Lead, since they share
1847/// implementations.
1848#[derive(
1849    Clone,
1850    Debug,
1851    Eq,
1852    PartialEq,
1853    Ord,
1854    PartialOrd,
1855    Serialize,
1856    Deserialize,
1857    Hash
1858)]
1859pub enum LagLeadType {
1860    Lag,
1861    Lead,
1862}
1863
1864#[derive(
1865    Clone,
1866    Debug,
1867    Eq,
1868    PartialEq,
1869    Ord,
1870    PartialOrd,
1871    Serialize,
1872    Deserialize,
1873    Hash
1874)]
1875pub enum AggregateFunc {
1876    MaxNumeric,
1877    MaxInt16,
1878    MaxInt32,
1879    MaxInt64,
1880    MaxUInt16,
1881    MaxUInt32,
1882    MaxUInt64,
1883    MaxMzTimestamp,
1884    MaxFloat32,
1885    MaxFloat64,
1886    MaxBool,
1887    MaxString,
1888    MaxDate,
1889    MaxTimestamp,
1890    MaxTimestampTz,
1891    MaxInterval,
1892    MaxTime,
1893    MinNumeric,
1894    MinInt16,
1895    MinInt32,
1896    MinInt64,
1897    MinUInt16,
1898    MinUInt32,
1899    MinUInt64,
1900    MinMzTimestamp,
1901    MinFloat32,
1902    MinFloat64,
1903    MinBool,
1904    MinString,
1905    MinDate,
1906    MinTimestamp,
1907    MinTimestampTz,
1908    MinInterval,
1909    MinTime,
1910    SumInt16,
1911    SumInt32,
1912    SumInt64,
1913    SumUInt16,
1914    SumUInt32,
1915    SumUInt64,
1916    SumFloat32,
1917    SumFloat64,
1918    SumNumeric,
1919    Count,
1920    Any,
1921    All,
1922    /// Accumulates `Datum::List`s whose first element is a JSON-typed `Datum`s
1923    /// into a JSON list. The other elements are columns used by `order_by`.
1924    ///
1925    /// WARNING: Unlike the `jsonb_agg` function that is exposed by the SQL
1926    /// layer, this function filters out `Datum::Null`, for consistency with
1927    /// the other aggregate functions.
1928    JsonbAgg {
1929        order_by: Vec<ColumnOrder>,
1930    },
1931    /// Zips `Datum::List`s whose first element is a JSON-typed `Datum`s into a
1932    /// JSON map. The other elements are columns used by `order_by`.
1933    ///
1934    /// WARNING: Unlike the `jsonb_object_agg` function that is exposed by the SQL
1935    /// layer, this function filters out `Datum::Null`, for consistency with
1936    /// the other aggregate functions.
1937    JsonbObjectAgg {
1938        order_by: Vec<ColumnOrder>,
1939    },
1940    /// Zips a `Datum::List` whose first element is a `Datum::List` guaranteed
1941    /// to be non-empty and whose len % 2 == 0 into a `Datum::Map`. The other
1942    /// elements are columns used by `order_by`.
1943    MapAgg {
1944        order_by: Vec<ColumnOrder>,
1945        value_type: SqlScalarType,
1946    },
1947    /// Accumulates `Datum::Array`s of `SqlScalarType::Record` whose first element is a `Datum::Array`
1948    /// into a single `Datum::Array` (the remaining fields are used by `order_by`).
1949    ArrayConcat {
1950        order_by: Vec<ColumnOrder>,
1951    },
1952    /// Accumulates `Datum::List`s of `SqlScalarType::Record` whose first field is a `Datum::List`
1953    /// into a single `Datum::List` (the remaining fields are used by `order_by`).
1954    ListConcat {
1955        order_by: Vec<ColumnOrder>,
1956    },
1957    StringAgg {
1958        order_by: Vec<ColumnOrder>,
1959    },
1960    RowNumber {
1961        order_by: Vec<ColumnOrder>,
1962    },
1963    Rank {
1964        order_by: Vec<ColumnOrder>,
1965    },
1966    DenseRank {
1967        order_by: Vec<ColumnOrder>,
1968    },
1969    LagLead {
1970        order_by: Vec<ColumnOrder>,
1971        lag_lead: LagLeadType,
1972        ignore_nulls: bool,
1973    },
1974    FirstValue {
1975        order_by: Vec<ColumnOrder>,
1976        window_frame: WindowFrame,
1977    },
1978    LastValue {
1979        order_by: Vec<ColumnOrder>,
1980        window_frame: WindowFrame,
1981    },
1982    /// Several value window functions fused into one function, to amortize overheads.
1983    FusedValueWindowFunc {
1984        funcs: Vec<AggregateFunc>,
1985        /// Currently, all the fused functions must have the same `order_by`. (We can later
1986        /// eliminate this limitation.)
1987        order_by: Vec<ColumnOrder>,
1988    },
1989    WindowAggregate {
1990        wrapped_aggregate: Box<AggregateFunc>,
1991        order_by: Vec<ColumnOrder>,
1992        window_frame: WindowFrame,
1993    },
1994    FusedWindowAggregate {
1995        wrapped_aggregates: Vec<AggregateFunc>,
1996        order_by: Vec<ColumnOrder>,
1997        window_frame: WindowFrame,
1998    },
1999    /// Accumulates any number of `Datum::Dummy`s into `Datum::Dummy`.
2000    ///
2001    /// Useful for removing an expensive aggregation while maintaining the shape
2002    /// of a reduce operator.
2003    Dummy,
2004}
2005
2006/// Expands an iterator of `(datum, diff)` into one `datum` per unit of `diff`.
2007///
2008/// A non-positive `diff` contributes no copies. This is used by aggregates that
2009/// are sensitive to multiplicity (e.g. `sum`), to recover a flat datum stream
2010/// from the count-aware surface.
2011fn expand_counts<'a, I>(datums: I) -> impl Iterator<Item = Datum<'a>>
2012where
2013    I: IntoIterator<Item = (Datum<'a>, Diff)>,
2014{
2015    datums.into_iter().flat_map(|(datum, diff)| {
2016        let copies = usize::try_from(diff.into_inner()).unwrap_or(0);
2017        std::iter::repeat(datum).take(copies)
2018    })
2019}
2020
2021impl AggregateFunc {
2022    /// Whether this aggregate's result is independent of the multiplicity of its
2023    /// inputs (e.g. `min`/`max`/`any`/`all`).
2024    ///
2025    /// Such aggregates can ignore the `diff` of each input, evaluating over the
2026    /// distinct datums rather than expanding by count. This keeps idempotent
2027    /// reductions linear in the number of distinct inputs.
2028    fn ignores_multiplicity(&self) -> bool {
2029        use AggregateFunc::*;
2030        matches!(
2031            self,
2032            MaxNumeric
2033                | MaxInt16
2034                | MaxInt32
2035                | MaxInt64
2036                | MaxUInt16
2037                | MaxUInt32
2038                | MaxUInt64
2039                | MaxMzTimestamp
2040                | MaxFloat32
2041                | MaxFloat64
2042                | MaxBool
2043                | MaxString
2044                | MaxDate
2045                | MaxTimestamp
2046                | MaxTimestampTz
2047                | MaxInterval
2048                | MaxTime
2049                | MinNumeric
2050                | MinInt16
2051                | MinInt32
2052                | MinInt64
2053                | MinUInt16
2054                | MinUInt32
2055                | MinUInt64
2056                | MinMzTimestamp
2057                | MinFloat32
2058                | MinFloat64
2059                | MinBool
2060                | MinString
2061                | MinDate
2062                | MinTimestamp
2063                | MinTimestampTz
2064                | MinInterval
2065                | MinTime
2066                | Any
2067                | All
2068        )
2069    }
2070
2071    /// Evaluates the aggregate over an iterator of `(datum, diff)` pairs.
2072    ///
2073    /// Each aggregate consumes the multiplicity (`diff`) in whatever way is most
2074    /// efficient: `count` sums the diffs, multiplicity-insensitive aggregates
2075    /// (see `AggregateFunc::ignores_multiplicity`) ignore them, and everything
2076    /// else expands each datum into `diff` copies (see `expand_counts`).
2077    pub fn eval<'a, I>(&self, datums: I, temp_storage: &'a RowArena) -> Datum<'a>
2078    where
2079        I: IntoIterator<Item = (Datum<'a>, Diff)>,
2080    {
2081        // Accumulable aggregates consume multiplicity directly rather than
2082        // expanding each `(datum, diff)` into `diff` copies. The cases handled
2083        // here mirror the dataflow's accumulable reduction (`build_accumulable`
2084        // in `mz_compute::render::reduce`) so that constant folding produces the
2085        // same result the dataflow would. Signed integer sums are folded here;
2086        // unsigned sums are not, because their negative-accumulation case is a
2087        // query error in the dataflow that this `Datum`-returning path cannot
2088        // signal. Floats and numerics use bespoke fixed-point/wide-decimal
2089        // accumulators in the dataflow that `expand_counts` does not reproduce.
2090        match self {
2091            AggregateFunc::Count => count(datums),
2092            AggregateFunc::SumInt16 | AggregateFunc::SumInt32 => {
2093                // `finalize_accum` narrows these to `i64` with wrapping.
2094                sum_signed_int_counted(datums, |accum| {
2095                    #[allow(clippy::as_conversions)]
2096                    let narrowed = accum as i64;
2097                    Datum::Int64(narrowed)
2098                })
2099            }
2100            AggregateFunc::SumInt64 => sum_signed_int_counted(datums, Datum::from),
2101            _ if self.ignores_multiplicity() => {
2102                self.eval_datums(datums.into_iter().map(|(datum, _diff)| datum), temp_storage)
2103            }
2104            _ => self.eval_datums(expand_counts(datums), temp_storage),
2105        }
2106    }
2107
2108    /// Evaluates the aggregate over a flat iterator of datums, ignoring multiplicity.
2109    fn eval_datums<'a, I>(&self, datums: I, temp_storage: &'a RowArena) -> Datum<'a>
2110    where
2111        I: IntoIterator<Item = Datum<'a>>,
2112    {
2113        match self {
2114            AggregateFunc::MaxNumeric => {
2115                max_datum::<'a, I, OrderedDecimal<numeric::Numeric>>(datums)
2116            }
2117            AggregateFunc::MaxInt16 => max_datum::<'a, I, i16>(datums),
2118            AggregateFunc::MaxInt32 => max_datum::<'a, I, i32>(datums),
2119            AggregateFunc::MaxInt64 => max_datum::<'a, I, i64>(datums),
2120            AggregateFunc::MaxUInt16 => max_datum::<'a, I, u16>(datums),
2121            AggregateFunc::MaxUInt32 => max_datum::<'a, I, u32>(datums),
2122            AggregateFunc::MaxUInt64 => max_datum::<'a, I, u64>(datums),
2123            AggregateFunc::MaxMzTimestamp => max_datum::<'a, I, mz_repr::Timestamp>(datums),
2124            AggregateFunc::MaxFloat32 => max_datum::<'a, I, OrderedFloat<f32>>(datums),
2125            AggregateFunc::MaxFloat64 => max_datum::<'a, I, OrderedFloat<f64>>(datums),
2126            AggregateFunc::MaxBool => max_datum::<'a, I, bool>(datums),
2127            AggregateFunc::MaxString => max_string(datums),
2128            AggregateFunc::MaxDate => max_datum::<'a, I, Date>(datums),
2129            AggregateFunc::MaxTimestamp => {
2130                max_datum::<'a, I, CheckedTimestamp<NaiveDateTime>>(datums)
2131            }
2132            AggregateFunc::MaxTimestampTz => {
2133                max_datum::<'a, I, CheckedTimestamp<DateTime<Utc>>>(datums)
2134            }
2135            AggregateFunc::MaxInterval => max_datum::<'a, I, Interval>(datums),
2136            AggregateFunc::MaxTime => max_datum::<'a, I, NaiveTime>(datums),
2137            AggregateFunc::MinNumeric => {
2138                min_datum::<'a, I, OrderedDecimal<numeric::Numeric>>(datums)
2139            }
2140            AggregateFunc::MinInt16 => min_datum::<'a, I, i16>(datums),
2141            AggregateFunc::MinInt32 => min_datum::<'a, I, i32>(datums),
2142            AggregateFunc::MinInt64 => min_datum::<'a, I, i64>(datums),
2143            AggregateFunc::MinUInt16 => min_datum::<'a, I, u16>(datums),
2144            AggregateFunc::MinUInt32 => min_datum::<'a, I, u32>(datums),
2145            AggregateFunc::MinUInt64 => min_datum::<'a, I, u64>(datums),
2146            AggregateFunc::MinMzTimestamp => min_datum::<'a, I, mz_repr::Timestamp>(datums),
2147            AggregateFunc::MinFloat32 => min_datum::<'a, I, OrderedFloat<f32>>(datums),
2148            AggregateFunc::MinFloat64 => min_datum::<'a, I, OrderedFloat<f64>>(datums),
2149            AggregateFunc::MinBool => min_datum::<'a, I, bool>(datums),
2150            AggregateFunc::MinString => min_string(datums),
2151            AggregateFunc::MinDate => min_datum::<'a, I, Date>(datums),
2152            AggregateFunc::MinTimestamp => {
2153                min_datum::<'a, I, CheckedTimestamp<NaiveDateTime>>(datums)
2154            }
2155            AggregateFunc::MinTimestampTz => {
2156                min_datum::<'a, I, CheckedTimestamp<DateTime<Utc>>>(datums)
2157            }
2158            AggregateFunc::MinInterval => min_datum::<'a, I, Interval>(datums),
2159            AggregateFunc::MinTime => min_datum::<'a, I, NaiveTime>(datums),
2160            AggregateFunc::SumInt16 => sum_datum::<'a, I, i16, i64>(datums),
2161            AggregateFunc::SumInt32 => sum_datum::<'a, I, i32, i64>(datums),
2162            AggregateFunc::SumInt64 => sum_datum::<'a, I, i64, i128>(datums),
2163            AggregateFunc::SumUInt16 => sum_datum::<'a, I, u16, u64>(datums),
2164            AggregateFunc::SumUInt32 => sum_datum::<'a, I, u32, u64>(datums),
2165            AggregateFunc::SumUInt64 => sum_datum::<'a, I, u64, u128>(datums),
2166            AggregateFunc::SumFloat32 => sum_datum::<'a, I, f32, f32>(datums),
2167            AggregateFunc::SumFloat64 => sum_datum::<'a, I, f64, f64>(datums),
2168            AggregateFunc::SumNumeric => sum_numeric(datums),
2169            AggregateFunc::Count => unreachable!("Count is handled in `eval`"),
2170            AggregateFunc::Any => any(datums),
2171            AggregateFunc::All => all(datums),
2172            AggregateFunc::JsonbAgg { order_by } => jsonb_agg(datums, temp_storage, order_by),
2173            AggregateFunc::MapAgg { order_by, .. } | AggregateFunc::JsonbObjectAgg { order_by } => {
2174                dict_agg(datums, temp_storage, order_by)
2175            }
2176            AggregateFunc::ArrayConcat { order_by } => array_concat(datums, temp_storage, order_by),
2177            AggregateFunc::ListConcat { order_by } => list_concat(datums, temp_storage, order_by),
2178            AggregateFunc::StringAgg { order_by } => string_agg(datums, temp_storage, order_by),
2179            AggregateFunc::RowNumber { order_by } => row_number(datums, temp_storage, order_by),
2180            AggregateFunc::Rank { order_by } => rank(datums, temp_storage, order_by),
2181            AggregateFunc::DenseRank { order_by } => dense_rank(datums, temp_storage, order_by),
2182            AggregateFunc::LagLead {
2183                order_by,
2184                lag_lead: lag_lead_type,
2185                ignore_nulls,
2186            } => lag_lead(datums, temp_storage, order_by, lag_lead_type, ignore_nulls),
2187            AggregateFunc::FirstValue {
2188                order_by,
2189                window_frame,
2190            } => first_value(datums, temp_storage, order_by, window_frame),
2191            AggregateFunc::LastValue {
2192                order_by,
2193                window_frame,
2194            } => last_value(datums, temp_storage, order_by, window_frame),
2195            AggregateFunc::WindowAggregate {
2196                wrapped_aggregate,
2197                order_by,
2198                window_frame,
2199            } => window_aggr::<_, NaiveOneByOneAggr>(
2200                datums,
2201                temp_storage,
2202                wrapped_aggregate,
2203                order_by,
2204                window_frame,
2205            ),
2206            AggregateFunc::FusedValueWindowFunc { funcs, order_by } => {
2207                fused_value_window_func(datums, temp_storage, funcs, order_by)
2208            }
2209            AggregateFunc::FusedWindowAggregate {
2210                wrapped_aggregates,
2211                order_by,
2212                window_frame,
2213            } => fused_window_aggr::<_, NaiveOneByOneAggr>(
2214                datums,
2215                temp_storage,
2216                wrapped_aggregates,
2217                order_by,
2218                window_frame,
2219            ),
2220            AggregateFunc::Dummy => Datum::Dummy,
2221        }
2222    }
2223
2224    /// Like `eval`, but it's given a [OneByOneAggr]. If `self` is a `WindowAggregate`, then
2225    /// the given [OneByOneAggr] will be used to evaluate the wrapped aggregate inside the
2226    /// `WindowAggregate`. If `self` is not a `WindowAggregate`, then it simply calls `eval`.
2227    pub fn eval_with_fast_window_agg<'a, I, W>(
2228        &self,
2229        datums: I,
2230        temp_storage: &'a RowArena,
2231    ) -> Datum<'a>
2232    where
2233        I: IntoIterator<Item = (Datum<'a>, Diff)>,
2234        W: OneByOneAggr,
2235    {
2236        match self {
2237            AggregateFunc::WindowAggregate {
2238                wrapped_aggregate,
2239                order_by,
2240                window_frame,
2241            } => window_aggr::<_, W>(
2242                expand_counts(datums),
2243                temp_storage,
2244                wrapped_aggregate,
2245                order_by,
2246                window_frame,
2247            ),
2248            AggregateFunc::FusedWindowAggregate {
2249                wrapped_aggregates,
2250                order_by,
2251                window_frame,
2252            } => fused_window_aggr::<_, W>(
2253                expand_counts(datums),
2254                temp_storage,
2255                wrapped_aggregates,
2256                order_by,
2257                window_frame,
2258            ),
2259            _ => self.eval(datums, temp_storage),
2260        }
2261    }
2262
2263    pub fn eval_with_unnest_list<'a, I, W>(
2264        &self,
2265        datums: I,
2266        temp_storage: &'a RowArena,
2267    ) -> impl Iterator<Item = Datum<'a>>
2268    where
2269        I: IntoIterator<Item = (Datum<'a>, Diff)>,
2270        W: OneByOneAggr,
2271    {
2272        // TODO: Use `enum_dispatch` to construct a unified iterator instead of `collect_vec`.
2273        assert!(self.can_fuse_with_unnest_list());
2274        // Window functions are sensitive to multiplicity, so expand counts.
2275        let datums = expand_counts(datums);
2276        match self {
2277            AggregateFunc::RowNumber { order_by } => {
2278                row_number_no_list(datums, temp_storage, order_by).collect_vec()
2279            }
2280            AggregateFunc::Rank { order_by } => {
2281                rank_no_list(datums, temp_storage, order_by).collect_vec()
2282            }
2283            AggregateFunc::DenseRank { order_by } => {
2284                dense_rank_no_list(datums, temp_storage, order_by).collect_vec()
2285            }
2286            AggregateFunc::LagLead {
2287                order_by,
2288                lag_lead: lag_lead_type,
2289                ignore_nulls,
2290            } => lag_lead_no_list(datums, temp_storage, order_by, lag_lead_type, ignore_nulls)
2291                .collect_vec(),
2292            AggregateFunc::FirstValue {
2293                order_by,
2294                window_frame,
2295            } => first_value_no_list(datums, temp_storage, order_by, window_frame).collect_vec(),
2296            AggregateFunc::LastValue {
2297                order_by,
2298                window_frame,
2299            } => last_value_no_list(datums, temp_storage, order_by, window_frame).collect_vec(),
2300            AggregateFunc::FusedValueWindowFunc { funcs, order_by } => {
2301                fused_value_window_func_no_list(datums, temp_storage, funcs, order_by).collect_vec()
2302            }
2303            AggregateFunc::WindowAggregate {
2304                wrapped_aggregate,
2305                order_by,
2306                window_frame,
2307            } => window_aggr_no_list::<_, W>(
2308                datums,
2309                temp_storage,
2310                wrapped_aggregate,
2311                order_by,
2312                window_frame,
2313            )
2314            .collect_vec(),
2315            AggregateFunc::FusedWindowAggregate {
2316                wrapped_aggregates,
2317                order_by,
2318                window_frame,
2319            } => fused_window_aggr_no_list::<_, W>(
2320                datums,
2321                temp_storage,
2322                wrapped_aggregates,
2323                order_by,
2324                window_frame,
2325            )
2326            .collect_vec(),
2327            _ => unreachable!("asserted above that `can_fuse_with_unnest_list`"),
2328        }
2329        .into_iter()
2330    }
2331
2332    /// Returns the output of the aggregation function when applied on an empty
2333    /// input relation.
2334    pub fn default(&self) -> Datum<'static> {
2335        match self {
2336            AggregateFunc::Count => Datum::Int64(0),
2337            AggregateFunc::Any => Datum::False,
2338            AggregateFunc::All => Datum::True,
2339            AggregateFunc::Dummy => Datum::Dummy,
2340            _ => Datum::Null,
2341        }
2342    }
2343
2344    /// Returns a datum whose inclusion in the aggregation will not change its
2345    /// result.
2346    pub fn identity_datum(&self) -> Datum<'static> {
2347        match self {
2348            AggregateFunc::Any => Datum::False,
2349            AggregateFunc::All => Datum::True,
2350            AggregateFunc::Dummy => Datum::Dummy,
2351            AggregateFunc::ArrayConcat { .. } => Datum::empty_array(),
2352            AggregateFunc::ListConcat { .. } => Datum::empty_list(),
2353            AggregateFunc::RowNumber { .. }
2354            | AggregateFunc::Rank { .. }
2355            | AggregateFunc::DenseRank { .. }
2356            | AggregateFunc::LagLead { .. }
2357            | AggregateFunc::FirstValue { .. }
2358            | AggregateFunc::LastValue { .. }
2359            | AggregateFunc::WindowAggregate { .. }
2360            | AggregateFunc::FusedValueWindowFunc { .. }
2361            | AggregateFunc::FusedWindowAggregate { .. } => Datum::empty_list(),
2362            AggregateFunc::MaxNumeric
2363            | AggregateFunc::MaxInt16
2364            | AggregateFunc::MaxInt32
2365            | AggregateFunc::MaxInt64
2366            | AggregateFunc::MaxUInt16
2367            | AggregateFunc::MaxUInt32
2368            | AggregateFunc::MaxUInt64
2369            | AggregateFunc::MaxMzTimestamp
2370            | AggregateFunc::MaxFloat32
2371            | AggregateFunc::MaxFloat64
2372            | AggregateFunc::MaxBool
2373            | AggregateFunc::MaxString
2374            | AggregateFunc::MaxDate
2375            | AggregateFunc::MaxTimestamp
2376            | AggregateFunc::MaxTimestampTz
2377            | AggregateFunc::MaxInterval
2378            | AggregateFunc::MaxTime
2379            | AggregateFunc::MinNumeric
2380            | AggregateFunc::MinInt16
2381            | AggregateFunc::MinInt32
2382            | AggregateFunc::MinInt64
2383            | AggregateFunc::MinUInt16
2384            | AggregateFunc::MinUInt32
2385            | AggregateFunc::MinUInt64
2386            | AggregateFunc::MinMzTimestamp
2387            | AggregateFunc::MinFloat32
2388            | AggregateFunc::MinFloat64
2389            | AggregateFunc::MinBool
2390            | AggregateFunc::MinString
2391            | AggregateFunc::MinDate
2392            | AggregateFunc::MinTimestamp
2393            | AggregateFunc::MinTimestampTz
2394            | AggregateFunc::MinInterval
2395            | AggregateFunc::MinTime
2396            | AggregateFunc::SumInt16
2397            | AggregateFunc::SumInt32
2398            | AggregateFunc::SumInt64
2399            | AggregateFunc::SumUInt16
2400            | AggregateFunc::SumUInt32
2401            | AggregateFunc::SumUInt64
2402            | AggregateFunc::SumFloat32
2403            | AggregateFunc::SumFloat64
2404            | AggregateFunc::SumNumeric
2405            | AggregateFunc::Count
2406            | AggregateFunc::JsonbAgg { .. }
2407            | AggregateFunc::JsonbObjectAgg { .. }
2408            | AggregateFunc::MapAgg { .. }
2409            | AggregateFunc::StringAgg { .. } => Datum::Null,
2410        }
2411    }
2412
2413    pub fn can_fuse_with_unnest_list(&self) -> bool {
2414        match self {
2415            AggregateFunc::RowNumber { .. }
2416            | AggregateFunc::Rank { .. }
2417            | AggregateFunc::DenseRank { .. }
2418            | AggregateFunc::LagLead { .. }
2419            | AggregateFunc::FirstValue { .. }
2420            | AggregateFunc::LastValue { .. }
2421            | AggregateFunc::WindowAggregate { .. }
2422            | AggregateFunc::FusedValueWindowFunc { .. }
2423            | AggregateFunc::FusedWindowAggregate { .. } => true,
2424            AggregateFunc::ArrayConcat { .. }
2425            | AggregateFunc::ListConcat { .. }
2426            | AggregateFunc::Any
2427            | AggregateFunc::All
2428            | AggregateFunc::Dummy
2429            | AggregateFunc::MaxNumeric
2430            | AggregateFunc::MaxInt16
2431            | AggregateFunc::MaxInt32
2432            | AggregateFunc::MaxInt64
2433            | AggregateFunc::MaxUInt16
2434            | AggregateFunc::MaxUInt32
2435            | AggregateFunc::MaxUInt64
2436            | AggregateFunc::MaxMzTimestamp
2437            | AggregateFunc::MaxFloat32
2438            | AggregateFunc::MaxFloat64
2439            | AggregateFunc::MaxBool
2440            | AggregateFunc::MaxString
2441            | AggregateFunc::MaxDate
2442            | AggregateFunc::MaxTimestamp
2443            | AggregateFunc::MaxTimestampTz
2444            | AggregateFunc::MaxInterval
2445            | AggregateFunc::MaxTime
2446            | AggregateFunc::MinNumeric
2447            | AggregateFunc::MinInt16
2448            | AggregateFunc::MinInt32
2449            | AggregateFunc::MinInt64
2450            | AggregateFunc::MinUInt16
2451            | AggregateFunc::MinUInt32
2452            | AggregateFunc::MinUInt64
2453            | AggregateFunc::MinMzTimestamp
2454            | AggregateFunc::MinFloat32
2455            | AggregateFunc::MinFloat64
2456            | AggregateFunc::MinBool
2457            | AggregateFunc::MinString
2458            | AggregateFunc::MinDate
2459            | AggregateFunc::MinTimestamp
2460            | AggregateFunc::MinTimestampTz
2461            | AggregateFunc::MinInterval
2462            | AggregateFunc::MinTime
2463            | AggregateFunc::SumInt16
2464            | AggregateFunc::SumInt32
2465            | AggregateFunc::SumInt64
2466            | AggregateFunc::SumUInt16
2467            | AggregateFunc::SumUInt32
2468            | AggregateFunc::SumUInt64
2469            | AggregateFunc::SumFloat32
2470            | AggregateFunc::SumFloat64
2471            | AggregateFunc::SumNumeric
2472            | AggregateFunc::Count
2473            | AggregateFunc::JsonbAgg { .. }
2474            | AggregateFunc::JsonbObjectAgg { .. }
2475            | AggregateFunc::MapAgg { .. }
2476            | AggregateFunc::StringAgg { .. } => false,
2477        }
2478    }
2479
2480    /// The output column type for the result of an aggregation.
2481    ///
2482    /// The output column type also contains nullability information, which
2483    /// is (without further information) true for aggregations that are not
2484    /// counts.
2485    pub fn output_sql_type(&self, input_type: SqlColumnType) -> SqlColumnType {
2486        let scalar_type = match self {
2487            AggregateFunc::Count => SqlScalarType::Int64,
2488            AggregateFunc::Any => SqlScalarType::Bool,
2489            AggregateFunc::All => SqlScalarType::Bool,
2490            AggregateFunc::JsonbAgg { .. } => SqlScalarType::Jsonb,
2491            AggregateFunc::JsonbObjectAgg { .. } => SqlScalarType::Jsonb,
2492            AggregateFunc::SumInt16 => SqlScalarType::Int64,
2493            AggregateFunc::SumInt32 => SqlScalarType::Int64,
2494            AggregateFunc::SumInt64 => SqlScalarType::Numeric {
2495                max_scale: Some(NumericMaxScale::ZERO),
2496            },
2497            AggregateFunc::SumUInt16 => SqlScalarType::UInt64,
2498            AggregateFunc::SumUInt32 => SqlScalarType::UInt64,
2499            AggregateFunc::SumUInt64 => SqlScalarType::Numeric {
2500                max_scale: Some(NumericMaxScale::ZERO),
2501            },
2502            AggregateFunc::MapAgg { value_type, .. } => SqlScalarType::Map {
2503                value_type: Box::new(value_type.clone()),
2504                custom_id: None,
2505            },
2506            AggregateFunc::ArrayConcat { .. } | AggregateFunc::ListConcat { .. } => {
2507                match input_type.scalar_type {
2508                    // The input is wrapped in a Record if there's an ORDER BY, so extract it out.
2509                    SqlScalarType::Record { ref fields, .. } => fields[0].1.scalar_type.clone(),
2510                    _ => unreachable!(),
2511                }
2512            }
2513            AggregateFunc::StringAgg { .. } => SqlScalarType::String,
2514            AggregateFunc::RowNumber { .. } => {
2515                AggregateFunc::output_type_ranking_window_funcs(&input_type, "?row_number?")
2516            }
2517            AggregateFunc::Rank { .. } => {
2518                AggregateFunc::output_type_ranking_window_funcs(&input_type, "?rank?")
2519            }
2520            AggregateFunc::DenseRank { .. } => {
2521                AggregateFunc::output_type_ranking_window_funcs(&input_type, "?dense_rank?")
2522            }
2523            AggregateFunc::LagLead { lag_lead: lag_lead_type, .. } => {
2524                // The input type for Lag is ((OriginalRow, EncodedArgs), OrderByExprs...)
2525                let fields = input_type.scalar_type.unwrap_record_element_type();
2526                let original_row_type = fields[0].unwrap_record_element_type()[0]
2527                    .clone()
2528                    .nullable(false);
2529                let encoded_args = fields[0].unwrap_record_element_type()[1];
2530                let output_type_inner =
2531                    Self::lag_lead_output_type_inner_from_encoded_args(encoded_args);
2532                let column_name = Self::lag_lead_result_column_name(lag_lead_type);
2533
2534                SqlScalarType::List {
2535                    element_type: Box::new(SqlScalarType::Record {
2536                        fields: [
2537                            (column_name, output_type_inner),
2538                            (ColumnName::from("?orig_row?"), original_row_type),
2539                        ].into(),
2540                        custom_id: None,
2541                    }),
2542                    custom_id: None,
2543                }
2544            }
2545            AggregateFunc::FirstValue { .. } => {
2546                // The input type for FirstValue is ((OriginalRow, Arg), OrderByExprs...)
2547                let fields = input_type.scalar_type.unwrap_record_element_type();
2548                let original_row_type = fields[0].unwrap_record_element_type()[0]
2549                    .clone()
2550                    .nullable(false);
2551                let value_type = fields[0].unwrap_record_element_type()[1]
2552                    .clone()
2553                    .nullable(true); // null when the partition is empty
2554
2555                SqlScalarType::List {
2556                    element_type: Box::new(SqlScalarType::Record {
2557                        fields: [
2558                            (ColumnName::from("?first_value?"), value_type),
2559                            (ColumnName::from("?orig_row?"), original_row_type),
2560                        ].into(),
2561                        custom_id: None,
2562                    }),
2563                    custom_id: None,
2564                }
2565            }
2566            AggregateFunc::LastValue { .. } => {
2567                // The input type for LastValue is ((OriginalRow, Arg), OrderByExprs...)
2568                let fields = input_type.scalar_type.unwrap_record_element_type();
2569                let original_row_type = fields[0].unwrap_record_element_type()[0]
2570                    .clone()
2571                    .nullable(false);
2572                let value_type = fields[0].unwrap_record_element_type()[1]
2573                    .clone()
2574                    .nullable(true); // null when the partition is empty
2575
2576                SqlScalarType::List {
2577                    element_type: Box::new(SqlScalarType::Record {
2578                        fields: [
2579                            (ColumnName::from("?last_value?"), value_type),
2580                            (ColumnName::from("?orig_row?"), original_row_type),
2581                        ].into(),
2582                        custom_id: None,
2583                    }),
2584                    custom_id: None,
2585                }
2586            }
2587            AggregateFunc::WindowAggregate {
2588                wrapped_aggregate, ..
2589            } => {
2590                // The input type for a window aggregate is ((OriginalRow, Arg), OrderByExprs...)
2591                let fields = input_type.scalar_type.unwrap_record_element_type();
2592                let original_row_type = fields[0].unwrap_record_element_type()[0]
2593                    .clone()
2594                    .nullable(false);
2595                let arg_type = fields[0].unwrap_record_element_type()[1]
2596                    .clone()
2597                    .nullable(true);
2598                let wrapped_aggr_out_type = wrapped_aggregate.output_sql_type(arg_type);
2599
2600                SqlScalarType::List {
2601                    element_type: Box::new(SqlScalarType::Record {
2602                        fields: [
2603                            (ColumnName::from("?window_agg?"), wrapped_aggr_out_type),
2604                            (ColumnName::from("?orig_row?"), original_row_type),
2605                        ].into(),
2606                        custom_id: None,
2607                    }),
2608                    custom_id: None,
2609                }
2610            }
2611            AggregateFunc::FusedWindowAggregate {
2612                wrapped_aggregates, ..
2613            } => {
2614                // The input type for a fused window aggregate is ((OriginalRow, Args), OrderByExprs...)
2615                // where `Args` is a record.
2616                let fields = input_type.scalar_type.unwrap_record_element_type();
2617                let original_row_type = fields[0].unwrap_record_element_type()[0]
2618                    .clone()
2619                    .nullable(false);
2620                let args_type = fields[0].unwrap_record_element_type()[1];
2621                let arg_types = args_type.unwrap_record_element_type();
2622                let out_fields = arg_types.iter().zip_eq(wrapped_aggregates).map(
2623                    |(arg_type, wrapped_agg)| {
2624                    (
2625                        ColumnName::from(wrapped_agg.name()),
2626                        wrapped_agg.output_sql_type((**arg_type).clone().nullable(true)),
2627                    )
2628                }).collect_vec();
2629
2630                SqlScalarType::List {
2631                    element_type: Box::new(SqlScalarType::Record {
2632                        fields: [
2633                            (ColumnName::from("?fused_window_agg?"), SqlScalarType::Record {
2634                                fields: out_fields.into(),
2635                                custom_id: None,
2636                            }.nullable(false)),
2637                            (ColumnName::from("?orig_row?"), original_row_type),
2638                        ].into(),
2639                        custom_id: None,
2640                    }),
2641                    custom_id: None,
2642                }
2643            }
2644            AggregateFunc::FusedValueWindowFunc { funcs, order_by: _ } => {
2645                // The input type is ((OriginalRow, EncodedArgs), OrderByExprs...)
2646                // where EncodedArgs is a record, where each element is the argument to one of the
2647                // function calls that got fused. This is a record for lag/lead, and a simple type
2648                // for first_value/last_value.
2649                let fields = input_type.scalar_type.unwrap_record_element_type();
2650                let original_row_type = fields[0].unwrap_record_element_type()[0]
2651                    .clone()
2652                    .nullable(false);
2653                let encoded_args_type = fields[0]
2654                    .unwrap_record_element_type()[1]
2655                    .unwrap_record_element_type();
2656
2657                SqlScalarType::List {
2658                    element_type: Box::new(SqlScalarType::Record {
2659                        fields: [
2660                            (
2661                                ColumnName::from("?fused_value_window_func?"),
2662                                SqlScalarType::Record {
2663                                fields: encoded_args_type.into_iter().zip_eq(funcs).map(
2664                                    |(arg_type, func)| {
2665                                    match func {
2666                                        AggregateFunc::LagLead {
2667                                            lag_lead: lag_lead_type, ..
2668                                        } => {
2669                                            let name = Self::lag_lead_result_column_name(
2670                                                lag_lead_type,
2671                                            );
2672                                            let ty = Self
2673                                                ::lag_lead_output_type_inner_from_encoded_args(
2674                                                    arg_type,
2675                                                );
2676                                            (name, ty)
2677                                        },
2678                                        AggregateFunc::FirstValue { .. } => {
2679                                            (
2680                                                ColumnName::from("?first_value?"),
2681                                                arg_type.clone().nullable(true),
2682                                            )
2683                                        }
2684                                        AggregateFunc::LastValue { .. } => {
2685                                            (
2686                                                ColumnName::from("?last_value?"),
2687                                                arg_type.clone().nullable(true),
2688                                            )
2689                                        }
2690                                        _ => panic!("FusedValueWindowFunc has an unknown function"),
2691                                    }
2692                                }).collect(),
2693                                custom_id: None,
2694                            }.nullable(false)),
2695                            (ColumnName::from("?orig_row?"), original_row_type),
2696                        ].into(),
2697                        custom_id: None,
2698                    }),
2699                    custom_id: None,
2700                }
2701            }
2702            AggregateFunc::Dummy
2703            | AggregateFunc::MaxNumeric
2704            | AggregateFunc::MaxInt16
2705            | AggregateFunc::MaxInt32
2706            | AggregateFunc::MaxInt64
2707            | AggregateFunc::MaxUInt16
2708            | AggregateFunc::MaxUInt32
2709            | AggregateFunc::MaxUInt64
2710            | AggregateFunc::MaxMzTimestamp
2711            | AggregateFunc::MaxFloat32
2712            | AggregateFunc::MaxFloat64
2713            | AggregateFunc::MaxBool
2714            // Note AggregateFunc::MaxString, MinString rely on returning input
2715            // type as output type to support the proper return type for
2716            // character input.
2717            | AggregateFunc::MaxString
2718            | AggregateFunc::MaxDate
2719            | AggregateFunc::MaxTimestamp
2720            | AggregateFunc::MaxTimestampTz
2721            | AggregateFunc::MaxInterval
2722            | AggregateFunc::MaxTime
2723            | AggregateFunc::MinNumeric
2724            | AggregateFunc::MinInt16
2725            | AggregateFunc::MinInt32
2726            | AggregateFunc::MinInt64
2727            | AggregateFunc::MinUInt16
2728            | AggregateFunc::MinUInt32
2729            | AggregateFunc::MinUInt64
2730            | AggregateFunc::MinMzTimestamp
2731            | AggregateFunc::MinFloat32
2732            | AggregateFunc::MinFloat64
2733            | AggregateFunc::MinBool
2734            | AggregateFunc::MinString
2735            | AggregateFunc::MinDate
2736            | AggregateFunc::MinTimestamp
2737            | AggregateFunc::MinTimestampTz
2738            | AggregateFunc::MinInterval
2739            | AggregateFunc::MinTime
2740            | AggregateFunc::SumFloat32
2741            | AggregateFunc::SumFloat64
2742            | AggregateFunc::SumNumeric => input_type.scalar_type.clone(),
2743        };
2744        // Count never produces null, and other aggregations only produce
2745        // null in the presence of null inputs.
2746        let nullable = match self {
2747            AggregateFunc::Count => false,
2748            // Use the nullability of the underlying column being aggregated, not the Records wrapping it
2749            AggregateFunc::StringAgg { .. } => match input_type.scalar_type {
2750                // The outer Record wraps the input in the first position, and any ORDER BY expressions afterwards
2751                SqlScalarType::Record { fields, .. } => match &fields[0].1.scalar_type {
2752                    // The inner Record is a (value, separator) tuple
2753                    SqlScalarType::Record { fields, .. } => fields[0].1.nullable,
2754                    _ => unreachable!(),
2755                },
2756                _ => unreachable!(),
2757            },
2758            _ => input_type.nullable,
2759        };
2760        scalar_type.nullable(nullable)
2761    }
2762
2763    /// Computes the representation type of this aggregate function.
2764    ///
2765    /// This is a wrapper around [`Self::output_sql_type`] that converts the result to a representation type.
2766    pub fn output_type(&self, input_type: ReprColumnType) -> ReprColumnType {
2767        ReprColumnType::from(&self.output_sql_type(SqlColumnType::from_repr(&input_type)))
2768    }
2769
2770    /// Compute output type for ROW_NUMBER, RANK, DENSE_RANK
2771    fn output_type_ranking_window_funcs(
2772        input_type: &SqlColumnType,
2773        col_name: &str,
2774    ) -> SqlScalarType {
2775        match input_type.scalar_type {
2776            SqlScalarType::Record { ref fields, .. } => SqlScalarType::List {
2777                element_type: Box::new(SqlScalarType::Record {
2778                    fields: [
2779                        (
2780                            ColumnName::from(col_name),
2781                            SqlScalarType::Int64.nullable(false),
2782                        ),
2783                        (ColumnName::from("?orig_row?"), {
2784                            let inner = match &fields[0].1.scalar_type {
2785                                SqlScalarType::List { element_type, .. } => element_type.clone(),
2786                                _ => unreachable!(),
2787                            };
2788                            inner.nullable(false)
2789                        }),
2790                    ]
2791                    .into(),
2792                    custom_id: None,
2793                }),
2794                custom_id: None,
2795            },
2796            _ => unreachable!(),
2797        }
2798    }
2799
2800    /// Given the `EncodedArgs` part of `((OriginalRow, EncodedArgs), OrderByExprs...)`,
2801    /// this computes the type of the first field of the output type. (The first field is the
2802    /// real result, the rest is the original row.)
2803    fn lag_lead_output_type_inner_from_encoded_args(
2804        encoded_args_type: &SqlScalarType,
2805    ) -> SqlColumnType {
2806        // lag/lead have 3 arguments, and the output type is
2807        // the same as the first of these, but always nullable. (It's null when the
2808        // lag/lead computation reaches over the bounds of the window partition.)
2809        encoded_args_type.unwrap_record_element_type()[0]
2810            .clone()
2811            .nullable(true)
2812    }
2813
2814    fn lag_lead_result_column_name(lag_lead_type: &LagLeadType) -> ColumnName {
2815        ColumnName::from(match lag_lead_type {
2816            LagLeadType::Lag => "?lag?",
2817            LagLeadType::Lead => "?lead?",
2818        })
2819    }
2820
2821    /// Returns true if the non-null constraint on the aggregation can be
2822    /// converted into a non-null constraint on its parameter expression, ie.
2823    /// whether the result of the aggregation is null if all the input values
2824    /// are null.
2825    pub fn propagates_nonnull_constraint(&self) -> bool {
2826        match self {
2827            AggregateFunc::MaxNumeric
2828            | AggregateFunc::MaxInt16
2829            | AggregateFunc::MaxInt32
2830            | AggregateFunc::MaxInt64
2831            | AggregateFunc::MaxUInt16
2832            | AggregateFunc::MaxUInt32
2833            | AggregateFunc::MaxUInt64
2834            | AggregateFunc::MaxMzTimestamp
2835            | AggregateFunc::MaxFloat32
2836            | AggregateFunc::MaxFloat64
2837            | AggregateFunc::MaxBool
2838            | AggregateFunc::MaxString
2839            | AggregateFunc::MaxDate
2840            | AggregateFunc::MaxTimestamp
2841            | AggregateFunc::MaxTimestampTz
2842            | AggregateFunc::MaxInterval
2843            | AggregateFunc::MaxTime
2844            | AggregateFunc::MinNumeric
2845            | AggregateFunc::MinInt16
2846            | AggregateFunc::MinInt32
2847            | AggregateFunc::MinInt64
2848            | AggregateFunc::MinUInt16
2849            | AggregateFunc::MinUInt32
2850            | AggregateFunc::MinUInt64
2851            | AggregateFunc::MinMzTimestamp
2852            | AggregateFunc::MinFloat32
2853            | AggregateFunc::MinFloat64
2854            | AggregateFunc::MinBool
2855            | AggregateFunc::MinString
2856            | AggregateFunc::MinDate
2857            | AggregateFunc::MinTimestamp
2858            | AggregateFunc::MinTimestampTz
2859            | AggregateFunc::MinInterval
2860            | AggregateFunc::MinTime
2861            | AggregateFunc::SumInt16
2862            | AggregateFunc::SumInt32
2863            | AggregateFunc::SumInt64
2864            | AggregateFunc::SumUInt16
2865            | AggregateFunc::SumUInt32
2866            | AggregateFunc::SumUInt64
2867            | AggregateFunc::SumFloat32
2868            | AggregateFunc::SumFloat64
2869            | AggregateFunc::SumNumeric
2870            | AggregateFunc::StringAgg { .. } => true,
2871            // Count is never null
2872            AggregateFunc::Count
2873            | AggregateFunc::Any
2874            | AggregateFunc::All
2875            | AggregateFunc::JsonbAgg { .. }
2876            | AggregateFunc::JsonbObjectAgg { .. }
2877            | AggregateFunc::MapAgg { .. }
2878            | AggregateFunc::ArrayConcat { .. }
2879            | AggregateFunc::ListConcat { .. }
2880            | AggregateFunc::RowNumber { .. }
2881            | AggregateFunc::Rank { .. }
2882            | AggregateFunc::DenseRank { .. }
2883            | AggregateFunc::LagLead { .. }
2884            | AggregateFunc::FirstValue { .. }
2885            | AggregateFunc::LastValue { .. }
2886            | AggregateFunc::FusedValueWindowFunc { .. }
2887            | AggregateFunc::WindowAggregate { .. }
2888            | AggregateFunc::FusedWindowAggregate { .. }
2889            | AggregateFunc::Dummy => false,
2890        }
2891    }
2892}
2893
2894fn jsonb_each<'a>(a: Datum<'a>) -> impl Iterator<Item = (Row, Diff)> + 'a {
2895    // First produce a map, so that a common iterator can be returned.
2896    let map = match a {
2897        Datum::Map(dict) => dict,
2898        _ => mz_repr::DatumMap::empty(),
2899    };
2900
2901    map.iter()
2902        .map(move |(k, v)| (Row::pack_slice(&[Datum::String(k), v]), Diff::ONE))
2903}
2904
2905fn jsonb_each_stringify<'a>(
2906    a: Datum<'a>,
2907    temp_storage: &'a RowArena,
2908) -> impl Iterator<Item = (Row, Diff)> + 'a {
2909    // First produce a map, so that a common iterator can be returned.
2910    let map = match a {
2911        Datum::Map(dict) => dict,
2912        _ => mz_repr::DatumMap::empty(),
2913    };
2914
2915    map.iter().map(move |(k, mut v)| {
2916        v = jsonb_stringify(v, temp_storage)
2917            .map(Datum::String)
2918            .unwrap_or(Datum::Null);
2919        (Row::pack_slice(&[Datum::String(k), v]), Diff::ONE)
2920    })
2921}
2922
2923fn jsonb_object_keys<'a>(a: Datum<'a>) -> impl Iterator<Item = (Row, Diff)> + 'a {
2924    let map = match a {
2925        Datum::Map(dict) => dict,
2926        _ => mz_repr::DatumMap::empty(),
2927    };
2928
2929    map.iter()
2930        .map(move |(k, _)| (Row::pack_slice(&[Datum::String(k)]), Diff::ONE))
2931}
2932
2933fn jsonb_array_elements<'a>(a: Datum<'a>) -> impl Iterator<Item = (Row, Diff)> + 'a {
2934    let list = match a {
2935        Datum::List(list) => list,
2936        _ => mz_repr::DatumList::empty(),
2937    };
2938    list.iter().map(move |e| (Row::pack_slice(&[e]), Diff::ONE))
2939}
2940
2941fn jsonb_array_elements_stringify<'a>(
2942    a: Datum<'a>,
2943    temp_storage: &'a RowArena,
2944) -> impl Iterator<Item = (Row, Diff)> + 'a {
2945    let list = match a {
2946        Datum::List(list) => list,
2947        _ => mz_repr::DatumList::empty(),
2948    };
2949    list.iter().map(move |mut e| {
2950        e = jsonb_stringify(e, temp_storage)
2951            .map(Datum::String)
2952            .unwrap_or(Datum::Null);
2953        (Row::pack_slice(&[e]), Diff::ONE)
2954    })
2955}
2956
2957fn regexp_extract(a: Datum, r: &AnalyzedRegex) -> Option<(Row, Diff)> {
2958    let r = r.inner();
2959    let a = a.unwrap_str();
2960    let captures = r.captures(a)?;
2961    let datums = captures
2962        .iter()
2963        .skip(1)
2964        .map(|m| Datum::from(m.map(|m| m.as_str())));
2965    Some((Row::pack(datums), Diff::ONE))
2966}
2967
2968fn regexp_matches<'a>(
2969    exprs: &[Datum<'a>],
2970) -> Result<impl Iterator<Item = (Row, Diff)> + 'a, EvalError> {
2971    // There are only two acceptable ways to call this function:
2972    // 1. regexp_matches(string, regex)
2973    // 2. regexp_matches(string, regex, flag)
2974    assert!(exprs.len() == 2 || exprs.len() == 3);
2975    let a = exprs[0].unwrap_str();
2976    let r = exprs[1].unwrap_str();
2977
2978    let (regex, opts) = if exprs.len() == 3 {
2979        let flag = exprs[2].unwrap_str();
2980        let opts = AnalyzedRegexOpts::from_str(flag)?;
2981        (AnalyzedRegex::new(r, opts)?, opts)
2982    } else {
2983        let opts = AnalyzedRegexOpts::default();
2984        (AnalyzedRegex::new(r, opts)?, opts)
2985    };
2986
2987    let regex = regex.inner().clone();
2988
2989    let iter = regex.captures_iter(a).map(move |captures| {
2990        let matches = captures
2991            .iter()
2992            // The first match is the *entire* match, we want the capture groups by themselves.
2993            .skip(1)
2994            .map(|m| Datum::from(m.map(|m| m.as_str())))
2995            .collect::<Vec<_>>();
2996
2997        let mut binding = SharedRow::get();
2998        let mut packer = binding.packer();
2999
3000        let dimension = ArrayDimension {
3001            lower_bound: 1,
3002            length: matches.len(),
3003        };
3004        packer
3005            .try_push_array(&[dimension], matches)
3006            .expect("generated dimensions above");
3007
3008        (binding.clone(), Diff::ONE)
3009    });
3010
3011    // This is slightly unfortunate, but we need to collect the captures into a
3012    // Vec before we can yield them, because we can't return a iter with a
3013    // reference to the local `regex` variable.
3014    // We attempt to minimize the cost of this by using a SmallVec.
3015    let out = iter.collect::<SmallVec<[_; 3]>>();
3016
3017    if opts.global {
3018        Ok(Either::Left(out.into_iter()))
3019    } else {
3020        Ok(Either::Right(out.into_iter().take(1)))
3021    }
3022}
3023
3024fn generate_series<N>(
3025    start: N,
3026    stop: N,
3027    step: N,
3028) -> Result<impl Iterator<Item = (Row, Diff)>, EvalError>
3029where
3030    N: Integer + Signed + CheckedAdd + Clone,
3031    Datum<'static>: From<N>,
3032{
3033    if step == N::zero() {
3034        return Err(EvalError::InvalidParameterValue(
3035            "step size cannot equal zero".into(),
3036        ));
3037    }
3038    Ok(num::range_step_inclusive(start, stop, step)
3039        .map(move |i| (Row::pack_slice(&[Datum::from(i)]), Diff::ONE)))
3040}
3041
3042/// Like
3043/// [`num::range_step_inclusive`](https://github.com/rust-num/num-iter/blob/ddb14c1e796d401014c6c7a727de61d8109ad986/src/lib.rs#L279),
3044/// but for our timestamp types using [`Interval`] for `step`.xwxw
3045#[derive(Clone)]
3046pub struct TimestampRangeStepInclusive<T> {
3047    state: CheckedTimestamp<T>,
3048    stop: CheckedTimestamp<T>,
3049    step: Interval,
3050    rev: bool,
3051    done: bool,
3052}
3053
3054impl<T: TimestampLike> Iterator for TimestampRangeStepInclusive<T> {
3055    type Item = CheckedTimestamp<T>;
3056
3057    #[inline]
3058    fn next(&mut self) -> Option<CheckedTimestamp<T>> {
3059        if !self.done
3060            && ((self.rev && self.state >= self.stop) || (!self.rev && self.state <= self.stop))
3061        {
3062            let result = self.state.clone();
3063            match add_timestamp_months(self.state.deref(), self.step.months) {
3064                Ok(state) => match state.checked_add_signed(self.step.duration_as_chrono()) {
3065                    Some(v) => match CheckedTimestamp::from_timestamplike(v) {
3066                        Ok(v) => {
3067                            // Advance only if the step makes progress toward `stop`. A mixed
3068                            // month/day step can reach an in-bounds fixed point (month addition
3069                            // saturates a short month back onto the start day), which would
3070                            // otherwise loop forever.
3071                            let progressed = if self.rev {
3072                                v < self.state
3073                            } else {
3074                                v > self.state
3075                            };
3076                            if progressed {
3077                                self.state = v
3078                            } else {
3079                                self.done = true
3080                            }
3081                        }
3082                        Err(_) => self.done = true,
3083                    },
3084                    None => self.done = true,
3085                },
3086                Err(..) => {
3087                    self.done = true;
3088                }
3089            }
3090
3091            Some(result)
3092        } else {
3093            None
3094        }
3095    }
3096}
3097
3098fn generate_series_ts<T: TimestampLike>(
3099    start: CheckedTimestamp<T>,
3100    stop: CheckedTimestamp<T>,
3101    step: Interval,
3102    conv: fn(CheckedTimestamp<T>) -> Datum<'static>,
3103) -> Result<impl Iterator<Item = (Row, Diff)>, EvalError> {
3104    let normalized_step = step.as_microseconds();
3105    if normalized_step == 0 {
3106        return Err(EvalError::InvalidParameterValue(
3107            "step size cannot equal zero".into(),
3108        ));
3109    }
3110    let rev = normalized_step < 0;
3111
3112    let trsi = TimestampRangeStepInclusive {
3113        state: start,
3114        stop,
3115        step,
3116        rev,
3117        done: false,
3118    };
3119
3120    Ok(trsi.map(move |i| (Row::pack_slice(&[conv(i)]), Diff::ONE)))
3121}
3122
3123fn generate_subscripts_array(
3124    a: Datum,
3125    dim: i32,
3126) -> Result<Box<dyn Iterator<Item = (Row, Diff)>>, EvalError> {
3127    if dim <= 0 {
3128        return Ok(Box::new(iter::empty()));
3129    }
3130
3131    match a.unwrap_array().dims().into_iter().nth(
3132        (dim - 1)
3133            .try_into()
3134            .map_err(|_| EvalError::Int32OutOfRange((dim - 1).to_string().into()))?,
3135    ) {
3136        Some(requested_dim) => {
3137            let lower_bound: i32 = requested_dim.lower_bound.try_into().map_err(|_| {
3138                EvalError::Int32OutOfRange(requested_dim.lower_bound.to_string().into())
3139            })?;
3140            // The subscripts run from the lower bound to the upper bound,
3141            // inclusive. The upper bound is `lower_bound + length - 1`.
3142            let length: i32 = requested_dim
3143                .length
3144                .try_into()
3145                .map_err(|_| EvalError::Int32OutOfRange(requested_dim.length.to_string().into()))?;
3146            let upper_bound = lower_bound.checked_add(length - 1).ok_or_else(|| {
3147                EvalError::Int32OutOfRange(requested_dim.length.to_string().into())
3148            })?;
3149            Ok(Box::new(generate_series::<i32>(
3150                lower_bound,
3151                upper_bound,
3152                1,
3153            )?))
3154        }
3155        None => Ok(Box::new(iter::empty())),
3156    }
3157}
3158
3159fn unnest_array<'a>(a: Datum<'a>) -> impl Iterator<Item = (Row, Diff)> + 'a {
3160    a.unwrap_array()
3161        .elements()
3162        .iter()
3163        .map(move |e| (Row::pack_slice(&[e]), Diff::ONE))
3164}
3165
3166fn unnest_list<'a>(a: Datum<'a>) -> impl Iterator<Item = (Row, Diff)> + 'a {
3167    a.unwrap_list()
3168        .iter()
3169        .map(move |e| (Row::pack_slice(&[e]), Diff::ONE))
3170}
3171
3172fn unnest_map<'a>(a: Datum<'a>) -> impl Iterator<Item = (Row, Diff)> + 'a {
3173    a.unwrap_map()
3174        .iter()
3175        .map(move |(k, v)| (Row::pack_slice(&[Datum::from(k), v]), Diff::ONE))
3176}
3177
3178impl AggregateFunc {
3179    /// The base function name without the `~[...]` suffix used when rendering
3180    /// variants that represent a parameterized function family.
3181    pub fn name(&self) -> &'static str {
3182        match self {
3183            Self::MaxNumeric => "max",
3184            Self::MaxInt16 => "max",
3185            Self::MaxInt32 => "max",
3186            Self::MaxInt64 => "max",
3187            Self::MaxUInt16 => "max",
3188            Self::MaxUInt32 => "max",
3189            Self::MaxUInt64 => "max",
3190            Self::MaxMzTimestamp => "max",
3191            Self::MaxFloat32 => "max",
3192            Self::MaxFloat64 => "max",
3193            Self::MaxBool => "max",
3194            Self::MaxString => "max",
3195            Self::MaxDate => "max",
3196            Self::MaxTimestamp => "max",
3197            Self::MaxTimestampTz => "max",
3198            Self::MaxInterval => "max",
3199            Self::MaxTime => "max",
3200            Self::MinNumeric => "min",
3201            Self::MinInt16 => "min",
3202            Self::MinInt32 => "min",
3203            Self::MinInt64 => "min",
3204            Self::MinUInt16 => "min",
3205            Self::MinUInt32 => "min",
3206            Self::MinUInt64 => "min",
3207            Self::MinMzTimestamp => "min",
3208            Self::MinFloat32 => "min",
3209            Self::MinFloat64 => "min",
3210            Self::MinBool => "min",
3211            Self::MinString => "min",
3212            Self::MinDate => "min",
3213            Self::MinTimestamp => "min",
3214            Self::MinTimestampTz => "min",
3215            Self::MinInterval => "min",
3216            Self::MinTime => "min",
3217            Self::SumInt16 => "sum",
3218            Self::SumInt32 => "sum",
3219            Self::SumInt64 => "sum",
3220            Self::SumUInt16 => "sum",
3221            Self::SumUInt32 => "sum",
3222            Self::SumUInt64 => "sum",
3223            Self::SumFloat32 => "sum",
3224            Self::SumFloat64 => "sum",
3225            Self::SumNumeric => "sum",
3226            Self::Count => "count",
3227            Self::Any => "any",
3228            Self::All => "all",
3229            Self::JsonbAgg { .. } => "jsonb_agg",
3230            Self::JsonbObjectAgg { .. } => "jsonb_object_agg",
3231            Self::MapAgg { .. } => "map_agg",
3232            Self::ArrayConcat { .. } => "array_agg",
3233            Self::ListConcat { .. } => "list_agg",
3234            Self::StringAgg { .. } => "string_agg",
3235            Self::RowNumber { .. } => "row_number",
3236            Self::Rank { .. } => "rank",
3237            Self::DenseRank { .. } => "dense_rank",
3238            Self::LagLead {
3239                lag_lead: LagLeadType::Lag,
3240                ..
3241            } => "lag",
3242            Self::LagLead {
3243                lag_lead: LagLeadType::Lead,
3244                ..
3245            } => "lead",
3246            Self::FirstValue { .. } => "first_value",
3247            Self::LastValue { .. } => "last_value",
3248            Self::WindowAggregate { .. } => "window_agg",
3249            Self::FusedValueWindowFunc { .. } => "fused_value_window_func",
3250            Self::FusedWindowAggregate { .. } => "fused_window_agg",
3251            Self::Dummy => "dummy",
3252        }
3253    }
3254}
3255
3256impl<'a, M> fmt::Display for HumanizedExpr<'a, AggregateFunc, M>
3257where
3258    M: HumanizerMode,
3259{
3260    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
3261        use AggregateFunc::*;
3262        let name = self.expr.name();
3263        match self.expr {
3264            JsonbAgg { order_by }
3265            | JsonbObjectAgg { order_by }
3266            | MapAgg { order_by, .. }
3267            | ArrayConcat { order_by }
3268            | ListConcat { order_by }
3269            | StringAgg { order_by }
3270            | RowNumber { order_by }
3271            | Rank { order_by }
3272            | DenseRank { order_by } => {
3273                let order_by = order_by.iter().map(|col| self.child(col));
3274                write!(f, "{}[order_by=[{}]]", name, separated(", ", order_by))
3275            }
3276            LagLead {
3277                lag_lead: _,
3278                ignore_nulls,
3279                order_by,
3280            } => {
3281                let order_by = order_by.iter().map(|col| self.child(col));
3282                f.write_str(name)?;
3283                f.write_str("[")?;
3284                if *ignore_nulls {
3285                    f.write_str("ignore_nulls=true, ")?;
3286                }
3287                write!(f, "order_by=[{}]", separated(", ", order_by))?;
3288                f.write_str("]")
3289            }
3290            FirstValue {
3291                order_by,
3292                window_frame,
3293            } => {
3294                let order_by = order_by.iter().map(|col| self.child(col));
3295                f.write_str(name)?;
3296                f.write_str("[")?;
3297                write!(f, "order_by=[{}]", separated(", ", order_by))?;
3298                if *window_frame != WindowFrame::default() {
3299                    write!(f, " {}", window_frame)?;
3300                }
3301                f.write_str("]")
3302            }
3303            LastValue {
3304                order_by,
3305                window_frame,
3306            } => {
3307                let order_by = order_by.iter().map(|col| self.child(col));
3308                f.write_str(name)?;
3309                f.write_str("[")?;
3310                write!(f, "order_by=[{}]", separated(", ", order_by))?;
3311                if *window_frame != WindowFrame::default() {
3312                    write!(f, " {}", window_frame)?;
3313                }
3314                f.write_str("]")
3315            }
3316            WindowAggregate {
3317                wrapped_aggregate,
3318                order_by,
3319                window_frame,
3320            } => {
3321                let order_by = order_by.iter().map(|col| self.child(col));
3322                let wrapped_aggregate = self.child(wrapped_aggregate.deref());
3323                f.write_str(name)?;
3324                f.write_str("[")?;
3325                write!(f, "{} ", wrapped_aggregate)?;
3326                write!(f, "order_by=[{}]", separated(", ", order_by))?;
3327                if *window_frame != WindowFrame::default() {
3328                    write!(f, " {}", window_frame)?;
3329                }
3330                f.write_str("]")
3331            }
3332            FusedValueWindowFunc { funcs, order_by } => {
3333                let order_by = order_by.iter().map(|col| self.child(col));
3334                let funcs = separated(", ", funcs.iter().map(|func| self.child(func)));
3335                f.write_str(name)?;
3336                f.write_str("[")?;
3337                write!(f, "{} ", funcs)?;
3338                write!(f, "order_by=[{}]", separated(", ", order_by))?;
3339                f.write_str("]")
3340            }
3341            _ => f.write_str(name),
3342        }
3343    }
3344}
3345
3346#[derive(
3347    Clone,
3348    Debug,
3349    Eq,
3350    PartialEq,
3351    Ord,
3352    PartialOrd,
3353    Serialize,
3354    Deserialize,
3355    Hash
3356)]
3357pub struct CaptureGroupDesc {
3358    pub index: u32,
3359    pub name: Option<String>,
3360    pub nullable: bool,
3361}
3362
3363#[derive(
3364    Clone,
3365    Copy,
3366    Debug,
3367    Eq,
3368    PartialEq,
3369    Ord,
3370    PartialOrd,
3371    Serialize,
3372    Deserialize,
3373    Hash,
3374    Default
3375)]
3376pub struct AnalyzedRegexOpts {
3377    pub case_insensitive: bool,
3378    pub global: bool,
3379}
3380
3381impl FromStr for AnalyzedRegexOpts {
3382    type Err = EvalError;
3383
3384    fn from_str(s: &str) -> Result<Self, Self::Err> {
3385        let mut opts = AnalyzedRegexOpts::default();
3386        for c in s.chars() {
3387            match c {
3388                'i' => opts.case_insensitive = true,
3389                'g' => opts.global = true,
3390                _ => return Err(EvalError::InvalidRegexFlag(c)),
3391            }
3392        }
3393        Ok(opts)
3394    }
3395}
3396
3397#[derive(
3398    Clone,
3399    Debug,
3400    Eq,
3401    PartialEq,
3402    Ord,
3403    PartialOrd,
3404    Serialize,
3405    Deserialize,
3406    Hash
3407)]
3408pub struct AnalyzedRegex(ReprRegex, Vec<CaptureGroupDesc>, AnalyzedRegexOpts);
3409
3410impl AnalyzedRegex {
3411    pub fn new(s: &str, opts: AnalyzedRegexOpts) -> Result<Self, RegexCompilationError> {
3412        let r = ReprRegex::new(s, opts.case_insensitive)?;
3413        // TODO(benesch): remove potentially dangerous usage of `as`.
3414        #[allow(clippy::as_conversions)]
3415        let descs: Vec<_> = r
3416            .capture_names()
3417            .enumerate()
3418            // The first capture is the entire matched string.
3419            // This will often not be useful, so skip it.
3420            // If people want it they can just surround their
3421            // entire regex in an explicit capture group.
3422            .skip(1)
3423            .map(|(i, name)| CaptureGroupDesc {
3424                index: i as u32,
3425                name: name.map(String::from),
3426                // TODO -- we can do better.
3427                // https://github.com/MaterializeInc/database-issues/issues/612
3428                nullable: true,
3429            })
3430            .collect();
3431        Ok(Self(r, descs, opts))
3432    }
3433    pub fn capture_groups_len(&self) -> usize {
3434        self.1.len()
3435    }
3436    pub fn capture_groups_iter(&self) -> impl Iterator<Item = &CaptureGroupDesc> {
3437        self.1.iter()
3438    }
3439    pub fn inner(&self) -> &Regex {
3440        &(self.0).regex
3441    }
3442    pub fn opts(&self) -> &AnalyzedRegexOpts {
3443        &self.2
3444    }
3445}
3446
3447pub fn csv_extract(a: Datum<'_>, n_cols: usize) -> impl Iterator<Item = (Row, Diff)> + '_ {
3448    let bytes = a.unwrap_str().as_bytes();
3449    let mut row = Row::default();
3450    let csv_reader = csv::ReaderBuilder::new()
3451        .has_headers(false)
3452        .from_reader(bytes);
3453    csv_reader.into_records().filter_map(move |res| match res {
3454        Ok(sr) if sr.len() == n_cols => {
3455            row.packer().extend(sr.iter().map(Datum::String));
3456            Some((row.clone(), Diff::ONE))
3457        }
3458        _ => None,
3459    })
3460}
3461
3462pub fn repeat_row(a: Datum) -> Option<(Row, Diff)> {
3463    let n = a.unwrap_int64();
3464    if n != 0 {
3465        Some((Row::default(), n.into()))
3466    } else {
3467        None
3468    }
3469}
3470
3471pub fn repeat_row_non_negative<'a>(
3472    a: Datum,
3473) -> Result<Box<dyn Iterator<Item = (Row, Diff)> + 'a>, EvalError> {
3474    let n = a.unwrap_int64();
3475    if n < 0 {
3476        Err(EvalError::InvalidParameterValue(
3477            format!("repeat_row_non_negative got {}", n).into(),
3478        ))
3479    } else if n == 0 {
3480        Ok(Box::new(iter::empty()))
3481    } else {
3482        // iterator with 1 element; n goes into the diff
3483        Ok(Box::new(iter::once((Row::default(), n.into()))))
3484    }
3485}
3486
3487fn wrap<'a>(datums: &'a [Datum<'a>], width: usize) -> impl Iterator<Item = (Row, Diff)> + 'a {
3488    datums
3489        .chunks(width)
3490        .map(|chunk| (Row::pack(chunk), Diff::ONE))
3491}
3492
3493fn acl_explode<'a>(
3494    acl_items: Datum<'a>,
3495    temp_storage: &'a RowArena,
3496) -> Result<impl Iterator<Item = (Row, Diff)> + 'a, EvalError> {
3497    let acl_items = acl_items.unwrap_array();
3498    let mut res = Vec::new();
3499    for acl_item in acl_items.elements().iter() {
3500        if acl_item.is_null() {
3501            return Err(EvalError::AclArrayNullElement);
3502        }
3503        let acl_item = acl_item.unwrap_acl_item();
3504        for privilege in acl_item.acl_mode.explode() {
3505            let row = [
3506                Datum::UInt32(acl_item.grantor.0),
3507                Datum::UInt32(acl_item.grantee.0),
3508                Datum::String(temp_storage.push_string(privilege.to_string())),
3509                // GRANT OPTION is not implemented, so we hardcode false.
3510                Datum::False,
3511            ];
3512            res.push((Row::pack_slice(&row), Diff::ONE));
3513        }
3514    }
3515    Ok(res.into_iter())
3516}
3517
3518fn mz_acl_explode<'a>(
3519    mz_acl_items: Datum<'a>,
3520    temp_storage: &'a RowArena,
3521) -> Result<impl Iterator<Item = (Row, Diff)> + 'a, EvalError> {
3522    let mz_acl_items = mz_acl_items.unwrap_array();
3523    let mut res = Vec::new();
3524    for mz_acl_item in mz_acl_items.elements().iter() {
3525        if mz_acl_item.is_null() {
3526            return Err(EvalError::MzAclArrayNullElement);
3527        }
3528        let mz_acl_item = mz_acl_item.unwrap_mz_acl_item();
3529        for privilege in mz_acl_item.acl_mode.explode() {
3530            let row = [
3531                Datum::String(temp_storage.push_string(mz_acl_item.grantor.to_string())),
3532                Datum::String(temp_storage.push_string(mz_acl_item.grantee.to_string())),
3533                Datum::String(temp_storage.push_string(privilege.to_string())),
3534                // GRANT OPTION is not implemented, so we hardcode false.
3535                Datum::False,
3536            ];
3537            res.push((Row::pack_slice(&row), Diff::ONE));
3538        }
3539    }
3540    Ok(res.into_iter())
3541}
3542
3543/// When adding a new `TableFunc` variant, please consider adding it to
3544/// `TableFunc::with_ordinality`!
3545#[derive(
3546    Clone,
3547    Debug,
3548    Eq,
3549    PartialEq,
3550    Ord,
3551    PartialOrd,
3552    Serialize,
3553    Deserialize,
3554    Hash
3555)]
3556pub enum TableFunc {
3557    AclExplode,
3558    MzAclExplode,
3559    JsonbEach,
3560    JsonbEachStringify,
3561    JsonbObjectKeys,
3562    JsonbArrayElements,
3563    JsonbArrayElementsStringify,
3564    RegexpExtract(AnalyzedRegex),
3565    CsvExtract(usize),
3566    GenerateSeriesInt32,
3567    GenerateSeriesInt64,
3568    /// An int64 `generate_series` that the optimizer promises to leave as an
3569    /// enumeration: no transform may match on this variant to replace its
3570    /// evaluation with a cardinality shortcut (compare the collapse of an
3571    /// unused `GenerateSeriesInt64` into `RepeatRowNonNegative`). Its
3572    /// *argument* expressions are still simplified like any other scalar.
3573    ///
3574    /// Exposed as `mz_unsafe.generate_series_unoptimized` for tests that rely
3575    /// on the work of enumeration actually happening (e.g. stress tests whose
3576    /// load would otherwise be optimized away). As with everything in
3577    /// `mz_unsafe`, it is not a supported surface: bug reports must not
3578    /// depend on it.
3579    GenerateSeriesUnoptimized,
3580    GenerateSeriesTimestamp,
3581    GenerateSeriesTimestampTz,
3582    /// Supplied with an input count,
3583    ///   1. Adds a column as if a typed subquery result,
3584    ///   2. Filters the row away if the count is only one,
3585    ///   3. Errors if the count is not exactly one.
3586    /// The intent is that this presents as if a subquery result with too many
3587    /// records contributing. The error column has the same type as the result
3588    /// should have, but we only produce it if the count exceeds one.
3589    ///
3590    /// This logic could nearly be achieved with map, filter, project logic,
3591    /// but has been challenging to do in a way that respects the vagaries of
3592    /// SQL and our semantics. If we reveal a constant value in the column we
3593    /// risk the optimizer pruning the branch; if we reveal that this will not
3594    /// produce rows we risk the optimizer pruning the branch; if we reveal that
3595    /// the only possible value is an error we risk the optimizer propagating that
3596    /// error without guards.
3597    ///
3598    /// Before replacing this by an `MirScalarExpr`, quadruple check that it
3599    /// would not result in misoptimizations due to expression evaluation order
3600    /// being utterly undefined, and predicate pushdown trimming any fragments
3601    /// that might produce columns that will not be needed.
3602    GuardSubquerySize {
3603        column_type: SqlScalarType,
3604    },
3605    /// Repeats the input row the given number of times. Can even repeat a negative number of times,
3606    /// which has some important consequences:
3607    /// - can lead to negative accumulations downstream;
3608    /// - can't be used in `WITH ORDINALITY` and other constructs that are implemented by
3609    ///   `TableFunc::WithOrdinality`, e.g., `ROWS FROM`;
3610    /// - output is non-monotonic.
3611    RepeatRow,
3612    /// Same as `RepeatRow`, but errors on a negative count, and thereby avoids the above
3613    /// peculiarities.
3614    RepeatRowNonNegative,
3615    UnnestArray {
3616        el_typ: SqlScalarType,
3617    },
3618    UnnestList {
3619        el_typ: SqlScalarType,
3620    },
3621    UnnestMap {
3622        value_type: SqlScalarType,
3623    },
3624    /// Given `n` input expressions, wraps them into `n / width` rows, each of
3625    /// `width` columns.
3626    ///
3627    /// This function is not intended to be called directly by end users, but
3628    /// is useful in the planning of e.g. VALUES clauses.
3629    Wrap {
3630        types: Vec<SqlColumnType>,
3631        width: usize,
3632    },
3633    GenerateSubscriptsArray,
3634    /// Execute some arbitrary scalar function as a table function.
3635    TabletizedScalar {
3636        name: String,
3637        relation: SqlRelationType,
3638    },
3639    RegexpMatches,
3640    /// Implements the WITH ORDINALITY clause.
3641    ///
3642    /// Don't construct `TableFunc::WithOrdinality` manually! Use the `with_ordinality` constructor
3643    /// function instead, which checks whether the given table function supports `WithOrdinality`.
3644    #[allow(private_interfaces)]
3645    WithOrdinality(WithOrdinality),
3646}
3647
3648/// Evaluates the inner table function, expands its results into unary (repeating each row as
3649/// many times as the diff indicates), and appends an integer corresponding to the ordinal
3650/// position (starting from 1). For example, it numbers the elements of a list when calling
3651/// `unnest_list`.
3652///
3653/// Private enum variant of `TableFunc`. Don't construct this directly, but use
3654/// `TableFunc::with_ordinality` instead.
3655#[derive(
3656    Clone,
3657    Debug,
3658    Eq,
3659    PartialEq,
3660    Ord,
3661    PartialOrd,
3662    Serialize,
3663    Deserialize,
3664    Hash
3665)]
3666struct WithOrdinality {
3667    inner: Box<TableFunc>,
3668}
3669
3670impl TableFunc {
3671    /// Adds `WITH ORDINALITY` to a table function if it's allowed on the given table function.
3672    pub fn with_ordinality(inner: TableFunc) -> Option<TableFunc> {
3673        match inner {
3674            TableFunc::AclExplode
3675            | TableFunc::MzAclExplode
3676            | TableFunc::JsonbEach
3677            | TableFunc::JsonbEachStringify
3678            | TableFunc::JsonbObjectKeys
3679            | TableFunc::JsonbArrayElements
3680            | TableFunc::JsonbArrayElementsStringify
3681            | TableFunc::RegexpExtract(_)
3682            | TableFunc::CsvExtract(_)
3683            | TableFunc::GenerateSeriesInt32
3684            | TableFunc::GenerateSeriesInt64
3685            | TableFunc::GenerateSeriesUnoptimized
3686            | TableFunc::GenerateSeriesTimestamp
3687            | TableFunc::GenerateSeriesTimestampTz
3688            | TableFunc::GuardSubquerySize { .. }
3689            | TableFunc::RepeatRowNonNegative
3690            | TableFunc::UnnestArray { .. }
3691            | TableFunc::UnnestList { .. }
3692            | TableFunc::UnnestMap { .. }
3693            | TableFunc::Wrap { .. }
3694            | TableFunc::GenerateSubscriptsArray
3695            | TableFunc::TabletizedScalar { .. }
3696            | TableFunc::RegexpMatches => Some(TableFunc::WithOrdinality(WithOrdinality {
3697                inner: Box::new(inner),
3698            })),
3699            // IMPORTANT: Before adding a new table function above, consider negative diffs:
3700            // `WithOrdinality::eval` will panic if the inner table function emits a negative diff.
3701            // (Note that negative diffs in the table function's _input_ don't matter. The table
3702            // function implementation doesn't see the input diffs, so the thing that matters here
3703            // is whether the table function itself can emit a negative diff.)
3704            TableFunc::RepeatRow // can produce negative diffs
3705            | TableFunc::WithOrdinality(_) => None, // no nesting of `WITH ORDINALITY` allowed
3706        }
3707    }
3708}
3709
3710impl TableFunc {
3711    /// Executes `self` on the given input row (`datums`).
3712    pub fn eval<'a>(
3713        &'a self,
3714        datums: &'a [Datum<'a>],
3715        temp_storage: &'a RowArena,
3716    ) -> Result<Box<dyn Iterator<Item = (Row, Diff)> + 'a>, EvalError> {
3717        if self.empty_on_null_input() && datums.iter().any(|d| d.is_null()) {
3718            return Ok(Box::new(vec![].into_iter()));
3719        }
3720        match self {
3721            TableFunc::AclExplode => Ok(Box::new(acl_explode(datums[0], temp_storage)?)),
3722            TableFunc::MzAclExplode => Ok(Box::new(mz_acl_explode(datums[0], temp_storage)?)),
3723            TableFunc::JsonbEach => Ok(Box::new(jsonb_each(datums[0]))),
3724            TableFunc::JsonbEachStringify => {
3725                Ok(Box::new(jsonb_each_stringify(datums[0], temp_storage)))
3726            }
3727            TableFunc::JsonbObjectKeys => Ok(Box::new(jsonb_object_keys(datums[0]))),
3728            TableFunc::JsonbArrayElements => Ok(Box::new(jsonb_array_elements(datums[0]))),
3729            TableFunc::JsonbArrayElementsStringify => Ok(Box::new(jsonb_array_elements_stringify(
3730                datums[0],
3731                temp_storage,
3732            ))),
3733            TableFunc::RegexpExtract(a) => Ok(Box::new(regexp_extract(datums[0], a).into_iter())),
3734            TableFunc::CsvExtract(n_cols) => Ok(Box::new(csv_extract(datums[0], *n_cols))),
3735            TableFunc::GenerateSeriesInt32 => {
3736                let res = generate_series(
3737                    datums[0].unwrap_int32(),
3738                    datums[1].unwrap_int32(),
3739                    datums[2].unwrap_int32(),
3740                )?;
3741                Ok(Box::new(res))
3742            }
3743            TableFunc::GenerateSeriesInt64 | TableFunc::GenerateSeriesUnoptimized => {
3744                let res = generate_series(
3745                    datums[0].unwrap_int64(),
3746                    datums[1].unwrap_int64(),
3747                    datums[2].unwrap_int64(),
3748                )?;
3749                Ok(Box::new(res))
3750            }
3751            TableFunc::GenerateSeriesTimestamp => {
3752                fn pass_through<'a>(d: CheckedTimestamp<NaiveDateTime>) -> Datum<'a> {
3753                    Datum::from(d)
3754                }
3755                let res = generate_series_ts(
3756                    datums[0].unwrap_timestamp(),
3757                    datums[1].unwrap_timestamp(),
3758                    datums[2].unwrap_interval(),
3759                    pass_through,
3760                )?;
3761                Ok(Box::new(res))
3762            }
3763            TableFunc::GenerateSeriesTimestampTz => {
3764                fn gen_ts_tz<'a>(d: CheckedTimestamp<DateTime<Utc>>) -> Datum<'a> {
3765                    Datum::from(d)
3766                }
3767                let res = generate_series_ts(
3768                    datums[0].unwrap_timestamptz(),
3769                    datums[1].unwrap_timestamptz(),
3770                    datums[2].unwrap_interval(),
3771                    gen_ts_tz,
3772                )?;
3773                Ok(Box::new(res))
3774            }
3775            TableFunc::GenerateSubscriptsArray => {
3776                generate_subscripts_array(datums[0], datums[1].unwrap_int32())
3777            }
3778            TableFunc::GuardSubquerySize { column_type: _ } => {
3779                // A subquery used as an expression may return at most one row.
3780                // For 0 or 1 we emit no rows and let the subquery's own output
3781                // flow through. Zero can't come directly from the count that
3782                // lowering plants (an MIR `count(true)`, at least 1 per group),
3783                // but over a provably empty subquery body the optimizer may
3784                // vacuously rewrite the counted expression to `null`, and a
3785                // count of `null` over a group is 0. Later simplifications can
3786                // surface that 0 as a literal argument that is evaluated
3787                // during optimization. Emitting no rows is also the correct
3788                // semantics: the empty subquery decorrelates to NULL via the
3789                // outer lookup.
3790                let count = datums[0].unwrap_int64();
3791                if count > 1 {
3792                    Err(EvalError::MultipleRowsFromSubquery)
3793                } else if count < 0 {
3794                    // Would require negative multiplicities to reach the guard.
3795                    Err(EvalError::NegativeRowsFromSubquery)
3796                } else {
3797                    Ok(Box::new([].into_iter()))
3798                }
3799            }
3800            TableFunc::RepeatRow => Ok(Box::new(repeat_row(datums[0]).into_iter())),
3801            TableFunc::RepeatRowNonNegative => repeat_row_non_negative(datums[0]),
3802            TableFunc::UnnestArray { .. } => Ok(Box::new(unnest_array(datums[0]))),
3803            TableFunc::UnnestList { .. } => Ok(Box::new(unnest_list(datums[0]))),
3804            TableFunc::UnnestMap { .. } => Ok(Box::new(unnest_map(datums[0]))),
3805            TableFunc::Wrap { width, .. } => Ok(Box::new(wrap(datums, *width))),
3806            TableFunc::TabletizedScalar { .. } => {
3807                let r = Row::pack_slice(datums);
3808                Ok(Box::new(std::iter::once((r, Diff::ONE))))
3809            }
3810            TableFunc::RegexpMatches => Ok(Box::new(regexp_matches(datums)?)),
3811            TableFunc::WithOrdinality(func_with_ordinality) => {
3812                func_with_ordinality.eval(datums, temp_storage)
3813            }
3814        }
3815    }
3816
3817    pub fn output_sql_type(&self) -> SqlRelationType {
3818        let (column_types, keys) = match self {
3819            TableFunc::AclExplode => {
3820                let column_types = vec![
3821                    SqlScalarType::Oid.nullable(false),
3822                    SqlScalarType::Oid.nullable(false),
3823                    SqlScalarType::String.nullable(false),
3824                    SqlScalarType::Bool.nullable(false),
3825                ];
3826                let keys = vec![];
3827                (column_types, keys)
3828            }
3829            TableFunc::MzAclExplode => {
3830                let column_types = vec![
3831                    SqlScalarType::String.nullable(false),
3832                    SqlScalarType::String.nullable(false),
3833                    SqlScalarType::String.nullable(false),
3834                    SqlScalarType::Bool.nullable(false),
3835                ];
3836                let keys = vec![];
3837                (column_types, keys)
3838            }
3839            TableFunc::JsonbEach => {
3840                let column_types = vec![
3841                    SqlScalarType::String.nullable(false),
3842                    SqlScalarType::Jsonb.nullable(false),
3843                ];
3844                let keys = vec![];
3845                (column_types, keys)
3846            }
3847            TableFunc::JsonbEachStringify => {
3848                let column_types = vec![
3849                    SqlScalarType::String.nullable(false),
3850                    SqlScalarType::String.nullable(true),
3851                ];
3852                let keys = vec![];
3853                (column_types, keys)
3854            }
3855            TableFunc::JsonbObjectKeys => {
3856                let column_types = vec![SqlScalarType::String.nullable(false)];
3857                let keys = vec![];
3858                (column_types, keys)
3859            }
3860            TableFunc::JsonbArrayElements => {
3861                let column_types = vec![SqlScalarType::Jsonb.nullable(false)];
3862                let keys = vec![];
3863                (column_types, keys)
3864            }
3865            TableFunc::JsonbArrayElementsStringify => {
3866                let column_types = vec![SqlScalarType::String.nullable(true)];
3867                let keys = vec![];
3868                (column_types, keys)
3869            }
3870            TableFunc::RegexpExtract(a) => {
3871                let column_types = a
3872                    .capture_groups_iter()
3873                    .map(|cg| SqlScalarType::String.nullable(cg.nullable))
3874                    .collect();
3875                let keys = vec![];
3876                (column_types, keys)
3877            }
3878            TableFunc::CsvExtract(n_cols) => {
3879                let column_types = iter::repeat(SqlScalarType::String.nullable(false))
3880                    .take(*n_cols)
3881                    .collect();
3882                let keys = vec![];
3883                (column_types, keys)
3884            }
3885            TableFunc::GenerateSeriesInt32 => {
3886                let column_types = vec![SqlScalarType::Int32.nullable(false)];
3887                let keys = vec![vec![0]];
3888                (column_types, keys)
3889            }
3890            TableFunc::GenerateSeriesInt64 | TableFunc::GenerateSeriesUnoptimized => {
3891                let column_types = vec![SqlScalarType::Int64.nullable(false)];
3892                let keys = vec![vec![0]];
3893                (column_types, keys)
3894            }
3895            TableFunc::GenerateSeriesTimestamp => {
3896                let column_types =
3897                    vec![SqlScalarType::Timestamp { precision: None }.nullable(false)];
3898                let keys = vec![vec![0]];
3899                (column_types, keys)
3900            }
3901            TableFunc::GenerateSeriesTimestampTz => {
3902                let column_types =
3903                    vec![SqlScalarType::TimestampTz { precision: None }.nullable(false)];
3904                let keys = vec![vec![0]];
3905                (column_types, keys)
3906            }
3907            TableFunc::GenerateSubscriptsArray => {
3908                let column_types = vec![SqlScalarType::Int32.nullable(false)];
3909                let keys = vec![vec![0]];
3910                (column_types, keys)
3911            }
3912            TableFunc::GuardSubquerySize { column_type } => {
3913                let column_types = vec![column_type.clone().nullable(false)];
3914                let keys = vec![];
3915                (column_types, keys)
3916            }
3917            TableFunc::RepeatRow | TableFunc::RepeatRowNonNegative => {
3918                let column_types = vec![];
3919                let keys = vec![];
3920                (column_types, keys)
3921            }
3922            TableFunc::UnnestArray { el_typ } => {
3923                let column_types = vec![el_typ.clone().nullable(true)];
3924                let keys = vec![];
3925                (column_types, keys)
3926            }
3927            TableFunc::UnnestList { el_typ } => {
3928                let column_types = vec![el_typ.clone().nullable(true)];
3929                let keys = vec![];
3930                (column_types, keys)
3931            }
3932            TableFunc::UnnestMap { value_type } => {
3933                let column_types = vec![
3934                    SqlScalarType::String.nullable(false),
3935                    value_type.clone().nullable(true),
3936                ];
3937                let keys = vec![vec![0]];
3938                (column_types, keys)
3939            }
3940            TableFunc::Wrap { types, .. } => {
3941                let column_types = types.clone();
3942                let keys = vec![];
3943                (column_types, keys)
3944            }
3945            TableFunc::TabletizedScalar { relation, .. } => {
3946                return relation.clone();
3947            }
3948            TableFunc::RegexpMatches => {
3949                let column_types =
3950                    vec![SqlScalarType::Array(Box::new(SqlScalarType::String)).nullable(false)];
3951                let keys = vec![];
3952
3953                (column_types, keys)
3954            }
3955            TableFunc::WithOrdinality(WithOrdinality { inner }) => {
3956                let mut typ = inner.output_sql_type();
3957                // Add the ordinality column.
3958                typ.column_types.push(SqlScalarType::Int64.nullable(false));
3959                // The ordinality column is always a key.
3960                typ.keys.push(vec![typ.column_types.len() - 1]);
3961                (typ.column_types, typ.keys)
3962            }
3963        };
3964
3965        soft_assert_eq_no_log!(column_types.len(), self.output_arity());
3966
3967        if !keys.is_empty() {
3968            SqlRelationType::new(column_types).with_keys(keys)
3969        } else {
3970            SqlRelationType::new(column_types)
3971        }
3972    }
3973
3974    /// Computes the representation type of this table function.
3975    ///
3976    /// This is a wrapper around [`Self::output_sql_type`] that converts the result to a representation type.
3977    pub fn output_type(&self) -> ReprRelationType {
3978        ReprRelationType::from(&self.output_sql_type())
3979    }
3980
3981    pub fn output_arity(&self) -> usize {
3982        match self {
3983            TableFunc::AclExplode => 4,
3984            TableFunc::MzAclExplode => 4,
3985            TableFunc::JsonbEach => 2,
3986            TableFunc::JsonbEachStringify => 2,
3987            TableFunc::JsonbObjectKeys => 1,
3988            TableFunc::JsonbArrayElements => 1,
3989            TableFunc::JsonbArrayElementsStringify => 1,
3990            TableFunc::RegexpExtract(a) => a.capture_groups_len(),
3991            TableFunc::CsvExtract(n_cols) => *n_cols,
3992            TableFunc::GenerateSeriesInt32 => 1,
3993            TableFunc::GenerateSeriesInt64 => 1,
3994            TableFunc::GenerateSeriesUnoptimized => 1,
3995            TableFunc::GenerateSeriesTimestamp => 1,
3996            TableFunc::GenerateSeriesTimestampTz => 1,
3997            TableFunc::GenerateSubscriptsArray => 1,
3998            TableFunc::GuardSubquerySize { .. } => 1,
3999            TableFunc::RepeatRow => 0,
4000            TableFunc::RepeatRowNonNegative => 0,
4001            TableFunc::UnnestArray { .. } => 1,
4002            TableFunc::UnnestList { .. } => 1,
4003            TableFunc::UnnestMap { .. } => 2,
4004            TableFunc::Wrap { width, .. } => *width,
4005            TableFunc::TabletizedScalar { relation, .. } => relation.column_types.len(),
4006            TableFunc::RegexpMatches => 1,
4007            TableFunc::WithOrdinality(WithOrdinality { inner }) => inner.output_arity() + 1,
4008        }
4009    }
4010
4011    pub fn empty_on_null_input(&self) -> bool {
4012        match self {
4013            TableFunc::AclExplode
4014            | TableFunc::MzAclExplode
4015            | TableFunc::JsonbEach
4016            | TableFunc::JsonbEachStringify
4017            | TableFunc::JsonbObjectKeys
4018            | TableFunc::JsonbArrayElements
4019            | TableFunc::JsonbArrayElementsStringify
4020            | TableFunc::GenerateSeriesInt32
4021            | TableFunc::GenerateSeriesInt64
4022            | TableFunc::GenerateSeriesUnoptimized
4023            | TableFunc::GenerateSeriesTimestamp
4024            | TableFunc::GenerateSeriesTimestampTz
4025            | TableFunc::GenerateSubscriptsArray
4026            | TableFunc::RegexpExtract(_)
4027            | TableFunc::CsvExtract(_)
4028            | TableFunc::RepeatRow
4029            | TableFunc::RepeatRowNonNegative
4030            | TableFunc::UnnestArray { .. }
4031            | TableFunc::UnnestList { .. }
4032            | TableFunc::UnnestMap { .. }
4033            | TableFunc::RegexpMatches => true,
4034            TableFunc::GuardSubquerySize { .. } => false,
4035            TableFunc::Wrap { .. } => false,
4036            TableFunc::TabletizedScalar { .. } => false,
4037            TableFunc::WithOrdinality(WithOrdinality { inner }) => inner.empty_on_null_input(),
4038        }
4039    }
4040
4041    /// True iff the table function preserves the append-only property of its input.
4042    pub fn preserves_monotonicity(&self) -> bool {
4043        // Most variants preserve monotonicity, but all variants are enumerated to
4044        // ensure that added variants at least check that this is the case.
4045        match self {
4046            TableFunc::AclExplode => false,
4047            TableFunc::MzAclExplode => false,
4048            TableFunc::JsonbEach => true,
4049            TableFunc::JsonbEachStringify => true,
4050            TableFunc::JsonbObjectKeys => true,
4051            TableFunc::JsonbArrayElements => true,
4052            TableFunc::JsonbArrayElementsStringify => true,
4053            TableFunc::RegexpExtract(_) => true,
4054            TableFunc::CsvExtract(_) => true,
4055            TableFunc::GenerateSeriesInt32 => true,
4056            TableFunc::GenerateSeriesInt64 => true,
4057            TableFunc::GenerateSeriesUnoptimized => true,
4058            TableFunc::GenerateSeriesTimestamp => true,
4059            TableFunc::GenerateSeriesTimestampTz => true,
4060            TableFunc::GenerateSubscriptsArray => true,
4061            TableFunc::RepeatRow => false,
4062            TableFunc::RepeatRowNonNegative => true,
4063            TableFunc::UnnestArray { .. } => true,
4064            TableFunc::UnnestList { .. } => true,
4065            TableFunc::UnnestMap { .. } => true,
4066            TableFunc::Wrap { .. } => true,
4067            TableFunc::TabletizedScalar { .. } => true,
4068            TableFunc::RegexpMatches => true,
4069            TableFunc::GuardSubquerySize { .. } => false,
4070            TableFunc::WithOrdinality(WithOrdinality { inner }) => inner.preserves_monotonicity(),
4071        }
4072    }
4073}
4074
4075impl fmt::Display for TableFunc {
4076    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
4077        match self {
4078            TableFunc::AclExplode => f.write_str("aclexplode"),
4079            TableFunc::MzAclExplode => f.write_str("mz_aclexplode"),
4080            TableFunc::JsonbEach => f.write_str("jsonb_each"),
4081            TableFunc::JsonbEachStringify => f.write_str("jsonb_each_text"),
4082            TableFunc::JsonbObjectKeys => f.write_str("jsonb_object_keys"),
4083            TableFunc::JsonbArrayElements => f.write_str("jsonb_array_elements"),
4084            TableFunc::JsonbArrayElementsStringify => f.write_str("jsonb_array_elements_text"),
4085            TableFunc::RegexpExtract(a) => write!(f, "regexp_extract({:?}, _)", a.0),
4086            TableFunc::CsvExtract(n_cols) => write!(f, "csv_extract({}, _)", n_cols),
4087            TableFunc::GenerateSeriesInt32 => f.write_str("generate_series"),
4088            TableFunc::GenerateSeriesInt64 => f.write_str("generate_series"),
4089            TableFunc::GenerateSeriesUnoptimized => f.write_str("generate_series_unoptimized"),
4090            TableFunc::GenerateSeriesTimestamp => f.write_str("generate_series"),
4091            TableFunc::GenerateSeriesTimestampTz => f.write_str("generate_series"),
4092            TableFunc::GenerateSubscriptsArray => f.write_str("generate_subscripts"),
4093            TableFunc::GuardSubquerySize { .. } => f.write_str("guard_subquery_size"),
4094            TableFunc::RepeatRow => f.write_str(REPEAT_ROW_NAME),
4095            TableFunc::RepeatRowNonNegative => f.write_str("repeat_row_non_negative"),
4096            TableFunc::UnnestArray { .. } => f.write_str("unnest_array"),
4097            TableFunc::UnnestList { .. } => f.write_str("unnest_list"),
4098            TableFunc::UnnestMap { .. } => f.write_str("unnest_map"),
4099            TableFunc::Wrap { width, .. } => write!(f, "wrap{}", width),
4100            TableFunc::TabletizedScalar { name, .. } => f.write_str(name),
4101            TableFunc::RegexpMatches => write!(f, "regexp_matches(_, _, _)"),
4102            TableFunc::WithOrdinality(WithOrdinality { inner }) => {
4103                write!(f, "{}[with_ordinality]", inner)
4104            }
4105        }
4106    }
4107}
4108
4109impl WithOrdinality {
4110    /// Executes the `self.inner` table function on the given input row (`datums`), and zips
4111    /// 1, 2, 3, ... to the result as a new column. We need to expand rows with non-1 diffs into the
4112    /// corresponding number of rows with unit diffs, because the ordinality column will have
4113    /// different values for each copy.
4114    ///
4115    /// # Panics
4116    ///
4117    /// Panics if the `inner` table function emits a negative diff.
4118    fn eval<'a>(
4119        &'a self,
4120        datums: &'a [Datum<'a>],
4121        temp_storage: &'a RowArena,
4122    ) -> Result<Box<dyn Iterator<Item = (Row, Diff)> + 'a>, EvalError> {
4123        let mut next_ordinal: i64 = 1;
4124        let it = self
4125            .inner
4126            .eval(datums, temp_storage)?
4127            .flat_map(move |(mut row, diff)| {
4128                let diff = diff.into_inner();
4129                // WITH ORDINALITY is not well-defined for negative diffs. This is ok, and
4130                // `TableFunc::with_ordinality` refuses to wrap such table functions in
4131                // `WithOrdinality` that can emit negative diffs, e.g., `repeat_row`.
4132                //
4133                // (Note that we don't need to worry about negative diffs in FlatMap's input,
4134                // because the diff of the input of the FlatMap is factored in after we return from
4135                // here.)
4136                assert!(diff >= 0);
4137                // The ordinals that will be associated with this row.
4138                let mut ordinals = next_ordinal..(next_ordinal + diff);
4139                next_ordinal += diff;
4140                // The maximum byte capacity we need for the original row and its ordinal.
4141                let cap = row.data_len() + datum_size(&Datum::Int64(next_ordinal));
4142                iter::from_fn(move || {
4143                    let ordinal = ordinals.next()?;
4144                    let mut row = if ordinals.is_empty() {
4145                        // This is the last row, so no need to clone. (Most table functions emit
4146                        // only 1 diffs, so this completely avoids cloning in most cases.)
4147                        std::mem::take(&mut row)
4148                    } else {
4149                        let mut new_row = Row::with_capacity(cap);
4150                        new_row.clone_from(&row);
4151                        new_row
4152                    };
4153                    RowPacker::for_existing_row(&mut row).push(Datum::Int64(ordinal));
4154                    Some((row, Diff::ONE))
4155                })
4156            });
4157        Ok(Box::new(it))
4158    }
4159}
4160
4161pub const REPEAT_ROW_NAME: &str = "repeat_row";
4162
4163#[cfg(test)]
4164mod tests {
4165    use mz_repr::{Datum, RowArena, SqlScalarType};
4166
4167    use super::TableFunc;
4168    use crate::EvalError;
4169
4170    /// 0 and 1 are valid (no guard rows), >1 errors with
4171    /// `MultipleRowsFromSubquery`, <0 with `NegativeRowsFromSubquery`. Zero is
4172    /// legitimate, not "can't happen": the optimizer can turn an empty
4173    /// subquery's count into a literal `0` that is evaluated during
4174    /// optimization (see the comment in `eval`), so it must not panic (exposed
4175    /// by #37049).
4176    #[mz_ore::test]
4177    fn guard_subquery_size_accepts_zero_and_one() {
4178        let func = TableFunc::GuardSubquerySize {
4179            column_type: SqlScalarType::Int64,
4180        };
4181        let temp_storage = RowArena::new();
4182
4183        for count in [0_i64, 1] {
4184            let rows = func
4185                .eval(&[Datum::Int64(count)], &temp_storage)
4186                .unwrap_or_else(|e| panic!("count {count} should be accepted, got {e:?}"))
4187                .count();
4188            assert_eq!(rows, 0, "count {count} should emit no guard rows");
4189        }
4190
4191        assert_eq!(
4192            func.eval(&[Datum::Int64(2)], &temp_storage).err(),
4193            Some(EvalError::MultipleRowsFromSubquery),
4194        );
4195        assert_eq!(
4196            func.eval(&[Datum::Int64(-1)], &temp_storage).err(),
4197            Some(EvalError::NegativeRowsFromSubquery),
4198        );
4199    }
4200}