Skip to main content

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}