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

mz_timely_util/containers/
stack.rs
