1#![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
52fn 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
133fn 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 accum = accum.wrapping_add(value.wrapping_mul(i128::from(diff.into_inner())));
166 non_nulls += diff;
167 }
168 if accum == 0 && non_nulls.is_zero() {
169 Datum::Null
170 } else {
171 narrow(accum)
172 }
173}
174
175fn sum_numeric<'a, I>(datums: I) -> Datum<'a>
176where
177 I: IntoIterator<Item = Datum<'a>>,
178{
179 let mut cx = numeric::cx_datum();
180 let mut sum = Numeric::zero();
181 let mut empty = true;
182 for d in datums {
183 if !d.is_null() {
184 empty = false;
185 cx.add(&mut sum, &d.unwrap_numeric().0);
186 }
187 }
188 match empty {
189 true => Datum::Null,
190 false => Datum::from(sum),
191 }
192}
193
194fn count<'a, I>(datums: I) -> Datum<'a>
195where
196 I: IntoIterator<Item = (Datum<'a>, Diff)>,
197{
198 let mut count = Diff::ZERO;
204 for (datum, diff) in datums {
205 if !datum.is_null() {
206 count += diff;
207 }
208 }
209 Datum::from(count.into_inner())
210}
211
212fn any<'a, I>(datums: I) -> Datum<'a>
213where
214 I: IntoIterator<Item = Datum<'a>>,
215{
216 datums
217 .into_iter()
218 .fold(Datum::False, |state, next| match (state, next) {
219 (Datum::True, _) | (_, Datum::True) => Datum::True,
220 (Datum::Null, _) | (_, Datum::Null) => Datum::Null,
221 _ => Datum::False,
222 })
223}
224
225fn all<'a, I>(datums: I) -> Datum<'a>
226where
227 I: IntoIterator<Item = Datum<'a>>,
228{
229 datums
230 .into_iter()
231 .fold(Datum::True, |state, next| match (state, next) {
232 (Datum::False, _) | (_, Datum::False) => Datum::False,
233 (Datum::Null, _) | (_, Datum::Null) => Datum::Null,
234 _ => Datum::True,
235 })
236}
237
238fn string_agg<'a, I>(datums: I, temp_storage: &'a RowArena, order_by: &[ColumnOrder]) -> Datum<'a>
239where
240 I: IntoIterator<Item = Datum<'a>>,
241{
242 const EMPTY_SEP: &str = "";
243
244 let datums = order_aggregate_datums(datums, order_by);
245 let mut sep_value_pairs = datums.into_iter().filter_map(|d| {
246 if d.is_null() {
247 return None;
248 }
249 let mut value_sep = d.unwrap_list().iter();
250 match (value_sep.next().unwrap(), value_sep.next().unwrap()) {
251 (Datum::Null, _) => None,
252 (Datum::String(val), Datum::Null) => Some((EMPTY_SEP, val)),
253 (Datum::String(val), Datum::String(sep)) => Some((sep, val)),
254 _ => unreachable!(),
255 }
256 });
257
258 let mut s = String::default();
259 match sep_value_pairs.next() {
260 Some((_, value)) => s.push_str(value),
262 None => return Datum::Null,
264 }
265
266 for (sep, value) in sep_value_pairs {
267 s.push_str(sep);
268 s.push_str(value);
269 }
270
271 Datum::String(temp_storage.push_string(s))
272}
273
274fn jsonb_agg<'a, I>(datums: I, temp_storage: &'a RowArena, order_by: &[ColumnOrder]) -> Datum<'a>
275where
276 I: IntoIterator<Item = Datum<'a>>,
277{
278 let datums = order_aggregate_datums(datums, order_by);
279 temp_storage.make_datum(|packer| {
280 packer.push_list(datums.into_iter().filter(|d| !d.is_null()));
281 })
282}
283
284fn dict_agg<'a, I>(datums: I, temp_storage: &'a RowArena, order_by: &[ColumnOrder]) -> Datum<'a>
285where
286 I: IntoIterator<Item = Datum<'a>>,
287{
288 let datums = order_aggregate_datums(datums, order_by);
289 temp_storage.make_datum(|packer| {
290 let mut datums: Vec<_> = datums
291 .into_iter()
292 .filter_map(|d| {
293 if d.is_null() {
294 return None;
295 }
296 let mut list = d.unwrap_list().iter();
297 let key = list.next().unwrap();
298 let val = list.next().unwrap();
299 if key.is_null() {
300 None
303 } else {
304 Some((key.unwrap_str(), val))
305 }
306 })
307 .collect();
308 datums.sort_by_key(|(k, _v)| *k);
314 datums.reverse();
315 datums.dedup_by_key(|(k, _v)| *k);
316 datums.reverse();
317 packer.push_dict(datums);
318 })
319}
320
321pub fn order_aggregate_datums<'a: 'b, 'b, I>(
331 datums: I,
332 order_by: &[ColumnOrder],
333) -> impl Iterator<Item = Datum<'b>>
334where
335 I: IntoIterator<Item = Datum<'a>>,
336{
337 order_aggregate_datums_with_rank_inner(datums, order_by)
338 .into_iter()
339 .map(|(payload, _order_datums)| payload)
341}
342
343fn order_aggregate_datums_with_rank<'a, I>(
346 datums: I,
347 order_by: &[ColumnOrder],
348) -> impl Iterator<Item = (Datum<'a>, Row)>
349where
350 I: IntoIterator<Item = Datum<'a>>,
351{
352 order_aggregate_datums_with_rank_inner(datums, order_by)
353 .into_iter()
354 .map(|(payload, order_by_datums)| (payload, Row::pack(order_by_datums)))
355}
356
357fn order_aggregate_datums_with_rank_inner<'a, I>(
358 datums: I,
359 order_by: &[ColumnOrder],
360) -> Vec<(Datum<'a>, Vec<Datum<'a>>)>
361where
362 I: IntoIterator<Item = Datum<'a>>,
363{
364 let mut decoded: Vec<(Datum, Vec<Datum>)> = datums
365 .into_iter()
366 .map(|d| {
367 let list = d.unwrap_list();
368 let mut list_it = list.iter();
369 let payload = list_it.next().unwrap();
370
371 let mut order_by_datums = Vec::with_capacity(order_by.len());
381 for _ in 0..order_by.len() {
382 order_by_datums.push(
383 list_it
384 .next()
385 .expect("must have exactly the same number of Datums as `order_by`"),
386 );
387 }
388
389 (payload, order_by_datums)
390 })
391 .collect();
392
393 let mut sort_by =
394 |(payload_left, left_order_by_datums): &(Datum, Vec<Datum>),
395 (payload_right, right_order_by_datums): &(Datum, Vec<Datum>)| {
396 compare_columns(
397 order_by,
398 left_order_by_datums,
399 right_order_by_datums,
400 || payload_left.cmp(payload_right),
401 )
402 };
403 decoded.sort_unstable_by(&mut sort_by);
408 decoded
409}
410
411fn array_concat<'a, I>(datums: I, temp_storage: &'a RowArena, order_by: &[ColumnOrder]) -> Datum<'a>
412where
413 I: IntoIterator<Item = Datum<'a>>,
414{
415 let datums = order_aggregate_datums(datums, order_by);
416 let datums: Vec<_> = datums
417 .into_iter()
418 .map(|d| d.unwrap_array().elements().iter())
419 .flatten()
420 .collect();
421 let dims = ArrayDimension {
422 lower_bound: 1,
423 length: datums.len(),
424 };
425 temp_storage.make_datum(|packer| {
426 packer.try_push_array(&[dims], datums).unwrap();
427 })
428}
429
430fn list_concat<'a, I>(datums: I, temp_storage: &'a RowArena, order_by: &[ColumnOrder]) -> Datum<'a>
431where
432 I: IntoIterator<Item = Datum<'a>>,
433{
434 let datums = order_aggregate_datums(datums, order_by);
435 temp_storage.make_datum(|packer| {
436 packer.push_list(datums.into_iter().map(|d| d.unwrap_list().iter()).flatten());
437 })
438}
439
440fn row_number<'a, I>(
444 datums: I,
445 callers_temp_storage: &'a RowArena,
446 order_by: &[ColumnOrder],
447) -> Datum<'a>
448where
449 I: IntoIterator<Item = Datum<'a>>,
450{
451 let temp_storage = RowArena::new();
455 let datums = row_number_no_list(datums, &temp_storage, order_by);
456
457 callers_temp_storage.make_datum(|packer| {
458 packer.push_list(datums);
459 })
460}
461
462fn row_number_no_list<'a: 'b, 'b, I>(
465 datums: I,
466 callers_temp_storage: &'b RowArena,
467 order_by: &[ColumnOrder],
468) -> impl Iterator<Item = Datum<'b>>
469where
470 I: IntoIterator<Item = Datum<'a>>,
471{
472 let datums = order_aggregate_datums(datums, order_by);
473
474 callers_temp_storage.reserve(datums.size_hint().0);
475 #[allow(clippy::disallowed_methods)]
476 datums
477 .into_iter()
478 .map(|d| d.unwrap_list().iter())
479 .flatten()
480 .zip(1i64..)
481 .map(|(d, i)| {
482 callers_temp_storage.make_datum(|packer| {
483 packer.push_list_with(|packer| {
484 packer.push(Datum::Int64(i));
485 packer.push(d);
486 });
487 })
488 })
489}
490
491fn rank<'a, I>(datums: I, callers_temp_storage: &'a RowArena, order_by: &[ColumnOrder]) -> Datum<'a>
495where
496 I: IntoIterator<Item = Datum<'a>>,
497{
498 let temp_storage = RowArena::new();
499 let datums = rank_no_list(datums, &temp_storage, order_by);
500
501 callers_temp_storage.make_datum(|packer| {
502 packer.push_list(datums);
503 })
504}
505
506fn rank_no_list<'a: 'b, 'b, I>(
509 datums: I,
510 callers_temp_storage: &'b RowArena,
511 order_by: &[ColumnOrder],
512) -> impl Iterator<Item = Datum<'b>>
513where
514 I: IntoIterator<Item = Datum<'a>>,
515{
516 let datums = order_aggregate_datums_with_rank(datums, order_by);
518
519 let mut datums = datums
520 .into_iter()
521 .map(|(d0, order_row)| {
522 d0.unwrap_list()
523 .iter()
524 .map(move |d1| (d1, order_row.clone()))
525 })
526 .flatten();
527
528 callers_temp_storage.reserve(datums.size_hint().0);
529 datums
530 .next()
531 .map_or(vec![], |(first_datum, first_order_row)| {
532 datums.fold(
535 (first_order_row, 1, 1, vec![(first_datum, 1)]),
536 |mut acc, (next_datum, next_order_row)| {
537 let (ref mut acc_row, ref mut acc_rank, ref mut acc_row_num, ref mut output) = acc;
538 *acc_row_num += 1;
539 if *acc_row != next_order_row {
541 *acc_rank = *acc_row_num;
542 *acc_row = next_order_row;
543 }
544
545 (*output).push((next_datum, *acc_rank));
546 acc
547 })
548 }.3).into_iter().map(|(d, i)| {
549 callers_temp_storage.make_datum(|packer| {
550 packer.push_list_with(|packer| {
551 packer.push(Datum::Int64(i));
552 packer.push(d);
553 });
554 })
555 })
556}
557
558fn dense_rank<'a, I>(
562 datums: I,
563 callers_temp_storage: &'a RowArena,
564 order_by: &[ColumnOrder],
565) -> Datum<'a>
566where
567 I: IntoIterator<Item = Datum<'a>>,
568{
569 let temp_storage = RowArena::new();
570 let datums = dense_rank_no_list(datums, &temp_storage, order_by);
571
572 callers_temp_storage.make_datum(|packer| {
573 packer.push_list(datums);
574 })
575}
576
577fn dense_rank_no_list<'a: 'b, 'b, I>(
580 datums: I,
581 callers_temp_storage: &'b RowArena,
582 order_by: &[ColumnOrder],
583) -> impl Iterator<Item = Datum<'b>>
584where
585 I: IntoIterator<Item = Datum<'a>>,
586{
587 let datums = order_aggregate_datums_with_rank(datums, order_by);
589
590 let mut datums = datums
591 .into_iter()
592 .map(|(d0, order_row)| {
593 d0.unwrap_list()
594 .iter()
595 .map(move |d1| (d1, order_row.clone()))
596 })
597 .flatten();
598
599 callers_temp_storage.reserve(datums.size_hint().0);
600 datums
601 .next()
602 .map_or(vec![], |(first_datum, first_order_row)| {
603 datums.fold(
606 (first_order_row, 1, vec![(first_datum, 1)]),
607 |mut acc, (next_datum, next_order_row)| {
608 let (ref mut acc_row, ref mut acc_rank, ref mut output) = acc;
609 if *acc_row != next_order_row {
611 *acc_rank += 1;
612 *acc_row = next_order_row;
613 }
614
615 (*output).push((next_datum, *acc_rank));
616 acc
617 })
618 }.2).into_iter().map(|(d, i)| {
619 callers_temp_storage.make_datum(|packer| {
620 packer.push_list_with(|packer| {
621 packer.push(Datum::Int64(i));
622 packer.push(d);
623 });
624 })
625 })
626}
627
628fn lag_lead<'a, I>(
650 datums: I,
651 callers_temp_storage: &'a RowArena,
652 order_by: &[ColumnOrder],
653 lag_lead_type: &LagLeadType,
654 ignore_nulls: &bool,
655) -> Datum<'a>
656where
657 I: IntoIterator<Item = Datum<'a>>,
658{
659 let temp_storage = RowArena::new();
660 let iter = lag_lead_no_list(datums, &temp_storage, order_by, lag_lead_type, ignore_nulls);
661 callers_temp_storage.make_datum(|packer| {
662 packer.push_list(iter);
663 })
664}
665
666fn lag_lead_no_list<'a: 'b, 'b, I>(
669 datums: I,
670 callers_temp_storage: &'b RowArena,
671 order_by: &[ColumnOrder],
672 lag_lead_type: &LagLeadType,
673 ignore_nulls: &bool,
674) -> impl Iterator<Item = Datum<'b>>
675where
676 I: IntoIterator<Item = Datum<'a>>,
677{
678 let datums = order_aggregate_datums(datums, order_by);
680
681 let (orig_rows, unwrapped_args): (Vec<_>, Vec<_>) = datums
685 .into_iter()
686 .map(|d| {
687 let mut iter = d.unwrap_list().iter();
688 let original_row = iter.next().unwrap();
689 let (input_value, offset, default_value) =
690 unwrap_lag_lead_encoded_args(iter.next().unwrap());
691 (original_row, (input_value, offset, default_value))
692 })
693 .unzip();
694
695 let result = lag_lead_inner(unwrapped_args, lag_lead_type, ignore_nulls);
696
697 callers_temp_storage.reserve(result.len());
698 result
699 .into_iter()
700 .zip_eq(orig_rows)
701 .map(|(result_value, original_row)| {
702 callers_temp_storage.make_datum(|packer| {
703 packer.push_list_with(|packer| {
704 packer.push(result_value);
705 packer.push(original_row);
706 });
707 })
708 })
709}
710
711fn unwrap_lag_lead_encoded_args(encoded_args: Datum) -> (Datum, Datum, Datum) {
713 let mut encoded_args_iter = encoded_args.unwrap_list().iter();
714 let (input_value, offset, default_value) = (
715 encoded_args_iter.next().unwrap(),
716 encoded_args_iter.next().unwrap(),
717 encoded_args_iter.next().unwrap(),
718 );
719 (input_value, offset, default_value)
720}
721
722fn lag_lead_inner<'a>(
725 args: Vec<(Datum<'a>, Datum<'a>, Datum<'a>)>,
726 lag_lead_type: &LagLeadType,
727 ignore_nulls: &bool,
728) -> Vec<Datum<'a>> {
729 if *ignore_nulls {
730 lag_lead_inner_ignore_nulls(args, lag_lead_type)
731 } else {
732 lag_lead_inner_respect_nulls(args, lag_lead_type)
733 }
734}
735
736fn lag_lead_inner_respect_nulls<'a>(
737 args: Vec<(Datum<'a>, Datum<'a>, Datum<'a>)>,
738 lag_lead_type: &LagLeadType,
739) -> Vec<Datum<'a>> {
740 let mut result: Vec<Datum> = Vec::with_capacity(args.len());
741 for (idx, (_, offset, default_value)) in args.iter().enumerate() {
742 if offset.is_null() {
744 result.push(Datum::Null);
745 continue;
746 }
747
748 let idx = i64::try_from(idx).expect("Array index does not fit in i64");
749 let offset = i64::from(offset.unwrap_int32());
750 let offset = match lag_lead_type {
751 LagLeadType::Lag => -offset,
752 LagLeadType::Lead => offset,
753 };
754
755 let datums_get = |i: i64| -> Option<Datum> {
757 match u64::try_from(i) {
758 Ok(i) => args
759 .get(usize::cast_from(i))
760 .map(|d| Some(d.0)) .unwrap_or(None), Err(_) => None, }
764 };
765
766 let lagged_value = datums_get(idx + offset).unwrap_or(*default_value);
767
768 result.push(lagged_value);
769 }
770
771 result
772}
773
774#[allow(clippy::as_conversions)]
778fn lag_lead_inner_ignore_nulls<'a>(
779 args: Vec<(Datum<'a>, Datum<'a>, Datum<'a>)>,
780 lag_lead_type: &LagLeadType,
781) -> Vec<Datum<'a>> {
782 if i64::try_from(args.len()).is_err() {
785 panic!("window partition way too big")
786 }
787 let mut skip_nulls_backward = vec![None; args.len()];
790 let mut last_non_null: i64 = -1;
791 let pairs = args
792 .iter()
793 .enumerate()
794 .zip_eq(skip_nulls_backward.iter_mut());
795 for ((i, (d, _, _)), slot) in pairs {
796 if d.is_null() {
797 *slot = Some(last_non_null);
798 } else {
799 last_non_null = i as i64;
800 }
801 }
802 let mut skip_nulls_forward = vec![None; args.len()];
803 let mut last_non_null: i64 = args.len() as i64;
804 let pairs = args
805 .iter()
806 .enumerate()
807 .rev()
808 .zip_eq(skip_nulls_forward.iter_mut().rev());
809 for ((i, (d, _, _)), slot) in pairs {
810 if d.is_null() {
811 *slot = Some(last_non_null);
812 } else {
813 last_non_null = i as i64;
814 }
815 }
816
817 let mut result: Vec<Datum> = Vec::with_capacity(args.len());
819 for (idx, (_, offset, default_value)) in args.iter().enumerate() {
820 if offset.is_null() {
822 result.push(Datum::Null);
823 continue;
824 }
825
826 let idx = idx as i64; let offset = i64::cast_from(offset.unwrap_int32());
828 let offset = match lag_lead_type {
829 LagLeadType::Lag => -offset,
830 LagLeadType::Lead => offset,
831 };
832 let increment = offset.signum();
833
834 let datums_get = |i: i64| -> Option<Datum> {
836 match u64::try_from(i) {
837 Ok(i) => args
838 .get(usize::cast_from(i))
839 .map(|d| Some(d.0)) .unwrap_or(None), Err(_) => None, }
843 };
844
845 let lagged_value = if increment != 0 {
846 let mut j = idx;
856 for _ in 0..num::abs(offset) {
857 j += increment;
858 if datums_get(j).is_some_and(|d| d.is_null()) {
860 let ju = j as usize; if increment > 0 {
862 j = skip_nulls_forward[ju].expect("checked above that it's null");
863 } else {
864 j = skip_nulls_backward[ju].expect("checked above that it's null");
865 }
866 }
867 if datums_get(j).is_none() {
868 break;
869 }
870 }
871 match datums_get(j) {
872 Some(datum) => datum,
873 None => *default_value,
874 }
875 } else {
876 assert_eq!(offset, 0);
877 let datum = datums_get(idx).expect("known to exist");
878 if !datum.is_null() {
879 datum
880 } else {
881 panic!("0 offset in lag/lead IGNORE NULLS");
886 }
887 };
888
889 result.push(lagged_value);
890 }
891
892 result
893}
894
895fn first_value<'a, I>(
897 datums: I,
898 callers_temp_storage: &'a RowArena,
899 order_by: &[ColumnOrder],
900 window_frame: &WindowFrame,
901) -> Datum<'a>
902where
903 I: IntoIterator<Item = Datum<'a>>,
904{
905 let temp_storage = RowArena::new();
906 let iter = first_value_no_list(datums, &temp_storage, order_by, window_frame);
907 callers_temp_storage.make_datum(|packer| {
908 packer.push_list(iter);
909 })
910}
911
912fn first_value_no_list<'a: 'b, 'b, I>(
915 datums: I,
916 callers_temp_storage: &'b RowArena,
917 order_by: &[ColumnOrder],
918 window_frame: &WindowFrame,
919) -> impl Iterator<Item = Datum<'b>>
920where
921 I: IntoIterator<Item = Datum<'a>>,
922{
923 let datums = order_aggregate_datums(datums, order_by);
925
926 let (orig_rows, args): (Vec<_>, Vec<_>) = datums
928 .into_iter()
929 .map(|d| {
930 let mut iter = d.unwrap_list().iter();
931 let original_row = iter.next().unwrap();
932 let arg = iter.next().unwrap();
933
934 (original_row, arg)
935 })
936 .unzip();
937
938 let results = first_value_inner(args, window_frame);
939
940 callers_temp_storage.reserve(results.len());
941 results
942 .into_iter()
943 .zip_eq(orig_rows)
944 .map(|(result_value, original_row)| {
945 callers_temp_storage.make_datum(|packer| {
946 packer.push_list_with(|packer| {
947 packer.push(result_value);
948 packer.push(original_row);
949 });
950 })
951 })
952}
953
954fn first_value_inner<'a>(datums: Vec<Datum<'a>>, window_frame: &WindowFrame) -> Vec<Datum<'a>> {
955 let length = datums.len();
956 let mut result: Vec<Datum> = Vec::with_capacity(length);
957 for (idx, current_datum) in datums.iter().enumerate() {
958 let first_value = match &window_frame.start_bound {
959 WindowFrameBound::CurrentRow => *current_datum,
961 WindowFrameBound::UnboundedPreceding => {
962 if let WindowFrameBound::OffsetPreceding(end_offset) = &window_frame.end_bound {
963 let end_offset = usize::cast_from(*end_offset);
964
965 if idx < end_offset {
967 Datum::Null
968 } else {
969 datums[0]
970 }
971 } else {
972 datums[0]
973 }
974 }
975 WindowFrameBound::OffsetPreceding(offset) => {
976 let start_offset = usize::cast_from(*offset);
977 let start_idx = idx.saturating_sub(start_offset);
978 if let WindowFrameBound::OffsetPreceding(end_offset) = &window_frame.end_bound {
979 let end_offset = usize::cast_from(*end_offset);
980
981 if start_offset < end_offset || idx < end_offset {
983 Datum::Null
984 } else {
985 datums[start_idx]
986 }
987 } else {
988 datums[start_idx]
989 }
990 }
991 WindowFrameBound::OffsetFollowing(offset) => {
992 let start_offset = usize::cast_from(*offset);
993 let start_idx = idx.saturating_add(start_offset);
994 if let WindowFrameBound::OffsetFollowing(end_offset) = &window_frame.end_bound {
995 if offset > end_offset || start_idx >= length {
997 Datum::Null
998 } else {
999 datums[start_idx]
1000 }
1001 } else {
1002 datums
1003 .get(start_idx)
1004 .map(|d| d.clone())
1005 .unwrap_or(Datum::Null)
1006 }
1007 }
1008 WindowFrameBound::UnboundedFollowing => unreachable!(),
1010 };
1011 result.push(first_value);
1012 }
1013 result
1014}
1015
1016fn last_value<'a, I>(
1018 datums: I,
1019 callers_temp_storage: &'a RowArena,
1020 order_by: &[ColumnOrder],
1021 window_frame: &WindowFrame,
1022) -> Datum<'a>
1023where
1024 I: IntoIterator<Item = Datum<'a>>,
1025{
1026 let temp_storage = RowArena::new();
1027 let iter = last_value_no_list(datums, &temp_storage, order_by, window_frame);
1028 callers_temp_storage.make_datum(|packer| {
1029 packer.push_list(iter);
1030 })
1031}
1032
1033fn last_value_no_list<'a: 'b, 'b, I>(
1036 datums: I,
1037 callers_temp_storage: &'b RowArena,
1038 order_by: &[ColumnOrder],
1039 window_frame: &WindowFrame,
1040) -> impl Iterator<Item = Datum<'b>>
1041where
1042 I: IntoIterator<Item = Datum<'a>>,
1043{
1044 let datums = order_aggregate_datums_with_rank(datums, order_by);
1047
1048 let size_hint = datums.size_hint().0;
1050 let mut args = Vec::with_capacity(size_hint);
1051 let mut original_rows = Vec::with_capacity(size_hint);
1052 let mut order_by_rows = Vec::with_capacity(size_hint);
1053 for (d, order_by_row) in datums.into_iter() {
1054 let mut iter = d.unwrap_list().iter();
1055 let original_row = iter.next().unwrap();
1056 let arg = iter.next().unwrap();
1057 order_by_rows.push(order_by_row);
1058 original_rows.push(original_row);
1059 args.push(arg);
1060 }
1061
1062 let results = last_value_inner(args, &order_by_rows, window_frame);
1063
1064 callers_temp_storage.reserve(results.len());
1065 results
1066 .into_iter()
1067 .zip_eq(original_rows)
1068 .map(|(result_value, original_row)| {
1069 callers_temp_storage.make_datum(|packer| {
1070 packer.push_list_with(|packer| {
1071 packer.push(result_value);
1072 packer.push(original_row);
1073 });
1074 })
1075 })
1076}
1077
1078fn last_value_inner<'a>(
1079 args: Vec<Datum<'a>>,
1080 order_by_rows: &Vec<Row>,
1081 window_frame: &WindowFrame,
1082) -> Vec<Datum<'a>> {
1083 let length = args.len();
1084 let mut results: Vec<Datum> = Vec::with_capacity(length);
1085 for (idx, (current_datum, order_by_row)) in args.iter().zip_eq(order_by_rows).enumerate() {
1086 let last_value = match &window_frame.end_bound {
1087 WindowFrameBound::CurrentRow => match &window_frame.units {
1088 WindowFrameUnits::Rows => *current_datum,
1090 WindowFrameUnits::Range => {
1091 let target_idx = order_by_rows[idx..]
1096 .iter()
1097 .enumerate()
1098 .take_while(|(_, row)| *row == order_by_row)
1099 .last()
1100 .unwrap()
1101 .0
1102 + idx;
1103 args[target_idx]
1104 }
1105 WindowFrameUnits::Groups => unreachable!(),
1107 },
1108 WindowFrameBound::UnboundedFollowing => {
1109 if let WindowFrameBound::OffsetFollowing(start_offset) = &window_frame.start_bound {
1110 let start_offset = usize::cast_from(*start_offset);
1111
1112 if idx + start_offset > length - 1 {
1114 Datum::Null
1115 } else {
1116 args[length - 1]
1117 }
1118 } else {
1119 args[length - 1]
1120 }
1121 }
1122 WindowFrameBound::OffsetFollowing(offset) => {
1123 let end_offset = usize::cast_from(*offset);
1124 let end_idx = idx.saturating_add(end_offset);
1125 if let WindowFrameBound::OffsetFollowing(start_offset) = &window_frame.start_bound {
1126 let start_offset = usize::cast_from(*start_offset);
1127 let start_idx = idx.saturating_add(start_offset);
1128
1129 if end_offset < start_offset || start_idx >= length {
1131 Datum::Null
1132 } else {
1133 args.get(end_idx).unwrap_or(&args[length - 1]).clone()
1135 }
1136 } else {
1137 args.get(end_idx).unwrap_or(&args[length - 1]).clone()
1138 }
1139 }
1140 WindowFrameBound::OffsetPreceding(offset) => {
1141 let end_offset = usize::cast_from(*offset);
1142 let end_idx = idx.saturating_sub(end_offset);
1143 if idx < end_offset {
1144 Datum::Null
1146 } else if let WindowFrameBound::OffsetPreceding(start_offset) =
1147 &window_frame.start_bound
1148 {
1149 if offset > start_offset {
1151 Datum::Null
1152 } else {
1153 args[end_idx]
1154 }
1155 } else {
1156 args[end_idx]
1157 }
1158 }
1159 WindowFrameBound::UnboundedPreceding => unreachable!(),
1161 };
1162 results.push(last_value);
1163 }
1164 results
1165}
1166
1167fn fused_value_window_func<'a, I>(
1173 input_datums: I,
1174 callers_temp_storage: &'a RowArena,
1175 funcs: &Vec<AggregateFunc>,
1176 order_by: &Vec<ColumnOrder>,
1177) -> Datum<'a>
1178where
1179 I: IntoIterator<Item = Datum<'a>>,
1180{
1181 let temp_storage = RowArena::new();
1182 let iter = fused_value_window_func_no_list(input_datums, &temp_storage, funcs, order_by);
1183 callers_temp_storage.make_datum(|packer| {
1184 packer.push_list(iter);
1185 })
1186}
1187
1188fn fused_value_window_func_no_list<'a: 'b, 'b, I>(
1191 input_datums: I,
1192 callers_temp_storage: &'b RowArena,
1193 funcs: &Vec<AggregateFunc>,
1194 order_by: &Vec<ColumnOrder>,
1195) -> impl Iterator<Item = Datum<'b>>
1196where
1197 I: IntoIterator<Item = Datum<'a>>,
1198{
1199 let has_last_value = funcs
1200 .iter()
1201 .any(|f| matches!(f, AggregateFunc::LastValue { .. }));
1202
1203 let input_datums_with_ranks = order_aggregate_datums_with_rank(input_datums, order_by);
1204
1205 let size_hint = input_datums_with_ranks.size_hint().0;
1206 let mut encoded_argsss = vec![Vec::with_capacity(size_hint); funcs.len()];
1207 let mut original_rows = Vec::with_capacity(size_hint);
1208 let mut order_by_rows = Vec::with_capacity(size_hint);
1209 for (d, order_by_row) in input_datums_with_ranks {
1210 let mut iter = d.unwrap_list().iter();
1211 let original_row = iter.next().unwrap();
1212 original_rows.push(original_row);
1213 let mut argss_iter = iter.next().unwrap().unwrap_list().iter();
1214 for i in 0..funcs.len() {
1215 let encoded_args = argss_iter.next().unwrap();
1216 encoded_argsss[i].push(encoded_args);
1217 }
1218 if has_last_value {
1219 order_by_rows.push(order_by_row);
1220 }
1221 }
1222
1223 let mut results_per_row = vec![Vec::with_capacity(funcs.len()); original_rows.len()];
1224 for (func, encoded_argss) in funcs.iter().zip_eq(encoded_argsss) {
1225 let results = match func {
1226 AggregateFunc::LagLead {
1227 order_by: inner_order_by,
1228 lag_lead,
1229 ignore_nulls,
1230 } => {
1231 assert_eq!(order_by, inner_order_by);
1232 let unwrapped_argss = encoded_argss
1233 .into_iter()
1234 .map(|encoded_args| unwrap_lag_lead_encoded_args(encoded_args))
1235 .collect();
1236 lag_lead_inner(unwrapped_argss, lag_lead, ignore_nulls)
1237 }
1238 AggregateFunc::FirstValue {
1239 order_by: inner_order_by,
1240 window_frame,
1241 } => {
1242 assert_eq!(order_by, inner_order_by);
1243 first_value_inner(encoded_argss, window_frame)
1246 }
1247 AggregateFunc::LastValue {
1248 order_by: inner_order_by,
1249 window_frame,
1250 } => {
1251 assert_eq!(order_by, inner_order_by);
1252 last_value_inner(encoded_argss, &order_by_rows, window_frame)
1255 }
1256 _ => panic!("unknown window function in FusedValueWindowFunc"),
1257 };
1258 for (results, result) in results_per_row.iter_mut().zip_eq(results) {
1259 results.push(result);
1260 }
1261 }
1262
1263 callers_temp_storage.reserve(2 * original_rows.len());
1264 results_per_row
1265 .into_iter()
1266 .enumerate()
1267 .map(move |(i, results)| {
1268 callers_temp_storage.make_datum(|packer| {
1269 packer.push_list_with(|packer| {
1270 packer
1271 .push(callers_temp_storage.make_datum(|packer| packer.push_list(results)));
1272 packer.push(original_rows[i]);
1273 });
1274 })
1275 })
1276}
1277
1278fn window_aggr<'a, I, A>(
1287 input_datums: I,
1288 callers_temp_storage: &'a RowArena,
1289 wrapped_aggregate: &AggregateFunc,
1290 order_by: &[ColumnOrder],
1291 window_frame: &WindowFrame,
1292) -> Datum<'a>
1293where
1294 I: IntoIterator<Item = Datum<'a>>,
1295 A: OneByOneAggr,
1296{
1297 let temp_storage = RowArena::new();
1298 let iter = window_aggr_no_list::<I, A>(
1299 input_datums,
1300 &temp_storage,
1301 wrapped_aggregate,
1302 order_by,
1303 window_frame,
1304 );
1305 callers_temp_storage.make_datum(|packer| {
1306 packer.push_list(iter);
1307 })
1308}
1309
1310fn window_aggr_no_list<'a: 'b, 'b, I, A>(
1313 input_datums: I,
1314 callers_temp_storage: &'b RowArena,
1315 wrapped_aggregate: &AggregateFunc,
1316 order_by: &[ColumnOrder],
1317 window_frame: &WindowFrame,
1318) -> impl Iterator<Item = Datum<'b>>
1319where
1320 I: IntoIterator<Item = Datum<'a>>,
1321 A: OneByOneAggr,
1322{
1323 let datums = order_aggregate_datums_with_rank(input_datums, order_by);
1326
1327 let size_hint = datums.size_hint().0;
1329 let mut args: Vec<Datum> = Vec::with_capacity(size_hint);
1330 let mut original_rows: Vec<Datum> = Vec::with_capacity(size_hint);
1331 let mut order_by_rows = Vec::with_capacity(size_hint);
1332 for (d, order_by_row) in datums.into_iter() {
1333 let mut iter = d.unwrap_list().iter();
1334 let original_row = iter.next().unwrap();
1335 let arg = iter.next().unwrap();
1336 order_by_rows.push(order_by_row);
1337 original_rows.push(original_row);
1338 args.push(arg);
1339 }
1340
1341 let results = window_aggr_inner::<A>(
1342 args,
1343 &order_by_rows,
1344 wrapped_aggregate,
1345 order_by,
1346 window_frame,
1347 callers_temp_storage,
1348 );
1349
1350 callers_temp_storage.reserve(results.len());
1351 results
1352 .into_iter()
1353 .zip_eq(original_rows)
1354 .map(|(result_value, original_row)| {
1355 callers_temp_storage.make_datum(|packer| {
1356 packer.push_list_with(|packer| {
1357 packer.push(result_value);
1358 packer.push(original_row);
1359 });
1360 })
1361 })
1362}
1363
1364fn window_aggr_inner<'a, A>(
1365 mut args: Vec<Datum<'a>>,
1366 order_by_rows: &Vec<Row>,
1367 wrapped_aggregate: &AggregateFunc,
1368 order_by: &[ColumnOrder],
1369 window_frame: &WindowFrame,
1370 temp_storage: &'a RowArena,
1371) -> Vec<Datum<'a>>
1372where
1373 A: OneByOneAggr,
1374{
1375 let length = args.len();
1376 let mut result: Vec<Datum> = Vec::with_capacity(length);
1377
1378 soft_assert_or_log!(
1384 !((matches!(window_frame.units, WindowFrameUnits::Groups)
1385 || matches!(window_frame.units, WindowFrameUnits::Range))
1386 && !window_frame.includes_current_row()),
1387 "window frame without current row"
1388 );
1389
1390 if (matches!(
1391 window_frame.start_bound,
1392 WindowFrameBound::UnboundedPreceding
1393 ) && matches!(window_frame.end_bound, WindowFrameBound::UnboundedFollowing))
1394 || (order_by.is_empty()
1395 && (matches!(window_frame.units, WindowFrameUnits::Groups)
1396 || matches!(window_frame.units, WindowFrameUnits::Range))
1397 && window_frame.includes_current_row())
1398 {
1399 let result_value =
1406 wrapped_aggregate.eval(args.into_iter().map(|d| (d, Diff::ONE)), temp_storage);
1407 for _ in 0..length {
1409 result.push(result_value);
1410 }
1411 } else {
1412 fn rows_between_unbounded_preceding_and_current_row<'a, A>(
1413 args: Vec<Datum<'a>>,
1414 result: &mut Vec<Datum<'a>>,
1415 mut one_by_one_aggr: A,
1416 temp_storage: &'a RowArena,
1417 ) where
1418 A: OneByOneAggr,
1419 {
1420 for current_arg in args.into_iter() {
1421 one_by_one_aggr.give(¤t_arg);
1422 let result_value = one_by_one_aggr.get_current_aggregate(temp_storage);
1423 result.push(result_value);
1424 }
1425 }
1426
1427 fn groups_between_unbounded_preceding_and_current_row<'a, A>(
1428 args: Vec<Datum<'a>>,
1429 order_by_rows: &Vec<Row>,
1430 result: &mut Vec<Datum<'a>>,
1431 mut one_by_one_aggr: A,
1432 temp_storage: &'a RowArena,
1433 ) where
1434 A: OneByOneAggr,
1435 {
1436 let mut peer_group_start = 0;
1437 while peer_group_start < args.len() {
1438 let mut peer_group_end = peer_group_start + 1;
1442 while peer_group_end < args.len()
1443 && order_by_rows[peer_group_start] == order_by_rows[peer_group_end]
1444 {
1445 peer_group_end += 1;
1447 }
1448 for current_arg in args[peer_group_start..peer_group_end].iter() {
1451 one_by_one_aggr.give(current_arg);
1452 }
1453 let agg_for_peer_group = one_by_one_aggr.get_current_aggregate(temp_storage);
1454 for _ in args[peer_group_start..peer_group_end].iter() {
1456 result.push(agg_for_peer_group);
1457 }
1458 peer_group_start = peer_group_end;
1460 }
1461 }
1462
1463 fn rows_between_offset_and_offset<'a>(
1464 args: Vec<Datum<'a>>,
1465 result: &mut Vec<Datum<'a>>,
1466 wrapped_aggregate: &AggregateFunc,
1467 temp_storage: &'a RowArena,
1468 offset_start: i64,
1469 offset_end: i64,
1470 ) {
1471 let len = args
1472 .len()
1473 .to_i64()
1474 .expect("window partition's len should fit into i64");
1475 for i in 0..len {
1476 let i = i.to_i64().expect("window partition shouldn't be super big");
1477 let frame_start = max(i + offset_start, 0)
1480 .to_usize()
1481 .expect("The max made sure it's not negative");
1482 let frame_end = min(i + offset_end, len - 1).to_usize();
1485 match frame_end {
1486 Some(frame_end) => {
1487 if frame_start <= frame_end {
1488 let frame_values = args[frame_start..=frame_end]
1499 .iter()
1500 .map(|d| (*d, Diff::ONE));
1501 let result_value = wrapped_aggregate.eval(frame_values, temp_storage);
1502 result.push(result_value);
1503 } else {
1504 let result_value = wrapped_aggregate.default();
1506 result.push(result_value);
1507 }
1508 }
1509 None => {
1510 let result_value = wrapped_aggregate.default();
1512 result.push(result_value);
1513 }
1514 }
1515 }
1516 }
1517
1518 match (
1519 &window_frame.units,
1520 &window_frame.start_bound,
1521 &window_frame.end_bound,
1522 ) {
1523 (Rows, UnboundedPreceding, CurrentRow) => {
1528 rows_between_unbounded_preceding_and_current_row::<A>(
1529 args,
1530 &mut result,
1531 A::new(wrapped_aggregate, false),
1532 temp_storage,
1533 );
1534 }
1535 (Rows, CurrentRow, UnboundedFollowing) => {
1536 args.reverse();
1538 rows_between_unbounded_preceding_and_current_row::<A>(
1539 args,
1540 &mut result,
1541 A::new(wrapped_aggregate, true),
1542 temp_storage,
1543 );
1544 result.reverse();
1545 }
1546 (Range, UnboundedPreceding, CurrentRow) => {
1547 groups_between_unbounded_preceding_and_current_row::<A>(
1550 args,
1551 order_by_rows,
1552 &mut result,
1553 A::new(wrapped_aggregate, false),
1554 temp_storage,
1555 );
1556 }
1557 (Rows, OffsetPreceding(start_prec), OffsetPreceding(end_prec)) => {
1561 let start_prec = start_prec.to_i64().expect(
1562 "window frame start OFFSET shouldn't be super big (the planning ensured this)",
1563 );
1564 let end_prec = end_prec.to_i64().expect(
1565 "window frame end OFFSET shouldn't be super big (the planning ensured this)",
1566 );
1567 rows_between_offset_and_offset(
1568 args,
1569 &mut result,
1570 wrapped_aggregate,
1571 temp_storage,
1572 -start_prec,
1573 -end_prec,
1574 );
1575 }
1576 (Rows, OffsetPreceding(start_prec), OffsetFollowing(end_fol)) => {
1577 let start_prec = start_prec.to_i64().expect(
1578 "window frame start OFFSET shouldn't be super big (the planning ensured this)",
1579 );
1580 let end_fol = end_fol.to_i64().expect(
1581 "window frame end OFFSET shouldn't be super big (the planning ensured this)",
1582 );
1583 rows_between_offset_and_offset(
1584 args,
1585 &mut result,
1586 wrapped_aggregate,
1587 temp_storage,
1588 -start_prec,
1589 end_fol,
1590 );
1591 }
1592 (Rows, OffsetFollowing(start_fol), OffsetFollowing(end_fol)) => {
1593 let start_fol = start_fol.to_i64().expect(
1594 "window frame start OFFSET shouldn't be super big (the planning ensured this)",
1595 );
1596 let end_fol = end_fol.to_i64().expect(
1597 "window frame end OFFSET shouldn't be super big (the planning ensured this)",
1598 );
1599 rows_between_offset_and_offset(
1600 args,
1601 &mut result,
1602 wrapped_aggregate,
1603 temp_storage,
1604 start_fol,
1605 end_fol,
1606 );
1607 }
1608 (Rows, OffsetFollowing(_), OffsetPreceding(_)) => {
1609 unreachable!() }
1611 (Rows, OffsetPreceding(start_prec), CurrentRow) => {
1612 let start_prec = start_prec.to_i64().expect(
1613 "window frame start OFFSET shouldn't be super big (the planning ensured this)",
1614 );
1615 let end_fol = 0;
1616 rows_between_offset_and_offset(
1617 args,
1618 &mut result,
1619 wrapped_aggregate,
1620 temp_storage,
1621 -start_prec,
1622 end_fol,
1623 );
1624 }
1625 (Rows, CurrentRow, OffsetFollowing(end_fol)) => {
1626 let start_fol = 0;
1627 let end_fol = end_fol.to_i64().expect(
1628 "window frame end OFFSET shouldn't be super big (the planning ensured this)",
1629 );
1630 rows_between_offset_and_offset(
1631 args,
1632 &mut result,
1633 wrapped_aggregate,
1634 temp_storage,
1635 start_fol,
1636 end_fol,
1637 );
1638 }
1639 (Rows, CurrentRow, CurrentRow) => {
1640 let start_fol = 0;
1643 let end_fol = 0;
1644 rows_between_offset_and_offset(
1645 args,
1646 &mut result,
1647 wrapped_aggregate,
1648 temp_storage,
1649 start_fol,
1650 end_fol,
1651 );
1652 }
1653 (Rows, CurrentRow, OffsetPreceding(_))
1654 | (Rows, UnboundedFollowing, _)
1655 | (Rows, _, UnboundedPreceding)
1656 | (Rows, OffsetFollowing(..), CurrentRow) => {
1657 unreachable!() }
1659 (Rows, UnboundedPreceding, UnboundedFollowing) => {
1660 unreachable!()
1663 }
1664 (Rows, UnboundedPreceding, OffsetPreceding(_))
1665 | (Rows, UnboundedPreceding, OffsetFollowing(_))
1666 | (Rows, OffsetPreceding(..), UnboundedFollowing)
1667 | (Rows, OffsetFollowing(..), UnboundedFollowing) => {
1668 unreachable!()
1671 }
1672 (Range, _, _) => {
1673 unreachable!()
1680 }
1681 (Groups, _, _) => {
1682 unreachable!()
1686 }
1687 }
1688 }
1689
1690 result
1691}
1692
1693fn fused_window_aggr<'a, I, A>(
1697 input_datums: I,
1698 callers_temp_storage: &'a RowArena,
1699 wrapped_aggregates: &Vec<AggregateFunc>,
1700 order_by: &Vec<ColumnOrder>,
1701 window_frame: &WindowFrame,
1702) -> Datum<'a>
1703where
1704 I: IntoIterator<Item = Datum<'a>>,
1705 A: OneByOneAggr,
1706{
1707 let temp_storage = RowArena::new();
1708 let iter = fused_window_aggr_no_list::<_, A>(
1709 input_datums,
1710 &temp_storage,
1711 wrapped_aggregates,
1712 order_by,
1713 window_frame,
1714 );
1715 callers_temp_storage.make_datum(|packer| {
1716 packer.push_list(iter);
1717 })
1718}
1719
1720fn fused_window_aggr_no_list<'a: 'b, 'b, I, A>(
1723 input_datums: I,
1724 callers_temp_storage: &'b RowArena,
1725 wrapped_aggregates: &Vec<AggregateFunc>,
1726 order_by: &Vec<ColumnOrder>,
1727 window_frame: &WindowFrame,
1728) -> impl Iterator<Item = Datum<'b>>
1729where
1730 I: IntoIterator<Item = Datum<'a>>,
1731 A: OneByOneAggr,
1732{
1733 let datums = order_aggregate_datums_with_rank(input_datums, order_by);
1736
1737 let size_hint = datums.size_hint().0;
1738 let mut argss = vec![Vec::with_capacity(size_hint); wrapped_aggregates.len()];
1739 let mut original_rows = Vec::with_capacity(size_hint);
1740 let mut order_by_rows = Vec::with_capacity(size_hint);
1741 for (d, order_by_row) in datums {
1742 let mut iter = d.unwrap_list().iter();
1743 let original_row = iter.next().unwrap();
1744 original_rows.push(original_row);
1745 let args_iter = iter.next().unwrap().unwrap_list().iter();
1746 for (args, arg) in argss.iter_mut().zip_eq(args_iter) {
1748 args.push(arg);
1749 }
1750 order_by_rows.push(order_by_row);
1751 }
1752
1753 let mut results_per_row =
1754 vec![Vec::with_capacity(wrapped_aggregates.len()); original_rows.len()];
1755 for (wrapped_aggr, args) in wrapped_aggregates.iter().zip_eq(argss) {
1756 let results = window_aggr_inner::<A>(
1757 args,
1758 &order_by_rows,
1759 wrapped_aggr,
1760 order_by,
1761 window_frame,
1762 callers_temp_storage,
1763 );
1764 for (results, result) in results_per_row.iter_mut().zip_eq(results) {
1765 results.push(result);
1766 }
1767 }
1768
1769 callers_temp_storage.reserve(2 * original_rows.len());
1770 results_per_row
1771 .into_iter()
1772 .enumerate()
1773 .map(move |(i, results)| {
1774 callers_temp_storage.make_datum(|packer| {
1775 packer.push_list_with(|packer| {
1776 packer
1777 .push(callers_temp_storage.make_datum(|packer| packer.push_list(results)));
1778 packer.push(original_rows[i]);
1779 });
1780 })
1781 })
1782}
1783
1784pub trait OneByOneAggr {
1788 fn new(agg: &AggregateFunc, reverse: bool) -> Self;
1793 fn give(&mut self, d: &Datum);
1795 fn get_current_aggregate<'a>(&self, temp_storage: &'a RowArena) -> Datum<'a>;
1797}
1798
1799#[derive(Debug)]
1805pub struct NaiveOneByOneAggr {
1806 agg: AggregateFunc,
1807 input: Vec<Row>,
1808 reverse: bool,
1809}
1810
1811impl OneByOneAggr for NaiveOneByOneAggr {
1812 fn new(agg: &AggregateFunc, reverse: bool) -> Self {
1813 NaiveOneByOneAggr {
1814 agg: agg.clone(),
1815 input: Vec::new(),
1816 reverse,
1817 }
1818 }
1819
1820 fn give(&mut self, d: &Datum) {
1821 let mut row = Row::default();
1822 row.packer().push(d);
1823 self.input.push(row);
1824 }
1825
1826 fn get_current_aggregate<'a>(&self, temp_storage: &'a RowArena) -> Datum<'a> {
1827 temp_storage.make_datum(|packer| {
1828 packer.push(if !self.reverse {
1829 self.agg.eval(
1830 self.input.iter().map(|r| (r.unpack_first(), Diff::ONE)),
1831 temp_storage,
1832 )
1833 } else {
1834 self.agg.eval(
1835 self.input
1836 .iter()
1837 .rev()
1838 .map(|r| (r.unpack_first(), Diff::ONE)),
1839 temp_storage,
1840 )
1841 });
1842 })
1843 }
1844}
1845
1846#[derive(
1849 Clone,
1850 Debug,
1851 Eq,
1852 PartialEq,
1853 Ord,
1854 PartialOrd,
1855 Serialize,
1856 Deserialize,
1857 Hash
1858)]
1859pub enum LagLeadType {
1860 Lag,
1861 Lead,
1862}
1863
1864#[derive(
1865 Clone,
1866 Debug,
1867 Eq,
1868 PartialEq,
1869 Ord,
1870 PartialOrd,
1871 Serialize,
1872 Deserialize,
1873 Hash
1874)]
1875pub enum AggregateFunc {
1876 MaxNumeric,
1877 MaxInt16,
1878 MaxInt32,
1879 MaxInt64,
1880 MaxUInt16,
1881 MaxUInt32,
1882 MaxUInt64,
1883 MaxMzTimestamp,
1884 MaxFloat32,
1885 MaxFloat64,
1886 MaxBool,
1887 MaxString,
1888 MaxDate,
1889 MaxTimestamp,
1890 MaxTimestampTz,
1891 MaxInterval,
1892 MaxTime,
1893 MinNumeric,
1894 MinInt16,
1895 MinInt32,
1896 MinInt64,
1897 MinUInt16,
1898 MinUInt32,
1899 MinUInt64,
1900 MinMzTimestamp,
1901 MinFloat32,
1902 MinFloat64,
1903 MinBool,
1904 MinString,
1905 MinDate,
1906 MinTimestamp,
1907 MinTimestampTz,
1908 MinInterval,
1909 MinTime,
1910 SumInt16,
1911 SumInt32,
1912 SumInt64,
1913 SumUInt16,
1914 SumUInt32,
1915 SumUInt64,
1916 SumFloat32,
1917 SumFloat64,
1918 SumNumeric,
1919 Count,
1920 Any,
1921 All,
1922 JsonbAgg {
1929 order_by: Vec<ColumnOrder>,
1930 },
1931 JsonbObjectAgg {
1938 order_by: Vec<ColumnOrder>,
1939 },
1940 MapAgg {
1944 order_by: Vec<ColumnOrder>,
1945 value_type: SqlScalarType,
1946 },
1947 ArrayConcat {
1950 order_by: Vec<ColumnOrder>,
1951 },
1952 ListConcat {
1955 order_by: Vec<ColumnOrder>,
1956 },
1957 StringAgg {
1958 order_by: Vec<ColumnOrder>,
1959 },
1960 RowNumber {
1961 order_by: Vec<ColumnOrder>,
1962 },
1963 Rank {
1964 order_by: Vec<ColumnOrder>,
1965 },
1966 DenseRank {
1967 order_by: Vec<ColumnOrder>,
1968 },
1969 LagLead {
1970 order_by: Vec<ColumnOrder>,
1971 lag_lead: LagLeadType,
1972 ignore_nulls: bool,
1973 },
1974 FirstValue {
1975 order_by: Vec<ColumnOrder>,
1976 window_frame: WindowFrame,
1977 },
1978 LastValue {
1979 order_by: Vec<ColumnOrder>,
1980 window_frame: WindowFrame,
1981 },
1982 FusedValueWindowFunc {
1984 funcs: Vec<AggregateFunc>,
1985 order_by: Vec<ColumnOrder>,
1988 },
1989 WindowAggregate {
1990 wrapped_aggregate: Box<AggregateFunc>,
1991 order_by: Vec<ColumnOrder>,
1992 window_frame: WindowFrame,
1993 },
1994 FusedWindowAggregate {
1995 wrapped_aggregates: Vec<AggregateFunc>,
1996 order_by: Vec<ColumnOrder>,
1997 window_frame: WindowFrame,
1998 },
1999 Dummy,
2004}
2005
2006fn expand_counts<'a, I>(datums: I) -> impl Iterator<Item = Datum<'a>>
2012where
2013 I: IntoIterator<Item = (Datum<'a>, Diff)>,
2014{
2015 datums.into_iter().flat_map(|(datum, diff)| {
2016 let copies = usize::try_from(diff.into_inner()).unwrap_or(0);
2017 std::iter::repeat(datum).take(copies)
2018 })
2019}
2020
2021impl AggregateFunc {
2022 fn ignores_multiplicity(&self) -> bool {
2029 use AggregateFunc::*;
2030 matches!(
2031 self,
2032 MaxNumeric
2033 | MaxInt16
2034 | MaxInt32
2035 | MaxInt64
2036 | MaxUInt16
2037 | MaxUInt32
2038 | MaxUInt64
2039 | MaxMzTimestamp
2040 | MaxFloat32
2041 | MaxFloat64
2042 | MaxBool
2043 | MaxString
2044 | MaxDate
2045 | MaxTimestamp
2046 | MaxTimestampTz
2047 | MaxInterval
2048 | MaxTime
2049 | MinNumeric
2050 | MinInt16
2051 | MinInt32
2052 | MinInt64
2053 | MinUInt16
2054 | MinUInt32
2055 | MinUInt64
2056 | MinMzTimestamp
2057 | MinFloat32
2058 | MinFloat64
2059 | MinBool
2060 | MinString
2061 | MinDate
2062 | MinTimestamp
2063 | MinTimestampTz
2064 | MinInterval
2065 | MinTime
2066 | Any
2067 | All
2068 )
2069 }
2070
2071 pub fn eval<'a, I>(&self, datums: I, temp_storage: &'a RowArena) -> Datum<'a>
2078 where
2079 I: IntoIterator<Item = (Datum<'a>, Diff)>,
2080 {
2081 match self {
2091 AggregateFunc::Count => count(datums),
2092 AggregateFunc::SumInt16 | AggregateFunc::SumInt32 => {
2093 sum_signed_int_counted(datums, |accum| {
2095 #[allow(clippy::as_conversions)]
2096 let narrowed = accum as i64;
2097 Datum::Int64(narrowed)
2098 })
2099 }
2100 AggregateFunc::SumInt64 => sum_signed_int_counted(datums, Datum::from),
2101 _ if self.ignores_multiplicity() => {
2102 self.eval_datums(datums.into_iter().map(|(datum, _diff)| datum), temp_storage)
2103 }
2104 _ => self.eval_datums(expand_counts(datums), temp_storage),
2105 }
2106 }
2107
2108 fn eval_datums<'a, I>(&self, datums: I, temp_storage: &'a RowArena) -> Datum<'a>
2110 where
2111 I: IntoIterator<Item = Datum<'a>>,
2112 {
2113 match self {
2114 AggregateFunc::MaxNumeric => {
2115 max_datum::<'a, I, OrderedDecimal<numeric::Numeric>>(datums)
2116 }
2117 AggregateFunc::MaxInt16 => max_datum::<'a, I, i16>(datums),
2118 AggregateFunc::MaxInt32 => max_datum::<'a, I, i32>(datums),
2119 AggregateFunc::MaxInt64 => max_datum::<'a, I, i64>(datums),
2120 AggregateFunc::MaxUInt16 => max_datum::<'a, I, u16>(datums),
2121 AggregateFunc::MaxUInt32 => max_datum::<'a, I, u32>(datums),
2122 AggregateFunc::MaxUInt64 => max_datum::<'a, I, u64>(datums),
2123 AggregateFunc::MaxMzTimestamp => max_datum::<'a, I, mz_repr::Timestamp>(datums),
2124 AggregateFunc::MaxFloat32 => max_datum::<'a, I, OrderedFloat<f32>>(datums),
2125 AggregateFunc::MaxFloat64 => max_datum::<'a, I, OrderedFloat<f64>>(datums),
2126 AggregateFunc::MaxBool => max_datum::<'a, I, bool>(datums),
2127 AggregateFunc::MaxString => max_string(datums),
2128 AggregateFunc::MaxDate => max_datum::<'a, I, Date>(datums),
2129 AggregateFunc::MaxTimestamp => {
2130 max_datum::<'a, I, CheckedTimestamp<NaiveDateTime>>(datums)
2131 }
2132 AggregateFunc::MaxTimestampTz => {
2133 max_datum::<'a, I, CheckedTimestamp<DateTime<Utc>>>(datums)
2134 }
2135 AggregateFunc::MaxInterval => max_datum::<'a, I, Interval>(datums),
2136 AggregateFunc::MaxTime => max_datum::<'a, I, NaiveTime>(datums),
2137 AggregateFunc::MinNumeric => {
2138 min_datum::<'a, I, OrderedDecimal<numeric::Numeric>>(datums)
2139 }
2140 AggregateFunc::MinInt16 => min_datum::<'a, I, i16>(datums),
2141 AggregateFunc::MinInt32 => min_datum::<'a, I, i32>(datums),
2142 AggregateFunc::MinInt64 => min_datum::<'a, I, i64>(datums),
2143 AggregateFunc::MinUInt16 => min_datum::<'a, I, u16>(datums),
2144 AggregateFunc::MinUInt32 => min_datum::<'a, I, u32>(datums),
2145 AggregateFunc::MinUInt64 => min_datum::<'a, I, u64>(datums),
2146 AggregateFunc::MinMzTimestamp => min_datum::<'a, I, mz_repr::Timestamp>(datums),
2147 AggregateFunc::MinFloat32 => min_datum::<'a, I, OrderedFloat<f32>>(datums),
2148 AggregateFunc::MinFloat64 => min_datum::<'a, I, OrderedFloat<f64>>(datums),
2149 AggregateFunc::MinBool => min_datum::<'a, I, bool>(datums),
2150 AggregateFunc::MinString => min_string(datums),
2151 AggregateFunc::MinDate => min_datum::<'a, I, Date>(datums),
2152 AggregateFunc::MinTimestamp => {
2153 min_datum::<'a, I, CheckedTimestamp<NaiveDateTime>>(datums)
2154 }
2155 AggregateFunc::MinTimestampTz => {
2156 min_datum::<'a, I, CheckedTimestamp<DateTime<Utc>>>(datums)
2157 }
2158 AggregateFunc::MinInterval => min_datum::<'a, I, Interval>(datums),
2159 AggregateFunc::MinTime => min_datum::<'a, I, NaiveTime>(datums),
2160 AggregateFunc::SumInt16 => sum_datum::<'a, I, i16, i64>(datums),
2161 AggregateFunc::SumInt32 => sum_datum::<'a, I, i32, i64>(datums),
2162 AggregateFunc::SumInt64 => sum_datum::<'a, I, i64, i128>(datums),
2163 AggregateFunc::SumUInt16 => sum_datum::<'a, I, u16, u64>(datums),
2164 AggregateFunc::SumUInt32 => sum_datum::<'a, I, u32, u64>(datums),
2165 AggregateFunc::SumUInt64 => sum_datum::<'a, I, u64, u128>(datums),
2166 AggregateFunc::SumFloat32 => sum_datum::<'a, I, f32, f32>(datums),
2167 AggregateFunc::SumFloat64 => sum_datum::<'a, I, f64, f64>(datums),
2168 AggregateFunc::SumNumeric => sum_numeric(datums),
2169 AggregateFunc::Count => unreachable!("Count is handled in `eval`"),
2170 AggregateFunc::Any => any(datums),
2171 AggregateFunc::All => all(datums),
2172 AggregateFunc::JsonbAgg { order_by } => jsonb_agg(datums, temp_storage, order_by),
2173 AggregateFunc::MapAgg { order_by, .. } | AggregateFunc::JsonbObjectAgg { order_by } => {
2174 dict_agg(datums, temp_storage, order_by)
2175 }
2176 AggregateFunc::ArrayConcat { order_by } => array_concat(datums, temp_storage, order_by),
2177 AggregateFunc::ListConcat { order_by } => list_concat(datums, temp_storage, order_by),
2178 AggregateFunc::StringAgg { order_by } => string_agg(datums, temp_storage, order_by),
2179 AggregateFunc::RowNumber { order_by } => row_number(datums, temp_storage, order_by),
2180 AggregateFunc::Rank { order_by } => rank(datums, temp_storage, order_by),
2181 AggregateFunc::DenseRank { order_by } => dense_rank(datums, temp_storage, order_by),
2182 AggregateFunc::LagLead {
2183 order_by,
2184 lag_lead: lag_lead_type,
2185 ignore_nulls,
2186 } => lag_lead(datums, temp_storage, order_by, lag_lead_type, ignore_nulls),
2187 AggregateFunc::FirstValue {
2188 order_by,
2189 window_frame,
2190 } => first_value(datums, temp_storage, order_by, window_frame),
2191 AggregateFunc::LastValue {
2192 order_by,
2193 window_frame,
2194 } => last_value(datums, temp_storage, order_by, window_frame),
2195 AggregateFunc::WindowAggregate {
2196 wrapped_aggregate,
2197 order_by,
2198 window_frame,
2199 } => window_aggr::<_, NaiveOneByOneAggr>(
2200 datums,
2201 temp_storage,
2202 wrapped_aggregate,
2203 order_by,
2204 window_frame,
2205 ),
2206 AggregateFunc::FusedValueWindowFunc { funcs, order_by } => {
2207 fused_value_window_func(datums, temp_storage, funcs, order_by)
2208 }
2209 AggregateFunc::FusedWindowAggregate {
2210 wrapped_aggregates,
2211 order_by,
2212 window_frame,
2213 } => fused_window_aggr::<_, NaiveOneByOneAggr>(
2214 datums,
2215 temp_storage,
2216 wrapped_aggregates,
2217 order_by,
2218 window_frame,
2219 ),
2220 AggregateFunc::Dummy => Datum::Dummy,
2221 }
2222 }
2223
2224 pub fn eval_with_fast_window_agg<'a, I, W>(
2228 &self,
2229 datums: I,
2230 temp_storage: &'a RowArena,
2231 ) -> Datum<'a>
2232 where
2233 I: IntoIterator<Item = (Datum<'a>, Diff)>,
2234 W: OneByOneAggr,
2235 {
2236 match self {
2237 AggregateFunc::WindowAggregate {
2238 wrapped_aggregate,
2239 order_by,
2240 window_frame,
2241 } => window_aggr::<_, W>(
2242 expand_counts(datums),
2243 temp_storage,
2244 wrapped_aggregate,
2245 order_by,
2246 window_frame,
2247 ),
2248 AggregateFunc::FusedWindowAggregate {
2249 wrapped_aggregates,
2250 order_by,
2251 window_frame,
2252 } => fused_window_aggr::<_, W>(
2253 expand_counts(datums),
2254 temp_storage,
2255 wrapped_aggregates,
2256 order_by,
2257 window_frame,
2258 ),
2259 _ => self.eval(datums, temp_storage),
2260 }
2261 }
2262
2263 pub fn eval_with_unnest_list<'a, I, W>(
2264 &self,
2265 datums: I,
2266 temp_storage: &'a RowArena,
2267 ) -> impl Iterator<Item = Datum<'a>>
2268 where
2269 I: IntoIterator<Item = (Datum<'a>, Diff)>,
2270 W: OneByOneAggr,
2271 {
2272 assert!(self.can_fuse_with_unnest_list());
2274 let datums = expand_counts(datums);
2276 match self {
2277 AggregateFunc::RowNumber { order_by } => {
2278 row_number_no_list(datums, temp_storage, order_by).collect_vec()
2279 }
2280 AggregateFunc::Rank { order_by } => {
2281 rank_no_list(datums, temp_storage, order_by).collect_vec()
2282 }
2283 AggregateFunc::DenseRank { order_by } => {
2284 dense_rank_no_list(datums, temp_storage, order_by).collect_vec()
2285 }
2286 AggregateFunc::LagLead {
2287 order_by,
2288 lag_lead: lag_lead_type,
2289 ignore_nulls,
2290 } => lag_lead_no_list(datums, temp_storage, order_by, lag_lead_type, ignore_nulls)
2291 .collect_vec(),
2292 AggregateFunc::FirstValue {
2293 order_by,
2294 window_frame,
2295 } => first_value_no_list(datums, temp_storage, order_by, window_frame).collect_vec(),
2296 AggregateFunc::LastValue {
2297 order_by,
2298 window_frame,
2299 } => last_value_no_list(datums, temp_storage, order_by, window_frame).collect_vec(),
2300 AggregateFunc::FusedValueWindowFunc { funcs, order_by } => {
2301 fused_value_window_func_no_list(datums, temp_storage, funcs, order_by).collect_vec()
2302 }
2303 AggregateFunc::WindowAggregate {
2304 wrapped_aggregate,
2305 order_by,
2306 window_frame,
2307 } => window_aggr_no_list::<_, W>(
2308 datums,
2309 temp_storage,
2310 wrapped_aggregate,
2311 order_by,
2312 window_frame,
2313 )
2314 .collect_vec(),
2315 AggregateFunc::FusedWindowAggregate {
2316 wrapped_aggregates,
2317 order_by,
2318 window_frame,
2319 } => fused_window_aggr_no_list::<_, W>(
2320 datums,
2321 temp_storage,
2322 wrapped_aggregates,
2323 order_by,
2324 window_frame,
2325 )
2326 .collect_vec(),
2327 _ => unreachable!("asserted above that `can_fuse_with_unnest_list`"),
2328 }
2329 .into_iter()
2330 }
2331
2332 pub fn default(&self) -> Datum<'static> {
2335 match self {
2336 AggregateFunc::Count => Datum::Int64(0),
2337 AggregateFunc::Any => Datum::False,
2338 AggregateFunc::All => Datum::True,
2339 AggregateFunc::Dummy => Datum::Dummy,
2340 _ => Datum::Null,
2341 }
2342 }
2343
2344 pub fn identity_datum(&self) -> Datum<'static> {
2347 match self {
2348 AggregateFunc::Any => Datum::False,
2349 AggregateFunc::All => Datum::True,
2350 AggregateFunc::Dummy => Datum::Dummy,
2351 AggregateFunc::ArrayConcat { .. } => Datum::empty_array(),
2352 AggregateFunc::ListConcat { .. } => Datum::empty_list(),
2353 AggregateFunc::RowNumber { .. }
2354 | AggregateFunc::Rank { .. }
2355 | AggregateFunc::DenseRank { .. }
2356 | AggregateFunc::LagLead { .. }
2357 | AggregateFunc::FirstValue { .. }
2358 | AggregateFunc::LastValue { .. }
2359 | AggregateFunc::WindowAggregate { .. }
2360 | AggregateFunc::FusedValueWindowFunc { .. }
2361 | AggregateFunc::FusedWindowAggregate { .. } => Datum::empty_list(),
2362 AggregateFunc::MaxNumeric
2363 | AggregateFunc::MaxInt16
2364 | AggregateFunc::MaxInt32
2365 | AggregateFunc::MaxInt64
2366 | AggregateFunc::MaxUInt16
2367 | AggregateFunc::MaxUInt32
2368 | AggregateFunc::MaxUInt64
2369 | AggregateFunc::MaxMzTimestamp
2370 | AggregateFunc::MaxFloat32
2371 | AggregateFunc::MaxFloat64
2372 | AggregateFunc::MaxBool
2373 | AggregateFunc::MaxString
2374 | AggregateFunc::MaxDate
2375 | AggregateFunc::MaxTimestamp
2376 | AggregateFunc::MaxTimestampTz
2377 | AggregateFunc::MaxInterval
2378 | AggregateFunc::MaxTime
2379 | AggregateFunc::MinNumeric
2380 | AggregateFunc::MinInt16
2381 | AggregateFunc::MinInt32
2382 | AggregateFunc::MinInt64
2383 | AggregateFunc::MinUInt16
2384 | AggregateFunc::MinUInt32
2385 | AggregateFunc::MinUInt64
2386 | AggregateFunc::MinMzTimestamp
2387 | AggregateFunc::MinFloat32
2388 | AggregateFunc::MinFloat64
2389 | AggregateFunc::MinBool
2390 | AggregateFunc::MinString
2391 | AggregateFunc::MinDate
2392 | AggregateFunc::MinTimestamp
2393 | AggregateFunc::MinTimestampTz
2394 | AggregateFunc::MinInterval
2395 | AggregateFunc::MinTime
2396 | AggregateFunc::SumInt16
2397 | AggregateFunc::SumInt32
2398 | AggregateFunc::SumInt64
2399 | AggregateFunc::SumUInt16
2400 | AggregateFunc::SumUInt32
2401 | AggregateFunc::SumUInt64
2402 | AggregateFunc::SumFloat32
2403 | AggregateFunc::SumFloat64
2404 | AggregateFunc::SumNumeric
2405 | AggregateFunc::Count
2406 | AggregateFunc::JsonbAgg { .. }
2407 | AggregateFunc::JsonbObjectAgg { .. }
2408 | AggregateFunc::MapAgg { .. }
2409 | AggregateFunc::StringAgg { .. } => Datum::Null,
2410 }
2411 }
2412
2413 pub fn can_fuse_with_unnest_list(&self) -> bool {
2414 match self {
2415 AggregateFunc::RowNumber { .. }
2416 | AggregateFunc::Rank { .. }
2417 | AggregateFunc::DenseRank { .. }
2418 | AggregateFunc::LagLead { .. }
2419 | AggregateFunc::FirstValue { .. }
2420 | AggregateFunc::LastValue { .. }
2421 | AggregateFunc::WindowAggregate { .. }
2422 | AggregateFunc::FusedValueWindowFunc { .. }
2423 | AggregateFunc::FusedWindowAggregate { .. } => true,
2424 AggregateFunc::ArrayConcat { .. }
2425 | AggregateFunc::ListConcat { .. }
2426 | AggregateFunc::Any
2427 | AggregateFunc::All
2428 | AggregateFunc::Dummy
2429 | AggregateFunc::MaxNumeric
2430 | AggregateFunc::MaxInt16
2431 | AggregateFunc::MaxInt32
2432 | AggregateFunc::MaxInt64
2433 | AggregateFunc::MaxUInt16
2434 | AggregateFunc::MaxUInt32
2435 | AggregateFunc::MaxUInt64
2436 | AggregateFunc::MaxMzTimestamp
2437 | AggregateFunc::MaxFloat32
2438 | AggregateFunc::MaxFloat64
2439 | AggregateFunc::MaxBool
2440 | AggregateFunc::MaxString
2441 | AggregateFunc::MaxDate
2442 | AggregateFunc::MaxTimestamp
2443 | AggregateFunc::MaxTimestampTz
2444 | AggregateFunc::MaxInterval
2445 | AggregateFunc::MaxTime
2446 | AggregateFunc::MinNumeric
2447 | AggregateFunc::MinInt16
2448 | AggregateFunc::MinInt32
2449 | AggregateFunc::MinInt64
2450 | AggregateFunc::MinUInt16
2451 | AggregateFunc::MinUInt32
2452 | AggregateFunc::MinUInt64
2453 | AggregateFunc::MinMzTimestamp
2454 | AggregateFunc::MinFloat32
2455 | AggregateFunc::MinFloat64
2456 | AggregateFunc::MinBool
2457 | AggregateFunc::MinString
2458 | AggregateFunc::MinDate
2459 | AggregateFunc::MinTimestamp
2460 | AggregateFunc::MinTimestampTz
2461 | AggregateFunc::MinInterval
2462 | AggregateFunc::MinTime
2463 | AggregateFunc::SumInt16
2464 | AggregateFunc::SumInt32
2465 | AggregateFunc::SumInt64
2466 | AggregateFunc::SumUInt16
2467 | AggregateFunc::SumUInt32
2468 | AggregateFunc::SumUInt64
2469 | AggregateFunc::SumFloat32
2470 | AggregateFunc::SumFloat64
2471 | AggregateFunc::SumNumeric
2472 | AggregateFunc::Count
2473 | AggregateFunc::JsonbAgg { .. }
2474 | AggregateFunc::JsonbObjectAgg { .. }
2475 | AggregateFunc::MapAgg { .. }
2476 | AggregateFunc::StringAgg { .. } => false,
2477 }
2478 }
2479
2480 pub fn output_sql_type(&self, input_type: SqlColumnType) -> SqlColumnType {
2486 let scalar_type = match self {
2487 AggregateFunc::Count => SqlScalarType::Int64,
2488 AggregateFunc::Any => SqlScalarType::Bool,
2489 AggregateFunc::All => SqlScalarType::Bool,
2490 AggregateFunc::JsonbAgg { .. } => SqlScalarType::Jsonb,
2491 AggregateFunc::JsonbObjectAgg { .. } => SqlScalarType::Jsonb,
2492 AggregateFunc::SumInt16 => SqlScalarType::Int64,
2493 AggregateFunc::SumInt32 => SqlScalarType::Int64,
2494 AggregateFunc::SumInt64 => SqlScalarType::Numeric {
2495 max_scale: Some(NumericMaxScale::ZERO),
2496 },
2497 AggregateFunc::SumUInt16 => SqlScalarType::UInt64,
2498 AggregateFunc::SumUInt32 => SqlScalarType::UInt64,
2499 AggregateFunc::SumUInt64 => SqlScalarType::Numeric {
2500 max_scale: Some(NumericMaxScale::ZERO),
2501 },
2502 AggregateFunc::MapAgg { value_type, .. } => SqlScalarType::Map {
2503 value_type: Box::new(value_type.clone()),
2504 custom_id: None,
2505 },
2506 AggregateFunc::ArrayConcat { .. } | AggregateFunc::ListConcat { .. } => {
2507 match input_type.scalar_type {
2508 SqlScalarType::Record { ref fields, .. } => fields[0].1.scalar_type.clone(),
2510 _ => unreachable!(),
2511 }
2512 }
2513 AggregateFunc::StringAgg { .. } => SqlScalarType::String,
2514 AggregateFunc::RowNumber { .. } => {
2515 AggregateFunc::output_type_ranking_window_funcs(&input_type, "?row_number?")
2516 }
2517 AggregateFunc::Rank { .. } => {
2518 AggregateFunc::output_type_ranking_window_funcs(&input_type, "?rank?")
2519 }
2520 AggregateFunc::DenseRank { .. } => {
2521 AggregateFunc::output_type_ranking_window_funcs(&input_type, "?dense_rank?")
2522 }
2523 AggregateFunc::LagLead { lag_lead: lag_lead_type, .. } => {
2524 let fields = input_type.scalar_type.unwrap_record_element_type();
2526 let original_row_type = fields[0].unwrap_record_element_type()[0]
2527 .clone()
2528 .nullable(false);
2529 let encoded_args = fields[0].unwrap_record_element_type()[1];
2530 let output_type_inner =
2531 Self::lag_lead_output_type_inner_from_encoded_args(encoded_args);
2532 let column_name = Self::lag_lead_result_column_name(lag_lead_type);
2533
2534 SqlScalarType::List {
2535 element_type: Box::new(SqlScalarType::Record {
2536 fields: [
2537 (column_name, output_type_inner),
2538 (ColumnName::from("?orig_row?"), original_row_type),
2539 ].into(),
2540 custom_id: None,
2541 }),
2542 custom_id: None,
2543 }
2544 }
2545 AggregateFunc::FirstValue { .. } => {
2546 let fields = input_type.scalar_type.unwrap_record_element_type();
2548 let original_row_type = fields[0].unwrap_record_element_type()[0]
2549 .clone()
2550 .nullable(false);
2551 let value_type = fields[0].unwrap_record_element_type()[1]
2552 .clone()
2553 .nullable(true); SqlScalarType::List {
2556 element_type: Box::new(SqlScalarType::Record {
2557 fields: [
2558 (ColumnName::from("?first_value?"), value_type),
2559 (ColumnName::from("?orig_row?"), original_row_type),
2560 ].into(),
2561 custom_id: None,
2562 }),
2563 custom_id: None,
2564 }
2565 }
2566 AggregateFunc::LastValue { .. } => {
2567 let fields = input_type.scalar_type.unwrap_record_element_type();
2569 let original_row_type = fields[0].unwrap_record_element_type()[0]
2570 .clone()
2571 .nullable(false);
2572 let value_type = fields[0].unwrap_record_element_type()[1]
2573 .clone()
2574 .nullable(true); SqlScalarType::List {
2577 element_type: Box::new(SqlScalarType::Record {
2578 fields: [
2579 (ColumnName::from("?last_value?"), value_type),
2580 (ColumnName::from("?orig_row?"), original_row_type),
2581 ].into(),
2582 custom_id: None,
2583 }),
2584 custom_id: None,
2585 }
2586 }
2587 AggregateFunc::WindowAggregate {
2588 wrapped_aggregate, ..
2589 } => {
2590 let fields = input_type.scalar_type.unwrap_record_element_type();
2592 let original_row_type = fields[0].unwrap_record_element_type()[0]
2593 .clone()
2594 .nullable(false);
2595 let arg_type = fields[0].unwrap_record_element_type()[1]
2596 .clone()
2597 .nullable(true);
2598 let wrapped_aggr_out_type = wrapped_aggregate.output_sql_type(arg_type);
2599
2600 SqlScalarType::List {
2601 element_type: Box::new(SqlScalarType::Record {
2602 fields: [
2603 (ColumnName::from("?window_agg?"), wrapped_aggr_out_type),
2604 (ColumnName::from("?orig_row?"), original_row_type),
2605 ].into(),
2606 custom_id: None,
2607 }),
2608 custom_id: None,
2609 }
2610 }
2611 AggregateFunc::FusedWindowAggregate {
2612 wrapped_aggregates, ..
2613 } => {
2614 let fields = input_type.scalar_type.unwrap_record_element_type();
2617 let original_row_type = fields[0].unwrap_record_element_type()[0]
2618 .clone()
2619 .nullable(false);
2620 let args_type = fields[0].unwrap_record_element_type()[1];
2621 let arg_types = args_type.unwrap_record_element_type();
2622 let out_fields = arg_types.iter().zip_eq(wrapped_aggregates).map(
2623 |(arg_type, wrapped_agg)| {
2624 (
2625 ColumnName::from(wrapped_agg.name()),
2626 wrapped_agg.output_sql_type((**arg_type).clone().nullable(true)),
2627 )
2628 }).collect_vec();
2629
2630 SqlScalarType::List {
2631 element_type: Box::new(SqlScalarType::Record {
2632 fields: [
2633 (ColumnName::from("?fused_window_agg?"), SqlScalarType::Record {
2634 fields: out_fields.into(),
2635 custom_id: None,
2636 }.nullable(false)),
2637 (ColumnName::from("?orig_row?"), original_row_type),
2638 ].into(),
2639 custom_id: None,
2640 }),
2641 custom_id: None,
2642 }
2643 }
2644 AggregateFunc::FusedValueWindowFunc { funcs, order_by: _ } => {
2645 let fields = input_type.scalar_type.unwrap_record_element_type();
2650 let original_row_type = fields[0].unwrap_record_element_type()[0]
2651 .clone()
2652 .nullable(false);
2653 let encoded_args_type = fields[0]
2654 .unwrap_record_element_type()[1]
2655 .unwrap_record_element_type();
2656
2657 SqlScalarType::List {
2658 element_type: Box::new(SqlScalarType::Record {
2659 fields: [
2660 (
2661 ColumnName::from("?fused_value_window_func?"),
2662 SqlScalarType::Record {
2663 fields: encoded_args_type.into_iter().zip_eq(funcs).map(
2664 |(arg_type, func)| {
2665 match func {
2666 AggregateFunc::LagLead {
2667 lag_lead: lag_lead_type, ..
2668 } => {
2669 let name = Self::lag_lead_result_column_name(
2670 lag_lead_type,
2671 );
2672 let ty = Self
2673 ::lag_lead_output_type_inner_from_encoded_args(
2674 arg_type,
2675 );
2676 (name, ty)
2677 },
2678 AggregateFunc::FirstValue { .. } => {
2679 (
2680 ColumnName::from("?first_value?"),
2681 arg_type.clone().nullable(true),
2682 )
2683 }
2684 AggregateFunc::LastValue { .. } => {
2685 (
2686 ColumnName::from("?last_value?"),
2687 arg_type.clone().nullable(true),
2688 )
2689 }
2690 _ => panic!("FusedValueWindowFunc has an unknown function"),
2691 }
2692 }).collect(),
2693 custom_id: None,
2694 }.nullable(false)),
2695 (ColumnName::from("?orig_row?"), original_row_type),
2696 ].into(),
2697 custom_id: None,
2698 }),
2699 custom_id: None,
2700 }
2701 }
2702 AggregateFunc::Dummy
2703 | AggregateFunc::MaxNumeric
2704 | AggregateFunc::MaxInt16
2705 | AggregateFunc::MaxInt32
2706 | AggregateFunc::MaxInt64
2707 | AggregateFunc::MaxUInt16
2708 | AggregateFunc::MaxUInt32
2709 | AggregateFunc::MaxUInt64
2710 | AggregateFunc::MaxMzTimestamp
2711 | AggregateFunc::MaxFloat32
2712 | AggregateFunc::MaxFloat64
2713 | AggregateFunc::MaxBool
2714 | AggregateFunc::MaxString
2718 | AggregateFunc::MaxDate
2719 | AggregateFunc::MaxTimestamp
2720 | AggregateFunc::MaxTimestampTz
2721 | AggregateFunc::MaxInterval
2722 | AggregateFunc::MaxTime
2723 | AggregateFunc::MinNumeric
2724 | AggregateFunc::MinInt16
2725 | AggregateFunc::MinInt32
2726 | AggregateFunc::MinInt64
2727 | AggregateFunc::MinUInt16
2728 | AggregateFunc::MinUInt32
2729 | AggregateFunc::MinUInt64
2730 | AggregateFunc::MinMzTimestamp
2731 | AggregateFunc::MinFloat32
2732 | AggregateFunc::MinFloat64
2733 | AggregateFunc::MinBool
2734 | AggregateFunc::MinString
2735 | AggregateFunc::MinDate
2736 | AggregateFunc::MinTimestamp
2737 | AggregateFunc::MinTimestampTz
2738 | AggregateFunc::MinInterval
2739 | AggregateFunc::MinTime
2740 | AggregateFunc::SumFloat32
2741 | AggregateFunc::SumFloat64
2742 | AggregateFunc::SumNumeric => input_type.scalar_type.clone(),
2743 };
2744 let nullable = match self {
2747 AggregateFunc::Count => false,
2748 AggregateFunc::StringAgg { .. } => match input_type.scalar_type {
2750 SqlScalarType::Record { fields, .. } => match &fields[0].1.scalar_type {
2752 SqlScalarType::Record { fields, .. } => fields[0].1.nullable,
2754 _ => unreachable!(),
2755 },
2756 _ => unreachable!(),
2757 },
2758 _ => input_type.nullable,
2759 };
2760 scalar_type.nullable(nullable)
2761 }
2762
2763 pub fn output_type(&self, input_type: ReprColumnType) -> ReprColumnType {
2767 ReprColumnType::from(&self.output_sql_type(SqlColumnType::from_repr(&input_type)))
2768 }
2769
2770 fn output_type_ranking_window_funcs(
2772 input_type: &SqlColumnType,
2773 col_name: &str,
2774 ) -> SqlScalarType {
2775 match input_type.scalar_type {
2776 SqlScalarType::Record { ref fields, .. } => SqlScalarType::List {
2777 element_type: Box::new(SqlScalarType::Record {
2778 fields: [
2779 (
2780 ColumnName::from(col_name),
2781 SqlScalarType::Int64.nullable(false),
2782 ),
2783 (ColumnName::from("?orig_row?"), {
2784 let inner = match &fields[0].1.scalar_type {
2785 SqlScalarType::List { element_type, .. } => element_type.clone(),
2786 _ => unreachable!(),
2787 };
2788 inner.nullable(false)
2789 }),
2790 ]
2791 .into(),
2792 custom_id: None,
2793 }),
2794 custom_id: None,
2795 },
2796 _ => unreachable!(),
2797 }
2798 }
2799
2800 fn lag_lead_output_type_inner_from_encoded_args(
2804 encoded_args_type: &SqlScalarType,
2805 ) -> SqlColumnType {
2806 encoded_args_type.unwrap_record_element_type()[0]
2810 .clone()
2811 .nullable(true)
2812 }
2813
2814 fn lag_lead_result_column_name(lag_lead_type: &LagLeadType) -> ColumnName {
2815 ColumnName::from(match lag_lead_type {
2816 LagLeadType::Lag => "?lag?",
2817 LagLeadType::Lead => "?lead?",
2818 })
2819 }
2820
2821 pub fn propagates_nonnull_constraint(&self) -> bool {
2826 match self {
2827 AggregateFunc::MaxNumeric
2828 | AggregateFunc::MaxInt16
2829 | AggregateFunc::MaxInt32
2830 | AggregateFunc::MaxInt64
2831 | AggregateFunc::MaxUInt16
2832 | AggregateFunc::MaxUInt32
2833 | AggregateFunc::MaxUInt64
2834 | AggregateFunc::MaxMzTimestamp
2835 | AggregateFunc::MaxFloat32
2836 | AggregateFunc::MaxFloat64
2837 | AggregateFunc::MaxBool
2838 | AggregateFunc::MaxString
2839 | AggregateFunc::MaxDate
2840 | AggregateFunc::MaxTimestamp
2841 | AggregateFunc::MaxTimestampTz
2842 | AggregateFunc::MaxInterval
2843 | AggregateFunc::MaxTime
2844 | AggregateFunc::MinNumeric
2845 | AggregateFunc::MinInt16
2846 | AggregateFunc::MinInt32
2847 | AggregateFunc::MinInt64
2848 | AggregateFunc::MinUInt16
2849 | AggregateFunc::MinUInt32
2850 | AggregateFunc::MinUInt64
2851 | AggregateFunc::MinMzTimestamp
2852 | AggregateFunc::MinFloat32
2853 | AggregateFunc::MinFloat64
2854 | AggregateFunc::MinBool
2855 | AggregateFunc::MinString
2856 | AggregateFunc::MinDate
2857 | AggregateFunc::MinTimestamp
2858 | AggregateFunc::MinTimestampTz
2859 | AggregateFunc::MinInterval
2860 | AggregateFunc::MinTime
2861 | AggregateFunc::SumInt16
2862 | AggregateFunc::SumInt32
2863 | AggregateFunc::SumInt64
2864 | AggregateFunc::SumUInt16
2865 | AggregateFunc::SumUInt32
2866 | AggregateFunc::SumUInt64
2867 | AggregateFunc::SumFloat32
2868 | AggregateFunc::SumFloat64
2869 | AggregateFunc::SumNumeric
2870 | AggregateFunc::StringAgg { .. } => true,
2871 AggregateFunc::Count
2873 | AggregateFunc::Any
2874 | AggregateFunc::All
2875 | AggregateFunc::JsonbAgg { .. }
2876 | AggregateFunc::JsonbObjectAgg { .. }
2877 | AggregateFunc::MapAgg { .. }
2878 | AggregateFunc::ArrayConcat { .. }
2879 | AggregateFunc::ListConcat { .. }
2880 | AggregateFunc::RowNumber { .. }
2881 | AggregateFunc::Rank { .. }
2882 | AggregateFunc::DenseRank { .. }
2883 | AggregateFunc::LagLead { .. }
2884 | AggregateFunc::FirstValue { .. }
2885 | AggregateFunc::LastValue { .. }
2886 | AggregateFunc::FusedValueWindowFunc { .. }
2887 | AggregateFunc::WindowAggregate { .. }
2888 | AggregateFunc::FusedWindowAggregate { .. }
2889 | AggregateFunc::Dummy => false,
2890 }
2891 }
2892}
2893
2894fn jsonb_each<'a>(a: Datum<'a>) -> impl Iterator<Item = (Row, Diff)> + 'a {
2895 let map = match a {
2897 Datum::Map(dict) => dict,
2898 _ => mz_repr::DatumMap::empty(),
2899 };
2900
2901 map.iter()
2902 .map(move |(k, v)| (Row::pack_slice(&[Datum::String(k), v]), Diff::ONE))
2903}
2904
2905fn jsonb_each_stringify<'a>(
2906 a: Datum<'a>,
2907 temp_storage: &'a RowArena,
2908) -> impl Iterator<Item = (Row, Diff)> + 'a {
2909 let map = match a {
2911 Datum::Map(dict) => dict,
2912 _ => mz_repr::DatumMap::empty(),
2913 };
2914
2915 map.iter().map(move |(k, mut v)| {
2916 v = jsonb_stringify(v, temp_storage)
2917 .map(Datum::String)
2918 .unwrap_or(Datum::Null);
2919 (Row::pack_slice(&[Datum::String(k), v]), Diff::ONE)
2920 })
2921}
2922
2923fn jsonb_object_keys<'a>(a: Datum<'a>) -> impl Iterator<Item = (Row, Diff)> + 'a {
2924 let map = match a {
2925 Datum::Map(dict) => dict,
2926 _ => mz_repr::DatumMap::empty(),
2927 };
2928
2929 map.iter()
2930 .map(move |(k, _)| (Row::pack_slice(&[Datum::String(k)]), Diff::ONE))
2931}
2932
2933fn jsonb_array_elements<'a>(a: Datum<'a>) -> impl Iterator<Item = (Row, Diff)> + 'a {
2934 let list = match a {
2935 Datum::List(list) => list,
2936 _ => mz_repr::DatumList::empty(),
2937 };
2938 list.iter().map(move |e| (Row::pack_slice(&[e]), Diff::ONE))
2939}
2940
2941fn jsonb_array_elements_stringify<'a>(
2942 a: Datum<'a>,
2943 temp_storage: &'a RowArena,
2944) -> impl Iterator<Item = (Row, Diff)> + 'a {
2945 let list = match a {
2946 Datum::List(list) => list,
2947 _ => mz_repr::DatumList::empty(),
2948 };
2949 list.iter().map(move |mut e| {
2950 e = jsonb_stringify(e, temp_storage)
2951 .map(Datum::String)
2952 .unwrap_or(Datum::Null);
2953 (Row::pack_slice(&[e]), Diff::ONE)
2954 })
2955}
2956
2957fn regexp_extract(a: Datum, r: &AnalyzedRegex) -> Option<(Row, Diff)> {
2958 let r = r.inner();
2959 let a = a.unwrap_str();
2960 let captures = r.captures(a)?;
2961 let datums = captures
2962 .iter()
2963 .skip(1)
2964 .map(|m| Datum::from(m.map(|m| m.as_str())));
2965 Some((Row::pack(datums), Diff::ONE))
2966}
2967
2968fn regexp_matches<'a>(
2969 exprs: &[Datum<'a>],
2970) -> Result<impl Iterator<Item = (Row, Diff)> + 'a, EvalError> {
2971 assert!(exprs.len() == 2 || exprs.len() == 3);
2975 let a = exprs[0].unwrap_str();
2976 let r = exprs[1].unwrap_str();
2977
2978 let (regex, opts) = if exprs.len() == 3 {
2979 let flag = exprs[2].unwrap_str();
2980 let opts = AnalyzedRegexOpts::from_str(flag)?;
2981 (AnalyzedRegex::new(r, opts)?, opts)
2982 } else {
2983 let opts = AnalyzedRegexOpts::default();
2984 (AnalyzedRegex::new(r, opts)?, opts)
2985 };
2986
2987 let regex = regex.inner().clone();
2988
2989 let iter = regex.captures_iter(a).map(move |captures| {
2990 let matches = captures
2991 .iter()
2992 .skip(1)
2994 .map(|m| Datum::from(m.map(|m| m.as_str())))
2995 .collect::<Vec<_>>();
2996
2997 let mut binding = SharedRow::get();
2998 let mut packer = binding.packer();
2999
3000 let dimension = ArrayDimension {
3001 lower_bound: 1,
3002 length: matches.len(),
3003 };
3004 packer
3005 .try_push_array(&[dimension], matches)
3006 .expect("generated dimensions above");
3007
3008 (binding.clone(), Diff::ONE)
3009 });
3010
3011 let out = iter.collect::<SmallVec<[_; 3]>>();
3016
3017 if opts.global {
3018 Ok(Either::Left(out.into_iter()))
3019 } else {
3020 Ok(Either::Right(out.into_iter().take(1)))
3021 }
3022}
3023
3024fn generate_series<N>(
3025 start: N,
3026 stop: N,
3027 step: N,
3028) -> Result<impl Iterator<Item = (Row, Diff)>, EvalError>
3029where
3030 N: Integer + Signed + CheckedAdd + Clone,
3031 Datum<'static>: From<N>,
3032{
3033 if step == N::zero() {
3034 return Err(EvalError::InvalidParameterValue(
3035 "step size cannot equal zero".into(),
3036 ));
3037 }
3038 Ok(num::range_step_inclusive(start, stop, step)
3039 .map(move |i| (Row::pack_slice(&[Datum::from(i)]), Diff::ONE)))
3040}
3041
3042#[derive(Clone)]
3046pub struct TimestampRangeStepInclusive<T> {
3047 state: CheckedTimestamp<T>,
3048 stop: CheckedTimestamp<T>,
3049 step: Interval,
3050 rev: bool,
3051 done: bool,
3052}
3053
3054impl<T: TimestampLike> Iterator for TimestampRangeStepInclusive<T> {
3055 type Item = CheckedTimestamp<T>;
3056
3057 #[inline]
3058 fn next(&mut self) -> Option<CheckedTimestamp<T>> {
3059 if !self.done
3060 && ((self.rev && self.state >= self.stop) || (!self.rev && self.state <= self.stop))
3061 {
3062 let result = self.state.clone();
3063 match add_timestamp_months(self.state.deref(), self.step.months) {
3064 Ok(state) => match state.checked_add_signed(self.step.duration_as_chrono()) {
3065 Some(v) => match CheckedTimestamp::from_timestamplike(v) {
3066 Ok(v) => {
3067 let progressed = if self.rev {
3072 v < self.state
3073 } else {
3074 v > self.state
3075 };
3076 if progressed {
3077 self.state = v
3078 } else {
3079 self.done = true
3080 }
3081 }
3082 Err(_) => self.done = true,
3083 },
3084 None => self.done = true,
3085 },
3086 Err(..) => {
3087 self.done = true;
3088 }
3089 }
3090
3091 Some(result)
3092 } else {
3093 None
3094 }
3095 }
3096}
3097
3098fn generate_series_ts<T: TimestampLike>(
3099 start: CheckedTimestamp<T>,
3100 stop: CheckedTimestamp<T>,
3101 step: Interval,
3102 conv: fn(CheckedTimestamp<T>) -> Datum<'static>,
3103) -> Result<impl Iterator<Item = (Row, Diff)>, EvalError> {
3104 let normalized_step = step.as_microseconds();
3105 if normalized_step == 0 {
3106 return Err(EvalError::InvalidParameterValue(
3107 "step size cannot equal zero".into(),
3108 ));
3109 }
3110 let rev = normalized_step < 0;
3111
3112 let trsi = TimestampRangeStepInclusive {
3113 state: start,
3114 stop,
3115 step,
3116 rev,
3117 done: false,
3118 };
3119
3120 Ok(trsi.map(move |i| (Row::pack_slice(&[conv(i)]), Diff::ONE)))
3121}
3122
3123fn generate_subscripts_array(
3124 a: Datum,
3125 dim: i32,
3126) -> Result<Box<dyn Iterator<Item = (Row, Diff)>>, EvalError> {
3127 if dim <= 0 {
3128 return Ok(Box::new(iter::empty()));
3129 }
3130
3131 match a.unwrap_array().dims().into_iter().nth(
3132 (dim - 1)
3133 .try_into()
3134 .map_err(|_| EvalError::Int32OutOfRange((dim - 1).to_string().into()))?,
3135 ) {
3136 Some(requested_dim) => {
3137 let lower_bound: i32 = requested_dim.lower_bound.try_into().map_err(|_| {
3138 EvalError::Int32OutOfRange(requested_dim.lower_bound.to_string().into())
3139 })?;
3140 let length: i32 = requested_dim
3143 .length
3144 .try_into()
3145 .map_err(|_| EvalError::Int32OutOfRange(requested_dim.length.to_string().into()))?;
3146 let upper_bound = lower_bound.checked_add(length - 1).ok_or_else(|| {
3147 EvalError::Int32OutOfRange(requested_dim.length.to_string().into())
3148 })?;
3149 Ok(Box::new(generate_series::<i32>(
3150 lower_bound,
3151 upper_bound,
3152 1,
3153 )?))
3154 }
3155 None => Ok(Box::new(iter::empty())),
3156 }
3157}
3158
3159fn unnest_array<'a>(a: Datum<'a>) -> impl Iterator<Item = (Row, Diff)> + 'a {
3160 a.unwrap_array()
3161 .elements()
3162 .iter()
3163 .map(move |e| (Row::pack_slice(&[e]), Diff::ONE))
3164}
3165
3166fn unnest_list<'a>(a: Datum<'a>) -> impl Iterator<Item = (Row, Diff)> + 'a {
3167 a.unwrap_list()
3168 .iter()
3169 .map(move |e| (Row::pack_slice(&[e]), Diff::ONE))
3170}
3171
3172fn unnest_map<'a>(a: Datum<'a>) -> impl Iterator<Item = (Row, Diff)> + 'a {
3173 a.unwrap_map()
3174 .iter()
3175 .map(move |(k, v)| (Row::pack_slice(&[Datum::from(k), v]), Diff::ONE))
3176}
3177
3178impl AggregateFunc {
3179 pub fn name(&self) -> &'static str {
3182 match self {
3183 Self::MaxNumeric => "max",
3184 Self::MaxInt16 => "max",
3185 Self::MaxInt32 => "max",
3186 Self::MaxInt64 => "max",
3187 Self::MaxUInt16 => "max",
3188 Self::MaxUInt32 => "max",
3189 Self::MaxUInt64 => "max",
3190 Self::MaxMzTimestamp => "max",
3191 Self::MaxFloat32 => "max",
3192 Self::MaxFloat64 => "max",
3193 Self::MaxBool => "max",
3194 Self::MaxString => "max",
3195 Self::MaxDate => "max",
3196 Self::MaxTimestamp => "max",
3197 Self::MaxTimestampTz => "max",
3198 Self::MaxInterval => "max",
3199 Self::MaxTime => "max",
3200 Self::MinNumeric => "min",
3201 Self::MinInt16 => "min",
3202 Self::MinInt32 => "min",
3203 Self::MinInt64 => "min",
3204 Self::MinUInt16 => "min",
3205 Self::MinUInt32 => "min",
3206 Self::MinUInt64 => "min",
3207 Self::MinMzTimestamp => "min",
3208 Self::MinFloat32 => "min",
3209 Self::MinFloat64 => "min",
3210 Self::MinBool => "min",
3211 Self::MinString => "min",
3212 Self::MinDate => "min",
3213 Self::MinTimestamp => "min",
3214 Self::MinTimestampTz => "min",
3215 Self::MinInterval => "min",
3216 Self::MinTime => "min",
3217 Self::SumInt16 => "sum",
3218 Self::SumInt32 => "sum",
3219 Self::SumInt64 => "sum",
3220 Self::SumUInt16 => "sum",
3221 Self::SumUInt32 => "sum",
3222 Self::SumUInt64 => "sum",
3223 Self::SumFloat32 => "sum",
3224 Self::SumFloat64 => "sum",
3225 Self::SumNumeric => "sum",
3226 Self::Count => "count",
3227 Self::Any => "any",
3228 Self::All => "all",
3229 Self::JsonbAgg { .. } => "jsonb_agg",
3230 Self::JsonbObjectAgg { .. } => "jsonb_object_agg",
3231 Self::MapAgg { .. } => "map_agg",
3232 Self::ArrayConcat { .. } => "array_agg",
3233 Self::ListConcat { .. } => "list_agg",
3234 Self::StringAgg { .. } => "string_agg",
3235 Self::RowNumber { .. } => "row_number",
3236 Self::Rank { .. } => "rank",
3237 Self::DenseRank { .. } => "dense_rank",
3238 Self::LagLead {
3239 lag_lead: LagLeadType::Lag,
3240 ..
3241 } => "lag",
3242 Self::LagLead {
3243 lag_lead: LagLeadType::Lead,
3244 ..
3245 } => "lead",
3246 Self::FirstValue { .. } => "first_value",
3247 Self::LastValue { .. } => "last_value",
3248 Self::WindowAggregate { .. } => "window_agg",
3249 Self::FusedValueWindowFunc { .. } => "fused_value_window_func",
3250 Self::FusedWindowAggregate { .. } => "fused_window_agg",
3251 Self::Dummy => "dummy",
3252 }
3253 }
3254}
3255
3256impl<'a, M> fmt::Display for HumanizedExpr<'a, AggregateFunc, M>
3257where
3258 M: HumanizerMode,
3259{
3260 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
3261 use AggregateFunc::*;
3262 let name = self.expr.name();
3263 match self.expr {
3264 JsonbAgg { order_by }
3265 | JsonbObjectAgg { order_by }
3266 | MapAgg { order_by, .. }
3267 | ArrayConcat { order_by }
3268 | ListConcat { order_by }
3269 | StringAgg { order_by }
3270 | RowNumber { order_by }
3271 | Rank { order_by }
3272 | DenseRank { order_by } => {
3273 let order_by = order_by.iter().map(|col| self.child(col));
3274 write!(f, "{}[order_by=[{}]]", name, separated(", ", order_by))
3275 }
3276 LagLead {
3277 lag_lead: _,
3278 ignore_nulls,
3279 order_by,
3280 } => {
3281 let order_by = order_by.iter().map(|col| self.child(col));
3282 f.write_str(name)?;
3283 f.write_str("[")?;
3284 if *ignore_nulls {
3285 f.write_str("ignore_nulls=true, ")?;
3286 }
3287 write!(f, "order_by=[{}]", separated(", ", order_by))?;
3288 f.write_str("]")
3289 }
3290 FirstValue {
3291 order_by,
3292 window_frame,
3293 } => {
3294 let order_by = order_by.iter().map(|col| self.child(col));
3295 f.write_str(name)?;
3296 f.write_str("[")?;
3297 write!(f, "order_by=[{}]", separated(", ", order_by))?;
3298 if *window_frame != WindowFrame::default() {
3299 write!(f, " {}", window_frame)?;
3300 }
3301 f.write_str("]")
3302 }
3303 LastValue {
3304 order_by,
3305 window_frame,
3306 } => {
3307 let order_by = order_by.iter().map(|col| self.child(col));
3308 f.write_str(name)?;
3309 f.write_str("[")?;
3310 write!(f, "order_by=[{}]", separated(", ", order_by))?;
3311 if *window_frame != WindowFrame::default() {
3312 write!(f, " {}", window_frame)?;
3313 }
3314 f.write_str("]")
3315 }
3316 WindowAggregate {
3317 wrapped_aggregate,
3318 order_by,
3319 window_frame,
3320 } => {
3321 let order_by = order_by.iter().map(|col| self.child(col));
3322 let wrapped_aggregate = self.child(wrapped_aggregate.deref());
3323 f.write_str(name)?;
3324 f.write_str("[")?;
3325 write!(f, "{} ", wrapped_aggregate)?;
3326 write!(f, "order_by=[{}]", separated(", ", order_by))?;
3327 if *window_frame != WindowFrame::default() {
3328 write!(f, " {}", window_frame)?;
3329 }
3330 f.write_str("]")
3331 }
3332 FusedValueWindowFunc { funcs, order_by } => {
3333 let order_by = order_by.iter().map(|col| self.child(col));
3334 let funcs = separated(", ", funcs.iter().map(|func| self.child(func)));
3335 f.write_str(name)?;
3336 f.write_str("[")?;
3337 write!(f, "{} ", funcs)?;
3338 write!(f, "order_by=[{}]", separated(", ", order_by))?;
3339 f.write_str("]")
3340 }
3341 _ => f.write_str(name),
3342 }
3343 }
3344}
3345
3346#[derive(
3347 Clone,
3348 Debug,
3349 Eq,
3350 PartialEq,
3351 Ord,
3352 PartialOrd,
3353 Serialize,
3354 Deserialize,
3355 Hash
3356)]
3357pub struct CaptureGroupDesc {
3358 pub index: u32,
3359 pub name: Option<String>,
3360 pub nullable: bool,
3361}
3362
3363#[derive(
3364 Clone,
3365 Copy,
3366 Debug,
3367 Eq,
3368 PartialEq,
3369 Ord,
3370 PartialOrd,
3371 Serialize,
3372 Deserialize,
3373 Hash,
3374 Default
3375)]
3376pub struct AnalyzedRegexOpts {
3377 pub case_insensitive: bool,
3378 pub global: bool,
3379}
3380
3381impl FromStr for AnalyzedRegexOpts {
3382 type Err = EvalError;
3383
3384 fn from_str(s: &str) -> Result<Self, Self::Err> {
3385 let mut opts = AnalyzedRegexOpts::default();
3386 for c in s.chars() {
3387 match c {
3388 'i' => opts.case_insensitive = true,
3389 'g' => opts.global = true,
3390 _ => return Err(EvalError::InvalidRegexFlag(c)),
3391 }
3392 }
3393 Ok(opts)
3394 }
3395}
3396
3397#[derive(
3398 Clone,
3399 Debug,
3400 Eq,
3401 PartialEq,
3402 Ord,
3403 PartialOrd,
3404 Serialize,
3405 Deserialize,
3406 Hash
3407)]
3408pub struct AnalyzedRegex(ReprRegex, Vec<CaptureGroupDesc>, AnalyzedRegexOpts);
3409
3410impl AnalyzedRegex {
3411 pub fn new(s: &str, opts: AnalyzedRegexOpts) -> Result<Self, RegexCompilationError> {
3412 let r = ReprRegex::new(s, opts.case_insensitive)?;
3413 #[allow(clippy::as_conversions)]
3415 let descs: Vec<_> = r
3416 .capture_names()
3417 .enumerate()
3418 .skip(1)
3423 .map(|(i, name)| CaptureGroupDesc {
3424 index: i as u32,
3425 name: name.map(String::from),
3426 nullable: true,
3429 })
3430 .collect();
3431 Ok(Self(r, descs, opts))
3432 }
3433 pub fn capture_groups_len(&self) -> usize {
3434 self.1.len()
3435 }
3436 pub fn capture_groups_iter(&self) -> impl Iterator<Item = &CaptureGroupDesc> {
3437 self.1.iter()
3438 }
3439 pub fn inner(&self) -> &Regex {
3440 &(self.0).regex
3441 }
3442 pub fn opts(&self) -> &AnalyzedRegexOpts {
3443 &self.2
3444 }
3445}
3446
3447pub fn csv_extract(a: Datum<'_>, n_cols: usize) -> impl Iterator<Item = (Row, Diff)> + '_ {
3448 let bytes = a.unwrap_str().as_bytes();
3449 let mut row = Row::default();
3450 let csv_reader = csv::ReaderBuilder::new()
3451 .has_headers(false)
3452 .from_reader(bytes);
3453 csv_reader.into_records().filter_map(move |res| match res {
3454 Ok(sr) if sr.len() == n_cols => {
3455 row.packer().extend(sr.iter().map(Datum::String));
3456 Some((row.clone(), Diff::ONE))
3457 }
3458 _ => None,
3459 })
3460}
3461
3462pub fn repeat_row(a: Datum) -> Option<(Row, Diff)> {
3463 let n = a.unwrap_int64();
3464 if n != 0 {
3465 Some((Row::default(), n.into()))
3466 } else {
3467 None
3468 }
3469}
3470
3471pub fn repeat_row_non_negative<'a>(
3472 a: Datum,
3473) -> Result<Box<dyn Iterator<Item = (Row, Diff)> + 'a>, EvalError> {
3474 let n = a.unwrap_int64();
3475 if n < 0 {
3476 Err(EvalError::InvalidParameterValue(
3477 format!("repeat_row_non_negative got {}", n).into(),
3478 ))
3479 } else if n == 0 {
3480 Ok(Box::new(iter::empty()))
3481 } else {
3482 Ok(Box::new(iter::once((Row::default(), n.into()))))
3484 }
3485}
3486
3487fn wrap<'a>(datums: &'a [Datum<'a>], width: usize) -> impl Iterator<Item = (Row, Diff)> + 'a {
3488 datums
3489 .chunks(width)
3490 .map(|chunk| (Row::pack(chunk), Diff::ONE))
3491}
3492
3493fn acl_explode<'a>(
3494 acl_items: Datum<'a>,
3495 temp_storage: &'a RowArena,
3496) -> Result<impl Iterator<Item = (Row, Diff)> + 'a, EvalError> {
3497 let acl_items = acl_items.unwrap_array();
3498 let mut res = Vec::new();
3499 for acl_item in acl_items.elements().iter() {
3500 if acl_item.is_null() {
3501 return Err(EvalError::AclArrayNullElement);
3502 }
3503 let acl_item = acl_item.unwrap_acl_item();
3504 for privilege in acl_item.acl_mode.explode() {
3505 let row = [
3506 Datum::UInt32(acl_item.grantor.0),
3507 Datum::UInt32(acl_item.grantee.0),
3508 Datum::String(temp_storage.push_string(privilege.to_string())),
3509 Datum::False,
3511 ];
3512 res.push((Row::pack_slice(&row), Diff::ONE));
3513 }
3514 }
3515 Ok(res.into_iter())
3516}
3517
3518fn mz_acl_explode<'a>(
3519 mz_acl_items: Datum<'a>,
3520 temp_storage: &'a RowArena,
3521) -> Result<impl Iterator<Item = (Row, Diff)> + 'a, EvalError> {
3522 let mz_acl_items = mz_acl_items.unwrap_array();
3523 let mut res = Vec::new();
3524 for mz_acl_item in mz_acl_items.elements().iter() {
3525 if mz_acl_item.is_null() {
3526 return Err(EvalError::MzAclArrayNullElement);
3527 }
3528 let mz_acl_item = mz_acl_item.unwrap_mz_acl_item();
3529 for privilege in mz_acl_item.acl_mode.explode() {
3530 let row = [
3531 Datum::String(temp_storage.push_string(mz_acl_item.grantor.to_string())),
3532 Datum::String(temp_storage.push_string(mz_acl_item.grantee.to_string())),
3533 Datum::String(temp_storage.push_string(privilege.to_string())),
3534 Datum::False,
3536 ];
3537 res.push((Row::pack_slice(&row), Diff::ONE));
3538 }
3539 }
3540 Ok(res.into_iter())
3541}
3542
3543#[derive(
3546 Clone,
3547 Debug,
3548 Eq,
3549 PartialEq,
3550 Ord,
3551 PartialOrd,
3552 Serialize,
3553 Deserialize,
3554 Hash
3555)]
3556pub enum TableFunc {
3557 AclExplode,
3558 MzAclExplode,
3559 JsonbEach,
3560 JsonbEachStringify,
3561 JsonbObjectKeys,
3562 JsonbArrayElements,
3563 JsonbArrayElementsStringify,
3564 RegexpExtract(AnalyzedRegex),
3565 CsvExtract(usize),
3566 GenerateSeriesInt32,
3567 GenerateSeriesInt64,
3568 GenerateSeriesUnoptimized,
3580 GenerateSeriesTimestamp,
3581 GenerateSeriesTimestampTz,
3582 GuardSubquerySize {
3603 column_type: SqlScalarType,
3604 },
3605 RepeatRow,
3612 RepeatRowNonNegative,
3615 UnnestArray {
3616 el_typ: SqlScalarType,
3617 },
3618 UnnestList {
3619 el_typ: SqlScalarType,
3620 },
3621 UnnestMap {
3622 value_type: SqlScalarType,
3623 },
3624 Wrap {
3630 types: Vec<SqlColumnType>,
3631 width: usize,
3632 },
3633 GenerateSubscriptsArray,
3634 TabletizedScalar {
3636 name: String,
3637 relation: SqlRelationType,
3638 },
3639 RegexpMatches,
3640 #[allow(private_interfaces)]
3645 WithOrdinality(WithOrdinality),
3646}
3647
3648#[derive(
3656 Clone,
3657 Debug,
3658 Eq,
3659 PartialEq,
3660 Ord,
3661 PartialOrd,
3662 Serialize,
3663 Deserialize,
3664 Hash
3665)]
3666struct WithOrdinality {
3667 inner: Box<TableFunc>,
3668}
3669
3670impl TableFunc {
3671 pub fn with_ordinality(inner: TableFunc) -> Option<TableFunc> {
3673 match inner {
3674 TableFunc::AclExplode
3675 | TableFunc::MzAclExplode
3676 | TableFunc::JsonbEach
3677 | TableFunc::JsonbEachStringify
3678 | TableFunc::JsonbObjectKeys
3679 | TableFunc::JsonbArrayElements
3680 | TableFunc::JsonbArrayElementsStringify
3681 | TableFunc::RegexpExtract(_)
3682 | TableFunc::CsvExtract(_)
3683 | TableFunc::GenerateSeriesInt32
3684 | TableFunc::GenerateSeriesInt64
3685 | TableFunc::GenerateSeriesUnoptimized
3686 | TableFunc::GenerateSeriesTimestamp
3687 | TableFunc::GenerateSeriesTimestampTz
3688 | TableFunc::GuardSubquerySize { .. }
3689 | TableFunc::RepeatRowNonNegative
3690 | TableFunc::UnnestArray { .. }
3691 | TableFunc::UnnestList { .. }
3692 | TableFunc::UnnestMap { .. }
3693 | TableFunc::Wrap { .. }
3694 | TableFunc::GenerateSubscriptsArray
3695 | TableFunc::TabletizedScalar { .. }
3696 | TableFunc::RegexpMatches => Some(TableFunc::WithOrdinality(WithOrdinality {
3697 inner: Box::new(inner),
3698 })),
3699 TableFunc::RepeatRow | TableFunc::WithOrdinality(_) => None, }
3707 }
3708}
3709
3710impl TableFunc {
3711 pub fn eval<'a>(
3713 &'a self,
3714 datums: &'a [Datum<'a>],
3715 temp_storage: &'a RowArena,
3716 ) -> Result<Box<dyn Iterator<Item = (Row, Diff)> + 'a>, EvalError> {
3717 if self.empty_on_null_input() && datums.iter().any(|d| d.is_null()) {
3718 return Ok(Box::new(vec![].into_iter()));
3719 }
3720 match self {
3721 TableFunc::AclExplode => Ok(Box::new(acl_explode(datums[0], temp_storage)?)),
3722 TableFunc::MzAclExplode => Ok(Box::new(mz_acl_explode(datums[0], temp_storage)?)),
3723 TableFunc::JsonbEach => Ok(Box::new(jsonb_each(datums[0]))),
3724 TableFunc::JsonbEachStringify => {
3725 Ok(Box::new(jsonb_each_stringify(datums[0], temp_storage)))
3726 }
3727 TableFunc::JsonbObjectKeys => Ok(Box::new(jsonb_object_keys(datums[0]))),
3728 TableFunc::JsonbArrayElements => Ok(Box::new(jsonb_array_elements(datums[0]))),
3729 TableFunc::JsonbArrayElementsStringify => Ok(Box::new(jsonb_array_elements_stringify(
3730 datums[0],
3731 temp_storage,
3732 ))),
3733 TableFunc::RegexpExtract(a) => Ok(Box::new(regexp_extract(datums[0], a).into_iter())),
3734 TableFunc::CsvExtract(n_cols) => Ok(Box::new(csv_extract(datums[0], *n_cols))),
3735 TableFunc::GenerateSeriesInt32 => {
3736 let res = generate_series(
3737 datums[0].unwrap_int32(),
3738 datums[1].unwrap_int32(),
3739 datums[2].unwrap_int32(),
3740 )?;
3741 Ok(Box::new(res))
3742 }
3743 TableFunc::GenerateSeriesInt64 | TableFunc::GenerateSeriesUnoptimized => {
3744 let res = generate_series(
3745 datums[0].unwrap_int64(),
3746 datums[1].unwrap_int64(),
3747 datums[2].unwrap_int64(),
3748 )?;
3749 Ok(Box::new(res))
3750 }
3751 TableFunc::GenerateSeriesTimestamp => {
3752 fn pass_through<'a>(d: CheckedTimestamp<NaiveDateTime>) -> Datum<'a> {
3753 Datum::from(d)
3754 }
3755 let res = generate_series_ts(
3756 datums[0].unwrap_timestamp(),
3757 datums[1].unwrap_timestamp(),
3758 datums[2].unwrap_interval(),
3759 pass_through,
3760 )?;
3761 Ok(Box::new(res))
3762 }
3763 TableFunc::GenerateSeriesTimestampTz => {
3764 fn gen_ts_tz<'a>(d: CheckedTimestamp<DateTime<Utc>>) -> Datum<'a> {
3765 Datum::from(d)
3766 }
3767 let res = generate_series_ts(
3768 datums[0].unwrap_timestamptz(),
3769 datums[1].unwrap_timestamptz(),
3770 datums[2].unwrap_interval(),
3771 gen_ts_tz,
3772 )?;
3773 Ok(Box::new(res))
3774 }
3775 TableFunc::GenerateSubscriptsArray => {
3776 generate_subscripts_array(datums[0], datums[1].unwrap_int32())
3777 }
3778 TableFunc::GuardSubquerySize { column_type: _ } => {
3779 let count = datums[0].unwrap_int64();
3791 if count > 1 {
3792 Err(EvalError::MultipleRowsFromSubquery)
3793 } else if count < 0 {
3794 Err(EvalError::NegativeRowsFromSubquery)
3796 } else {
3797 Ok(Box::new([].into_iter()))
3798 }
3799 }
3800 TableFunc::RepeatRow => Ok(Box::new(repeat_row(datums[0]).into_iter())),
3801 TableFunc::RepeatRowNonNegative => repeat_row_non_negative(datums[0]),
3802 TableFunc::UnnestArray { .. } => Ok(Box::new(unnest_array(datums[0]))),
3803 TableFunc::UnnestList { .. } => Ok(Box::new(unnest_list(datums[0]))),
3804 TableFunc::UnnestMap { .. } => Ok(Box::new(unnest_map(datums[0]))),
3805 TableFunc::Wrap { width, .. } => Ok(Box::new(wrap(datums, *width))),
3806 TableFunc::TabletizedScalar { .. } => {
3807 let r = Row::pack_slice(datums);
3808 Ok(Box::new(std::iter::once((r, Diff::ONE))))
3809 }
3810 TableFunc::RegexpMatches => Ok(Box::new(regexp_matches(datums)?)),
3811 TableFunc::WithOrdinality(func_with_ordinality) => {
3812 func_with_ordinality.eval(datums, temp_storage)
3813 }
3814 }
3815 }
3816
3817 pub fn output_sql_type(&self) -> SqlRelationType {
3818 let (column_types, keys) = match self {
3819 TableFunc::AclExplode => {
3820 let column_types = vec![
3821 SqlScalarType::Oid.nullable(false),
3822 SqlScalarType::Oid.nullable(false),
3823 SqlScalarType::String.nullable(false),
3824 SqlScalarType::Bool.nullable(false),
3825 ];
3826 let keys = vec![];
3827 (column_types, keys)
3828 }
3829 TableFunc::MzAclExplode => {
3830 let column_types = vec![
3831 SqlScalarType::String.nullable(false),
3832 SqlScalarType::String.nullable(false),
3833 SqlScalarType::String.nullable(false),
3834 SqlScalarType::Bool.nullable(false),
3835 ];
3836 let keys = vec![];
3837 (column_types, keys)
3838 }
3839 TableFunc::JsonbEach => {
3840 let column_types = vec![
3841 SqlScalarType::String.nullable(false),
3842 SqlScalarType::Jsonb.nullable(false),
3843 ];
3844 let keys = vec![];
3845 (column_types, keys)
3846 }
3847 TableFunc::JsonbEachStringify => {
3848 let column_types = vec![
3849 SqlScalarType::String.nullable(false),
3850 SqlScalarType::String.nullable(true),
3851 ];
3852 let keys = vec![];
3853 (column_types, keys)
3854 }
3855 TableFunc::JsonbObjectKeys => {
3856 let column_types = vec![SqlScalarType::String.nullable(false)];
3857 let keys = vec![];
3858 (column_types, keys)
3859 }
3860 TableFunc::JsonbArrayElements => {
3861 let column_types = vec![SqlScalarType::Jsonb.nullable(false)];
3862 let keys = vec![];
3863 (column_types, keys)
3864 }
3865 TableFunc::JsonbArrayElementsStringify => {
3866 let column_types = vec![SqlScalarType::String.nullable(true)];
3867 let keys = vec![];
3868 (column_types, keys)
3869 }
3870 TableFunc::RegexpExtract(a) => {
3871 let column_types = a
3872 .capture_groups_iter()
3873 .map(|cg| SqlScalarType::String.nullable(cg.nullable))
3874 .collect();
3875 let keys = vec![];
3876 (column_types, keys)
3877 }
3878 TableFunc::CsvExtract(n_cols) => {
3879 let column_types = iter::repeat(SqlScalarType::String.nullable(false))
3880 .take(*n_cols)
3881 .collect();
3882 let keys = vec![];
3883 (column_types, keys)
3884 }
3885 TableFunc::GenerateSeriesInt32 => {
3886 let column_types = vec![SqlScalarType::Int32.nullable(false)];
3887 let keys = vec![vec![0]];
3888 (column_types, keys)
3889 }
3890 TableFunc::GenerateSeriesInt64 | TableFunc::GenerateSeriesUnoptimized => {
3891 let column_types = vec![SqlScalarType::Int64.nullable(false)];
3892 let keys = vec![vec![0]];
3893 (column_types, keys)
3894 }
3895 TableFunc::GenerateSeriesTimestamp => {
3896 let column_types =
3897 vec![SqlScalarType::Timestamp { precision: None }.nullable(false)];
3898 let keys = vec![vec![0]];
3899 (column_types, keys)
3900 }
3901 TableFunc::GenerateSeriesTimestampTz => {
3902 let column_types =
3903 vec![SqlScalarType::TimestampTz { precision: None }.nullable(false)];
3904 let keys = vec![vec![0]];
3905 (column_types, keys)
3906 }
3907 TableFunc::GenerateSubscriptsArray => {
3908 let column_types = vec![SqlScalarType::Int32.nullable(false)];
3909 let keys = vec![vec![0]];
3910 (column_types, keys)
3911 }
3912 TableFunc::GuardSubquerySize { column_type } => {
3913 let column_types = vec![column_type.clone().nullable(false)];
3914 let keys = vec![];
3915 (column_types, keys)
3916 }
3917 TableFunc::RepeatRow | TableFunc::RepeatRowNonNegative => {
3918 let column_types = vec![];
3919 let keys = vec![];
3920 (column_types, keys)
3921 }
3922 TableFunc::UnnestArray { el_typ } => {
3923 let column_types = vec![el_typ.clone().nullable(true)];
3924 let keys = vec![];
3925 (column_types, keys)
3926 }
3927 TableFunc::UnnestList { el_typ } => {
3928 let column_types = vec![el_typ.clone().nullable(true)];
3929 let keys = vec![];
3930 (column_types, keys)
3931 }
3932 TableFunc::UnnestMap { value_type } => {
3933 let column_types = vec![
3934 SqlScalarType::String.nullable(false),
3935 value_type.clone().nullable(true),
3936 ];
3937 let keys = vec![vec![0]];
3938 (column_types, keys)
3939 }
3940 TableFunc::Wrap { types, .. } => {
3941 let column_types = types.clone();
3942 let keys = vec![];
3943 (column_types, keys)
3944 }
3945 TableFunc::TabletizedScalar { relation, .. } => {
3946 return relation.clone();
3947 }
3948 TableFunc::RegexpMatches => {
3949 let column_types =
3950 vec![SqlScalarType::Array(Box::new(SqlScalarType::String)).nullable(false)];
3951 let keys = vec![];
3952
3953 (column_types, keys)
3954 }
3955 TableFunc::WithOrdinality(WithOrdinality { inner }) => {
3956 let mut typ = inner.output_sql_type();
3957 typ.column_types.push(SqlScalarType::Int64.nullable(false));
3959 typ.keys.push(vec![typ.column_types.len() - 1]);
3961 (typ.column_types, typ.keys)
3962 }
3963 };
3964
3965 soft_assert_eq_no_log!(column_types.len(), self.output_arity());
3966
3967 if !keys.is_empty() {
3968 SqlRelationType::new(column_types).with_keys(keys)
3969 } else {
3970 SqlRelationType::new(column_types)
3971 }
3972 }
3973
3974 pub fn output_type(&self) -> ReprRelationType {
3978 ReprRelationType::from(&self.output_sql_type())
3979 }
3980
3981 pub fn output_arity(&self) -> usize {
3982 match self {
3983 TableFunc::AclExplode => 4,
3984 TableFunc::MzAclExplode => 4,
3985 TableFunc::JsonbEach => 2,
3986 TableFunc::JsonbEachStringify => 2,
3987 TableFunc::JsonbObjectKeys => 1,
3988 TableFunc::JsonbArrayElements => 1,
3989 TableFunc::JsonbArrayElementsStringify => 1,
3990 TableFunc::RegexpExtract(a) => a.capture_groups_len(),
3991 TableFunc::CsvExtract(n_cols) => *n_cols,
3992 TableFunc::GenerateSeriesInt32 => 1,
3993 TableFunc::GenerateSeriesInt64 => 1,
3994 TableFunc::GenerateSeriesUnoptimized => 1,
3995 TableFunc::GenerateSeriesTimestamp => 1,
3996 TableFunc::GenerateSeriesTimestampTz => 1,
3997 TableFunc::GenerateSubscriptsArray => 1,
3998 TableFunc::GuardSubquerySize { .. } => 1,
3999 TableFunc::RepeatRow => 0,
4000 TableFunc::RepeatRowNonNegative => 0,
4001 TableFunc::UnnestArray { .. } => 1,
4002 TableFunc::UnnestList { .. } => 1,
4003 TableFunc::UnnestMap { .. } => 2,
4004 TableFunc::Wrap { width, .. } => *width,
4005 TableFunc::TabletizedScalar { relation, .. } => relation.column_types.len(),
4006 TableFunc::RegexpMatches => 1,
4007 TableFunc::WithOrdinality(WithOrdinality { inner }) => inner.output_arity() + 1,
4008 }
4009 }
4010
4011 pub fn empty_on_null_input(&self) -> bool {
4012 match self {
4013 TableFunc::AclExplode
4014 | TableFunc::MzAclExplode
4015 | TableFunc::JsonbEach
4016 | TableFunc::JsonbEachStringify
4017 | TableFunc::JsonbObjectKeys
4018 | TableFunc::JsonbArrayElements
4019 | TableFunc::JsonbArrayElementsStringify
4020 | TableFunc::GenerateSeriesInt32
4021 | TableFunc::GenerateSeriesInt64
4022 | TableFunc::GenerateSeriesUnoptimized
4023 | TableFunc::GenerateSeriesTimestamp
4024 | TableFunc::GenerateSeriesTimestampTz
4025 | TableFunc::GenerateSubscriptsArray
4026 | TableFunc::RegexpExtract(_)
4027 | TableFunc::CsvExtract(_)
4028 | TableFunc::RepeatRow
4029 | TableFunc::RepeatRowNonNegative
4030 | TableFunc::UnnestArray { .. }
4031 | TableFunc::UnnestList { .. }
4032 | TableFunc::UnnestMap { .. }
4033 | TableFunc::RegexpMatches => true,
4034 TableFunc::GuardSubquerySize { .. } => false,
4035 TableFunc::Wrap { .. } => false,
4036 TableFunc::TabletizedScalar { .. } => false,
4037 TableFunc::WithOrdinality(WithOrdinality { inner }) => inner.empty_on_null_input(),
4038 }
4039 }
4040
4041 pub fn preserves_monotonicity(&self) -> bool {
4043 match self {
4046 TableFunc::AclExplode => false,
4047 TableFunc::MzAclExplode => false,
4048 TableFunc::JsonbEach => true,
4049 TableFunc::JsonbEachStringify => true,
4050 TableFunc::JsonbObjectKeys => true,
4051 TableFunc::JsonbArrayElements => true,
4052 TableFunc::JsonbArrayElementsStringify => true,
4053 TableFunc::RegexpExtract(_) => true,
4054 TableFunc::CsvExtract(_) => true,
4055 TableFunc::GenerateSeriesInt32 => true,
4056 TableFunc::GenerateSeriesInt64 => true,
4057 TableFunc::GenerateSeriesUnoptimized => true,
4058 TableFunc::GenerateSeriesTimestamp => true,
4059 TableFunc::GenerateSeriesTimestampTz => true,
4060 TableFunc::GenerateSubscriptsArray => true,
4061 TableFunc::RepeatRow => false,
4062 TableFunc::RepeatRowNonNegative => true,
4063 TableFunc::UnnestArray { .. } => true,
4064 TableFunc::UnnestList { .. } => true,
4065 TableFunc::UnnestMap { .. } => true,
4066 TableFunc::Wrap { .. } => true,
4067 TableFunc::TabletizedScalar { .. } => true,
4068 TableFunc::RegexpMatches => true,
4069 TableFunc::GuardSubquerySize { .. } => false,
4070 TableFunc::WithOrdinality(WithOrdinality { inner }) => inner.preserves_monotonicity(),
4071 }
4072 }
4073}
4074
4075impl fmt::Display for TableFunc {
4076 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
4077 match self {
4078 TableFunc::AclExplode => f.write_str("aclexplode"),
4079 TableFunc::MzAclExplode => f.write_str("mz_aclexplode"),
4080 TableFunc::JsonbEach => f.write_str("jsonb_each"),
4081 TableFunc::JsonbEachStringify => f.write_str("jsonb_each_text"),
4082 TableFunc::JsonbObjectKeys => f.write_str("jsonb_object_keys"),
4083 TableFunc::JsonbArrayElements => f.write_str("jsonb_array_elements"),
4084 TableFunc::JsonbArrayElementsStringify => f.write_str("jsonb_array_elements_text"),
4085 TableFunc::RegexpExtract(a) => write!(f, "regexp_extract({:?}, _)", a.0),
4086 TableFunc::CsvExtract(n_cols) => write!(f, "csv_extract({}, _)", n_cols),
4087 TableFunc::GenerateSeriesInt32 => f.write_str("generate_series"),
4088 TableFunc::GenerateSeriesInt64 => f.write_str("generate_series"),
4089 TableFunc::GenerateSeriesUnoptimized => f.write_str("generate_series_unoptimized"),
4090 TableFunc::GenerateSeriesTimestamp => f.write_str("generate_series"),
4091 TableFunc::GenerateSeriesTimestampTz => f.write_str("generate_series"),
4092 TableFunc::GenerateSubscriptsArray => f.write_str("generate_subscripts"),
4093 TableFunc::GuardSubquerySize { .. } => f.write_str("guard_subquery_size"),
4094 TableFunc::RepeatRow => f.write_str(REPEAT_ROW_NAME),
4095 TableFunc::RepeatRowNonNegative => f.write_str("repeat_row_non_negative"),
4096 TableFunc::UnnestArray { .. } => f.write_str("unnest_array"),
4097 TableFunc::UnnestList { .. } => f.write_str("unnest_list"),
4098 TableFunc::UnnestMap { .. } => f.write_str("unnest_map"),
4099 TableFunc::Wrap { width, .. } => write!(f, "wrap{}", width),
4100 TableFunc::TabletizedScalar { name, .. } => f.write_str(name),
4101 TableFunc::RegexpMatches => write!(f, "regexp_matches(_, _, _)"),
4102 TableFunc::WithOrdinality(WithOrdinality { inner }) => {
4103 write!(f, "{}[with_ordinality]", inner)
4104 }
4105 }
4106 }
4107}
4108
4109impl WithOrdinality {
4110 fn eval<'a>(
4119 &'a self,
4120 datums: &'a [Datum<'a>],
4121 temp_storage: &'a RowArena,
4122 ) -> Result<Box<dyn Iterator<Item = (Row, Diff)> + 'a>, EvalError> {
4123 let mut next_ordinal: i64 = 1;
4124 let it = self
4125 .inner
4126 .eval(datums, temp_storage)?
4127 .flat_map(move |(mut row, diff)| {
4128 let diff = diff.into_inner();
4129 assert!(diff >= 0);
4137 let mut ordinals = next_ordinal..(next_ordinal + diff);
4139 next_ordinal += diff;
4140 let cap = row.data_len() + datum_size(&Datum::Int64(next_ordinal));
4142 iter::from_fn(move || {
4143 let ordinal = ordinals.next()?;
4144 let mut row = if ordinals.is_empty() {
4145 std::mem::take(&mut row)
4148 } else {
4149 let mut new_row = Row::with_capacity(cap);
4150 new_row.clone_from(&row);
4151 new_row
4152 };
4153 RowPacker::for_existing_row(&mut row).push(Datum::Int64(ordinal));
4154 Some((row, Diff::ONE))
4155 })
4156 });
4157 Ok(Box::new(it))
4158 }
4159}
4160
4161pub const REPEAT_ROW_NAME: &str = "repeat_row";
4162
4163#[cfg(test)]
4164mod tests {
4165 use mz_repr::{Datum, RowArena, SqlScalarType};
4166
4167 use super::TableFunc;
4168 use crate::EvalError;
4169
4170 #[mz_ore::test]
4177 fn guard_subquery_size_accepts_zero_and_one() {
4178 let func = TableFunc::GuardSubquerySize {
4179 column_type: SqlScalarType::Int64,
4180 };
4181 let temp_storage = RowArena::new();
4182
4183 for count in [0_i64, 1] {
4184 let rows = func
4185 .eval(&[Datum::Int64(count)], &temp_storage)
4186 .unwrap_or_else(|e| panic!("count {count} should be accepted, got {e:?}"))
4187 .count();
4188 assert_eq!(rows, 0, "count {count} should emit no guard rows");
4189 }
4190
4191 assert_eq!(
4192 func.eval(&[Datum::Int64(2)], &temp_storage).err(),
4193 Some(EvalError::MultipleRowsFromSubquery),
4194 );
4195 assert_eq!(
4196 func.eval(&[Datum::Int64(-1)], &temp_storage).err(),
4197 Some(EvalError::NegativeRowsFromSubquery),
4198 );
4199 }
4200}