From 47385cd1f593d5a6bdb07aef4d5d1714cb98c6d4 Mon Sep 17 00:00:00 2001 From: Frank McSherry Date: Tue, 8 Sep 2026 15:29:22 -0400 Subject: [PATCH 1/3] compute: store accumulable reduce accumulators in columnar form The accumulable reduce keeps one `Accum` per aggregate in the diff of its input arrangement. `Accum` is an enum sized for its `Numeric` variant, so every slot costs 112 bytes even when it holds a 24 byte `count` or integer `sum`. Deriving `Columnar` for `Accum` and storing the arrangement's diffs in differential's `Coltainer` lays the accumulators out by variant, so each one pays only for its own fields. Supporting changes: `OrderedNumericAgg` in `mz_repr` hosts the `Columnar` impl the orphan rule forbids on `OrderedDecimal`; `Overflowing: Columnar` reuses `T`'s container so `Overflowing` works; `RowLayout`, `RowSpine`, `RowBuilder` and `RowAgent` take a defaulted diff-container parameter; a `HeapSize` trait in `mz_timely_util` lets arrangement size logging cover both container kinds. Measured on 2M keys with count, three integer sums and one float sum, the accumulable arrangement shrinks from 1.21 GB to 436 MB (608 to 221 bytes per record). The batcher still stages `Accum` in columnation chunks, so its `Columnation` impl remains. Co-Authored-By: Claude Fable 5.1 --- src/compute/src/extensions/arrange.rs | 7 +- src/compute/src/render/reduce.rs | 27 ++- src/compute/src/typedefs.rs | 4 +- src/ore/src/overflowing.rs | 50 ++++- src/repr/src/adt/numeric.rs | 211 +++++++++++++++++++- src/row-spine/src/lib.rs | 36 +++- src/timely-util/src/containers.rs | 3 + src/timely-util/src/containers/heap_size.rs | 39 ++++ 8 files changed, 344 insertions(+), 33 deletions(-) create mode 100644 src/timely-util/src/containers/heap_size.rs diff --git a/src/compute/src/extensions/arrange.rs b/src/compute/src/extensions/arrange.rs index b06e1a56790b7..ad3febd5dc9d5 100644 --- a/src/compute/src/extensions/arrange.rs +++ b/src/compute/src/extensions/arrange.rs @@ -15,12 +15,14 @@ use differential_dataflow::difference::Semigroup; use differential_dataflow::lattice::Lattice; use differential_dataflow::operators::arrange::arrangement::arrange_core; use differential_dataflow::operators::arrange::{Arranged, TraceAgent}; +use differential_dataflow::trace::implementations::BatchContainer; use differential_dataflow::trace::implementations::spine_fueled::Spine; use differential_dataflow::trace::{Batch, Batcher, Builder, Trace, TraceReader}; use differential_dataflow::{Collection, Data, ExchangeData, Hashable, VecCollection}; use mz_compute_types::dyncfgs::{ENABLE_COLUMN_PAGED_BATCHER, ENABLE_COLUMNAR_MERGE_BATCHER}; use mz_dyncfg::ConfigSet; use mz_row_spine::ArcBatch; +use mz_timely_util::containers::HeapSize; use timely::Container; use timely::container::{ContainerBuilder, PushInto}; use timely::dataflow::Stream; @@ -497,10 +499,11 @@ where } } -impl<'scope, T, R> ArrangementSize for Arranged<'scope, RowAgent> +impl<'scope, T, R, DC> ArrangementSize for Arranged<'scope, RowAgent> where T: MzTimestamp, - R: Semigroup + Ord + MzArrangeData + 'static, + R: Semigroup + Ord + 'static, + DC: BatchContainer + HeapSize, { fn log_arrangement_size(self) -> Self { log_arrangement_size_inner(self, |batch| { diff --git a/src/compute/src/render/reduce.rs b/src/compute/src/render/reduce.rs index a5e8afcbb36fa..a30b796a14af0 100644 --- a/src/compute/src/render/reduce.rs +++ b/src/compute/src/render/reduce.rs @@ -13,10 +13,11 @@ use std::collections::BTreeMap; +use columnar::Columnar; use columnation::{Columnation, CopyRegion}; -use dec::OrderedDecimal; use differential_dataflow::Diff as _; use differential_dataflow::collection::AsCollection; +use differential_dataflow::columnar::layout::Coltainer; use differential_dataflow::consolidation::ConsolidatingContainerBuilder; use differential_dataflow::difference::{IsZero, Multiply, Semigroup}; use differential_dataflow::hashable::Hashable; @@ -35,7 +36,7 @@ use mz_compute_types::plan::reduce::{ use mz_compute_types::plan::scalar::LirScalarExpr; use mz_expr::{AggregateFunc, EvalError, SafeMfpPlan}; use mz_ore::cast::CastLossy; -use mz_repr::adt::numeric::{self, Numeric, NumericAgg}; +use mz_repr::adt::numeric::{self, Numeric, NumericAgg, OrderedNumericAgg}; use mz_repr::fixed_length::ExtendDatums; use mz_repr::{Datum, DatumVec, Diff, Row, RowArena, SharedRow}; use mz_timely_util::columnation::ColumnationChunker; @@ -1482,12 +1483,14 @@ impl<'scope, T: RenderTimestamp> Context<'scope, T> { let error_logger = self.error_logger(); let err_full_aggrs = full_aggrs.clone(); + // The diffs are stored columnar so that each `Accum` occupies only its own + // variant's columns, rather than the footprint of the largest variant. let arranged = collection .mz_arrange::< ColumnationChunker<_>, RowBatcher<_, _>, - RowBuilder<_, _>, - RowSpine<_, (Vec, Diff)>, + RowBuilder<_, _, Coltainer<_>>, + RowSpine<_, (Vec, Diff), Coltainer<_>>, >( "ArrangeAccumulable [val: empty]", ); @@ -1622,7 +1625,7 @@ fn accumulable_zero(aggr_func: &AggregateFunc) -> Accum { non_nulls: Diff::ZERO, }, AggregateFunc::SumNumeric => Accum::Numeric { - accum: OrderedDecimal(NumericAgg::zero()), + accum: OrderedNumericAgg(NumericAgg::zero()), pos_infs: Diff::ZERO, neg_infs: Diff::ZERO, nans: Diff::ZERO, @@ -1778,7 +1781,7 @@ fn datum_to_accumulator(aggregate_func: &AggregateFunc, datum: Datum) -> Accum { }; Accum::Numeric { - accum: OrderedDecimal(accum), + accum: OrderedNumericAgg(accum), pos_infs, neg_infs, nans, @@ -1786,7 +1789,7 @@ fn datum_to_accumulator(aggregate_func: &AggregateFunc, datum: Datum) -> Accum { } } Datum::Null => Accum::Numeric { - accum: OrderedDecimal(NumericAgg::zero()), + accum: OrderedNumericAgg(NumericAgg::zero()), pos_infs: Diff::ZERO, neg_infs: Diff::ZERO, nans: Diff::ZERO, @@ -2018,8 +2021,10 @@ type AccumCount = mz_ore::Overflowing; PartialOrd, Ord, Serialize, - Deserialize + Deserialize, + Columnar )] +#[columnar(derive(PartialEq, Eq, PartialOrd, Ord))] enum Accum { /// Accumulates boolean values. Bool { @@ -2052,7 +2057,7 @@ enum Accum { /// Accumulates arbitrary precision decimals. Numeric { /// Accumulates non-special values - accum: OrderedDecimal, + accum: OrderedNumericAgg, /// Counts +inf pos_infs: Diff, /// Counts -inf @@ -2254,7 +2259,7 @@ impl Multiply for Accum { // http://speleotrove.com/decimal/dncont.html assert!(!cx.status().rounded(), "Accum::Numeric multiply overflow"); Accum::Numeric { - accum: OrderedDecimal(f), + accum: OrderedNumericAgg(f), pos_infs: pos_infs * factor, neg_infs: neg_infs * factor, nans: nans * factor, @@ -2265,6 +2270,8 @@ impl Multiply for Accum { } } +// The batcher stages updates in columnation chunks before they reach the arrangement, +// which stores `Accum` in its columnar form. impl Columnation for Accum { type InnerRegion = CopyRegion; } diff --git a/src/compute/src/typedefs.rs b/src/compute/src/typedefs.rs index 0c71a624f5fbb..3a0883281dbdc 100644 --- a/src/compute/src/typedefs.rs +++ b/src/compute/src/typedefs.rs @@ -18,7 +18,7 @@ use differential_dataflow::trace::implementations::merge_batcher::MergeBatcher; use differential_dataflow::trace::wrappers::enter::TraceEnter; use differential_dataflow::trace::wrappers::frontier::TraceFrontier; use mz_repr::Diff; -use mz_timely_util::columnation::ColInternalMerger; +use mz_timely_util::columnation::{ColInternalMerger, ColumnationStack}; use mz_row_spine::RowValBuilder; @@ -96,7 +96,7 @@ pub type RowRowAgent = TraceAgent>; pub type RowRowArrangement<'scope, T> = Arranged<'scope, RowRowAgent>; pub type RowRowEnter = TraceEnter>, TEnter>; // Row specialized spines and agents. -pub type RowAgent = TraceAgent>; +pub type RowAgent> = TraceAgent>; pub type RowArrangement<'scope, T> = Arranged<'scope, RowAgent>; pub type RowEnter = TraceEnter>, TEnter>; diff --git a/src/ore/src/overflowing.rs b/src/ore/src/overflowing.rs index 93c5128d205d7..e8b6e0f2374c4 100644 --- a/src/ore/src/overflowing.rs +++ b/src/ore/src/overflowing.rs @@ -98,6 +98,7 @@ impl std::fmt::Display for Overflowing { #[cfg(feature = "columnar")] mod columnar { use crate::overflowing::Overflowing; + use columnar::bytes::indexed::DecodedStore; use columnar::common::PushIndexAs; use columnar::{ AsBytes, Borrow, Clear, Columnar, Container, FromBytes, Index, IndexAs, Len, Push, @@ -105,16 +106,15 @@ mod columnar { use serde::{Deserialize, Serialize}; use std::ops::Range; - impl Columnar for Overflowing + impl> + Copy + Send> Columnar for Overflowing where - for<'a> &'a [T]: AsBytes<'a> + FromBytes<'a>, Overflowing: From, { #[inline(always)] fn into_owned(other: columnar::Ref<'_, Self>) -> Self { other } - type Container = Overflows; + type Container = Overflows; #[inline(always)] fn reborrow<'b, 'a: 'b>(thing: columnar::Ref<'a, Self>) -> columnar::Ref<'b, Self> where @@ -124,9 +124,11 @@ mod columnar { } } - /// Columnar container for [`Overflowing`]. + /// Columnar container for [`Overflowing`], delegating to `T`'s own container `TC`, so + /// `Overflowing` uses columnar's byte-encoded `i128` store rather than requiring + /// `&[i128]` to be castable to bytes. #[derive(Copy, Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] - pub struct Overflows>(TC, std::marker::PhantomData); + pub struct Overflows(TC, std::marker::PhantomData); impl Default for Overflows { #[inline(always)] @@ -201,6 +203,13 @@ mod columnar { fn from_bytes(bytes: &mut impl Iterator) -> Self { Self(TC::from_bytes(bytes), std::marker::PhantomData) } + #[inline(always)] + fn from_store(store: &DecodedStore<'a>, offset: &mut usize) -> Self { + Self(TC::from_store(store, offset), std::marker::PhantomData) + } + fn element_sizes(sizes: &mut Vec) -> Result<(), String> { + TC::element_sizes(sizes) + } } impl Len for Overflows { @@ -235,10 +244,10 @@ mod columnar { } } - impl> Push<&Overflowing> for Overflows { + impl Push<&'a T>> Push<&Overflowing> for Overflows { #[inline(always)] fn push(&mut self, item: &Overflowing) { - self.0.push(item.0); + self.0.push(&item.0); } } } @@ -768,4 +777,31 @@ mod test { let result = Overflowing::::MAX.checked_add(Overflowing::::ONE); assert_eq!(result, None); } + + #[cfg(feature = "columnar")] + #[crate::test] + fn test_columnar_i128_round_trip() { + use ::columnar::{AsBytes, Borrow, BorrowedOf, Columnar, FromBytes, Index, Len}; + + let values = [ + Overflowing::::MIN, + Overflowing(-7), + Overflowing::::ZERO, + Overflowing::::ONE, + Overflowing::::MAX, + ]; + let container = Overflowing::::as_columns(values.iter()); + assert_eq!(container.len(), values.len()); + let borrowed = container.borrow(); + for (index, value) in values.iter().enumerate() { + assert_eq!(borrowed.get(index), *value); + } + + let bytes: Vec<&[u8]> = borrowed.as_bytes().map(|(_align, bytes)| bytes).collect(); + let decoded = BorrowedOf::>::from_bytes(&mut bytes.into_iter()); + assert_eq!(decoded.len(), values.len()); + for (index, value) in values.iter().enumerate() { + assert_eq!(decoded.get(index), *value); + } + } } diff --git a/src/repr/src/adt/numeric.rs b/src/repr/src/adt/numeric.rs index 81ba5dbfb681c..c93223fa55688 100644 --- a/src/repr/src/adt/numeric.rs +++ b/src/repr/src/adt/numeric.rs @@ -12,12 +12,13 @@ //! //! [`rust-dec`]: https://github.com/MaterializeInc/rust-dec/ +use std::cmp::Ordering; use std::error::Error; use std::fmt; use std::sync::LazyLock; use anyhow::bail; -use dec::{Context, Decimal}; +use dec::{Context, Decimal, OrderedDecimal}; use mz_ore::cast; use mz_persist_types::columnar::FixedSizeCodec; use mz_proto::{ProtoType, RustType, TryFromProtoError}; @@ -53,6 +54,35 @@ pub const NUMERIC_AGG_MAX_PRECISION: u8 = NUMERIC_AGG_WIDTH * 3; /// A double-width version of [`Numeric`] for use in aggregations. pub type NumericAgg = Decimal; +/// A [`NumericAgg`] with the total order of [`OrderedDecimal`], storable in columnar form. +/// +/// Equality and ordering are those of `OrderedDecimal`, so NaN equals NaN and every +/// value has a defined position. This crate cannot implement `Columnar` for +/// `OrderedDecimal`, as both the trait and the type are foreign, so the +/// newtype hosts that impl. +#[derive(Clone, Copy, Debug, Serialize, Deserialize)] +pub struct OrderedNumericAgg(pub NumericAgg); + +impl PartialEq for OrderedNumericAgg { + fn eq(&self, other: &Self) -> bool { + OrderedDecimal(self.0) == OrderedDecimal(other.0) + } +} + +impl Eq for OrderedNumericAgg {} + +impl PartialOrd for OrderedNumericAgg { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +impl Ord for OrderedNumericAgg { + fn cmp(&self, other: &Self) -> Ordering { + OrderedDecimal(self.0).cmp(&OrderedDecimal(other.0)) + } +} + static CX_DATUM: LazyLock> = LazyLock::new(|| { let mut cx = Context::::default(); cx.set_max_exponent(isize::from(NUMERIC_DATUM_MAX_PRECISION - 1)) @@ -861,6 +891,149 @@ impl FixedSizeCodec for PackedNumeric { } } +mod columnar_impls { + use std::ops::Range; + + use columnar::bytes::indexed::DecodedStore; + use columnar::{AsBytes, Borrow, Clear, Columnar, Container, FromBytes, Index, Len, Push}; + + use super::{NUMERIC_AGG_WIDTH_USIZE, NumericAgg, OrderedNumericAgg}; + + /// The raw parts of a [`NumericAgg`], in the order `Decimal::to_raw_parts` returns them. + type Parts = (u32, i32, u8, [u16; NUMERIC_AGG_WIDTH_USIZE]); + /// One column per raw part. The coefficient units use `Vec<[u16; N]>` directly rather + /// than the array's own columnar container, which would add per-element offsets to a + /// fixed-width value. + type PartsContainer = ( + Vec, + Vec, + Vec, + Vec<[u16; NUMERIC_AGG_WIDTH_USIZE]>, + ); + type PartsBorrowed<'a> = ::Borrowed<'a>; + + impl Columnar for OrderedNumericAgg { + #[inline(always)] + fn into_owned(other: columnar::Ref<'_, Self>) -> Self { + other + } + type Container = OrderedNumericAggs; + #[inline(always)] + fn reborrow<'b, 'a: 'b>(thing: columnar::Ref<'a, Self>) -> columnar::Ref<'b, Self> + where + Self: 'a, + { + thing + } + } + + /// Columnar container for [`OrderedNumericAgg`]. + /// + /// References are owned values rebuilt from the part columns, so comparisons on + /// references use the decimal order rather than the raw parts' lexicographic order. + #[derive(Copy, Clone, Debug, Default)] + pub struct OrderedNumericAggs(TC); + + impl Borrow for OrderedNumericAggs { + type Ref<'a> = OrderedNumericAgg; + type Borrowed<'a> = OrderedNumericAggs>; + #[inline(always)] + fn borrow<'a>(&'a self) -> Self::Borrowed<'a> { + OrderedNumericAggs(self.0.borrow()) + } + #[inline(always)] + fn reborrow<'b, 'a: 'b>(item: Self::Borrowed<'a>) -> Self::Borrowed<'b> + where + Self: 'a, + { + OrderedNumericAggs(::reborrow(item.0)) + } + #[inline(always)] + fn reborrow_ref<'b, 'a: 'b>(item: Self::Ref<'a>) -> Self::Ref<'b> + where + Self: 'a, + { + item + } + } + + impl Container for OrderedNumericAggs { + #[inline(always)] + fn extend_from_self(&mut self, other: Self::Borrowed<'_>, range: Range) { + self.0.extend_from_self(other.0, range); + } + #[inline(always)] + fn reserve_for<'a, I>(&mut self, selves: I) + where + Self: 'a, + I: Iterator> + Clone, + { + self.0.reserve_for(selves.map(|s| s.0)); + } + } + + impl Len for OrderedNumericAggs { + #[inline(always)] + fn len(&self) -> usize { + self.0.len() + } + } + + impl Clear for OrderedNumericAggs { + #[inline(always)] + fn clear(&mut self) { + self.0.clear(); + } + } + + impl<'a> Index for OrderedNumericAggs> { + type Ref = OrderedNumericAgg; + #[inline(always)] + fn get(&self, index: usize) -> Self::Ref { + let (digits, exponent, bits, lsu) = self.0.get(index); + OrderedNumericAgg(NumericAgg::from_raw_parts(*digits, *exponent, *bits, *lsu)) + } + } + + impl Push for OrderedNumericAggs { + #[inline(always)] + fn push(&mut self, item: OrderedNumericAgg) { + let parts: Parts = item.0.to_raw_parts(); + self.0.push(parts); + } + } + + impl Push<&OrderedNumericAgg> for OrderedNumericAggs { + #[inline(always)] + fn push(&mut self, item: &OrderedNumericAgg) { + self.push(*item); + } + } + + impl<'a, TC: AsBytes<'a>> AsBytes<'a> for OrderedNumericAggs { + const SLICE_COUNT: usize = TC::SLICE_COUNT; + #[inline(always)] + fn get_byte_slice(&self, index: usize) -> (u64, &'a [u8]) { + self.0.get_byte_slice(index) + } + } + + impl<'a, TC: FromBytes<'a>> FromBytes<'a> for OrderedNumericAggs { + const SLICE_COUNT: usize = TC::SLICE_COUNT; + #[inline(always)] + fn from_bytes(bytes: &mut impl Iterator) -> Self { + OrderedNumericAggs(TC::from_bytes(bytes)) + } + #[inline(always)] + fn from_store(store: &DecodedStore<'a>, offset: &mut usize) -> Self { + OrderedNumericAggs(TC::from_store(store, offset)) + } + fn element_sizes(sizes: &mut Vec) -> Result<(), String> { + TC::element_sizes(sizes) + } + } +} + #[cfg(test)] mod tests { use mz_ore::assert_ok; @@ -961,4 +1134,40 @@ mod tests { insta::assert_debug_snapshot!(all_numerics); } + + #[mz_ore::test] + fn ordered_numeric_agg_columnar_round_trip() { + use columnar::{AsBytes, Borrow, BorrowedOf, Columnar, FromBytes, Index, Len}; + + let mut cx = cx_agg(); + let values: Vec = [ + "0", + "-0", + "1", + "-12345.678", + "9e39", + "9e-39", + "123456789012345678901234567890123456789012345678901234567890", + "NaN", + "Infinity", + "-Infinity", + ] + .into_iter() + .map(|s| OrderedNumericAgg(cx.parse(s).unwrap())) + .collect(); + + let container = OrderedNumericAgg::as_columns(values.iter()); + assert_eq!(container.len(), values.len()); + let borrowed = container.borrow(); + for (index, value) in values.iter().enumerate() { + assert_eq!(borrowed.get(index), *value); + } + + let bytes: Vec<&[u8]> = borrowed.as_bytes().map(|(_align, bytes)| bytes).collect(); + let decoded = BorrowedOf::::from_bytes(&mut bytes.into_iter()); + assert_eq!(decoded.len(), values.len()); + for (index, value) in values.iter().enumerate() { + assert_eq!(decoded.get(index), *value); + } + } } diff --git a/src/row-spine/src/lib.rs b/src/row-spine/src/lib.rs index a64795a1c3008..dba8a51d71ef0 100644 --- a/src/row-spine/src/lib.rs +++ b/src/row-spine/src/lib.rs @@ -34,6 +34,7 @@ pub static DICTIONARY_COMPRESSION: std::sync::atomic::AtomicBool = /// Spines specialized to contain `Row` types in keys and values. mod spines { use columnation::Columnation; + use differential_dataflow::trace::implementations::BatchContainer; use differential_dataflow::trace::implementations::Layout; use differential_dataflow::trace::implementations::Update; use differential_dataflow::trace::implementations::Vector; @@ -76,9 +77,12 @@ mod spines { pub type RowValBuilder = ArcBuilder>; - pub type RowSpine = Spine>>>; + /// Key-only `Row` spine. `DC` is the diff container, see [`RowLayout`]. + pub type RowSpine> = + Spine>>>; pub type RowBatcher = KeyBatcher; - pub type RowBuilder = ArcBuilder>; + pub type RowBuilder> = + ArcBuilder>; pub type ValRowSpine = Spine>>>; pub type ValRowBatcher = KeyValBatcher; @@ -116,8 +120,13 @@ mod spines { pub struct RowValLayout> { phantom: std::marker::PhantomData, } - pub struct RowLayout> { - phantom: std::marker::PhantomData, + /// Layout for key-only `Row` updates. `DC` is the diff `BatchContainer`, a + /// columnation stack by default. + pub struct RowLayout::Diff>> + where + U: Update, + { + phantom: std::marker::PhantomData<(U, DC)>, } /// Mirror of [`RowValLayout`] with the roles swapped: arbitrary `Columnation` /// keys with `Row` values stored as packed bytes in a [`DatumContainer`]. @@ -148,15 +157,15 @@ mod spines { type DiffContainer = ColumnationStack; type OffsetContainer = OffsetOptimized; } - impl> Layout for RowLayout + impl, DC> Layout for RowLayout where U::Time: Columnation, - U::Diff: Columnation, + DC: BatchContainer, { type KeyContainer = DatumContainer; type ValContainer = ColumnationStack<()>; type TimeContainer = ColumnationStack; - type DiffContainer = ColumnationStack; + type DiffContainer = DC; type OffsetContainer = OffsetOptimized; } impl> Layout for ValRowLayout @@ -903,6 +912,7 @@ mod dictionary { use differential_dataflow::lattice::Lattice; use differential_dataflow::trace::Builder; use differential_dataflow::trace::Description; + use differential_dataflow::trace::implementations::BatchContainer; use differential_dataflow::trace::implementations::ord_neu::{OrdKeyBatch, OrdKeyBuilder}; use differential_dataflow::trace::implementations::ord_neu::{OrdValBatch, OrdValBuilder}; use mz_timely_util::columnar::Column; @@ -1081,16 +1091,20 @@ mod dictionary { pub struct RowBuilder< T: Lattice + Timestamp + Columnation, R: Ord + Semigroup + Columnation + 'static, + DC: BatchContainer = TimelyStack, > { - inner: OrdKeyBuilder, TimelyStack<((Row, ()), T, R)>>, + inner: OrdKeyBuilder, TimelyStack<((Row, ()), T, R)>>, } - impl - Builder for RowBuilder + impl Builder for RowBuilder + where + T: Lattice + Timestamp + Columnation, + R: Ord + Semigroup + Columnation + 'static, + DC: BatchContainer, { type Input = TimelyStack<((Row, ()), T, R)>; type Time = T; - type Output = OrdKeyBatch>; + type Output = OrdKeyBatch>; fn with_capacity(keys: usize, vals: usize, upds: usize) -> Self { Self { diff --git a/src/timely-util/src/containers.rs b/src/timely-util/src/containers.rs index 00dc8b7aafeaa..e6cb9845ef2e1 100644 --- a/src/timely-util/src/containers.rs +++ b/src/timely-util/src/containers.rs @@ -15,4 +15,7 @@ //! Reusable containers. +pub mod heap_size; pub mod stack; + +pub use heap_size::HeapSize; diff --git a/src/timely-util/src/containers/heap_size.rs b/src/timely-util/src/containers/heap_size.rs new file mode 100644 index 0000000000000..036d7d77ebf76 --- /dev/null +++ b/src/timely-util/src/containers/heap_size.rs @@ -0,0 +1,39 @@ +// Copyright Materialize, Inc. and contributors. All rights reserved. +// +// Use of this software is governed by the Business Source License +// included in the LICENSE file. +// +// As of the Change Date specified in that file, in accordance with +// the Business Source License, use of this software will be governed +// by the Apache License, Version 2.0. + +//! Heap size accounting for the containers that back arrangement batches. + +use columnar::{AsBytes, Borrow, Columnar}; +use columnation::Columnation; +use differential_dataflow::columnar::layout::Coltainer; + +use crate::columnation::ColumnationStack; + +/// A container that can report the heap allocations backing it. +pub trait HeapSize { + /// Calls `callback(size, capacity)`, in bytes, once per allocation backing `self`. + fn heap_size(&self, callback: impl FnMut(usize, usize)); +} + +impl HeapSize for ColumnationStack { + fn heap_size(&self, callback: impl FnMut(usize, usize)) { + ColumnationStack::heap_size(self, callback) + } +} + +impl HeapSize for Coltainer { + fn heap_size(&self, mut callback: impl FnMut(usize, usize)) { + // Columnar containers expose their contents as byte slices but not their spare + // capacity, so each slice reports its length as both size and capacity. The + // capacity is therefore a lower bound. + for (_align, bytes) in self.container.borrow().as_bytes() { + callback(bytes.len(), bytes.len()); + } + } +} From 14ed0f694232fb8d59fa27e5be2608c7c16d29df Mon Sep 17 00:00:00 2001 From: Frank McSherry Date: Wed, 9 Sep 2026 09:28:35 -0400 Subject: [PATCH 2/3] compute: address review of the columnar Accum arrangement Unbracket the `RowLayout` reference in `RowSpine`'s doc, which is a private item and failed the rustdoc lint. Write the `RowSpine` arrangement-size impl over the diff container alone, since the diff type is its `Owned` type. Widen the numeric accumulator's coefficient column from 27 to 28 units so each element is a whole number of `u64` words. columnar's indexed byte store pads columns to words, and its array decoder cannot strip padding whose width does not divide the element width, so the 54 byte column panicked on decode. The round-trip tests now also decode through the indexed store, which fails without the widening. Add a test that drives `Accum` of every variant through its derived container and through `Coltainer<(Vec, Diff)>`, checking round trips, byte decoding, and that reference ordering agrees with owned ordering, which the columnar diff container relies on. Note that agreement at the derive, and note what `HeapSize` reports for columnar containers. Co-Authored-By: Claude Fable 5.1 --- src/compute/src/extensions/arrange.rs | 5 +- src/compute/src/render/reduce.rs | 108 ++++++++++++++++++++ src/ore/src/overflowing.rs | 11 ++ src/repr/src/adt/numeric.rs | 38 +++++-- src/row-spine/src/lib.rs | 2 +- src/timely-util/src/containers/heap_size.rs | 7 +- 6 files changed, 154 insertions(+), 17 deletions(-) diff --git a/src/compute/src/extensions/arrange.rs b/src/compute/src/extensions/arrange.rs index ad3febd5dc9d5..73cc2a9589608 100644 --- a/src/compute/src/extensions/arrange.rs +++ b/src/compute/src/extensions/arrange.rs @@ -499,11 +499,10 @@ where } } -impl<'scope, T, R, DC> ArrangementSize for Arranged<'scope, RowAgent> +impl<'scope, T, DC> ArrangementSize for Arranged<'scope, RowAgent> where T: MzTimestamp, - R: Semigroup + Ord + 'static, - DC: BatchContainer + HeapSize, + DC: BatchContainer + HeapSize, { fn log_arrangement_size(self) -> Self { log_arrangement_size_inner(self, |batch| { diff --git a/src/compute/src/render/reduce.rs b/src/compute/src/render/reduce.rs index a30b796a14af0..8b491b28a7bcd 100644 --- a/src/compute/src/render/reduce.rs +++ b/src/compute/src/render/reduce.rs @@ -2024,6 +2024,8 @@ type AccumCount = mz_ore::Overflowing; Deserialize, Columnar )] +// The columnar container orders references with this derived `Ord`, which must agree with +// the owned `Ord`. It does because every field's reference type is its owned type. #[columnar(derive(PartialEq, Eq, PartialOrd, Ord))] enum Accum { /// Accumulates boolean values. @@ -2729,4 +2731,110 @@ mod tests { let datum = finalize_accum(&func, &acc, Diff::from(2_i64)); assert_eq!(datum, Datum::from(0.0_f64)); } + + /// Accumulators of every variant, in zero, accumulated, and negated states. + fn sample_accums() -> Vec { + let mut cx = numeric::cx_datum(); + let mut numeric = |s: &str| Datum::from(cx.parse(s).unwrap()); + let cases: Vec<(AggregateFunc, Vec)> = vec![ + (AggregateFunc::Count, vec![Datum::Null, Datum::Int64(5)]), + ( + AggregateFunc::SumInt64, + vec![Datum::Int64(-7), Datum::Int64(i64::MAX)], + ), + ( + AggregateFunc::SumUInt16, + vec![Datum::UInt16(3), Datum::Null], + ), + ( + AggregateFunc::Any, + vec![Datum::True, Datum::False, Datum::Null], + ), + ( + AggregateFunc::SumFloat64, + vec![ + Datum::from(1.5_f64), + Datum::from(f64::NAN), + Datum::from(f64::NEG_INFINITY), + ], + ), + ( + AggregateFunc::SumNumeric, + vec![ + numeric("-12345.678"), + numeric("9e39"), + numeric("NaN"), + numeric("Infinity"), + Datum::Null, + ], + ), + ]; + let mut accums = Vec::new(); + for (func, datums) in cases { + let mut sum = accumulable_zero(&func); + accums.push(sum); + for datum in datums { + let accum = datum_to_accumulator(&func, datum); + sum.plus_equals(&accum); + accums.push(accum); + accums.push(accum.multiply(&Diff::from(-1_i64))); + } + accums.push(sum); + } + accums + } + + #[mz_ore::test] + fn accum_columnar_round_trip() { + use columnar::bytes::indexed::{DecodedStore, encode}; + use columnar::{AsBytes, Borrow, BorrowedOf, FromBytes, Index, Len}; + use differential_dataflow::trace::implementations::BatchContainer; + + let accums = sample_accums(); + let container = Accum::as_columns(accums.iter()); + assert_eq!(container.len(), accums.len()); + let borrowed = container.borrow(); + for (index, accum) in accums.iter().enumerate() { + assert_eq!(Accum::into_owned(borrowed.get(index)), *accum); + } + for (i, a) in accums.iter().enumerate() { + for (j, b) in accums.iter().enumerate() { + assert_eq!(borrowed.get(i).cmp(&borrowed.get(j)), a.cmp(b)); + } + } + + let bytes: Vec<&[u8]> = borrowed.as_bytes().map(|(_align, bytes)| bytes).collect(); + let decoded = BorrowedOf::::from_bytes(&mut bytes.into_iter()); + for (index, accum) in accums.iter().enumerate() { + assert_eq!(Accum::into_owned(decoded.get(index)), *accum); + } + // NOTE: the `i128` columns cannot be `validate`d, see the `Overflowing` test in + // `mz_ore`, so this only decodes. + let mut words = Vec::new(); + encode(&mut words, &borrowed); + let decoded = BorrowedOf::::from_store(&DecodedStore::new(&words), &mut 0); + for (index, accum) in accums.iter().enumerate() { + assert_eq!(Accum::into_owned(decoded.get(index)), *accum); + } + + // The arrangement's diff container, holding whole `(Vec, Diff)` diffs. + let diffs: Vec<(Vec, Diff)> = accums + .chunks(3) + .map(|chunk| (chunk.to_vec(), Diff::ONE)) + .collect(); + let mut coltainer = Coltainer::<(Vec, Diff)>::default(); + for diff in &diffs { + coltainer.push_own(diff); + } + assert_eq!(coltainer.len(), diffs.len()); + for (index, diff) in diffs.iter().enumerate() { + assert_eq!( + , Diff)>>::into_owned(coltainer.index(index)), + *diff + ); + } + let mut sum = , Diff)>>::into_owned(coltainer.index(0)); + sum.plus_equals(&sum.clone().multiply(&Diff::from(-1_i64))); + assert!(sum.is_zero()); + } } diff --git a/src/ore/src/overflowing.rs b/src/ore/src/overflowing.rs index e8b6e0f2374c4..0a15a39c07555 100644 --- a/src/ore/src/overflowing.rs +++ b/src/ore/src/overflowing.rs @@ -803,5 +803,16 @@ mod test { for (index, value) in values.iter().enumerate() { assert_eq!(decoded.get(index), *value); } + + // NOTE: columnar's `i128` store does not implement `element_sizes`, so the indexed + // store can be decoded but not `validate`d for this type. + let mut words = Vec::new(); + ::columnar::bytes::indexed::encode(&mut words, &borrowed); + let store = ::columnar::bytes::indexed::DecodedStore::new(&words); + let decoded = BorrowedOf::>::from_store(&store, &mut 0); + assert_eq!(decoded.len(), values.len()); + for (index, value) in values.iter().enumerate() { + assert_eq!(decoded.get(index), *value); + } } } diff --git a/src/repr/src/adt/numeric.rs b/src/repr/src/adt/numeric.rs index c93223fa55688..341d5af9bf25c 100644 --- a/src/repr/src/adt/numeric.rs +++ b/src/repr/src/adt/numeric.rs @@ -899,17 +899,18 @@ mod columnar_impls { use super::{NUMERIC_AGG_WIDTH_USIZE, NumericAgg, OrderedNumericAgg}; - /// The raw parts of a [`NumericAgg`], in the order `Decimal::to_raw_parts` returns them. - type Parts = (u32, i32, u8, [u16; NUMERIC_AGG_WIDTH_USIZE]); + /// Width of the coefficient column. One unit wider than the decimal's coefficient so + /// that each element is 56 bytes, a whole number of `u64` words. The indexed byte + /// store pads every column to whole words and cannot decode an element width that + /// does not divide that padding. + const LSU_COLUMN_WIDTH: usize = NUMERIC_AGG_WIDTH_USIZE + 1; + /// The raw parts of a [`NumericAgg`], in the order `Decimal::to_raw_parts` returns + /// them, with the coefficient units zero-padded to [`LSU_COLUMN_WIDTH`]. + type Parts = (u32, i32, u8, [u16; LSU_COLUMN_WIDTH]); /// One column per raw part. The coefficient units use `Vec<[u16; N]>` directly rather /// than the array's own columnar container, which would add per-element offsets to a /// fixed-width value. - type PartsContainer = ( - Vec, - Vec, - Vec, - Vec<[u16; NUMERIC_AGG_WIDTH_USIZE]>, - ); + type PartsContainer = (Vec, Vec, Vec, Vec<[u16; LSU_COLUMN_WIDTH]>); type PartsBorrowed<'a> = ::Borrowed<'a>; impl Columnar for OrderedNumericAgg { @@ -991,14 +992,19 @@ mod columnar_impls { #[inline(always)] fn get(&self, index: usize) -> Self::Ref { let (digits, exponent, bits, lsu) = self.0.get(index); - OrderedNumericAgg(NumericAgg::from_raw_parts(*digits, *exponent, *bits, *lsu)) + let mut units = [0u16; NUMERIC_AGG_WIDTH_USIZE]; + units.copy_from_slice(&lsu[..NUMERIC_AGG_WIDTH_USIZE]); + OrderedNumericAgg(NumericAgg::from_raw_parts(*digits, *exponent, *bits, units)) } } impl Push for OrderedNumericAggs { #[inline(always)] fn push(&mut self, item: OrderedNumericAgg) { - let parts: Parts = item.0.to_raw_parts(); + let (digits, exponent, bits, units) = item.0.to_raw_parts(); + let mut lsu = [0u16; LSU_COLUMN_WIDTH]; + lsu[..NUMERIC_AGG_WIDTH_USIZE].copy_from_slice(&units); + let parts: Parts = (digits, exponent, bits, lsu); self.0.push(parts); } } @@ -1169,5 +1175,17 @@ mod tests { for (index, value) in values.iter().enumerate() { assert_eq!(decoded.get(index), *value); } + + // The indexed store pads each column to whole words, so decoding through it + // also checks that the column layout tolerates that padding. + let mut words = Vec::new(); + columnar::bytes::indexed::encode(&mut words, &borrowed); + columnar::bytes::indexed::validate::>(&words).unwrap(); + let store = columnar::bytes::indexed::DecodedStore::new(&words); + let decoded = BorrowedOf::::from_store(&store, &mut 0); + assert_eq!(decoded.len(), values.len()); + for (index, value) in values.iter().enumerate() { + assert_eq!(decoded.get(index), *value); + } } } diff --git a/src/row-spine/src/lib.rs b/src/row-spine/src/lib.rs index dba8a51d71ef0..80ca39a875511 100644 --- a/src/row-spine/src/lib.rs +++ b/src/row-spine/src/lib.rs @@ -77,7 +77,7 @@ mod spines { pub type RowValBuilder = ArcBuilder>; - /// Key-only `Row` spine. `DC` is the diff container, see [`RowLayout`]. + /// Key-only `Row` spine. `DC` is the diff container, see `RowLayout`. pub type RowSpine> = Spine>>>; pub type RowBatcher = KeyBatcher; diff --git a/src/timely-util/src/containers/heap_size.rs b/src/timely-util/src/containers/heap_size.rs index 036d7d77ebf76..d997e107b624e 100644 --- a/src/timely-util/src/containers/heap_size.rs +++ b/src/timely-util/src/containers/heap_size.rs @@ -29,9 +29,10 @@ impl HeapSize for ColumnationStack { impl HeapSize for Coltainer { fn heap_size(&self, mut callback: impl FnMut(usize, usize)) { - // Columnar containers expose their contents as byte slices but not their spare - // capacity, so each slice reports its length as both size and capacity. The - // capacity is therefore a lower bound. + // Columnar containers expose their contents as byte slices, one per column, and each + // non-empty column is one `Vec` allocation, so the callback count is right. They do + // not expose spare capacity, so each slice reports its length as both size and + // capacity, and the capacity is a lower bound. for (_align, bytes) in self.container.borrow().as_bytes() { callback(bytes.len(), bytes.len()); } From 4778bde0f4be016489a37ee2103749169afd9393 Mon Sep 17 00:00:00 2001 From: Frank McSherry Date: Thu, 10 Sep 2026 16:11:34 -0400 Subject: [PATCH 3/3] compute: gate the columnar accumulable diff behind a replica-scoped dyncfg `enable_columnar_accumulable_diff` selects the diff container of the accumulable reduce's arrangement at render time. It defaults off in production, keeping the columnation stack, and on in CI, where it is also randomized, so both layouts stay exercised. The choice is one `if` at the arrange site. The two reduce operators that consume the arrangement move into a helper generic over the diff container, so they are rendered once and the types at the operator boundaries do not change. Co-Authored-By: Claude Fable 5.1 --- misc/python/materialize/mzcompose/__init__.py | 5 ++ .../materialize/parallel_workload/action.py | 1 + src/compute-types/src/dyncfgs.rs | 17 +++++ src/compute/src/render/reduce.rs | 63 ++++++++++++++----- .../mzcompose.py | 1 + 5 files changed, 73 insertions(+), 14 deletions(-) diff --git a/misc/python/materialize/mzcompose/__init__.py b/misc/python/materialize/mzcompose/__init__.py index a2e156a9e627b..042251f16b068 100644 --- a/misc/python/materialize/mzcompose/__init__.py +++ b/misc/python/materialize/mzcompose/__init__.py @@ -284,6 +284,11 @@ def get_variable_system_parameters( VariableSystemParameter( "enable_columnar_merge_batcher", "true", ["true", "false"] ), + # On by default so CI exercises the columnar accumulable diff layout, which + # is off in production while it earns trust. + VariableSystemParameter( + "enable_columnar_accumulable_diff", "true", ["true", "false"] + ), VariableSystemParameter( "compute_peek_response_stash_threshold_bytes", # 1 MiB, an in-between value diff --git a/misc/python/materialize/parallel_workload/action.py b/misc/python/materialize/parallel_workload/action.py index 6a008fc235834..f2d61ed41340a 100644 --- a/misc/python/materialize/parallel_workload/action.py +++ b/misc/python/materialize/parallel_workload/action.py @@ -3141,6 +3141,7 @@ def __init__( self.flags_with_values["enable_compute_sync_mv_sink"] = BOOLEAN_FLAG_VALUES self.flags_with_values["enable_column_paged_batcher"] = BOOLEAN_FLAG_VALUES self.flags_with_values["enable_columnar_merge_batcher"] = BOOLEAN_FLAG_VALUES + self.flags_with_values["enable_columnar_accumulable_diff"] = BOOLEAN_FLAG_VALUES self.flags_with_values["enable_column_paged_batcher_spill"] = ( BOOLEAN_FLAG_VALUES ) diff --git a/src/compute-types/src/dyncfgs.rs b/src/compute-types/src/dyncfgs.rs index 01cf2f5292d9c..1d74ae1ba03f0 100644 --- a/src/compute-types/src/dyncfgs.rs +++ b/src/compute-types/src/dyncfgs.rs @@ -90,6 +90,22 @@ pub const ENABLE_COLUMNAR_MERGE_BATCHER: Config = Config::new( ParameterScope::Replica, ); +/// Store the accumulable reduce's accumulators in columnar form. +/// +/// The accumulable reduce keeps one accumulator per aggregate in the diff of its +/// input arrangement. When `true`, that arrangement holds its diffs in a columnar +/// container, which lays the accumulators out by variant so each pays only for its +/// own fields. When `false` (the default), the diffs live in a columnation stack at +/// the width of the largest variant. Read at operator construction time, so flips +/// take effect on dataflows created after the change. +pub const ENABLE_COLUMNAR_ACCUMULABLE_DIFF: Config = Config::new( + "enable_columnar_accumulable_diff", + false, + "Store the accumulable reduce's accumulators in a columnar arrangement diff, laid out by \ + variant, instead of a columnation stack of fixed-width accumulator enums.", + ParameterScope::Replica, +); + /// Allow the column-paged batcher's pager to evict chunks under memory /// pressure. Only meaningful when [`ENABLE_COLUMN_PAGED_BATCHER`] is `true`. /// With the spill flag off the pager keeps every chunk resident regardless of @@ -844,6 +860,7 @@ pub fn all_dyncfgs(configs: ConfigSet) -> ConfigSet { .add(&MV_SINK_ADVANCE_PERSIST_FRONTIERS) .add(&ENABLE_COLUMN_PAGED_BATCHER) .add(&ENABLE_COLUMNAR_MERGE_BATCHER) + .add(&ENABLE_COLUMNAR_ACCUMULABLE_DIFF) .add(&ENABLE_COLUMN_PAGED_BATCHER_SPILL) .add(&COLUMN_PAGED_BATCHER_BUDGET_FRACTION) .add(&COLUMN_PAGED_BATCHER_LZ4) diff --git a/src/compute/src/render/reduce.rs b/src/compute/src/render/reduce.rs index 8b491b28a7bcd..3d896f2c4eed3 100644 --- a/src/compute/src/render/reduce.rs +++ b/src/compute/src/render/reduce.rs @@ -27,7 +27,9 @@ use differential_dataflow::trace::implementations::BatchContainer; use differential_dataflow::trace::{Builder, Cursor, Navigable, Trace}; use differential_dataflow::{Data, VecCollection}; use itertools::Itertools; -use mz_compute_types::dyncfgs::{ENABLE_COMPUTE_TEMPORAL_BUCKETING, TEMPORAL_BUCKETING_SUMMARY}; +use mz_compute_types::dyncfgs::{ + ENABLE_COLUMNAR_ACCUMULABLE_DIFF, ENABLE_COMPUTE_TEMPORAL_BUCKETING, TEMPORAL_BUCKETING_SUMMARY, +}; use mz_compute_types::plan::ArrangementStrategy; use mz_compute_types::plan::reduce::{ AccumulablePlan, BasicPlan, BucketedPlan, HierarchicalPlan, KeyValPlan, LirAggregateExpr, @@ -55,8 +57,8 @@ use crate::render::errors::MaybeValidatingRow; use crate::render::reduce::monoids::{ReductionMonoid, get_monoid}; use crate::render::{ArrangementFlavor, Pairer, RenderTimestamp}; use crate::typedefs::{ - ErrBatcher, ErrBuilder, KeyBatcher, RowErrBuilder, RowErrSpine, RowRowAgent, RowRowArrangement, - RowRowSpine, RowSpine, RowValSpine, + ErrBatcher, ErrBuilder, KeyBatcher, RowAgent, RowErrBuilder, RowErrSpine, RowRowAgent, + RowRowArrangement, RowRowSpine, RowSpine, RowValSpine, }; use mz_row_spine::{ DatumContainer, DatumSeq, RowBatcher, RowBuilder, RowRowBatcher, RowRowBuilder, RowValBatcher, @@ -1474,6 +1476,50 @@ impl<'scope, T: RenderTimestamp> Context<'scope, T> { differential_dataflow::collection::concatenate(collection_scope, to_aggregate) }; + // The accumulators travel in the arrangement's diffs. A columnar diff container + // lays each `Accum` out by variant, so it occupies only its own variant's + // columns rather than the footprint of the largest variant. Both layouts feed + // the same reduce operators. + if ENABLE_COLUMNAR_ACCUMULABLE_DIFF.get(&self.config_set) { + let arranged = collection + .mz_arrange::< + ColumnationChunker<_>, + RowBatcher<_, _>, + RowBuilder<_, _, Coltainer<_>>, + RowSpine<_, (Vec, Diff), Coltainer<_>>, + >( + "ArrangeAccumulable [val: empty]", + ); + self.reduce_accumulable(arranged, full_aggrs, mfp_after) + } else { + let arranged = collection + .mz_arrange::< + ColumnationChunker<_>, + RowBatcher<_, _>, + RowBuilder<_, _>, + RowSpine<_, (Vec, Diff)>, + >( + "ArrangeAccumulable [val: empty]", + ); + self.reduce_accumulable(arranged, full_aggrs, mfp_after) + } + } + + /// Reduces arranged accumulators to output rows, and to the errors the accumulated + /// values can reveal. Generic over the container holding the diffs, so both diff + /// layouts share one rendering of the reduce operators. + fn reduce_accumulable<'s, DC>( + &self, + arranged: Arranged<'s, RowAgent, Diff), DC>>, + full_aggrs: Vec, + mfp_after: Option>, + ) -> ( + RowRowArrangement<'s, T>, + VecCollection<'s, T, DataflowErrorSer, Diff>, + ) + where + DC: BatchContainer, Diff)>, + { // Allocations for the two closures. let mut datums1 = DatumVec::new(); let mut datums2 = DatumVec::new(); @@ -1483,17 +1529,6 @@ impl<'scope, T: RenderTimestamp> Context<'scope, T> { let error_logger = self.error_logger(); let err_full_aggrs = full_aggrs.clone(); - // The diffs are stored columnar so that each `Accum` occupies only its own - // variant's columns, rather than the footprint of the largest variant. - let arranged = collection - .mz_arrange::< - ColumnationChunker<_>, - RowBatcher<_, _>, - RowBuilder<_, _, Coltainer<_>>, - RowSpine<_, (Vec, Diff), Coltainer<_>>, - >( - "ArrangeAccumulable [val: empty]", - ); let arranged_output = arranged .clone() .mz_reduce_abelian::<_, RowRowBuilder<_, _>, RowRowSpine<_, _>, _>( diff --git a/test/launchdarkly-flag-consistency/mzcompose.py b/test/launchdarkly-flag-consistency/mzcompose.py index 618411c9e9835..dca6979751586 100644 --- a/test/launchdarkly-flag-consistency/mzcompose.py +++ b/test/launchdarkly-flag-consistency/mzcompose.py @@ -246,6 +246,7 @@ enable_statement_arrival_logging enable_binary_date_bin enable_coalesce_case_transform + enable_columnar_accumulable_diff enable_columnar_merge_batcher enable_compute_half_join2 enable_compute_index_peek_offload