mz_compute_types/plan/reduce.rs
1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10//! Reduction execution planning and dataflow construction.
11
12//! We build `ReducePlan`s to manage the complexity of planning the generated dataflow for a
13//! given reduce expression. The intent here is that each creating a `ReducePlan` should capture
14//! all of the decision making about what kind of dataflow do we need to render and what each
15//! operator needs to do, and then actually rendering the plan can be a relatively simple application
16//! of (as much as possible) straight line code.
17//!
18//! Materialize needs to be able to maintain reductions incrementally (roughly, using
19//! time proportional to the number of changes in the input) and ideally, with a
20//! memory footprint proportional to the number of reductions being computed. We have to employ
21//! several tricks to achieve that, and these tricks constitute most of the complexity involved
22//! with planning and rendering reduce expressions. There's some additional complexity involved
23//! in handling aggregations with `DISTINCT` correctly so that we can efficiently suppress
24//! duplicate updates.
25//!
26//! In order to optimize the performance of our rendered dataflow, we divide all aggregations
27//! into three distinct types. Each type gets rendered separately, with its own specialized plan
28//! and dataflow. The three types are as follows:
29//!
30//! 1. Accumulable:
31//! Accumulable reductions can be computed inline in a Differential update's `difference`
32//! field because they basically boil down to tracking counts of things. `sum()` is an
33//! example of an accumulable reduction, and when some element `x` is removed from the set
34//! of elements being summed, we can introduce `-x` to incrementally maintain the sum. More
35//! formally, accumulable reductions correspond to instances of commutative Abelian groups.
36//! 2. Hierarchical:
37//! Hierarchical reductions don't have a meaningful negation like accumulable reductions do, but
38//! they are still commutative and associative, which lets us compute the reduction over subsets
39//! of the input, and then compute the reduction again on those results. For example:
40//! `min[2, 5, 1, 10]` is the same as `min[ min[2, 5], min[1, 10]]`. When we compute hierarchical
41//! reductions this way, we can maintain the computation in sublinear time with respect to
42//! the overall input. `min` and `max` are two examples of hierarchical reductions. More formally,
43//! hierarchical reductions correspond to instances of semigroups, in that they are associative,
44//! but in order to benefit from being computed hierarchically, they need to have some reduction
45//! in data size as well. A function like "concat-everything-to-a-string" wouldn't benefit from
46//! hierarchical evaluation.
47//!
48//! When the input is append-only, or monotonic, reductions that would otherwise have to be computed
49//! hierarchically can instead be computed in-place, because we only need to keep the value that's
50//! better than the "best" (minimal or maximal for min and max) seen so far.
51//! 3. Basic:
52//! Basic reductions are a bit like the Hufflepuffs of this trifecta. They are neither accumulable nor
53//! hierarchical (most likely they are associative but don't involve any data reduction) and so for these
54//! we can't do much more than just defer to Differential's reduce operator and eat a large maintenance cost.
55//!
56//! When we render these reductions we want to limit the number of arrangements we produce. When we build a
57//! dataflow for a reduction containing multiple types of reductions, we have no choice but to divide up the
58//! requested aggregations by type, render each type separately and then take those results and collate them
59//! back in the requested output order. However, if we only need to perform aggregations of a single reduction
60//! type, we can specialize and render the dataflow to compute those aggregations in the correct order, and
61//! return the output arrangement directly and avoid the extra collation arrangement.
62
63use std::fmt::Display;
64
65use mz_expr::explain::{HumanizeDisplay, HumanizedExpr, HumanizerMode};
66use mz_expr::{
67 AggregateExpr, AggregateFunc, MapFilterProject, MirScalarExpr, UnmaterializableFunc,
68 permutation_for_arrangement,
69};
70use mz_ore::soft_assert_or_log;
71use serde::{Deserialize, Serialize};
72
73use crate::plan::scalar::LirScalarExpr;
74use crate::plan::{AvailableCollections, bucketing_of_expected_group_size};
75
76/// This enum represents the three potential types of aggregations.
77#[derive(
78 Copy,
79 Clone,
80 Debug,
81 Deserialize,
82 Eq,
83 Hash,
84 Ord,
85 PartialEq,
86 PartialOrd,
87 Serialize
88)]
89pub enum ReductionType {
90 /// Accumulable functions can be subtracted from (are invertible), and associative.
91 /// We can compute these results by moving some data to the diff field under arbitrary
92 /// changes to inputs. Examples include sum or count.
93 Accumulable,
94 /// Hierarchical functions are associative, which means we can split up the work of
95 /// computing them across subsets. Note that hierarchical reductions should also
96 /// reduce the data in some way, as otherwise rendering them hierarchically is not
97 /// worth it. Examples include min or max.
98 Hierarchical,
99 /// Basic, for lack of a better word, are functions that are neither accumulable
100 /// nor hierarchical. Examples include jsonb_agg.
101 Basic,
102}
103
104impl TryFrom<&ReducePlan> for ReductionType {
105 type Error = ();
106
107 fn try_from(plan: &ReducePlan) -> Result<Self, Self::Error> {
108 match plan {
109 ReducePlan::Hierarchical(_) => Ok(ReductionType::Hierarchical),
110 ReducePlan::Accumulable(_) => Ok(ReductionType::Accumulable),
111 ReducePlan::Basic(_) => Ok(ReductionType::Basic),
112 ReducePlan::Distinct => Err(()),
113 }
114 }
115}
116
117/// A `ReducePlan` provides a concise description for how we will
118/// execute a given reduce expression.
119///
120/// The provided reduce expression can have no
121/// aggregations, in which case its just a `Distinct` and otherwise
122/// it's composed of a combination of accumulable, hierarchical and
123/// basic aggregations.
124///
125/// We want to try to centralize as much decision making about the
126/// shape / general computation of the rendered dataflow graph
127/// in this plan, and then make actually rendering the graph
128/// be as simple (and compiler verifiable) as possible.
129#[derive(Clone, Debug, Serialize, Deserialize, Eq, PartialEq, Ord, PartialOrd)]
130pub enum ReducePlan {
131 /// Plan for not computing any aggregations, just determining the set of
132 /// distinct keys.
133 Distinct,
134 /// Plan for computing only accumulable aggregations.
135 Accumulable(AccumulablePlan),
136 /// Plan for computing only hierarchical aggregations.
137 Hierarchical(HierarchicalPlan),
138 /// Plan for computing only basic aggregations.
139 Basic(BasicPlan),
140}
141
142/// All reduce plans depend on a notion of aggregation.
143///
144/// We could use `mz_expr::AggregateExpr`, but it explicitly names `MirScalarExpr`
145/// and derives `MzReflect` (which cannot accommodate type parameters).
146///
147/// We don't build a separate `AggregateFunc`, since we'd only eliminate one variant
148/// and need to duplicate the evaluation code.
149#[derive(
150 Clone,
151 Debug,
152 Eq,
153 PartialEq,
154 Ord,
155 PartialOrd,
156 Serialize,
157 Deserialize,
158 Hash
159)]
160pub struct LirAggregateExpr {
161 /// Names the aggregation function.
162 pub func: AggregateFunc,
163 /// An expression which extracts from each row the input to `func`.
164 pub expr: LirScalarExpr,
165 /// Should the aggregation be applied only to distinct results in each group.
166 #[serde(default)]
167 pub distinct: bool,
168}
169
170impl LirAggregateExpr {
171 /// Translates an aggregate from MIR to LIR.
172 ///
173 /// Panics on unmaterializable functions.
174 pub fn from_mir(mir: AggregateExpr) -> Self {
175 Self::try_from(mir).expect("no unmaterializable functions in aggregates")
176 }
177
178 /// Determines whether this aggregate is `COUNT(*)`.
179 pub fn is_count_asterisk(&self) -> bool {
180 self.func == AggregateFunc::Count && self.expr.is_literal_true() && !self.distinct
181 }
182}
183
184impl TryFrom<AggregateExpr> for LirAggregateExpr {
185 type Error = Vec<UnmaterializableFunc>;
186
187 fn try_from(mir: AggregateExpr) -> Result<Self, Self::Error> {
188 let func = mir.func;
189 let expr = LirScalarExpr::try_from(&mir.expr)?;
190 let distinct = mir.distinct;
191
192 Ok(LirAggregateExpr {
193 func,
194 expr,
195 distinct,
196 })
197 }
198}
199
200impl HumanizeDisplay for LirAggregateExpr {
201 fn humanize<'a, M: HumanizerMode>(
202 e: &HumanizedExpr<'a, Self, M>,
203 f: &mut std::fmt::Formatter<'_>,
204 ) -> std::fmt::Result {
205 if e.expr.is_count_asterisk() {
206 return write!(f, "count(*)");
207 }
208
209 write!(
210 f,
211 "{}({}",
212 e.child(&e.expr.func),
213 if e.expr.distinct { "distinct " } else { "" }
214 )?;
215
216 e.child(&e.expr.expr).fmt(f)?;
217 write!(f, ")")
218 }
219}
220
221/// Plan for computing a set of accumulable aggregations.
222///
223/// We fuse all of the accumulable aggregations together
224/// and compute them with one dataflow fragment. We need to
225/// be careful to separate out the aggregations that
226/// apply only to the distinct set of values. We need
227/// to apply a distinct operator to those before we
228/// combine them with everything else.
229#[derive(Clone, Debug, Serialize, Deserialize, Eq, PartialEq, Ord, PartialOrd)]
230pub struct AccumulablePlan {
231 /// All of the aggregations we were asked to compute, stored
232 /// in order.
233 pub full_aggrs: Vec<LirAggregateExpr>,
234 /// All of the non-distinct accumulable aggregates.
235 /// Each element represents:
236 /// (index of the datum among inputs, aggregation expr)
237 /// These will all be rendered together in one dataflow fragment.
238 pub simple_aggrs: Vec<(usize, LirAggregateExpr)>,
239 /// Same as above but for all of the `DISTINCT` accumulable aggregations.
240 pub distinct_aggrs: Vec<(usize, LirAggregateExpr)>,
241}
242
243/// Plan for computing a set of hierarchical aggregations.
244///
245/// In the append-only setting we can render them in-place
246/// with monotonic plans, but otherwise, we need to render
247/// them with a reduction tree that splits the inputs into
248/// small, and then progressively larger, buckets
249#[derive(Clone, Debug, Serialize, Deserialize, Eq, PartialEq, Ord, PartialOrd)]
250pub enum HierarchicalPlan {
251 /// Plan hierarchical aggregations under monotonic inputs.
252 Monotonic(MonotonicPlan),
253 /// Plan for hierarchical aggregations under non-monotonic inputs.
254 Bucketed(BucketedPlan),
255}
256
257impl HierarchicalPlan {
258 /// Returns the set of aggregations computed by this plan.
259 pub fn aggr_funcs(&self) -> &[AggregateFunc] {
260 match self {
261 HierarchicalPlan::Monotonic(plan) => &plan.aggr_funcs,
262 HierarchicalPlan::Bucketed(plan) => &plan.aggr_funcs,
263 }
264 }
265
266 /// Upgrades from a bucketed plan to a monotonic plan, if necessary,
267 /// and sets consolidation requirements.
268 pub fn as_monotonic(&mut self, must_consolidate: bool) {
269 match self {
270 HierarchicalPlan::Bucketed(bucketed) => {
271 // TODO: ideally we would not have the `clone()` but ownership
272 // seems fraught here as we are behind a `&mut self` reference.
273 *self =
274 HierarchicalPlan::Monotonic(bucketed.clone().into_monotonic(must_consolidate));
275 }
276 HierarchicalPlan::Monotonic(monotonic) => {
277 monotonic.must_consolidate = must_consolidate;
278 }
279 }
280 }
281}
282
283/// Plan for computing a set of hierarchical aggregations with a
284/// monotonic input.
285///
286/// Here, the aggregations will be rendered in place. We don't
287/// need to worry about retractions because the inputs are
288/// append only, so we can change our computation to
289/// only retain the "best" value in the diff field, instead
290/// of holding onto all values.
291#[derive(Clone, Debug, Serialize, Deserialize, Eq, PartialEq, Ord, PartialOrd)]
292pub struct MonotonicPlan {
293 /// All of the aggregations we were asked to compute.
294 pub aggr_funcs: Vec<AggregateFunc>,
295 /// True if the input is not physically monotonic, and the operator must perform
296 /// consolidation to remove potential negations. The operator implementation is
297 /// free to consolidate as late as possible while ensuring correctness, so it is
298 /// not a requirement that the input be directly subjected to consolidation.
299 /// More details in the monotonic one-shot `SELECT`s design doc.[^1]
300 ///
301 /// [^1]: <https://github.com/MaterializeInc/materialize/blob/main/doc/developer/design/20230421_stabilize_monotonic_select.md>
302 pub must_consolidate: bool,
303}
304
305/// Plan for computing a set of hierarchical aggregations
306/// with non-monotonic inputs.
307///
308/// To perform hierarchical aggregations with stable runtimes
309/// under updates we'll subdivide the group key into buckets, compute
310/// the reduction in each of those subdivided buckets and then combine
311/// the results into a coarser bucket (one that represents a larger
312/// fraction of the original input) and redo the reduction in another
313/// layer. Effectively, we'll construct a min / max heap out of a series
314/// of reduce operators (each one is a separate layer).
315#[derive(Clone, Debug, Serialize, Deserialize, Eq, PartialEq, Ord, PartialOrd)]
316pub struct BucketedPlan {
317 /// All of the aggregations we were asked to compute.
318 pub aggr_funcs: Vec<AggregateFunc>,
319 /// The number of buckets in each layer of the reduction tree. Should
320 /// be decreasing, and ideally, a power of two so that we can easily
321 /// distribute values to buckets with `value.hashed() % buckets[layer]`.
322 pub buckets: Vec<u64>,
323}
324
325impl BucketedPlan {
326 /// Convert to a monotonic plan, indicate whether the operator must apply
327 /// consolidation to its input.
328 fn into_monotonic(self, must_consolidate: bool) -> MonotonicPlan {
329 MonotonicPlan {
330 aggr_funcs: self.aggr_funcs,
331 must_consolidate,
332 }
333 }
334}
335
336/// Plan for computing a set of basic aggregations.
337///
338/// There's much less complexity when rendering basic aggregations.
339/// Each aggregation corresponds to one Differential reduce operator.
340/// That's it. However, we still want to present one final arrangement
341/// so basic aggregations present results with the same interface
342/// (one arrangement containing a row with all results) that accumulable
343/// and hierarchical aggregations do. To provide that, we render an
344/// additional reduce operator whenever we have multiple reduce aggregates
345/// to combine and present results in the appropriate order. If we
346/// were only asked to compute a single aggregation, we can skip
347/// that step and return the arrangement provided by computing the aggregation
348/// directly.
349#[derive(Clone, Debug, Serialize, Deserialize, Eq, PartialEq, Ord, PartialOrd)]
350pub enum BasicPlan {
351 /// Plan for rendering a single basic aggregation.
352 Single(SingleBasicPlan),
353 /// Plan for rendering multiple basic aggregations.
354 /// These need to then be collated together in an additional
355 /// reduction. Each element represents the:
356 /// `(index of the set of the input we are aggregating over,
357 /// the aggregation function)`
358 Multiple(Vec<LirAggregateExpr>),
359}
360
361/// Plan for rendering a single basic aggregation, with possibly fusing a `FlatMap UnnestList` with
362/// this aggregation.
363#[derive(Clone, Debug, Serialize, Deserialize, Eq, PartialEq, Ord, PartialOrd)]
364pub struct SingleBasicPlan {
365 /// The aggregation that we should perform.
366 pub expr: LirAggregateExpr,
367 /// Whether we fused a `FlatMap UnnestList` with this aggregation.
368 pub fused_unnest_list: bool,
369}
370
371/// Plan for collating the results of computing multiple aggregation
372/// types.
373///
374/// TODO: could we express this as a delta join
375#[derive(
376 Clone,
377 Debug,
378 Default,
379 Serialize,
380 Deserialize,
381 Eq,
382 PartialEq,
383 Ord,
384 PartialOrd
385)]
386pub struct CollationPlan {
387 /// Accumulable aggregation results to collate, if any.
388 pub accumulable: Option<AccumulablePlan>,
389 /// Hierarchical aggregation results to collate, if any.
390 pub hierarchical: Option<HierarchicalPlan>,
391 /// Basic aggregation results to collate, if any.
392 pub basic: Option<BasicPlan>,
393 /// When we get results back from each of the different
394 /// aggregation types, they will be subsequences of
395 /// the sequence aggregations in the original reduce expression.
396 /// We keep a map from output position -> reduction type
397 /// to easily merge results back into the requested order.
398 pub aggregate_types: Vec<ReductionType>,
399}
400
401impl CollationPlan {
402 /// Upgrades the hierarchical component of the collation plan to monotonic, if necessary,
403 /// and sets consolidation requirements.
404 pub fn as_monotonic(&mut self, must_consolidate: bool) {
405 self.hierarchical
406 .as_mut()
407 .map(|plan| plan.as_monotonic(must_consolidate));
408 }
409}
410
411impl ReducePlan {
412 /// Generate a plan for computing the supplied aggregations.
413 ///
414 /// The resulting plan summarizes what the dataflow to be created
415 /// and how the aggregations will be executed.
416 pub fn create_from(
417 aggregates: Vec<AggregateExpr>,
418 monotonic: bool,
419 expected_group_size: Option<u64>,
420 fused_unnest_list: bool,
421 ) -> Self {
422 // We need to make sure that all aggregates have the same type.
423 let mut aggregates_list = Vec::with_capacity(aggregates.len());
424 let mut aggregates = aggregates.into_iter();
425 if let Some(aggregate) = aggregates.next() {
426 let typ = reduction_type(&aggregate.func);
427 aggregates_list.push(LirAggregateExpr::from_mir(aggregate));
428
429 for aggregate in aggregates {
430 assert_eq!(
431 typ,
432 reduction_type(&aggregate.func),
433 "Multiple reduction types detected"
434 );
435 aggregates_list.push(LirAggregateExpr::from_mir(aggregate));
436 }
437 ReducePlan::create_inner(
438 typ,
439 aggregates_list,
440 monotonic,
441 expected_group_size,
442 fused_unnest_list,
443 )
444 } else {
445 // If we don't have any aggregations we are just computing a distinct.
446 ReducePlan::Distinct
447 }
448 }
449
450 /// Generate a plan for computing the specified type of aggregations.
451 ///
452 /// This function assumes that all of the supplied aggregates are
453 /// actually of the correct reduction type.
454 fn create_inner(
455 typ: ReductionType,
456 aggregates_list: Vec<LirAggregateExpr>,
457 monotonic: bool,
458 expected_group_size: Option<u64>,
459 fused_unnest_list: bool,
460 ) -> Self {
461 if fused_unnest_list {
462 assert!(matches!(typ, ReductionType::Basic) && aggregates_list.len() == 1);
463 }
464 assert!(
465 aggregates_list.len() > 0,
466 "error: tried to render a reduce dataflow with no aggregates"
467 );
468
469 match typ {
470 ReductionType::Accumulable => {
471 let mut simple_aggrs = vec![];
472 let mut distinct_aggrs = vec![];
473 let full_aggrs = aggregates_list.clone();
474 for (datum_index, aggr) in aggregates_list.into_iter().enumerate() {
475 // Accumulable aggregations need to do extra per-aggregate work
476 // for aggregations with the distinct bit set, so we'll separate
477 // those out now.
478 if aggr.distinct {
479 distinct_aggrs.push((datum_index, aggr));
480 } else {
481 simple_aggrs.push((datum_index, aggr));
482 };
483 }
484 ReducePlan::Accumulable(AccumulablePlan {
485 full_aggrs,
486 simple_aggrs,
487 distinct_aggrs,
488 })
489 }
490 ReductionType::Hierarchical => {
491 let aggr_funcs = aggregates_list
492 .iter()
493 .map(|aggr| aggr.func.clone())
494 .collect();
495
496 if monotonic {
497 let monotonic = MonotonicPlan {
498 aggr_funcs,
499 must_consolidate: false,
500 };
501 ReducePlan::Hierarchical(HierarchicalPlan::Monotonic(monotonic))
502 } else {
503 let buckets = bucketing_of_expected_group_size(expected_group_size);
504 let bucketed = BucketedPlan {
505 aggr_funcs,
506 buckets,
507 };
508
509 ReducePlan::Hierarchical(HierarchicalPlan::Bucketed(bucketed))
510 }
511 }
512 ReductionType::Basic => match <_ as TryInto<[_; 1]>>::try_into(aggregates_list) {
513 Ok([expr]) => ReducePlan::Basic(BasicPlan::Single(SingleBasicPlan {
514 expr,
515 fused_unnest_list,
516 })),
517 Err(aggregates_list) => ReducePlan::Basic(BasicPlan::Multiple(aggregates_list)),
518 },
519 }
520 }
521
522 /// Reports all keys of produced arrangements.
523 ///
524 /// This is likely either an empty vector, for no arrangement,
525 /// or a singleton vector containing the list of expressions
526 /// that key a single arrangement.
527 pub fn keys(&self, key_arity: usize, arity: usize) -> AvailableCollections {
528 let key = (0..key_arity)
529 .map(LirScalarExpr::column)
530 .collect::<Vec<_>>();
531 let (permutation, thinning) = permutation_for_arrangement(&key, arity);
532 AvailableCollections::new_arranged(vec![(key, permutation, thinning)])
533 }
534
535 /// Extracts a fusable MFP for the reduction from the given `mfp` along with a residual
536 /// non-fusable MFP and potentially revised output arity. The provided `mfp` must be the
537 /// one sitting on top of the reduction.
538 ///
539 /// Non-fusable parts include temporal predicates or any other parts that cannot be
540 /// conservatively asserted to not increase the memory requirements of the output
541 /// arrangement for the reduction. Either the fusable or non-fusable parts may end up
542 /// being the identity MFP.
543 pub fn extract_mfp_after(
544 &self,
545 mut mfp: MapFilterProject,
546 key_arity: usize,
547 ) -> (MapFilterProject, MapFilterProject, usize) {
548 // Extract temporal predicates, as we cannot push them into `Reduce`.
549 let temporal_mfp = mfp.extract_temporal();
550 let non_temporal = mfp;
551 mfp = temporal_mfp;
552
553 // We ensure we do not attempt to project away the key, as we cannot accomplish
554 // this. This is done by a simple analysis of the non-temporal part of `mfp` to
555 // check if can be directly absorbed; if it can't, we then default to a general
556 // strategy that unpacks the MFP to absorb only the filter and supporting map
557 // parts, followed by a post-MFP step.
558 let input_arity = non_temporal.input_arity;
559 let key = Vec::from_iter(0..key_arity);
560 let mut mfp_push;
561 let output_arity;
562
563 if non_temporal.projection.len() <= input_arity
564 && non_temporal.projection.iter().all(|c| *c < input_arity)
565 && non_temporal.projection.starts_with(&key)
566 {
567 // Special case: The key is preserved as a prefix and the projection is only
568 // of output fields from the reduction. So we know that: (a) We can process the
569 // fused MFP per-key; (b) The MFP application gets rid of all mapped columns;
570 // and (c) The output projection is at most as wide as the output that would be
571 // produced by the reduction, so we are sure to never regress the memory
572 // requirements of the output arrangement.
573 // Note that this strategy may change the arity of the output arrangement.
574 output_arity = non_temporal.projection.len();
575 mfp_push = non_temporal;
576 } else {
577 // General strategy: Unpack MFP as MF followed by P' that removes all M
578 // columns, then MP afterwards.
579 // Note that this strategy does not result in any changes to the arity of
580 // the output arrangement.
581 let (m, f, p) = non_temporal.into_map_filter_project();
582 mfp_push = MapFilterProject::new(input_arity)
583 .map(m.clone())
584 .filter(f)
585 .project(0..input_arity);
586 output_arity = input_arity;
587
588 // We still need to perform the map and projection for the actual output.
589 let mfp_left = MapFilterProject::new(input_arity).map(m).project(p);
590
591 // Compose the non-pushed MFP components.
592 mfp = MapFilterProject::compose(mfp_left, mfp);
593 }
594 mfp_push.optimize();
595 mfp.optimize();
596 (mfp_push, mfp, output_arity)
597 }
598}
599
600/// Plan for extracting keys and values in preparation for a reduction.
601#[derive(Clone, Debug, Serialize, Deserialize, Eq, PartialEq, Ord, PartialOrd)]
602pub struct KeyValPlan {
603 /// Extracts the columns used as the key.
604 pub key_plan: mz_expr::SafeMfpPlan<LirScalarExpr>,
605 /// Extracts the columns used to feed the aggregations.
606 pub val_plan: mz_expr::SafeMfpPlan<LirScalarExpr>,
607}
608
609impl KeyValPlan {
610 /// Create a new [KeyValPlan] from aggregation arguments.
611 pub fn new(
612 input_arity: usize,
613 group_key: &[MirScalarExpr],
614 aggregates: &[AggregateExpr],
615 input_permutation_and_new_arity: Option<(Vec<usize>, usize)>,
616 ) -> Self {
617 // Form an operator for evaluating key expressions.
618 let mut key_mfp = MapFilterProject::new(input_arity)
619 .map(group_key.iter().cloned())
620 .project(input_arity..(input_arity + group_key.len()));
621 if let Some((input_permutation, new_arity)) = input_permutation_and_new_arity.clone() {
622 key_mfp.permute_fn(|c| input_permutation[c], new_arity);
623 }
624
625 // Form an operator for evaluating value expressions.
626 let mut val_mfp = MapFilterProject::new(input_arity)
627 .map(aggregates.iter().map(|a| a.expr.clone()))
628 .project(input_arity..(input_arity + aggregates.len()));
629 if let Some((input_permutation, new_arity)) = input_permutation_and_new_arity {
630 val_mfp.permute_fn(|c| input_permutation[c], new_arity);
631 }
632
633 key_mfp.optimize();
634 let key_plan = crate::plan::scalar::safe_mfp_mir_to_lir(
635 key_mfp.into_plan().unwrap().into_nontemporal().unwrap(),
636 );
637 val_mfp.optimize();
638 let val_plan = crate::plan::scalar::safe_mfp_mir_to_lir(
639 val_mfp.into_plan().unwrap().into_nontemporal().unwrap(),
640 );
641
642 Self { key_plan, val_plan }
643 }
644
645 /// The arity of the key plan
646 pub fn key_arity(&self) -> usize {
647 self.key_plan.projection.len()
648 }
649}
650
651/// Transforms a vector containing indexes of needed columns into one containing
652/// the "skips" an iterator over a Row would need to perform to see those values.
653///
654/// This function requires that all of the elements in `indexes` are strictly
655/// increasing.
656///
657/// # Examples
658///
659/// ```
660/// use mz_compute_types::plan::reduce::convert_indexes_to_skips;
661/// assert_eq!(convert_indexes_to_skips(vec![3, 6, 10, 15]), [3, 2, 3, 4])
662/// ```
663pub fn convert_indexes_to_skips(mut indexes: Vec<usize>) -> Vec<usize> {
664 for i in 1..indexes.len() {
665 soft_assert_or_log!(
666 indexes[i - 1] < indexes[i],
667 "convert_indexes_to_skip needs indexes to be strictly increasing. Received: {:?}",
668 indexes,
669 );
670 }
671
672 for i in (1..indexes.len()).rev() {
673 indexes[i] -= indexes[i - 1];
674 indexes[i] -= 1;
675 }
676
677 indexes
678}
679
680/// Determines whether a function can be accumulated in an update's "difference" field,
681/// and whether it can be subjected to recursive (hierarchical) aggregation.
682///
683/// Accumulable aggregations will be packed into differential dataflow's "difference" field,
684/// which can be accumulated in-place using the addition operation on the type. Aggregations
685/// that indicate they are accumulable will still need to provide an action that takes their
686/// data and introduces it as a difference, and the post-processing when the accumulated value
687/// is presented as data.
688///
689/// Hierarchical aggregations will be subjected to repeated aggregation on initially small but
690/// increasingly large subsets of each key. This has the intended property that no invocation
691/// is on a significantly large set of values (and so, no incremental update needs to reform
692/// significant input data). Hierarchical aggregates can be rendered more efficiently if the
693/// input stream is append-only as then we only need to retain the "currently winning" value.
694/// Every hierarchical aggregate needs to supply a corresponding ReductionMonoid implementation.
695pub fn reduction_type(func: &AggregateFunc) -> ReductionType {
696 match func {
697 AggregateFunc::SumInt16
698 | AggregateFunc::SumInt32
699 | AggregateFunc::SumInt64
700 | AggregateFunc::SumUInt16
701 | AggregateFunc::SumUInt32
702 | AggregateFunc::SumUInt64
703 | AggregateFunc::SumFloat32
704 | AggregateFunc::SumFloat64
705 | AggregateFunc::SumNumeric
706 | AggregateFunc::Count
707 | AggregateFunc::Any
708 | AggregateFunc::All
709 | AggregateFunc::Dummy => ReductionType::Accumulable,
710 AggregateFunc::MaxNumeric
711 | AggregateFunc::MaxInt16
712 | AggregateFunc::MaxInt32
713 | AggregateFunc::MaxInt64
714 | AggregateFunc::MaxUInt16
715 | AggregateFunc::MaxUInt32
716 | AggregateFunc::MaxUInt64
717 | AggregateFunc::MaxMzTimestamp
718 | AggregateFunc::MaxFloat32
719 | AggregateFunc::MaxFloat64
720 | AggregateFunc::MaxBool
721 | AggregateFunc::MaxString
722 | AggregateFunc::MaxDate
723 | AggregateFunc::MaxTimestamp
724 | AggregateFunc::MaxTimestampTz
725 | AggregateFunc::MaxInterval
726 | AggregateFunc::MaxTime
727 | AggregateFunc::MinNumeric
728 | AggregateFunc::MinInt16
729 | AggregateFunc::MinInt32
730 | AggregateFunc::MinInt64
731 | AggregateFunc::MinUInt16
732 | AggregateFunc::MinUInt32
733 | AggregateFunc::MinUInt64
734 | AggregateFunc::MinMzTimestamp
735 | AggregateFunc::MinInterval
736 | AggregateFunc::MinFloat32
737 | AggregateFunc::MinFloat64
738 | AggregateFunc::MinBool
739 | AggregateFunc::MinString
740 | AggregateFunc::MinDate
741 | AggregateFunc::MinTimestamp
742 | AggregateFunc::MinTimestampTz
743 | AggregateFunc::MinTime => ReductionType::Hierarchical,
744 AggregateFunc::JsonbAgg { .. }
745 | AggregateFunc::JsonbObjectAgg { .. }
746 | AggregateFunc::MapAgg { .. }
747 | AggregateFunc::ArrayConcat { .. }
748 | AggregateFunc::ListConcat { .. }
749 | AggregateFunc::StringAgg { .. }
750 | AggregateFunc::RowNumber { .. }
751 | AggregateFunc::Rank { .. }
752 | AggregateFunc::DenseRank { .. }
753 | AggregateFunc::LagLead { .. }
754 | AggregateFunc::FirstValue { .. }
755 | AggregateFunc::LastValue { .. }
756 | AggregateFunc::WindowAggregate { .. }
757 | AggregateFunc::FusedValueWindowFunc { .. }
758 | AggregateFunc::FusedWindowAggregate { .. } => ReductionType::Basic,
759 }
760}