feat(hash agg): do state cleaning by watermarks in hash agg (close #6112)

Author: Liang Zhao

src/storage/src/table/streaming_table/state_table.rs

@@ -27,7 +27,7 @@ use itertools::{izip, Itertools};
 use risingwave_common::array::{Op, Row, RowDeserializer, StreamChunk, Vis};
 use risingwave_common::buffer::Bitmap;
 use risingwave_common::catalog::{ColumnDesc, TableId, TableOption};
-use risingwave_common::types::VirtualNode;
+use risingwave_common::types::{ScalarImpl, VirtualNode};
 use risingwave_common::util::epoch::EpochPair;
 use risingwave_common::util::ordered::OrderedRowSerde;
 use risingwave_common::util::sort_util::OrderType;
@@ -583,18 +583,25 @@ impl<S: StateStore> StateTable<S> {
         self.epoch = Some(new_epoch);
     }
 
-    pub async fn commit(&mut self, new_epoch: EpochPair) -> StorageResult<()> {
+    pub async fn commit(
+        &mut self,
+        new_epoch: EpochPair,
+        watermark: Option<&ScalarImpl>,
+    ) -> StorageResult<()> {
         assert_eq!(self.epoch(), new_epoch.prev);
         let mem_table = std::mem::take(&mut self.mem_table).into_parts();
-        self.batch_write_rows(mem_table, new_epoch.prev).await?;
+        self.batch_write_rows(mem_table, new_epoch.prev, watermark)
+            .await?;
         self.update_epoch(new_epoch);
         Ok(())
     }
 
     /// used for unit test, and do not need to assert epoch.
+    #[cfg(any(test, feature = "test"))]
     pub async fn commit_for_test(&mut self, new_epoch: EpochPair) -> StorageResult<()> {
         let mem_table = std::mem::take(&mut self.mem_table).into_parts();
-        self.batch_write_rows(mem_table, new_epoch.prev).await?;
+        self.batch_write_rows(mem_table, new_epoch.prev, None)
+            .await?;
         self.update_epoch(new_epoch);
         Ok(())
     }
@@ -612,6 +619,7 @@ impl<S: StateStore> StateTable<S> {
         &mut self,
         buffer: BTreeMap<Vec<u8>, RowOp>,
         epoch: u64,
+        watermark: Option<&ScalarImpl>,
     ) -> StorageResult<()> {
         let mut write_batch = self.keyspace.start_write_batch(WriteOptions {
             epoch,
@@ -643,6 +651,17 @@ impl<S: StateStore> StateTable<S> {
                 }
             }
         }
+        if let Some(watermark) = watermark {
+            let prefix_serializer = self.pk_serde.prefix(1);
+            let encoded_prefix =
+                serialize_pk(&Row::new(vec![Some(watermark.clone())]), &prefix_serializer);
+            for vnode in self.vnodes.ones() {
+                let vnode_bytes = vnode.to_be_bytes().to_vec();
+                let mut range_end = vnode_bytes.clone();
+                range_end.extend(&encoded_prefix);
+                write_batch.delete_range(vnode_bytes, range_end);
+            }
+        }
         write_batch.ingest().await?;
         Ok(())
     }

src/stream/src/executor/dynamic_filter.rs

@@ -390,7 +390,7 @@ impl<S: StateStore> DynamicFilterExecutor<S> {
                                 self.right_table.insert(row.clone());
                                 last_committed_epoch_row = Some(row.clone());
                             }
-                            self.right_table.commit(barrier.epoch).await?;
+                            self.right_table.commit(barrier.epoch, None).await?;
                         } else {
                             self.right_table.commit_no_data_expected(barrier.epoch);
                         }

src/stream/src/executor/global_simple_agg.rs

@@ -214,7 +214,7 @@ impl<S: StateStore> GlobalSimpleAggExecutor<S> {
 
             // Commit all state tables except for result table.
             futures::future::try_join_all(
-                iter_table_storage(storages).map(|state_table| state_table.commit(epoch)),
+                iter_table_storage(storages).map(|state_table| state_table.commit(epoch, None)),
             )
             .await?;
 
@@ -248,7 +248,7 @@ impl<S: StateStore> GlobalSimpleAggExecutor<S> {
             } else {
                 result_table.insert(result_row);
             }
-            result_table.commit(epoch).await?;
+            result_table.commit(epoch, None).await?;
 
             let columns = builders
                 .into_iter()

src/stream/src/executor/hash_agg.rs

@@ -395,13 +395,12 @@ impl<K: HashKey, S: StateStore> HashAggExecutor<K, S> {
             ref total_lookup_count,
             ref metrics,
             ref chunk_size,
-            buffered_watermarks: ref _buffered_watermarks,
+            ref buffered_watermarks,
             ..
         }: &'a mut HashAggExecutorExtra<K, S>,
         agg_groups: &'a mut AggGroupMap<K, S>,
         epoch: EpochPair,
     ) {
-        // TODO("https://github.com/risingwavelabs/risingwave/issues/6112"): use buffered_watermarks[0] do some state cleaning
         let actor_id_str = ctx.id.to_string();
         metrics
             .agg_lookup_miss_count
@@ -481,12 +480,17 @@ impl<K: HashKey, S: StateStore> HashAggExecutor<K, S> {
                 yield chunk;
             }
 
+            let state_clean_watermark = buffered_watermarks
+                .first()
+                .and_then(|opt_watermark| opt_watermark.as_ref().map(|watermark| &watermark.val));
+
             // Commit all state tables.
             futures::future::try_join_all(
-                iter_table_storage(storages).map(|state_table| state_table.commit(epoch)),
+                iter_table_storage(storages)
+                    .map(|state_table| state_table.commit(epoch, state_clean_watermark)),
             )
             .await?;
-            result_table.commit(epoch).await?;
+            result_table.commit(epoch, state_clean_watermark).await?;
 
             // Evict cache to target capacity.
             agg_groups.evict();

src/stream/src/executor/managed_state/dynamic_filter.rs

@@ -203,7 +203,7 @@ impl<S: StateStore> RangeCache<S> {
     /// Flush writes to the `StateTable` from the in-memory buffer.
     pub async fn flush(&mut self, epoch: EpochPair) -> StreamExecutorResult<()> {
         // self.metrics.flush();
-        self.state_table.commit(epoch).await?;
+        self.state_table.commit(epoch, None).await?;
         Ok(())
     }
 }

src/stream/src/executor/managed_state/join/mod.rs

@@ -477,8 +477,8 @@ impl<K: HashKey, S: StateStore> JoinHashMap<K, S> {
 
     pub async fn flush(&mut self, epoch: EpochPair) -> StreamExecutorResult<()> {
         self.metrics.flush();
-        self.state.table.commit(epoch).await?;
-        self.degree_state.table.commit(epoch).await?;
+        self.state.table.commit(epoch, None).await?;
+        self.degree_state.table.commit(epoch, None).await?;
         Ok(())
     }
 

src/stream/src/executor/managed_state/top_n/top_n_state.rs

@@ -280,7 +280,7 @@ impl<S: StateStore> ManagedTopNState<S> {
     }
 
     pub async fn flush(&mut self, epoch: EpochPair) -> StreamExecutorResult<()> {
-        self.state_table.commit(epoch).await?;
+        self.state_table.commit(epoch, None).await?;
         Ok(())
     }
 }

src/stream/src/executor/mview/materialize.rs

@@ -143,7 +143,7 @@ impl<S: StateStore> MaterializeExecutor<S> {
                     Message::Chunk(chunk)
                 }
                 Message::Barrier(b) => {
-                    self.state_table.commit(b.epoch).await?;
+                    self.state_table.commit(b.epoch, None).await?;
 
                     // Update the vnode bitmap for the state table if asked.
                     if let Some(vnode_bitmap) = b.as_update_vnode_bitmap(self.actor_context.id) {

src/stream/src/executor/sort.rs

@@ -216,7 +216,7 @@ impl<S: StateStore> SortExecutor<S> {
                             }
                         }
                         // Commit the epoch.
-                        self.state_table.commit(barrier.epoch).await?;
+                        self.state_table.commit(barrier.epoch, None).await?;
                     } else {
                         // If the barrier is not a checkpoint, then there is no actual data to
                         // commit. Therefore, we simply update the epoch of state table.

src/stream/src/executor/source/source_executor.rs

@@ -215,7 +215,10 @@ impl<S: StateStore> SourceExecutor<S> {
             self.split_state_store.take_snapshot(cache).await?
         }
         // commit anyway, even if no message saved
-        self.split_state_store.state_store.commit(epoch).await?;
+        self.split_state_store
+            .state_store
+            .commit(epoch, None)
+            .await?;
 
         Ok(())
     }

src/stream/src/executor/source/state_table_handler.rs

@@ -176,7 +176,7 @@ pub(crate) mod tests {
 
         state_table.init_epoch(init_epoch);
         state_table.insert(Row::new(vec![a.clone(), b.clone()]));
-        state_table.commit(next_epoch).await.unwrap();
+        state_table.commit(next_epoch, None).await.unwrap();
 
         let a: Arc<str> = String::from("a").into();
         let a: Datum = Some(ScalarImpl::Utf8(a.as_ref().into()));
@@ -202,9 +202,15 @@ pub(crate) mod tests {
         state_table_handler
             .take_snapshot(vec![split_impl.clone()])
             .await?;
-        state_table_handler.state_store.commit(epoch_2).await?;
+        state_table_handler
+            .state_store
+            .commit(epoch_2, None)
+            .await?;
 
-        state_table_handler.state_store.commit(epoch_3).await?;
+        state_table_handler
+            .state_store
+            .commit(epoch_3, None)
+            .await?;
 
         match state_table_handler
             .try_recover_from_state_store(&split_impl)

src/stream/src/executor/watermark_filter.rs

@@ -212,7 +212,7 @@ impl<S: StateStore> WatermarkFilterExecutor<S> {
                             // FIXME(yuhao): use upsert.
                             table.insert(row);
                         }
-                        table.commit(barrier.epoch).await?;
+                        table.commit(barrier.epoch, None).await?;
                     } else {
                         table.commit_no_data_expected(barrier.epoch);
                     }