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

mz_timely_util/
columnar.rs
