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