Distributor trait

Introduce a distributor that knows how to partition containers across
multiple pushers. This moves the partition logic from container builders
into a bespoke trait and implementation instead of mixing two different
concepts, building containers, and partitioning them.

It allows us to implement a distributor in the future that only partitions
by key, for example when the container doesn't have easy access by row.

Removes the old `partition` function from container builders.

Signed-off-by: Moritz Hoffmann <[email protected]>

Author: antiguru

container/src/lib.rs

@@ -113,19 +113,6 @@ pub trait ContainerBuilder: Default + 'static {
     /// be called repeatedly until it returns `None`.
     #[must_use]
     fn finish(&mut self) -> Option<&mut Self::Container>;
-    /// Partitions `container` among `builders`, using the function `index` to direct items.
-    fn partition<I>(container: &mut Self::Container, builders: &mut [Self], mut index: I)
-    where
-        Self: for<'a> PushInto<<Self::Container as Container>::Item<'a>>,
-        I: for<'a> FnMut(&<Self::Container as Container>::Item<'a>) -> usize,
-    {
-        for datum in container.drain() {
-            let index = index(&datum);
-            builders[index].push_into(datum);
-        }
-        container.clear();
-    }
-
     /// Indicates a good moment to release resources.
     ///
     /// By default, does nothing. Callers first needs to drain the contents using [`Self::finish`]

timely/src/dataflow/channels/pact.rs

@@ -14,6 +14,7 @@ use crate::{Container, container::{ContainerBuilder, LengthPreservingContainerBu
 use crate::communication::allocator::thread::{ThreadPusher, ThreadPuller};
 use crate::communication::{Push, Pull};
 use crate::dataflow::channels::pushers::Exchange as ExchangePusher;
+use crate::dataflow::channels::pushers::exchange::DrainContainerDistributor;
 use crate::dataflow::channels::Message;
 use crate::logging::{TimelyLogger as Logger, MessagesEvent};
 use crate::progress::Timestamp;
@@ -79,20 +80,25 @@ where
 }
 
 // Exchange uses a `Box<Pushable>` because it cannot know what type of pushable will return from the allocator.
-impl<T: Timestamp, CB, H: 'static> ParallelizationContract<T, CB::Container> for ExchangeCore<CB, H>
+impl<T: Timestamp, CB, H> ParallelizationContract<T, CB::Container> for ExchangeCore<CB, H>
 where
-    CB: ContainerBuilder,
-    CB: for<'a> PushInto<<CB::Container as Container>::Item<'a>>,
+    CB: ContainerBuilder + for<'a> PushInto<<CB::Container as Container>::Item<'a>>,
     CB::Container: Data + Send + crate::dataflow::channels::ContainerBytes,
-    for<'a> H: FnMut(&<CB::Container as Container>::Item<'a>) -> u64
+    for<'a> H: FnMut(&<CB::Container as Container>::Item<'a>) -> u64 + 'static,
 {
-    type Pusher = ExchangePusher<T, CB, LogPusher<Box<dyn Push<Message<T, CB::Container>>>>, H>;
+    type Pusher = ExchangePusher<
+        T,
+        CB::Container,
+        LogPusher<Box<dyn Push<Message<T, CB::Container>>>>,
+        DrainContainerDistributor<CB, T, LogPusher<Box<dyn Push<Message<T, CB::Container>>>>, H>
+    >;
     type Puller = LogPuller<Box<dyn Pull<Message<T, CB::Container>>>>;
 
     fn connect<A: AsWorker>(self, allocator: &mut A, identifier: usize, address: Rc<[usize]>, logging: Option<Logger>) -> (Self::Pusher, Self::Puller) {
         let (senders, receiver) = allocator.allocate::<Message<T, CB::Container>>(identifier, address);
         let senders = senders.into_iter().enumerate().map(|(i,x)| LogPusher::new(x, allocator.index(), i, identifier, logging.clone())).collect::<Vec<_>>();
-        (ExchangePusher::new(senders, self.hash_func), LogPuller::new(receiver, allocator.index(), identifier, logging.clone()))
+        let distributor = DrainContainerDistributor::new(self.hash_func);
+        (ExchangePusher::new(senders, distributor), LogPuller::new(receiver, allocator.index(), identifier, logging.clone()))
     }
 }
 

timely/src/dataflow/channels/pushers/exchange.rs

@@ -5,55 +5,122 @@ use crate::container::{ContainerBuilder, PushInto};
 use crate::dataflow::channels::Message;
 use crate::{Container, Data};
 
-// TODO : Software write combining
-/// Distributes records among target pushees according to a distribution function.
-pub struct Exchange<T, CB, P, H>
+/// Distribute containers to several pushers.
+///
+/// A distributor sits behind an exchange pusher, and partitions containers at a given time
+/// into several pushers. It can use [`Message::push_at`] to push its outputs at the desired
+/// pusher.
+///
+/// It needs to uphold progress tracking requirements. The count of the input container
+/// must be preserved across the output containers, from the first call to `partition` until the
+/// call to `flush` for a specific time stamp.
+pub trait Distributor<T, C, P> {
+    /// Partition the contents of `container` at `time` into the `pushers`.
+    fn partition(&mut self, container: &mut C, time: &T, pushers: &mut [P]);
+    /// Flush any remaining contents into the `pushers` at time `time`.
+    fn flush(&mut self, time: &T, pushers: &mut [P]);
+    /// Optionally release resources, such as memory.
+    fn relax(&mut self);
+}
+
+/// A distributor creating containers from a drainable container based
+/// on a hash function of the container's item.
+pub struct DrainContainerDistributor<CB, T, P, H> {
+    builders: Vec<CB>,
+    hash_func: H,
+    _phantom: std::marker::PhantomData<(T, P)>,
+}
+
+impl<CB: Default, T, P, H> DrainContainerDistributor<CB, T, P, H> {
+    /// Allocates a new `DrainContainerPartitioner` with the given pusher count and hash function.
+    pub fn new(hash_func: H) -> Self {
+        DrainContainerDistributor {
+            builders: Vec::new(),
+            hash_func,
+            _phantom: std::marker::PhantomData,
+        }
+    }
+}
+
+impl<CB, T, P, H> Distributor<T, CB::Container, P> for DrainContainerDistributor<CB, T, P, H>
 where
-    CB: ContainerBuilder,
+    CB: ContainerBuilder + for<'a> PushInto<<CB::Container as Container>::Item<'a>>,
+    CB::Container: Container,
+    T: Clone,
     P: Push<Message<T, CB::Container>>,
-    for<'a> H: FnMut(&<CB::Container as Container>::Item<'a>) -> u64
+    for<'a> H: FnMut(&<CB::Container as Container>::Item<'a>) -> u64,
 {
+    fn partition(&mut self, container: &mut CB::Container, time: &T, pushers: &mut [P]) {
+        self.builders.resize_with(pushers.len(), Default::default);
+        if pushers.len().is_power_of_two() {
+            let mask = (pushers.len() - 1) as u64;
+            for datum in container.drain() {
+                let index = ((self.hash_func)(&datum) & mask) as usize;
+                self.builders[index].push_into(datum);
+                while let Some(produced) = self.builders[index].extract() {
+                    Message::push_at(produced, time.clone(), &mut pushers[index]);
+                }
+            }
+        }
+        else {
+            let num_pushers = pushers.len() as u64;
+            for datum in container.drain() {
+                let index = ((self.hash_func)(&datum) % num_pushers) as usize;
+                self.builders[index].push_into(datum);
+                while let Some(produced) = self.builders[index].extract() {
+                    Message::push_at(produced, time.clone(), &mut pushers[index]);
+                }
+            }
+        }
+    }
+
+    fn flush(&mut self, time: &T, pushers: &mut [P]) {
+        for (builder, pusher) in self.builders.iter_mut().zip(pushers.iter_mut()) {
+            while let Some(container) = builder.finish() {
+                Message::push_at(container, time.clone(), pusher);
+            }
+        }
+    }
+
+    fn relax(&mut self) {
+        for builder in &mut self.builders {
+            builder.relax();
+        }
+    }
+}
+
+// TODO : Software write combining
+/// Distributes records among target pushees according to a distributor.
+pub struct Exchange<T, C, P, D> {
     pushers: Vec<P>,
-    builders: Vec<CB>,
     current: Option<T>,
-    hash_func: H,
+    distributor: D,
+    _phantom: std::marker::PhantomData<C>,
 }
 
-impl<T: Clone, CB, P, H>  Exchange<T, CB, P, H>
+impl<T: Clone, C, P, D>  Exchange<T, C, P, D>
 where
-    CB: ContainerBuilder,
-    P: Push<Message<T, CB::Container>>,
-    for<'a> H: FnMut(&<CB::Container as Container>::Item<'a>) -> u64
+    P: Push<Message<T, C>>,
+    D: Distributor<T, C, P>,
 {
-    /// Allocates a new `Exchange` from a supplied set of pushers and a distribution function.
-    pub fn new(pushers: Vec<P>, key: H) -> Exchange<T, CB, P, H> {
-        let builders = std::iter::repeat_with(Default::default).take(pushers.len()).collect();
+    /// Allocates a new `Exchange` from a supplied set of pushers and a distributor.
+    pub fn new(pushers: Vec<P>, distributor: D) -> Exchange<T, C, P, D> {
         Exchange {
             pushers,
-            hash_func: key,
-            builders,
             current: None,
-        }
-    }
-    #[inline]
-    fn flush(&mut self, index: usize) {
-        while let Some(container) = self.builders[index].finish() {
-            if let Some(ref time) = self.current {
-                Message::push_at(container, time.clone(), &mut self.pushers[index]);
-            }
+            distributor,
+            _phantom: std::marker::PhantomData,
         }
     }
 }
 
-impl<T: Eq+Data, CB, P, H> Push<Message<T, CB::Container>> for Exchange<T, CB, P, H>
+impl<T: Eq+Data, C, P, D> Push<Message<T, C>> for Exchange<T, C, P, D>
 where
-    CB: ContainerBuilder,
-    CB: for<'a> PushInto<<CB::Container as Container>::Item<'a>>,
-    P: Push<Message<T, CB::Container>>,
-    for<'a> H: FnMut(&<CB::Container as Container>::Item<'a>) -> u64
+    P: Push<Message<T, C>>,
+    D: Distributor<T, C, P>,
 {
     #[inline(never)]
-    fn push(&mut self, message: &mut Option<Message<T, CB::Container>>) {
+    fn push(&mut self, message: &mut Option<Message<T, C>>) {
         // if only one pusher, no exchange
         if self.pushers.len() == 1 {
             self.pushers[0].push(message);
@@ -65,35 +132,19 @@ where
 
             // if the time isn't right, flush everything.
             if self.current.as_ref().is_some_and(|x| x != time) {
-                for index in 0..self.pushers.len() {
-                    self.flush(index);
-                }
+                self.distributor.flush(time, &mut self.pushers);
             }
             self.current = Some(time.clone());
 
-            let hash_func = &mut self.hash_func;
-
-            // if the number of pushers is a power of two, use a mask
-            if self.pushers.len().is_power_of_two() {
-                let mask = (self.pushers.len() - 1) as u64;
-                CB::partition(data, &mut self.builders, |datum| ((hash_func)(datum) & mask) as usize);
-            }
-            // as a last resort, use mod (%)
-            else {
-                let num_pushers = self.pushers.len() as u64;
-                CB::partition(data, &mut self.builders, |datum| ((hash_func)(datum) % num_pushers) as usize);
-            }
-            for (buffer, pusher) in self.builders.iter_mut().zip(self.pushers.iter_mut()) {
-                while let Some(container) = buffer.extract() {
-                    Message::push_at(container, time.clone(), pusher);
-                }
-            }
+            self.distributor.partition(data, time, &mut self.pushers);
         }
         else {
             // flush
+            if let Some(time) = self.current.as_ref() {
+                self.distributor.flush(time, &mut self.pushers);
+            }
+            self.distributor.relax();
             for index in 0..self.pushers.len() {
-                self.flush(index);
-                self.builders[index].relax();
                 self.pushers[index].push(&mut None);
             }
         }