Repository: risingwavelabs/risingwave
File: src/stream/src/common/table/state_table.rs (Rust)

Code before:
        key_bytes: TableKey<Bytes>,
        prefix_hint: Option<Bytes>,
    ) -> StreamExecutorResult<Option<OwnedRow>> {
        if let Some(m) = &self.metrics {
            m.get_count.inc();
        }
        if let Some(rows) = &self.all_rows {
            let (vnode, key) = key_bytes.split_vnode_bytes();
            return Ok(rows.get(&vnode).expect("covered vnode").get(&key).cloned());
        }

        // Try to prune using vnode statistics
        let should_prune = if let Some(stats) = &self.vnode_stats
            && let (vnode, key) = key_bytes.split_vnode_bytes()
            && let Some(vnode_stat) = stats.get(&vnode)
            && vnode_stat.can_prune(&key)
        {
            if let Some(m) = &self.metrics {
                m.get_vnode_pruned_count.inc();
            }
            true
        } else {
            false
        };

        // In dry-run mode, we verify pruning correctness but don't apply it
        if should_prune && self.enable_state_table_vnode_stats_pruning {
            return Ok(None);
        }

        let read_options = ReadOptions {
            prefix_hint,
            cache_policy: CachePolicy::Fill(Hint::Normal),
            ..Default::default()
        };

        let result = self
            .state_store
            .on_key_value(key_bytes, read_options, move |_, value| {
                let row = self.row_serde.deserialize(value)?;
                Ok(OwnedRow::new(row))
            })
            .await
            .map_err(Into::<StreamExecutorError>::into)?;

        // In dry-run mode, verify that pruning would have been correct
        if should_prune && result.is_some() {
            tracing::warn!(
                table_id = %self.table_id,
                "vnode stats pruning dry run fails for get. This will not affect correctness."
            );

Review comment (refactor):
I'm quite confused why this branch has a dry-run document. It hits the prune logic and returns directly.


Code after:
        key_bytes: TableKey<Bytes>,
        prefix_hint: Option<Bytes>,
    ) -> StreamExecutorResult<Option<OwnedRow>> {
        if let Some(m) = &self.metrics {
            m.get_count.inc();
        }
        if let Some(rows) = &self.all_rows {
            let (vnode, key) = key_bytes.split_vnode_bytes();
            return Ok(rows.get(&vnode).expect("covered vnode").get(&key).cloned());
        }

        // Try to prune using vnode statistics
        let should_prune = if let Some(stats) = &self.vnode_stats
            && let (vnode, key) = key_bytes.split_vnode_bytes()
            && let Some(vnode_stat) = stats.get(&vnode)
            && vnode_stat.can_prune(&key)
        {
            if let Some(m) = &self.metrics {
                m.get_vnode_pruned_count.inc();
            }
            true
        } else {
            false
        };

        if should_prune && self.enable_state_table_vnode_stats_pruning {
            return Ok(None);
        }

        let read_options = ReadOptions {
            prefix_hint,
            cache_policy: CachePolicy::Fill(Hint::Normal),
            ..Default::default()
        };

        let result = self
            .state_store
            .on_key_value(key_bytes, read_options, move |_, value| {
                let row = self.row_serde.deserialize(value)?;
                Ok(OwnedRow::new(row))
            })
            .await
            .map_err(Into::<StreamExecutorError>::into)?;

        // In dry-run mode, verify that pruning would have been correct
        if should_prune && result.is_some() {
            tracing::warn!(
                table_id = %self.table_id,
                "vnode stats pruning dry run fails for get. This will not affect correctness."
            );
        }