columnar: Refactor PartEncoder, add decode_into(...) to PartDecoder (MaterializeInc#26669)

This PR was split out from
MaterializeInc#26605



### Motivation

This refactors `PartEncoder` to take ownership of the columns it's
encoding into, rather than just mutable references. More detail is
provided in MaterializeInc#26605 but
the tl;dr is we only want to downcast `dyn Array`s once, and there isn't
any lifetime that we can associate with the borrows to achieve this.

It also renames the existing `PartDecoder::decode(idx, &mut V)` method
to `PartDecoder::decode_into(...)` and adds a new
`PartDecoder::decode(idx) -> V` method. The goal of `decode_into(...)`
is it allows you to re-use allocations where possible, but that isn't
currently applicable when decoding from our structured columnar data.
When decoding with `Codec` we go from `&[u8]` -> `ProtoRow` -> `Row`, so
it's very helpful to retain the intermediate `ProtoRow` to reduce
allocations. But with columnar data we go directly from `dyn
arrow::Array` -> `Row`, so there isn't any intermediate step to retain.
Also we return the `Row`s we decode as part of an iterator, so there
isn't any way to reclaim them for re-use.

In practice without `PartDecoder::decode(idx)` we end up with a pattern
like:
```
let mut val = V::default();
decoder.decode_into(idx, &mut val);
val
```
Instead of requiring our generic parameters `K` and `V` to implement
`Default`, it was easier to add the new method.

### Tips for reviewer

This PR is broken up into two commits:

1. Refactor to `PartEncoder`
2. New methods on `PartDecoder`

### Checklist

- [ ] This PR has adequate test coverage / QA involvement has been duly
considered. ([trigger-ci for additional test/nightly
runs](https://trigger-ci.dev.materialize.com/))
- [ ] This PR has an associated up-to-date [design
doc](https://github.com/MaterializeInc/materialize/blob/main/doc/developer/design/README.md),
is a design doc
([template](https://github.com/MaterializeInc/materialize/blob/main/doc/developer/design/00000000_template.md)),
or is sufficiently small to not require a design.
  <!-- Reference the design in the description. -->
- [ ] If this PR evolves [an existing `$T ⇔ Proto$T`
mapping](https://github.com/MaterializeInc/materialize/blob/main/doc/developer/command-and-response-binary-encoding.md)
(possibly in a backwards-incompatible way), then it is tagged with a
`T-proto` label.
- [ ] If this PR will require changes to cloud orchestration or tests,
there is a companion cloud PR to account for those changes that is
tagged with the release-blocker label
([example](MaterializeInc/cloud#5021)).
<!-- Ask in #team-cloud on Slack if you need help preparing the cloud
PR. -->
- [x] This PR includes the following [user-facing behavior
changes](https://github.com/MaterializeInc/materialize/blob/main/doc/developer/guide-changes.md#what-changes-require-a-release-note):
  - N/a

Author: ParkMyCar

src/persist-cli/src/maelstrom/txn_list_append_single.rs

@@ -46,7 +46,7 @@ use crate::maelstrom::txn_list_append_single::codec_impls::{
 use crate::maelstrom::Args;
 
 /// Key of the persist shard used by [Transactor]
-#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
+#[derive(Debug, Clone, Default, PartialEq, Eq, PartialOrd, Ord, Hash)]
 pub struct MaelstromKey(u64);
 
 /// Val of the persist shard used by [Transactor]
@@ -749,8 +749,7 @@ mod codec_impls {
     pub struct MaelstromKeySchema;
 
     impl Schema<MaelstromKey> for MaelstromKeySchema {
-        type Encoder<'a> = SimpleEncoder<'a, MaelstromKey, u64>;
-
+        type Encoder = SimpleEncoder<MaelstromKey, u64>;
         type Decoder = SimpleDecoder<MaelstromKey, u64>;
 
         fn columns(&self) -> DynStructCfg {
@@ -761,7 +760,7 @@ mod codec_impls {
             SimpleSchema::<MaelstromKey, u64>::decoder(cols, |val, ret| ret.0 = val)
         }
 
-        fn encoder<'a>(&self, cols: ColumnsMut<'a>) -> Result<Self::Encoder<'a>, String> {
+        fn encoder(&self, cols: ColumnsMut) -> Result<Self::Encoder, String> {
             SimpleSchema::<MaelstromKey, u64>::encoder(cols, |val| val.0)
         }
     }
@@ -793,8 +792,7 @@ mod codec_impls {
     pub struct MaelstromValSchema;
 
     impl Schema<MaelstromVal> for MaelstromValSchema {
-        type Encoder<'a> = SimpleEncoder<'a, MaelstromVal, Vec<u8>>;
-
+        type Encoder = SimpleEncoder<MaelstromVal, Vec<u8>>;
         type Decoder = SimpleDecoder<MaelstromVal, Vec<u8>>;
 
         fn columns(&self) -> DynStructCfg {
@@ -807,7 +805,7 @@ mod codec_impls {
             })
         }
 
-        fn encoder<'a>(&self, cols: ColumnsMut<'a>) -> Result<Self::Encoder<'a>, String> {
+        fn encoder(&self, cols: ColumnsMut) -> Result<Self::Encoder, String> {
             SimpleSchema::<MaelstromVal, Vec<u8>>::push_encoder(cols, |col, val| {
                 let mut buf = Vec::new();
                 MaelstromVal::encode(val, &mut buf);

src/persist-client/src/fetch.rs

@@ -629,8 +629,8 @@ pub(crate) struct EncodedPart<T> {
 
 impl<K, V, T, D> FetchedPart<K, V, T, D>
 where
-    K: Debug + Codec,
-    V: Debug + Codec,
+    K: Debug + Codec + Default,
+    V: Debug + Codec + Default,
     T: Timestamp + Lattice + Codec64,
     D: Semigroup + Codec64 + Send + Sync,
 {
@@ -694,8 +694,8 @@ where
 
 impl<K, V, T, D> Iterator for FetchedPart<K, V, T, D>
 where
-    K: Debug + Codec,
-    V: Debug + Codec,
+    K: Debug + Codec + Default,
+    V: Debug + Codec + Default,
     T: Timestamp + Lattice + Codec64,
     D: Semigroup + Codec64 + Send + Sync,
 {

src/persist-client/src/read.rs

@@ -147,7 +147,15 @@ where
             }
         }
     }
+}
 
+impl<K, V, T, D> Subscribe<K, V, T, D>
+where
+    K: Debug + Codec + Default,
+    V: Debug + Codec + Default,
+    T: Timestamp + Lattice + Codec64,
+    D: Semigroup + Codec64 + Send + Sync,
+{
     /// Equivalent to `next`, but rather than returning a [`LeasedBatchPart`],
     /// fetches and returns the data from within it.
     #[instrument(level = "debug", fields(shard = %self.listen.handle.machine.shard_id()))]
@@ -184,7 +192,15 @@ where
     pub async fn fetch_batch_part(&mut self, part: LeasedBatchPart<T>) -> FetchedPart<K, V, T, D> {
         self.listen.fetch_batch_part(part).await
     }
+}
 
+impl<K, V, T, D> Subscribe<K, V, T, D>
+where
+    K: Debug + Codec,
+    V: Debug + Codec,
+    T: Timestamp + Lattice + Codec64,
+    D: Semigroup + Codec64 + Send + Sync,
+{
     /// Takes a [`SerdeLeasedBatchPart`] into a [`LeasedBatchPart`].
     pub fn leased_part_from_exchangeable(&self, x: SerdeLeasedBatchPart) -> LeasedBatchPart<T> {
         self.listen
@@ -269,19 +285,6 @@ where
         }
     }
 
-    /// Convert listener into futures::Stream
-    pub fn into_stream(
-        mut self,
-    ) -> impl Stream<Item = ListenEvent<T, ((Result<K, String>, Result<V, String>), T, D)>> {
-        async_stream::stream!({
-            loop {
-                for msg in self.fetch_next().await {
-                    yield msg;
-                }
-            }
-        })
-    }
-
     /// An exclusive upper bound on the progress of this Listen.
     pub fn frontier(&self) -> &Antichain<T> {
         &self.frontier
@@ -383,7 +386,15 @@ where
 
         (parts, self.frontier.clone())
     }
+}
 
+impl<K, V, T, D> Listen<K, V, T, D>
+where
+    K: Debug + Codec + Default,
+    V: Debug + Codec + Default,
+    T: Timestamp + Lattice + Codec64,
+    D: Semigroup + Codec64 + Send + Sync,
+{
     /// Attempt to pull out the next values of this subscription.
     ///
     /// The updates received in [ListenEvent::Updates] should be assumed to be in arbitrary order
@@ -411,6 +422,54 @@ where
         ret
     }
 
+    /// Convert listener into futures::Stream
+    pub fn into_stream(
+        mut self,
+    ) -> impl Stream<Item = ListenEvent<T, ((Result<K, String>, Result<V, String>), T, D)>> {
+        async_stream::stream!({
+            loop {
+                for msg in self.fetch_next().await {
+                    yield msg;
+                }
+            }
+        })
+    }
+
+    /// Test helper to read from the listener until the given frontier is
+    /// reached. Because compaction can arbitrarily combine batches, we only
+    /// return the final progress info.
+    #[cfg(test)]
+    #[track_caller]
+    pub async fn read_until(
+        &mut self,
+        ts: &T,
+    ) -> (
+        Vec<((Result<K, String>, Result<V, String>), T, D)>,
+        Antichain<T>,
+    ) {
+        let mut updates = Vec::new();
+        let mut frontier = Antichain::from_elem(T::minimum());
+        while self.frontier.less_than(ts) {
+            for event in self.fetch_next().await {
+                match event {
+                    ListenEvent::Updates(mut x) => updates.append(&mut x),
+                    ListenEvent::Progress(x) => frontier = x,
+                }
+            }
+        }
+        // Unlike most tests, intentionally don't consolidate updates here
+        // because Listen replays them at the original fidelity.
+        (updates, frontier)
+    }
+}
+
+impl<K, V, T, D> Listen<K, V, T, D>
+where
+    K: Debug + Codec,
+    V: Debug + Codec,
+    T: Timestamp + Lattice + Codec64,
+    D: Semigroup + Codec64 + Send + Sync,
+{
     /// Fetches the contents of `part` and returns its lease.
     ///
     /// This is broken out into its own function to provide a trivial means for
@@ -441,33 +500,6 @@ where
     pub async fn expire(self) {
         self.handle.expire().await
     }
-
-    /// Test helper to read from the listener until the given frontier is
-    /// reached. Because compaction can arbitrarily combine batches, we only
-    /// return the final progress info.
-    #[cfg(test)]
-    #[track_caller]
-    pub async fn read_until(
-        &mut self,
-        ts: &T,
-    ) -> (
-        Vec<((Result<K, String>, Result<V, String>), T, D)>,
-        Antichain<T>,
-    ) {
-        let mut updates = Vec::new();
-        let mut frontier = Antichain::from_elem(T::minimum());
-        while self.frontier.less_than(ts) {
-            for event in self.fetch_next().await {
-                match event {
-                    ListenEvent::Updates(mut x) => updates.append(&mut x),
-                    ListenEvent::Progress(x) => frontier = x,
-                }
-            }
-        }
-        // Unlike most tests, intentionally don't consolidate updates here
-        // because Listen replays them at the original fidelity.
-        (updates, frontier)
-    }
 }
 
 #[derive(Clone, Debug)]
@@ -1100,7 +1132,15 @@ where
             parts,
         })
     }
+}
 
+impl<K, V, T, D> ReadHandle<K, V, T, D>
+where
+    K: Debug + Codec + Ord + Default,
+    V: Debug + Codec + Ord + Default,
+    T: Timestamp + Lattice + Codec64,
+    D: Semigroup + Codec64 + Send + Sync,
+{
     /// Generates a [Self::snapshot], and streams out all of the updates
     /// it contains in bounded memory.
     ///

src/persist-client/src/stats.rs

@@ -14,7 +14,6 @@ use std::sync::Arc;
 
 use mz_dyncfg::{Config, ConfigSet};
 use mz_persist::indexed::columnar::ColumnarRecords;
-use mz_persist_types::columnar::{PartEncoder, Schema};
 use mz_persist_types::part::PartBuilder;
 use mz_persist_types::stats::PartStats;
 use mz_persist_types::Codec;
@@ -129,24 +128,20 @@ pub(crate) fn part_stats_for_legacy_part<K: Codec, V: Codec>(
     // This is a laughably inefficient placeholder implementation of stats
     // on the old part format. We don't intend to make this fast, rather we
     // intend to compute stats on the new part format.
-    let mut new_format = PartBuilder::new(schemas.key.as_ref(), schemas.val.as_ref());
-    let mut builder = new_format.get_mut();
-    let mut key = schemas.key.encoder(builder.key)?;
-    let mut val = schemas.val.encoder(builder.val)?;
+    let mut builder = PartBuilder::new(schemas.key.as_ref(), schemas.val.as_ref())?;
     for x in part {
         for ((k, v), t, d) in x.iter() {
             let k = K::decode(k)?;
             let v = V::decode(v)?;
-            key.encode(&k);
-            val.encode(&v);
-            builder.ts.push(i64::from_le_bytes(t));
-            builder.diff.push(i64::from_le_bytes(d));
+            let t = i64::from_le_bytes(t);
+            let d = i64::from_le_bytes(d);
+
+            builder.push(&k, &v, t, d);
         }
     }
-    drop(key);
-    drop(val);
-    let new_format = new_format.finish()?;
-    PartStats::new(&new_format)
+    let part = builder.finish();
+
+    PartStats::new(&part)
 }
 
 /// Statistics about the contents of a shard as_of some time.

src/persist-txn/src/txn_read.rs

@@ -180,8 +180,8 @@ impl<T: Timestamp + Lattice + TotalOrder + Codec64> DataSnapshot<T> {
         data_read: &mut ReadHandle<K, V, T, D>,
     ) -> Result<impl Stream<Item = ((Result<K, String>, Result<V, String>), T, D)>, Since<T>>
     where
-        K: Debug + Codec + Ord,
-        V: Debug + Codec + Ord,
+        K: Debug + Codec + Ord + Default,
+        V: Debug + Codec + Ord + Default,
         D: Semigroup + Codec64 + Send + Sync,
     {
         let data_write = WriteHandle::from_read(data_read, "unblock_read");