Promote columnar Column/ColumnBuilder into the container crate

container/Cargo.toml

@@ -8,3 +8,7 @@ rust-version.workspace = true
 
 [lints]
 workspace = true
+
+[dependencies]
+columnar = { workspace = true }
+timely_bytes = { path = "../bytes", version = "0.30" }

container/src/columnar.rs

@@ -0,0 +1,151 @@
+//! A container storing [`columnar`] data, either as a typed container or as
+//! serialized bytes.
+//!
+//! [`Column`] is [`columnar`]'s [`Stash`] fixed to timely's [`Bytes`] type. It
+//! implements the container traits so that columnar data can be produced,
+//! exchanged, and consumed like any other container; its contents can be borrowed
+//! whether they are still typed or have been serialized to bytes.
+
+use std::collections::VecDeque;
+
+use columnar::bytes::{indexed, stash::Stash};
+use columnar::common::IterOwn;
+use columnar::{Columnar, Index, Len};
+
+use timely_bytes::arc::Bytes;
+
+use crate::{
+    Accountable, ContainerBuilder, DrainContainer, LengthPreservingContainerBuilder, PushInto,
+    SizableContainer,
+};
+
+/// A columnar container, held either as a typed container or as serialized bytes.
+///
+/// This is [`columnar`]'s [`Stash`] with the byte representation fixed to timely's
+/// [`Bytes`]. Received data are borrowed in place (no per-record deserialization);
+/// owning an element requires `into_owned`.
+pub type Column<C> = Stash<C, Bytes>;
+
+/// The [`Column`] type appropriate for a [`Columnar`] row type `T`.
+pub type ColumnOf<T> = Column<<T as Columnar>::Container>;
+
+/// A built [`Column`] is considered full once its serialized size reaches this many
+/// bytes. This sets the columnar batching granularity, which is much coarser than
+/// the `Vec` default (see [`crate::buffer::BUFFER_SIZE_BYTES`]).
+pub const COLUMN_CAPACITY_BYTES: usize = 1 << 20;
+
+impl<C: columnar::ContainerBytes> Accountable for Column<C> {
+    #[inline]
+    fn record_count(&self) -> i64 {
+        i64::try_from(self.borrow().len()).unwrap()
+    }
+    #[inline]
+    fn is_empty(&self) -> bool {
+        self.borrow().is_empty()
+    }
+}
+
+impl<C: columnar::ContainerBytes> DrainContainer for Column<C> {
+    type Item<'a> = C::Ref<'a>;
+    type DrainIter<'a> = IterOwn<C::Borrowed<'a>>;
+    fn drain(&mut self) -> Self::DrainIter<'_> {
+        self.borrow().into_index_iter()
+    }
+}
+
+impl<C: columnar::ContainerBytes> SizableContainer for Column<C> {
+    fn at_capacity(&self) -> bool {
+        match self {
+            Stash::Typed(t) => 8 * indexed::length_in_words(&t.borrow()) >= COLUMN_CAPACITY_BYTES,
+            Stash::Bytes(_) | Stash::Align(_) => true,
+        }
+    }
+    fn ensure_capacity(&mut self, _stash: &mut Option<Self>) {}
+}
+
+impl<C: columnar::Container + columnar::ContainerBytes, T> PushInto<T> for Column<C>
+where
+    C: columnar::Push<T>,
+{
+    #[inline]
+    fn push_into(&mut self, item: T) {
+        use columnar::Push;
+        self.push(item)
+    }
+}
+
+/// A container builder for [`Column`].
+///
+/// Items are accumulated into a typed columnar container. When that container's
+/// serialized size approaches a target, it is encoded into an aligned byte buffer
+/// which is enqueued for extraction, and the typed container is cleared for reuse.
+///
+/// Encoding eagerly, rather than deferring encoding until data are sent, keeps the
+/// (comparatively expensive to rebuild) typed accumulator resident in the builder:
+/// only the flat encoded bytes travel onward, so the columns need not be rebuilt
+/// each batch. This matters on the exchange path, where a sealed container is sent
+/// to a peer and never returns, so there is no opportunity to recycle it.
+#[derive(Default)]
+pub struct ColumnBuilder<C> {
+    /// Container into which we accumulate pushed items.
+    current: C,
+    /// An empty allocation, retained for reuse across extractions.
+    empty: Option<Column<C>>,
+    /// Completed containers, pending extraction.
+    pending: VecDeque<Column<C>>,
+}
+
+impl<C: columnar::ContainerBytes, T> PushInto<T> for ColumnBuilder<C>
+where
+    C: columnar::Push<T>,
+{
+    #[inline]
+    fn push_into(&mut self, item: T) {
+        self.current.push(item);
+        // If there is less than 10% slop with a 2MB-aligned backing allocation, mint
+        // a container by encoding into an exactly-sized aligned buffer.
+        let words = indexed::length_in_words(&self.current.borrow());
+        let round = (words + ((1 << 18) - 1)) & !((1 << 18) - 1);
+        if round - words < round / 10 {
+            let mut alloc = Vec::with_capacity(round);
+            indexed::encode(&mut alloc, &self.current.borrow());
+            self.pending.push_back(Stash::Align(alloc.into_boxed_slice().into()));
+            self.current.clear();
+        }
+    }
+}
+
+impl<C: columnar::ContainerBytes> ContainerBuilder for ColumnBuilder<C> {
+    type Container = Column<C>;
+
+    #[inline]
+    fn extract(&mut self) -> Option<&mut Self::Container> {
+        if let Some(container) = self.pending.pop_front() {
+            self.empty = Some(container);
+            self.empty.as_mut()
+        } else {
+            None
+        }
+    }
+
+    #[inline]
+    fn finish(&mut self) -> Option<&mut Self::Container> {
+        if !self.current.is_empty() {
+            self.pending
+                .push_back(Stash::Typed(std::mem::take(&mut self.current)));
+        }
+        self.empty = self.pending.pop_front();
+        self.empty.as_mut()
+    }
+
+    #[inline]
+    fn relax(&mut self) {
+        // The caller is responsible for draining all contents; assert that we are empty.
+        // The assertion is not strictly necessary, but it helps catch bugs.
+        assert!(self.current.is_empty());
+        assert!(self.pending.is_empty());
+        *self = Self::default();
+    }
+}
+
+impl<C: columnar::ContainerBytes> LengthPreservingContainerBuilder for ColumnBuilder<C> {}

container/src/lib.rs

@@ -4,6 +4,9 @@
 
 use std::collections::VecDeque;
 
+pub mod columnar;
+pub use columnar::{Column, ColumnBuilder, ColumnOf, COLUMN_CAPACITY_BYTES};
+
 /// A type containing a number of records accounted for by progress tracking.
 ///
 /// The object stores a number of updates and thus is able to describe it count

timely/examples/columnar.rs

@@ -4,7 +4,7 @@ use std::collections::HashMap;
 
 use columnar::Index;
 use timely::Accountable;
-use timely::container::CapacityContainerBuilder;
+use timely::container::{CapacityContainerBuilder, Column, ColumnBuilder};
 use timely::dataflow::channels::pact::{ExchangeCore, Pipeline};
 use timely::dataflow::InputHandle;
 use timely::dataflow::operators::{InspectCore, Operator, Probe};
@@ -124,127 +124,3 @@ fn main() {
     })
     .unwrap();
 }
-
-
-pub use container::Column;
-mod container {
-
-    use columnar::bytes::stash::Stash;
-
-    #[derive(Clone, Default)]
-    pub struct Column<C> { pub stash: Stash<C, timely_bytes::arc::Bytes> }
-
-    use columnar::{Len, Index};
-    use columnar::bytes::indexed;
-    use columnar::common::IterOwn;
-
-    impl<C: columnar::ContainerBytes> Column<C> {
-        /// Borrows the contents no matter their representation.
-        #[inline(always)] pub fn borrow(&self) -> C::Borrowed<'_> { self.stash.borrow() }
-    }
-
-    impl<C: columnar::ContainerBytes> timely::Accountable for Column<C> {
-        #[inline] fn record_count(&self) -> i64 { i64::try_from(self.borrow().len()).unwrap() }
-        #[inline] fn is_empty(&self) -> bool { self.borrow().is_empty() }
-    }
-    impl<C: columnar::ContainerBytes> timely::container::DrainContainer for Column<C> {
-        type Item<'a> = C::Ref<'a>;
-        type DrainIter<'a> = IterOwn<C::Borrowed<'a>>;
-        fn drain<'a>(&'a mut self) -> Self::DrainIter<'a> { self.borrow().into_index_iter() }
-    }
-
-    impl<C: columnar::ContainerBytes> timely::container::SizableContainer for Column<C> {
-        fn at_capacity(&self) -> bool {
-            match &self.stash {
-                Stash::Typed(t) => {
-                    let length_in_bytes = 8 * indexed::length_in_words(&t.borrow());
-                    length_in_bytes >= (1 << 20)
-                },
-                Stash::Bytes(_) => true,
-                Stash::Align(_) => true,
-            }
-        }
-        fn ensure_capacity(&mut self, _stash: &mut Option<Self>) { }
-    }
-
-    impl<C: columnar::Container + columnar::ContainerBytes, T> timely::container::PushInto<T> for Column<C> where C: columnar::Push<T> {
-        #[inline] fn push_into(&mut self, item: T) { use columnar::Push; self.stash.push(item) }
-    }
-
-    impl<C: columnar::ContainerBytes> timely::dataflow::channels::ContainerBytes for Column<C> {
-        fn from_bytes(bytes: timely::bytes::arc::Bytes) -> Self { Self { stash: Stash::try_from_bytes(bytes).expect("valid columnar data") } }
-        fn length_in_bytes(&self) -> usize { self.stash.length_in_bytes() }
-        fn into_bytes<W: ::std::io::Write>(&self, writer: &mut W) { self.stash.write_bytes(writer).expect("write failed") }
-    }
-}
-
-
-use builder::ColumnBuilder;
-mod builder {
-
-    use std::collections::VecDeque;
-    use columnar::bytes::{indexed, stash::Stash};
-    use super::Column;
-
-    /// A container builder for `Column<C>`.
-    #[derive(Default)]
-    pub struct ColumnBuilder<C> {
-        /// Container that we're writing to.
-        current: C,
-        /// Empty allocation.
-        empty: Option<Column<C>>,
-        /// Completed containers pending to be sent.
-        pending: VecDeque<Column<C>>,
-    }
-
-    impl<C: columnar::ContainerBytes, T> timely::container::PushInto<T> for ColumnBuilder<C> where C: columnar::Push<T> {
-        #[inline]
-        fn push_into(&mut self, item: T) {
-            self.current.push(item);
-            // If there is less than 10% slop with 2MB backing allocations, mint a container.
-            let words = indexed::length_in_words(&self.current.borrow());
-            let round = (words + ((1 << 18) - 1)) & !((1 << 18) - 1);
-            if round - words < round / 10 {
-                let mut alloc = Vec::with_capacity(round);
-                indexed::encode(&mut alloc, &self.current.borrow());
-                self.pending.push_back(Column { stash: Stash::Align(alloc.into_boxed_slice().into()) });
-                self.current.clear();
-            }
-        }
-    }
-
-    use timely::container::{ContainerBuilder, LengthPreservingContainerBuilder};
-    impl<C: columnar::ContainerBytes> ContainerBuilder for ColumnBuilder<C> {
-        type Container = Column<C>;
-
-        #[inline]
-        fn extract(&mut self) -> Option<&mut Self::Container> {
-            if let Some(container) = self.pending.pop_front() {
-                self.empty = Some(container);
-                self.empty.as_mut()
-            } else {
-                None
-            }
-        }
-
-        #[inline]
-        fn finish(&mut self) -> Option<&mut Self::Container> {
-            if !self.current.is_empty() {
-                self.pending.push_back(Column { stash: Stash::Typed(std::mem::take(&mut self.current)) });
-            }
-            self.empty = self.pending.pop_front();
-            self.empty.as_mut()
-        }
-
-        #[inline]
-        fn relax(&mut self) {
-            // The caller is responsible for draining all contents; assert that we are empty.
-            // The assertion is not strictly necessary, but it helps catch bugs.
-            assert!(self.current.is_empty());
-            assert!(self.pending.is_empty());
-            *self = Self::default();
-        }
-    }
-
-    impl<C: columnar::ContainerBytes> LengthPreservingContainerBuilder for ColumnBuilder<C> { }
-}

timely/src/dataflow/channels/mod.rs

@@ -117,6 +117,22 @@ mod implementations {
     use serde::{Serialize, Deserialize};
     use crate::dataflow::channels::ContainerBytes;
 
+    // The columnar container (`timely_container::Column`) serializes to and borrows
+    // from its own aligned byte representation; received bytes are borrowed in place
+    // rather than deserialized. The container traits live in `timely_container`; the
+    // wire encoding is this crate's `ContainerBytes` trait, so it is implemented here.
+    impl<C: columnar::ContainerBytes> ContainerBytes for crate::container::Column<C> {
+        fn from_bytes(bytes: crate::bytes::arc::Bytes) -> Self {
+            columnar::bytes::stash::Stash::try_from_bytes(bytes).expect("invalid columnar data")
+        }
+        fn length_in_bytes(&self) -> usize {
+            columnar::bytes::stash::Stash::length_in_bytes(self)
+        }
+        fn into_bytes<W: Write>(&self, writer: &mut W) {
+            columnar::bytes::stash::Stash::write_bytes(self, writer).expect("write failed")
+        }
+    }
+
     impl<T: Serialize + for<'a> Deserialize<'a>> ContainerBytes for Vec<T> {
         fn from_bytes(bytes: crate::bytes::arc::Bytes) -> Self {
             ::bincode::deserialize(&bytes[..]).expect("bincode::deserialize() failed")