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