mz_compute/render/
threshold.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//! Threshold execution logic.
11//!
12//! Consult [ThresholdPlan] documentation for details.
13
14use differential_dataflow::operators::arrange::Arranged;
15use differential_dataflow::trace::Cursor;
16use differential_dataflow::trace::cursor::BatchCursor;
17use mz_compute_types::plan::scalar::LirScalarExpr;
18use mz_compute_types::plan::threshold::{BasicThresholdPlan, ThresholdPlan};
19use mz_repr::{Diff, Row, Timestamp};
20use mz_row_spine::{DatumSeq, RowRowBuilder};
21
22use crate::extensions::arrange::{KeyCollection, MzArrange};
23use crate::extensions::reduce::MzReduce;
24use crate::render::RenderTimestamp;
25use crate::render::context::{ArrangementFlavor, CollectionBundle, Context};
26use crate::typedefs::{ErrBatcher, RowRowAgent, RowRowEnter, RowRowSpine};
27
28/// Thresholds a dataflow-local ok arrangement, keeping rows with a positive count.
29///
30/// The reduce goes through a concrete-spine helper because `reduce_abelian`'s higher-ranked
31/// output-key bound only normalizes when the input and output trace types are concrete through a
32/// function signature.
33fn threshold_local<'scope, T: RenderTimestamp>(
34    arrangement: Arranged<'scope, RowRowAgent<T, Diff>>,
35    name: &str,
36) -> Arranged<'scope, RowRowAgent<T, Diff>> {
37    arrangement.mz_reduce_abelian::<_, RowRowBuilder<_, _>, RowRowSpine<_, _>, _>(
38        name,
39        move |_key, s, t| {
40            for (record, count) in s.iter() {
41                if count.is_positive() {
42                    t.push((
43                        <BatchCursor<RowRowSpine<T, Diff>> as Cursor>::owned_val(*record),
44                        *count,
45                    ));
46                }
47            }
48        },
49    )
50}
51
52/// Like [`threshold_local`] but over an imported trace's ok arrangement.
53fn threshold_trace<'scope, T: RenderTimestamp>(
54    arrangement: Arranged<'scope, RowRowEnter<Timestamp, Diff, T>>,
55    name: &str,
56) -> Arranged<'scope, RowRowAgent<T, Diff>> {
57    let logic = move |_key: DatumSeq<'_>, s: &[(DatumSeq<'_>, Diff)], t: &mut Vec<(Row, Diff)>| {
58        for (record, count) in s.iter() {
59            if count.is_positive() {
60                t.push((
61                    <BatchCursor<RowRowSpine<T, Diff>> as Cursor>::owned_val(*record),
62                    *count,
63                ));
64            }
65        }
66    };
67    arrangement.mz_reduce_abelian::<_, RowRowBuilder<T, Diff>, RowRowSpine<T, Diff>, _>(name, logic)
68}
69
70/// Build a dataflow to threshold the input data.
71///
72/// This implementation maintains rows in the output, i.e. all rows that have a count greater than
73/// zero. It returns a [CollectionBundle] populated from a local arrangement.
74pub fn build_threshold_basic<'scope, T: RenderTimestamp>(
75    input: CollectionBundle<'scope, T>,
76    key: Vec<LirScalarExpr>,
77) -> CollectionBundle<'scope, T> {
78    let arrangement = input
79        .arrangement(&key)
80        .expect("Arrangement ensured to exist");
81    match arrangement {
82        ArrangementFlavor::Local(oks, errs) => {
83            let oks = threshold_local(oks, "Threshold local");
84            CollectionBundle::from_expressions(key, ArrangementFlavor::Local(oks, errs))
85        }
86        ArrangementFlavor::Trace(_, oks, errs) => {
87            let oks = threshold_trace(oks, "Threshold trace");
88            let errs: KeyCollection<_, _, _> = errs.as_collection(|k, _| k.clone()).into();
89            let errs = errs.mz_arrange::<ErrBatcher<_, _>, _>("Arrange threshold basic err");
90            CollectionBundle::from_expressions(key, ArrangementFlavor::Local(oks, errs))
91        }
92    }
93}
94
95impl<'scope, T: RenderTimestamp> Context<'scope, T> {
96    pub(crate) fn render_threshold(
97        &self,
98        input: CollectionBundle<'scope, T>,
99        threshold_plan: ThresholdPlan,
100    ) -> CollectionBundle<'scope, T> {
101        match threshold_plan {
102            ThresholdPlan::Basic(BasicThresholdPlan {
103                ensure_arrangement: (key, _, _),
104            }) => {
105                // We do not need to apply the permutation here,
106                // since threshold doesn't inspect the values, but only
107                // their counts.
108                build_threshold_basic(input, key)
109            }
110        }
111    }
112}
mz_compute/render/threshold.rs

mz_compute/render/
threshold.rs
