Skip to main content

mz_timely_util/columnar/
builder.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 container builder for columns.
17
18use std::collections::VecDeque;
19
20use columnar::bytes::indexed;
21use columnar::{Clear, Columnar, Len, Push};
22use timely::container::PushInto;
23use timely::container::{ContainerBuilder, LengthPreservingContainerBuilder};
24
25use crate::columnar::Column;
26
27/// A container builder for `Column<C>`.
28pub struct ColumnBuilder<C: Columnar> {
29    /// Container that we're writing to.
30    current: C::Container,
31    /// Finished container that we presented to callers of extract/finish.
32    ///
33    /// We don't recycle the column because for extract, it's not typed, and after calls
34    /// to finish it'll be `None`.
35    finished: Option<Column<C>>,
36    /// Completed containers pending to be sent.
37    pending: VecDeque<Column<C>>,
38}
39
40impl<C: Columnar, T> PushInto<T> for ColumnBuilder<C>
41where
42    C::Container: Push<T>,
43{
44    #[inline]
45    fn push_into(&mut self, item: T) {
46        self.current.push(item);
47        // Mint a container once the serialized size reaches the ship threshold.
48        use columnar::Borrow;
49        if crate::columnar::at_serialized_capacity(&self.current.borrow()) {
50            /// Move the contents from `current` to a `Vec<u64>` allocation built via
51            /// `indexed::encode` (so no zero-init pre-pass), and push it to `pending`.
52            #[cold]
53            fn outlined_align<C>(current: &mut C::Container, pending: &mut VecDeque<Column<C>>)
54            where
55                C: Columnar,
56            {
57                use columnar::Borrow;
58                let words = indexed::length_in_words(&current.borrow());
59                let mut alloc: Vec<u64> = Vec::with_capacity(words);
60                indexed::encode(&mut alloc, &current.borrow());
61                pending.push_back(Column::Align(alloc));
62                current.clear();
63            }
64
65            outlined_align(&mut self.current, &mut self.pending);
66        }
67    }
68}
69
70impl<C: Columnar> Default for ColumnBuilder<C> {
71    #[inline]
72    fn default() -> Self {
73        ColumnBuilder {
74            current: Default::default(),
75            finished: None,
76            pending: Default::default(),
77        }
78    }
79}
80
81impl<C: Columnar> ContainerBuilder for ColumnBuilder<C>
82where
83    C::Container: Clone,
84{
85    type Container = Column<C>;
86
87    #[inline]
88    fn extract(&mut self) -> Option<&mut Self::Container> {
89        if let Some(container) = self.pending.pop_front() {
90            self.finished = Some(container);
91            self.finished.as_mut()
92        } else {
93            None
94        }
95    }
96
97    #[inline]
98    fn finish(&mut self) -> Option<&mut Self::Container> {
99        if !self.current.is_empty() {
100            self.pending
101                .push_back(Column::Typed(std::mem::take(&mut self.current)));
102        }
103        self.finished = self.pending.pop_front();
104        self.finished.as_mut()
105    }
106}
107
108impl<C: Columnar> LengthPreservingContainerBuilder for ColumnBuilder<C> where C::Container: Clone {}