persist: Add columnar format to HollowBatchPart

misc/python/materialize/mzcompose/__init__.py

@@ -104,6 +104,7 @@
     "storage_source_decode_fuel": "100000",
     "timestamp_oracle": "postgres",
     "wait_catalog_consolidation_on_startup": "true",
+    "persist_batch_record_part_format": "true",
 }
 
 DEFAULT_CRDB_ENVIRONMENT = [

misc/python/materialize/parallel_workload/action.py

@@ -933,6 +933,7 @@ def __init__(
         )
         self.flags_with_values["enable_eager_delta_joins"] = BOOLEAN_FLAG_VALUES
         self.flags_with_values["persist_batch_columnar_format"] = ["row", "both"]
+        self.flags_with_values["persist_batch_record_part_format"] = BOOLEAN_FLAG_VALUES
 
     def run(self, exe: Executor) -> bool:
         flag_name = self.rng.choice(list(self.flags_with_values.keys()))

src/persist-client/src/batch.rs

@@ -334,6 +334,7 @@ pub struct BatchBuilderConfig {
     pub(crate) batch_delete_enabled: bool,
     pub(crate) batch_builder_max_outstanding_parts: usize,
     pub(crate) batch_columnar_format: BatchColumnarFormat,
+    pub(crate) batch_record_part_format: bool,
     pub(crate) inline_writes_single_max_bytes: usize,
     pub(crate) stats_collection_enabled: bool,
     pub(crate) stats_budget: usize,
@@ -351,7 +352,13 @@ pub(crate) const BATCH_DELETE_ENABLED: Config<bool> = Config::new(
 pub(crate) const BATCH_COLUMNAR_FORMAT: Config<&'static str> = Config::new(
     "persist_batch_columnar_format",
     BatchColumnarFormat::default().as_str(),
-    "Columnar format for a batch written to Persist, either 'row' or 'both' (Materialize).",
+    "Columnar format for a batch written to Persist, either 'row', 'both', or 'both_v1' (Materialize).",
+);
+
+pub(crate) const BATCH_RECORD_PART_FORMAT: Config<bool> = Config::new(
+    "persist_batch_record_part_format",
+    false,
+    "Wether we record the format of the Part in state (Materialize).",
 );
 
 /// A target maximum size of blob payloads in bytes. If a logical "batch" is
@@ -392,6 +399,7 @@ impl BatchBuilderConfig {
                 .dynamic
                 .batch_builder_max_outstanding_parts(),
             batch_columnar_format: BatchColumnarFormat::from_str(&BATCH_COLUMNAR_FORMAT.get(value)),
+            batch_record_part_format: BATCH_RECORD_PART_FORMAT.get(value),
             inline_writes_single_max_bytes: INLINE_WRITES_SINGLE_MAX_BYTES.get(value),
             stats_collection_enabled: STATS_COLLECTION_ENABLED.get(value),
             stats_budget: STATS_BUDGET_BYTES.get(value),
@@ -1123,6 +1131,11 @@ impl<T: Timestamp + Codec64> BatchParts<T> {
             }
             stats
         });
+        let format = if cfg.batch_record_part_format {
+            Some(cfg.batch_columnar_format)
+        } else {
+            None
+        };
 
         BatchPart::Hollow(HollowBatchPart {
             key: partial_key,
@@ -1131,6 +1144,7 @@ impl<T: Timestamp + Codec64> BatchParts<T> {
             stats,
             ts_rewrite,
             diffs_sum: cfg.write_diffs_sum.then_some(diffs_sum),
+            format,
         })
     }
 

src/persist-client/src/cfg.rs

@@ -324,6 +324,7 @@ pub fn all_dyncfgs(configs: ConfigSet) -> ConfigSet {
     mz_persist::cfg::all_dyn_configs(configs)
         .add(&crate::batch::BATCH_DELETE_ENABLED)
         .add(&crate::batch::BATCH_COLUMNAR_FORMAT)
+        .add(&crate::batch::BATCH_RECORD_PART_FORMAT)
         .add(&crate::batch::BLOB_TARGET_SIZE)
         .add(&crate::batch::INLINE_WRITES_TOTAL_MAX_BYTES)
         .add(&crate::batch::INLINE_WRITES_SINGLE_MAX_BYTES)

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

@@ -111,6 +111,7 @@ impl<'a> DirectiveArgs<'a> {
                         stats: None,
                         ts_rewrite: None,
                         diffs_sum: None,
+                        format: None,
                     })
                 })
                 .collect(),

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

@@ -17,9 +17,10 @@ use std::sync::Arc;
 use bytes::{Buf, Bytes};
 use differential_dataflow::lattice::Lattice;
 use differential_dataflow::trace::Description;
+use mz_ore::cast::CastInto;
 use mz_ore::halt;
 use mz_persist::indexed::columnar::ColumnarRecords;
-use mz_persist::indexed::encoding::{BlobTraceBatchPart, BlobTraceUpdates};
+use mz_persist::indexed::encoding::{BatchColumnarFormat, BlobTraceBatchPart, BlobTraceUpdates};
 use mz_persist::location::{SeqNo, VersionedData};
 use mz_persist::metrics::ColumnarMetrics;
 use mz_persist_types::stats::{PartStats, ProtoStructStats};
@@ -1258,6 +1259,7 @@ impl<T: Timestamp + Codec64> RustType<ProtoHollowBatch> for HollowBatch<T> {
                 stats: None,
                 ts_rewrite: None,
                 diffs_sum: None,
+                format: None,
             })
         }));
         Ok(HollowBatch {
@@ -1279,6 +1281,7 @@ impl<T: Timestamp + Codec64> RustType<ProtoHollowBatchPart> for BatchPart<T> {
                 key_stats: x.stats.into_proto(),
                 ts_rewrite: x.ts_rewrite.as_ref().map(|x| x.into_proto()),
                 diffs_sum: x.diffs_sum.as_ref().map(|x| i64::from_le_bytes(*x)),
+                format: x.format.map(|f| f.into_proto()),
             },
             BatchPart::Inline {
                 updates,
@@ -1290,6 +1293,7 @@ impl<T: Timestamp + Codec64> RustType<ProtoHollowBatchPart> for BatchPart<T> {
                 key_stats: None,
                 ts_rewrite: ts_rewrite.as_ref().map(|x| x.into_proto()),
                 diffs_sum: None,
+                format: None,
             },
         }
     }
@@ -1308,6 +1312,7 @@ impl<T: Timestamp + Codec64> RustType<ProtoHollowBatchPart> for BatchPart<T> {
                     stats: proto.key_stats.into_rust()?,
                     ts_rewrite,
                     diffs_sum: proto.diffs_sum.map(i64::to_le_bytes),
+                    format: proto.format.map(|f| f.into_rust()).transpose()?,
                 }))
             }
             Some(proto_hollow_batch_part::Kind::Inline(x)) => {
@@ -1328,6 +1333,27 @@ impl<T: Timestamp + Codec64> RustType<ProtoHollowBatchPart> for BatchPart<T> {
     }
 }
 
+impl RustType<proto_hollow_batch_part::Format> for BatchColumnarFormat {
+    fn into_proto(&self) -> proto_hollow_batch_part::Format {
+        match self {
+            BatchColumnarFormat::Row => proto_hollow_batch_part::Format::Row(()),
+            BatchColumnarFormat::Both(version) => {
+                proto_hollow_batch_part::Format::RowAndColumnar((*version).cast_into())
+            }
+        }
+    }
+
+    fn from_proto(proto: proto_hollow_batch_part::Format) -> Result<Self, TryFromProtoError> {
+        let format = match proto {
+            proto_hollow_batch_part::Format::Row(_) => BatchColumnarFormat::Row,
+            proto_hollow_batch_part::Format::RowAndColumnar(version) => {
+                BatchColumnarFormat::Both(version.cast_into())
+            }
+        };
+        Ok(format)
+    }
+}
+
 /// Aggregate statistics about data contained in a part.
 ///
 /// These are "lazy" in the sense that we don't decode them (or even validate
@@ -1620,6 +1646,7 @@ mod tests {
                 stats: None,
                 ts_rewrite: None,
                 diffs_sum: None,
+                format: None,
             })],
             4,
             vec![],
@@ -1642,6 +1669,7 @@ mod tests {
             stats: None,
             ts_rewrite: None,
             diffs_sum: None,
+            format: None,
         }));
         assert_eq!(<HollowBatch<u64>>::from_proto(old).unwrap(), expected);
     }

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

@@ -16,6 +16,7 @@ syntax = "proto3";
 
 package mz_persist_client.internal.state;
 
+import "google/protobuf/empty.proto";
 import "persist/src/persist.proto";
 
 message ProtoU64Antichain {
@@ -39,6 +40,11 @@ message ProtoHollowBatchPart {
     uint64 encoded_size_bytes = 2;
     bytes key_lower = 3;
     optional int64 diffs_sum = 6;
+    oneof format {
+        google.protobuf.Empty row = 7;
+        uint64 row_and_columnar = 8;
+    }
+
     optional bytes key_stats = 536870906;
 
     reserved 536870907 to 536870911;

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

@@ -21,6 +21,7 @@ use mz_dyncfg::Config;
 use mz_ore::cast::CastFrom;
 use mz_ore::now::EpochMillis;
 use mz_ore::vec::PartialOrdVecExt;
+use mz_persist::indexed::encoding::BatchColumnarFormat;
 use mz_persist::location::SeqNo;
 use mz_persist_types::{Codec, Codec64, Opaque};
 use proptest_derive::Arbitrary;
@@ -318,6 +319,13 @@ pub struct HollowBatchPart<T> {
     /// we later decide that's of some benefit.
     #[serde(serialize_with = "serialize_diffs_sum")]
     pub diffs_sum: Option<[u8; 8]>,
+    /// Columnar format that this batch was written in.
+    ///
+    /// This is `None` if this part was written before we started writing structured
+    /// columnar data, or if the [`BATCH_RECORD_PART_FORMAT`] dyncfg is off.
+    ///
+    /// [`BATCH_RECORD_PART_FORMAT`]: crate::batch::BATCH_RECORD_PART_FORMAT
+    pub format: Option<BatchColumnarFormat>,
 }
 
 /// A [Batch] but with the updates themselves stored externally.
@@ -584,6 +592,7 @@ impl<T: Ord> Ord for HollowBatchPart<T> {
             stats: self_stats,
             ts_rewrite: self_ts_rewrite,
             diffs_sum: self_diffs_sum,
+            format: self_format,
         } = self;
         let HollowBatchPart {
             key: other_key,
@@ -592,6 +601,7 @@ impl<T: Ord> Ord for HollowBatchPart<T> {
             stats: other_stats,
             ts_rewrite: other_ts_rewrite,
             diffs_sum: other_diffs_sum,
+            format: other_format,
         } = other;
         (
             self_key,
@@ -600,6 +610,7 @@ impl<T: Ord> Ord for HollowBatchPart<T> {
             self_stats,
             self_ts_rewrite.as_ref().map(|x| x.elements()),
             self_diffs_sum,
+            self_format,
         )
             .cmp(&(
                 other_key,
@@ -608,6 +619,7 @@ impl<T: Ord> Ord for HollowBatchPart<T> {
                 other_stats,
                 other_ts_rewrite.as_ref().map(|x| x.elements()),
                 other_diffs_sum,
+                other_format,
             ))
     }
 }
@@ -1926,14 +1938,18 @@ pub(crate) mod tests {
                 any_some_lazy_part_stats(),
                 any::<Option<T>>(),
                 any::<[u8; 8]>(),
+                any::<Option<BatchColumnarFormat>>(),
             ),
-            |(key, encoded_size_bytes, key_lower, stats, ts_rewrite, diffs_sum)| HollowBatchPart {
-                key,
-                encoded_size_bytes,
-                key_lower,
-                stats,
-                ts_rewrite: ts_rewrite.map(Antichain::from_elem),
-                diffs_sum: Some(diffs_sum),
+            |(key, encoded_size_bytes, key_lower, stats, ts_rewrite, diffs_sum, format)| {
+                HollowBatchPart {
+                    key,
+                    encoded_size_bytes,
+                    key_lower,
+                    stats,
+                    ts_rewrite: ts_rewrite.map(Antichain::from_elem),
+                    diffs_sum: Some(diffs_sum),
+                    format,
+                }
             },
         )
     }
@@ -2083,6 +2099,7 @@ pub(crate) mod tests {
                         stats: None,
                         ts_rewrite: None,
                         diffs_sum: None,
+                        format: None,
                     })
                 })
                 .collect(),