feat: build Payload in encode_batch for LogBatchEncoder (apache#209)

analytic_engine/src/instance/alter.rs

@@ -117,12 +117,10 @@ impl Instance {
                     region_id: table_data.wal_region_id(),
                 })?;
 
-        let log_batch = log_batch_encoder
-            .encode(&[payload])
-            .context(EncodePayloads {
-                table: &table_data.name,
-                region_id: table_data.wal_region_id(),
-            })?;
+        let log_batch = log_batch_encoder.encode(&payload).context(EncodePayloads {
+            table: &table_data.name,
+            region_id: table_data.wal_region_id(),
+        })?;
 
         // Write log batch
         let write_ctx = WriteContext::default();
@@ -290,12 +288,10 @@ impl Instance {
                     region_id: table_data.wal_region_id(),
                 })?;
 
-        let log_batch = log_batch_encoder
-            .encode(&[payload])
-            .context(EncodePayloads {
-                table: &table_data.name,
-                region_id: table_data.wal_region_id(),
-            })?;
+        let log_batch = log_batch_encoder.encode(&payload).context(EncodePayloads {
+            table: &table_data.name,
+            region_id: table_data.wal_region_id(),
+        })?;
 
         // Write log batch
         let write_ctx = WriteContext::default();

analytic_engine/src/instance/write.rs

@@ -374,12 +374,10 @@ impl Instance {
                     table: &table_data.name,
                     region_id: table_data.wal_region_id(),
                 })?;
-        let log_batch = log_batch_encoder
-            .encode(&[payload])
-            .context(EncodePayloads {
-                table: &table_data.name,
-                region_id: table_data.wal_region_id(),
-            })?;
+        let log_batch = log_batch_encoder.encode(&payload).context(EncodePayloads {
+            table: &table_data.name,
+            region_id: table_data.wal_region_id(),
+        })?;
 
         // Write to wal manager
         let write_ctx = WriteContext::default();

analytic_engine/src/meta/details.rs

@@ -191,7 +191,7 @@ impl ManifestImpl {
             .encoder(region_id)
             .context(GetLogBatchEncoder { region_id })?;
         let log_batch = log_batch_encoder
-            .encode(&[payload])
+            .encode(&payload)
             .context(EncodePayloads { region_id })?;
 
         let write_ctx = WriteContext::default();
@@ -331,21 +331,13 @@ impl MetaUpdateLogStore for RegionWal {
     }
 
     async fn store(&self, log_entries: &[MetaUpdateLogEntry]) -> Result<()> {
-        let mut payload_batch = Vec::with_capacity(log_entries.len());
-
-        // TODO(ygf11): maybe we can build payload in encode loop.
-        for entry in log_entries {
-            let payload = MetaUpdatePayload::from(entry);
-            payload_batch.push(payload);
-        }
-
         let region_id = self.region_id;
         let log_batch_encoder = self
             .wal_manager
             .encoder(self.region_id)
             .context(GetLogBatchEncoder { region_id })?;
         let log_batch = log_batch_encoder
-            .encode(&payload_batch)
+            .encode_batch::<MetaUpdatePayload, MetaUpdateLogEntry>(log_entries)
             .context(EncodePayloads { region_id })?;
 
         let write_ctx = WriteContext::default();

analytic_engine/src/payload.rs

@@ -131,6 +131,12 @@ impl<'a> Payload for WritePayload<'a> {
     }
 }
 
+impl<'a> From<&'a table_requests::WriteRequest> for WritePayload<'a> {
+    fn from(write_request: &'a table_requests::WriteRequest) -> Self {
+        Self::Write(write_request)
+    }
+}
+
 /// Payload decoded from wal
 #[derive(Debug)]
 pub enum ReadPayload {

benchmarks/src/util.rs

@@ -192,3 +192,9 @@ impl<'a> Payload for WritePayload<'a> {
         Ok(())
     }
 }
+
+impl<'a> From<&'a Vec<u8>> for WritePayload<'a> {
+    fn from(data: &'a Vec<u8>) -> Self {
+        Self(data)
+    }
+}

benchmarks/src/wal_write_bench.rs

@@ -74,16 +74,11 @@ impl WalWriteBench {
             .expect("should succeed to open WalNamespaceImpl(Memory)");
 
             let values = self.build_value_vec();
-            let payloads = values
-                .iter()
-                .map(|value| WritePayload(value))
-                .collect::<Vec<_>>();
-
             let wal_encoder = wal
                 .encoder(1)
                 .expect("should succeed to create wal encoder");
             let log_batch = wal_encoder
-                .encode(&payloads)
+                .encode_batch::<WritePayload, Vec<u8>>(values.as_slice())
                 .expect("should succeed to encode payload batch");
 
             // Write to wal manager

wal/src/kv_encoder.rs

@@ -527,19 +527,45 @@ pub struct LogBatchEncoder {
 }
 
 impl LogBatchEncoder {
+    /// Create LogBatchEncoder with specific region_id.
     pub fn create(region_id: RegionId) -> Self {
         Self {
             region_id,
             log_encoding: LogEncoding::newest(),
         }
     }
 
-    pub fn encode<P: Payload>(self, payload_batch: &[P]) -> manager::Result<LogWriteBatch> {
+    /// Consume LogBatchEncoder and encode single payload to LogWriteBatch.
+    pub fn encode(self, payload: &impl Payload) -> manager::Result<LogWriteBatch> {
         let mut write_batch = LogWriteBatch::new(self.region_id);
         let mut buf = BytesMut::new();
-        for payload in payload_batch.iter() {
+        self.log_encoding
+            .encode_value(&mut buf, payload)
+            .map_err(|e| Box::new(e) as _)
+            .context(Encoding)?;
+
+        write_batch.push(LogWriteEntry {
+            payload: buf.to_vec(),
+        });
+
+        Ok(write_batch)
+    }
+
+    /// Consume LogBatchEncoder and encode raw payload batch to LogWriteBatch.
+    /// Note: To build payload from raw payload in `encode_batch`, raw payload
+    /// need implement From trait.
+    pub fn encode_batch<'a, P: Payload, I>(
+        self,
+        raw_payload_batch: &'a [I],
+    ) -> manager::Result<LogWriteBatch>
+    where
+        &'a I: Into<P>,
+    {
+        let mut write_batch = LogWriteBatch::new(self.region_id);
+        let mut buf = BytesMut::new();
+        for raw_payload in raw_payload_batch.iter() {
             self.log_encoding
-                .encode_value(&mut buf, payload)
+                .encode_value(&mut buf, &raw_payload.into())
                 .map_err(|e| Box::new(e) as _)
                 .context(Encoding)?;
 

wal/src/table_kv_impl/namespace.rs

@@ -1699,9 +1699,10 @@ mod tests {
             payload_batch.push(payload);
         }
 
+        let log_entries = (start_sequence..end_sequence).collect::<Vec<_>>();
         let wal_encoder = LogBatchEncoder::create(region_id);
         let log_batch = wal_encoder
-            .encode(&payload_batch)
+            .encode_batch::<TestPayload, u32>(&log_entries)
             .expect("should succeed to encode payload batch");
         let write_ctx = manager::WriteContext::default();
         namespace

wal/src/tests/util.rs

@@ -148,16 +148,17 @@ impl<B: WalBuilder> TestEnv<B> {
         start: u32,
         end: u32,
     ) -> (Vec<TestPayload>, LogWriteBatch) {
-        let payload_batch = self.build_payload_batch(start, end);
+        let log_entries = (start..end).collect::<Vec<_>>();
 
         let log_batch_encoder = wal
             .encoder(region_id)
             .expect("should succeed to create log batch encoder");
 
         let log_batch = log_batch_encoder
-            .encode(&payload_batch)
+            .encode_batch::<TestPayload, u32>(&log_entries)
             .expect("should succeed to encode payloads");
 
+        let payload_batch = self.build_payload_batch(start, end);
         (payload_batch, log_batch)
     }
 
@@ -219,6 +220,12 @@ impl Payload for TestPayload {
     }
 }
 
+impl From<&u32> for TestPayload {
+    fn from(v: &u32) -> Self {
+        Self { val: *v }
+    }
+}
+
 pub struct TestPayloadDecoder;
 
 impl PayloadDecoder for TestPayloadDecoder {