mz_timely_util/containers/
stack.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
16//! A chunked columnar container based on the columnation library. It stores the local
17//! portion in region-allocated data, too, which is different to the `TimelyStack` type.
18
19use std::cell::Cell;
20
21use differential_dataflow::containers::{Columnation, TimelyStack};
22use timely::container::{ContainerBuilder, PushInto};
23
24/// A Stacked container builder that keep track of container memory usage.
25#[derive(Default)]
26pub struct AccountedStackBuilder<CB> {
27    pub bytes: Cell<usize>,
28    pub builder: CB,
29}
30
31impl<T, CB> ContainerBuilder for AccountedStackBuilder<CB>
32where
33    T: Clone + Columnation + 'static,
34    CB: ContainerBuilder<Container = TimelyStack<T>>,
35{
36    type Container = TimelyStack<T>;
37
38    fn extract(&mut self) -> Option<&mut Self::Container> {
39        let container = self.builder.extract()?;
40        let mut new_bytes = 0;
41        container.heap_size(|_, cap| new_bytes += cap);
42        self.bytes.set(self.bytes.get() + new_bytes);
43        Some(container)
44    }
45
46    fn finish(&mut self) -> Option<&mut Self::Container> {
47        let container = self.builder.finish()?;
48        let mut new_bytes = 0;
49        container.heap_size(|_, cap| new_bytes += cap);
50        self.bytes.set(self.bytes.get() + new_bytes);
51        Some(container)
52    }
53}
54
55impl<T, CB: PushInto<T>> PushInto<T> for AccountedStackBuilder<CB> {
56    #[inline]
57    fn push_into(&mut self, item: T) {
58        self.builder.push_into(item);
59    }
60}
mz_timely_util/containers/stack.rs

mz_timely_util/containers/
stack.rs
