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

Code before:
    let table = gen_pbtable(
        TEST_TABLE_ID,
        column_descs.clone(),
        order_types.clone(),
        pk_index.clone(),
        read_prefix_len_hint,
    );

    test_env.register_table(table.clone()).await;

    let registry = Registry::new();
    let metrics = StreamingMetrics::new(&registry, MetricLevel::Debug).new_state_table_metrics(
        TEST_TABLE_ID,
        ActorId::new(2),
        FragmentId::new(2),
    );

    let vnode_bitmap = {
        let mut builder = BitmapBuilder::zeroed(VirtualNode::COUNT_FOR_TEST);
        for i in 0..VirtualNode::COUNT_FOR_TEST {
            builder.set(i, true);
        }
        builder.finish()
    };
    let mut config = StreamingConfig::default();
    config.developer.enable_state_table_vnode_stats_prunning = true;
    let mut state_table: StateTable<HummockStorage> = StateTableBuilder::new(
        &table,
        test_env.storage.clone(),
        Some(Arc::new(vnode_bitmap)),
    )
    .enable_vnode_key_stats(true, &Arc::new(config))
    .with_metrics(metrics)
    .forbid_preload_all_rows()
    .build()
    .await;

    let mut epoch = EpochPair::new_test_epoch(test_epoch(1));
    test_env
        .storage
        .start_epoch(epoch.curr, HashSet::from_iter([TEST_TABLE_ID]));
    state_table.init_epoch(epoch).await.unwrap();

    // Insert rows with two-column PKs
    // (1, 10), (1, 20), (1, 30)
    // (5, 50), (5, 60)
    // (9, 90)
    state_table.insert(OwnedRow::new(vec![
        Some(1_i32.into()),
        Some(10_i32.into()),
        Some(100_i32.into()),

Review comment (suggestion):
Spelling error: "prunning" should be "pruning" (single 'n'). This affects the field access which should be `enable_state_table_vnode_stats_pruning` instead of `enable_state_table_vnode_stats_prunning`.
```suggestion
    config.developer.enable_state_table_vnode_stats_pruning = true;
```

Code after:
    pin_mut!(pruned_iter);
    assert!(pruned_iter.next().await.is_none());
    assert_eq!(
        state_table.metrics().unwrap().iter_vnode_pruned_count.get(),
        1
    );

    // iteration over another pruned prefix should immediately return empty and increment metric
    let pruned_prefix = OwnedRow::new(vec![Some(10_i32.into())]);
    let pruned_iter = state_table
        .iter_with_prefix(pruned_prefix, sub_range, Default::default())
        .await
        .unwrap();
    pin_mut!(pruned_iter);
    assert!(pruned_iter.next().await.is_none());
    assert_eq!(
        state_table.metrics().unwrap().iter_vnode_pruned_count.get(),
        2
    );
}

#[tokio::test]
async fn test_state_table_pruned_key_range_with_two_pk_columns() {
    use prometheus::Registry;
    use risingwave_common::config::MetricLevel;

    use crate::common::table::state_table::StateTableBuilder;
    use crate::executor::monitor::StreamingMetrics;
    use crate::task::{ActorId, FragmentId};

    const TEST_TABLE_ID: TableId = TableId::new(234);
    let test_env = prepare_hummock_test_env().await;

    // Create a table with two primary key columns
    let column_descs = vec![
        ColumnDesc::unnamed(ColumnId::from(0), DataType::Int32),
        ColumnDesc::unnamed(ColumnId::from(1), DataType::Int32),
        ColumnDesc::unnamed(ColumnId::from(2), DataType::Int32),
    ];
    let order_types = vec![OrderType::ascending(), OrderType::ascending()];
    let pk_index = vec![0_usize, 1_usize];
    let read_prefix_len_hint = 0;
    let table = gen_pbtable(
        TEST_TABLE_ID,
        column_descs.clone(),
        order_types.clone(),
        pk_index.clone(),
        read_prefix_len_hint,
    );

    test_env.register_table(table.clone()).await;