compute: store accumulable reduce accumulators in columnar form

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

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
         )

src/compute-types/src/dyncfgs.rs

@@ -90,6 +90,22 @@ pub const ENABLE_COLUMNAR_MERGE_BATCHER: Config<bool> = 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<bool> = 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)

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,10 @@ where
     }
 }
 
-impl<'scope, T, R> ArrangementSize for Arranged<'scope, RowAgent<T, R>>
+impl<'scope, T, DC> ArrangementSize for Arranged<'scope, RowAgent<T, DC::Owned, DC>>
 where
     T: MzTimestamp,
-    R: Semigroup + Ord + MzArrangeData + 'static,
+    DC: BatchContainer<Owned: Semigroup + 'static> + HeapSize,
 {
     fn log_arrangement_size(self) -> Self {
         log_arrangement_size_inner(self, |batch| {

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;
@@ -26,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,
@@ -35,7 +38,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;
@@ -54,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,
@@ -1473,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<Accum>, Diff), Coltainer<_>>,
+                >(
+                    "ArrangeAccumulable [val: empty]",
+                );
+            self.reduce_accumulable(arranged, full_aggrs, mfp_after)
+        } else {
+            let arranged = collection
+                .mz_arrange::<
+                    ColumnationChunker<_>,
+                    RowBatcher<_, _>,
+                    RowBuilder<_, _>,
+                    RowSpine<_, (Vec<Accum>, 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<T, (Vec<Accum>, Diff), DC>>,
+        full_aggrs: Vec<LirAggregateExpr>,
+        mfp_after: Option<SafeMfpPlan<LirScalarExpr>>,
+    ) -> (
+        RowRowArrangement<'s, T>,
+        VecCollection<'s, T, DataflowErrorSer, Diff>,
+    )
+    where
+        DC: BatchContainer<Owned = (Vec<Accum>, Diff)>,
+    {
         // Allocations for the two closures.
         let mut datums1 = DatumVec::new();
         let mut datums2 = DatumVec::new();
@@ -1482,15 +1529,6 @@ impl<'scope, T: RenderTimestamp> Context<'scope, T> {
 
         let error_logger = self.error_logger();
         let err_full_aggrs = full_aggrs.clone();
-        let arranged = collection
-            .mz_arrange::<
-                ColumnationChunker<_>,
-                RowBatcher<_, _>,
-                RowBuilder<_, _>,
-                RowSpine<_, (Vec<Accum>, Diff)>,
-            >(
-                "ArrangeAccumulable [val: empty]",
-            );
         let arranged_output = arranged
             .clone()
             .mz_reduce_abelian::<_, RowRowBuilder<_, _>, RowRowSpine<_, _>, _>(
@@ -1622,7 +1660,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,15 +1816,15 @@ fn datum_to_accumulator(aggregate_func: &AggregateFunc, datum: Datum) -> Accum {
                 };
 
                 Accum::Numeric {
-                    accum: OrderedDecimal(accum),
+                    accum: OrderedNumericAgg(accum),
                     pos_infs,
                     neg_infs,
                     nans,
                     non_nulls: Diff::ONE,
                 }
             }
             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 +2056,12 @@ type AccumCount = mz_ore::Overflowing<i128>;
     PartialOrd,
     Ord,
     Serialize,
-    Deserialize
+    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.
     Bool {
@@ -2052,7 +2094,7 @@ enum Accum {
     /// Accumulates arbitrary precision decimals.
     Numeric {
         /// Accumulates non-special values
-        accum: OrderedDecimal<NumericAgg>,
+        accum: OrderedNumericAgg,
         /// Counts +inf
         pos_infs: Diff,
         /// Counts -inf
@@ -2254,7 +2296,7 @@ impl Multiply<Diff> 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 +2307,8 @@ impl Multiply<Diff> 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<Self>;
 }
@@ -2722,4 +2766,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<Accum> {
+        let mut cx = numeric::cx_datum();
+        let mut numeric = |s: &str| Datum::from(cx.parse(s).unwrap());
+        let cases: Vec<(AggregateFunc, Vec<Datum>)> = 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::<Accum>::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<i128>` test in
+        // `mz_ore`, so this only decodes.
+        let mut words = Vec::new();
+        encode(&mut words, &borrowed);
+        let decoded = BorrowedOf::<Accum>::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<Accum>, Diff)` diffs.
+        let diffs: Vec<(Vec<Accum>, Diff)> = accums
+            .chunks(3)
+            .map(|chunk| (chunk.to_vec(), Diff::ONE))
+            .collect();
+        let mut coltainer = Coltainer::<(Vec<Accum>, 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!(
+                <Coltainer<(Vec<Accum>, Diff)>>::into_owned(coltainer.index(index)),
+                *diff
+            );
+        }
+        let mut sum = <Coltainer<(Vec<Accum>, Diff)>>::into_owned(coltainer.index(0));
+        sum.plus_equals(&sum.clone().multiply(&Diff::from(-1_i64)));
+        assert!(sum.is_zero());
+    }
 }

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<T, R> = TraceAgent<RowRowSpine<T, R>>;
 pub type RowRowArrangement<'scope, T> = Arranged<'scope, RowRowAgent<T, Diff>>;
 pub type RowRowEnter<T, R, TEnter> = TraceEnter<TraceFrontier<RowRowAgent<T, R>>, TEnter>;
 // Row specialized spines and agents.
-pub type RowAgent<T, R> = TraceAgent<RowSpine<T, R>>;
+pub type RowAgent<T, R, DC = ColumnationStack<R>> = TraceAgent<RowSpine<T, R, DC>>;
 pub type RowArrangement<'scope, T> = Arranged<'scope, RowAgent<T, Diff>>;
 pub type RowEnter<T, R, TEnter> = TraceEnter<TraceFrontier<RowAgent<T, R>>, TEnter>;