refactor: unify additional column id logic using max_column_id (#17336

)

Signed-off-by: tabVersion <tabversion@bupt.icu>
Signed-off-by: xxchan <xxchan22f@gmail.com>
Co-authored-by: tabVersion <tabversion@bupt.icu>
Co-authored-by: xxchan <xxchan22f@gmail.com>

src/common/src/catalog/column.rs

@@ -451,6 +451,20 @@ pub fn debug_assert_column_ids_distinct(columns: &[ColumnCatalog]) {
     );
 }
 
+/// FIXME: perhapts we should use sth like `ColumnIdGenerator::new_alter`,
+/// However, the `SourceVersion` is problematic: It doesn't contain `next_col_id`.
+/// (But for now this isn't a large problem, since drop column is not allowed for source yet..)
+///
+/// Besides, the logic of column id handling is a mess.
+/// In some places, we use `ColumnId::placeholder()`, and use `col_id_gen` to fill it at the end;
+/// In other places, we create column id ad-hoc.
+pub fn max_column_id(columns: &[ColumnCatalog]) -> ColumnId {
+    // XXX: should we check the column IDs of struct fields here?
+    columns
+        .iter()
+        .fold(ColumnId::first_user_column(), |a, b| a.max(b.column_id()))
+}
+
 #[cfg(test)]
 pub mod tests {
     use risingwave_pb::plan_common::PbColumnDesc;

src/connector/src/parser/additional_columns.rs

@@ -16,7 +16,7 @@ use std::collections::{HashMap, HashSet};
 use std::sync::LazyLock;
 
 use risingwave_common::bail;
-use risingwave_common::catalog::{ColumnCatalog, ColumnDesc, ColumnId};
+use risingwave_common::catalog::{max_column_id, ColumnCatalog, ColumnDesc, ColumnId};
 use risingwave_common::types::{DataType, StructType};
 use risingwave_pb::data::data_type::TypeName;
 use risingwave_pb::data::DataType as PbDataType;
@@ -280,11 +280,7 @@ pub fn source_add_partition_offset_cols(
     connector_name: &str,
 ) -> ([bool; 2], [ColumnCatalog; 2]) {
     let mut columns_exist = [false; 2];
-    let mut last_column_id = columns
-        .iter()
-        .map(|c| c.column_desc.column_id)
-        .max()
-        .unwrap_or(ColumnId::placeholder());
+    let mut last_column_id = max_column_id(columns);
 
     let additional_columns: Vec<_> = {
         let compat_col_types = COMPATIBLE_ADDITIONAL_COLUMNS

src/frontend/src/handler/alter_source_column.rs

@@ -14,7 +14,7 @@
 
 use itertools::Itertools;
 use pgwire::pg_response::{PgResponse, StatementType};
-use risingwave_common::catalog::{ColumnCatalog, ColumnId};
+use risingwave_common::catalog::max_column_id;
 use risingwave_connector::source::{extract_source_struct, SourceEncode, SourceStruct};
 use risingwave_sqlparser::ast::{
     AlterSourceOperation, ColumnDef, CreateSourceStatement, ObjectName, Statement,
@@ -144,20 +144,6 @@ pub fn alter_definition_add_column(definition: &str, column: ColumnDef) -> Resul
     Ok(stmt.to_string())
 }
 
-/// FIXME: perhapts we should use sth like `ColumnIdGenerator::new_alter`,
-/// However, the `SourceVersion` is problematic: It doesn't contain `next_col_id`.
-/// (But for now this isn't a large problem, since drop column is not allowed for source yet..)
-///
-/// Besides, the logic of column id handling is a mess.
-/// In some places, we use `ColumnId::placeholder()`, and use `col_id_gen` to fill it at the end;
-/// In other places, we create column id ad-hoc.
-pub fn max_column_id(columns: &Vec<ColumnCatalog>) -> ColumnId {
-    // XXX: should we check the column IDs of struct fields here?
-    columns
-        .iter()
-        .fold(ColumnId::first_user_column(), |a, b| a.max(b.column_id()))
-}
-
 #[cfg(test)]
 pub mod tests {
     use std::collections::HashMap;

src/frontend/src/handler/alter_source_with_sr.rs

@@ -18,7 +18,7 @@ use anyhow::Context;
 use itertools::Itertools;
 use pgwire::pg_response::StatementType;
 use risingwave_common::bail_not_implemented;
-use risingwave_common::catalog::ColumnCatalog;
+use risingwave_common::catalog::{max_column_id, ColumnCatalog};
 use risingwave_connector::WithPropertiesExt;
 use risingwave_pb::catalog::StreamSourceInfo;
 use risingwave_pb::plan_common::{EncodeType, FormatType};
@@ -28,7 +28,6 @@ use risingwave_sqlparser::ast::{
 };
 use risingwave_sqlparser::parser::Parser;
 
-use super::alter_source_column::max_column_id;
 use super::alter_table_column::schema_has_schema_registry;
 use super::create_source::{bind_columns_from_source, validate_compatibility};
 use super::util::SourceSchemaCompatExt;