Repository: risingwavelabs/risingwave
File: src/frontend/src/optimizer/plan_node/stream_group_topn.rs (Rust)

Code before:
use crate::optimizer::property::{MonotonicityMap, Order};
use crate::stream_fragmenter::BuildFragmentGraphState;

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct StreamGroupTopN {
    pub base: PlanBase<Stream>,
    core: generic::TopN<PlanRef>,
    /// an optional column index which is the vnode of each row computed by the input's consistent
    /// hash distribution
    vnode_col_idx: Option<usize>,
}

impl StreamGroupTopN {
    pub fn new(core: generic::TopN<PlanRef>, vnode_col_idx: Option<usize>) -> Result<Self> {
        assert!(!core.group_key.is_empty());
        assert!(core.limit_attr.limit() > 0);
        reject_upsert_input!(core.input);

        let input = &core.input;

        // FIXME(rc): Actually only watermark messages on the first group-by column are propagated
        // acccoring to the current GroupTopN implementation. This should be fixed.
        let watermark_columns = if input.append_only() {
            input.watermark_columns().clone()
        } else {
            input.watermark_columns().retain_clone(&core.group_key)
        };

        let mut stream_key = core
            .stream_key()
            .expect("logical node should have stream key here");
        if let Some(vnode_col_idx) = vnode_col_idx
            && stream_key.len() > 1
        {
            // The output stream key of `GroupTopN` is a union of group key and input stream key,
            // while vnode is calculated from a subset of input stream key. So we can safely remove
            // the vnode column from output stream key. While at meanwhile we cannot leave the stream key
            // as empty, so we only remove it when stream key length is > 1.
            stream_key.remove(stream_key.iter().position(|i| *i == vnode_col_idx).unwrap());
        }

        let base = PlanBase::new_stream(
            core.ctx(),
            core.schema(),
            Some(stream_key),
            core.functional_dependency(),
            input.distribution().clone(),
            StreamKind::Retract,
            // TODO: https://github.com/risingwavelabs/risingwave/issues/8348
            false,
            watermark_columns,

Review comment (performance):
The comment mentions forwarding watermarks on "ascending ORDER BY columns on append-only input" as a TODO. However, this comment could be clearer. The executor only checks if the watermark column index matches the first group-by column (group_by[0]), it doesn't check ORDER BY columns or whether input is append-only. Consider clarifying that the TODO refers to a potential future enhancement where the executor behavior would also be changed, not just the optimizer.
```suggestion
        // TODO: Conceptually it's also safe to forward watermarks on
        // - other group key columns, and
        // - ascending ORDER BY columns on append-only input,
        // but supporting these would require extending the executor behavior as well,
        // not just changing the optimizer. Until then, we conservatively forward only
        // the first group key watermark.
```

Code after:
use crate::optimizer::property::{MonotonicityMap, Order};
use crate::stream_fragmenter::BuildFragmentGraphState;

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct StreamGroupTopN {
    pub base: PlanBase<Stream>,
    core: generic::TopN<PlanRef>,
    /// an optional column index which is the vnode of each row computed by the input's consistent
    /// hash distribution
    vnode_col_idx: Option<usize>,
}

impl StreamGroupTopN {
    pub fn new(core: generic::TopN<PlanRef>, vnode_col_idx: Option<usize>) -> Result<Self> {
        assert!(!core.group_key.is_empty());
        assert!(core.limit_attr.limit() > 0);
        reject_upsert_input!(core.input);

        let input = &core.input;

        // NOTE: The executor only forwards watermarks on the first group-by column.
        // Keep the optimizer in sync so EOWC won't be enabled on unsupported plans.
        //
        // TODO: Actually it's also safe to forward watermarks on
        // - other group key columns,
        // - ascending ORDER BY columns on append-only input.
        let watermark_columns = input.watermark_columns().retain_clone(&[core.group_key[0]]);

        let mut stream_key = core
            .stream_key()
            .expect("logical node should have stream key here");
        if let Some(vnode_col_idx) = vnode_col_idx
            && stream_key.len() > 1
        {
            // The output stream key of `GroupTopN` is a union of group key and input stream key,
            // while vnode is calculated from a subset of input stream key. So we can safely remove
            // the vnode column from output stream key. While at meanwhile we cannot leave the stream key
            // as empty, so we only remove it when stream key length is > 1.
            stream_key.remove(stream_key.iter().position(|i| *i == vnode_col_idx).unwrap());
        }

        let base = PlanBase::new_stream(
            core.ctx(),
            core.schema(),
            Some(stream_key),
            core.functional_dependency(),
            input.distribution().clone(),
            StreamKind::Retract,
            // TODO: https://github.com/risingwavelabs/risingwave/issues/8348
            false,
            watermark_columns,