From 8fb97ea1d8de342e99cc34b988c0c843ec3dfd12 Mon Sep 17 00:00:00 2001
From: dennis zhuang <killme2008@gmail.com>
Date: Tue, 21 Mar 2023 11:19:43 +0800
Subject: [PATCH] fix: losing region numbers after altering table (#1209)

---
 src/mito/src/engine/tests.rs | 2 ++
 src/table/src/metadata.rs    | 7 +++++++
 2 files changed, 9 insertions(+)

diff --git a/src/mito/src/engine/tests.rs b/src/mito/src/engine/tests.rs
index 1368dc6a29..417ecf98ee 100644
--- a/src/mito/src/engine/tests.rs
+++ b/src/mito/src/engine/tests.rs
@@ -523,6 +523,7 @@ async fn test_alter_table_add_column() {
     assert_eq!(new_schema.timestamp_column(), old_schema.timestamp_column());
     assert_eq!(new_schema.version(), old_schema.version() + 1);
     assert_eq!(new_meta.next_column_id, old_meta.next_column_id + 2);
+    assert_eq!(new_meta.region_numbers, old_meta.region_numbers);
 }
 
 #[tokio::test]
@@ -572,6 +573,7 @@ async fn test_alter_table_remove_column() {
     assert_eq!(&[1, 2], &new_meta.value_indices[..]);
     assert_eq!(new_schema.timestamp_column(), old_schema.timestamp_column());
     assert_eq!(new_schema.version(), old_schema.version() + 1);
+    assert_eq!(new_meta.region_numbers, old_meta.region_numbers);
 }
 
 #[tokio::test]
diff --git a/src/table/src/metadata.rs b/src/table/src/metadata.rs
index 55f660ca76..c95b0d2c8e 100644
--- a/src/table/src/metadata.rs
+++ b/src/table/src/metadata.rs
@@ -89,6 +89,9 @@ pub struct TableIdent {
     pub version: TableVersion,
 }
 
+/// The table medatadata
+/// Note: if you add new fields to this struct, please ensure 'new_meta_builder' function works.
+/// TODO(dennis): find a better way to ensure 'new_meta_builder' works when adding new fields.
 #[derive(Clone, Debug, Builder, PartialEq, Eq)]
 #[builder(pattern = "mutable")]
 pub struct TableMeta {
@@ -197,6 +200,7 @@ impl TableMeta {
             .engine_options(self.engine_options.clone())
             .options(self.options.clone())
             .created_on(self.created_on)
+            .region_numbers(self.region_numbers.clone())
             .next_column_id(self.next_column_id);
 
         builder
@@ -572,6 +576,7 @@ mod tests {
             .unwrap();
 
         let new_meta = add_columns_to_meta(&meta);
+        assert_eq!(meta.region_numbers, new_meta.region_numbers);
 
         let names: Vec<String> = new_meta
             .schema
@@ -606,6 +611,8 @@ mod tests {
             .build()
             .unwrap();
 
+        assert_eq!(meta.region_numbers, new_meta.region_numbers);
+
         let names: Vec<String> = new_meta
             .schema
             .column_schemas()