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

mz_compute/render/
threshold.rs
