Stash the encoded schemas in Batch

Author: bkirwi

src/persist-client/build.rs

@@ -52,6 +52,7 @@ fn main() {
         .type_attribute(".", "#[allow(missing_docs)]")
         .btree_map(["."])
         .bytes([
+            ".mz_persist_client.batch.ProtoBatch",
             ".mz_persist_client.internal.diff.ProtoStateFieldDiffs",
             ".mz_persist_client.internal.service.ProtoPushDiff",
             ".mz_persist_client.internal.state.ProtoEncodedSchemas",

src/persist-client/src/batch.proto

@@ -21,6 +21,9 @@ import "persist-client/src/internal/state.proto";
 message ProtoBatch {
   string shard_id = 1;
   string version = 2;
+  bytes key_schema = 5;
+  bytes val_schema = 6;
+
   mz_persist_client.internal.state.ProtoHollowBatch batch = 3;
 
   reserved 4;

src/persist-client/src/batch.rs

@@ -81,6 +81,9 @@ pub struct Batch<K, V, T, D> {
     /// The version of Materialize which wrote this batch.
     pub(crate) version: Version,
 
+    /// The encoded schemas of the data in the batch.
+    pub(crate) schemas: (Bytes, Bytes),
+
     /// A handle to the data represented by this batch.
     pub(crate) batch: HollowBatch<T>,
 
@@ -121,13 +124,15 @@ where
         blob: Arc<dyn Blob>,
         shard_metrics: Arc<ShardMetrics>,
         version: Version,
+        schemas: (Bytes, Bytes),
         batch: HollowBatch<T>,
     ) -> Self {
         Self {
             batch_delete_enabled,
             metrics,
             shard_metrics,
             version,
+            schemas,
             batch,
             blob,
             _phantom: PhantomData,
@@ -209,6 +214,8 @@ where
             shard_id: self.shard_metrics.shard_id.into_proto(),
             version: self.version.to_string(),
             batch: Some(self.batch.into_proto()),
+            key_schema: self.schemas.0.clone(),
+            val_schema: self.schemas.1.clone(),
         };
         self.mark_consumed();
         ret
@@ -729,6 +736,10 @@ where
             self.blob,
             shard_metrics,
             self.version,
+            (
+                K::encode_schema(&*self.write_schemas.key),
+                V::encode_schema(&*self.write_schemas.val),
+            ),
             HollowBatch::new(desc, run_parts, total_updates, run_meta, run_splits),
         );
 

src/persist-client/src/internal/machine.rs

@@ -1538,6 +1538,10 @@ pub mod datadriven {
                 Arc::clone(&self.client.blob),
                 self.client.metrics.shards.shard(&self.shard_id, "test"),
                 self.client.cfg.build_version.clone(),
+                (
+                    <String>::encode_schema(&*SCHEMAS.key),
+                    <()>::encode_schema(&*SCHEMAS.val),
+                ),
                 hollow,
             )
         }

src/persist-client/src/lib.rs

@@ -603,6 +603,7 @@ impl PersistClient {
             metrics: Arc::clone(&self.metrics),
             shard_metrics,
             version: Version::parse(&batch.version).expect("valid transmittable batch"),
+            schemas: (batch.key_schema, batch.val_schema),
             batch: batch
                 .batch
                 .into_rust_if_some("ProtoBatch::batch")

src/persist-client/src/write.rs

@@ -822,6 +822,7 @@ where
             metrics: Arc::clone(&self.metrics),
             shard_metrics: Arc::clone(&self.machine.applier.shard_metrics),
             version: Version::parse(&batch.version).expect("valid transmittable batch"),
+            schemas: (batch.key_schema, batch.val_schema),
             batch: batch
                 .batch
                 .into_rust_if_some("ProtoBatch::batch")