mz_timely_util/

columnar.rs

// Copyright Materialize, Inc. and contributors. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License in the LICENSE file at the
// root of this repository, or online at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Container for columnar data.

#![deny(missing_docs)]

pub mod batcher;
pub mod body;
pub mod builder;
pub mod builder_input;
pub mod chunk;
pub mod consolidate;
pub mod merge_batcher;
pub mod unload;

use std::hash::{BuildHasher, Hash, Hasher};
use std::sync::LazyLock;

use columnar::Borrow;
use columnar::bytes::indexed;
use columnar::common::IterOwn;
use columnar::{Clear, FromBytes, Index, Len};
use columnar::{Columnar, Ref};
use differential_dataflow::Hashable;
use differential_dataflow::collection::containers::{Enter, ResultsIn};
use differential_dataflow::trace::implementations::merge_batcher::MergeBatcher;
use itertools::Itertools;
use timely::Accountable;
use timely::bytes::arc::Bytes;
use timely::container::{DrainContainer, PushInto, SizableContainer};
use timely::dataflow::channels::ContainerBytes;
use timely::progress::Timestamp;
use timely::progress::timestamp::Refines;

use crate::columnation::ColInternalMerger;

/// A batcher for columnar storage.
///
/// The chunker is supplied to the arrange operator separately. Callers pass
/// it explicitly: [`ColumnationChunker`](crate::columnation::ColumnationChunker)
/// for `Vec<_>` input, or [`batcher::Chunker`] (over a `ColumnationStack<_>`) for
/// [`Column`] input.
pub type Col2ValBatcher<K, V, T, R> = MergeBatcher<ColInternalMerger<(K, V), T, R>>;
/// A batcher for columnar storage with unit values.
pub type Col2KeyBatcher<K, T, R> = Col2ValBatcher<K, (), T, R>;

/// Pageable counterpart to [`Col2ValBatcher`]. Routes every chunk produced
/// by chunking, merging, or extract through a [`crate::column_pager::ColumnPager`],
/// so memory pressure can spill chains to a backing store without touching
/// the merge / extract bodies.
///
/// Drop-in shape at the type level: both aliases take `(K, V, T, R)` and
/// produce a `Batcher<Input = Column<((K, V), T, R)>, Output = Column<((K,
/// V), T, R)>>`. Call sites can swap with `cargo fix`–style renaming once
/// downstream `Trace`/`Builder` impls have been wired up. The pager itself
/// defaults to [`crate::column_pager::ColumnPager::disabled`]; inject a
/// real one via [`merge_batcher::ColumnMergeBatcher::set_pager`].
pub type Col2ValPagedBatcher<K, V, T, R> = merge_batcher::ColumnMergeBatcher<(K, V), T, R>;

/// Columnar-native counterpart to [`Col2ValBatcher`], holding
/// [`body::ColumnBody`] chunks rather than columnation stacks and merging
/// them through [`batcher::ColumnMerger`].
///
/// Pairs with [`batcher::ColumnChunker`] and any builder whose `Input` is
/// `ColumnBody<((K, V), T, R)>`. Unlike [`Col2ValPagedBatcher`] the chains
/// stay resident, so this arm carries no pager and no spill budget.
pub type Col2ValColBatcher<K, V, T, R> = MergeBatcher<batcher::ColumnMerger<(K, V), T, R>>;

/// A container based on a columnar store, encoded in aligned bytes: the container on dataflow
/// edges.
///
/// The type can represent typed data, bytes from Timely, or an aligned allocation. The name
/// is singular to express that the preferred format is [`Column::Align`]. The [`Column::Typed`]
/// variant is used to construct the container, and it owns potentially multiple columns of data.
///
/// Data at rest behind an edge, in a merge batcher's chains, a spill-backed chunk, or a batch
/// builder, is a [`body::ColumnBody`] instead, which has no channel-bytes form.
pub enum Column<C: Columnar> {
    /// The typed variant of the container.
    Typed(C::Container),
    /// The binary variant of the container.
    Bytes(Bytes),
    /// Relocated, aligned binary data, if `Bytes` doesn't work for some reason.
    ///
    /// Reasons could include misalignment, cloning of data, or wanting
    /// to release the `Bytes` as a scarce resource.
    ///
    /// `Vec<u64>` guarantees `u64` alignment for the contained bytes.
    Align(Vec<u64>),
}

impl<C: Columnar> Column<C> {
    /// Empties the column, retaining the `Typed` variant's allocation so the
    /// caller can refill it.
    ///
    /// [`columnar::Clear`] clears the typed container in place without
    /// releasing its capacity. The serialized variants (`Bytes`/`Align`) own
    /// no reusable typed buffer, so they are reset to an empty `Typed`.
    #[inline]
    pub fn clear(&mut self) {
        match self {
            Column::Typed(t) => t.clear(),
            Column::Bytes(_) | Column::Align(_) => *self = Default::default(),
        }
    }

    /// True when the column holds no records.
    ///
    /// The `Typed` variant answers from the container itself. The serialized
    /// variants have to reconstruct their borrowed view to reach a length, so
    /// there this costs as much as [`Column::borrow`].
    #[inline]
    pub fn is_empty(&self) -> bool {
        match self {
            Column::Typed(t) => t.is_empty(),
            Column::Bytes(_) | Column::Align(_) => self.borrow().is_empty(),
        }
    }

    /// Borrows the container as a reference.
    // NOTE: `inline(always)` measured 7.1 ns to 5.8 ns per call for a per-element borrow of an
    // `Align` column. The header decode stays per call, so loops still hoist the borrow.
    #[inline(always)]
    pub fn borrow(&self) -> <C::Container as Borrow>::Borrowed<'_> {
        match self {
            Column::Typed(t) => t.borrow(),
            Column::Bytes(b) => <<C::Container as Borrow>::Borrowed<'_>>::from_bytes(
                &mut indexed::decode(bytemuck::cast_slice(b)),
            ),
            Column::Align(a) => {
                <<C::Container as Borrow>::Borrowed<'_>>::from_bytes(&mut indexed::decode(a))
            }
        }
    }
}

impl<C: Columnar> Default for Column<C> {
    fn default() -> Self {
        Self::Typed(Default::default())
    }
}

impl<C: Columnar> Clone for Column<C>
where
    C::Container: Clone,
{
    fn clone(&self) -> Self {
        match self {
            // Typed stays typed, although we would have the option to move to aligned data.
            // If we did it might be confusing why we couldn't push into a cloned column.
            Column::Typed(t) => Column::Typed(t.clone()),
            Column::Bytes(b) => {
                assert_eq!(b.len() % 8, 0);
                Self::Align(bytemuck::allocation::pod_collect_to_vec(b))
            }
            Column::Align(a) => Column::Align(a.clone()),
        }
    }
}

impl<C: Columnar> Accountable for Column<C> {
    #[inline]
    fn record_count(&self) -> i64 {
        self.borrow().len().try_into().expect("Must fit")
    }
}
impl<C: Columnar> DrainContainer for Column<C> {
    type Item<'a> = Ref<'a, C>;
    type DrainIter<'a> = IterOwn<<C::Container as Borrow>::Borrowed<'a>>;
    #[inline]
    fn drain(&mut self) -> Self::DrainIter<'_> {
        self.borrow().into_index_iter()
    }
}

impl<C: Columnar, T> PushInto<T> for Column<C>
where
    C::Container: columnar::Push<T>,
{
    #[inline]
    fn push_into(&mut self, item: T) {
        use columnar::Push;
        match self {
            Column::Typed(t) => t.push(item),
            Column::Align(_) | Column::Bytes(_) => {
                // We really oughtn't be calling this in this case.
                // We could convert to owned, but need more constraints on `C`.
                unimplemented!("Pushing into Column::Bytes without first clearing");
            }
        }
    }
}

/// Re-encodes a column's times into a refining timestamp, so a columnar collection can enter
/// an iterative scope.
impl<D, T1, T2, R> Enter<T1, T2> for Column<(D, T1, R)>
where
    D: Columnar,
    T1: Columnar + Timestamp,
    T2: Columnar + Refines<T1>,
    R: Columnar,
    (D, T1, R): Columnar<Container = (D::Container, T1::Container, R::Container)>,
    (D, T2, R): Columnar<Container = (D::Container, T2::Container, R::Container)>,
    for<'a> D::Container: columnar::Push<Ref<'a, D>>,
    for<'a> T2::Container: columnar::Push<&'a T2>,
    for<'a> R::Container: columnar::Push<Ref<'a, R>>,
{
    type InnerContainer = Column<(D, T2, R)>;

    fn enter(self) -> Self::InnerContainer {
        use columnar::Push;
        match self {
            // Only the times change, so the data and diff columns move across whole.
            Column::Typed((data, times, diffs)) => {
                let mut inner = T2::Container::default();
                for time in times.borrow().into_index_iter() {
                    inner.push(&T2::to_inner(T1::into_owned(time)));
                }
                Column::Typed((data, inner, diffs))
            }
            // A serialized column owns no typed sub-containers, so only the times are
            // materialized. The data and diff columns go from their borrowed views
            // straight into the output allocation, a copy per column rather than a
            // decode and re-encode per record.
            serialized => {
                let (borrowed_data, borrowed_times, borrowed_diffs) = serialized.borrow();
                let mut times = T2::Container::default();
                for time in borrowed_times.into_index_iter() {
                    times.push(&T2::to_inner(T1::into_owned(time)));
                }
                let view = (borrowed_data, times.borrow(), borrowed_diffs);
                let words = indexed::length_in_words(&view);
                let mut alloc: Vec<u64> = Vec::with_capacity(words);
                indexed::encode(&mut alloc, &view);
                Column::Align(alloc)
            }
        }
    }
}

/// Advances a column's times by `step`, so a columnar collection can be a loop variable.
///
/// A time that the summary does not advance drops its record, so the data and diff columns
/// cannot move across whole the way [`Enter`] moves them: a record's position in one column
/// has to keep matching its position in the others. The times are read and advanced first,
/// and the rest is rebuilt only when the step actually dropped something.
impl<D, T, R> ResultsIn<T::Summary> for Column<(D, T, R)>
where
    D: Columnar,
    T: Columnar + Timestamp,
    R: Columnar,
    (D, T, R): Columnar<Container = (D::Container, T::Container, R::Container)>,
    for<'a> D::Container: columnar::Push<Ref<'a, D>>,
    for<'a> T::Container: columnar::Push<&'a T>,
    for<'a> R::Container: columnar::Push<Ref<'a, R>>,
{
    fn results_in(self, step: &T::Summary) -> Self {
        use columnar::Push;
        use timely::progress::PathSummary;

        // Advance the times first. Whether the step dropped any record decides whether the
        // data and diff columns still line up with the new times, which is what lets them
        // be reused rather than rebuilt.
        let (times, kept) = {
            let (_, borrowed_times, _) = self.borrow();
            let mut times = T::Container::default();
            let mut time = T::minimum();
            let kept: Vec<bool> = borrowed_times
                .into_index_iter()
                .map(|reference| {
                    time.copy_from(reference);
                    match step.results_in(&time) {
                        Some(time) => {
                            times.push(&time);
                            true
                        }
                        None => false,
                    }
                })
                .collect();
            (times, kept)
        };

        if kept.iter().all(|kept| *kept) {
            match self {
                // Only the times changed, so the data and diff columns move across whole.
                Column::Typed((data, _, diffs)) => Column::Typed((data, times, diffs)),
                // A serialized column owns no typed sub-containers, so its data and diff
                // columns go from their borrowed views straight into the output allocation.
                serialized => {
                    let (borrowed_data, _, borrowed_diffs) = serialized.borrow();
                    let view = (borrowed_data, times.borrow(), borrowed_diffs);
                    let words = indexed::length_in_words(&view);
                    let mut alloc: Vec<u64> = Vec::with_capacity(words);
                    indexed::encode(&mut alloc, &view);
                    Column::Align(alloc)
                }
            }
        } else {
            let (borrowed_data, _, borrowed_diffs) = self.borrow();
            let mut data = D::Container::default();
            let mut diffs = R::Container::default();
            let records = borrowed_data
                .into_index_iter()
                .zip_eq(borrowed_diffs.into_index_iter())
                .zip_eq(kept.iter());
            for ((datum, diff), _) in records.filter(|(_, kept)| **kept) {
                data.push(datum);
                diffs.push(diff);
            }
            Column::Typed((data, times, diffs))
        }
    }
}

/// Words per 2 MiB. `length_in_words` returns serialized size in `u64` units,
/// so this is the page count we round up to. Picked to match
/// [`builder::ColumnBuilder`]'s output granularity so chunks shipped from the
/// merger and chunks shipped from the builder are sized comparably.
const SHIP_WORDS: usize = 1 << 18;

/// Returns true once the serialized size of `borrow` reaches 10% under
/// `SHIP_WORDS`.
///
/// Monotone in size, deliberately not a window below the boundary. A single
/// record wider than a window steps clear over it, and a ship signal that
/// un-fires past the boundary lets a chunk grow until it exceeds the buffer
/// pool's largest size class, past which a spilled body degrades to
/// permanently resident. The same heuristic as [`builder::ColumnBuilder`]'s
/// ship point, lifted out so the builder, the merger, and the
/// `SizableContainer` impl agree on the signal.
#[inline]
pub(crate) fn at_serialized_capacity<'a, A>(borrow: &A) -> bool
where
    A: columnar::AsBytes<'a>,
{
    indexed::length_in_words(borrow) >= SHIP_WORDS - SHIP_WORDS / 10
}

impl<C: Columnar> SizableContainer for Column<C> {
    fn at_capacity(&self) -> bool {
        // Match `ColumnBuilder`'s ship heuristic: serialized size at the
        // 2 MiB ship threshold. Aligns chunk-size choices across the two
        // paths and keeps recipients dealing with a single granularity.
        //
        // Serialized chunks (`Bytes` / `Align`) have no typed builder to push
        // into, so they're trivially "at capacity" — there's no further work
        // they can absorb.
        match self {
            Column::Typed(c) => at_serialized_capacity(&c.borrow()),
            Column::Bytes(_) | Column::Align(_) => true,
        }
    }

    fn ensure_capacity(&mut self, _stash: &mut Option<Self>) {
        // No pre-reservation: chunks are recycled by the merge framework, so
        // leaf capacities settle to steady-state after the first round and
        // there is nothing useful to reserve up front. The `SizableContainer`
        // impl exists so `at_capacity` is callable on result chunks during
        // `Merger::merge` orchestration; `ensure_capacity` is a required
        // method on the trait but has no work to do here.
    }
}

impl<C: Columnar> ContainerBytes for Column<C> {
    #[inline]
    fn from_bytes(bytes: Bytes) -> Self {
        // Our expectation / hope is that `bytes` is `u64` aligned and sized.
        // If the alignment is borked, we can relocate. If the size is borked,
        // not sure what we do in that case. An incorrect size indicates a problem
        // of `into_bytes`, or a failure of the communication layer, both of which
        // are unrecoverable.
        assert_eq!(bytes.len() % 8, 0);
        if let Ok(_) = bytemuck::try_cast_slice::<_, u64>(&bytes) {
            Self::Bytes(bytes)
        } else {
            // We failed to cast the slice, so we'll reallocate. `Vec<u64>`
            // is u64-aligned by construction.
            Self::Align(bytemuck::allocation::pod_collect_to_vec(&bytes[..]))
        }
    }

    #[inline]
    fn length_in_bytes(&self) -> usize {
        match self {
            Column::Typed(t) => indexed::length_in_bytes(&t.borrow()),
            Column::Bytes(b) => b.len(),
            Column::Align(a) => 8 * a.len(),
        }
    }

    #[inline]
    fn into_bytes<W: ::std::io::Write>(&self, writer: &mut W) {
        match self {
            Column::Typed(t) => indexed::write(writer, &t.borrow()).unwrap(),
            Column::Bytes(b) => writer.write_all(b).unwrap(),
            Column::Align(a) => writer.write_all(bytemuck::cast_slice(a)).unwrap(),
        }
    }
}

/// An exchange function for columnar tuples of the form `((K, V), T, D)`. Rust has a hard
/// time to figure out the lifetimes of the elements when specified as a closure, so we rather
/// specify it as a function.
#[inline(always)]
pub fn columnar_exchange<K, V, T, D>(((k, _), _, _): &Ref<'_, ((K, V), T, D)>) -> u64
where
    K: Columnar,
    for<'a> Ref<'a, K>: Hash,
    V: Columnar,
    D: Columnar,
    T: Columnar,
{
    k.hashed()
}

/// Routes a `(D, T, R)` column by the hash of its data column.
///
/// Counterpart to [`columnar_exchange`] for collections whose data is not a
/// key/value pair. Consolidation sites want [`columnar_consolidate_exchange`]
/// instead.
pub fn columnar_exchange_data<D, T, R>((d, _, _): &Ref<'_, (D, T, R)>) -> u64
where
    D: Columnar,
    for<'a> Ref<'a, D>: Hash,
    T: Columnar,
    R: Columnar,
{
    d.hashed()
}

/// Routes a `(D, T, R)` column for consolidation, by a fixed-seed AHash of its data column.
///
/// Worker assignment is `hash % workers`, so the low bits alone decide the split and the
/// [`Hashable`] default (FNV) the other exchange functions use diffuses them poorly. The
/// seed is fixed, so routing is identical across builds and replicas.
///
/// Spelled as a function rather than a closure over the hasher state, because the argument
/// is higher-ranked in its lifetime and closure inference cannot express that.
pub fn columnar_consolidate_exchange<D, T, R>((d, _, _): &Ref<'_, (D, T, R)>) -> u64
where
    D: Columnar,
    for<'a> Ref<'a, D>: Hash,
    T: Columnar,
    R: Columnar,
{
    static STATE: LazyLock<ahash::RandomState> = LazyLock::new(crate::hash::fixed_state);
    let mut hasher = STATE.build_hasher();
    d.hash(&mut hasher);
    hasher.finish()
}

#[cfg(test)]
mod tests {
    use timely::bytes::arc::BytesMut;
    use timely::container::PushInto;
    use timely::dataflow::channels::ContainerBytes;

    use super::*;

    /// Produce some bytes that are in columnar format.
    fn raw_columnar_bytes() -> Vec<u8> {
        let mut raw = Vec::new();
        raw.extend(16_u64.to_le_bytes()); // offsets
        raw.extend(28_u64.to_le_bytes()); // length
        raw.extend(1_i32.to_le_bytes());
        raw.extend(2_i32.to_le_bytes());
        raw.extend(3_i32.to_le_bytes());
        raw.extend([0, 0, 0, 0]); // padding
        raw
    }

    #[mz_ore::test]
    fn test_column_clone() {
        let columns = Columnar::as_columns([1, 2, 3].iter());
        let column_typed: Column<i32> = Column::Typed(columns);
        let column_typed2 = column_typed.clone();

        assert_eq!(
            column_typed2.borrow().into_index_iter().collect::<Vec<_>>(),
            vec![&1, &2, &3]
        );

        let bytes = BytesMut::from(raw_columnar_bytes()).freeze();
        let column_bytes: Column<i32> = Column::Bytes(bytes);
        let column_bytes2 = column_bytes.clone();

        assert_eq!(
            column_bytes2.borrow().into_index_iter().collect::<Vec<_>>(),
            vec![&1, &2, &3]
        );

        let raw = raw_columnar_bytes();
        let mut region: Vec<u64> = vec![0; raw.len() / 8];
        let region_bytes = bytemuck::cast_slice_mut(&mut region[..]);
        region_bytes[..raw.len()].copy_from_slice(&raw);
        let column_align: Column<i32> = Column::Align(region);
        let column_align2 = column_align.clone();

        assert_eq!(
            column_align2.borrow().into_index_iter().collect::<Vec<_>>(),
            vec![&1, &2, &3]
        );
    }

    /// Assert the desired contents of raw_columnar_bytes so that diagnosing test failures is
    /// easier.
    #[mz_ore::test]
    fn test_column_known_bytes() {
        let mut column: Column<i32> = Default::default();
        column.push_into(1);
        column.push_into(2);
        column.push_into(3);
        let mut data = Vec::new();
        column.into_bytes(&mut std::io::Cursor::new(&mut data));
        assert_eq!(data, raw_columnar_bytes());
    }

    #[mz_ore::test]
    fn test_column_from_bytes() {
        let raw = raw_columnar_bytes();

        let buf = vec![0; raw.len() + 8];
        let align = buf.as_ptr().align_offset(std::mem::size_of::<u64>());
        let mut bytes_mut = BytesMut::from(buf);
        let _ = bytes_mut.extract_to(align);
        bytes_mut[..raw.len()].copy_from_slice(&raw);
        let aligned_bytes = bytes_mut.extract_to(raw.len());

        let column: Column<i32> = Column::from_bytes(aligned_bytes);
        assert!(matches!(column, Column::Bytes(_)));
        assert_eq!(
            column.borrow().into_index_iter().collect::<Vec<_>>(),
            vec![&1, &2, &3]
        );

        let buf = vec![0; raw.len() + 8];
        let align = buf.as_ptr().align_offset(std::mem::size_of::<u64>());
        let mut bytes_mut = BytesMut::from(buf);
        let _ = bytes_mut.extract_to(align + 1);
        bytes_mut[..raw.len()].copy_from_slice(&raw);
        let unaligned_bytes = bytes_mut.extract_to(raw.len());

        let column: Column<i32> = Column::from_bytes(unaligned_bytes);
        assert!(matches!(column, Column::Align(_)));
        assert_eq!(
            column.borrow().into_index_iter().collect::<Vec<_>>(),
            vec![&1, &2, &3]
        );
    }

    /// The ship signal is monotone: once it fires it stays fired, even when
    /// a single wide record steps far past the 2 MiB boundary in one push.
    #[mz_ore::test]
    #[cfg_attr(miri, ignore)] // too slow
    fn ship_threshold_monotone() {
        use columnar::Push;
        let mut container = <Vec<u64> as Columnar>::Container::default();
        // Wider than 10% of any boundary a 25 MiB run can reach.
        let wide: Vec<u64> = vec![0u64; 50_000];
        let mut fired = false;
        for pushes in 1..=64 {
            container.push(&wide);
            let now = at_serialized_capacity(&container.borrow());
            if fired {
                assert!(now, "ship signal un-fired at {pushes} records");
            }
            fired = fired || now;
        }
        assert!(fired, "ship signal never fired");
    }
}