feat(stream): lazy emit for HashAggExecutor

src/compute/tests/integration_tests.rs

@@ -427,7 +427,7 @@ async fn test_row_seq_scan() -> Result<()> {
         ColumnDesc::unnamed(ColumnId::from(2), schema[2].data_type.clone()),
     ];
 
-    let mut state = StateTable::new_without_distribution(
+    let mut state = StateTable::<_>::new_without_distribution(
         memory_state_store.clone(),
         TableId::from(0x42),
         column_descs.clone(),

src/stream/src/common/table/mod.rs

@@ -13,7 +13,7 @@
 // limitations under the License.
 
 pub mod state_table;
-mod watermark;
+pub mod watermark;
 
 #[cfg(test)]
 pub mod test_state_table;

src/stream/src/common/table/state_table.rs

@@ -51,18 +51,15 @@ use risingwave_storage::table::{compute_chunk_vnode, compute_vnode, Distribution
 use risingwave_storage::StateStore;
 use tracing::trace;
 
-use super::watermark::{WatermarkBufferByEpoch, WatermarkBufferStrategy};
+use super::watermark::{WatermarkBufferStrategy, WatermarkBufferStrategyByEpochDefault};
 use crate::executor::{StreamExecutorError, StreamExecutorResult};
 
-/// This num is arbitrary and we may want to improve this choice in the future.
-const STATE_CLEANING_PERIOD_EPOCH: usize = 5;
-
 /// `StateTable` is the interface accessing relational data in KV(`StateStore`) with
 /// row-based encoding.
 #[derive(Clone)]
 pub struct StateTable<
     S: StateStore,
-    W: WatermarkBufferStrategy = WatermarkBufferByEpoch<STATE_CLEANING_PERIOD_EPOCH>,
+    W: WatermarkBufferStrategy = WatermarkBufferStrategyByEpochDefault,
 > {
     /// Id for this table.
     table_id: TableId,
@@ -508,7 +505,7 @@ impl<S: StateStore, W: WatermarkBufferStrategy> StateTable<S, W> {
 const ENABLE_SANITY_CHECK: bool = cfg!(debug_assertions);
 
 // point get
-impl<S: StateStore> StateTable<S> {
+impl<S: StateStore, W: WatermarkBufferStrategy> StateTable<S, W> {
     /// Get a single row from state table.
     pub async fn get_row(&self, pk: impl Row) -> StreamExecutorResult<Option<OwnedRow>> {
         let compacted_row: Option<CompactedRow> = self.get_compacted_row(pk).await?;
@@ -600,7 +597,7 @@ impl<S: StateStore> StateTable<S> {
 }
 
 // write
-impl<S: StateStore> StateTable<S> {
+impl<S: StateStore, W: WatermarkBufferStrategy> StateTable<S, W> {
     #[expect(clippy::boxed_local)]
     fn handle_mem_table_error(&self, e: Box<MemTableError>) {
         match *e {

src/stream/src/common/table/test_state_table.rs

@@ -50,7 +50,7 @@ async fn test_state_table_update_insert() {
 
     test_env.register_table(table.clone()).await;
     let mut state_table =
-        StateTable::from_table_catalog_no_sanity_check(&table, test_env.storage.clone(), None)
+        StateTable::<_>::from_table_catalog_no_sanity_check(&table, test_env.storage.clone(), None)
             .await;
 
     let mut epoch = EpochPair::new_test_epoch(1);
@@ -228,7 +228,7 @@ async fn test_state_table_iter_with_prefix() {
 
     test_env.register_table(table.clone()).await;
     let mut state_table =
-        StateTable::from_table_catalog_no_sanity_check(&table, test_env.storage.clone(), None)
+        StateTable::<_>::from_table_catalog_no_sanity_check(&table, test_env.storage.clone(), None)
             .await;
 
     let mut epoch = EpochPair::new_test_epoch(1);
@@ -353,7 +353,7 @@ async fn test_state_table_iter_with_pk_range() {
 
     test_env.register_table(table.clone()).await;
     let mut state_table =
-        StateTable::from_table_catalog_no_sanity_check(&table, test_env.storage.clone(), None)
+        StateTable::<_>::from_table_catalog_no_sanity_check(&table, test_env.storage.clone(), None)
             .await;
 
     let mut epoch = EpochPair::new_test_epoch(1);
@@ -487,7 +487,7 @@ async fn test_mem_table_assertion() {
 
     test_env.register_table(table.clone()).await;
     let mut state_table =
-        StateTable::from_table_catalog_no_sanity_check(&table, test_env.storage.clone(), None)
+        StateTable::<_>::from_table_catalog_no_sanity_check(&table, test_env.storage.clone(), None)
             .await;
 
     let epoch = EpochPair::new_test_epoch(1);
@@ -530,7 +530,7 @@ async fn test_state_table_iter_with_value_indices() {
 
     test_env.register_table(table.clone()).await;
     let mut state_table =
-        StateTable::from_table_catalog_no_sanity_check(&table, test_env.storage.clone(), None)
+        StateTable::<_>::from_table_catalog_no_sanity_check(&table, test_env.storage.clone(), None)
             .await;
 
     let mut epoch = EpochPair::new_test_epoch(1);
@@ -691,7 +691,7 @@ async fn test_state_table_iter_with_shuffle_value_indices() {
 
     test_env.register_table(table.clone()).await;
     let mut state_table =
-        StateTable::from_table_catalog_no_sanity_check(&table, test_env.storage.clone(), None)
+        StateTable::<_>::from_table_catalog_no_sanity_check(&table, test_env.storage.clone(), None)
             .await;
 
     let mut epoch = EpochPair::new_test_epoch(1);
@@ -926,7 +926,7 @@ async fn test_state_table_write_chunk() {
 
     test_env.register_table(table.clone()).await;
     let mut state_table =
-        StateTable::from_table_catalog_no_sanity_check(&table, test_env.storage.clone(), None)
+        StateTable::<_>::from_table_catalog_no_sanity_check(&table, test_env.storage.clone(), None)
             .await;
 
     let epoch = EpochPair::new_test_epoch(1);
@@ -1055,7 +1055,7 @@ async fn test_state_table_write_chunk_visibility() {
 
     test_env.register_table(table.clone()).await;
     let mut state_table =
-        StateTable::from_table_catalog_no_sanity_check(&table, test_env.storage.clone(), None)
+        StateTable::<_>::from_table_catalog_no_sanity_check(&table, test_env.storage.clone(), None)
             .await;
 
     let epoch = EpochPair::new_test_epoch(1);
@@ -1182,7 +1182,7 @@ async fn test_state_table_write_chunk_value_indices() {
 
     test_env.register_table(table.clone()).await;
     let mut state_table =
-        StateTable::from_table_catalog_no_sanity_check(&table, test_env.storage.clone(), None)
+        StateTable::<_>::from_table_catalog_no_sanity_check(&table, test_env.storage.clone(), None)
             .await;
 
     let epoch = EpochPair::new_test_epoch(1);

src/stream/src/common/table/test_storage_table.rs

@@ -64,7 +64,7 @@ async fn test_storage_table_value_indices() {
 
     test_env.register_table(table.clone()).await;
     let mut state =
-        StateTable::from_table_catalog_no_sanity_check(&table, test_env.storage.clone(), None)
+        StateTable::<_>::from_table_catalog_no_sanity_check(&table, test_env.storage.clone(), None)
             .await;
 
     let table = StorageTable::for_test(
@@ -192,7 +192,7 @@ async fn test_shuffled_column_id_for_storage_table_get_row() {
 
     test_env.register_table(table.clone()).await;
     let mut state =
-        StateTable::from_table_catalog_no_sanity_check(&table, test_env.storage.clone(), None)
+        StateTable::<_>::from_table_catalog_no_sanity_check(&table, test_env.storage.clone(), None)
             .await;
 
     let mut epoch = EpochPair::new_test_epoch(1);
@@ -295,7 +295,7 @@ async fn test_row_based_storage_table_point_get_in_batch_mode() {
 
     test_env.register_table(table.clone()).await;
     let mut state =
-        StateTable::from_table_catalog_no_sanity_check(&table, test_env.storage.clone(), None)
+        StateTable::<_>::from_table_catalog_no_sanity_check(&table, test_env.storage.clone(), None)
             .await;
 
     let column_ids_partial = vec![ColumnId::from(1), ColumnId::from(2)];
@@ -402,7 +402,7 @@ async fn test_batch_scan_with_value_indices() {
 
     test_env.register_table(table.clone()).await;
     let mut state =
-        StateTable::from_table_catalog_no_sanity_check(&table, test_env.storage.clone(), None)
+        StateTable::<_>::from_table_catalog_no_sanity_check(&table, test_env.storage.clone(), None)
             .await;
 
     let column_ids_partial = vec![ColumnId::from(1), ColumnId::from(2)];

src/stream/src/common/table/watermark.rs

@@ -13,7 +13,7 @@
 // limitations under the License.
 
 /// Strategy to decide how to buffer the watermarks, used for state cleaning.
-pub trait WatermarkBufferStrategy: Default {
+pub trait WatermarkBufferStrategy: Default + Send + Sync + 'static {
     /// Trigger when a epoch is committed.
     fn tick(&mut self);
 
@@ -58,3 +58,9 @@ impl<const PERIOD: usize> WatermarkBufferStrategy for WatermarkBufferByEpoch<PER
         }
     }
 }
+
+/// This num is arbitrary and we may want to improve this choice in the future.
+const STATE_CLEANING_PERIOD_EPOCH: usize = 5;
+
+pub type WatermarkBufferStrategyByEpochDefault =
+    WatermarkBufferByEpoch<STATE_CLEANING_PERIOD_EPOCH>;

src/stream/src/executor/aggregation/agg_group.rs

@@ -27,6 +27,7 @@ use risingwave_storage::StateStore;
 use super::agg_state::{AggState, AggStateStorage};
 use super::AggCall;
 use crate::common::table::state_table::StateTable;
+use crate::common::table::watermark::WatermarkBufferStrategy;
 use crate::executor::error::StreamExecutorResult;
 use crate::executor::PkIndices;
 
@@ -47,7 +48,7 @@ impl<S: StateStore> Debug for AggGroup<S> {
         f.debug_struct("AggGroup")
             .field("group_key", &self.group_key)
             .field("prev_outputs", &self.prev_outputs)
-            .finish()
+            .finish_non_exhaustive()
     }
 }
 
@@ -67,11 +68,11 @@ pub struct AggChangesInfo {
 impl<S: StateStore> AggGroup<S> {
     /// Create [`AggGroup`] for the given [`AggCall`]s and `group_key`.
     /// For [`crate::executor::GlobalSimpleAggExecutor`], the `group_key` should be `None`.
-    pub async fn create(
+    pub async fn create<W: WatermarkBufferStrategy>(
         group_key: Option<OwnedRow>,
         agg_calls: &[AggCall],
         storages: &[AggStateStorage<S>],
-        result_table: &StateTable<S>,
+        result_table: &StateTable<S, W>,
         pk_indices: &PkIndices,
         extreme_cache_size: usize,
         input_schema: &Schema,

src/stream/src/executor/aggregation/minput.rs

@@ -325,7 +325,7 @@ mod tests {
             .collect_vec();
         let mapping = StateTableColumnMapping::new(upstream_columns, None);
         let pk_len = order_types.len();
-        let table = StateTable::new_without_distribution(
+        let table = StateTable::<_>::new_without_distribution(
             MemoryStateStore::new(),
             table_id,
             columns,

src/stream/src/executor/aggregation/mod.rs

@@ -39,9 +39,11 @@ mod state_cache;
 mod table;
 mod value;
 
-/// Generate [`crate::executor::HashAggExecutor`]'s schema from `input`, `agg_calls` and
-/// `group_key_indices`. For [`crate::executor::HashAggExecutor`], the group key indices should
-/// be provided.
+/// Generate [`HashAggExecutor`][HashAggExecutor]'s schema from `input`, `agg_calls` and
+/// `group_key_indices`. For [`HashAggExecutor`][HashAggExecutor], the group key indices should be
+/// provided.
+///
+/// [HashAggExecutor]:crate::executor::HashAggExecutor
 pub fn generate_agg_schema(
     input: &dyn Executor,
     agg_calls: &[AggCall],

src/stream/src/executor/dynamic_filter.rs

@@ -492,15 +492,15 @@ mod tests {
     ) -> (StateTable<MemoryStateStore>, StateTable<MemoryStateStore>) {
         let column_descs = ColumnDesc::unnamed(ColumnId::new(0), DataType::Int64);
         // TODO: enable sanity check for dynamic filter <https://github.com/risingwavelabs/risingwave/issues/3893>
-        let state_table_l = StateTable::new_without_distribution_no_sanity_check(
+        let state_table_l = StateTable::<_>::new_without_distribution_no_sanity_check(
             mem_state.clone(),
             TableId::new(0),
             vec![column_descs.clone()],
             vec![OrderType::Ascending],
             vec![0],
         )
         .await;
-        let state_table_r = StateTable::new_without_distribution_no_sanity_check(
+        let state_table_r = StateTable::<_>::new_without_distribution_no_sanity_check(
             mem_state,
             TableId::new(1),
             vec![column_descs],