Merge branch 'main' of https://github.com/risingwavelabs/risingwave i…

…nto li0k/storage_fix_correct_commit_ssts

src/batch/src/executor/join/distributed_lookup_join.rs

@@ -17,8 +17,9 @@ use std::mem::swap;
 
 use futures::pin_mut;
 use itertools::Itertools;
+use risingwave_common::bitmap::Bitmap;
 use risingwave_common::catalog::{ColumnDesc, ColumnId, Field, Schema};
-use risingwave_common::hash::{HashKey, HashKeyDispatcher};
+use risingwave_common::hash::{HashKey, HashKeyDispatcher, VirtualNode};
 use risingwave_common::memory::MemoryContext;
 use risingwave_common::row::OwnedRow;
 use risingwave_common::types::{DataType, Datum};
@@ -30,7 +31,7 @@ use risingwave_pb::batch_plan::plan_node::NodeBody;
 use risingwave_pb::common::BatchQueryEpoch;
 use risingwave_storage::store::PrefetchOptions;
 use risingwave_storage::table::batch_table::storage_table::StorageTable;
-use risingwave_storage::table::{TableDistribution, TableIter};
+use risingwave_storage::table::TableIter;
 use risingwave_storage::{dispatch_state_store, StateStore};
 
 use crate::error::Result;
@@ -194,7 +195,8 @@ impl BoxedExecutorBuilder for DistributedLookupJoinExecutorBuilder {
             .collect();
 
         // Lookup Join always contains distribution key, so we don't need vnode bitmap
-        let vnodes = Some(TableDistribution::all_vnodes());
+        // TODO(var-vnode): use vnode count from table desc
+        let vnodes = Some(Bitmap::ones(VirtualNode::COUNT).into());
         dispatch_state_store!(source.context().state_store(), state_store, {
             let table = StorageTable::new_partial(state_store, column_ids, vnodes, table_desc);
             let inner_side_builder = InnerSideExecutorBuilder::new(

src/batch/src/executor/join/local_lookup_join.rs

@@ -17,7 +17,7 @@ use std::marker::PhantomData;
 
 use anyhow::Context;
 use itertools::Itertools;
-use risingwave_common::bitmap::BitmapBuilder;
+use risingwave_common::bitmap::{Bitmap, BitmapBuilder};
 use risingwave_common::catalog::{ColumnDesc, Field, Schema};
 use risingwave_common::hash::table_distribution::TableDistribution;
 use risingwave_common::hash::{
@@ -408,12 +408,11 @@ impl BoxedExecutorBuilder for LocalLookupJoinExecutorBuilder {
             })
             .collect();
 
+        // TODO(var-vnode): use vnode count from table desc
+        let vnodes = Some(Bitmap::ones(VirtualNode::COUNT).into());
         let inner_side_builder = InnerSideExecutorBuilder {
             table_desc: table_desc.clone(),
-            table_distribution: TableDistribution::new_from_storage_table_desc(
-                Some(TableDistribution::all_vnodes()),
-                table_desc,
-            ),
+            table_distribution: TableDistribution::new_from_storage_table_desc(vnodes, table_desc),
             vnode_mapping,
             outer_side_key_types,
             inner_side_schema,

src/batch/src/executor/log_row_seq_scan.rs

@@ -22,13 +22,14 @@ use prometheus::Histogram;
 use risingwave_common::array::{DataChunk, Op};
 use risingwave_common::bitmap::Bitmap;
 use risingwave_common::catalog::{ColumnId, Field, Schema};
+use risingwave_common::hash::VirtualNode;
 use risingwave_common::row::{Row, RowExt};
 use risingwave_common::types::ScalarImpl;
 use risingwave_pb::batch_plan::plan_node::NodeBody;
 use risingwave_pb::common::{batch_query_epoch, BatchQueryEpoch};
 use risingwave_pb::plan_common::StorageTableDesc;
 use risingwave_storage::table::batch_table::storage_table::StorageTable;
-use risingwave_storage::table::{collect_data_chunk, TableDistribution};
+use risingwave_storage::table::collect_data_chunk;
 use risingwave_storage::{dispatch_state_store, StateStore};
 
 use super::{BoxedDataChunkStream, BoxedExecutor, BoxedExecutorBuilder, Executor, ExecutorBuilder};
@@ -106,7 +107,8 @@ impl BoxedExecutorBuilder for LogStoreRowSeqScanExecutorBuilder {
             Some(vnodes) => Some(Bitmap::from(vnodes).into()),
             // This is possible for dml. vnode_bitmap is not filled by scheduler.
             // Or it's single distribution, e.g., distinct agg. We scan in a single executor.
-            None => Some(TableDistribution::all_vnodes()),
+            // TODO(var-vnode): use vnode count from table desc
+            None => Some(Bitmap::ones(VirtualNode::COUNT).into()),
         };
 
         let chunk_size = source.context.get_config().developer.chunk_size as u32;

src/batch/src/executor/row_seq_scan.rs

@@ -21,6 +21,7 @@ use prometheus::Histogram;
 use risingwave_common::array::DataChunk;
 use risingwave_common::bitmap::Bitmap;
 use risingwave_common::catalog::{ColumnId, Schema};
+use risingwave_common::hash::VirtualNode;
 use risingwave_common::row::{OwnedRow, Row};
 use risingwave_common::types::{DataType, Datum};
 use risingwave_common::util::chunk_coalesce::DataChunkBuilder;
@@ -32,7 +33,6 @@ use risingwave_pb::plan_common::as_of::AsOfType;
 use risingwave_pb::plan_common::{as_of, PbAsOf, StorageTableDesc};
 use risingwave_storage::store::PrefetchOptions;
 use risingwave_storage::table::batch_table::storage_table::StorageTable;
-use risingwave_storage::table::TableDistribution;
 use risingwave_storage::{dispatch_state_store, StateStore};
 
 use crate::error::{BatchError, Result};
@@ -210,7 +210,8 @@ impl BoxedExecutorBuilder for RowSeqScanExecutorBuilder {
             Some(vnodes) => Some(Bitmap::from(vnodes).into()),
             // This is possible for dml. vnode_bitmap is not filled by scheduler.
             // Or it's single distribution, e.g., distinct agg. We scan in a single executor.
-            None => Some(TableDistribution::all_vnodes()),
+            // TODO(var-vnode): use vnode count from table desc
+            None => Some(Bitmap::ones(VirtualNode::COUNT).into()),
         };
 
         let scan_ranges = {

src/common/src/hash/consistent_hash/bitmap.rs

@@ -15,6 +15,7 @@
 use std::ops::RangeInclusive;
 
 use crate::bitmap::Bitmap;
+use crate::hash::table_distribution::SINGLETON_VNODE;
 use crate::hash::VirtualNode;
 
 /// An extension trait for `Bitmap` to support virtual node operations.
@@ -36,4 +37,17 @@ impl Bitmap {
         self.high_ranges()
             .map(|r| (VirtualNode::from_index(*r.start())..=VirtualNode::from_index(*r.end())))
     }
+
+    /// Returns whether only the [`SINGLETON_VNODE`] is set in the bitmap.
+    ///
+    /// Note that this method returning `true` does not imply that the bitmap was created by
+    /// [`VnodeBitmapExt::singleton`], or that the bitmap has length 1.
+    pub fn is_singleton(&self) -> bool {
+        self.count_ones() == 1 && self.iter_vnodes().next().unwrap() == SINGLETON_VNODE
+    }
+
+    /// Creates a bitmap with length 1 and the single bit set.
+    pub fn singleton() -> Self {
+        Self::ones(1)
+    }
 }

src/common/src/hash/consistent_hash/vnode.rs

@@ -114,6 +114,11 @@ impl VirtualNode {
     }
 }
 
+impl VirtualNode {
+    pub const COUNT_FOR_TEST: usize = Self::COUNT;
+    pub const MAX_FOR_TEST: VirtualNode = Self::MAX;
+}
+
 impl VirtualNode {
     // `compute_chunk` is used to calculate the `VirtualNode` for the columns in the
     // chunk. When only one column is provided and its type is `Serial`, we consider the column to