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 `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}