fix(stream): fix pk indices of GroupTopN executors (#12304)

Signed-off-by: Richard Chien <[email protected]>

e2e_test/streaming/group_top_n/group_top_1.slt

@@ -0,0 +1,45 @@
+statement ok
+SET RW_IMPLICIT_FLUSH TO true;
+
+# https:/risingwavelabs/risingwave/issues/12282
+
+statement ok
+create table t(a int, b int, c int);
+
+statement ok
+create materialized view mv as SELECT * FROM (
+ SELECT
+ *,
+ row_number() OVER (PARTITION BY a ORDER BY b) AS rank
+ FROM t
+) WHERE rank <= 1;
+
+statement ok
+insert into t values (1, 1, 1);
+
+query iiiI
+select * from mv;
+----
+1 1 1 1
+
+statement ok
+insert into t values (1, 0, 1);
+
+query iiiI
+select * from mv;
+----
+1 0 1 1
+
+statement ok
+insert into t values (1, 0, 1);
+
+query iiiI
+select * from mv;
+----
+1 0 1 1
+
+statement ok
+drop materialized view mv;
+
+statement ok
+drop table t;

src/stream/src/executor/top_n/group_top_n.rs

@@ -51,6 +51,7 @@ impl<K: HashKey, S: StateStore, const WITH_TIES: bool> GroupTopNExecutor<K, S, W
  group_by: Vec<usize>,
  state_table: StateTable<S>,
  watermark_epoch: AtomicU64Ref,
+ pk_indices: PkIndices,
  ) -> StreamResult<Self> {
  let info = input.info();
  Ok(TopNExecutorWrapper {
@@ -66,6 +67,7 @@ impl<K: HashKey, S: StateStore, const WITH_TIES: bool> GroupTopNExecutor<K, S, W
  state_table,
  watermark_epoch,
  ctx,
+ pk_indices,
  )?,
  })
  }
@@ -109,9 +111,11 @@ impl<K: HashKey, S: StateStore, const WITH_TIES: bool> InnerGroupTopNExecutor<K,
  state_table: StateTable<S>,
  watermark_epoch: AtomicU64Ref,
  ctx: ActorContextRef,
+ pk_indices: PkIndices,
  ) -> StreamResult<Self> {
  let ExecutorInfo {
- pk_indices, schema, ..
+ schema: input_schema,
+ ..
  } = input_info;
 
  let metrics_info = MetricsInfo::new(
@@ -121,12 +125,13 @@ impl<K: HashKey, S: StateStore, const WITH_TIES: bool> InnerGroupTopNExecutor<K,
  "GroupTopN",
  );
 
- let cache_key_serde = create_cache_key_serde(&storage_key, &schema, &order_by, &group_by);
+ let cache_key_serde =
+ create_cache_key_serde(&storage_key, &input_schema, &order_by, &group_by);
  let managed_state = ManagedTopNState::<S>::new(state_table, cache_key_serde.clone());
 
  Ok(Self {
  info: ExecutorInfo {
- schema,
+ schema: input_schema,
  pk_indices,
  identity: format!("GroupTopNExecutor {:X}", executor_id),
  },
@@ -408,6 +413,7 @@ mod tests {
  vec![1],
  state_table,
  Arc::new(AtomicU64::new(0)),
+ pk_indices(),
  )
  .unwrap();
  let top_n_executor = Box::new(a);
@@ -505,6 +511,7 @@ mod tests {
  vec![1],
  state_table,
  Arc::new(AtomicU64::new(0)),
+ pk_indices(),
  )
  .unwrap(),
  );
@@ -595,6 +602,7 @@ mod tests {
  vec![1, 2],
  state_table,
  Arc::new(AtomicU64::new(0)),
+ pk_indices(),
  )
  .unwrap(),
  );

src/stream/src/executor/top_n/group_top_n_appendonly.rs

@@ -69,6 +69,7 @@ impl<K: HashKey, S: StateStore, const WITH_TIES: bool>
  group_by: Vec<usize>,
  state_table: StateTable<S>,
  watermark_epoch: AtomicU64Ref,
+ pk_indices: PkIndices,
  ) -> StreamResult<Self> {
  let info = input.info();
  Ok(TopNExecutorWrapper {
@@ -84,6 +85,7 @@ impl<K: HashKey, S: StateStore, const WITH_TIES: bool>
  state_table,
  watermark_epoch,
  ctx,
+ pk_indices,
  )?,
  })
  }
@@ -129,9 +131,11 @@ impl<K: HashKey, S: StateStore, const WITH_TIES: bool>
  state_table: StateTable<S>,
  watermark_epoch: AtomicU64Ref,
  ctx: ActorContextRef,
+ pk_indices: PkIndices,
  ) -> StreamResult<Self> {
  let ExecutorInfo {
- pk_indices, schema, ..
+ schema: input_schema,
+ ..
  } = input_info;
 
  let metrics_info = MetricsInfo::new(
@@ -141,12 +145,13 @@ impl<K: HashKey, S: StateStore, const WITH_TIES: bool>
  "GroupTopN",
  );
 
- let cache_key_serde = create_cache_key_serde(&storage_key, &schema, &order_by, &group_by);
+ let cache_key_serde =
+ create_cache_key_serde(&storage_key, &input_schema, &order_by, &group_by);
  let managed_state = ManagedTopNState::<S>::new(state_table, cache_key_serde.clone());
 
  Ok(Self {
  info: ExecutorInfo {
- schema,
+ schema: input_schema,
  pk_indices,
  identity: format!("AppendOnlyGroupTopNExecutor {:X}", executor_id),
  },

src/stream/src/from_proto/group_top_n.rs

@@ -21,7 +21,7 @@ use risingwave_pb::stream_plan::GroupTopNNode;
 
 use super::*;
 use crate::common::table::state_table::StateTable;
-use crate::executor::{ActorContextRef, AppendOnlyGroupTopNExecutor, GroupTopNExecutor};
+use crate::executor::{ActorContextRef, AppendOnlyGroupTopNExecutor, GroupTopNExecutor, PkIndices};
 use crate::task::AtomicU64Ref;
 
 pub struct GroupTopNExecutorBuilder<const APPEND_ONLY: bool>;
@@ -80,6 +80,7 @@ impl<const APPEND_ONLY: bool> ExecutorBuilder for GroupTopNExecutorBuilder<APPEN
  state_table,
  watermark_epoch: stream.get_watermark_epoch(),
  group_key_types,
+ pk_indices: params.pk_indices,
 
  with_ties: node.with_ties,
  append_only: APPEND_ONLY,
@@ -99,6 +100,7 @@ struct GroupTopNExecutorDispatcherArgs<S: StateStore> {
  state_table: StateTable<S>,
  watermark_epoch: AtomicU64Ref,
  group_key_types: Vec<DataType>,
+ pk_indices: PkIndices,
 
  with_ties: bool,
  append_only: bool,
@@ -120,6 +122,7 @@ impl<S: StateStore> HashKeyDispatcher for GroupTopNExecutorDispatcherArgs<S> {
  self.group_by,
  self.state_table,
  self.watermark_epoch,
+ self.pk_indices,
  )?
  .boxed())
  };