mz_timely_util/

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

//! Common operator transformations on timely streams and differential collections.

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

use columnation::Columnation;
use differential_dataflow::consolidation::ConsolidatingContainerBuilder;
use differential_dataflow::difference::{Multiply, Semigroup};
use differential_dataflow::lattice::Lattice;
use differential_dataflow::trace::Batcher;
use differential_dataflow::{AsCollection, Collection, Hashable, VecCollection};
use timely::container::{DrainContainer, PushInto};
use timely::dataflow::channels::pact::{Exchange, ParallelizationContract, Pipeline};
use timely::dataflow::operators::Capability;
use timely::dataflow::operators::generic::builder_rc::{
    OperatorBuilder as OperatorBuilderRc, OperatorBuilder,
};
use timely::dataflow::operators::generic::operator::{self, Operator};
use timely::dataflow::operators::generic::{
    InputHandleCore, OperatorInfo, OutputBuilder, OutputBuilderSession,
};
use timely::dataflow::{Scope, Stream, StreamVec};
use timely::progress::operate::FrontierInterest;
use timely::progress::{Antichain, Timestamp};
use timely::{Container, ContainerBuilder, PartialOrder};

use crate::columnation::{ColumnationChunker, ColumnationStack};

/// Extension methods for timely [`Stream`]s.
pub trait StreamExt<'scope, T, C1>
where
    T: Timestamp,
    C1: Container + DrainContainer + Clone + 'static,
{
    /// Like `timely::dataflow::operators::generic::operator::Operator::unary`,
    /// but the logic function can handle failures.
    ///
    /// Creates a new dataflow operator that partitions its input stream by a
    /// parallelization strategy `pact` and repeatedly invokes `logic`, the
    /// function returned by the function passed as `constructor`. The `logic`
    /// function can read to the input stream and write to either of two output
    /// streams, where the first output stream represents successful
    /// computations and the second output stream represents failed
    /// computations.
    fn unary_fallible<DCB, ECB, B, P>(
        self,
        pact: P,
        name: &str,
        constructor: B,
    ) -> (
        Stream<'scope, T, DCB::Container>,
        Stream<'scope, T, ECB::Container>,
    )
    where
        DCB: ContainerBuilder,
        ECB: ContainerBuilder,
        B: FnOnce(
            Capability<T>,
            OperatorInfo,
        ) -> Box<
            dyn FnMut(
                    &mut InputHandleCore<T, C1, P::Puller>,
                    &mut OutputBuilderSession<'_, T, DCB>,
                    &mut OutputBuilderSession<'_, T, ECB>,
                ) + 'static,
        >,
        P: ParallelizationContract<T, C1>;

    /// Like [`timely::dataflow::operators::vec::Map::flat_map`], but `logic`
    /// is allowed to fail. The first returned stream will contain the
    /// successful applications of `logic`, while the second returned stream
    /// will contain the failed applications.
    fn flat_map_fallible<DCB, ECB, D2, E, I, L>(
        self,
        name: &str,
        logic: L,
    ) -> (
        Stream<'scope, T, DCB::Container>,
        Stream<'scope, T, ECB::Container>,
    )
    where
        DCB: ContainerBuilder + PushInto<D2>,
        ECB: ContainerBuilder + PushInto<E>,
        I: IntoIterator<Item = Result<D2, E>>,
        L: for<'a> FnMut(C1::Item<'a>) -> I + 'static;

    /// Routes each record to one of two output streams by a per-record predicate.
    ///
    /// Records for which `predicate` returns `true` go to the first output and all
    /// others to the second. Both outputs are sent under the input capability. That
    /// capability is only a lower bound on the times of the records it carries, so a
    /// split that depends on a record's time must inspect the record, and cannot be
    /// decided once per container.
    fn partition_by<CB, L>(
        self,
        name: &str,
        predicate: L,
    ) -> (
        Stream<'scope, T, CB::Container>,
        Stream<'scope, T, CB::Container>,
    )
    where
        CB: ContainerBuilder + for<'a> PushInto<C1::Item<'a>>,
        L: for<'a> FnMut(&C1::Item<'a>) -> bool + 'static;

    /// Block progress of the frontier at `expiration` time
    fn expire_stream_at(self, name: &str, expiration: T) -> Stream<'scope, T, C1>;
}

/// Extension methods for differential [`Collection`]s.
pub trait CollectionExt<'scope, T, D1, R>: Sized
where
    T: Timestamp,
    R: Semigroup,
{
    /// Creates a new empty collection in `scope`.
    fn empty(scope: Scope<'scope, T>) -> VecCollection<'scope, T, D1, R>;

    /// Like [`Collection::map`], but `logic` is allowed to fail. The first
    /// returned collection will contain successful applications of `logic`,
    /// while the second returned collection will contain the failed
    /// applications.
    ///
    /// Callers need to specify the following type parameters:
    /// * `DCB`: The container builder for the `Ok` output.
    /// * `ECB`: The container builder for the `Err` output.
    fn map_fallible<DCB, ECB, D2, E, L>(
        self,
        name: &str,
        mut logic: L,
    ) -> (
        VecCollection<'scope, T, D2, R>,
        VecCollection<'scope, T, E, R>,
    )
    where
        DCB: ContainerBuilder<Container = Vec<(D2, T, R)>> + PushInto<(D2, T, R)>,
        ECB: ContainerBuilder<Container = Vec<(E, T, R)>> + PushInto<(E, T, R)>,
        D2: Clone + 'static,
        E: Clone + 'static,
        L: FnMut(D1) -> Result<D2, E> + 'static,
    {
        self.flat_map_fallible::<DCB, ECB, _, _, _, _>(name, move |record| Some(logic(record)))
    }

    /// Like [`Collection::flat_map`], but `logic` is allowed to fail. The first
    /// returned collection will contain the successful applications of `logic`,
    /// while the second returned collection will contain the failed
    /// applications.
    fn flat_map_fallible<DCB, ECB, D2, E, I, L>(
        self,
        name: &str,
        logic: L,
    ) -> (
        Collection<'scope, T, DCB::Container>,
        Collection<'scope, T, ECB::Container>,
    )
    where
        DCB: ContainerBuilder + PushInto<(D2, T, R)>,
        ECB: ContainerBuilder + PushInto<(E, T, R)>,
        D2: Clone + 'static,
        E: Clone + 'static,
        I: IntoIterator<Item = Result<D2, E>>,
        L: FnMut(D1) -> I + 'static;

    /// Block progress of the frontier at `expiration` time.
    fn expire_collection_at(self, name: &str, expiration: T) -> VecCollection<'scope, T, D1, R>;

    /// Replaces each record with another, with a new difference type.
    ///
    /// This method is most commonly used to take records containing aggregatable data (e.g. numbers to be summed)
    /// and move the data into the difference component. This will allow differential dataflow to update in-place.
    fn explode_one<D2, R2, L>(
        self,
        logic: L,
    ) -> VecCollection<'scope, T, D2, <R2 as Multiply<R>>::Output>
    where
        D2: differential_dataflow::Data,
        R2: Semigroup + Multiply<R>,
        <R2 as Multiply<R>>::Output: Clone + 'static + Semigroup,
        L: FnMut(D1) -> (D2, R2) + 'static,
        T: Lattice;

    /// Partitions the input into a monotonic collection and
    /// non-monotone exceptions, with respect to differences.
    ///
    /// The exceptions are transformed by `into_err`.
    fn ensure_monotonic<E, IE>(
        self,
        into_err: IE,
    ) -> (
        VecCollection<'scope, T, D1, R>,
        VecCollection<'scope, T, E, R>,
    )
    where
        E: Clone + 'static,
        IE: Fn(D1, R) -> (E, R) + 'static,
        R: num_traits::sign::Signed;

    /// Consolidates the collection if `must_consolidate` is `true` and leaves it
    /// untouched otherwise.
    fn consolidate_named_if<Ba>(self, must_consolidate: bool, name: &str) -> Self
    where
        D1: differential_dataflow::ExchangeData + Hash + Columnation,
        R: Semigroup + differential_dataflow::ExchangeData + Columnation,
        T: Lattice + Columnation,
        Ba: Batcher<Time = T, Output = ColumnationStack<((D1, ()), T, R)>> + 'static;

    /// Consolidates the collection.
    fn consolidate_named<Ba>(self, name: &str) -> Self
    where
        D1: differential_dataflow::ExchangeData + Hash + Columnation,
        R: Semigroup + differential_dataflow::ExchangeData + Columnation,
        T: Lattice + Columnation,
        Ba: Batcher<Time = T, Output = ColumnationStack<((D1, ()), T, R)>> + 'static;
}

impl<'scope, T, C1> StreamExt<'scope, T, C1> for Stream<'scope, T, C1>
where
    T: Timestamp,
    C1: Container + DrainContainer + Clone + 'static,
{
    fn unary_fallible<DCB, ECB, B, P>(
        self,
        pact: P,
        name: &str,
        constructor: B,
    ) -> (
        Stream<'scope, T, DCB::Container>,
        Stream<'scope, T, ECB::Container>,
    )
    where
        DCB: ContainerBuilder,
        ECB: ContainerBuilder,
        B: FnOnce(
            Capability<T>,
            OperatorInfo,
        ) -> Box<
            dyn FnMut(
                    &mut InputHandleCore<T, C1, P::Puller>,
                    &mut OutputBuilderSession<'_, T, DCB>,
                    &mut OutputBuilderSession<'_, T, ECB>,
                ) + 'static,
        >,
        P: ParallelizationContract<T, C1>,
    {
        let mut builder = OperatorBuilderRc::new(name.into(), self.scope());

        let operator_info = builder.operator_info();

        let mut input = builder.new_input(self.clone(), pact);
        builder.set_notify_for(0, FrontierInterest::Never);
        let (ok_output, ok_stream) = builder.new_output();
        let mut ok_output = OutputBuilder::from(ok_output);
        let (err_output, err_stream) = builder.new_output();
        let mut err_output = OutputBuilder::from(err_output);

        builder.build(move |mut capabilities| {
            // `capabilities` should be a single-element vector.
            let capability = capabilities.pop().unwrap();
            let mut logic = constructor(capability, operator_info);
            move |_frontiers| {
                let mut ok_output_handle = ok_output.activate();
                let mut err_output_handle = err_output.activate();
                logic(&mut input, &mut ok_output_handle, &mut err_output_handle);
            }
        });

        (ok_stream, err_stream)
    }

    // XXX(guswynn): file an minimization bug report for the logic flat_map
    // false positive here
    // TODO(guswynn): remove this after https://github.com/rust-lang/rust-clippy/issues/8098 is
    // resolved. The `logic` `FnMut` needs to be borrowed in the `flat_map` call, not moved in
    // so the simple `|d1| logic(d1)` closure is load-bearing
    #[allow(clippy::redundant_closure)]
    fn flat_map_fallible<DCB, ECB, D2, E, I, L>(
        self,
        name: &str,
        mut logic: L,
    ) -> (
        Stream<'scope, T, DCB::Container>,
        Stream<'scope, T, ECB::Container>,
    )
    where
        DCB: ContainerBuilder + PushInto<D2>,
        ECB: ContainerBuilder + PushInto<E>,
        I: IntoIterator<Item = Result<D2, E>>,
        L: for<'a> FnMut(C1::Item<'a>) -> I + 'static,
    {
        self.unary_fallible::<DCB, ECB, _, _>(Pipeline, name, move |_, _| {
            Box::new(move |input, ok_output, err_output| {
                input.for_each_time(|time, data| {
                    let mut ok_session = ok_output.session_with_builder(&time);
                    let mut err_session = err_output.session_with_builder(&time);
                    for r in data
                        .flat_map(DrainContainer::drain)
                        .flat_map(|d1| logic(d1))
                    {
                        match r {
                            Ok(d2) => ok_session.give(d2),
                            Err(e) => err_session.give(e),
                        }
                    }
                })
            })
        })
    }

    fn partition_by<CB, L>(
        self,
        name: &str,
        mut predicate: L,
    ) -> (
        Stream<'scope, T, CB::Container>,
        Stream<'scope, T, CB::Container>,
    )
    where
        CB: ContainerBuilder + for<'a> PushInto<C1::Item<'a>>,
        L: for<'a> FnMut(&C1::Item<'a>) -> bool + 'static,
    {
        self.unary_fallible::<CB, CB, _, _>(Pipeline, name, move |_, _| {
            Box::new(move |input, matching_output, rest_output| {
                input.for_each_time(|time, data| {
                    let mut matching = matching_output.session_with_builder(&time);
                    let mut rest = rest_output.session_with_builder(&time);
                    for item in data.flat_map(DrainContainer::drain) {
                        if predicate(&item) {
                            matching.give(item);
                        } else {
                            rest.give(item);
                        }
                    }
                })
            })
        })
    }

    fn expire_stream_at(self, name: &str, expiration: T) -> Stream<'scope, T, C1> {
        let name = format!("expire_stream_at({name})");
        self.unary_frontier(Pipeline, &name.clone(), move |cap, _| {
            // Retain a capability for the expiration time, which we'll only drop if the token
            // is dropped. Else, block progress at the expiration time to prevent downstream
            // operators from making any statement about expiration time or any following time.
            let cap = Some(cap.delayed(&expiration));
            let mut warned = false;
            move |(input, frontier), output| {
                let _ = &cap;
                let frontier = frontier.frontier();
                if !frontier.less_than(&expiration) && !warned {
                    // Here, we print a warning, not an error. The state is only a liveness
                    // concern, but not relevant for correctness. Additionally, a race between
                    // shutting down the dataflow and dropping the token can cause the dataflow
                    // to shut down before we drop the token.  This can happen when dropping
                    // the last remaining capability on a different worker.  We do not want to
                    // log an error every time this happens.

                    tracing::warn!(
                        name = name,
                        frontier = ?frontier,
                        expiration = ?expiration,
                        "frontier not less than expiration"
                    );
                    warned = true;
                }
                input.for_each(|time, data| {
                    let mut session = output.session(&time);
                    session.give_container(data);
                });
            }
        })
    }
}

impl<'scope, T, D1, R> CollectionExt<'scope, T, D1, R> for VecCollection<'scope, T, D1, R>
where
    T: Timestamp + Clone + 'static,
    D1: Clone + 'static,
    R: Semigroup + 'static,
{
    fn empty(scope: Scope<'scope, T>) -> VecCollection<'scope, T, D1, R> {
        operator::empty(scope).as_collection()
    }

    fn flat_map_fallible<DCB, ECB, D2, E, I, L>(
        self,
        name: &str,
        mut logic: L,
    ) -> (
        Collection<'scope, T, DCB::Container>,
        Collection<'scope, T, ECB::Container>,
    )
    where
        DCB: ContainerBuilder + PushInto<(D2, T, R)>,
        ECB: ContainerBuilder + PushInto<(E, T, R)>,
        D2: Clone + 'static,
        E: Clone + 'static,
        I: IntoIterator<Item = Result<D2, E>>,
        L: FnMut(D1) -> I + 'static,
    {
        let (ok_stream, err_stream) =
            self.inner
                .flat_map_fallible::<DCB, ECB, _, _, _, _>(name, move |(d1, t, r)| {
                    logic(d1).into_iter().map(move |res| match res {
                        Ok(d2) => Ok((d2, t.clone(), r.clone())),
                        Err(e) => Err((e, t.clone(), r.clone())),
                    })
                });
        (ok_stream.as_collection(), err_stream.as_collection())
    }

    fn expire_collection_at(self, name: &str, expiration: T) -> VecCollection<'scope, T, D1, R> {
        self.inner
            .expire_stream_at(name, expiration)
            .as_collection()
    }

    fn explode_one<D2, R2, L>(
        self,
        mut logic: L,
    ) -> VecCollection<'scope, T, D2, <R2 as Multiply<R>>::Output>
    where
        D2: differential_dataflow::Data,
        R2: Semigroup + Multiply<R>,
        <R2 as Multiply<R>>::Output: Clone + 'static + Semigroup,
        L: FnMut(D1) -> (D2, R2) + 'static,
        T: Lattice,
    {
        self.inner
            .clone()
            .unary::<ConsolidatingContainerBuilder<_>, _, _, _>(
                Pipeline,
                "ExplodeOne",
                move |_, _| {
                    move |input, output| {
                        input.for_each(|time, data| {
                            output
                                .session_with_builder(&time)
                                .give_iterator(data.drain(..).map(|(x, t, d)| {
                                    let (x, d2) = logic(x);
                                    (x, t, d2.multiply(&d))
                                }));
                        });
                    }
                },
            )
            .as_collection()
    }

    fn ensure_monotonic<E, IE>(
        self,
        into_err: IE,
    ) -> (
        VecCollection<'scope, T, D1, R>,
        VecCollection<'scope, T, E, R>,
    )
    where
        E: Clone + 'static,
        IE: Fn(D1, R) -> (E, R) + 'static,
        R: num_traits::sign::Signed,
    {
        let (oks, errs) = self
            .inner
            .unary_fallible(Pipeline, "EnsureMonotonic", move |_, _| {
                Box::new(move |input, ok_output, err_output| {
                    input.for_each(|time, data| {
                        let mut ok_session = ok_output.session(&time);
                        let mut err_session = err_output.session(&time);
                        for (x, t, d) in data.drain(..) {
                            if d.is_positive() {
                                ok_session.give((x, t, d))
                            } else {
                                let (e, d2) = into_err(x, d);
                                err_session.give((e, t, d2))
                            }
                        }
                    })
                })
            });
        (oks.as_collection(), errs.as_collection())
    }

    fn consolidate_named_if<Ba>(self, must_consolidate: bool, name: &str) -> Self
    where
        D1: differential_dataflow::ExchangeData + Hash + Columnation,
        R: Semigroup + differential_dataflow::ExchangeData + Columnation,
        T: Lattice + Ord + Columnation,
        Ba: Batcher<Time = T, Output = ColumnationStack<((D1, ()), T, R)>> + 'static,
    {
        if must_consolidate {
            // We employ AHash below instead of the default hasher in DD to obtain
            // a better distribution of data to workers. AHash claims empirically
            // both speed and high quality, according to
            // https://github.com/tkaitchuck/aHash/blob/master/compare/readme.md.
            // TODO(vmarcos): Consider here if it is worth it to spend the time to
            // implement twisted tabulation hashing as proposed in Mihai Patrascu,
            // Mikkel Thorup: Twisted Tabulation Hashing. SODA 2013: 209-228, available
            // at https://epubs.siam.org/doi/epdf/10.1137/1.9781611973105.16. The latter
            // would provide good bounds for balls-into-bins problems when the number of
            // bins is small (as is our case), so we'd have a theoretical guarantee.
            // The seeds are fixed for determinism across builds; see
            // [`crate::hash::fixed_state`].
            let random_state = crate::hash::fixed_state();
            let exchange = Exchange::new(move |update: &((D1, _), T, R)| {
                let data = &(update.0).0;
                let mut h = random_state.build_hasher();
                data.hash(&mut h);
                h.finish()
            });
            consolidate_pact::<ColumnationChunker<((D1, ()), T, R)>, Ba, _, _>(
                self.map(|k| (k, ())).inner,
                exchange,
                name,
            )
            .unary(Pipeline, "unpack consolidated", |_, _| {
                |input, output| {
                    input.for_each(|time, data| {
                        let mut session = output.session(&time);
                        for ((k, ()), t, d) in data.iter().flatten().flat_map(|chunk| chunk.iter())
                        {
                            session.give((k.clone(), t.clone(), d.clone()))
                        }
                    })
                }
            })
            .as_collection()
        } else {
            self
        }
    }

    fn consolidate_named<Ba>(self, name: &str) -> Self
    where
        D1: differential_dataflow::ExchangeData + Hash + Columnation,
        R: Semigroup + differential_dataflow::ExchangeData + Columnation,
        T: Lattice + Ord + Columnation,
        Ba: Batcher<Time = T, Output = ColumnationStack<((D1, ()), T, R)>> + 'static,
    {
        let exchange = Exchange::new(move |update: &((D1, ()), T, R)| (update.0).0.hashed());

        consolidate_pact::<ColumnationChunker<((D1, ()), T, R)>, Ba, _, _>(
            self.map(|k| (k, ())).inner,
            exchange,
            name,
        )
        .unary(Pipeline, &format!("Unpack {name}"), |_, _| {
            |input, output| {
                input.for_each(|time, data| {
                    let mut session = output.session(&time);
                    for ((k, ()), t, d) in data.iter().flatten().flat_map(|chunk| chunk.iter()) {
                        session.give((k.clone(), t.clone(), d.clone()))
                    }
                })
            }
        })
        .as_collection()
    }
}

/// Aggregates the weights of equal records into at most one record.
///
/// Produces a stream of chains of records, partitioned according to `pact`. The
/// data is sorted according to `Ba`. For each timestamp, it produces at most one chain.
///
/// The data are accumulated in place, each held back until their timestamp has completed.
pub fn consolidate_pact<'scope, Chu, Ba, C, P>(
    stream: Stream<'scope, Ba::Time, C>,
    pact: P,
    name: &str,
) -> StreamVec<'scope, Ba::Time, Vec<Ba::Output>>
where
    Ba: Batcher + 'static,
    Chu: ContainerBuilder<Container = Ba::Output> + for<'a> PushInto<&'a mut C> + 'static,
    C: Container + Clone + 'static,
    Ba::Output: Clone,
    P: ParallelizationContract<Ba::Time, C>,
{
    let logger = stream
        .scope()
        .worker()
        .logger_for("differential/arrange")
        .map(Into::into);
    stream.unary_frontier(pact, name, |_cap, info| {
        // Acquire a logger for arrange events.

        let mut batcher = Ba::new(logger, info.global_id);
        // The chunker consolidates raw input containers into the chunks the
        // batcher consumes.
        let mut chunker = Chu::default();
        // Capabilities for the lower envelope of updates in `batcher`.
        let mut capabilities = Antichain::<Capability<Ba::Time>>::new();
        let mut prev_frontier = Antichain::from_elem(Ba::Time::minimum());

        move |(input, frontier), output| {
            input.for_each(|cap, data| {
                capabilities.insert(cap.retain(0));
                chunker.push_into(data);
                while let Some(chunk) = chunker.extract() {
                    batcher.push_into(std::mem::take(chunk));
                }
            });

            if prev_frontier.borrow() != frontier.frontier() {
                // Flush any data the chunker is still accumulating into the
                // batcher before we seal.
                while let Some(chunk) = chunker.finish() {
                    batcher.push_into(std::mem::take(chunk));
                }

                if capabilities
                    .elements()
                    .iter()
                    .any(|c| !frontier.less_equal(c.time()))
                {
                    let mut upper = Antichain::new(); // re-used allocation for sealing batches.

                    // For each capability not in advance of the input frontier ...
                    for (index, capability) in capabilities.elements().iter().enumerate() {
                        if !frontier.less_equal(capability.time()) {
                            // Assemble the upper bound on times we can commit with this capabilities.
                            // We must respect the input frontier, and *subsequent* capabilities, as
                            // we are pretending to retire the capability changes one by one.
                            upper.clear();
                            for time in frontier.frontier().iter() {
                                upper.insert(time.clone());
                            }
                            for other_capability in &capabilities.elements()[(index + 1)..] {
                                upper.insert(other_capability.time().clone());
                            }

                            // send the batch to downstream consumers, empty or not.
                            let mut session = output.session(&capabilities.elements()[index]);
                            // Extract updates not in advance of `upper`.
                            let (chain, _description) = batcher.seal(upper.clone());
                            session.give(chain);
                        }
                    }

                    // Having extracted and sent batches between each capability and the input frontier,
                    // we should downgrade all capabilities to match the batcher's lower update frontier.
                    // This may involve discarding capabilities, which is fine as any new updates arrive
                    // in messages with new capabilities.

                    let mut new_capabilities = Antichain::new();
                    for time in batcher.frontier().iter() {
                        if let Some(capability) = capabilities
                            .elements()
                            .iter()
                            .find(|c| c.time().less_equal(time))
                        {
                            new_capabilities.insert(capability.delayed(time));
                        } else {
                            panic!("failed to find capability");
                        }
                    }

                    capabilities = new_capabilities;
                }

                prev_frontier.clear();
                prev_frontier.extend(frontier.frontier().iter().cloned());
            }
        }
    })
}

/// Merge the contents of multiple streams and combine the containers using a container builder.
pub trait ConcatenateFlatten<'scope, T: Timestamp, C: Container + DrainContainer> {
    /// Merge the contents of multiple streams and use the provided container builder to form
    /// output containers.
    ///
    /// # Examples
    /// ```
    /// use timely::container::CapacityContainerBuilder;
    /// use timely::dataflow::operators::{ToStream, Inspect};
    /// use mz_timely_util::operator::ConcatenateFlatten;
    ///
    /// timely::example(|scope| {
    ///
    ///     let streams: Vec<timely::dataflow::StreamVec<_, i32>> =
    ///         vec![(0..10).to_stream(scope),
    ///              (0..10).to_stream(scope),
    ///              (0..10).to_stream(scope)];
    ///
    ///     scope.concatenate_flatten::<_, CapacityContainerBuilder<Vec<i32>>>(streams)
    ///          .inspect(|x| println!("seen: {:?}", x));
    /// });
    /// ```
    fn concatenate_flatten<I, CB>(&self, sources: I) -> Stream<'scope, T, CB::Container>
    where
        I: IntoIterator<Item = Stream<'scope, T, C>>,
        CB: ContainerBuilder + for<'a> PushInto<C::Item<'a>>;
}

impl<'scope, T, C> ConcatenateFlatten<'scope, T, C> for Stream<'scope, T, C>
where
    T: Timestamp,
    C: Container + DrainContainer + Clone + 'static,
{
    fn concatenate_flatten<I, CB>(&self, sources: I) -> Stream<'scope, T, CB::Container>
    where
        I: IntoIterator<Item = Stream<'scope, T, C>>,
        CB: ContainerBuilder + for<'a> PushInto<C::Item<'a>>,
    {
        self.scope()
            .concatenate_flatten::<_, CB>(Some(Clone::clone(self)).into_iter().chain(sources))
    }
}

impl<'scope, T, C> ConcatenateFlatten<'scope, T, C> for Scope<'scope, T>
where
    T: Timestamp,
    C: Container + DrainContainer,
{
    fn concatenate_flatten<I, CB>(&self, sources: I) -> Stream<'scope, T, CB::Container>
    where
        I: IntoIterator<Item = Stream<'scope, T, C>>,
        CB: ContainerBuilder + for<'a> PushInto<C::Item<'a>>,
    {
        let mut builder = OperatorBuilder::new("ConcatenateFlatten".to_string(), self.clone());

        // create new input handles for each input stream.
        let mut handles = sources
            .into_iter()
            .map(|s| builder.new_input(s, Pipeline))
            .collect::<Vec<_>>();
        for i in 0..handles.len() {
            builder.set_notify_for(i, FrontierInterest::Never);
        }

        // create one output handle for the concatenated results.
        let (output, result) = builder.new_output::<CB::Container>();
        let mut output = OutputBuilder::<_, CB>::from(output);

        builder.build(move |_capability| {
            move |_frontier| {
                let mut output = output.activate();
                for handle in handles.iter_mut() {
                    handle.for_each_time(|time, data| {
                        output
                            .session_with_builder(&time)
                            .give_iterator(data.flat_map(DrainContainer::drain));
                    })
                }
            }
        });

        result
    }
}

/// A trait for containers that can be cleared.
pub trait ClearContainer {
    /// Clear the contents of the container.
    fn clear(&mut self);
}

impl<T> ClearContainer for Vec<T> {
    fn clear(&mut self) {
        Vec::clear(self)
    }
}

#[cfg(test)]
mod tests {
    use timely::container::CapacityContainerBuilder;
    use timely::dataflow::operators::Capture;
    use timely::dataflow::operators::capture::Extract;
    use timely::dataflow::operators::core::to_stream::ToStreamBuilder;
    use timely::dataflow::operators::vec::ToStream;

    use crate::columnar::Column;
    use crate::columnar::builder::ColumnBuilder;

    use super::*;

    /// Updates whose own times straddle `SPLIT`, all carried under the single input
    /// capability at time zero. A split decided per container would route them together.
    const UPDATES: [(u64, u64, i64); 4] = [(0, 0, 1), (1, 3, 1), (2, 5, -1), (3, 7, 1)];
    const SPLIT: u64 = 5;
    const MATCHING: [(u64, u64, i64); 2] = [(2, 5, -1), (3, 7, 1)];
    const REST: [(u64, u64, i64); 2] = [(0, 0, 1), (1, 3, 1)];

    #[mz_ore::test]
    fn partition_by_routes_vec_records_by_their_own_time() {
        let (matching, rest) = timely::execute_directly(|worker| {
            worker.dataflow::<u64, _, _>(|scope| {
                let (matching, rest) = UPDATES
                    .to_vec()
                    .to_stream(scope)
                    .partition_by::<CapacityContainerBuilder<Vec<_>>, _>("Test", |(_, time, _)| {
                        *time >= SPLIT
                    });
                (matching.capture(), rest.capture())
            })
        });
        let flatten = |captured: Vec<(u64, Vec<(u64, u64, i64)>)>| {
            captured
                .into_iter()
                .flat_map(|(capability, updates)| {
                    assert_eq!(capability, 0, "outputs keep the input capability");
                    updates
                })
                .collect::<Vec<_>>()
        };
        assert_eq!(flatten(matching.extract()), MATCHING);
        assert_eq!(flatten(rest.extract()), REST);
    }

    #[mz_ore::test]
    fn partition_by_routes_columnar_records_by_their_own_time() {
        let (matching, rest) = timely::execute_directly(|worker| {
            worker.dataflow::<u64, _, _>(|scope| {
                let (matching, rest) = UPDATES
                    .to_vec()
                    .to_stream_with_builder::<_, ColumnBuilder<(u64, u64, i64)>>(scope)
                    .partition_by::<ColumnBuilder<(u64, u64, i64)>, _>("Test", |(_, time, _)| {
                        **time >= SPLIT
                    });
                (matching.capture(), rest.capture())
            })
        });
        let flatten = |captured: Vec<(u64, Column<(u64, u64, i64)>)>| {
            captured
                .into_iter()
                .flat_map(|(capability, mut column)| {
                    assert_eq!(capability, 0, "outputs keep the input capability");
                    column
                        .drain()
                        .map(|(data, time, diff)| (*data, *time, *diff))
                        .collect::<Vec<_>>()
                })
                .collect::<Vec<_>>()
        };
        assert_eq!(flatten(matching.extract()), MATCHING);
        assert_eq!(flatten(rest.extract()), REST);
    }
}