Add a metadata map to the hollow batch part

Author: bkirwi

src/persist-client/src/batch.rs

@@ -51,7 +51,9 @@ use crate::async_runtime::IsolatedRuntime;
 use crate::cfg::{BATCH_BUILDER_MAX_OUTSTANDING_PARTS, MiB};
 use crate::error::InvalidUsage;
 use crate::internal::compact::{CompactConfig, Compactor};
-use crate::internal::encoding::{LazyInlineBatchPart, LazyPartStats, LazyProto, Schemas};
+use crate::internal::encoding::{
+    LazyInlineBatchPart, LazyPartStats, LazyProto, MetadataMap, Schemas,
+};
 use crate::internal::machine::retry_external;
 use crate::internal::merge::{MergeTree, Pending};
 use crate::internal::metrics::{BatchWriteMetrics, Metrics, RetryMetrics, ShardMetrics};
@@ -1279,8 +1281,10 @@ impl<T: Timestamp + Codec64> BatchParts<T> {
             stats
         });
 
+        let meta = MetadataMap::default();
         BatchPart::Hollow(HollowBatchPart {
             key: partial_key,
+            meta,
             encoded_size_bytes: payload_len,
             key_lower,
             structured_key_lower,

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

@@ -104,6 +104,7 @@ impl<'a> DirectiveArgs<'a> {
                 .map(|x| {
                     RunPart::Single(BatchPart::Hollow(HollowBatchPart {
                         key: PartialBatchKey((*x).to_owned()),
+                        meta: Default::default(),
                         encoded_size_bytes: 0,
                         key_lower: vec![],
                         structured_key_lower: None,

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

@@ -1444,6 +1444,7 @@ impl<T: Timestamp + Codec64> RustType<ProtoHollowBatch> for HollowBatch<T> {
         parts.extend(proto.deprecated_keys.into_iter().map(|key| {
             RunPart::Single(BatchPart::Hollow(HollowBatchPart {
                 key: PartialBatchKey(key),
+                meta: Default::default(),
                 encoded_size_bytes: 0,
                 key_lower: vec![],
                 structured_key_lower: None,
@@ -1554,6 +1555,7 @@ impl<T: Timestamp + Codec64> RustType<ProtoHollowBatchPart> for HollowRunRef<T>
             schema_id: None,
             structured_key_lower: self.structured_key_lower.into_proto(),
             deprecated_schema_id: None,
+            metadata: BTreeMap::default(),
         };
         part
     }
@@ -1591,6 +1593,7 @@ impl<T: Timestamp + Codec64> RustType<ProtoHollowBatchPart> for BatchPart<T> {
                 format: x.format.map(|f| f.into_proto()),
                 schema_id: x.schema_id.into_proto(),
                 deprecated_schema_id: x.deprecated_schema_id.into_proto(),
+                metadata: BTreeMap::default(),
             },
             BatchPart::Inline {
                 updates,
@@ -1608,6 +1611,7 @@ impl<T: Timestamp + Codec64> RustType<ProtoHollowBatchPart> for BatchPart<T> {
                 format: None,
                 schema_id: schema_id.into_proto(),
                 deprecated_schema_id: deprecated_schema_id.into_proto(),
+                metadata: BTreeMap::default(),
             },
         }
     }
@@ -1623,6 +1627,7 @@ impl<T: Timestamp + Codec64> RustType<ProtoHollowBatchPart> for BatchPart<T> {
             Some(proto_hollow_batch_part::Kind::Key(key)) => {
                 Ok(BatchPart::Hollow(HollowBatchPart {
                     key: key.into_rust()?,
+                    meta: proto.metadata.into_rust()?,
                     encoded_size_bytes: proto.encoded_size_bytes.into_rust()?,
                     key_lower: proto.key_lower.into(),
                     structured_key_lower: proto.structured_key_lower.into_rust()?,
@@ -2008,6 +2013,7 @@ mod tests {
             ),
             vec![RunPart::Single(BatchPart::Hollow(HollowBatchPart {
                 key: PartialBatchKey("a".into()),
+                meta: Default::default(),
                 encoded_size_bytes: 5,
                 key_lower: vec![],
                 structured_key_lower: None,
@@ -2035,6 +2041,7 @@ mod tests {
             .parts
             .push(RunPart::Single(BatchPart::Hollow(HollowBatchPart {
                 key: PartialBatchKey("b".into()),
+                meta: Default::default(),
                 encoded_size_bytes: 0,
                 key_lower: vec![],
                 structured_key_lower: None,

src/persist-client/src/internal/state.proto

@@ -60,6 +60,8 @@ message ProtoHollowBatchPart {
   }
   optional uint64 schema_id = 12;
 
+  map<string, bytes> metadata = 14;
+
   optional bytes key_stats = 536870906;
 
   reserved 536870907 to 536870911;

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

@@ -56,7 +56,9 @@ use uuid::Uuid;
 
 use crate::critical::CriticalReaderId;
 use crate::error::InvalidUsage;
-use crate::internal::encoding::{LazyInlineBatchPart, LazyPartStats, LazyProto, parse_id};
+use crate::internal::encoding::{
+    LazyInlineBatchPart, LazyPartStats, LazyProto, MetadataMap, parse_id,
+};
 use crate::internal::gc::GcReq;
 use crate::internal::machine::retry_external;
 use crate::internal::paths::{BlobKey, PartId, PartialBatchKey, PartialRollupKey, WriterKey};
@@ -825,6 +827,9 @@ pub struct RunMeta {
 pub struct HollowBatchPart<T> {
     /// Pointer usable to retrieve the updates.
     pub key: PartialBatchKey,
+    /// Miscellaneous metadata.
+    #[serde(skip_serializing_if = "MetadataMap::is_empty")]
+    pub meta: MetadataMap,
     /// The encoded size of this part.
     pub encoded_size_bytes: usize,
     /// A lower bound on the keys in the part. (By default, this the minimum
@@ -1213,6 +1218,7 @@ impl<T: Ord> Ord for HollowBatchPart<T> {
         // are added.
         let HollowBatchPart {
             key: self_key,
+            meta: self_meta,
             encoded_size_bytes: self_encoded_size_bytes,
             key_lower: self_key_lower,
             structured_key_lower: self_structured_key_lower,
@@ -1225,6 +1231,7 @@ impl<T: Ord> Ord for HollowBatchPart<T> {
         } = self;
         let HollowBatchPart {
             key: other_key,
+            meta: other_meta,
             encoded_size_bytes: other_encoded_size_bytes,
             key_lower: other_key_lower,
             structured_key_lower: other_structured_key_lower,
@@ -1237,6 +1244,7 @@ impl<T: Ord> Ord for HollowBatchPart<T> {
         } = other;
         (
             self_key,
+            self_meta,
             self_encoded_size_bytes,
             self_key_lower,
             self_structured_key_lower,
@@ -1249,6 +1257,7 @@ impl<T: Ord> Ord for HollowBatchPart<T> {
         )
             .cmp(&(
                 other_key,
+                other_meta,
                 other_encoded_size_bytes,
                 other_key_lower,
                 other_structured_key_lower,
@@ -2998,6 +3007,7 @@ pub(crate) mod tests {
             )| {
                 HollowBatchPart {
                     key,
+                    meta: Default::default(),
                     encoded_size_bytes,
                     key_lower,
                     structured_key_lower: None,
@@ -3171,6 +3181,7 @@ pub(crate) mod tests {
                 .map(|x| {
                     RunPart::Single(BatchPart::Hollow(HollowBatchPart {
                         key: PartialBatchKey((*x).to_owned()),
+                        meta: Default::default(),
                         encoded_size_bytes: 0,
                         key_lower: vec![],
                         structured_key_lower: None,

src/persist-client/src/iter.rs

@@ -1273,6 +1273,7 @@ mod tests {
                             key: PartialBatchKey(
                                 "n0000000/p00000000-0000-0000-0000-000000000000".into(),
                             ),
+                            meta: Default::default(),
                             encoded_size_bytes,
                             key_lower: vec![],
                             structured_key_lower: None,