mz_timely_util/
columnar.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//! Container for columnar data.
17
18#![deny(missing_docs)]
19
20pub mod batcher;
21pub mod builder;
22pub mod builder_input;
23pub mod consolidate;
24pub mod merge_batcher;
25pub mod unload;
26
27use std::hash::Hash;
28
29use columnar::Borrow;
30use columnar::bytes::indexed;
31use columnar::common::IterOwn;
32use columnar::{Clear, FromBytes, Index, Len};
33use columnar::{Columnar, Ref};
34use differential_dataflow::Hashable;
35use differential_dataflow::trace::implementations::merge_batcher::MergeBatcher;
36use timely::Accountable;
37use timely::bytes::arc::Bytes;
38use timely::container::{DrainContainer, PushInto, SizableContainer};
39use timely::dataflow::channels::ContainerBytes;
40
41use crate::columnation::ColInternalMerger;
42
43/// A batcher for columnar storage.
44///
45/// The chunker is supplied to the arrange operator separately. Callers pass
46/// it explicitly: [`ColumnationChunker`](crate::columnation::ColumnationChunker)
47/// for `Vec<_>` input, or [`batcher::Chunker`] (over a `ColumnationStack<_>`) for
48/// [`Column`] input.
49pub type Col2ValBatcher<K, V, T, R> = MergeBatcher<ColInternalMerger<(K, V), T, R>>;
50/// A batcher for columnar storage with unit values.
51pub type Col2KeyBatcher<K, T, R> = Col2ValBatcher<K, (), T, R>;
52
53/// Pageable counterpart to [`Col2ValBatcher`]. Routes every chunk produced
54/// by chunking, merging, or extract through a [`crate::column_pager::ColumnPager`],
55/// so memory pressure can spill chains to a backing store without touching
56/// the merge / extract bodies.
57///
58/// Drop-in shape at the type level: both aliases take `(K, V, T, R)` and
59/// produce a `Batcher<Input = Column<((K, V), T, R)>, Output = Column<((K,
60/// V), T, R)>>`. Call sites can swap with `cargo fix`–style renaming once
61/// downstream `Trace`/`Builder` impls have been wired up. The pager itself
62/// defaults to [`crate::column_pager::ColumnPager::disabled`]; inject a
63/// real one via [`merge_batcher::ColumnMergeBatcher::set_pager`].
64pub type Col2ValPagedBatcher<K, V, T, R> = merge_batcher::ColumnMergeBatcher<(K, V), T, R>;
65
66/// A container based on a columnar store, encoded in aligned bytes.
67///
68/// The type can represent typed data, bytes from Timely, or an aligned allocation. The name
69/// is singular to express that the preferred format is [`Column::Align`]. The [`Column::Typed`]
70/// variant is used to construct the container, and it owns potentially multiple columns of data.
71pub enum Column<C: Columnar> {
72    /// The typed variant of the container.
73    Typed(C::Container),
74    /// The binary variant of the container.
75    Bytes(Bytes),
76    /// Relocated, aligned binary data, if `Bytes` doesn't work for some reason.
77    ///
78    /// Reasons could include misalignment, cloning of data, or wanting
79    /// to release the `Bytes` as a scarce resource.
80    ///
81    /// `Vec<u64>` guarantees `u64` alignment for the contained bytes.
82    Align(Vec<u64>),
83}
84
85impl<C: Columnar> Column<C> {
86    /// Empties the column, retaining the `Typed` variant's allocation so the
87    /// caller can refill it.
88    ///
89    /// [`columnar::Clear`] clears the typed container in place without
90    /// releasing its capacity. The serialized variants (`Bytes`/`Align`) own
91    /// no reusable typed buffer, so they are reset to an empty `Typed`.
92    #[inline]
93    pub fn clear(&mut self) {
94        match self {
95            Column::Typed(t) => t.clear(),
96            Column::Bytes(_) | Column::Align(_) => *self = Default::default(),
97        }
98    }
99
100    /// Borrows the container as a reference.
101    #[inline]
102    pub fn borrow(&self) -> <C::Container as Borrow>::Borrowed<'_> {
103        match self {
104            Column::Typed(t) => t.borrow(),
105            Column::Bytes(b) => <<C::Container as Borrow>::Borrowed<'_>>::from_bytes(
106                &mut indexed::decode(bytemuck::cast_slice(b)),
107            ),
108            Column::Align(a) => {
109                <<C::Container as Borrow>::Borrowed<'_>>::from_bytes(&mut indexed::decode(a))
110            }
111        }
112    }
113}
114
115impl<C: Columnar> Default for Column<C> {
116    fn default() -> Self {
117        Self::Typed(Default::default())
118    }
119}
120
121impl<C: Columnar> Clone for Column<C>
122where
123    C::Container: Clone,
124{
125    fn clone(&self) -> Self {
126        match self {
127            // Typed stays typed, although we would have the option to move to aligned data.
128            // If we did it might be confusing why we couldn't push into a cloned column.
129            Column::Typed(t) => Column::Typed(t.clone()),
130            Column::Bytes(b) => {
131                assert_eq!(b.len() % 8, 0);
132                Self::Align(bytemuck::allocation::pod_collect_to_vec(b))
133            }
134            Column::Align(a) => Column::Align(a.clone()),
135        }
136    }
137}
138
139impl<C: Columnar> Accountable for Column<C> {
140    #[inline]
141    fn record_count(&self) -> i64 {
142        self.borrow().len().try_into().expect("Must fit")
143    }
144}
145impl<C: Columnar> DrainContainer for Column<C> {
146    type Item<'a> = Ref<'a, C>;
147    type DrainIter<'a> = IterOwn<<C::Container as Borrow>::Borrowed<'a>>;
148    #[inline]
149    fn drain(&mut self) -> Self::DrainIter<'_> {
150        self.borrow().into_index_iter()
151    }
152}
153
154impl<C: Columnar, T> PushInto<T> for Column<C>
155where
156    C::Container: columnar::Push<T>,
157{
158    #[inline]
159    fn push_into(&mut self, item: T) {
160        use columnar::Push;
161        match self {
162            Column::Typed(t) => t.push(item),
163            Column::Align(_) | Column::Bytes(_) => {
164                // We really oughtn't be calling this in this case.
165                // We could convert to owned, but need more constraints on `C`.
166                unimplemented!("Pushing into Column::Bytes without first clearing");
167            }
168        }
169    }
170}
171
172/// Words per 2 MiB. `length_in_words` returns serialized size in `u64` units,
173/// so this is the page count we round up to. Picked to match
174/// [`builder::ColumnBuilder`]'s output granularity so chunks shipped from the
175/// merger and chunks shipped from the builder are sized comparably.
176const SHIP_WORDS: usize = 1 << 18;
177
178/// Returns true once the serialized size of `borrow` is within 10% of the next
179/// `SHIP_WORDS` boundary.
180///
181/// Same heuristic as `ColumnBuilder::push_into`; lifted out so the merger and
182/// the `SizableContainer` impl agree on the ship signal.
183#[inline]
184pub(crate) fn at_serialized_capacity<'a, A>(borrow: &A) -> bool
185where
186    A: columnar::AsBytes<'a>,
187{
188    let words = indexed::length_in_words(borrow);
189    let round = (words + (SHIP_WORDS - 1)) & !(SHIP_WORDS - 1);
190    round - words < round / 10
191}
192
193impl<C: Columnar> SizableContainer for Column<C> {
194    fn at_capacity(&self) -> bool {
195        // Match `ColumnBuilder`'s ship heuristic: serialized size within 10%
196        // of the next 2 MiB. Aligns chunk-size choices across the two paths
197        // and keeps recipients dealing with a single granularity.
198        //
199        // Serialized chunks (`Bytes` / `Align`) have no typed builder to push
200        // into, so they're trivially "at capacity" — there's no further work
201        // they can absorb.
202        match self {
203            Column::Typed(c) => at_serialized_capacity(&c.borrow()),
204            Column::Bytes(_) | Column::Align(_) => true,
205        }
206    }
207
208    fn ensure_capacity(&mut self, _stash: &mut Option<Self>) {
209        // No pre-reservation: chunks are recycled by the merge framework, so
210        // leaf capacities settle to steady-state after the first round and
211        // there is nothing useful to reserve up front. The `SizableContainer`
212        // impl exists so `at_capacity` is callable on result chunks during
213        // `Merger::merge` orchestration; `ensure_capacity` is a required
214        // method on the trait but has no work to do here.
215    }
216}
217
218impl<C: Columnar> ContainerBytes for Column<C> {
219    #[inline]
220    fn from_bytes(bytes: Bytes) -> Self {
221        // Our expectation / hope is that `bytes` is `u64` aligned and sized.
222        // If the alignment is borked, we can relocate. If the size is borked,
223        // not sure what we do in that case. An incorrect size indicates a problem
224        // of `into_bytes`, or a failure of the communication layer, both of which
225        // are unrecoverable.
226        assert_eq!(bytes.len() % 8, 0);
227        if let Ok(_) = bytemuck::try_cast_slice::<_, u64>(&bytes) {
228            Self::Bytes(bytes)
229        } else {
230            // We failed to cast the slice, so we'll reallocate. `Vec<u64>`
231            // is u64-aligned by construction.
232            Self::Align(bytemuck::allocation::pod_collect_to_vec(&bytes[..]))
233        }
234    }
235
236    #[inline]
237    fn length_in_bytes(&self) -> usize {
238        match self {
239            Column::Typed(t) => indexed::length_in_bytes(&t.borrow()),
240            Column::Bytes(b) => b.len(),
241            Column::Align(a) => 8 * a.len(),
242        }
243    }
244
245    #[inline]
246    fn into_bytes<W: ::std::io::Write>(&self, writer: &mut W) {
247        match self {
248            Column::Typed(t) => indexed::write(writer, &t.borrow()).unwrap(),
249            Column::Bytes(b) => writer.write_all(b).unwrap(),
250            Column::Align(a) => writer.write_all(bytemuck::cast_slice(a)).unwrap(),
251        }
252    }
253}
254
255/// An exchange function for columnar tuples of the form `((K, V), T, D)`. Rust has a hard
256/// time to figure out the lifetimes of the elements when specified as a closure, so we rather
257/// specify it as a function.
258#[inline(always)]
259pub fn columnar_exchange<K, V, T, D>(((k, _), _, _): &Ref<'_, ((K, V), T, D)>) -> u64
260where
261    K: Columnar,
262    for<'a> Ref<'a, K>: Hash,
263    V: Columnar,
264    D: Columnar,
265    T: Columnar,
266{
267    k.hashed()
268}
269
270#[cfg(test)]
271mod tests {
272    use timely::bytes::arc::BytesMut;
273    use timely::container::PushInto;
274    use timely::dataflow::channels::ContainerBytes;
275
276    use super::*;
277
278    /// Produce some bytes that are in columnar format.
279    fn raw_columnar_bytes() -> Vec<u8> {
280        let mut raw = Vec::new();
281        raw.extend(16_u64.to_le_bytes()); // offsets
282        raw.extend(28_u64.to_le_bytes()); // length
283        raw.extend(1_i32.to_le_bytes());
284        raw.extend(2_i32.to_le_bytes());
285        raw.extend(3_i32.to_le_bytes());
286        raw.extend([0, 0, 0, 0]); // padding
287        raw
288    }
289
290    #[mz_ore::test]
291    fn test_column_clone() {
292        let columns = Columnar::as_columns([1, 2, 3].iter());
293        let column_typed: Column<i32> = Column::Typed(columns);
294        let column_typed2 = column_typed.clone();
295
296        assert_eq!(
297            column_typed2.borrow().into_index_iter().collect::<Vec<_>>(),
298            vec![&1, &2, &3]
299        );
300
301        let bytes = BytesMut::from(raw_columnar_bytes()).freeze();
302        let column_bytes: Column<i32> = Column::Bytes(bytes);
303        let column_bytes2 = column_bytes.clone();
304
305        assert_eq!(
306            column_bytes2.borrow().into_index_iter().collect::<Vec<_>>(),
307            vec![&1, &2, &3]
308        );
309
310        let raw = raw_columnar_bytes();
311        let mut region: Vec<u64> = vec![0; raw.len() / 8];
312        let region_bytes = bytemuck::cast_slice_mut(&mut region[..]);
313        region_bytes[..raw.len()].copy_from_slice(&raw);
314        let column_align: Column<i32> = Column::Align(region);
315        let column_align2 = column_align.clone();
316
317        assert_eq!(
318            column_align2.borrow().into_index_iter().collect::<Vec<_>>(),
319            vec![&1, &2, &3]
320        );
321    }
322
323    /// Assert the desired contents of raw_columnar_bytes so that diagnosing test failures is
324    /// easier.
325    #[mz_ore::test]
326    fn test_column_known_bytes() {
327        let mut column: Column<i32> = Default::default();
328        column.push_into(1);
329        column.push_into(2);
330        column.push_into(3);
331        let mut data = Vec::new();
332        column.into_bytes(&mut std::io::Cursor::new(&mut data));
333        assert_eq!(data, raw_columnar_bytes());
334    }
335
336    #[mz_ore::test]
337    fn test_column_from_bytes() {
338        let raw = raw_columnar_bytes();
339
340        let buf = vec![0; raw.len() + 8];
341        let align = buf.as_ptr().align_offset(std::mem::size_of::<u64>());
342        let mut bytes_mut = BytesMut::from(buf);
343        let _ = bytes_mut.extract_to(align);
344        bytes_mut[..raw.len()].copy_from_slice(&raw);
345        let aligned_bytes = bytes_mut.extract_to(raw.len());
346
347        let column: Column<i32> = Column::from_bytes(aligned_bytes);
348        assert!(matches!(column, Column::Bytes(_)));
349        assert_eq!(
350            column.borrow().into_index_iter().collect::<Vec<_>>(),
351            vec![&1, &2, &3]
352        );
353
354        let buf = vec![0; raw.len() + 8];
355        let align = buf.as_ptr().align_offset(std::mem::size_of::<u64>());
356        let mut bytes_mut = BytesMut::from(buf);
357        let _ = bytes_mut.extract_to(align + 1);
358        bytes_mut[..raw.len()].copy_from_slice(&raw);
359        let unaligned_bytes = bytes_mut.extract_to(raw.len());
360
361        let column: Column<i32> = Column::from_bytes(unaligned_bytes);
362        assert!(matches!(column, Column::Align(_)));
363        assert_eq!(
364            column.borrow().into_index_iter().collect::<Vec<_>>(),
365            vec![&1, &2, &3]
366        );
367    }
368}
mz_timely_util/columnar.rs

mz_timely_util/
columnar.rs
