Switch some dataflow fragments to columnar

Change the dataflow fragment for `specialized_arrange` and in the linear
join preparation phase to use columnar data on edges.

Depends on:
* TimelyDataflow/differential-dataflow#562
* TimelyDataflow/timely-dataflow#630

Signed-off-by: Moritz Hoffmann <mh@materialize.com>

Author: antiguru

Cargo.lock

@@ -285,6 +285,8 @@ incremental = true
 # merged), after which point it becomes impossible to build that historical
 # version of Materialize.
 [patch.crates-io]
+differential-dataflow = { path = "../differential-dataflow" }
+timely = { path = "../timely-dataflow/timely" }
 # Waiting on https://github.com/sfackler/rust-postgres/pull/752.
 postgres = { git = "https://github.com/MaterializeInc/rust-postgres" }
 tokio-postgres = { git = "https://github.com/MaterializeInc/rust-postgres" }

Cargo.toml

@@ -41,7 +41,7 @@ pub const LINEAR_JOIN_YIELDING: Config<&str> = Config::new(
 /// Enable lgalloc for columnation.
 pub const ENABLE_COLUMNATION_LGALLOC: Config<bool> = Config::new(
     "enable_columnation_lgalloc",
-    false,
+    true,
     "Enable allocating regions from lgalloc.",
 );
 

src/compute-types/src/dyncfgs.rs

@@ -110,6 +110,7 @@ use std::rc::{Rc, Weak};
 use std::sync::Arc;
 use std::task::Poll;
 
+use columnar::Columnar;
 use differential_dataflow::dynamic::pointstamp::PointStamp;
 use differential_dataflow::lattice::Lattice;
 use differential_dataflow::operators::arrange::Arranged;
@@ -498,6 +499,7 @@ impl<'g, G, T> Context<Child<'g, G, T>>
 where
     G: Scope<Timestamp = mz_repr::Timestamp>,
     T: Refines<G::Timestamp> + RenderTimestamp,
+    <T as Columnar>::Container: Clone + Send,
 {
     pub(crate) fn import_index(
         &mut self,
@@ -647,6 +649,7 @@ impl<'g, G, T> Context<Child<'g, G, T>>
 where
     G: Scope<Timestamp = mz_repr::Timestamp>,
     T: RenderTimestamp,
+    <T as Columnar>::Container: Clone + Send,
 {
     pub(crate) fn export_index_iterative(
         &self,
@@ -889,6 +892,8 @@ impl<G> Context<G>
 where
     G: Scope,
     G::Timestamp: RenderTimestamp,
+    <G::Timestamp as Columnar>::Container: Clone + Send,
+    for<'a> <G::Timestamp as Columnar>::Ref<'a>: Ord + Copy,
 {
     /// Renders a non-recursive plan to a differential dataflow, producing the collection of
     /// results.
@@ -1293,7 +1298,11 @@ where
 
 #[allow(dead_code)] // Some of the methods on this trait are unused, but useful to have.
 /// A timestamp type that can be used for operations within MZ's dataflow layer.
-pub trait RenderTimestamp: Timestamp + Lattice + Refines<mz_repr::Timestamp> + Columnation {
+pub trait RenderTimestamp:
+    Timestamp + Lattice + Refines<mz_repr::Timestamp> + Columnation + Columnar
+where
+    <Self as Columnar>::Container: Clone + Send,
+{
     /// The system timestamp component of the timestamp.
     ///
     /// This is useful for manipulating the system time, as when delaying
@@ -1429,6 +1438,7 @@ impl<S> WithStartSignal for MzArrangementImport<S>
 where
     S: Scope,
     S::Timestamp: RenderTimestamp,
+    <S::Timestamp as Columnar>::Container: Clone + Send,
 {
     fn with_start_signal(self, signal: StartSignal) -> Self {
         match self {
@@ -1443,6 +1453,7 @@ impl<S, Tr> WithStartSignal for Arranged<S, Tr>
 where
     S: Scope,
     S::Timestamp: RenderTimestamp,
+    <S::Timestamp as Columnar>::Container: Clone + Send,
     Tr: TraceReader + Clone,
 {
     fn with_start_signal(self, signal: StartSignal) -> Self {

src/compute/src/render.rs

@@ -14,37 +14,40 @@ use std::collections::BTreeMap;
 use std::rc::Weak;
 use std::sync::mpsc;
 
+use columnar::Columnar;
 use differential_dataflow::consolidation::ConsolidatingContainerBuilder;
 use differential_dataflow::lattice::Lattice;
 use differential_dataflow::operators::arrange::Arranged;
 use differential_dataflow::trace::cursor::IntoOwned;
 use differential_dataflow::trace::{BatchReader, Cursor, TraceReader};
-use differential_dataflow::{Collection, Data};
+use differential_dataflow::{AsCollection, Collection, Data};
 use mz_compute_types::dataflows::DataflowDescription;
 use mz_compute_types::plan::{AvailableCollections, LirId};
 use mz_expr::{Id, MapFilterProject, MirScalarExpr};
-use mz_repr::fixed_length::{FromDatumIter, ToDatumIter};
+use mz_repr::fixed_length::ToDatumIter;
 use mz_repr::{DatumVec, DatumVecBorrow, Diff, GlobalId, Row, RowArena, SharedRow};
 use mz_storage_types::controller::CollectionMetadata;
 use mz_storage_types::errors::DataflowError;
+use mz_timely_util::containers::{columnar_exchange, Col2ValBatcher, ColumnBuilder};
 use mz_timely_util::operator::{CollectionExt, StreamExt};
 use timely::container::columnation::Columnation;
 use timely::container::CapacityContainerBuilder;
-use timely::dataflow::channels::pact::Pipeline;
+use timely::dataflow::channels::pact::{ExchangeCore, Pipeline};
 use timely::dataflow::operators::generic::OutputHandleCore;
 use timely::dataflow::operators::Capability;
 use timely::dataflow::scopes::Child;
 use timely::dataflow::{Scope, ScopeParent};
 use timely::progress::timestamp::Refines;
 use timely::progress::{Antichain, Timestamp};
+use timely::Container;
 use tracing::error;
 
 use crate::arrangement::manager::SpecializedTraceHandle;
 use crate::compute_state::{ComputeState, HydrationEvent};
-use crate::extensions::arrange::{KeyCollection, MzArrange};
+use crate::extensions::arrange::{KeyCollection, MzArrange, MzArrangeCore};
 use crate::render::errors::ErrorLogger;
 use crate::render::{LinearJoinSpec, RenderTimestamp};
-use crate::row_spine::{RowRowBatcher, RowRowBuilder};
+use crate::row_spine::RowRowBuilder;
 use crate::typedefs::{
     ErrAgent, ErrBatcher, ErrBuilder, ErrEnter, ErrSpine, RowRowAgent, RowRowEnter, RowRowSpine,
 };
@@ -919,10 +922,11 @@ where
 
 impl<S, T> CollectionBundle<S, T>
 where
-    T: timely::progress::Timestamp + Lattice + Columnation,
+    T: Timestamp + Lattice + Columnation,
     S: Scope,
-    S::Timestamp:
-        Refines<T> + Lattice + timely::progress::Timestamp + crate::render::RenderTimestamp,
+    S::Timestamp: Refines<T> + RenderTimestamp,
+    <S::Timestamp as Columnar>::Container: Clone + Send,
+    for<'a> <S::Timestamp as Columnar>::Ref<'a>: Ord + Copy,
 {
     /// Presents `self` as a stream of updates, having been subjected to `mfp`.
     ///
@@ -1046,7 +1050,8 @@ where
                     .collection
                     .clone()
                     .expect("Collection constructed above");
-                let (oks, errs_keyed) = Self::specialized_arrange(&name, oks, &key, &thinning);
+                let (oks, errs_keyed) =
+                    Self::specialized_arrange(&name, oks, key.clone(), thinning.clone());
                 let errs: KeyCollection<_, _, _> = errs.concat(&errs_keyed).into();
                 let errs = errs.mz_arrange::<ErrBatcher<_, _>, ErrBuilder<_, _>, ErrSpine<_, _>>(
                     &format!("{}-errors", name),
@@ -1064,44 +1069,48 @@ where
     fn specialized_arrange(
         name: &String,
         oks: Collection<S, Row, i64>,
-        key: &Vec<MirScalarExpr>,
-        thinning: &Vec<usize>,
+        key: Vec<MirScalarExpr>,
+        thinning: Vec<usize>,
     ) -> (MzArrangement<S>, Collection<S, DataflowError, i64>) {
-        // Catch-all: Just use RowRow.
-        let (oks, errs) = oks
-            .map_fallible::<CapacityContainerBuilder<_>, CapacityContainerBuilder<_>, _, _, _>(
-                "FormArrangementKey",
-                specialized_arrangement_key(key.clone(), thinning.clone()),
+        let (oks, errs) =
+            oks.inner
+                .unary_fallible::<ColumnBuilder<((Row, Row), S::Timestamp, i64)>, _, _, _>(
+                    Pipeline,
+                    "FormArrangementKey",
+                    move |_, _| {
+                        Box::new(move |input, ok, err| {
+                            let mut key_buf = Row::default();
+                            let mut val_buf = Row::default();
+                            let mut datums = DatumVec::new();
+                            let temp_storage = RowArena::new();
+                            while let Some((time, data)) = input.next() {
+                                let mut ok_session = ok.session_with_builder(&time);
+                                let mut err_session = err.session(&time);
+                                for (row, time, diff) in data.iter() {
+                                    temp_storage.clear();
+                                    let datums = datums.borrow_with(&row);
+                                    let val_datum_iter = thinning.iter().map(|c| datums[*c]);
+                                    match key_buf.packer().try_extend(
+                                        key.iter().map(|k| k.eval(&datums, &temp_storage)),
+                                    ) {
+                                        Ok(()) => {
+                                            val_buf.packer().extend(val_datum_iter);
+                                            ok_session.give(((&*key_buf, &*val_buf), time, diff));
+                                        }
+                                        Err(e) => {
+                                            err_session.give((e.into(), time.clone(), *diff));
+                                        }
+                                    }
+                                }
+                            }
+                        })
+                    },
+                );
+        let oks = oks
+            .mz_arrange_core::<_, Col2ValBatcher<_, _,_, _>, RowRowBuilder<_, _>, RowRowSpine<_, _>>(
+                ExchangeCore::<ColumnBuilder<_>, _>::new_core(columnar_exchange::<Row, Row, S::Timestamp, Diff>),name
             );
-        let oks =
-            oks.mz_arrange::<RowRowBatcher<_, _>, RowRowBuilder<_, _>, RowRowSpine<_, _>>(name);
-        (MzArrangement::RowRow(oks), errs)
-    }
-}
-
-/// Obtains a function that maps input rows to (key, value) pairs according to
-/// the given key and thinning expressions. This function allows for specialization
-/// of key and value types and is intended to use to form arrangement keys.
-fn specialized_arrangement_key<K, V>(
-    key: Vec<MirScalarExpr>,
-    thinning: Vec<usize>,
-) -> impl FnMut(Row) -> Result<(K, V), DataflowError>
-where
-    K: Columnation + Data + FromDatumIter,
-    V: Columnation + Data + FromDatumIter,
-{
-    let mut key_buf = K::default();
-    let mut val_buf = V::default();
-    let mut datums = DatumVec::new();
-    move |row| {
-        // TODO: Consider reusing the `row` allocation; probably in *next* invocation.
-        let datums = datums.borrow_with(&row);
-        let temp_storage = RowArena::new();
-        let val_datum_iter = thinning.iter().map(|c| datums[*c]);
-        Ok::<(K, V), DataflowError>((
-            key_buf.try_from_datum_iter(key.iter().map(|k| k.eval(&datums, &temp_storage)))?,
-            val_buf.from_datum_iter(val_datum_iter),
-        ))
+        (MzArrangement::RowRow(oks), errs.as_collection())
     }
 }
 

src/compute/src/render/context.rs

@@ -7,6 +7,7 @@
 // the Business Source License, use of this software will be governed
 // by the Apache License, Version 2.0.
 
+use columnar::Columnar;
 use differential_dataflow::consolidation::ConsolidatingContainerBuilder;
 use mz_expr::MfpPlan;
 use mz_expr::{MapFilterProject, MirScalarExpr, TableFunc};
@@ -26,6 +27,7 @@ impl<G> Context<G>
 where
     G: Scope,
     G::Timestamp: crate::render::RenderTimestamp,
+    <G::Timestamp as Columnar>::Container: Clone + Send,
 {
     /// Applies a `TableFunc` to every row, followed by an `mfp`.
     pub fn render_flat_map(
@@ -130,6 +132,7 @@ fn drain_through_mfp<T>(
     >,
 ) where
     T: crate::render::RenderTimestamp,
+    <T as Columnar>::Container: Clone + Send,
 {
     let temp_storage = RowArena::new();
     let binding = SharedRow::get();

src/compute/src/render/flat_map.rs

@@ -15,6 +15,7 @@
 
 use std::collections::{BTreeMap, BTreeSet};
 
+use columnar::Columnar;
 use differential_dataflow::consolidation::ConsolidatingContainerBuilder;
 use differential_dataflow::lattice::Lattice;
 use differential_dataflow::operators::arrange::Arranged;
@@ -45,7 +46,9 @@ use crate::typedefs::{RowRowAgent, RowRowEnter};
 impl<G> Context<G>
 where
     G: Scope,
-    G::Timestamp: crate::render::RenderTimestamp,
+    G::Timestamp: RenderTimestamp,
+    <G::Timestamp as Columnar>::Container: Clone + Send,
+    for<'a> <G::Timestamp as Columnar>::Ref<'a>: Ord + Copy,
 {
     /// Renders `MirRelationExpr:Join` using dogs^3 delta query dataflows.
     ///
@@ -326,7 +329,9 @@ fn dispatch_build_halfjoin_local<G, CF>(
 )
 where
     G: Scope,
-    G::Timestamp: crate::render::RenderTimestamp,
+    G::Timestamp: RenderTimestamp,
+    <G::Timestamp as Columnar>::Container: Clone + Send,
+    for<'a> <G::Timestamp as Columnar>::Ref<'a>: Ord + Copy,
     CF: Fn(&G::Timestamp, &G::Timestamp) -> bool + 'static,
 {
     match trace {
@@ -358,7 +363,9 @@ fn dispatch_build_halfjoin_trace<G, T, CF>(
 where
     G: Scope,
     T: Timestamp + Lattice + Columnation,
-    G::Timestamp: Lattice + crate::render::RenderTimestamp + Refines<T> + Columnation,
+    G::Timestamp: RenderTimestamp + Refines<T>,
+    <G::Timestamp as Columnar>::Container: Clone + Send,
+    for<'a> <G::Timestamp as Columnar>::Ref<'a>: Ord + Copy,
     CF: Fn(&G::Timestamp, &G::Timestamp) -> bool + 'static,
 {
     match trace {
@@ -398,7 +405,9 @@ fn build_halfjoin<G, Tr, K, CF>(
 )
 where
     G: Scope,
-    G::Timestamp: crate::render::RenderTimestamp,
+    G::Timestamp: RenderTimestamp,
+    <G::Timestamp as Columnar>::Container: Clone + Send,
+    for<'a> <G::Timestamp as Columnar>::Ref<'a>: Ord + Copy,
     Tr: TraceReader<Time = G::Timestamp, Diff = Diff> + Clone + 'static,
     K: ExchangeData + Hashable + Default + FromDatumIter + ToDatumIter,
     for<'a> Tr::Key<'a>: IntoOwned<'a, Owned = K>,
@@ -531,7 +540,9 @@ fn dispatch_build_update_stream_local<G>(
 ) -> (Collection<G, Row, Diff>, Collection<G, DataflowError, Diff>)
 where
     G: Scope,
-    G::Timestamp: crate::render::RenderTimestamp,
+    G::Timestamp: RenderTimestamp,
+    <G::Timestamp as Columnar>::Container: Clone + Send,
+    for<'a> <G::Timestamp as Columnar>::Ref<'a>: Ord + Copy,
 {
     match trace {
         MzArrangement::RowRow(inner) => build_update_stream::<_, RowRowAgent<_, _>>(
@@ -553,7 +564,9 @@ fn dispatch_build_update_stream_trace<G, T>(
 where
     G: Scope,
     T: Timestamp + Lattice + Columnation,
-    G::Timestamp: Lattice + crate::render::RenderTimestamp + Refines<T> + Columnation,
+    G::Timestamp: Lattice + RenderTimestamp + Refines<T>,
+    <G::Timestamp as Columnar>::Container: Clone + Send,
+    for<'a> <G::Timestamp as Columnar>::Ref<'a>: Ord + Copy,
 {
     match trace {
         MzArrangementImport::RowRow(inner) => build_update_stream::<_, RowRowEnter<_, _, _>>(
@@ -578,7 +591,9 @@ fn build_update_stream<G, Tr>(
 ) -> (Collection<G, Row, Diff>, Collection<G, DataflowError, Diff>)
 where
     G: Scope,
-    G::Timestamp: crate::render::RenderTimestamp,
+    G::Timestamp: RenderTimestamp,
+    <G::Timestamp as Columnar>::Container: Clone + Send,
+    for<'a> <G::Timestamp as Columnar>::Ref<'a>: Ord + Copy,
     for<'a, 'b> &'a G::Timestamp: PartialEq<Tr::TimeGat<'b>>,
     Tr: for<'a> TraceReader<Time = G::Timestamp, Diff = Diff> + Clone + 'static,
     for<'a> Tr::Key<'a>: ToDatumIter,