diff --git a/src/common/src/hash/consistent_hash/compat.rs b/src/common/src/hash/consistent_hash/compat.rs index 12f65631f3e88..c659da2fc6059 100644 --- a/src/common/src/hash/consistent_hash/compat.rs +++ b/src/common/src/hash/consistent_hash/compat.rs @@ -38,6 +38,13 @@ impl VnodeCount { VnodeCount::Set(NonZeroUsize::new(v).unwrap()) } + /// Creates a `VnodeCount` set to the value for testing. + /// + /// Equivalent to `VnodeCount::set(VirtualNode::COUNT_FOR_TEST)`. + pub fn for_test() -> Self { + Self::set(VirtualNode::COUNT_FOR_TEST) + } + /// Converts to protobuf representation for `maybe_vnode_count`. pub fn to_protobuf(self) -> Option { match self { diff --git a/src/frontend/src/handler/alter_table_column.rs b/src/frontend/src/handler/alter_table_column.rs index e9ecb5713cb10..88e886ad667bf 100644 --- a/src/frontend/src/handler/alter_table_column.rs +++ b/src/frontend/src/handler/alter_table_column.rs @@ -19,6 +19,7 @@ use anyhow::{anyhow, Context}; use itertools::Itertools; use pgwire::pg_response::{PgResponse, StatementType}; use risingwave_common::catalog::ColumnCatalog; +use risingwave_common::hash::VnodeCount; use risingwave_common::types::DataType; use risingwave_common::util::column_index_mapping::ColIndexMapping; use risingwave_common::{bail, bail_not_implemented}; @@ -244,7 +245,7 @@ pub async fn get_replace_table_plan( // Set some fields ourselves so that the meta service does not need to maintain them. let mut table = table; table.incoming_sinks = incoming_sink_ids.iter().copied().collect(); - table.maybe_vnode_count = Some(original_catalog.vnode_count() as _); + table.maybe_vnode_count = VnodeCount::set(original_catalog.vnode_count()).to_protobuf(); Ok((source, table, graph, col_index_mapping, job_type)) } diff --git a/src/frontend/src/test_utils.rs b/src/frontend/src/test_utils.rs index ac8784fde4dd3..14befbaeb7357 100644 --- a/src/frontend/src/test_utils.rs +++ b/src/frontend/src/test_utils.rs @@ -28,6 +28,7 @@ use risingwave_common::catalog::{ FunctionId, IndexId, TableId, DEFAULT_DATABASE_NAME, DEFAULT_SCHEMA_NAME, DEFAULT_SUPER_USER, DEFAULT_SUPER_USER_ID, NON_RESERVED_USER_ID, PG_CATALOG_SCHEMA_NAME, RW_CATALOG_SCHEMA_NAME, }; +use risingwave_common::hash::{VirtualNode, VnodeCount, VnodeCountCompat}; use risingwave_common::session_config::SessionConfig; use risingwave_common::system_param::reader::SystemParamsReader; use risingwave_common::util::cluster_limit::ClusterLimit; @@ -281,6 +282,7 @@ impl CatalogWriter for MockCatalogWriter { ) -> Result<()> { table.id = self.gen_id(); table.stream_job_status = PbStreamJobStatus::Created as _; + table.maybe_vnode_count = VnodeCount::for_test().to_protobuf(); self.catalog.write().create_table(&table); self.add_table_or_source_id(table.id, table.schema_id, table.database_id); self.hummock_snapshot_manager @@ -320,6 +322,7 @@ impl CatalogWriter for MockCatalogWriter { _job_type: TableJobType, ) -> Result<()> { table.stream_job_status = PbStreamJobStatus::Created as _; + assert_eq!(table.vnode_count(), VirtualNode::COUNT_FOR_TEST); self.catalog.write().update_table(&table); Ok(()) } @@ -353,6 +356,7 @@ impl CatalogWriter for MockCatalogWriter { ) -> Result<()> { index_table.id = self.gen_id(); index_table.stream_job_status = PbStreamJobStatus::Created as _; + index_table.maybe_vnode_count = VnodeCount::for_test().to_protobuf(); self.catalog.write().create_table(&index_table); self.add_table_or_index_id( index_table.id, diff --git a/src/meta/src/controller/fragment.rs b/src/meta/src/controller/fragment.rs index dfdf91372ff0b..a79b890cade20 100644 --- a/src/meta/src/controller/fragment.rs +++ b/src/meta/src/controller/fragment.rs @@ -1627,7 +1627,7 @@ mod tests { .values() .flat_map(|m| m.keys().map(|x| *x as _)) .collect(), - maybe_vnode_count: VnodeCount::set(VirtualNode::COUNT_FOR_TEST).to_protobuf(), + maybe_vnode_count: VnodeCount::for_test().to_protobuf(), }; let pb_actor_status = (0..actor_count)