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 sum_interval<'a, I>(datums: I) -> Datum<'a>
198where
199 I: IntoIterator<Item = Datum<'a>>,
200{
201 sum_interval_counted(datums.into_iter().map(|datum| (datum, Diff::ONE)))
202}
203
204fn sum_interval_counted<'a, I>(datums: I) -> Datum<'a>
226where
227 I: IntoIterator<Item = (Datum<'a>, Diff)>,
228{
229 let (mut months, mut days, mut micros) = (0i128, 0i128, 0i128);
230 let mut non_nulls = Diff::ZERO;
231 for (datum, diff) in datums {
232 if datum.is_null() {
233 continue;
234 }
235 let interval = datum.unwrap_interval();
236 let scale = i128::from(diff.into_inner());
241 months = months.wrapping_add(i128::from(interval.months).wrapping_mul(scale));
242 days = days.wrapping_add(i128::from(interval.days).wrapping_mul(scale));
243 micros = micros.wrapping_add(i128::from(interval.micros).wrapping_mul(scale));
244 non_nulls += diff;
245 }
246 if months == 0 && days == 0 && micros == 0 && non_nulls.is_zero() {
247 return Datum::Null;
248 }
249 match Interval::try_new(months, days, micros) {
250 Some(interval) => Datum::Interval(interval),
251 None => Datum::Null,
252 }
253}
254
255fn count<'a, I>(datums: I) -> Datum<'a>
256where
257 I: IntoIterator<Item = (Datum<'a>, Diff)>,
258{
259 let mut count = Diff::ZERO;
265 for (datum, diff) in datums {
266 if !datum.is_null() {
267 count += diff;
268 }
269 }
270 Datum::from(count.into_inner())
271}
272
273fn any<'a, I>(datums: I) -> Datum<'a>
274where
275 I: IntoIterator<Item = Datum<'a>>,
276{
277 datums
278 .into_iter()
279 .fold(Datum::False, |state, next| match (state, next) {
280 (Datum::True, _) | (_, Datum::True) => Datum::True,
281 (Datum::Null, _) | (_, Datum::Null) => Datum::Null,
282 _ => Datum::False,
283 })
284}
285
286fn all<'a, I>(datums: I) -> Datum<'a>
287where
288 I: IntoIterator<Item = Datum<'a>>,
289{
290 datums
291 .into_iter()
292 .fold(Datum::True, |state, next| match (state, next) {
293 (Datum::False, _) | (_, Datum::False) => Datum::False,
294 (Datum::Null, _) | (_, Datum::Null) => Datum::Null,
295 _ => Datum::True,
296 })
297}
298
299fn string_agg<'a, I>(datums: I, temp_storage: &'a RowArena, order_by: &[ColumnOrder]) -> Datum<'a>
300where
301 I: IntoIterator<Item = Datum<'a>>,
302{
303 const EMPTY_SEP: &str = "";
304
305 let datums = order_aggregate_datums(datums, order_by);
306 let mut sep_value_pairs = datums.into_iter().filter_map(|d| {
307 if d.is_null() {
308 return None;
309 }
310 let mut value_sep = d.unwrap_list().iter();
311 match (value_sep.next().unwrap(), value_sep.next().unwrap()) {
312 (Datum::Null, _) => None,
313 (Datum::String(val), Datum::Null) => Some((EMPTY_SEP, val)),
314 (Datum::String(val), Datum::String(sep)) => Some((sep, val)),
315 _ => unreachable!(),
316 }
317 });
318
319 let mut s = String::default();
320 match sep_value_pairs.next() {
321 Some((_, value)) => s.push_str(value),
323 None => return Datum::Null,
325 }
326
327 for (sep, value) in sep_value_pairs {
328 s.push_str(sep);
329 s.push_str(value);
330 }
331
332 Datum::String(temp_storage.push_string(s))
333}
334
335fn jsonb_agg<'a, I>(datums: I, temp_storage: &'a RowArena, order_by: &[ColumnOrder]) -> Datum<'a>
336where
337 I: IntoIterator<Item = Datum<'a>>,
338{
339 let datums = order_aggregate_datums(datums, order_by);
340 temp_storage.make_datum(|packer| {
341 packer.push_list(datums.into_iter().filter(|d| !d.is_null()));
342 })
343}
344
345fn dict_agg<'a, I>(datums: I, temp_storage: &'a RowArena, order_by: &[ColumnOrder]) -> Datum<'a>
346where
347 I: IntoIterator<Item = Datum<'a>>,
348{
349 let datums = order_aggregate_datums(datums, order_by);
350 temp_storage.make_datum(|packer| {
351 let mut datums: Vec<_> = datums
352 .into_iter()
353 .filter_map(|d| {
354 if d.is_null() {
355 return None;
356 }
357 let mut list = d.unwrap_list().iter();
358 let key = list.next().unwrap();
359 let val = list.next().unwrap();
360 if key.is_null() {
361 None
364 } else {
365 Some((key.unwrap_str(), val))
366 }
367 })
368 .collect();
369 datums.sort_by_key(|(k, _v)| *k);
375 datums.reverse();
376 datums.dedup_by_key(|(k, _v)| *k);
377 datums.reverse();
378 packer.push_dict(datums);
379 })
380}
381
382pub fn order_aggregate_datums<'a: 'b, 'b, I>(
392 datums: I,
393 order_by: &[ColumnOrder],
394) -> impl Iterator<Item = Datum<'b>>
395where
396 I: IntoIterator<Item = Datum<'a>>,
397{
398 order_aggregate_datums_with_rank_inner(datums, order_by)
399 .into_iter()
400 .map(|(payload, _order_datums)| payload)
402}
403
404fn order_aggregate_datums_with_rank<'a, I>(
407 datums: I,
408 order_by: &[ColumnOrder],
409) -> impl Iterator<Item = (Datum<'a>, Row)>
410where
411 I: IntoIterator<Item = Datum<'a>>,
412{
413 order_aggregate_datums_with_rank_inner(datums, order_by)
414 .into_iter()
415 .map(|(payload, order_by_datums)| (payload, Row::pack(order_by_datums)))
416}
417
418fn order_aggregate_datums_with_rank_inner<'a, I>(
419 datums: I,
420 order_by: &[ColumnOrder],
421) -> Vec<(Datum<'a>, Vec<Datum<'a>>)>
422where
423 I: IntoIterator<Item = Datum<'a>>,
424{
425 let mut decoded: Vec<(Datum, Vec<Datum>)> = datums
426 .into_iter()
427 .map(|d| {
428 let list = d.unwrap_list();
429 let mut list_it = list.iter();
430 let payload = list_it.next().unwrap();
431
432 let mut order_by_datums = Vec::with_capacity(order_by.len());
442 for _ in 0..order_by.len() {
443 order_by_datums.push(
444 list_it
445 .next()
446 .expect("must have exactly the same number of Datums as `order_by`"),
447 );
448 }
449
450 (payload, order_by_datums)
451 })
452 .collect();
453
454 let mut sort_by =
455 |(payload_left, left_order_by_datums): &(Datum, Vec<Datum>),
456 (payload_right, right_order_by_datums): &(Datum, Vec<Datum>)| {
457 compare_columns(
458 order_by,
459 left_order_by_datums,
460 right_order_by_datums,
461 || payload_left.cmp(payload_right),
462 )
463 };
464 decoded.sort_unstable_by(&mut sort_by);
469 decoded
470}
471
472fn array_concat<'a, I>(datums: I, temp_storage: &'a RowArena, order_by: &[ColumnOrder]) -> Datum<'a>
473where
474 I: IntoIterator<Item = Datum<'a>>,
475{
476 let datums = order_aggregate_datums(datums, order_by);
477 let datums: Vec<_> = datums
478 .into_iter()
479 .map(|d| d.unwrap_array().elements().iter())
480 .flatten()
481 .collect();
482 let dims = ArrayDimension {
483 lower_bound: 1,
484 length: datums.len(),
485 };
486 temp_storage.make_datum(|packer| {
487 packer.try_push_array(&[dims], datums).unwrap();
488 })
489}
490
491fn list_concat<'a, I>(datums: I, temp_storage: &'a RowArena, order_by: &[ColumnOrder]) -> Datum<'a>
492where
493 I: IntoIterator<Item = Datum<'a>>,
494{
495 let datums = order_aggregate_datums(datums, order_by);
496 temp_storage.make_datum(|packer| {
497 packer.push_list(datums.into_iter().map(|d| d.unwrap_list().iter()).flatten());
498 })
499}
500
501fn row_number<'a, I>(
505 datums: I,
506 callers_temp_storage: &'a RowArena,
507 order_by: &[ColumnOrder],
508) -> Datum<'a>
509where
510 I: IntoIterator<Item = Datum<'a>>,
511{
512 let temp_storage = RowArena::new();
516 let datums = row_number_no_list(datums, &temp_storage, order_by);
517
518 callers_temp_storage.make_datum(|packer| {
519 packer.push_list(datums);
520 })
521}
522
523fn row_number_no_list<'a: 'b, 'b, I>(
526 datums: I,
527 callers_temp_storage: &'b RowArena,
528 order_by: &[ColumnOrder],
529) -> impl Iterator<Item = Datum<'b>>
530where
531 I: IntoIterator<Item = Datum<'a>>,
532{
533 let datums = order_aggregate_datums(datums, order_by);
534
535 callers_temp_storage.reserve(datums.size_hint().0);
536 #[allow(clippy::disallowed_methods)]
537 datums
538 .into_iter()
539 .map(|d| d.unwrap_list().iter())
540 .flatten()
541 .zip(1i64..)
542 .map(|(d, i)| {
543 callers_temp_storage.make_datum(|packer| {
544 packer.push_list_with(|packer| {
545 packer.push(Datum::Int64(i));
546 packer.push(d);
547 });
548 })
549 })
550}
551
552fn rank<'a, I>(datums: I, callers_temp_storage: &'a RowArena, order_by: &[ColumnOrder]) -> Datum<'a>
556where
557 I: IntoIterator<Item = Datum<'a>>,
558{
559 let temp_storage = RowArena::new();
560 let datums = rank_no_list(datums, &temp_storage, order_by);
561
562 callers_temp_storage.make_datum(|packer| {
563 packer.push_list(datums);
564 })
565}
566
567fn rank_no_list<'a: 'b, 'b, I>(
570 datums: I,
571 callers_temp_storage: &'b RowArena,
572 order_by: &[ColumnOrder],
573) -> impl Iterator<Item = Datum<'b>>
574where
575 I: IntoIterator<Item = Datum<'a>>,
576{
577 let datums = order_aggregate_datums_with_rank(datums, order_by);
579
580 let mut datums = datums
581 .into_iter()
582 .map(|(d0, order_row)| {
583 d0.unwrap_list()
584 .iter()
585 .map(move |d1| (d1, order_row.clone()))
586 })
587 .flatten();
588
589 callers_temp_storage.reserve(datums.size_hint().0);
590 datums
591 .next()
592 .map_or(vec![], |(first_datum, first_order_row)| {
593 datums.fold(
596 (first_order_row, 1, 1, vec![(first_datum, 1)]),
597 |mut acc, (next_datum, next_order_row)| {
598 let (ref mut acc_row, ref mut acc_rank, ref mut acc_row_num, ref mut output) = acc;
599 *acc_row_num += 1;
600 if *acc_row != next_order_row {
602 *acc_rank = *acc_row_num;
603 *acc_row = next_order_row;
604 }
605
606 (*output).push((next_datum, *acc_rank));
607 acc
608 })
609 }.3).into_iter().map(|(d, i)| {
610 callers_temp_storage.make_datum(|packer| {
611 packer.push_list_with(|packer| {
612 packer.push(Datum::Int64(i));
613 packer.push(d);
614 });
615 })
616 })
617}
618
619fn dense_rank<'a, I>(
623 datums: I,
624 callers_temp_storage: &'a RowArena,
625 order_by: &[ColumnOrder],
626) -> Datum<'a>
627where
628 I: IntoIterator<Item = Datum<'a>>,
629{
630 let temp_storage = RowArena::new();
631 let datums = dense_rank_no_list(datums, &temp_storage, order_by);
632
633 callers_temp_storage.make_datum(|packer| {
634 packer.push_list(datums);
635 })
636}
637
638fn dense_rank_no_list<'a: 'b, 'b, I>(
641 datums: I,
642 callers_temp_storage: &'b RowArena,
643 order_by: &[ColumnOrder],
644) -> impl Iterator<Item = Datum<'b>>
645where
646 I: IntoIterator<Item = Datum<'a>>,
647{
648 let datums = order_aggregate_datums_with_rank(datums, order_by);
650
651 let mut datums = datums
652 .into_iter()
653 .map(|(d0, order_row)| {
654 d0.unwrap_list()
655 .iter()
656 .map(move |d1| (d1, order_row.clone()))
657 })
658 .flatten();
659
660 callers_temp_storage.reserve(datums.size_hint().0);
661 datums
662 .next()
663 .map_or(vec![], |(first_datum, first_order_row)| {
664 datums.fold(
667 (first_order_row, 1, vec![(first_datum, 1)]),
668 |mut acc, (next_datum, next_order_row)| {
669 let (ref mut acc_row, ref mut acc_rank, ref mut output) = acc;
670 if *acc_row != next_order_row {
672 *acc_rank += 1;
673 *acc_row = next_order_row;
674 }
675
676 (*output).push((next_datum, *acc_rank));
677 acc
678 })
679 }.2).into_iter().map(|(d, i)| {
680 callers_temp_storage.make_datum(|packer| {
681 packer.push_list_with(|packer| {
682 packer.push(Datum::Int64(i));
683 packer.push(d);
684 });
685 })
686 })
687}
688
689fn lag_lead<'a, I>(
711 datums: I,
712 callers_temp_storage: &'a RowArena,
713 order_by: &[ColumnOrder],
714 lag_lead_type: &LagLeadType,
715 ignore_nulls: &bool,
716) -> Datum<'a>
717where
718 I: IntoIterator<Item = Datum<'a>>,
719{
720 let temp_storage = RowArena::new();
721 let iter = lag_lead_no_list(datums, &temp_storage, order_by, lag_lead_type, ignore_nulls);
722 callers_temp_storage.make_datum(|packer| {
723 packer.push_list(iter);
724 })
725}
726
727fn lag_lead_no_list<'a: 'b, 'b, I>(
730 datums: I,
731 callers_temp_storage: &'b RowArena,
732 order_by: &[ColumnOrder],
733 lag_lead_type: &LagLeadType,
734 ignore_nulls: &bool,
735) -> impl Iterator<Item = Datum<'b>>
736where
737 I: IntoIterator<Item = Datum<'a>>,
738{
739 let datums = order_aggregate_datums(datums, order_by);
741
742 let (orig_rows, unwrapped_args): (Vec<_>, Vec<_>) = datums
746 .into_iter()
747 .map(|d| {
748 let mut iter = d.unwrap_list().iter();
749 let original_row = iter.next().unwrap();
750 let (input_value, offset, default_value) =
751 unwrap_lag_lead_encoded_args(iter.next().unwrap());
752 (original_row, (input_value, offset, default_value))
753 })
754 .unzip();
755
756 let result = lag_lead_inner(unwrapped_args, lag_lead_type, ignore_nulls);
757
758 callers_temp_storage.reserve(result.len());
759 result
760 .into_iter()
761 .zip_eq(orig_rows)
762 .map(|(result_value, original_row)| {
763 callers_temp_storage.make_datum(|packer| {
764 packer.push_list_with(|packer| {
765 packer.push(result_value);
766 packer.push(original_row);
767 });
768 })
769 })
770}
771
772fn unwrap_lag_lead_encoded_args(encoded_args: Datum) -> (Datum, Datum, Datum) {
774 let mut encoded_args_iter = encoded_args.unwrap_list().iter();
775 let (input_value, offset, default_value) = (
776 encoded_args_iter.next().unwrap(),
777 encoded_args_iter.next().unwrap(),
778 encoded_args_iter.next().unwrap(),
779 );
780 (input_value, offset, default_value)
781}
782
783fn lag_lead_inner<'a>(
786 args: Vec<(Datum<'a>, Datum<'a>, Datum<'a>)>,
787 lag_lead_type: &LagLeadType,
788 ignore_nulls: &bool,
789) -> Vec<Datum<'a>> {
790 if *ignore_nulls {
791 lag_lead_inner_ignore_nulls(args, lag_lead_type)
792 } else {
793 lag_lead_inner_respect_nulls(args, lag_lead_type)
794 }
795}
796
797fn lag_lead_inner_respect_nulls<'a>(
798 args: Vec<(Datum<'a>, Datum<'a>, Datum<'a>)>,
799 lag_lead_type: &LagLeadType,
800) -> Vec<Datum<'a>> {
801 let mut result: Vec<Datum> = Vec::with_capacity(args.len());
802 for (idx, (_, offset, default_value)) in args.iter().enumerate() {
803 if offset.is_null() {
805 result.push(Datum::Null);
806 continue;
807 }
808
809 let idx = i64::try_from(idx).expect("Array index does not fit in i64");
810 let offset = i64::from(offset.unwrap_int32());
811 let offset = match lag_lead_type {
812 LagLeadType::Lag => -offset,
813 LagLeadType::Lead => offset,
814 };
815
816 let datums_get = |i: i64| -> Option<Datum> {
818 match u64::try_from(i) {
819 Ok(i) => args
820 .get(usize::cast_from(i))
821 .map(|d| Some(d.0)) .unwrap_or(None), Err(_) => None, }
825 };
826
827 let lagged_value = datums_get(idx + offset).unwrap_or(*default_value);
828
829 result.push(lagged_value);
830 }
831
832 result
833}
834
835#[allow(clippy::as_conversions)]
839fn lag_lead_inner_ignore_nulls<'a>(
840 args: Vec<(Datum<'a>, Datum<'a>, Datum<'a>)>,
841 lag_lead_type: &LagLeadType,
842) -> Vec<Datum<'a>> {
843 if i64::try_from(args.len()).is_err() {
846 panic!("window partition way too big")
847 }
848 let mut skip_nulls_backward = vec![None; args.len()];
851 let mut last_non_null: i64 = -1;
852 let pairs = args
853 .iter()
854 .enumerate()
855 .zip_eq(skip_nulls_backward.iter_mut());
856 for ((i, (d, _, _)), slot) in pairs {
857 if d.is_null() {
858 *slot = Some(last_non_null);
859 } else {
860 last_non_null = i as i64;
861 }
862 }
863 let mut skip_nulls_forward = vec![None; args.len()];
864 let mut last_non_null: i64 = args.len() as i64;
865 let pairs = args
866 .iter()
867 .enumerate()
868 .rev()
869 .zip_eq(skip_nulls_forward.iter_mut().rev());
870 for ((i, (d, _, _)), slot) in pairs {
871 if d.is_null() {
872 *slot = Some(last_non_null);
873 } else {
874 last_non_null = i as i64;
875 }
876 }
877
878 let mut result: Vec<Datum> = Vec::with_capacity(args.len());
880 for (idx, (_, offset, default_value)) in args.iter().enumerate() {
881 if offset.is_null() {
883 result.push(Datum::Null);
884 continue;
885 }
886
887 let idx = idx as i64; let offset = i64::cast_from(offset.unwrap_int32());
889 let offset = match lag_lead_type {
890 LagLeadType::Lag => -offset,
891 LagLeadType::Lead => offset,
892 };
893 let increment = offset.signum();
894
895 let datums_get = |i: i64| -> Option<Datum> {
897 match u64::try_from(i) {
898 Ok(i) => args
899 .get(usize::cast_from(i))
900 .map(|d| Some(d.0)) .unwrap_or(None), Err(_) => None, }
904 };
905
906 let lagged_value = if increment != 0 {
907 let mut j = idx;
917 for _ in 0..num::abs(offset) {
918 j += increment;
919 if datums_get(j).is_some_and(|d| d.is_null()) {
921 let ju = j as usize; if increment > 0 {
923 j = skip_nulls_forward[ju].expect("checked above that it's null");
924 } else {
925 j = skip_nulls_backward[ju].expect("checked above that it's null");
926 }
927 }
928 if datums_get(j).is_none() {
929 break;
930 }
931 }
932 match datums_get(j) {
933 Some(datum) => datum,
934 None => *default_value,
935 }
936 } else {
937 assert_eq!(offset, 0);
938 let datum = datums_get(idx).expect("known to exist");
939 if !datum.is_null() {
940 datum
941 } else {
942 panic!("0 offset in lag/lead IGNORE NULLS");
947 }
948 };
949
950 result.push(lagged_value);
951 }
952
953 result
954}
955
956fn first_value<'a, I>(
958 datums: I,
959 callers_temp_storage: &'a RowArena,
960 order_by: &[ColumnOrder],
961 window_frame: &WindowFrame,
962) -> Datum<'a>
963where
964 I: IntoIterator<Item = Datum<'a>>,
965{
966 let temp_storage = RowArena::new();
967 let iter = first_value_no_list(datums, &temp_storage, order_by, window_frame);
968 callers_temp_storage.make_datum(|packer| {
969 packer.push_list(iter);
970 })
971}
972
973fn first_value_no_list<'a: 'b, 'b, I>(
976 datums: I,
977 callers_temp_storage: &'b RowArena,
978 order_by: &[ColumnOrder],
979 window_frame: &WindowFrame,
980) -> impl Iterator<Item = Datum<'b>>
981where
982 I: IntoIterator<Item = Datum<'a>>,
983{
984 let datums = order_aggregate_datums(datums, order_by);
986
987 let (orig_rows, args): (Vec<_>, Vec<_>) = datums
989 .into_iter()
990 .map(|d| {
991 let mut iter = d.unwrap_list().iter();
992 let original_row = iter.next().unwrap();
993 let arg = iter.next().unwrap();
994
995 (original_row, arg)
996 })
997 .unzip();
998
999 let results = first_value_inner(args, window_frame);
1000
1001 callers_temp_storage.reserve(results.len());
1002 results
1003 .into_iter()
1004 .zip_eq(orig_rows)
1005 .map(|(result_value, original_row)| {
1006 callers_temp_storage.make_datum(|packer| {
1007 packer.push_list_with(|packer| {
1008 packer.push(result_value);
1009 packer.push(original_row);
1010 });
1011 })
1012 })
1013}
1014
1015fn first_value_inner<'a>(datums: Vec<Datum<'a>>, window_frame: &WindowFrame) -> Vec<Datum<'a>> {
1016 let length = datums.len();
1017 let mut result: Vec<Datum> = Vec::with_capacity(length);
1018 for (idx, current_datum) in datums.iter().enumerate() {
1019 let first_value = match &window_frame.start_bound {
1020 WindowFrameBound::CurrentRow => *current_datum,
1022 WindowFrameBound::UnboundedPreceding => {
1023 if let WindowFrameBound::OffsetPreceding(end_offset) = &window_frame.end_bound {
1024 let end_offset = usize::cast_from(*end_offset);
1025
1026 if idx < end_offset {
1028 Datum::Null
1029 } else {
1030 datums[0]
1031 }
1032 } else {
1033 datums[0]
1034 }
1035 }
1036 WindowFrameBound::OffsetPreceding(offset) => {
1037 let start_offset = usize::cast_from(*offset);
1038 let start_idx = idx.saturating_sub(start_offset);
1039 if let WindowFrameBound::OffsetPreceding(end_offset) = &window_frame.end_bound {
1040 let end_offset = usize::cast_from(*end_offset);
1041
1042 if start_offset < end_offset || idx < end_offset {
1044 Datum::Null
1045 } else {
1046 datums[start_idx]
1047 }
1048 } else {
1049 datums[start_idx]
1050 }
1051 }
1052 WindowFrameBound::OffsetFollowing(offset) => {
1053 let start_offset = usize::cast_from(*offset);
1054 let start_idx = idx.saturating_add(start_offset);
1055 if let WindowFrameBound::OffsetFollowing(end_offset) = &window_frame.end_bound {
1056 if offset > end_offset || start_idx >= length {
1058 Datum::Null
1059 } else {
1060 datums[start_idx]
1061 }
1062 } else {
1063 datums
1064 .get(start_idx)
1065 .map(|d| d.clone())
1066 .unwrap_or(Datum::Null)
1067 }
1068 }
1069 WindowFrameBound::UnboundedFollowing => unreachable!(),
1071 };
1072 result.push(first_value);
1073 }
1074 result
1075}
1076
1077fn last_value<'a, I>(
1079 datums: I,
1080 callers_temp_storage: &'a RowArena,
1081 order_by: &[ColumnOrder],
1082 window_frame: &WindowFrame,
1083) -> Datum<'a>
1084where
1085 I: IntoIterator<Item = Datum<'a>>,
1086{
1087 let temp_storage = RowArena::new();
1088 let iter = last_value_no_list(datums, &temp_storage, order_by, window_frame);
1089 callers_temp_storage.make_datum(|packer| {
1090 packer.push_list(iter);
1091 })
1092}
1093
1094fn last_value_no_list<'a: 'b, 'b, I>(
1097 datums: I,
1098 callers_temp_storage: &'b RowArena,
1099 order_by: &[ColumnOrder],
1100 window_frame: &WindowFrame,
1101) -> impl Iterator<Item = Datum<'b>>
1102where
1103 I: IntoIterator<Item = Datum<'a>>,
1104{
1105 let datums = order_aggregate_datums_with_rank(datums, order_by);
1108
1109 let size_hint = datums.size_hint().0;
1111 let mut args = Vec::with_capacity(size_hint);
1112 let mut original_rows = Vec::with_capacity(size_hint);
1113 let mut order_by_rows = Vec::with_capacity(size_hint);
1114 for (d, order_by_row) in datums.into_iter() {
1115 let mut iter = d.unwrap_list().iter();
1116 let original_row = iter.next().unwrap();
1117 let arg = iter.next().unwrap();
1118 order_by_rows.push(order_by_row);
1119 original_rows.push(original_row);
1120 args.push(arg);
1121 }
1122
1123 let results = last_value_inner(args, &order_by_rows, window_frame);
1124
1125 callers_temp_storage.reserve(results.len());
1126 results
1127 .into_iter()
1128 .zip_eq(original_rows)
1129 .map(|(result_value, original_row)| {
1130 callers_temp_storage.make_datum(|packer| {
1131 packer.push_list_with(|packer| {
1132 packer.push(result_value);
1133 packer.push(original_row);
1134 });
1135 })
1136 })
1137}
1138
1139fn last_value_inner<'a>(
1140 args: Vec<Datum<'a>>,
1141 order_by_rows: &Vec<Row>,
1142 window_frame: &WindowFrame,
1143) -> Vec<Datum<'a>> {
1144 let length = args.len();
1145 let mut results: Vec<Datum> = Vec::with_capacity(length);
1146 for (idx, (current_datum, order_by_row)) in args.iter().zip_eq(order_by_rows).enumerate() {
1147 let last_value = match &window_frame.end_bound {
1148 WindowFrameBound::CurrentRow => match &window_frame.units {
1149 WindowFrameUnits::Rows => *current_datum,
1151 WindowFrameUnits::Range => {
1152 let target_idx = order_by_rows[idx..]
1157 .iter()
1158 .enumerate()
1159 .take_while(|(_, row)| *row == order_by_row)
1160 .last()
1161 .unwrap()
1162 .0
1163 + idx;
1164 args[target_idx]
1165 }
1166 WindowFrameUnits::Groups => unreachable!(),
1168 },
1169 WindowFrameBound::UnboundedFollowing => {
1170 if let WindowFrameBound::OffsetFollowing(start_offset) = &window_frame.start_bound {
1171 let start_offset = usize::cast_from(*start_offset);
1172
1173 if idx + start_offset > length - 1 {
1175 Datum::Null
1176 } else {
1177 args[length - 1]
1178 }
1179 } else {
1180 args[length - 1]
1181 }
1182 }
1183 WindowFrameBound::OffsetFollowing(offset) => {
1184 let end_offset = usize::cast_from(*offset);
1185 let end_idx = idx.saturating_add(end_offset);
1186 if let WindowFrameBound::OffsetFollowing(start_offset) = &window_frame.start_bound {
1187 let start_offset = usize::cast_from(*start_offset);
1188 let start_idx = idx.saturating_add(start_offset);
1189
1190 if end_offset < start_offset || start_idx >= length {
1192 Datum::Null
1193 } else {
1194 args.get(end_idx).unwrap_or(&args[length - 1]).clone()
1196 }
1197 } else {
1198 args.get(end_idx).unwrap_or(&args[length - 1]).clone()
1199 }
1200 }
1201 WindowFrameBound::OffsetPreceding(offset) => {
1202 let end_offset = usize::cast_from(*offset);
1203 let end_idx = idx.saturating_sub(end_offset);
1204 if idx < end_offset {
1205 Datum::Null
1207 } else if let WindowFrameBound::OffsetPreceding(start_offset) =
1208 &window_frame.start_bound
1209 {
1210 if offset > start_offset {
1212 Datum::Null
1213 } else {
1214 args[end_idx]
1215 }
1216 } else {
1217 args[end_idx]
1218 }
1219 }
1220 WindowFrameBound::UnboundedPreceding => unreachable!(),
1222 };
1223 results.push(last_value);
1224 }
1225 results
1226}
1227
1228fn fused_value_window_func<'a, I>(
1234 input_datums: I,
1235 callers_temp_storage: &'a RowArena,
1236 funcs: &Vec<AggregateFunc>,
1237 order_by: &Vec<ColumnOrder>,
1238) -> Datum<'a>
1239where
1240 I: IntoIterator<Item = Datum<'a>>,
1241{
1242 let temp_storage = RowArena::new();
1243 let iter = fused_value_window_func_no_list(input_datums, &temp_storage, funcs, order_by);
1244 callers_temp_storage.make_datum(|packer| {
1245 packer.push_list(iter);
1246 })
1247}
1248
1249fn fused_value_window_func_no_list<'a: 'b, 'b, I>(
1252 input_datums: I,
1253 callers_temp_storage: &'b RowArena,
1254 funcs: &Vec<AggregateFunc>,
1255 order_by: &Vec<ColumnOrder>,
1256) -> impl Iterator<Item = Datum<'b>>
1257where
1258 I: IntoIterator<Item = Datum<'a>>,
1259{
1260 let has_last_value = funcs
1261 .iter()
1262 .any(|f| matches!(f, AggregateFunc::LastValue { .. }));
1263
1264 let input_datums_with_ranks = order_aggregate_datums_with_rank(input_datums, order_by);
1265
1266 let size_hint = input_datums_with_ranks.size_hint().0;
1267 let mut encoded_argsss = vec![Vec::with_capacity(size_hint); funcs.len()];
1268 let mut original_rows = Vec::with_capacity(size_hint);
1269 let mut order_by_rows = Vec::with_capacity(size_hint);
1270 for (d, order_by_row) in input_datums_with_ranks {
1271 let mut iter = d.unwrap_list().iter();
1272 let original_row = iter.next().unwrap();
1273 original_rows.push(original_row);
1274 let mut argss_iter = iter.next().unwrap().unwrap_list().iter();
1275 for i in 0..funcs.len() {
1276 let encoded_args = argss_iter.next().unwrap();
1277 encoded_argsss[i].push(encoded_args);
1278 }
1279 if has_last_value {
1280 order_by_rows.push(order_by_row);
1281 }
1282 }
1283
1284 let mut results_per_row = vec![Vec::with_capacity(funcs.len()); original_rows.len()];
1285 for (func, encoded_argss) in funcs.iter().zip_eq(encoded_argsss) {
1286 let results = match func {
1287 AggregateFunc::LagLead {
1288 order_by: inner_order_by,
1289 lag_lead,
1290 ignore_nulls,
1291 } => {
1292 assert_eq!(order_by, inner_order_by);
1293 let unwrapped_argss = encoded_argss
1294 .into_iter()
1295 .map(|encoded_args| unwrap_lag_lead_encoded_args(encoded_args))
1296 .collect();
1297 lag_lead_inner(unwrapped_argss, lag_lead, ignore_nulls)
1298 }
1299 AggregateFunc::FirstValue {
1300 order_by: inner_order_by,
1301 window_frame,
1302 } => {
1303 assert_eq!(order_by, inner_order_by);
1304 first_value_inner(encoded_argss, window_frame)
1307 }
1308 AggregateFunc::LastValue {
1309 order_by: inner_order_by,
1310 window_frame,
1311 } => {
1312 assert_eq!(order_by, inner_order_by);
1313 last_value_inner(encoded_argss, &order_by_rows, window_frame)
1316 }
1317 _ => panic!("unknown window function in FusedValueWindowFunc"),
1318 };
1319 for (results, result) in results_per_row.iter_mut().zip_eq(results) {
1320 results.push(result);
1321 }
1322 }
1323
1324 callers_temp_storage.reserve(2 * original_rows.len());
1325 results_per_row
1326 .into_iter()
1327 .enumerate()
1328 .map(move |(i, results)| {
1329 callers_temp_storage.make_datum(|packer| {
1330 packer.push_list_with(|packer| {
1331 packer
1332 .push(callers_temp_storage.make_datum(|packer| packer.push_list(results)));
1333 packer.push(original_rows[i]);
1334 });
1335 })
1336 })
1337}
1338
1339fn window_aggr<'a, I, A>(
1348 input_datums: I,
1349 callers_temp_storage: &'a RowArena,
1350 wrapped_aggregate: &AggregateFunc,
1351 order_by: &[ColumnOrder],
1352 window_frame: &WindowFrame,
1353) -> Datum<'a>
1354where
1355 I: IntoIterator<Item = Datum<'a>>,
1356 A: OneByOneAggr,
1357{
1358 let temp_storage = RowArena::new();
1359 let iter = window_aggr_no_list::<I, A>(
1360 input_datums,
1361 &temp_storage,
1362 wrapped_aggregate,
1363 order_by,
1364 window_frame,
1365 );
1366 callers_temp_storage.make_datum(|packer| {
1367 packer.push_list(iter);
1368 })
1369}
1370
1371fn window_aggr_no_list<'a: 'b, 'b, I, A>(
1374 input_datums: I,
1375 callers_temp_storage: &'b RowArena,
1376 wrapped_aggregate: &AggregateFunc,
1377 order_by: &[ColumnOrder],
1378 window_frame: &WindowFrame,
1379) -> impl Iterator<Item = Datum<'b>>
1380where
1381 I: IntoIterator<Item = Datum<'a>>,
1382 A: OneByOneAggr,
1383{
1384 let datums = order_aggregate_datums_with_rank(input_datums, order_by);
1387
1388 let size_hint = datums.size_hint().0;
1390 let mut args: Vec<Datum> = Vec::with_capacity(size_hint);
1391 let mut original_rows: Vec<Datum> = Vec::with_capacity(size_hint);
1392 let mut order_by_rows = Vec::with_capacity(size_hint);
1393 for (d, order_by_row) in datums.into_iter() {
1394 let mut iter = d.unwrap_list().iter();
1395 let original_row = iter.next().unwrap();
1396 let arg = iter.next().unwrap();
1397 order_by_rows.push(order_by_row);
1398 original_rows.push(original_row);
1399 args.push(arg);
1400 }
1401
1402 let results = window_aggr_inner::<A>(
1403 args,
1404 &order_by_rows,
1405 wrapped_aggregate,
1406 order_by,
1407 window_frame,
1408 callers_temp_storage,
1409 );
1410
1411 callers_temp_storage.reserve(results.len());
1412 results
1413 .into_iter()
1414 .zip_eq(original_rows)
1415 .map(|(result_value, original_row)| {
1416 callers_temp_storage.make_datum(|packer| {
1417 packer.push_list_with(|packer| {
1418 packer.push(result_value);
1419 packer.push(original_row);
1420 });
1421 })
1422 })
1423}
1424
1425fn window_aggr_inner<'a, A>(
1426 mut args: Vec<Datum<'a>>,
1427 order_by_rows: &Vec<Row>,
1428 wrapped_aggregate: &AggregateFunc,
1429 order_by: &[ColumnOrder],
1430 window_frame: &WindowFrame,
1431 temp_storage: &'a RowArena,
1432) -> Vec<Datum<'a>>
1433where
1434 A: OneByOneAggr,
1435{
1436 let length = args.len();
1437 let mut result: Vec<Datum> = Vec::with_capacity(length);
1438
1439 soft_assert_or_log!(
1445 !((matches!(window_frame.units, WindowFrameUnits::Groups)
1446 || matches!(window_frame.units, WindowFrameUnits::Range))
1447 && !window_frame.includes_current_row()),
1448 "window frame without current row"
1449 );
1450
1451 if (matches!(
1452 window_frame.start_bound,
1453 WindowFrameBound::UnboundedPreceding
1454 ) && matches!(window_frame.end_bound, WindowFrameBound::UnboundedFollowing))
1455 || (order_by.is_empty()
1456 && (matches!(window_frame.units, WindowFrameUnits::Groups)
1457 || matches!(window_frame.units, WindowFrameUnits::Range))
1458 && window_frame.includes_current_row())
1459 {
1460 let result_value =
1467 wrapped_aggregate.eval(args.into_iter().map(|d| (d, Diff::ONE)), temp_storage);
1468 for _ in 0..length {
1470 result.push(result_value);
1471 }
1472 } else {
1473 fn rows_between_unbounded_preceding_and_current_row<'a, A>(
1474 args: Vec<Datum<'a>>,
1475 result: &mut Vec<Datum<'a>>,
1476 mut one_by_one_aggr: A,
1477 temp_storage: &'a RowArena,
1478 ) where
1479 A: OneByOneAggr,
1480 {
1481 for current_arg in args.into_iter() {
1482 one_by_one_aggr.give(¤t_arg);
1483 let result_value = one_by_one_aggr.get_current_aggregate(temp_storage);
1484 result.push(result_value);
1485 }
1486 }
1487
1488 fn groups_between_unbounded_preceding_and_current_row<'a, A>(
1489 args: Vec<Datum<'a>>,
1490 order_by_rows: &Vec<Row>,
1491 result: &mut Vec<Datum<'a>>,
1492 mut one_by_one_aggr: A,
1493 temp_storage: &'a RowArena,
1494 ) where
1495 A: OneByOneAggr,
1496 {
1497 let mut peer_group_start = 0;
1498 while peer_group_start < args.len() {
1499 let mut peer_group_end = peer_group_start + 1;
1503 while peer_group_end < args.len()
1504 && order_by_rows[peer_group_start] == order_by_rows[peer_group_end]
1505 {
1506 peer_group_end += 1;
1508 }
1509 for current_arg in args[peer_group_start..peer_group_end].iter() {
1512 one_by_one_aggr.give(current_arg);
1513 }
1514 let agg_for_peer_group = one_by_one_aggr.get_current_aggregate(temp_storage);
1515 for _ in args[peer_group_start..peer_group_end].iter() {
1517 result.push(agg_for_peer_group);
1518 }
1519 peer_group_start = peer_group_end;
1521 }
1522 }
1523
1524 fn rows_between_offset_and_offset<'a>(
1525 args: Vec<Datum<'a>>,
1526 result: &mut Vec<Datum<'a>>,
1527 wrapped_aggregate: &AggregateFunc,
1528 temp_storage: &'a RowArena,
1529 offset_start: i64,
1530 offset_end: i64,
1531 ) {
1532 let len = args
1533 .len()
1534 .to_i64()
1535 .expect("window partition's len should fit into i64");
1536 for i in 0..len {
1537 let i = i.to_i64().expect("window partition shouldn't be super big");
1538 let frame_start = max(i + offset_start, 0)
1541 .to_usize()
1542 .expect("The max made sure it's not negative");
1543 let frame_end = min(i + offset_end, len - 1).to_usize();
1546 match frame_end {
1547 Some(frame_end) => {
1548 if frame_start <= frame_end {
1549 let frame_values = args[frame_start..=frame_end]
1560 .iter()
1561 .map(|d| (*d, Diff::ONE));
1562 let result_value = wrapped_aggregate.eval(frame_values, temp_storage);
1563 result.push(result_value);
1564 } else {
1565 let result_value = wrapped_aggregate.default();
1567 result.push(result_value);
1568 }
1569 }
1570 None => {
1571 let result_value = wrapped_aggregate.default();
1573 result.push(result_value);
1574 }
1575 }
1576 }
1577 }
1578
1579 match (
1580 &window_frame.units,
1581 &window_frame.start_bound,
1582 &window_frame.end_bound,
1583 ) {
1584 (Rows, UnboundedPreceding, CurrentRow) => {
1589 rows_between_unbounded_preceding_and_current_row::<A>(
1590 args,
1591 &mut result,
1592 A::new(wrapped_aggregate, false),
1593 temp_storage,
1594 );
1595 }
1596 (Rows, CurrentRow, UnboundedFollowing) => {
1597 args.reverse();
1599 rows_between_unbounded_preceding_and_current_row::<A>(
1600 args,
1601 &mut result,
1602 A::new(wrapped_aggregate, true),
1603 temp_storage,
1604 );
1605 result.reverse();
1606 }
1607 (Range, UnboundedPreceding, CurrentRow) => {
1608 groups_between_unbounded_preceding_and_current_row::<A>(
1611 args,
1612 order_by_rows,
1613 &mut result,
1614 A::new(wrapped_aggregate, false),
1615 temp_storage,
1616 );
1617 }
1618 (Rows, OffsetPreceding(start_prec), OffsetPreceding(end_prec)) => {
1622 let start_prec = start_prec.to_i64().expect(
1623 "window frame start OFFSET shouldn't be super big (the planning ensured this)",
1624 );
1625 let end_prec = end_prec.to_i64().expect(
1626 "window frame end OFFSET shouldn't be super big (the planning ensured this)",
1627 );
1628 rows_between_offset_and_offset(
1629 args,
1630 &mut result,
1631 wrapped_aggregate,
1632 temp_storage,
1633 -start_prec,
1634 -end_prec,
1635 );
1636 }
1637 (Rows, OffsetPreceding(start_prec), OffsetFollowing(end_fol)) => {
1638 let start_prec = start_prec.to_i64().expect(
1639 "window frame start OFFSET shouldn't be super big (the planning ensured this)",
1640 );
1641 let end_fol = end_fol.to_i64().expect(
1642 "window frame end OFFSET shouldn't be super big (the planning ensured this)",
1643 );
1644 rows_between_offset_and_offset(
1645 args,
1646 &mut result,
1647 wrapped_aggregate,
1648 temp_storage,
1649 -start_prec,
1650 end_fol,
1651 );
1652 }
1653 (Rows, OffsetFollowing(start_fol), OffsetFollowing(end_fol)) => {
1654 let start_fol = start_fol.to_i64().expect(
1655 "window frame start OFFSET shouldn't be super big (the planning ensured this)",
1656 );
1657 let end_fol = end_fol.to_i64().expect(
1658 "window frame end OFFSET shouldn't be super big (the planning ensured this)",
1659 );
1660 rows_between_offset_and_offset(
1661 args,
1662 &mut result,
1663 wrapped_aggregate,
1664 temp_storage,
1665 start_fol,
1666 end_fol,
1667 );
1668 }
1669 (Rows, OffsetFollowing(_), OffsetPreceding(_)) => {
1670 unreachable!() }
1672 (Rows, OffsetPreceding(start_prec), CurrentRow) => {
1673 let start_prec = start_prec.to_i64().expect(
1674 "window frame start OFFSET shouldn't be super big (the planning ensured this)",
1675 );
1676 let end_fol = 0;
1677 rows_between_offset_and_offset(
1678 args,
1679 &mut result,
1680 wrapped_aggregate,
1681 temp_storage,
1682 -start_prec,
1683 end_fol,
1684 );
1685 }
1686 (Rows, CurrentRow, OffsetFollowing(end_fol)) => {
1687 let start_fol = 0;
1688 let end_fol = end_fol.to_i64().expect(
1689 "window frame end OFFSET shouldn't be super big (the planning ensured this)",
1690 );
1691 rows_between_offset_and_offset(
1692 args,
1693 &mut result,
1694 wrapped_aggregate,
1695 temp_storage,
1696 start_fol,
1697 end_fol,
1698 );
1699 }
1700 (Rows, CurrentRow, CurrentRow) => {
1701 let start_fol = 0;
1704 let end_fol = 0;
1705 rows_between_offset_and_offset(
1706 args,
1707 &mut result,
1708 wrapped_aggregate,
1709 temp_storage,
1710 start_fol,
1711 end_fol,
1712 );
1713 }
1714 (Rows, CurrentRow, OffsetPreceding(_))
1715 | (Rows, UnboundedFollowing, _)
1716 | (Rows, _, UnboundedPreceding)
1717 | (Rows, OffsetFollowing(..), CurrentRow) => {
1718 unreachable!() }
1720 (Rows, UnboundedPreceding, UnboundedFollowing) => {
1721 unreachable!()
1724 }
1725 (Rows, UnboundedPreceding, OffsetPreceding(_))
1726 | (Rows, UnboundedPreceding, OffsetFollowing(_))
1727 | (Rows, OffsetPreceding(..), UnboundedFollowing)
1728 | (Rows, OffsetFollowing(..), UnboundedFollowing) => {
1729 unreachable!()
1732 }
1733 (Range, _, _) => {
1734 unreachable!()
1741 }
1742 (Groups, _, _) => {
1743 unreachable!()
1747 }
1748 }
1749 }
1750
1751 result
1752}
1753
1754fn fused_window_aggr<'a, I, A>(
1758 input_datums: I,
1759 callers_temp_storage: &'a RowArena,
1760 wrapped_aggregates: &Vec<AggregateFunc>,
1761 order_by: &Vec<ColumnOrder>,
1762 window_frame: &WindowFrame,
1763) -> Datum<'a>
1764where
1765 I: IntoIterator<Item = Datum<'a>>,
1766 A: OneByOneAggr,
1767{
1768 let temp_storage = RowArena::new();
1769 let iter = fused_window_aggr_no_list::<_, A>(
1770 input_datums,
1771 &temp_storage,
1772 wrapped_aggregates,
1773 order_by,
1774 window_frame,
1775 );
1776 callers_temp_storage.make_datum(|packer| {
1777 packer.push_list(iter);
1778 })
1779}
1780
1781fn fused_window_aggr_no_list<'a: 'b, 'b, I, A>(
1784 input_datums: I,
1785 callers_temp_storage: &'b RowArena,
1786 wrapped_aggregates: &Vec<AggregateFunc>,
1787 order_by: &Vec<ColumnOrder>,
1788 window_frame: &WindowFrame,
1789) -> impl Iterator<Item = Datum<'b>>
1790where
1791 I: IntoIterator<Item = Datum<'a>>,
1792 A: OneByOneAggr,
1793{
1794 let datums = order_aggregate_datums_with_rank(input_datums, order_by);
1797
1798 let size_hint = datums.size_hint().0;
1799 let mut argss = vec![Vec::with_capacity(size_hint); wrapped_aggregates.len()];
1800 let mut original_rows = Vec::with_capacity(size_hint);
1801 let mut order_by_rows = Vec::with_capacity(size_hint);
1802 for (d, order_by_row) in datums {
1803 let mut iter = d.unwrap_list().iter();
1804 let original_row = iter.next().unwrap();
1805 original_rows.push(original_row);
1806 let args_iter = iter.next().unwrap().unwrap_list().iter();
1807 for (args, arg) in argss.iter_mut().zip_eq(args_iter) {
1809 args.push(arg);
1810 }
1811 order_by_rows.push(order_by_row);
1812 }
1813
1814 let mut results_per_row =
1815 vec![Vec::with_capacity(wrapped_aggregates.len()); original_rows.len()];
1816 for (wrapped_aggr, args) in wrapped_aggregates.iter().zip_eq(argss) {
1817 let results = window_aggr_inner::<A>(
1818 args,
1819 &order_by_rows,
1820 wrapped_aggr,
1821 order_by,
1822 window_frame,
1823 callers_temp_storage,
1824 );
1825 for (results, result) in results_per_row.iter_mut().zip_eq(results) {
1826 results.push(result);
1827 }
1828 }
1829
1830 callers_temp_storage.reserve(2 * original_rows.len());
1831 results_per_row
1832 .into_iter()
1833 .enumerate()
1834 .map(move |(i, results)| {
1835 callers_temp_storage.make_datum(|packer| {
1836 packer.push_list_with(|packer| {
1837 packer
1838 .push(callers_temp_storage.make_datum(|packer| packer.push_list(results)));
1839 packer.push(original_rows[i]);
1840 });
1841 })
1842 })
1843}
1844
1845pub trait OneByOneAggr {
1849 fn new(agg: &AggregateFunc, reverse: bool) -> Self;
1854 fn give(&mut self, d: &Datum);
1856 fn get_current_aggregate<'a>(&self, temp_storage: &'a RowArena) -> Datum<'a>;
1858}
1859
1860#[derive(Debug)]
1866pub struct NaiveOneByOneAggr {
1867 agg: AggregateFunc,
1868 input: Vec<Row>,
1869 reverse: bool,
1870}
1871
1872impl OneByOneAggr for NaiveOneByOneAggr {
1873 fn new(agg: &AggregateFunc, reverse: bool) -> Self {
1874 NaiveOneByOneAggr {
1875 agg: agg.clone(),
1876 input: Vec::new(),
1877 reverse,
1878 }
1879 }
1880
1881 fn give(&mut self, d: &Datum) {
1882 let mut row = Row::default();
1883 row.packer().push(d);
1884 self.input.push(row);
1885 }
1886
1887 fn get_current_aggregate<'a>(&self, temp_storage: &'a RowArena) -> Datum<'a> {
1888 temp_storage.make_datum(|packer| {
1889 packer.push(if !self.reverse {
1890 self.agg.eval(
1891 self.input.iter().map(|r| (r.unpack_first(), Diff::ONE)),
1892 temp_storage,
1893 )
1894 } else {
1895 self.agg.eval(
1896 self.input
1897 .iter()
1898 .rev()
1899 .map(|r| (r.unpack_first(), Diff::ONE)),
1900 temp_storage,
1901 )
1902 });
1903 })
1904 }
1905}
1906
1907#[derive(
1910 Clone,
1911 Debug,
1912 Eq,
1913 PartialEq,
1914 Ord,
1915 PartialOrd,
1916 Serialize,
1917 Deserialize,
1918 Hash
1919)]
1920pub enum LagLeadType {
1921 Lag,
1922 Lead,
1923}
1924
1925#[derive(
1926 Clone,
1927 Debug,
1928 Eq,
1929 PartialEq,
1930 Ord,
1931 PartialOrd,
1932 Serialize,
1933 Deserialize,
1934 Hash
1935)]
1936pub enum AggregateFunc {
1937 MaxNumeric,
1938 MaxInt16,
1939 MaxInt32,
1940 MaxInt64,
1941 MaxUInt16,
1942 MaxUInt32,
1943 MaxUInt64,
1944 MaxMzTimestamp,
1945 MaxFloat32,
1946 MaxFloat64,
1947 MaxBool,
1948 MaxString,
1949 MaxDate,
1950 MaxTimestamp,
1951 MaxTimestampTz,
1952 MaxInterval,
1953 MaxTime,
1954 MaxArray,
1955 MaxList,
1956 MinNumeric,
1957 MinInt16,
1958 MinInt32,
1959 MinInt64,
1960 MinUInt16,
1961 MinUInt32,
1962 MinUInt64,
1963 MinMzTimestamp,
1964 MinFloat32,
1965 MinFloat64,
1966 MinBool,
1967 MinString,
1968 MinDate,
1969 MinTimestamp,
1970 MinTimestampTz,
1971 MinInterval,
1972 MinTime,
1973 MinArray,
1974 MinList,
1975 SumInt16,
1976 SumInt32,
1977 SumInt64,
1978 SumUInt16,
1979 SumUInt32,
1980 SumUInt64,
1981 SumFloat32,
1982 SumFloat64,
1983 SumNumeric,
1984 SumInterval,
1985 Count,
1986 Any,
1987 All,
1988 JsonbAgg {
1995 order_by: Vec<ColumnOrder>,
1996 },
1997 JsonbObjectAgg {
2004 order_by: Vec<ColumnOrder>,
2005 },
2006 MapAgg {
2010 order_by: Vec<ColumnOrder>,
2011 value_type: SqlScalarType,
2012 },
2013 ArrayConcat {
2016 order_by: Vec<ColumnOrder>,
2017 },
2018 ListConcat {
2021 order_by: Vec<ColumnOrder>,
2022 },
2023 StringAgg {
2024 order_by: Vec<ColumnOrder>,
2025 },
2026 RowNumber {
2027 order_by: Vec<ColumnOrder>,
2028 },
2029 Rank {
2030 order_by: Vec<ColumnOrder>,
2031 },
2032 DenseRank {
2033 order_by: Vec<ColumnOrder>,
2034 },
2035 LagLead {
2036 order_by: Vec<ColumnOrder>,
2037 lag_lead: LagLeadType,
2038 ignore_nulls: bool,
2039 },
2040 FirstValue {
2041 order_by: Vec<ColumnOrder>,
2042 window_frame: WindowFrame,
2043 },
2044 LastValue {
2045 order_by: Vec<ColumnOrder>,
2046 window_frame: WindowFrame,
2047 },
2048 FusedValueWindowFunc {
2050 funcs: Vec<AggregateFunc>,
2051 order_by: Vec<ColumnOrder>,
2054 },
2055 WindowAggregate {
2056 wrapped_aggregate: Box<AggregateFunc>,
2057 order_by: Vec<ColumnOrder>,
2058 window_frame: WindowFrame,
2059 },
2060 FusedWindowAggregate {
2061 wrapped_aggregates: Vec<AggregateFunc>,
2062 order_by: Vec<ColumnOrder>,
2063 window_frame: WindowFrame,
2064 },
2065 Dummy,
2070}
2071
2072fn expand_counts<'a, I>(datums: I) -> impl Iterator<Item = Datum<'a>>
2078where
2079 I: IntoIterator<Item = (Datum<'a>, Diff)>,
2080{
2081 datums.into_iter().flat_map(|(datum, diff)| {
2082 let copies = usize::try_from(diff.into_inner()).unwrap_or(0);
2083 std::iter::repeat(datum).take(copies)
2084 })
2085}
2086
2087impl AggregateFunc {
2088 fn ignores_multiplicity(&self) -> bool {
2095 use AggregateFunc::*;
2096 matches!(
2097 self,
2098 MaxNumeric
2099 | MaxInt16
2100 | MaxInt32
2101 | MaxInt64
2102 | MaxUInt16
2103 | MaxUInt32
2104 | MaxUInt64
2105 | MaxMzTimestamp
2106 | MaxFloat32
2107 | MaxFloat64
2108 | MaxBool
2109 | MaxString
2110 | MaxDate
2111 | MaxTimestamp
2112 | MaxTimestampTz
2113 | MaxInterval
2114 | MaxTime
2115 | MaxArray
2116 | MaxList
2117 | MinNumeric
2118 | MinInt16
2119 | MinInt32
2120 | MinInt64
2121 | MinUInt16
2122 | MinUInt32
2123 | MinUInt64
2124 | MinMzTimestamp
2125 | MinFloat32
2126 | MinFloat64
2127 | MinBool
2128 | MinString
2129 | MinDate
2130 | MinTimestamp
2131 | MinTimestampTz
2132 | MinInterval
2133 | MinTime
2134 | MinArray
2135 | MinList
2136 | Any
2137 | All
2138 )
2139 }
2140
2141 pub fn eval<'a, I>(&self, datums: I, temp_storage: &'a RowArena) -> Datum<'a>
2148 where
2149 I: IntoIterator<Item = (Datum<'a>, Diff)>,
2150 {
2151 match self {
2162 AggregateFunc::Count => count(datums),
2163 AggregateFunc::SumInt16 | AggregateFunc::SumInt32 => {
2164 sum_signed_int_counted(datums, |accum| {
2166 #[allow(clippy::as_conversions)]
2167 let narrowed = accum as i64;
2168 Datum::Int64(narrowed)
2169 })
2170 }
2171 AggregateFunc::SumInt64 => sum_signed_int_counted(datums, Datum::from),
2172 AggregateFunc::SumInterval => sum_interval_counted(datums),
2173 _ if self.ignores_multiplicity() => {
2174 self.eval_datums(datums.into_iter().map(|(datum, _diff)| datum), temp_storage)
2175 }
2176 _ => self.eval_datums(expand_counts(datums), temp_storage),
2177 }
2178 }
2179
2180 fn eval_datums<'a, I>(&self, datums: I, temp_storage: &'a RowArena) -> Datum<'a>
2182 where
2183 I: IntoIterator<Item = Datum<'a>>,
2184 {
2185 match self {
2186 AggregateFunc::MaxNumeric => {
2187 max_datum::<'a, I, OrderedDecimal<numeric::Numeric>>(datums)
2188 }
2189 AggregateFunc::MaxInt16 => max_datum::<'a, I, i16>(datums),
2190 AggregateFunc::MaxInt32 => max_datum::<'a, I, i32>(datums),
2191 AggregateFunc::MaxInt64 => max_datum::<'a, I, i64>(datums),
2192 AggregateFunc::MaxUInt16 => max_datum::<'a, I, u16>(datums),
2193 AggregateFunc::MaxUInt32 => max_datum::<'a, I, u32>(datums),
2194 AggregateFunc::MaxUInt64 => max_datum::<'a, I, u64>(datums),
2195 AggregateFunc::MaxMzTimestamp => max_datum::<'a, I, mz_repr::Timestamp>(datums),
2196 AggregateFunc::MaxFloat32 => max_datum::<'a, I, OrderedFloat<f32>>(datums),
2197 AggregateFunc::MaxFloat64 => max_datum::<'a, I, OrderedFloat<f64>>(datums),
2198 AggregateFunc::MaxBool => max_datum::<'a, I, bool>(datums),
2199 AggregateFunc::MaxString => max_string(datums),
2200 AggregateFunc::MaxDate => max_datum::<'a, I, Date>(datums),
2201 AggregateFunc::MaxTimestamp => {
2202 max_datum::<'a, I, CheckedTimestamp<NaiveDateTime>>(datums)
2203 }
2204 AggregateFunc::MaxTimestampTz => {
2205 max_datum::<'a, I, CheckedTimestamp<DateTime<Utc>>>(datums)
2206 }
2207 AggregateFunc::MaxInterval => max_datum::<'a, I, Interval>(datums),
2208 AggregateFunc::MaxTime => max_datum::<'a, I, NaiveTime>(datums),
2209 AggregateFunc::MaxArray | AggregateFunc::MaxList => {
2210 max_datum::<'a, I, Datum<'a>>(datums)
2211 }
2212 AggregateFunc::MinNumeric => {
2213 min_datum::<'a, I, OrderedDecimal<numeric::Numeric>>(datums)
2214 }
2215 AggregateFunc::MinInt16 => min_datum::<'a, I, i16>(datums),
2216 AggregateFunc::MinInt32 => min_datum::<'a, I, i32>(datums),
2217 AggregateFunc::MinInt64 => min_datum::<'a, I, i64>(datums),
2218 AggregateFunc::MinUInt16 => min_datum::<'a, I, u16>(datums),
2219 AggregateFunc::MinUInt32 => min_datum::<'a, I, u32>(datums),
2220 AggregateFunc::MinUInt64 => min_datum::<'a, I, u64>(datums),
2221 AggregateFunc::MinMzTimestamp => min_datum::<'a, I, mz_repr::Timestamp>(datums),
2222 AggregateFunc::MinFloat32 => min_datum::<'a, I, OrderedFloat<f32>>(datums),
2223 AggregateFunc::MinFloat64 => min_datum::<'a, I, OrderedFloat<f64>>(datums),
2224 AggregateFunc::MinBool => min_datum::<'a, I, bool>(datums),
2225 AggregateFunc::MinString => min_string(datums),
2226 AggregateFunc::MinDate => min_datum::<'a, I, Date>(datums),
2227 AggregateFunc::MinTimestamp => {
2228 min_datum::<'a, I, CheckedTimestamp<NaiveDateTime>>(datums)
2229 }
2230 AggregateFunc::MinTimestampTz => {
2231 min_datum::<'a, I, CheckedTimestamp<DateTime<Utc>>>(datums)
2232 }
2233 AggregateFunc::MinInterval => min_datum::<'a, I, Interval>(datums),
2234 AggregateFunc::MinTime => min_datum::<'a, I, NaiveTime>(datums),
2235 AggregateFunc::MinArray | AggregateFunc::MinList => {
2236 min_datum::<'a, I, Datum<'a>>(datums)
2237 }
2238 AggregateFunc::SumInt16 => sum_datum::<'a, I, i16, i64>(datums),
2239 AggregateFunc::SumInt32 => sum_datum::<'a, I, i32, i64>(datums),
2240 AggregateFunc::SumInt64 => sum_datum::<'a, I, i64, i128>(datums),
2241 AggregateFunc::SumUInt16 => sum_datum::<'a, I, u16, u64>(datums),
2242 AggregateFunc::SumUInt32 => sum_datum::<'a, I, u32, u64>(datums),
2243 AggregateFunc::SumUInt64 => sum_datum::<'a, I, u64, u128>(datums),
2244 AggregateFunc::SumFloat32 => sum_datum::<'a, I, f32, f32>(datums),
2245 AggregateFunc::SumFloat64 => sum_datum::<'a, I, f64, f64>(datums),
2246 AggregateFunc::SumNumeric => sum_numeric(datums),
2247 AggregateFunc::SumInterval => sum_interval(datums),
2248 AggregateFunc::Count => unreachable!("Count is handled in `eval`"),
2249 AggregateFunc::Any => any(datums),
2250 AggregateFunc::All => all(datums),
2251 AggregateFunc::JsonbAgg { order_by } => jsonb_agg(datums, temp_storage, order_by),
2252 AggregateFunc::MapAgg { order_by, .. } | AggregateFunc::JsonbObjectAgg { order_by } => {
2253 dict_agg(datums, temp_storage, order_by)
2254 }
2255 AggregateFunc::ArrayConcat { order_by } => array_concat(datums, temp_storage, order_by),
2256 AggregateFunc::ListConcat { order_by } => list_concat(datums, temp_storage, order_by),
2257 AggregateFunc::StringAgg { order_by } => string_agg(datums, temp_storage, order_by),
2258 AggregateFunc::RowNumber { order_by } => row_number(datums, temp_storage, order_by),
2259 AggregateFunc::Rank { order_by } => rank(datums, temp_storage, order_by),
2260 AggregateFunc::DenseRank { order_by } => dense_rank(datums, temp_storage, order_by),
2261 AggregateFunc::LagLead {
2262 order_by,
2263 lag_lead: lag_lead_type,
2264 ignore_nulls,
2265 } => lag_lead(datums, temp_storage, order_by, lag_lead_type, ignore_nulls),
2266 AggregateFunc::FirstValue {
2267 order_by,
2268 window_frame,
2269 } => first_value(datums, temp_storage, order_by, window_frame),
2270 AggregateFunc::LastValue {
2271 order_by,
2272 window_frame,
2273 } => last_value(datums, temp_storage, order_by, window_frame),
2274 AggregateFunc::WindowAggregate {
2275 wrapped_aggregate,
2276 order_by,
2277 window_frame,
2278 } => window_aggr::<_, NaiveOneByOneAggr>(
2279 datums,
2280 temp_storage,
2281 wrapped_aggregate,
2282 order_by,
2283 window_frame,
2284 ),
2285 AggregateFunc::FusedValueWindowFunc { funcs, order_by } => {
2286 fused_value_window_func(datums, temp_storage, funcs, order_by)
2287 }
2288 AggregateFunc::FusedWindowAggregate {
2289 wrapped_aggregates,
2290 order_by,
2291 window_frame,
2292 } => fused_window_aggr::<_, NaiveOneByOneAggr>(
2293 datums,
2294 temp_storage,
2295 wrapped_aggregates,
2296 order_by,
2297 window_frame,
2298 ),
2299 AggregateFunc::Dummy => Datum::Dummy,
2300 }
2301 }
2302
2303 pub fn eval_with_fast_window_agg<'a, I, W>(
2307 &self,
2308 datums: I,
2309 temp_storage: &'a RowArena,
2310 ) -> Datum<'a>
2311 where
2312 I: IntoIterator<Item = (Datum<'a>, Diff)>,
2313 W: OneByOneAggr,
2314 {
2315 match self {
2316 AggregateFunc::WindowAggregate {
2317 wrapped_aggregate,
2318 order_by,
2319 window_frame,
2320 } => window_aggr::<_, W>(
2321 expand_counts(datums),
2322 temp_storage,
2323 wrapped_aggregate,
2324 order_by,
2325 window_frame,
2326 ),
2327 AggregateFunc::FusedWindowAggregate {
2328 wrapped_aggregates,
2329 order_by,
2330 window_frame,
2331 } => fused_window_aggr::<_, W>(
2332 expand_counts(datums),
2333 temp_storage,
2334 wrapped_aggregates,
2335 order_by,
2336 window_frame,
2337 ),
2338 _ => self.eval(datums, temp_storage),
2339 }
2340 }
2341
2342 pub fn eval_with_unnest_list<'a, I, W>(
2343 &self,
2344 datums: I,
2345 temp_storage: &'a RowArena,
2346 ) -> impl Iterator<Item = Datum<'a>>
2347 where
2348 I: IntoIterator<Item = (Datum<'a>, Diff)>,
2349 W: OneByOneAggr,
2350 {
2351 assert!(self.can_fuse_with_unnest_list());
2353 let datums = expand_counts(datums);
2355 match self {
2356 AggregateFunc::RowNumber { order_by } => {
2357 row_number_no_list(datums, temp_storage, order_by).collect_vec()
2358 }
2359 AggregateFunc::Rank { order_by } => {
2360 rank_no_list(datums, temp_storage, order_by).collect_vec()
2361 }
2362 AggregateFunc::DenseRank { order_by } => {
2363 dense_rank_no_list(datums, temp_storage, order_by).collect_vec()
2364 }
2365 AggregateFunc::LagLead {
2366 order_by,
2367 lag_lead: lag_lead_type,
2368 ignore_nulls,
2369 } => lag_lead_no_list(datums, temp_storage, order_by, lag_lead_type, ignore_nulls)
2370 .collect_vec(),
2371 AggregateFunc::FirstValue {
2372 order_by,
2373 window_frame,
2374 } => first_value_no_list(datums, temp_storage, order_by, window_frame).collect_vec(),
2375 AggregateFunc::LastValue {
2376 order_by,
2377 window_frame,
2378 } => last_value_no_list(datums, temp_storage, order_by, window_frame).collect_vec(),
2379 AggregateFunc::FusedValueWindowFunc { funcs, order_by } => {
2380 fused_value_window_func_no_list(datums, temp_storage, funcs, order_by).collect_vec()
2381 }
2382 AggregateFunc::WindowAggregate {
2383 wrapped_aggregate,
2384 order_by,
2385 window_frame,
2386 } => window_aggr_no_list::<_, W>(
2387 datums,
2388 temp_storage,
2389 wrapped_aggregate,
2390 order_by,
2391 window_frame,
2392 )
2393 .collect_vec(),
2394 AggregateFunc::FusedWindowAggregate {
2395 wrapped_aggregates,
2396 order_by,
2397 window_frame,
2398 } => fused_window_aggr_no_list::<_, W>(
2399 datums,
2400 temp_storage,
2401 wrapped_aggregates,
2402 order_by,
2403 window_frame,
2404 )
2405 .collect_vec(),
2406 _ => unreachable!("asserted above that `can_fuse_with_unnest_list`"),
2407 }
2408 .into_iter()
2409 }
2410
2411 pub fn default(&self) -> Datum<'static> {
2414 match self {
2415 AggregateFunc::Count => Datum::Int64(0),
2416 AggregateFunc::Any => Datum::False,
2417 AggregateFunc::All => Datum::True,
2418 AggregateFunc::Dummy => Datum::Dummy,
2419 _ => Datum::Null,
2420 }
2421 }
2422
2423 pub fn identity_datum(&self) -> Datum<'static> {
2426 match self {
2427 AggregateFunc::Any => Datum::False,
2428 AggregateFunc::All => Datum::True,
2429 AggregateFunc::Dummy => Datum::Dummy,
2430 AggregateFunc::ArrayConcat { .. } => Datum::empty_array(),
2431 AggregateFunc::ListConcat { .. } => Datum::empty_list(),
2432 AggregateFunc::RowNumber { .. }
2433 | AggregateFunc::Rank { .. }
2434 | AggregateFunc::DenseRank { .. }
2435 | AggregateFunc::LagLead { .. }
2436 | AggregateFunc::FirstValue { .. }
2437 | AggregateFunc::LastValue { .. }
2438 | AggregateFunc::WindowAggregate { .. }
2439 | AggregateFunc::FusedValueWindowFunc { .. }
2440 | AggregateFunc::FusedWindowAggregate { .. } => Datum::empty_list(),
2441 AggregateFunc::MaxNumeric
2442 | AggregateFunc::MaxInt16
2443 | AggregateFunc::MaxInt32
2444 | AggregateFunc::MaxInt64
2445 | AggregateFunc::MaxUInt16
2446 | AggregateFunc::MaxUInt32
2447 | AggregateFunc::MaxUInt64
2448 | AggregateFunc::MaxMzTimestamp
2449 | AggregateFunc::MaxFloat32
2450 | AggregateFunc::MaxFloat64
2451 | AggregateFunc::MaxBool
2452 | AggregateFunc::MaxString
2453 | AggregateFunc::MaxDate
2454 | AggregateFunc::MaxTimestamp
2455 | AggregateFunc::MaxTimestampTz
2456 | AggregateFunc::MaxInterval
2457 | AggregateFunc::MaxTime
2458 | AggregateFunc::MaxArray
2459 | AggregateFunc::MaxList
2460 | AggregateFunc::MinNumeric
2461 | AggregateFunc::MinInt16
2462 | AggregateFunc::MinInt32
2463 | AggregateFunc::MinInt64
2464 | AggregateFunc::MinUInt16
2465 | AggregateFunc::MinUInt32
2466 | AggregateFunc::MinUInt64
2467 | AggregateFunc::MinMzTimestamp
2468 | AggregateFunc::MinFloat32
2469 | AggregateFunc::MinFloat64
2470 | AggregateFunc::MinBool
2471 | AggregateFunc::MinString
2472 | AggregateFunc::MinDate
2473 | AggregateFunc::MinTimestamp
2474 | AggregateFunc::MinTimestampTz
2475 | AggregateFunc::MinInterval
2476 | AggregateFunc::MinTime
2477 | AggregateFunc::MinArray
2478 | AggregateFunc::MinList
2479 | AggregateFunc::SumInt16
2480 | AggregateFunc::SumInt32
2481 | AggregateFunc::SumInt64
2482 | AggregateFunc::SumUInt16
2483 | AggregateFunc::SumUInt32
2484 | AggregateFunc::SumUInt64
2485 | AggregateFunc::SumFloat32
2486 | AggregateFunc::SumFloat64
2487 | AggregateFunc::SumNumeric
2488 | AggregateFunc::SumInterval
2489 | AggregateFunc::Count
2490 | AggregateFunc::JsonbAgg { .. }
2491 | AggregateFunc::JsonbObjectAgg { .. }
2492 | AggregateFunc::MapAgg { .. }
2493 | AggregateFunc::StringAgg { .. } => Datum::Null,
2494 }
2495 }
2496
2497 pub fn can_fuse_with_unnest_list(&self) -> bool {
2498 match self {
2499 AggregateFunc::RowNumber { .. }
2500 | AggregateFunc::Rank { .. }
2501 | AggregateFunc::DenseRank { .. }
2502 | AggregateFunc::LagLead { .. }
2503 | AggregateFunc::FirstValue { .. }
2504 | AggregateFunc::LastValue { .. }
2505 | AggregateFunc::WindowAggregate { .. }
2506 | AggregateFunc::FusedValueWindowFunc { .. }
2507 | AggregateFunc::FusedWindowAggregate { .. } => true,
2508 AggregateFunc::ArrayConcat { .. }
2509 | AggregateFunc::ListConcat { .. }
2510 | AggregateFunc::Any
2511 | AggregateFunc::All
2512 | AggregateFunc::Dummy
2513 | AggregateFunc::MaxNumeric
2514 | AggregateFunc::MaxInt16
2515 | AggregateFunc::MaxInt32
2516 | AggregateFunc::MaxInt64
2517 | AggregateFunc::MaxUInt16
2518 | AggregateFunc::MaxUInt32
2519 | AggregateFunc::MaxUInt64
2520 | AggregateFunc::MaxMzTimestamp
2521 | AggregateFunc::MaxFloat32
2522 | AggregateFunc::MaxFloat64
2523 | AggregateFunc::MaxBool
2524 | AggregateFunc::MaxString
2525 | AggregateFunc::MaxDate
2526 | AggregateFunc::MaxTimestamp
2527 | AggregateFunc::MaxTimestampTz
2528 | AggregateFunc::MaxInterval
2529 | AggregateFunc::MaxTime
2530 | AggregateFunc::MaxArray
2531 | AggregateFunc::MaxList
2532 | AggregateFunc::MinNumeric
2533 | AggregateFunc::MinInt16
2534 | AggregateFunc::MinInt32
2535 | AggregateFunc::MinInt64
2536 | AggregateFunc::MinUInt16
2537 | AggregateFunc::MinUInt32
2538 | AggregateFunc::MinUInt64
2539 | AggregateFunc::MinMzTimestamp
2540 | AggregateFunc::MinFloat32
2541 | AggregateFunc::MinFloat64
2542 | AggregateFunc::MinBool
2543 | AggregateFunc::MinString
2544 | AggregateFunc::MinDate
2545 | AggregateFunc::MinTimestamp
2546 | AggregateFunc::MinTimestampTz
2547 | AggregateFunc::MinInterval
2548 | AggregateFunc::MinTime
2549 | AggregateFunc::MinArray
2550 | AggregateFunc::MinList
2551 | AggregateFunc::SumInt16
2552 | AggregateFunc::SumInt32
2553 | AggregateFunc::SumInt64
2554 | AggregateFunc::SumUInt16
2555 | AggregateFunc::SumUInt32
2556 | AggregateFunc::SumUInt64
2557 | AggregateFunc::SumFloat32
2558 | AggregateFunc::SumFloat64
2559 | AggregateFunc::SumNumeric
2560 | AggregateFunc::SumInterval
2561 | AggregateFunc::Count
2562 | AggregateFunc::JsonbAgg { .. }
2563 | AggregateFunc::JsonbObjectAgg { .. }
2564 | AggregateFunc::MapAgg { .. }
2565 | AggregateFunc::StringAgg { .. } => false,
2566 }
2567 }
2568
2569 pub fn output_sql_type(&self, input_type: SqlColumnType) -> SqlColumnType {
2575 let scalar_type = match self {
2576 AggregateFunc::Count => SqlScalarType::Int64,
2577 AggregateFunc::Any => SqlScalarType::Bool,
2578 AggregateFunc::All => SqlScalarType::Bool,
2579 AggregateFunc::JsonbAgg { .. } => SqlScalarType::Jsonb,
2580 AggregateFunc::JsonbObjectAgg { .. } => SqlScalarType::Jsonb,
2581 AggregateFunc::SumInt16 => SqlScalarType::Int64,
2582 AggregateFunc::SumInt32 => SqlScalarType::Int64,
2583 AggregateFunc::SumInt64 => SqlScalarType::Numeric {
2584 max_scale: Some(NumericMaxScale::ZERO),
2585 },
2586 AggregateFunc::SumUInt16 => SqlScalarType::UInt64,
2587 AggregateFunc::SumUInt32 => SqlScalarType::UInt64,
2588 AggregateFunc::SumUInt64 => SqlScalarType::Numeric {
2589 max_scale: Some(NumericMaxScale::ZERO),
2590 },
2591 AggregateFunc::MapAgg { value_type, .. } => SqlScalarType::Map {
2592 value_type: Box::new(value_type.clone()),
2593 custom_id: None,
2594 },
2595 AggregateFunc::ArrayConcat { .. } | AggregateFunc::ListConcat { .. } => {
2596 match input_type.scalar_type {
2597 SqlScalarType::Record { ref fields, .. } => fields[0].1.scalar_type.clone(),
2599 _ => unreachable!(),
2600 }
2601 }
2602 AggregateFunc::StringAgg { .. } => SqlScalarType::String,
2603 AggregateFunc::RowNumber { .. } => {
2604 AggregateFunc::output_type_ranking_window_funcs(&input_type, "?row_number?")
2605 }
2606 AggregateFunc::Rank { .. } => {
2607 AggregateFunc::output_type_ranking_window_funcs(&input_type, "?rank?")
2608 }
2609 AggregateFunc::DenseRank { .. } => {
2610 AggregateFunc::output_type_ranking_window_funcs(&input_type, "?dense_rank?")
2611 }
2612 AggregateFunc::LagLead { lag_lead: lag_lead_type, .. } => {
2613 let fields = input_type.scalar_type.unwrap_record_element_type();
2615 let original_row_type = fields[0].unwrap_record_element_type()[0]
2616 .clone()
2617 .nullable(false);
2618 let encoded_args = fields[0].unwrap_record_element_type()[1];
2619 let output_type_inner =
2620 Self::lag_lead_output_type_inner_from_encoded_args(encoded_args);
2621 let column_name = Self::lag_lead_result_column_name(lag_lead_type);
2622
2623 SqlScalarType::List {
2624 element_type: Box::new(SqlScalarType::Record {
2625 fields: [
2626 (column_name, output_type_inner),
2627 (ColumnName::from("?orig_row?"), original_row_type),
2628 ].into(),
2629 custom_id: None,
2630 }),
2631 custom_id: None,
2632 }
2633 }
2634 AggregateFunc::FirstValue { .. } => {
2635 let fields = input_type.scalar_type.unwrap_record_element_type();
2637 let original_row_type = fields[0].unwrap_record_element_type()[0]
2638 .clone()
2639 .nullable(false);
2640 let value_type = fields[0].unwrap_record_element_type()[1]
2641 .clone()
2642 .nullable(true); SqlScalarType::List {
2645 element_type: Box::new(SqlScalarType::Record {
2646 fields: [
2647 (ColumnName::from("?first_value?"), value_type),
2648 (ColumnName::from("?orig_row?"), original_row_type),
2649 ].into(),
2650 custom_id: None,
2651 }),
2652 custom_id: None,
2653 }
2654 }
2655 AggregateFunc::LastValue { .. } => {
2656 let fields = input_type.scalar_type.unwrap_record_element_type();
2658 let original_row_type = fields[0].unwrap_record_element_type()[0]
2659 .clone()
2660 .nullable(false);
2661 let value_type = fields[0].unwrap_record_element_type()[1]
2662 .clone()
2663 .nullable(true); SqlScalarType::List {
2666 element_type: Box::new(SqlScalarType::Record {
2667 fields: [
2668 (ColumnName::from("?last_value?"), value_type),
2669 (ColumnName::from("?orig_row?"), original_row_type),
2670 ].into(),
2671 custom_id: None,
2672 }),
2673 custom_id: None,
2674 }
2675 }
2676 AggregateFunc::WindowAggregate {
2677 wrapped_aggregate, ..
2678 } => {
2679 let fields = input_type.scalar_type.unwrap_record_element_type();
2681 let original_row_type = fields[0].unwrap_record_element_type()[0]
2682 .clone()
2683 .nullable(false);
2684 let arg_type = fields[0].unwrap_record_element_type()[1]
2685 .clone()
2686 .nullable(true);
2687 let wrapped_aggr_out_type = wrapped_aggregate.output_sql_type(arg_type);
2688
2689 SqlScalarType::List {
2690 element_type: Box::new(SqlScalarType::Record {
2691 fields: [
2692 (ColumnName::from("?window_agg?"), wrapped_aggr_out_type),
2693 (ColumnName::from("?orig_row?"), original_row_type),
2694 ].into(),
2695 custom_id: None,
2696 }),
2697 custom_id: None,
2698 }
2699 }
2700 AggregateFunc::FusedWindowAggregate {
2701 wrapped_aggregates, ..
2702 } => {
2703 let fields = input_type.scalar_type.unwrap_record_element_type();
2706 let original_row_type = fields[0].unwrap_record_element_type()[0]
2707 .clone()
2708 .nullable(false);
2709 let args_type = fields[0].unwrap_record_element_type()[1];
2710 let arg_types = args_type.unwrap_record_element_type();
2711 let out_fields = arg_types.iter().zip_eq(wrapped_aggregates).map(
2712 |(arg_type, wrapped_agg)| {
2713 (
2714 ColumnName::from(wrapped_agg.name()),
2715 wrapped_agg.output_sql_type((**arg_type).clone().nullable(true)),
2716 )
2717 }).collect_vec();
2718
2719 SqlScalarType::List {
2720 element_type: Box::new(SqlScalarType::Record {
2721 fields: [
2722 (ColumnName::from("?fused_window_agg?"), SqlScalarType::Record {
2723 fields: out_fields.into(),
2724 custom_id: None,
2725 }.nullable(false)),
2726 (ColumnName::from("?orig_row?"), original_row_type),
2727 ].into(),
2728 custom_id: None,
2729 }),
2730 custom_id: None,
2731 }
2732 }
2733 AggregateFunc::FusedValueWindowFunc { funcs, order_by: _ } => {
2734 let fields = input_type.scalar_type.unwrap_record_element_type();
2739 let original_row_type = fields[0].unwrap_record_element_type()[0]
2740 .clone()
2741 .nullable(false);
2742 let encoded_args_type = fields[0]
2743 .unwrap_record_element_type()[1]
2744 .unwrap_record_element_type();
2745
2746 SqlScalarType::List {
2747 element_type: Box::new(SqlScalarType::Record {
2748 fields: [
2749 (
2750 ColumnName::from("?fused_value_window_func?"),
2751 SqlScalarType::Record {
2752 fields: encoded_args_type.into_iter().zip_eq(funcs).map(
2753 |(arg_type, func)| {
2754 match func {
2755 AggregateFunc::LagLead {
2756 lag_lead: lag_lead_type, ..
2757 } => {
2758 let name = Self::lag_lead_result_column_name(
2759 lag_lead_type,
2760 );
2761 let ty = Self
2762 ::lag_lead_output_type_inner_from_encoded_args(
2763 arg_type,
2764 );
2765 (name, ty)
2766 },
2767 AggregateFunc::FirstValue { .. } => {
2768 (
2769 ColumnName::from("?first_value?"),
2770 arg_type.clone().nullable(true),
2771 )
2772 }
2773 AggregateFunc::LastValue { .. } => {
2774 (
2775 ColumnName::from("?last_value?"),
2776 arg_type.clone().nullable(true),
2777 )
2778 }
2779 _ => panic!("FusedValueWindowFunc has an unknown function"),
2780 }
2781 }).collect(),
2782 custom_id: None,
2783 }.nullable(false)),
2784 (ColumnName::from("?orig_row?"), original_row_type),
2785 ].into(),
2786 custom_id: None,
2787 }),
2788 custom_id: None,
2789 }
2790 }
2791 AggregateFunc::Dummy
2792 | AggregateFunc::MaxNumeric
2793 | AggregateFunc::MaxInt16
2794 | AggregateFunc::MaxInt32
2795 | AggregateFunc::MaxInt64
2796 | AggregateFunc::MaxUInt16
2797 | AggregateFunc::MaxUInt32
2798 | AggregateFunc::MaxUInt64
2799 | AggregateFunc::MaxMzTimestamp
2800 | AggregateFunc::MaxFloat32
2801 | AggregateFunc::MaxFloat64
2802 | AggregateFunc::MaxBool
2803 | AggregateFunc::MaxString
2807 | AggregateFunc::MaxDate
2808 | AggregateFunc::MaxTimestamp
2809 | AggregateFunc::MaxTimestampTz
2810 | AggregateFunc::MaxInterval
2811 | AggregateFunc::MaxTime
2812 | AggregateFunc::MaxArray
2813 | AggregateFunc::MaxList
2814 | AggregateFunc::MinNumeric
2815 | AggregateFunc::MinInt16
2816 | AggregateFunc::MinInt32
2817 | AggregateFunc::MinInt64
2818 | AggregateFunc::MinUInt16
2819 | AggregateFunc::MinUInt32
2820 | AggregateFunc::MinUInt64
2821 | AggregateFunc::MinMzTimestamp
2822 | AggregateFunc::MinFloat32
2823 | AggregateFunc::MinFloat64
2824 | AggregateFunc::MinBool
2825 | AggregateFunc::MinString
2826 | AggregateFunc::MinDate
2827 | AggregateFunc::MinTimestamp
2828 | AggregateFunc::MinTimestampTz
2829 | AggregateFunc::MinInterval
2830 | AggregateFunc::MinTime
2831 | AggregateFunc::MinArray
2832 | AggregateFunc::MinList
2833 | AggregateFunc::SumFloat32
2834 | AggregateFunc::SumFloat64
2835 | AggregateFunc::SumNumeric
2836 | AggregateFunc::SumInterval => input_type.scalar_type.clone(),
2837 };
2838 let nullable = match self {
2841 AggregateFunc::Count => false,
2842 AggregateFunc::StringAgg { .. } => match input_type.scalar_type {
2844 SqlScalarType::Record { fields, .. } => match &fields[0].1.scalar_type {
2846 SqlScalarType::Record { fields, .. } => fields[0].1.nullable,
2848 _ => unreachable!(),
2849 },
2850 _ => unreachable!(),
2851 },
2852 _ => input_type.nullable,
2853 };
2854 scalar_type.nullable(nullable)
2855 }
2856
2857 pub fn output_type(&self, input_type: ReprColumnType) -> ReprColumnType {
2861 ReprColumnType::from(&self.output_sql_type(SqlColumnType::from_repr(&input_type)))
2862 }
2863
2864 fn output_type_ranking_window_funcs(
2866 input_type: &SqlColumnType,
2867 col_name: &str,
2868 ) -> SqlScalarType {
2869 match input_type.scalar_type {
2870 SqlScalarType::Record { ref fields, .. } => SqlScalarType::List {
2871 element_type: Box::new(SqlScalarType::Record {
2872 fields: [
2873 (
2874 ColumnName::from(col_name),
2875 SqlScalarType::Int64.nullable(false),
2876 ),
2877 (ColumnName::from("?orig_row?"), {
2878 let inner = match &fields[0].1.scalar_type {
2879 SqlScalarType::List { element_type, .. } => element_type.clone(),
2880 _ => unreachable!(),
2881 };
2882 inner.nullable(false)
2883 }),
2884 ]
2885 .into(),
2886 custom_id: None,
2887 }),
2888 custom_id: None,
2889 },
2890 _ => unreachable!(),
2891 }
2892 }
2893
2894 fn lag_lead_output_type_inner_from_encoded_args(
2898 encoded_args_type: &SqlScalarType,
2899 ) -> SqlColumnType {
2900 encoded_args_type.unwrap_record_element_type()[0]
2904 .clone()
2905 .nullable(true)
2906 }
2907
2908 fn lag_lead_result_column_name(lag_lead_type: &LagLeadType) -> ColumnName {
2909 ColumnName::from(match lag_lead_type {
2910 LagLeadType::Lag => "?lag?",
2911 LagLeadType::Lead => "?lead?",
2912 })
2913 }
2914
2915 pub fn propagates_nonnull_constraint(&self) -> bool {
2920 match self {
2921 AggregateFunc::MaxNumeric
2922 | AggregateFunc::MaxInt16
2923 | AggregateFunc::MaxInt32
2924 | AggregateFunc::MaxInt64
2925 | AggregateFunc::MaxUInt16
2926 | AggregateFunc::MaxUInt32
2927 | AggregateFunc::MaxUInt64
2928 | AggregateFunc::MaxMzTimestamp
2929 | AggregateFunc::MaxFloat32
2930 | AggregateFunc::MaxFloat64
2931 | AggregateFunc::MaxBool
2932 | AggregateFunc::MaxString
2933 | AggregateFunc::MaxDate
2934 | AggregateFunc::MaxTimestamp
2935 | AggregateFunc::MaxTimestampTz
2936 | AggregateFunc::MaxInterval
2937 | AggregateFunc::MaxTime
2938 | AggregateFunc::MaxArray
2939 | AggregateFunc::MaxList
2940 | AggregateFunc::MinNumeric
2941 | AggregateFunc::MinInt16
2942 | AggregateFunc::MinInt32
2943 | AggregateFunc::MinInt64
2944 | AggregateFunc::MinUInt16
2945 | AggregateFunc::MinUInt32
2946 | AggregateFunc::MinUInt64
2947 | AggregateFunc::MinMzTimestamp
2948 | AggregateFunc::MinFloat32
2949 | AggregateFunc::MinFloat64
2950 | AggregateFunc::MinBool
2951 | AggregateFunc::MinString
2952 | AggregateFunc::MinDate
2953 | AggregateFunc::MinTimestamp
2954 | AggregateFunc::MinTimestampTz
2955 | AggregateFunc::MinInterval
2956 | AggregateFunc::MinTime
2957 | AggregateFunc::MinArray
2958 | AggregateFunc::MinList
2959 | AggregateFunc::SumInt16
2960 | AggregateFunc::SumInt32
2961 | AggregateFunc::SumInt64
2962 | AggregateFunc::SumUInt16
2963 | AggregateFunc::SumUInt32
2964 | AggregateFunc::SumUInt64
2965 | AggregateFunc::SumFloat32
2966 | AggregateFunc::SumFloat64
2967 | AggregateFunc::SumNumeric
2968 | AggregateFunc::SumInterval
2969 | AggregateFunc::StringAgg { .. } => true,
2970 AggregateFunc::Count
2972 | AggregateFunc::Any
2973 | AggregateFunc::All
2974 | AggregateFunc::JsonbAgg { .. }
2975 | AggregateFunc::JsonbObjectAgg { .. }
2976 | AggregateFunc::MapAgg { .. }
2977 | AggregateFunc::ArrayConcat { .. }
2978 | AggregateFunc::ListConcat { .. }
2979 | AggregateFunc::RowNumber { .. }
2980 | AggregateFunc::Rank { .. }
2981 | AggregateFunc::DenseRank { .. }
2982 | AggregateFunc::LagLead { .. }
2983 | AggregateFunc::FirstValue { .. }
2984 | AggregateFunc::LastValue { .. }
2985 | AggregateFunc::FusedValueWindowFunc { .. }
2986 | AggregateFunc::WindowAggregate { .. }
2987 | AggregateFunc::FusedWindowAggregate { .. }
2988 | AggregateFunc::Dummy => false,
2989 }
2990 }
2991}
2992
2993fn jsonb_each<'a>(a: Datum<'a>) -> impl Iterator<Item = (Row, Diff)> + 'a {
2994 let map = match a {
2996 Datum::Map(dict) => dict,
2997 _ => mz_repr::DatumMap::empty(),
2998 };
2999
3000 map.iter()
3001 .map(move |(k, v)| (Row::pack_slice(&[Datum::String(k), v]), Diff::ONE))
3002}
3003
3004fn jsonb_each_stringify<'a>(
3005 a: Datum<'a>,
3006 temp_storage: &'a RowArena,
3007) -> impl Iterator<Item = (Row, Diff)> + 'a {
3008 let map = match a {
3010 Datum::Map(dict) => dict,
3011 _ => mz_repr::DatumMap::empty(),
3012 };
3013
3014 map.iter().map(move |(k, mut v)| {
3015 v = jsonb_stringify(v, temp_storage)
3016 .map(Datum::String)
3017 .unwrap_or(Datum::Null);
3018 (Row::pack_slice(&[Datum::String(k), v]), Diff::ONE)
3019 })
3020}
3021
3022fn jsonb_object_keys<'a>(a: Datum<'a>) -> impl Iterator<Item = (Row, Diff)> + 'a {
3023 let map = match a {
3024 Datum::Map(dict) => dict,
3025 _ => mz_repr::DatumMap::empty(),
3026 };
3027
3028 map.iter()
3029 .map(move |(k, _)| (Row::pack_slice(&[Datum::String(k)]), Diff::ONE))
3030}
3031
3032fn jsonb_array_elements<'a>(a: Datum<'a>) -> impl Iterator<Item = (Row, Diff)> + 'a {
3033 let list = match a {
3034 Datum::List(list) => list,
3035 _ => mz_repr::DatumList::empty(),
3036 };
3037 list.iter().map(move |e| (Row::pack_slice(&[e]), Diff::ONE))
3038}
3039
3040fn jsonb_array_elements_stringify<'a>(
3041 a: Datum<'a>,
3042 temp_storage: &'a RowArena,
3043) -> impl Iterator<Item = (Row, Diff)> + 'a {
3044 let list = match a {
3045 Datum::List(list) => list,
3046 _ => mz_repr::DatumList::empty(),
3047 };
3048 list.iter().map(move |mut e| {
3049 e = jsonb_stringify(e, temp_storage)
3050 .map(Datum::String)
3051 .unwrap_or(Datum::Null);
3052 (Row::pack_slice(&[e]), Diff::ONE)
3053 })
3054}
3055
3056fn regexp_extract(a: Datum, r: &AnalyzedRegex) -> Option<(Row, Diff)> {
3057 let r = r.inner();
3058 let a = a.unwrap_str();
3059 let captures = r.captures(a)?;
3060 let datums = captures
3061 .iter()
3062 .skip(1)
3063 .map(|m| Datum::from(m.map(|m| m.as_str())));
3064 Some((Row::pack(datums), Diff::ONE))
3065}
3066
3067fn regexp_matches<'a>(
3068 exprs: &[Datum<'a>],
3069) -> Result<impl Iterator<Item = (Row, Diff)> + 'a, EvalError> {
3070 assert!(exprs.len() == 2 || exprs.len() == 3);
3074 let a = exprs[0].unwrap_str();
3075 let r = exprs[1].unwrap_str();
3076
3077 let (regex, opts) = if exprs.len() == 3 {
3078 let flag = exprs[2].unwrap_str();
3079 let opts = AnalyzedRegexOpts::from_str(flag)?;
3080 (AnalyzedRegex::new(r, opts)?, opts)
3081 } else {
3082 let opts = AnalyzedRegexOpts::default();
3083 (AnalyzedRegex::new(r, opts)?, opts)
3084 };
3085
3086 let regex = regex.inner().clone();
3087
3088 let iter = regex.captures_iter(a).map(move |captures| {
3089 let matches = captures
3090 .iter()
3091 .skip(1)
3093 .map(|m| Datum::from(m.map(|m| m.as_str())))
3094 .collect::<Vec<_>>();
3095
3096 let mut binding = SharedRow::get();
3097 let mut packer = binding.packer();
3098
3099 let dimension = ArrayDimension {
3100 lower_bound: 1,
3101 length: matches.len(),
3102 };
3103 packer
3104 .try_push_array(&[dimension], matches)
3105 .expect("generated dimensions above");
3106
3107 (binding.clone(), Diff::ONE)
3108 });
3109
3110 let out = iter.collect::<SmallVec<[_; 3]>>();
3115
3116 if opts.global {
3117 Ok(Either::Left(out.into_iter()))
3118 } else {
3119 Ok(Either::Right(out.into_iter().take(1)))
3120 }
3121}
3122
3123fn generate_series<N>(
3124 start: N,
3125 stop: N,
3126 step: N,
3127) -> Result<impl Iterator<Item = (Row, Diff)>, EvalError>
3128where
3129 N: Integer + Signed + CheckedAdd + Clone,
3130 Datum<'static>: From<N>,
3131{
3132 if step == N::zero() {
3133 return Err(EvalError::InvalidParameterValue(
3134 "step size cannot equal zero".into(),
3135 ));
3136 }
3137 Ok(num::range_step_inclusive(start, stop, step)
3138 .map(move |i| (Row::pack_slice(&[Datum::from(i)]), Diff::ONE)))
3139}
3140
3141#[derive(Clone)]
3145pub struct TimestampRangeStepInclusive<T> {
3146 state: CheckedTimestamp<T>,
3147 stop: CheckedTimestamp<T>,
3148 step: Interval,
3149 rev: bool,
3150 done: bool,
3151}
3152
3153impl<T: TimestampLike> Iterator for TimestampRangeStepInclusive<T> {
3154 type Item = CheckedTimestamp<T>;
3155
3156 #[inline]
3157 fn next(&mut self) -> Option<CheckedTimestamp<T>> {
3158 if !self.done
3159 && ((self.rev && self.state >= self.stop) || (!self.rev && self.state <= self.stop))
3160 {
3161 let result = self.state.clone();
3162 match add_timestamp_months(self.state.deref(), self.step.months) {
3163 Ok(state) => match state.checked_add_signed(self.step.duration_as_chrono()) {
3164 Some(v) => match CheckedTimestamp::from_timestamplike(v) {
3165 Ok(v) => {
3166 let progressed = if self.rev {
3171 v < self.state
3172 } else {
3173 v > self.state
3174 };
3175 if progressed {
3176 self.state = v
3177 } else {
3178 self.done = true
3179 }
3180 }
3181 Err(_) => self.done = true,
3182 },
3183 None => self.done = true,
3184 },
3185 Err(..) => {
3186 self.done = true;
3187 }
3188 }
3189
3190 Some(result)
3191 } else {
3192 None
3193 }
3194 }
3195}
3196
3197fn generate_series_ts<T: TimestampLike>(
3198 start: CheckedTimestamp<T>,
3199 stop: CheckedTimestamp<T>,
3200 step: Interval,
3201 conv: fn(CheckedTimestamp<T>) -> Datum<'static>,
3202) -> Result<impl Iterator<Item = (Row, Diff)>, EvalError> {
3203 let normalized_step = step.as_microseconds();
3204 if normalized_step == 0 {
3205 return Err(EvalError::InvalidParameterValue(
3206 "step size cannot equal zero".into(),
3207 ));
3208 }
3209 let rev = normalized_step < 0;
3210
3211 let trsi = TimestampRangeStepInclusive {
3212 state: start,
3213 stop,
3214 step,
3215 rev,
3216 done: false,
3217 };
3218
3219 Ok(trsi.map(move |i| (Row::pack_slice(&[conv(i)]), Diff::ONE)))
3220}
3221
3222fn generate_subscripts_array(
3223 a: Datum,
3224 dim: i32,
3225) -> Result<Box<dyn Iterator<Item = (Row, Diff)>>, EvalError> {
3226 if dim <= 0 {
3227 return Ok(Box::new(iter::empty()));
3228 }
3229
3230 match a.unwrap_array().dims().into_iter().nth(
3231 (dim - 1)
3232 .try_into()
3233 .map_err(|_| EvalError::Int32OutOfRange((dim - 1).to_string().into()))?,
3234 ) {
3235 Some(requested_dim) => {
3236 let lower_bound: i32 = requested_dim.lower_bound.try_into().map_err(|_| {
3237 EvalError::Int32OutOfRange(requested_dim.lower_bound.to_string().into())
3238 })?;
3239 let length: i32 = requested_dim
3242 .length
3243 .try_into()
3244 .map_err(|_| EvalError::Int32OutOfRange(requested_dim.length.to_string().into()))?;
3245 let upper_bound = lower_bound.checked_add(length - 1).ok_or_else(|| {
3246 EvalError::Int32OutOfRange(requested_dim.length.to_string().into())
3247 })?;
3248 Ok(Box::new(generate_series::<i32>(
3249 lower_bound,
3250 upper_bound,
3251 1,
3252 )?))
3253 }
3254 None => Ok(Box::new(iter::empty())),
3255 }
3256}
3257
3258fn unnest_array<'a>(a: Datum<'a>) -> impl Iterator<Item = (Row, Diff)> + 'a {
3259 a.unwrap_array()
3260 .elements()
3261 .iter()
3262 .map(move |e| (Row::pack_slice(&[e]), Diff::ONE))
3263}
3264
3265fn unnest_list<'a>(a: Datum<'a>) -> impl Iterator<Item = (Row, Diff)> + 'a {
3266 a.unwrap_list()
3267 .iter()
3268 .map(move |e| (Row::pack_slice(&[e]), Diff::ONE))
3269}
3270
3271fn unnest_map<'a>(a: Datum<'a>) -> impl Iterator<Item = (Row, Diff)> + 'a {
3272 a.unwrap_map()
3273 .iter()
3274 .map(move |(k, v)| (Row::pack_slice(&[Datum::from(k), v]), Diff::ONE))
3275}
3276
3277impl AggregateFunc {
3278 pub fn name(&self) -> &'static str {
3281 match self {
3282 Self::MaxNumeric => "max",
3283 Self::MaxInt16 => "max",
3284 Self::MaxInt32 => "max",
3285 Self::MaxInt64 => "max",
3286 Self::MaxUInt16 => "max",
3287 Self::MaxUInt32 => "max",
3288 Self::MaxUInt64 => "max",
3289 Self::MaxMzTimestamp => "max",
3290 Self::MaxFloat32 => "max",
3291 Self::MaxFloat64 => "max",
3292 Self::MaxBool => "max",
3293 Self::MaxString => "max",
3294 Self::MaxDate => "max",
3295 Self::MaxTimestamp => "max",
3296 Self::MaxTimestampTz => "max",
3297 Self::MaxInterval => "max",
3298 Self::MaxTime => "max",
3299 Self::MaxArray => "max",
3300 Self::MaxList => "max",
3301 Self::MinNumeric => "min",
3302 Self::MinInt16 => "min",
3303 Self::MinInt32 => "min",
3304 Self::MinInt64 => "min",
3305 Self::MinUInt16 => "min",
3306 Self::MinUInt32 => "min",
3307 Self::MinUInt64 => "min",
3308 Self::MinMzTimestamp => "min",
3309 Self::MinFloat32 => "min",
3310 Self::MinFloat64 => "min",
3311 Self::MinBool => "min",
3312 Self::MinString => "min",
3313 Self::MinDate => "min",
3314 Self::MinTimestamp => "min",
3315 Self::MinTimestampTz => "min",
3316 Self::MinInterval => "min",
3317 Self::MinTime => "min",
3318 Self::MinArray => "min",
3319 Self::MinList => "min",
3320 Self::SumInt16 => "sum",
3321 Self::SumInt32 => "sum",
3322 Self::SumInt64 => "sum",
3323 Self::SumUInt16 => "sum",
3324 Self::SumUInt32 => "sum",
3325 Self::SumUInt64 => "sum",
3326 Self::SumFloat32 => "sum",
3327 Self::SumFloat64 => "sum",
3328 Self::SumNumeric => "sum",
3329 Self::SumInterval => "sum",
3330 Self::Count => "count",
3331 Self::Any => "any",
3332 Self::All => "all",
3333 Self::JsonbAgg { .. } => "jsonb_agg",
3334 Self::JsonbObjectAgg { .. } => "jsonb_object_agg",
3335 Self::MapAgg { .. } => "map_agg",
3336 Self::ArrayConcat { .. } => "array_agg",
3337 Self::ListConcat { .. } => "list_agg",
3338 Self::StringAgg { .. } => "string_agg",
3339 Self::RowNumber { .. } => "row_number",
3340 Self::Rank { .. } => "rank",
3341 Self::DenseRank { .. } => "dense_rank",
3342 Self::LagLead {
3343 lag_lead: LagLeadType::Lag,
3344 ..
3345 } => "lag",
3346 Self::LagLead {
3347 lag_lead: LagLeadType::Lead,
3348 ..
3349 } => "lead",
3350 Self::FirstValue { .. } => "first_value",
3351 Self::LastValue { .. } => "last_value",
3352 Self::WindowAggregate { .. } => "window_agg",
3353 Self::FusedValueWindowFunc { .. } => "fused_value_window_func",
3354 Self::FusedWindowAggregate { .. } => "fused_window_agg",
3355 Self::Dummy => "dummy",
3356 }
3357 }
3358}
3359
3360impl<'a, M> fmt::Display for HumanizedExpr<'a, AggregateFunc, M>
3361where
3362 M: HumanizerMode,
3363{
3364 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
3365 use AggregateFunc::*;
3366 let name = self.expr.name();
3367 match self.expr {
3368 JsonbAgg { order_by }
3369 | JsonbObjectAgg { order_by }
3370 | MapAgg { order_by, .. }
3371 | ArrayConcat { order_by }
3372 | ListConcat { order_by }
3373 | StringAgg { order_by }
3374 | RowNumber { order_by }
3375 | Rank { order_by }
3376 | DenseRank { order_by } => {
3377 let order_by = order_by.iter().map(|col| self.child(col));
3378 write!(f, "{}[order_by=[{}]]", name, separated(", ", order_by))
3379 }
3380 LagLead {
3381 lag_lead: _,
3382 ignore_nulls,
3383 order_by,
3384 } => {
3385 let order_by = order_by.iter().map(|col| self.child(col));
3386 f.write_str(name)?;
3387 f.write_str("[")?;
3388 if *ignore_nulls {
3389 f.write_str("ignore_nulls=true, ")?;
3390 }
3391 write!(f, "order_by=[{}]", separated(", ", order_by))?;
3392 f.write_str("]")
3393 }
3394 FirstValue {
3395 order_by,
3396 window_frame,
3397 } => {
3398 let order_by = order_by.iter().map(|col| self.child(col));
3399 f.write_str(name)?;
3400 f.write_str("[")?;
3401 write!(f, "order_by=[{}]", separated(", ", order_by))?;
3402 if *window_frame != WindowFrame::default() {
3403 write!(f, " {}", window_frame)?;
3404 }
3405 f.write_str("]")
3406 }
3407 LastValue {
3408 order_by,
3409 window_frame,
3410 } => {
3411 let order_by = order_by.iter().map(|col| self.child(col));
3412 f.write_str(name)?;
3413 f.write_str("[")?;
3414 write!(f, "order_by=[{}]", separated(", ", order_by))?;
3415 if *window_frame != WindowFrame::default() {
3416 write!(f, " {}", window_frame)?;
3417 }
3418 f.write_str("]")
3419 }
3420 WindowAggregate {
3421 wrapped_aggregate,
3422 order_by,
3423 window_frame,
3424 } => {
3425 let order_by = order_by.iter().map(|col| self.child(col));
3426 let wrapped_aggregate = self.child(wrapped_aggregate.deref());
3427 f.write_str(name)?;
3428 f.write_str("[")?;
3429 write!(f, "{} ", wrapped_aggregate)?;
3430 write!(f, "order_by=[{}]", separated(", ", order_by))?;
3431 if *window_frame != WindowFrame::default() {
3432 write!(f, " {}", window_frame)?;
3433 }
3434 f.write_str("]")
3435 }
3436 FusedValueWindowFunc { funcs, order_by } => {
3437 let order_by = order_by.iter().map(|col| self.child(col));
3438 let funcs = separated(", ", funcs.iter().map(|func| self.child(func)));
3439 f.write_str(name)?;
3440 f.write_str("[")?;
3441 write!(f, "{} ", funcs)?;
3442 write!(f, "order_by=[{}]", separated(", ", order_by))?;
3443 f.write_str("]")
3444 }
3445 _ => f.write_str(name),
3446 }
3447 }
3448}
3449
3450#[derive(
3451 Clone,
3452 Debug,
3453 Eq,
3454 PartialEq,
3455 Ord,
3456 PartialOrd,
3457 Serialize,
3458 Deserialize,
3459 Hash
3460)]
3461pub struct CaptureGroupDesc {
3462 pub index: u32,
3463 pub name: Option<String>,
3464 pub nullable: bool,
3465}
3466
3467#[derive(
3468 Clone,
3469 Copy,
3470 Debug,
3471 Eq,
3472 PartialEq,
3473 Ord,
3474 PartialOrd,
3475 Serialize,
3476 Deserialize,
3477 Hash,
3478 Default
3479)]
3480pub struct AnalyzedRegexOpts {
3481 pub case_insensitive: bool,
3482 pub global: bool,
3483}
3484
3485impl FromStr for AnalyzedRegexOpts {
3486 type Err = EvalError;
3487
3488 fn from_str(s: &str) -> Result<Self, Self::Err> {
3489 let mut opts = AnalyzedRegexOpts::default();
3490 for c in s.chars() {
3491 match c {
3492 'i' => opts.case_insensitive = true,
3493 'g' => opts.global = true,
3494 _ => return Err(EvalError::InvalidRegexFlag(c)),
3495 }
3496 }
3497 Ok(opts)
3498 }
3499}
3500
3501#[derive(
3502 Clone,
3503 Debug,
3504 Eq,
3505 PartialEq,
3506 Ord,
3507 PartialOrd,
3508 Serialize,
3509 Deserialize,
3510 Hash
3511)]
3512pub struct AnalyzedRegex(ReprRegex, Vec<CaptureGroupDesc>, AnalyzedRegexOpts);
3513
3514impl AnalyzedRegex {
3515 pub fn new(s: &str, opts: AnalyzedRegexOpts) -> Result<Self, RegexCompilationError> {
3516 let r = ReprRegex::new(s, opts.case_insensitive)?;
3517 #[allow(clippy::as_conversions)]
3519 let descs: Vec<_> = r
3520 .capture_names()
3521 .enumerate()
3522 .skip(1)
3527 .map(|(i, name)| CaptureGroupDesc {
3528 index: i as u32,
3529 name: name.map(String::from),
3530 nullable: true,
3533 })
3534 .collect();
3535 Ok(Self(r, descs, opts))
3536 }
3537 pub fn capture_groups_len(&self) -> usize {
3538 self.1.len()
3539 }
3540 pub fn capture_groups_iter(&self) -> impl Iterator<Item = &CaptureGroupDesc> {
3541 self.1.iter()
3542 }
3543 pub fn inner(&self) -> &Regex {
3544 &(self.0).regex
3545 }
3546 pub fn opts(&self) -> &AnalyzedRegexOpts {
3547 &self.2
3548 }
3549}
3550
3551pub fn csv_extract(a: Datum<'_>, n_cols: usize) -> impl Iterator<Item = (Row, Diff)> + '_ {
3552 let bytes = a.unwrap_str().as_bytes();
3553 let mut row = Row::default();
3554 let csv_reader = csv::ReaderBuilder::new()
3555 .has_headers(false)
3556 .from_reader(bytes);
3557 csv_reader.into_records().filter_map(move |res| match res {
3558 Ok(sr) if sr.len() == n_cols => {
3559 row.packer().extend(sr.iter().map(Datum::String));
3560 Some((row.clone(), Diff::ONE))
3561 }
3562 _ => None,
3563 })
3564}
3565
3566pub fn repeat_row(a: Datum) -> Option<(Row, Diff)> {
3567 let n = a.unwrap_int64();
3568 if n != 0 {
3569 Some((Row::default(), n.into()))
3570 } else {
3571 None
3572 }
3573}
3574
3575pub fn repeat_row_non_negative<'a>(
3576 a: Datum,
3577) -> Result<Box<dyn Iterator<Item = (Row, Diff)> + 'a>, EvalError> {
3578 let n = a.unwrap_int64();
3579 if n < 0 {
3580 Err(EvalError::InvalidParameterValue(
3581 format!("repeat_row_non_negative got {}", n).into(),
3582 ))
3583 } else if n == 0 {
3584 Ok(Box::new(iter::empty()))
3585 } else {
3586 Ok(Box::new(iter::once((Row::default(), n.into()))))
3588 }
3589}
3590
3591fn wrap<'a>(datums: &'a [Datum<'a>], width: usize) -> impl Iterator<Item = (Row, Diff)> + 'a {
3592 datums
3593 .chunks(width)
3594 .map(|chunk| (Row::pack(chunk), Diff::ONE))
3595}
3596
3597fn acl_explode<'a>(
3598 acl_items: Datum<'a>,
3599 temp_storage: &'a RowArena,
3600) -> Result<impl Iterator<Item = (Row, Diff)> + 'a, EvalError> {
3601 let acl_items = acl_items.unwrap_array();
3602 let mut res = Vec::new();
3603 for acl_item in acl_items.elements().iter() {
3604 if acl_item.is_null() {
3605 return Err(EvalError::AclArrayNullElement);
3606 }
3607 let acl_item = acl_item.unwrap_acl_item();
3608 for privilege in acl_item.acl_mode.explode() {
3609 let row = [
3610 Datum::UInt32(acl_item.grantor.0),
3611 Datum::UInt32(acl_item.grantee.0),
3612 Datum::String(temp_storage.push_string(privilege.to_string())),
3613 Datum::False,
3615 ];
3616 res.push((Row::pack_slice(&row), Diff::ONE));
3617 }
3618 }
3619 Ok(res.into_iter())
3620}
3621
3622fn mz_acl_explode<'a>(
3623 mz_acl_items: Datum<'a>,
3624 temp_storage: &'a RowArena,
3625) -> Result<impl Iterator<Item = (Row, Diff)> + 'a, EvalError> {
3626 let mz_acl_items = mz_acl_items.unwrap_array();
3627 let mut res = Vec::new();
3628 for mz_acl_item in mz_acl_items.elements().iter() {
3629 if mz_acl_item.is_null() {
3630 return Err(EvalError::MzAclArrayNullElement);
3631 }
3632 let mz_acl_item = mz_acl_item.unwrap_mz_acl_item();
3633 for privilege in mz_acl_item.acl_mode.explode() {
3634 let row = [
3635 Datum::String(temp_storage.push_string(mz_acl_item.grantor.to_string())),
3636 Datum::String(temp_storage.push_string(mz_acl_item.grantee.to_string())),
3637 Datum::String(temp_storage.push_string(privilege.to_string())),
3638 Datum::False,
3640 ];
3641 res.push((Row::pack_slice(&row), Diff::ONE));
3642 }
3643 }
3644 Ok(res.into_iter())
3645}
3646
3647#[derive(
3650 Clone,
3651 Debug,
3652 Eq,
3653 PartialEq,
3654 Ord,
3655 PartialOrd,
3656 Serialize,
3657 Deserialize,
3658 Hash
3659)]
3660pub enum TableFunc {
3661 AclExplode,
3662 MzAclExplode,
3663 JsonbEach,
3664 JsonbEachStringify,
3665 JsonbObjectKeys,
3666 JsonbArrayElements,
3667 JsonbArrayElementsStringify,
3668 RegexpExtract(AnalyzedRegex),
3669 CsvExtract(usize),
3670 GenerateSeriesInt32,
3671 GenerateSeriesInt64,
3672 GenerateSeriesUnoptimized,
3684 GenerateSeriesTimestamp,
3685 GenerateSeriesTimestampTz,
3686 GuardSubquerySize {
3707 column_type: SqlScalarType,
3708 },
3709 RepeatRow,
3716 RepeatRowNonNegative,
3719 UnnestArray {
3720 el_typ: SqlScalarType,
3721 },
3722 UnnestList {
3723 el_typ: SqlScalarType,
3724 },
3725 UnnestMap {
3726 value_type: SqlScalarType,
3727 },
3728 Wrap {
3734 types: Vec<SqlColumnType>,
3735 width: usize,
3736 },
3737 GenerateSubscriptsArray,
3738 TabletizedScalar {
3740 name: String,
3741 relation: SqlRelationType,
3742 },
3743 RegexpMatches,
3744 #[allow(private_interfaces)]
3749 WithOrdinality(WithOrdinality),
3750}
3751
3752#[derive(
3760 Clone,
3761 Debug,
3762 Eq,
3763 PartialEq,
3764 Ord,
3765 PartialOrd,
3766 Serialize,
3767 Deserialize,
3768 Hash
3769)]
3770struct WithOrdinality {
3771 inner: Box<TableFunc>,
3772}
3773
3774impl TableFunc {
3775 pub fn with_ordinality(inner: TableFunc) -> Option<TableFunc> {
3777 match inner {
3778 TableFunc::AclExplode
3779 | TableFunc::MzAclExplode
3780 | TableFunc::JsonbEach
3781 | TableFunc::JsonbEachStringify
3782 | TableFunc::JsonbObjectKeys
3783 | TableFunc::JsonbArrayElements
3784 | TableFunc::JsonbArrayElementsStringify
3785 | TableFunc::RegexpExtract(_)
3786 | TableFunc::CsvExtract(_)
3787 | TableFunc::GenerateSeriesInt32
3788 | TableFunc::GenerateSeriesInt64
3789 | TableFunc::GenerateSeriesUnoptimized
3790 | TableFunc::GenerateSeriesTimestamp
3791 | TableFunc::GenerateSeriesTimestampTz
3792 | TableFunc::GuardSubquerySize { .. }
3793 | TableFunc::RepeatRowNonNegative
3794 | TableFunc::UnnestArray { .. }
3795 | TableFunc::UnnestList { .. }
3796 | TableFunc::UnnestMap { .. }
3797 | TableFunc::Wrap { .. }
3798 | TableFunc::GenerateSubscriptsArray
3799 | TableFunc::TabletizedScalar { .. }
3800 | TableFunc::RegexpMatches => Some(TableFunc::WithOrdinality(WithOrdinality {
3801 inner: Box::new(inner),
3802 })),
3803 TableFunc::RepeatRow | TableFunc::WithOrdinality(_) => None, }
3811 }
3812}
3813
3814impl TableFunc {
3815 pub fn eval<'a>(
3817 &'a self,
3818 datums: &'a [Datum<'a>],
3819 temp_storage: &'a RowArena,
3820 ) -> Result<Box<dyn Iterator<Item = (Row, Diff)> + 'a>, EvalError> {
3821 if self.empty_on_null_input() && datums.iter().any(|d| d.is_null()) {
3822 return Ok(Box::new(vec![].into_iter()));
3823 }
3824 match self {
3825 TableFunc::AclExplode => Ok(Box::new(acl_explode(datums[0], temp_storage)?)),
3826 TableFunc::MzAclExplode => Ok(Box::new(mz_acl_explode(datums[0], temp_storage)?)),
3827 TableFunc::JsonbEach => Ok(Box::new(jsonb_each(datums[0]))),
3828 TableFunc::JsonbEachStringify => {
3829 Ok(Box::new(jsonb_each_stringify(datums[0], temp_storage)))
3830 }
3831 TableFunc::JsonbObjectKeys => Ok(Box::new(jsonb_object_keys(datums[0]))),
3832 TableFunc::JsonbArrayElements => Ok(Box::new(jsonb_array_elements(datums[0]))),
3833 TableFunc::JsonbArrayElementsStringify => Ok(Box::new(jsonb_array_elements_stringify(
3834 datums[0],
3835 temp_storage,
3836 ))),
3837 TableFunc::RegexpExtract(a) => Ok(Box::new(regexp_extract(datums[0], a).into_iter())),
3838 TableFunc::CsvExtract(n_cols) => Ok(Box::new(csv_extract(datums[0], *n_cols))),
3839 TableFunc::GenerateSeriesInt32 => {
3840 let res = generate_series(
3841 datums[0].unwrap_int32(),
3842 datums[1].unwrap_int32(),
3843 datums[2].unwrap_int32(),
3844 )?;
3845 Ok(Box::new(res))
3846 }
3847 TableFunc::GenerateSeriesInt64 | TableFunc::GenerateSeriesUnoptimized => {
3848 let res = generate_series(
3849 datums[0].unwrap_int64(),
3850 datums[1].unwrap_int64(),
3851 datums[2].unwrap_int64(),
3852 )?;
3853 Ok(Box::new(res))
3854 }
3855 TableFunc::GenerateSeriesTimestamp => {
3856 fn pass_through<'a>(d: CheckedTimestamp<NaiveDateTime>) -> Datum<'a> {
3857 Datum::from(d)
3858 }
3859 let res = generate_series_ts(
3860 datums[0].unwrap_timestamp(),
3861 datums[1].unwrap_timestamp(),
3862 datums[2].unwrap_interval(),
3863 pass_through,
3864 )?;
3865 Ok(Box::new(res))
3866 }
3867 TableFunc::GenerateSeriesTimestampTz => {
3868 fn gen_ts_tz<'a>(d: CheckedTimestamp<DateTime<Utc>>) -> Datum<'a> {
3869 Datum::from(d)
3870 }
3871 let res = generate_series_ts(
3872 datums[0].unwrap_timestamptz(),
3873 datums[1].unwrap_timestamptz(),
3874 datums[2].unwrap_interval(),
3875 gen_ts_tz,
3876 )?;
3877 Ok(Box::new(res))
3878 }
3879 TableFunc::GenerateSubscriptsArray => {
3880 generate_subscripts_array(datums[0], datums[1].unwrap_int32())
3881 }
3882 TableFunc::GuardSubquerySize { column_type: _ } => {
3883 let count = datums[0].unwrap_int64();
3895 if count > 1 {
3896 Err(EvalError::MultipleRowsFromSubquery)
3897 } else if count < 0 {
3898 Err(EvalError::NegativeRowsFromSubquery)
3900 } else {
3901 Ok(Box::new([].into_iter()))
3902 }
3903 }
3904 TableFunc::RepeatRow => Ok(Box::new(repeat_row(datums[0]).into_iter())),
3905 TableFunc::RepeatRowNonNegative => repeat_row_non_negative(datums[0]),
3906 TableFunc::UnnestArray { .. } => Ok(Box::new(unnest_array(datums[0]))),
3907 TableFunc::UnnestList { .. } => Ok(Box::new(unnest_list(datums[0]))),
3908 TableFunc::UnnestMap { .. } => Ok(Box::new(unnest_map(datums[0]))),
3909 TableFunc::Wrap { width, .. } => Ok(Box::new(wrap(datums, *width))),
3910 TableFunc::TabletizedScalar { .. } => {
3911 let r = Row::pack_slice(datums);
3912 Ok(Box::new(std::iter::once((r, Diff::ONE))))
3913 }
3914 TableFunc::RegexpMatches => Ok(Box::new(regexp_matches(datums)?)),
3915 TableFunc::WithOrdinality(func_with_ordinality) => {
3916 func_with_ordinality.eval(datums, temp_storage)
3917 }
3918 }
3919 }
3920
3921 pub fn output_sql_type(&self) -> SqlRelationType {
3922 let (column_types, keys) = match self {
3923 TableFunc::AclExplode => {
3924 let column_types = vec![
3925 SqlScalarType::Oid.nullable(false),
3926 SqlScalarType::Oid.nullable(false),
3927 SqlScalarType::String.nullable(false),
3928 SqlScalarType::Bool.nullable(false),
3929 ];
3930 let keys = vec![];
3931 (column_types, keys)
3932 }
3933 TableFunc::MzAclExplode => {
3934 let column_types = vec![
3935 SqlScalarType::String.nullable(false),
3936 SqlScalarType::String.nullable(false),
3937 SqlScalarType::String.nullable(false),
3938 SqlScalarType::Bool.nullable(false),
3939 ];
3940 let keys = vec![];
3941 (column_types, keys)
3942 }
3943 TableFunc::JsonbEach => {
3944 let column_types = vec![
3945 SqlScalarType::String.nullable(false),
3946 SqlScalarType::Jsonb.nullable(false),
3947 ];
3948 let keys = vec![];
3949 (column_types, keys)
3950 }
3951 TableFunc::JsonbEachStringify => {
3952 let column_types = vec![
3953 SqlScalarType::String.nullable(false),
3954 SqlScalarType::String.nullable(true),
3955 ];
3956 let keys = vec![];
3957 (column_types, keys)
3958 }
3959 TableFunc::JsonbObjectKeys => {
3960 let column_types = vec![SqlScalarType::String.nullable(false)];
3961 let keys = vec![];
3962 (column_types, keys)
3963 }
3964 TableFunc::JsonbArrayElements => {
3965 let column_types = vec![SqlScalarType::Jsonb.nullable(false)];
3966 let keys = vec![];
3967 (column_types, keys)
3968 }
3969 TableFunc::JsonbArrayElementsStringify => {
3970 let column_types = vec![SqlScalarType::String.nullable(true)];
3971 let keys = vec![];
3972 (column_types, keys)
3973 }
3974 TableFunc::RegexpExtract(a) => {
3975 let column_types = a
3976 .capture_groups_iter()
3977 .map(|cg| SqlScalarType::String.nullable(cg.nullable))
3978 .collect();
3979 let keys = vec![];
3980 (column_types, keys)
3981 }
3982 TableFunc::CsvExtract(n_cols) => {
3983 let column_types = iter::repeat(SqlScalarType::String.nullable(false))
3984 .take(*n_cols)
3985 .collect();
3986 let keys = vec![];
3987 (column_types, keys)
3988 }
3989 TableFunc::GenerateSeriesInt32 => {
3990 let column_types = vec![SqlScalarType::Int32.nullable(false)];
3991 let keys = vec![vec![0]];
3992 (column_types, keys)
3993 }
3994 TableFunc::GenerateSeriesInt64 | TableFunc::GenerateSeriesUnoptimized => {
3995 let column_types = vec![SqlScalarType::Int64.nullable(false)];
3996 let keys = vec![vec![0]];
3997 (column_types, keys)
3998 }
3999 TableFunc::GenerateSeriesTimestamp => {
4000 let column_types =
4001 vec![SqlScalarType::Timestamp { precision: None }.nullable(false)];
4002 let keys = vec![vec![0]];
4003 (column_types, keys)
4004 }
4005 TableFunc::GenerateSeriesTimestampTz => {
4006 let column_types =
4007 vec![SqlScalarType::TimestampTz { precision: None }.nullable(false)];
4008 let keys = vec![vec![0]];
4009 (column_types, keys)
4010 }
4011 TableFunc::GenerateSubscriptsArray => {
4012 let column_types = vec![SqlScalarType::Int32.nullable(false)];
4013 let keys = vec![vec![0]];
4014 (column_types, keys)
4015 }
4016 TableFunc::GuardSubquerySize { column_type } => {
4017 let column_types = vec![column_type.clone().nullable(false)];
4018 let keys = vec![];
4019 (column_types, keys)
4020 }
4021 TableFunc::RepeatRow | TableFunc::RepeatRowNonNegative => {
4022 let column_types = vec![];
4023 let keys = vec![];
4024 (column_types, keys)
4025 }
4026 TableFunc::UnnestArray { el_typ } => {
4027 let column_types = vec![el_typ.clone().nullable(true)];
4028 let keys = vec![];
4029 (column_types, keys)
4030 }
4031 TableFunc::UnnestList { el_typ } => {
4032 let column_types = vec![el_typ.clone().nullable(true)];
4033 let keys = vec![];
4034 (column_types, keys)
4035 }
4036 TableFunc::UnnestMap { value_type } => {
4037 let column_types = vec![
4038 SqlScalarType::String.nullable(false),
4039 value_type.clone().nullable(true),
4040 ];
4041 let keys = vec![vec![0]];
4042 (column_types, keys)
4043 }
4044 TableFunc::Wrap { types, .. } => {
4045 let column_types = types.clone();
4046 let keys = vec![];
4047 (column_types, keys)
4048 }
4049 TableFunc::TabletizedScalar { relation, .. } => {
4050 return relation.clone();
4051 }
4052 TableFunc::RegexpMatches => {
4053 let column_types =
4054 vec![SqlScalarType::Array(Box::new(SqlScalarType::String)).nullable(false)];
4055 let keys = vec![];
4056
4057 (column_types, keys)
4058 }
4059 TableFunc::WithOrdinality(WithOrdinality { inner }) => {
4060 let mut typ = inner.output_sql_type();
4061 typ.column_types.push(SqlScalarType::Int64.nullable(false));
4063 typ.keys.push(vec![typ.column_types.len() - 1]);
4065 (typ.column_types, typ.keys)
4066 }
4067 };
4068
4069 soft_assert_eq_no_log!(column_types.len(), self.output_arity());
4070
4071 if !keys.is_empty() {
4072 SqlRelationType::new(column_types).with_keys(keys)
4073 } else {
4074 SqlRelationType::new(column_types)
4075 }
4076 }
4077
4078 pub fn output_type(&self) -> ReprRelationType {
4082 ReprRelationType::from(&self.output_sql_type())
4083 }
4084
4085 pub fn output_arity(&self) -> usize {
4086 match self {
4087 TableFunc::AclExplode => 4,
4088 TableFunc::MzAclExplode => 4,
4089 TableFunc::JsonbEach => 2,
4090 TableFunc::JsonbEachStringify => 2,
4091 TableFunc::JsonbObjectKeys => 1,
4092 TableFunc::JsonbArrayElements => 1,
4093 TableFunc::JsonbArrayElementsStringify => 1,
4094 TableFunc::RegexpExtract(a) => a.capture_groups_len(),
4095 TableFunc::CsvExtract(n_cols) => *n_cols,
4096 TableFunc::GenerateSeriesInt32 => 1,
4097 TableFunc::GenerateSeriesInt64 => 1,
4098 TableFunc::GenerateSeriesUnoptimized => 1,
4099 TableFunc::GenerateSeriesTimestamp => 1,
4100 TableFunc::GenerateSeriesTimestampTz => 1,
4101 TableFunc::GenerateSubscriptsArray => 1,
4102 TableFunc::GuardSubquerySize { .. } => 1,
4103 TableFunc::RepeatRow => 0,
4104 TableFunc::RepeatRowNonNegative => 0,
4105 TableFunc::UnnestArray { .. } => 1,
4106 TableFunc::UnnestList { .. } => 1,
4107 TableFunc::UnnestMap { .. } => 2,
4108 TableFunc::Wrap { width, .. } => *width,
4109 TableFunc::TabletizedScalar { relation, .. } => relation.column_types.len(),
4110 TableFunc::RegexpMatches => 1,
4111 TableFunc::WithOrdinality(WithOrdinality { inner }) => inner.output_arity() + 1,
4112 }
4113 }
4114
4115 pub fn empty_on_null_input(&self) -> bool {
4116 match self {
4117 TableFunc::AclExplode
4118 | TableFunc::MzAclExplode
4119 | TableFunc::JsonbEach
4120 | TableFunc::JsonbEachStringify
4121 | TableFunc::JsonbObjectKeys
4122 | TableFunc::JsonbArrayElements
4123 | TableFunc::JsonbArrayElementsStringify
4124 | TableFunc::GenerateSeriesInt32
4125 | TableFunc::GenerateSeriesInt64
4126 | TableFunc::GenerateSeriesUnoptimized
4127 | TableFunc::GenerateSeriesTimestamp
4128 | TableFunc::GenerateSeriesTimestampTz
4129 | TableFunc::GenerateSubscriptsArray
4130 | TableFunc::RegexpExtract(_)
4131 | TableFunc::CsvExtract(_)
4132 | TableFunc::RepeatRow
4133 | TableFunc::RepeatRowNonNegative
4134 | TableFunc::UnnestArray { .. }
4135 | TableFunc::UnnestList { .. }
4136 | TableFunc::UnnestMap { .. }
4137 | TableFunc::RegexpMatches => true,
4138 TableFunc::GuardSubquerySize { .. } => false,
4139 TableFunc::Wrap { .. } => false,
4140 TableFunc::TabletizedScalar { .. } => false,
4141 TableFunc::WithOrdinality(WithOrdinality { inner }) => inner.empty_on_null_input(),
4142 }
4143 }
4144
4145 pub fn preserves_monotonicity(&self) -> bool {
4147 match self {
4150 TableFunc::AclExplode => false,
4151 TableFunc::MzAclExplode => false,
4152 TableFunc::JsonbEach => true,
4153 TableFunc::JsonbEachStringify => true,
4154 TableFunc::JsonbObjectKeys => true,
4155 TableFunc::JsonbArrayElements => true,
4156 TableFunc::JsonbArrayElementsStringify => true,
4157 TableFunc::RegexpExtract(_) => true,
4158 TableFunc::CsvExtract(_) => true,
4159 TableFunc::GenerateSeriesInt32 => true,
4160 TableFunc::GenerateSeriesInt64 => true,
4161 TableFunc::GenerateSeriesUnoptimized => true,
4162 TableFunc::GenerateSeriesTimestamp => true,
4163 TableFunc::GenerateSeriesTimestampTz => true,
4164 TableFunc::GenerateSubscriptsArray => true,
4165 TableFunc::RepeatRow => false,
4166 TableFunc::RepeatRowNonNegative => true,
4167 TableFunc::UnnestArray { .. } => true,
4168 TableFunc::UnnestList { .. } => true,
4169 TableFunc::UnnestMap { .. } => true,
4170 TableFunc::Wrap { .. } => true,
4171 TableFunc::TabletizedScalar { .. } => true,
4172 TableFunc::RegexpMatches => true,
4173 TableFunc::GuardSubquerySize { .. } => false,
4174 TableFunc::WithOrdinality(WithOrdinality { inner }) => inner.preserves_monotonicity(),
4175 }
4176 }
4177}
4178
4179impl fmt::Display for TableFunc {
4180 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
4181 match self {
4182 TableFunc::AclExplode => f.write_str("aclexplode"),
4183 TableFunc::MzAclExplode => f.write_str("mz_aclexplode"),
4184 TableFunc::JsonbEach => f.write_str("jsonb_each"),
4185 TableFunc::JsonbEachStringify => f.write_str("jsonb_each_text"),
4186 TableFunc::JsonbObjectKeys => f.write_str("jsonb_object_keys"),
4187 TableFunc::JsonbArrayElements => f.write_str("jsonb_array_elements"),
4188 TableFunc::JsonbArrayElementsStringify => f.write_str("jsonb_array_elements_text"),
4189 TableFunc::RegexpExtract(a) => write!(f, "regexp_extract({:?}, _)", a.0),
4190 TableFunc::CsvExtract(n_cols) => write!(f, "csv_extract({}, _)", n_cols),
4191 TableFunc::GenerateSeriesInt32 => f.write_str("generate_series"),
4192 TableFunc::GenerateSeriesInt64 => f.write_str("generate_series"),
4193 TableFunc::GenerateSeriesUnoptimized => f.write_str("generate_series_unoptimized"),
4194 TableFunc::GenerateSeriesTimestamp => f.write_str("generate_series"),
4195 TableFunc::GenerateSeriesTimestampTz => f.write_str("generate_series"),
4196 TableFunc::GenerateSubscriptsArray => f.write_str("generate_subscripts"),
4197 TableFunc::GuardSubquerySize { .. } => f.write_str("guard_subquery_size"),
4198 TableFunc::RepeatRow => f.write_str(REPEAT_ROW_NAME),
4199 TableFunc::RepeatRowNonNegative => f.write_str("repeat_row_non_negative"),
4200 TableFunc::UnnestArray { .. } => f.write_str("unnest_array"),
4201 TableFunc::UnnestList { .. } => f.write_str("unnest_list"),
4202 TableFunc::UnnestMap { .. } => f.write_str("unnest_map"),
4203 TableFunc::Wrap { width, .. } => write!(f, "wrap{}", width),
4204 TableFunc::TabletizedScalar { name, .. } => f.write_str(name),
4205 TableFunc::RegexpMatches => write!(f, "regexp_matches(_, _, _)"),
4206 TableFunc::WithOrdinality(WithOrdinality { inner }) => {
4207 write!(f, "{}[with_ordinality]", inner)
4208 }
4209 }
4210 }
4211}
4212
4213impl WithOrdinality {
4214 fn eval<'a>(
4223 &'a self,
4224 datums: &'a [Datum<'a>],
4225 temp_storage: &'a RowArena,
4226 ) -> Result<Box<dyn Iterator<Item = (Row, Diff)> + 'a>, EvalError> {
4227 let mut next_ordinal: i64 = 1;
4228 let it = self
4229 .inner
4230 .eval(datums, temp_storage)?
4231 .flat_map(move |(mut row, diff)| {
4232 let diff = diff.into_inner();
4233 assert!(diff >= 0);
4241 let mut ordinals = next_ordinal..(next_ordinal + diff);
4243 next_ordinal += diff;
4244 let cap = row.data_len() + datum_size(&Datum::Int64(next_ordinal));
4246 iter::from_fn(move || {
4247 let ordinal = ordinals.next()?;
4248 let mut row = if ordinals.is_empty() {
4249 std::mem::take(&mut row)
4252 } else {
4253 let mut new_row = Row::with_capacity(cap);
4254 new_row.clone_from(&row);
4255 new_row
4256 };
4257 RowPacker::for_existing_row(&mut row).push(Datum::Int64(ordinal));
4258 Some((row, Diff::ONE))
4259 })
4260 });
4261 Ok(Box::new(it))
4262 }
4263}
4264
4265pub const REPEAT_ROW_NAME: &str = "repeat_row";
4266
4267#[cfg(test)]
4268mod tests {
4269 use mz_repr::{Datum, RowArena, SqlScalarType};
4270
4271 use super::TableFunc;
4272 use crate::EvalError;
4273
4274 #[mz_ore::test]
4281 fn guard_subquery_size_accepts_zero_and_one() {
4282 let func = TableFunc::GuardSubquerySize {
4283 column_type: SqlScalarType::Int64,
4284 };
4285 let temp_storage = RowArena::new();
4286
4287 for count in [0_i64, 1] {
4288 let rows = func
4289 .eval(&[Datum::Int64(count)], &temp_storage)
4290 .unwrap_or_else(|e| panic!("count {count} should be accepted, got {e:?}"))
4291 .count();
4292 assert_eq!(rows, 0, "count {count} should emit no guard rows");
4293 }
4294
4295 assert_eq!(
4296 func.eval(&[Datum::Int64(2)], &temp_storage).err(),
4297 Some(EvalError::MultipleRowsFromSubquery),
4298 );
4299 assert_eq!(
4300 func.eval(&[Datum::Int64(-1)], &temp_storage).err(),
4301 Some(EvalError::NegativeRowsFromSubquery),
4302 );
4303 }
4304}