mz_compute/extensions/
reduce.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License in the LICENSE file at the
6// root of this repository, or online at
7//
8//     http://www.apache.org/licenses/LICENSE-2.0
9//
10// Unless required by applicable law or agreed to in writing, software
11// distributed under the License is distributed on an "AS IS" BASIS,
12// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13// See the License for the specific language governing permissions and
14// limitations under the License.
15
16use columnation::Columnation;
17use differential_dataflow::Data;
18use differential_dataflow::difference::Abelian;
19use differential_dataflow::operators::arrange::{Arranged, TraceAgent};
20use differential_dataflow::trace::cursor::{BatchCursor, BatchDiff, BatchVal, BatchValOwn};
21use differential_dataflow::trace::implementations::BatchContainer;
22use differential_dataflow::trace::{Builder, Cursor, Navigable, Trace, TraceReader};
23use mz_timely_util::columnation::ColumnationStack;
24use timely::Container;
25use timely::container::PushInto;
26
27use crate::extensions::arrange::ArrangementSize;
28
29pub trait ClearContainer {
30    fn clear(&mut self);
31}
32
33impl<T> ClearContainer for Vec<T> {
34    fn clear(&mut self) {
35        Vec::clear(self)
36    }
37}
38
39impl<D, T, R> ClearContainer for ColumnationStack<(D, T, R)>
40where
41    D: Columnation + Clone + 'static,
42    T: Columnation + Clone + 'static,
43    R: Columnation + Clone + 'static,
44{
45    fn clear(&mut self) {
46        ColumnationStack::clear(self)
47    }
48}
49
50/// Extension trait for the `reduce_abelian` differential dataflow method.
51pub(crate) trait MzReduce<'scope, T1: TraceReader> {
52    /// Applies `reduce` to arranged data, and returns an arrangement of output data.
53    fn mz_reduce_abelian<L, Bu, T2, KC>(
54        self,
55        name: &str,
56        logic: L,
57    ) -> Arranged<'scope, TraceAgent<T2>>
58    where
59        // `KC` is a distinct parameter, not `T1::KeyContainer`: pinning it from the concrete input
60        // lets the compiler normalize the higher-ranked key equality on both cursors to
61        // `KC::ReadItem`, which a projection-to-projection bound would not.
62        T1: TraceReader<Batch: Navigable>,
63        KC: BatchContainer,
64        BatchCursor<T1>: Cursor<Time = T1::Time, KeyContainer = KC>,
65        for<'a> BatchCursor<T1>: Cursor<Key<'a> = KC::ReadItem<'a>>,
66        T2: Trace<Batch: Navigable, Time = T1::Time> + 'static,
67        BatchCursor<T2>: Cursor<Time = T2::Time, KeyContainer = KC>,
68        for<'a> BatchCursor<T2>:
69            Cursor<Key<'a> = KC::ReadItem<'a>, ValOwn: Data, Time = T2::Time, Diff: Abelian>,
70        Bu: Builder<Time = T1::Time, Output = T2::Batch> + 'static,
71        Bu::Input: Container
72            + Default
73            + ClearContainer
74            + PushInto<((KC::Owned, BatchValOwn<T2>), T2::Time, BatchDiff<T2>)>,
75        L: FnMut(
76                KC::ReadItem<'_>,
77                &[(BatchVal<'_, T1>, BatchDiff<T1>)],
78                &mut Vec<(BatchValOwn<T2>, BatchDiff<T2>)>,
79            ) + 'static,
80        Arranged<'scope, TraceAgent<T2>>: ArrangementSize;
81}
82
83impl<'scope, T1> MzReduce<'scope, T1> for Arranged<'scope, T1>
84where
85    T1: TraceReader + Clone + 'static,
86{
87    /// Applies `reduce` to arranged data, and returns an arrangement of output data.
88    fn mz_reduce_abelian<L, Bu, T2, KC>(
89        self,
90        name: &str,
91        logic: L,
92    ) -> Arranged<'scope, TraceAgent<T2>>
93    where
94        T1: TraceReader<Batch: Navigable>,
95        KC: BatchContainer,
96        BatchCursor<T1>: Cursor<Time = T1::Time, KeyContainer = KC>,
97        for<'a> BatchCursor<T1>: Cursor<Key<'a> = KC::ReadItem<'a>>,
98        T2: Trace<Batch: Navigable, Time = T1::Time> + 'static,
99        BatchCursor<T2>: Cursor<Time = T2::Time, KeyContainer = KC>,
100        for<'a> BatchCursor<T2>:
101            Cursor<Key<'a> = KC::ReadItem<'a>, ValOwn: Data, Time = T2::Time, Diff: Abelian>,
102        Bu: Builder<Time = T1::Time, Output = T2::Batch> + 'static,
103        Bu::Input: Container
104            + Default
105            + ClearContainer
106            + PushInto<((KC::Owned, BatchValOwn<T2>), T2::Time, BatchDiff<T2>)>,
107        L: FnMut(
108                KC::ReadItem<'_>,
109                &[(BatchVal<'_, T1>, BatchDiff<T1>)],
110                &mut Vec<(BatchValOwn<T2>, BatchDiff<T2>)>,
111            ) + 'static,
112        Arranged<'scope, TraceAgent<T2>>: ArrangementSize,
113    {
114        // The push closure clears its buffer between keys (the reduce operator cannot reset it,
115        // and failing to clear leaks one key's rows into the next), then stages each value update
116        // with the owned key prepended.
117        let push_closure =
118            |buf: &mut Bu::Input,
119             key: KC::ReadItem<'_>,
120             updates: &mut Vec<(BatchValOwn<T2>, T2::Time, BatchDiff<T2>)>| {
121                buf.clear();
122                let key_owned = KC::into_owned(key);
123                for (val, time, diff) in updates.drain(..) {
124                    buf.push_into(((key_owned.clone(), val), time, diff));
125                }
126            };
127
128        // Allow access to `reduce_abelian` since we're within Mz's wrapper and force arrangement
129        // size logging.
130        #[allow(clippy::disallowed_methods)]
131        Arranged::<_>::reduce_abelian::<_, Bu, T2, KC, _>(self, name, logic, push_closure)
132            .log_arrangement_size()
133    }
134}
mz_compute/extensions/reduce.rs

mz_compute/extensions/
reduce.rs
