mz_timely_util/columnar/

consolidate.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.

//! A `ContainerBuilder` that consolidates `(D, T, R)` updates and emits columnar containers.
//!
//! Two-level buffering:
//!
//! 1. AoS staging `Vec<(D, T, R)>` with a small cap so `consolidate_updates`' `n log n` cost
//!    stays bounded. Cancellations and same-key updates collapse here before reaching the
//!    column-shaped storage. A drain-multiple-of-half-cap trick keeps the leftover in staging,
//!    so cross-batch keys with the same `(D, T)` continue consolidating on the next sort.
//! 2. SoA accumulator: one sub-container per column (`D::Container`, `T::Container`,
//!    `R::Container`). Drains push in fixed-size chunks (`DRAIN_CHUNK_ROWS`) with three
//!    sequential per-column passes per chunk, so the inner pushes autovectorize. After each
//!    chunk, check the serialized size; once the accumulator reaches the flush threshold
//!    (90% of `OUTPUT_TARGET_WORDS`), serialize into an aligned `Vec<u64>` (no zero-fill —
//!    written by `indexed::encode`) and ship as `Column::Align`. Per-chunk granularity bounds
//!    overshoot to `K * row_words`. The trailing partial on `finish` ships as `Column::Typed`.
//!
//! Generic over `(D, T, R): Columnar` via the columnar tuple decomposition
//! `<(D, T, R) as Columnar>::Container = (D::Container, T::Container, R::Container)`.

use std::collections::VecDeque;

use columnar::bytes::indexed;
use columnar::{Borrow, Columnar, Push};
use differential_dataflow::Data;
use differential_dataflow::consolidation::consolidate_updates;
use differential_dataflow::difference::Semigroup;
use timely::container::{ContainerBuilder, PushInto};

use crate::columnar::Column;

/// Per-buffer byte budget for the staging cap. Matches the 8 KiB basis DD's
/// `ConsolidatingContainerBuilder` uses via `timely::container::buffer::default_capacity`.
const STAGING_BUFFER_BYTES: usize = 8 * 1024;

/// Default items per staging buffer: `2 * STAGING_BUFFER_BYTES / size_of::<(D, T, R)>()`,
/// matching DD's `ConsolidatingContainerBuilder`. Small enough that O(n log n) sort stays
/// cheap, large enough to amortize the sort + drain overhead across many pushes.
fn default_staging_cap<D, T, R>() -> usize {
    let elem = std::mem::size_of::<(D, T, R)>().max(1);
    // Floor at 2 so the half-cap drain grain is at least 1.
    (2 * STAGING_BUFFER_BYTES / elem).max(2)
}
/// Target serialized chunk size in u64 words (2 MiB). Bounded slop matters once we put these
/// behind huge pages.
const OUTPUT_TARGET_WORDS: usize = 1 << 18;
/// Flush when within 10% of `OUTPUT_TARGET_WORDS` — matches `ColumnBuilder`'s slop heuristic.
/// Computed at compile time so the hot loop is a single `cmp`/`jae` instead of a per-row
/// round-up + divide-by-10.
const FLUSH_THRESHOLD_WORDS: usize = OUTPUT_TARGET_WORDS - OUTPUT_TARGET_WORDS / 10;
/// Drain rows from staging in chunks of this size. Inside each chunk we do three sequential
/// per-column passes — long enough for autovectorization (4–8 vector iterations on NEON / SVE2
/// 128-bit / AVX2 / SVE 256-bit) — while the outer per-chunk size check bounds overshoot to
/// `K * row_words`. With a 1 KiB row (128 words) and K=16, worst-case overshoot is 2 KiB,
/// well under the 10% slop budget on a 2 MiB target.
const DRAIN_CHUNK_ROWS: usize = 16;

/// A container builder that consolidates `(D, T, R)` updates and emits `Column<(D, T, R)>`.
///
/// Stages updates in an AoS `Vec` for in-place consolidation, then drains consolidated rows
/// in `DRAIN_CHUNK_ROWS`-sized chunks (three sequential per-column passes per chunk) into SoA
/// sub-containers, flushing whenever the accumulator hits the flush threshold (90% of
/// `OUTPUT_TARGET_WORDS`). Flushed accumulators are written into aligned `Vec<u64>` (via
/// `indexed::encode`, no zero-fill) and queued as `Column::Align`. The trailing partial on
/// `finish` ships as `Column::Typed`.
///
/// Does **not** maintain FIFO ordering (consolidation reorders updates).
pub struct ConsolidatingColumnBuilder<D, T, R>
where
    D: Columnar,
    T: Columnar,
    R: Columnar,
{
    /// AoS staging buffer for in-place consolidation. Cap = [`Self::staging_cap`].
    staging: Vec<(D, T, R)>,
    /// Capacity of `staging`. Drain triggers when `staging.len()` hits this.
    staging_cap: usize,
    /// SoA accumulator, one sub-container per column.
    cur_d: D::Container,
    cur_t: T::Container,
    cur_r: R::Container,
    /// Number of `(D, T, R)` tuples currently in `cur_*`.
    cur_len: usize,
    /// Finished columns ready to ship.
    pending: VecDeque<Column<(D, T, R)>>,
    /// The currently extracted/finished column.
    finished: Option<Column<(D, T, R)>>,
}

impl<D, T, R> Default for ConsolidatingColumnBuilder<D, T, R>
where
    D: Columnar,
    T: Columnar,
    R: Columnar,
{
    fn default() -> Self {
        let cap = default_staging_cap::<D, T, R>();
        Self {
            // Pre-allocate so `push` is unconditional (no per-push capacity check or lazy
            // reserve branch).
            staging: Vec::with_capacity(cap),
            staging_cap: cap,
            cur_d: D::Container::default(),
            cur_t: T::Container::default(),
            cur_r: R::Container::default(),
            cur_len: 0,
            pending: VecDeque::new(),
            finished: None,
        }
    }
}

impl<D, T, R> ConsolidatingColumnBuilder<D, T, R>
where
    D: Data + Columnar,
    T: Data + Columnar,
    R: Semigroup + Columnar + 'static,
    (D, T, R): Columnar<Container = (D::Container, T::Container, R::Container)>,
{
    /// Sort and consolidate `staging`, then drain a multiple-of-`grain` prefix into the SoA
    /// accumulator. Pass `1` to drain everything (used by `finish`). Pushes in chunks of
    /// `DRAIN_CHUNK_ROWS` and flushes mid-drain whenever the accumulator hits
    /// `FLUSH_THRESHOLD_WORDS`, so a single drain can mint several aligned containers when the
    /// prefix is large.
    #[cold]
    fn consolidate_and_drain(&mut self, grain: usize) {
        consolidate_updates(&mut self.staging);
        let drain_n = (self.staging.len() / grain) * grain;
        if drain_n == 0 {
            return;
        }

        // Drain in chunks of `DRAIN_CHUNK_ROWS` rows. Three sequential per-column passes inside
        // the chunk give the compiler streamable loops it can autovectorize for primitive
        // containers; the per-chunk size check bounds overshoot of the 90%-of-2-MiB threshold to
        // ~`K * row_words`. After each chunk, check the serialized size (via `length_in_words`,
        // not item count, so output Columns stay bounded for variable-width `D` like `Row`).
        let mut consumed = 0;
        while consumed < drain_n {
            let take = (drain_n - consumed).min(DRAIN_CHUNK_ROWS);
            let head = &self.staging[consumed..consumed + take];
            for (d, _, _) in head {
                self.cur_d.push(d);
            }
            for (_, t, _) in head {
                self.cur_t.push(t);
            }
            for (_, _, r) in head {
                self.cur_r.push(r);
            }
            self.cur_len += take;
            consumed += take;

            let words = {
                let view = (
                    self.cur_d.borrow(),
                    self.cur_t.borrow(),
                    self.cur_r.borrow(),
                );
                indexed::length_in_words(&view)
            };
            if words >= FLUSH_THRESHOLD_WORDS {
                self.flush_aligned();
            }
        }
        self.staging.drain(..consumed);
    }

    /// Serialize the SoA accumulator into a `Column::Align` via `indexed::encode`, which
    /// builds the buffer with `Vec::push`/`extend_from_slice` so no memory is initialized
    /// twice.
    #[cold]
    fn flush_aligned(&mut self) {
        if self.cur_len == 0 {
            return;
        }
        let cur: <(D, T, R) as Columnar>::Container = (
            std::mem::take(&mut self.cur_d),
            std::mem::take(&mut self.cur_t),
            std::mem::take(&mut self.cur_r),
        );
        self.cur_len = 0;

        let mut buf: Vec<u64> = Vec::with_capacity(indexed::length_in_words(&cur.borrow()));
        indexed::encode(&mut buf, &cur.borrow());
        self.pending.push_back(Column::Align(buf));
    }
}

impl<D, T, R> PushInto<(D, T, R)> for ConsolidatingColumnBuilder<D, T, R>
where
    D: Data + Columnar,
    T: Data + Columnar,
    R: Semigroup + Columnar + 'static,
    (D, T, R): Columnar<Container = (D::Container, T::Container, R::Container)>,
{
    /// Push an element into the staging buffer; consolidate + drain when full.
    #[inline]
    fn push_into(&mut self, item: (D, T, R)) {
        self.staging.push(item);
        if self.staging.len() == self.staging_cap {
            self.consolidate_and_drain(self.staging_cap / 2);
        }
    }
}

impl<D, T, R> ContainerBuilder for ConsolidatingColumnBuilder<D, T, R>
where
    D: Data + Columnar,
    T: Data + Columnar,
    R: Semigroup + Columnar + 'static,
    (D, T, R): Columnar<Container = (D::Container, T::Container, R::Container)>,
    <(D, T, R) as Columnar>::Container: Clone,
{
    type Container = Column<(D, T, R)>;

    #[inline]
    fn extract(&mut self) -> Option<&mut Self::Container> {
        if let Some(c) = self.pending.pop_front() {
            self.finished = Some(c);
            self.finished.as_mut()
        } else {
            None
        }
    }

    #[inline]
    fn finish(&mut self) -> Option<&mut Self::Container> {
        if !self.staging.is_empty() {
            // `multiple = 1` so any remainder also leaves staging.
            self.consolidate_and_drain(1);
        }
        // Trailing partial: ship as `Column::Typed` (no extra serialize copy).
        if self.cur_len > 0 {
            let cur: <(D, T, R) as Columnar>::Container = (
                std::mem::take(&mut self.cur_d),
                std::mem::take(&mut self.cur_t),
                std::mem::take(&mut self.cur_r),
            );
            self.cur_len = 0;
            self.pending.push_back(Column::Typed(cur));
        }
        self.extract()
    }
}

#[cfg(test)]
mod tests {
    use columnar::Index;
    use columnar::Len;
    use timely::container::{ContainerBuilder, PushInto};

    use super::*;

    /// Collect every `(D, T, R)` row from a `Column<(u64, u64, i64)>`.
    fn rows(column: &Column<(u64, u64, i64)>) -> Vec<(u64, u64, i64)> {
        let borrow = column.borrow();
        (0..borrow.len())
            .map(|i| {
                let r = borrow.get(i);
                (*r.0, *r.1, *r.2)
            })
            .collect()
    }

    /// Drain a builder by repeatedly calling `extract` then `finish`.
    fn drain(mut builder: ConsolidatingColumnBuilder<u64, u64, i64>) -> Vec<(u64, u64, i64)> {
        let mut out: Vec<(u64, u64, i64)> = Vec::new();
        while let Some(c) = builder.extract() {
            for r in rows(c) {
                out.push(r);
            }
        }
        if let Some(c) = builder.finish() {
            for r in rows(c) {
                out.push(r);
            }
        }
        out
    }

    #[mz_ore::test]
    fn empty_finish_yields_none() {
        let mut builder: ConsolidatingColumnBuilder<u64, u64, i64> = Default::default();
        assert!(builder.extract().is_none());
        assert!(builder.finish().is_none());
    }

    #[mz_ore::test]
    fn single_push_finish_yields_one() {
        let mut builder: ConsolidatingColumnBuilder<u64, u64, i64> = Default::default();
        builder.push_into((1u64, 0u64, 1i64));
        let column = builder.finish().expect("one container");
        assert_eq!(rows(column), vec![(1, 0, 1)]);
        assert!(builder.finish().is_none());
    }

    #[mz_ore::test]
    fn consolidates_on_threshold() {
        let mut builder: ConsolidatingColumnBuilder<u64, u64, i64> = Default::default();
        // Push enough +1/-1 pairs to exceed the staging cap several times over and trigger
        // multiple consolidation cycles. Everything cancels in the staging buffer before
        // ever reaching the SoA accumulator.
        let cap = default_staging_cap::<u64, u64, i64>();
        for _ in 0..(cap * 4) {
            builder.push_into((7u64, 0u64, 1i64));
            builder.push_into((7u64, 0u64, -1i64));
        }
        assert!(drain(builder).is_empty());
    }

    #[mz_ore::test]
    fn cross_batch_consolidation() {
        // A single key pushed many times must collapse to one row even across many staging
        // refills. The drain-multiple-of-grain trick keeps the in-progress consolidated row
        // in staging so subsequent pushes merge into it.
        let mut builder: ConsolidatingColumnBuilder<u64, u64, i64> = Default::default();
        let n: i64 = 100_000;
        for _ in 0..n {
            builder.push_into((42u64, 0u64, 1i64));
        }
        let out = drain(builder);
        assert_eq!(out, vec![(42, 0, n)]);
    }

    #[mz_ore::test]
    fn multiple_distinct_keys() {
        let mut builder: ConsolidatingColumnBuilder<u64, u64, i64> = Default::default();
        builder.push_into((1u64, 0u64, 1i64));
        builder.push_into((2u64, 0u64, 1i64));
        builder.push_into((1u64, 0u64, 1i64));
        let mut out = drain(builder);
        out.sort();
        assert_eq!(out, vec![(1, 0, 2), (2, 0, 1)]);
    }

    #[mz_ore::test]
    #[cfg_attr(miri, ignore)] // too slow
    fn emits_multiple_containers() {
        let mut builder: ConsolidatingColumnBuilder<u64, u64, i64> = Default::default();
        // Enough distinct rows to fill the SoA accumulator past the output target multiple
        // times. Each row is 24 bytes; ~87k rows ≈ 2 MiB, so 300k pushes should mint at least
        // 2-3 aligned containers plus a typed partial on `finish`.
        let n: u64 = 300_000;
        for d in 0..n {
            builder.push_into((d, 0u64, 1i64));
        }

        let mut containers = 0;
        let mut out: Vec<(u64, u64, i64)> = Vec::new();
        while let Some(c) = builder.extract() {
            containers += 1;
            for r in rows(c) {
                out.push(r);
            }
        }
        if let Some(c) = builder.finish() {
            containers += 1;
            for r in rows(c) {
                out.push(r);
            }
        }
        assert!(
            containers > 1,
            "expected multiple containers, got {containers}"
        );
        out.sort();
        let expected: Vec<(u64, u64, i64)> = (0..n).map(|d| (d, 0u64, 1i64)).collect();
        assert_eq!(out, expected);
    }
}