feat(mito): backfill partition expr on region open (#6862)

* feat(mito): backfill partition expr on region open

Signed-off-by: Zhenchi <zhongzc_arch@outlook.com>

* tiny polish

Signed-off-by: Zhenchi <zhongzc_arch@outlook.com>

* only writable leader persists

Signed-off-by: Zhenchi <zhongzc_arch@outlook.com>

* catchup needs backfill

Signed-off-by: Zhenchi <zhongzc_arch@outlook.com>

* add log

Signed-off-by: Zhenchi <zhongzc_arch@outlook.com>

* handle distribute mode

Signed-off-by: Zhenchi <zhongzc_arch@outlook.com>

* only when transfer to set gracefully

Signed-off-by: Zhenchi <zhongzc_arch@outlook.com>

* fix fmt

Signed-off-by: Zhenchi <zhongzc_arch@outlook.com>

---------

Signed-off-by: Zhenchi <zhongzc_arch@outlook.com>

src/datanode/src/alive_keeper.rs

@@ -26,7 +26,7 @@ use common_meta::heartbeat::handler::{
 };
 use common_telemetry::{debug, error, info, trace, warn};
 use snafu::OptionExt;
-use store_api::region_engine::RegionRole;
+use store_api::region_engine::{RegionRole, SettableRegionRoleState};
 use store_api::region_request::{RegionCloseRequest, RegionRequest};
 use store_api::storage::RegionId;
 #[cfg(test)]
@@ -423,17 +423,27 @@ impl CountdownTask {
                             }
                         },
                         Some(CountdownCommand::Reset((role, deadline, extension_info))) => {
+                            let prev_role_leader = self.region_server.is_region_leader(self.region_id).unwrap_or(false);
                             if let Err(err) = self.region_server.set_region_role(self.region_id, role) {
                                 if err.status_code() == StatusCode::RegionNotFound {
                                     // Table metadata in metasrv is deleted after its regions are dropped.
                                     // The datanode may still receive lease renewal responses that depend on the metadata
                                     // during the short period before it is removed.
                                     warn!(err; "Failed to set region role to {role} for region {region_id}");
-                                }else{
+                                } else {
                                     error!(err; "Failed to set region role to {role} for region {region_id}");
                                 }
 
+                            // Finalize leadership: persist backfilled metadata.
+                            } else if !prev_role_leader && role == RegionRole::Leader
+                                && let Err(err) = self
+                                    .region_server
+                                    .set_region_role_state_gracefully(self.region_id, SettableRegionRoleState::Leader)
+                                    .await
+                            {
+                                error!(err; "Failed to set region role state gracefully to {role} for region {region_id}");
                             }
+
                             if let Some(ext_handler) = self.handler_ext.as_ref() {
                                 ext_handler.reset_deadline(
                                     &self.region_server,

src/datanode/src/datanode.rs

@@ -39,6 +39,7 @@ use meta_client::MetaClientRef;
 use metric_engine::engine::MetricEngine;
 use mito2::config::MitoConfig;
 use mito2::engine::{MitoEngine, MitoEngineBuilder};
+use mito2::region::opener::PartitionExprFetcherRef;
 use mito2::sst::file_ref::{FileReferenceManager, FileReferenceManagerRef};
 use object_store::manager::{ObjectStoreManager, ObjectStoreManagerRef};
 use object_store::util::normalize_dir;
@@ -48,7 +49,9 @@ use servers::export_metrics::ExportMetricsTask;
 use servers::server::ServerHandlers;
 use snafu::{OptionExt, ResultExt, ensure};
 use store_api::path_utils::WAL_DIR;
-use store_api::region_engine::{RegionEngineRef, RegionRole};
+use store_api::region_engine::{
+    RegionEngineRef, RegionRole, SetRegionRoleStateResponse, SettableRegionRoleState,
+};
 use tokio::fs;
 use tokio::sync::Notify;
 
@@ -64,6 +67,7 @@ use crate::event_listener::{
 };
 use crate::greptimedb_telemetry::get_greptimedb_telemetry_task;
 use crate::heartbeat::HeartbeatTask;
+use crate::partition_expr_fetcher::MetaPartitionExprFetcher;
 use crate::region_server::{DummyTableProviderFactory, RegionServer};
 use crate::store::{self, new_object_store_without_cache};
 use crate::utils::{RegionOpenRequests, build_region_open_requests};
@@ -448,12 +452,15 @@ impl DatanodeBuilder {
             }
         }
 
+        // Build a fetcher to backfill partition_expr on open.
+        let fetcher = Arc::new(MetaPartitionExprFetcher::new(self.kv_backend.clone()));
         let mito_engine = self
             .build_mito_engine(
                 object_store_manager.clone(),
                 mito_engine_config,
                 schema_metadata_manager.clone(),
                 file_ref_manager.clone(),
+                fetcher.clone(),
                 plugins.clone(),
             )
             .await?;
@@ -480,6 +487,7 @@ impl DatanodeBuilder {
         mut config: MitoConfig,
         schema_metadata_manager: SchemaMetadataManagerRef,
         file_ref_manager: FileReferenceManagerRef,
+        partition_expr_fetcher: PartitionExprFetcherRef,
         plugins: Plugins,
     ) -> Result<MitoEngine> {
         let opts = &self.opts;
@@ -502,6 +510,7 @@ impl DatanodeBuilder {
                     object_store_manager,
                     schema_metadata_manager,
                     file_ref_manager,
+                    partition_expr_fetcher.clone(),
                     plugins,
                 );
 
@@ -543,6 +552,7 @@ impl DatanodeBuilder {
                     object_store_manager,
                     schema_metadata_manager,
                     file_ref_manager,
+                    partition_expr_fetcher,
                     plugins,
                 );
 
@@ -655,12 +665,25 @@ async fn open_all_regions(
     }
 
     for region_id in open_regions {
-        if open_with_writable
-            && let Err(e) = region_server.set_region_role(region_id, RegionRole::Leader)
-        {
-            error!(
-                e; "failed to convert region {region_id} to leader"
-            );
+        if open_with_writable {
+            let res = region_server.set_region_role(region_id, RegionRole::Leader);
+            match res {
+                Ok(_) => {
+                    // Finalize leadership: persist backfilled metadata.
+                    if let SetRegionRoleStateResponse::InvalidTransition(err) = region_server
+                        .set_region_role_state_gracefully(
+                            region_id,
+                            SettableRegionRoleState::Leader,
+                        )
+                        .await?
+                    {
+                        error!(err; "failed to convert region {region_id} to leader");
+                    }
+                }
+                Err(e) => {
+                    error!(e; "failed to convert region {region_id} to leader");
+                }
+            }
         }
     }
 

src/datanode/src/lib.rs

@@ -23,6 +23,7 @@ pub mod event_listener;
 mod greptimedb_telemetry;
 pub mod heartbeat;
 pub mod metrics;
+mod partition_expr_fetcher;
 pub mod region_server;
 pub mod service;
 pub mod store;

src/datanode/src/partition_expr_fetcher.rs

@@ -0,0 +1,91 @@
+// Copyright 2023 Greptime Team
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use common_meta::key::table_route::TableRouteManager;
+use common_meta::kv_backend::KvBackendRef;
+use mito2::region::opener::PartitionExprFetcher;
+use store_api::storage::RegionId;
+
+/// A fetcher to fetch partition exprs from the table route.
+pub struct MetaPartitionExprFetcher {
+    table_route_manager: TableRouteManager,
+}
+
+impl MetaPartitionExprFetcher {
+    /// Creates a new [MetaPartitionExprFetcher].
+    pub fn new(kv: KvBackendRef) -> Self {
+        Self {
+            table_route_manager: TableRouteManager::new(kv),
+        }
+    }
+}
+
+#[async_trait::async_trait]
+impl PartitionExprFetcher for MetaPartitionExprFetcher {
+    async fn fetch_expr(&self, region_id: RegionId) -> Option<String> {
+        let table_id = region_id.table_id();
+        let Ok((_, route)) = self
+            .table_route_manager
+            .get_physical_table_route(table_id)
+            .await
+        else {
+            return None;
+        };
+        let region_number = region_id.region_number();
+        let rr = route
+            .region_routes
+            .iter()
+            .find(|r| r.region.id.region_number() == region_number)?;
+        Some(rr.region.partition_expr())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use common_meta::key::table_route::{TableRouteStorage, TableRouteValue};
+    use common_meta::kv_backend::TxnService;
+    use common_meta::kv_backend::memory::MemoryKvBackend;
+    use common_meta::rpc::router::{Region, RegionRoute};
+
+    use super::*;
+
+    #[tokio::test]
+    async fn test_fetch_expr_json() {
+        let kv = std::sync::Arc::new(MemoryKvBackend::new());
+        let storage = TableRouteStorage::new(kv.clone());
+        let table_id_u32 = 42;
+        let region_number = 7;
+        let region_id = RegionId::new(table_id_u32, region_number);
+        let expr_json =
+            r#"{"Expr":{"lhs":{"Column":"a"},"op":"GtEq","rhs":{"Value":{"UInt32":10}}}}"#;
+
+        let region_route = RegionRoute {
+            region: Region {
+                id: region_id,
+                name: "r".to_string(),
+                attrs: Default::default(),
+                partition: None,
+                partition_expr: expr_json.to_string(),
+            },
+            ..Default::default()
+        };
+        let trv = TableRouteValue::physical(vec![region_route]);
+        let (txn, _dec) = storage.build_create_txn(table_id_u32, &trv).unwrap();
+        kv.txn(txn).await.unwrap();
+
+        let fetcher = MetaPartitionExprFetcher::new(kv);
+        let fetched = fetcher.fetch_expr(region_id).await;
+        assert_eq!(fetched.as_deref(), Some(expr_json));
+    }
+}

src/mito2/src/engine.rs

@@ -111,6 +111,7 @@ use crate::metrics::HANDLE_REQUEST_ELAPSED;
 use crate::read::scan_region::{ScanRegion, Scanner};
 use crate::read::stream::ScanBatchStream;
 use crate::region::MitoRegionRef;
+use crate::region::opener::PartitionExprFetcherRef;
 use crate::request::{RegionEditRequest, WorkerRequest};
 use crate::sst::file::FileMeta;
 use crate::sst::file_ref::FileReferenceManagerRef;
@@ -129,12 +130,14 @@ pub struct MitoEngineBuilder<'a, S: LogStore> {
     object_store_manager: ObjectStoreManagerRef,
     schema_metadata_manager: SchemaMetadataManagerRef,
     file_ref_manager: FileReferenceManagerRef,
+    partition_expr_fetcher: PartitionExprFetcherRef,
     plugins: Plugins,
     #[cfg(feature = "enterprise")]
     extension_range_provider_factory: Option<BoxedExtensionRangeProviderFactory>,
 }
 
 impl<'a, S: LogStore> MitoEngineBuilder<'a, S> {
+    #[allow(clippy::too_many_arguments)]
     pub fn new(
         data_home: &'a str,
         config: MitoConfig,
@@ -142,6 +145,7 @@ impl<'a, S: LogStore> MitoEngineBuilder<'a, S> {
         object_store_manager: ObjectStoreManagerRef,
         schema_metadata_manager: SchemaMetadataManagerRef,
         file_ref_manager: FileReferenceManagerRef,
+        partition_expr_fetcher: PartitionExprFetcherRef,
         plugins: Plugins,
     ) -> Self {
         Self {
@@ -152,6 +156,7 @@ impl<'a, S: LogStore> MitoEngineBuilder<'a, S> {
             schema_metadata_manager,
             file_ref_manager,
             plugins,
+            partition_expr_fetcher,
             #[cfg(feature = "enterprise")]
             extension_range_provider_factory: None,
         }
@@ -179,6 +184,7 @@ impl<'a, S: LogStore> MitoEngineBuilder<'a, S> {
             self.object_store_manager,
             self.schema_metadata_manager,
             self.file_ref_manager,
+            self.partition_expr_fetcher.clone(),
             self.plugins,
         )
         .await?;
@@ -209,6 +215,7 @@ pub struct MitoEngine {
 
 impl MitoEngine {
     /// Returns a new [MitoEngine] with specific `config`, `log_store` and `object_store`.
+    #[allow(clippy::too_many_arguments)]
     pub async fn new<S: LogStore>(
         data_home: &str,
         config: MitoConfig,
@@ -216,6 +223,7 @@ impl MitoEngine {
         object_store_manager: ObjectStoreManagerRef,
         schema_metadata_manager: SchemaMetadataManagerRef,
         file_ref_manager: FileReferenceManagerRef,
+        partition_expr_fetcher: PartitionExprFetcherRef,
         plugins: Plugins,
     ) -> Result<MitoEngine> {
         let builder = MitoEngineBuilder::new(
@@ -225,6 +233,7 @@ impl MitoEngine {
             object_store_manager,
             schema_metadata_manager,
             file_ref_manager,
+            partition_expr_fetcher,
             plugins,
         );
         builder.try_build().await
@@ -946,6 +955,7 @@ impl MitoEngine {
         time_provider: crate::time_provider::TimeProviderRef,
         schema_metadata_manager: SchemaMetadataManagerRef,
         file_ref_manager: FileReferenceManagerRef,
+        partition_expr_fetcher: PartitionExprFetcherRef,
     ) -> Result<MitoEngine> {
         config.sanitize(data_home)?;
 
@@ -962,6 +972,7 @@ impl MitoEngine {
                     schema_metadata_manager,
                     file_ref_manager,
                     time_provider,
+                    partition_expr_fetcher,
                 )
                 .await?,
                 config,

src/mito2/src/engine/alter_test.rs

@@ -402,7 +402,7 @@ async fn test_alter_on_flushing() {
     let mut env = TestEnv::new().await;
     let listener = Arc::new(AlterFlushListener::default());
     let engine = env
-        .create_engine_with(MitoConfig::default(), None, Some(listener.clone()))
+        .create_engine_with(MitoConfig::default(), None, Some(listener.clone()), None)
         .await;
 
     let region_id = RegionId::new(1, 1);
@@ -506,7 +506,7 @@ async fn test_alter_column_fulltext_options() {
     let mut env = TestEnv::new().await;
     let listener = Arc::new(AlterFlushListener::default());
     let engine = env
-        .create_engine_with(MitoConfig::default(), None, Some(listener.clone()))
+        .create_engine_with(MitoConfig::default(), None, Some(listener.clone()), None)
         .await;
 
     let region_id = RegionId::new(1, 1);
@@ -629,7 +629,7 @@ async fn test_alter_column_set_inverted_index() {
     let mut env = TestEnv::new().await;
     let listener = Arc::new(AlterFlushListener::default());
     let engine = env
-        .create_engine_with(MitoConfig::default(), None, Some(listener.clone()))
+        .create_engine_with(MitoConfig::default(), None, Some(listener.clone()), None)
         .await;
 
     let region_id = RegionId::new(1, 1);
@@ -743,7 +743,7 @@ async fn test_alter_region_ttl_options() {
     let mut env = TestEnv::new().await;
     let listener = Arc::new(AlterFlushListener::default());
     let engine = env
-        .create_engine_with(MitoConfig::default(), None, Some(listener.clone()))
+        .create_engine_with(MitoConfig::default(), None, Some(listener.clone()), None)
         .await;
 
     let region_id = RegionId::new(1, 1);
@@ -793,7 +793,7 @@ async fn test_write_stall_on_altering() {
     let mut env = TestEnv::new().await;
     let listener = Arc::new(NotifyRegionChangeResultListener::default());
     let engine = env
-        .create_engine_with(MitoConfig::default(), None, Some(listener.clone()))
+        .create_engine_with(MitoConfig::default(), None, Some(listener.clone()), None)
         .await;
 
     let region_id = RegionId::new(1, 1);

src/mito2/src/engine/compaction_test.rs

@@ -519,6 +519,7 @@ async fn test_readonly_during_compaction() {
             },
             None,
             Some(listener.clone()),
+            None,
         )
         .await;
 

src/mito2/src/engine/drop_test.rs

@@ -38,7 +38,7 @@ async fn test_engine_drop_region() {
     let mut env = TestEnv::with_prefix("drop").await;
     let listener = Arc::new(DropListener::new(Duration::from_millis(100)));
     let engine = env
-        .create_engine_with(MitoConfig::default(), None, Some(listener.clone()))
+        .create_engine_with(MitoConfig::default(), None, Some(listener.clone()), None)
         .await;
 
     let region_id = RegionId::new(1, 1);

src/mito2/src/engine/edit_region_test.rs

@@ -148,6 +148,7 @@ async fn test_edit_region_fill_cache() {
             Some(Arc::new(EditRegionListener {
                 tx: Mutex::new(Some(tx)),
             })),
+            None,
         )
         .await;
 

src/mito2/src/engine/flush_test.rs

@@ -99,6 +99,7 @@ async fn test_flush_engine() {
             MitoConfig::default(),
             Some(write_buffer_manager.clone()),
             Some(listener.clone()),
+            None,
         )
         .await;
 
@@ -169,6 +170,7 @@ async fn test_write_stall() {
             MitoConfig::default(),
             Some(write_buffer_manager.clone()),
             Some(listener.clone()),
+            None,
         )
         .await;
 
@@ -243,6 +245,7 @@ async fn test_flush_empty() {
             MitoConfig::default(),
             Some(write_buffer_manager.clone()),
             None,
+            None,
         )
         .await;
 
@@ -478,6 +481,7 @@ async fn test_flush_workers() {
             },
             Some(write_buffer_manager.clone()),
             Some(listener.clone()),
+            None,
         )
         .await;
 
@@ -560,6 +564,7 @@ async fn test_update_topic_latest_entry_id(factory: Option<LogStoreFactory>) {
             MitoConfig::default(),
             Some(write_buffer_manager.clone()),
             Some(listener.clone()),
+            None,
         )
         .await;
     let region_id = RegionId::new(1, 1);

src/mito2/src/engine/open_test.rs

@@ -13,6 +13,8 @@
 // limitations under the License.
 
 use std::collections::HashMap;
+use std::sync::Arc;
+use std::sync::atomic::{AtomicUsize, Ordering};
 use std::time::Duration;
 
 use api::v1::Rows;
@@ -20,7 +22,7 @@ use common_error::ext::ErrorExt;
 use common_error::status_code::StatusCode;
 use common_recordbatch::RecordBatches;
 use either::Either;
-use store_api::region_engine::{RegionEngine, RegionRole};
+use store_api::region_engine::{RegionEngine, RegionRole, SettableRegionRoleState};
 use store_api::region_request::{
     PathType, RegionCloseRequest, RegionOpenRequest, RegionPutRequest, RegionRequest,
 };
@@ -30,6 +32,7 @@ use tokio::sync::oneshot;
 use crate::compaction::compactor::{OpenCompactionRegionRequest, open_compaction_region};
 use crate::config::MitoConfig;
 use crate::error;
+use crate::region::opener::{PartitionExprFetcher, PartitionExprFetcherRef};
 use crate::region::options::RegionOptions;
 use crate::test_util::{
     CreateRequestBuilder, TestEnv, build_rows, flush_region, put_rows, reopen_region, rows_schema,
@@ -503,3 +506,153 @@ async fn test_open_compaction_region() {
 
     assert_eq!(region_id, compaction_region.region_id);
 }
+
+/// Returns a dummy fetcher that returns the expr_json only once.
+fn new_dummy_fetcher(region_id: RegionId, expr_json: String) -> PartitionExprFetcherRef {
+    struct DummyFetcher {
+        region_id: RegionId,
+        expr_json: String,
+        fetch_count: AtomicUsize,
+    }
+    #[async_trait::async_trait]
+    impl PartitionExprFetcher for DummyFetcher {
+        async fn fetch_expr(&self, region_id: RegionId) -> Option<String> {
+            if region_id == self.region_id {
+                let count = self.fetch_count.fetch_add(1, Ordering::Relaxed);
+                (count == 0).then(|| self.expr_json.clone())
+            } else {
+                None
+            }
+        }
+    }
+
+    Arc::new(DummyFetcher {
+        region_id,
+        expr_json,
+        fetch_count: AtomicUsize::new(0),
+    })
+}
+
+#[tokio::test]
+async fn test_open_backfills_partition_expr_with_fetcher() {
+    let mut env = TestEnv::with_prefix("open-backfill-fetcher").await;
+
+    // prepare plugins with dummy fetcher
+    let region_id = RegionId::new(100, 1);
+    let expr_json = r#"{\"Expr\":{\"lhs\":{\"Column\":\"a\"},\"op\":\"GtEq\",\"rhs\":{\"Value\":{\"UInt32\":10}}}}"#.to_string();
+    let fetcher = new_dummy_fetcher(region_id, expr_json.clone());
+    let engine = env
+        .create_engine_with(MitoConfig::default(), None, None, Some(fetcher))
+        .await;
+
+    // create region with partition_expr=None (legacy style)
+    let request = CreateRequestBuilder::new()
+        .table_dir("t")
+        .partition_expr_json(None)
+        .build();
+    engine
+        .handle_request(region_id, RegionRequest::Create(request.clone()))
+        .await
+        .unwrap();
+
+    let meta = engine.get_region(region_id).unwrap().metadata();
+    assert!(meta.partition_expr.is_none());
+
+    // close and reopen to trigger backfill in opener
+    engine
+        .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {}))
+        .await
+        .unwrap();
+    engine
+        .handle_request(
+            region_id,
+            RegionRequest::Open(RegionOpenRequest {
+                engine: String::new(),
+                table_dir: request.table_dir.clone(),
+                path_type: PathType::Bare,
+                options: HashMap::default(),
+                skip_wal_replay: false,
+                checkpoint: None,
+            }),
+        )
+        .await
+        .unwrap();
+
+    // verify partition_expr is backfilled and persisted
+    let meta = engine.get_region(region_id).unwrap().metadata();
+    assert_eq!(meta.partition_expr.as_deref(), Some(expr_json.as_str()));
+
+    // Set leader and trigger backfill
+    engine
+        .set_region_role(region_id, RegionRole::Leader)
+        .unwrap();
+    engine
+        .set_region_role_state_gracefully(region_id, SettableRegionRoleState::Leader)
+        .await
+        .unwrap();
+
+    // reopen again to ensure no further changes and still Some
+    engine
+        .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {}))
+        .await
+        .unwrap();
+    let engine = env.reopen_engine(engine, MitoConfig::default()).await;
+    engine
+        .handle_request(
+            region_id,
+            RegionRequest::Open(RegionOpenRequest {
+                engine: String::new(),
+                table_dir: request.table_dir.clone(),
+                path_type: PathType::Bare,
+                options: HashMap::default(),
+                skip_wal_replay: false,
+                checkpoint: None,
+            }),
+        )
+        .await
+        .unwrap();
+    let meta = engine.get_region(region_id).unwrap().metadata();
+    assert_eq!(meta.partition_expr.as_deref(), Some(expr_json.as_str()));
+}
+
+#[tokio::test]
+async fn test_open_keeps_none_without_fetcher() {
+    let mut env = TestEnv::with_prefix("open-no-fetcher").await;
+    // engine without fetcher
+    let engine = env.create_engine(MitoConfig::default()).await;
+
+    let region_id = RegionId::new(200, 1);
+    let request = CreateRequestBuilder::new()
+        .table_dir("t2")
+        .partition_expr_json(None)
+        .build();
+    engine
+        .handle_request(region_id, RegionRequest::Create(request.clone()))
+        .await
+        .unwrap();
+    let meta = engine.get_region(region_id).unwrap().metadata();
+    assert!(meta.partition_expr.is_none());
+
+    engine
+        .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {}))
+        .await
+        .unwrap();
+    let engine = env.reopen_engine(engine, MitoConfig::default()).await;
+    engine
+        .handle_request(
+            region_id,
+            RegionRequest::Open(RegionOpenRequest {
+                engine: String::new(),
+                table_dir: request.table_dir.clone(),
+                path_type: PathType::Bare,
+                options: HashMap::default(),
+                skip_wal_replay: false,
+                checkpoint: None,
+            }),
+        )
+        .await
+        .unwrap();
+
+    let meta = engine.get_region(region_id).unwrap().metadata();
+    assert!(meta.partition_expr.is_none());
+}

src/mito2/src/engine/truncate_test.rs

@@ -304,6 +304,7 @@ async fn test_engine_truncate_during_flush() {
             MitoConfig::default(),
             Some(write_buffer_manager),
             Some(listener.clone()),
+            None,
         )
         .await;
 

src/mito2/src/region.rs

@@ -41,7 +41,9 @@ use crate::error::{
     FlushableRegionStateSnafu, RegionNotFoundSnafu, RegionStateSnafu, RegionTruncatedSnafu, Result,
     UpdateManifestSnafu,
 };
-use crate::manifest::action::{RegionManifest, RegionMetaAction, RegionMetaActionList};
+use crate::manifest::action::{
+    RegionChange, RegionManifest, RegionMetaAction, RegionMetaActionList,
+};
 use crate::manifest::manager::RegionManifestManager;
 use crate::memtable::MemtableBuilderRef;
 use crate::region::version::{VersionControlRef, VersionRef};
@@ -321,7 +323,7 @@ impl MitoRegion {
         &self,
         state: SettableRegionRoleState,
     ) -> Result<()> {
-        let _manager = self.manifest_ctx.manifest_manager.write().await;
+        let mut manager = self.manifest_ctx.manifest_manager.write().await;
         let current_state = self.state();
 
         match state {
@@ -423,6 +425,36 @@ impl MitoRegion {
             }
         }
 
+        // Hack(zhongzc): If we have just become leader (writable), persist any backfilled metadata.
+        if self.state() == RegionRoleState::Leader(RegionLeaderState::Writable) {
+            // Persist backfilled metadata if manifest is missing fields (e.g., partition_expr)
+            let manifest_meta = &manager.manifest().metadata;
+            let current_meta = &self.version().metadata;
+            if manifest_meta.partition_expr.is_none() && current_meta.partition_expr.is_some() {
+                let action = RegionMetaAction::Change(RegionChange {
+                    metadata: current_meta.clone(),
+                });
+                let result = manager
+                    .update(
+                        RegionMetaActionList::with_action(action),
+                        RegionRoleState::Leader(RegionLeaderState::Writable),
+                    )
+                    .await;
+
+                match result {
+                    Ok(version) => {
+                        info!(
+                            "Successfully persisted backfilled metadata for region {}, version: {}",
+                            self.region_id, version
+                        );
+                    }
+                    Err(e) => {
+                        warn!(e; "Failed to persist backfilled metadata for region {}", self.region_id);
+                    }
+                }
+            }
+        }
+
         Ok(())
     }
 

src/mito2/src/region/opener.rs

@@ -69,6 +69,18 @@ use crate::time_provider::TimeProviderRef;
 use crate::wal::entry_reader::WalEntryReader;
 use crate::wal::{EntryId, Wal};
 
+/// A fetcher to retrieve partition expr for a region.
+///
+/// Compatibility: older regions didn't persist `partition_expr` in engine metadata,
+/// while newer ones do. On open, we backfill it via this fetcher and persist it
+/// to the manifest so future opens don't need refetching.
+#[async_trait::async_trait]
+pub trait PartitionExprFetcher {
+    async fn fetch_expr(&self, region_id: RegionId) -> Option<String>;
+}
+
+pub type PartitionExprFetcherRef = Arc<dyn PartitionExprFetcher + Send + Sync>;
+
 /// Builder to create a new [MitoRegion] or open an existing one.
 pub(crate) struct RegionOpener {
     region_id: RegionId,
@@ -88,6 +100,7 @@ pub(crate) struct RegionOpener {
     wal_entry_reader: Option<Box<dyn WalEntryReader>>,
     replay_checkpoint: Option<u64>,
     file_ref_manager: FileReferenceManagerRef,
+    partition_expr_fetcher: PartitionExprFetcherRef,
 }
 
 impl RegionOpener {
@@ -105,6 +118,7 @@ impl RegionOpener {
         intermediate_manager: IntermediateManager,
         time_provider: TimeProviderRef,
         file_ref_manager: FileReferenceManagerRef,
+        partition_expr_fetcher: PartitionExprFetcherRef,
     ) -> RegionOpener {
         RegionOpener {
             region_id,
@@ -124,6 +138,7 @@ impl RegionOpener {
             wal_entry_reader: None,
             replay_checkpoint: None,
             file_ref_manager,
+            partition_expr_fetcher,
         }
     }
 
@@ -390,8 +405,18 @@ impl RegionOpener {
             return Ok(None);
         };
 
+        // Backfill `partition_expr` if missing. Use the backfilled metadata in-memory during this open.
         let manifest = manifest_manager.manifest();
-        let metadata = manifest.metadata.clone();
+        let metadata = if manifest.metadata.partition_expr.is_none()
+            && let Some(expr_json) = self.partition_expr_fetcher.fetch_expr(self.region_id).await
+        {
+            let metadata = manifest.metadata.as_ref().clone();
+            let mut builder = RegionMetadataBuilder::from_existing(metadata);
+            builder.partition_expr_json(Some(expr_json));
+            Arc::new(builder.build().context(InvalidMetadataSnafu)?)
+        } else {
+            manifest.metadata.clone()
+        };
 
         let region_id = self.region_id;
         let provider = self.provider::<S>(&region_options.wal_options)?;

src/mito2/src/test_util.rs

@@ -73,6 +73,7 @@ use crate::error::Result;
 use crate::flush::{WriteBufferManager, WriteBufferManagerRef};
 use crate::manifest::manager::{RegionManifestManager, RegionManifestOptions};
 use crate::read::{Batch, BatchBuilder, BatchReader};
+use crate::region::opener::{PartitionExprFetcher, PartitionExprFetcherRef};
 use crate::sst::file_purger::{FilePurgerRef, NoopFilePurger};
 use crate::sst::file_ref::{FileReferenceManager, FileReferenceManagerRef};
 use crate::sst::index::intermediate::IntermediateManager;
@@ -105,6 +106,19 @@ pub(crate) fn kafka_log_store_factory() -> Option<LogStoreFactory> {
     }))
 }
 
+pub(crate) fn noop_partition_expr_fetcher() -> PartitionExprFetcherRef {
+    struct NoopPartitionExprFetcher;
+
+    #[async_trait::async_trait]
+    impl PartitionExprFetcher for NoopPartitionExprFetcher {
+        async fn fetch_expr(&self, _region_id: RegionId) -> Option<String> {
+            None
+        }
+    }
+
+    Arc::new(NoopPartitionExprFetcher)
+}
+
 #[template]
 #[rstest]
 #[case::with_raft_engine(raft_engine_log_store_factory())]
@@ -212,6 +226,7 @@ pub struct TestEnv {
     schema_metadata_manager: SchemaMetadataManagerRef,
     file_ref_manager: FileReferenceManagerRef,
     kv_backend: KvBackendRef,
+    partition_expr_fetcher: PartitionExprFetcherRef,
 }
 
 impl TestEnv {
@@ -247,6 +262,7 @@ impl TestEnv {
             schema_metadata_manager,
             file_ref_manager: Arc::new(FileReferenceManager::new(None)),
             kv_backend,
+            partition_expr_fetcher: noop_partition_expr_fetcher(),
         }
     }
 
@@ -284,6 +300,7 @@ impl TestEnv {
                 zelf.object_store_manager.as_ref().unwrap().clone(),
                 zelf.schema_metadata_manager.clone(),
                 zelf.file_ref_manager.clone(),
+                zelf.partition_expr_fetcher.clone(),
                 Plugins::new(),
             )
             .await
@@ -318,6 +335,7 @@ impl TestEnv {
         config: MitoConfig,
         manager: Option<WriteBufferManagerRef>,
         listener: Option<EventListenerRef>,
+        partition_expr_fetcher: Option<PartitionExprFetcherRef>,
     ) -> MitoEngine {
         let (log_store, object_store_manager) = self.create_log_and_object_store_manager().await;
 
@@ -327,6 +345,10 @@ impl TestEnv {
 
         let data_home = self.data_home().display().to_string();
 
+        let partition_expr_fetcher =
+            partition_expr_fetcher.unwrap_or_else(noop_partition_expr_fetcher);
+        self.partition_expr_fetcher = partition_expr_fetcher;
+
         match log_store {
             LogStoreImpl::RaftEngine(log_store) => MitoEngine::new_for_test(
                 &data_home,
@@ -338,6 +360,7 @@ impl TestEnv {
                 Arc::new(StdTimeProvider),
                 self.schema_metadata_manager.clone(),
                 self.file_ref_manager.clone(),
+                self.partition_expr_fetcher.clone(),
             )
             .await
             .unwrap(),
@@ -351,6 +374,7 @@ impl TestEnv {
                 Arc::new(StdTimeProvider),
                 self.schema_metadata_manager.clone(),
                 self.file_ref_manager.clone(),
+                self.partition_expr_fetcher.clone(),
             )
             .await
             .unwrap(),
@@ -395,6 +419,7 @@ impl TestEnv {
                 Arc::new(StdTimeProvider),
                 self.schema_metadata_manager.clone(),
                 self.file_ref_manager.clone(),
+                self.partition_expr_fetcher.clone(),
             )
             .await
             .unwrap(),
@@ -408,6 +433,7 @@ impl TestEnv {
                 Arc::new(StdTimeProvider),
                 self.schema_metadata_manager.clone(),
                 self.file_ref_manager.clone(),
+                self.partition_expr_fetcher.clone(),
             )
             .await
             .unwrap(),
@@ -441,6 +467,7 @@ impl TestEnv {
                 time_provider.clone(),
                 self.schema_metadata_manager.clone(),
                 self.file_ref_manager.clone(),
+                self.partition_expr_fetcher.clone(),
             )
             .await
             .unwrap(),
@@ -454,6 +481,7 @@ impl TestEnv {
                 time_provider.clone(),
                 self.schema_metadata_manager.clone(),
                 self.file_ref_manager.clone(),
+                self.partition_expr_fetcher.clone(),
             )
             .await
             .unwrap(),
@@ -491,6 +519,7 @@ impl TestEnv {
                 Arc::new(object_store_manager),
                 self.schema_metadata_manager.clone(),
                 self.file_ref_manager.clone(),
+                self.partition_expr_fetcher.clone(),
                 Plugins::new(),
             )
             .await
@@ -501,6 +530,7 @@ impl TestEnv {
                 Arc::new(object_store_manager),
                 self.schema_metadata_manager.clone(),
                 self.file_ref_manager.clone(),
+                self.partition_expr_fetcher.clone(),
                 Plugins::new(),
             )
             .await

src/mito2/src/worker.rs

@@ -59,6 +59,7 @@ use crate::error::{self, CreateDirSnafu, JoinSnafu, Result, WorkerStoppedSnafu};
 use crate::flush::{FlushScheduler, WriteBufferManagerImpl, WriteBufferManagerRef};
 use crate::memtable::MemtableBuilderProvider;
 use crate::metrics::{REGION_COUNT, REQUEST_WAIT_TIME, WRITE_STALLING};
+use crate::region::opener::PartitionExprFetcherRef;
 use crate::region::{MitoRegionRef, OpeningRegions, OpeningRegionsRef, RegionMap, RegionMapRef};
 use crate::request::{
     BackgroundNotify, DdlRequest, SenderBulkRequest, SenderDdlRequest, SenderWriteRequest,
@@ -145,6 +146,7 @@ impl WorkerGroup {
         object_store_manager: ObjectStoreManagerRef,
         schema_metadata_manager: SchemaMetadataManagerRef,
         file_ref_manager: FileReferenceManagerRef,
+        partition_expr_fetcher: PartitionExprFetcherRef,
         plugins: Plugins,
     ) -> Result<WorkerGroup> {
         let (flush_sender, flush_receiver) = watch::channel(());
@@ -209,6 +211,7 @@ impl WorkerGroup {
                     plugins: plugins.clone(),
                     schema_metadata_manager: schema_metadata_manager.clone(),
                     file_ref_manager: file_ref_manager.clone(),
+                    partition_expr_fetcher: partition_expr_fetcher.clone(),
                 }
                 .start()
             })
@@ -306,6 +309,7 @@ impl WorkerGroup {
         schema_metadata_manager: SchemaMetadataManagerRef,
         file_ref_manager: FileReferenceManagerRef,
         time_provider: TimeProviderRef,
+        partition_expr_fetcher: PartitionExprFetcherRef,
     ) -> Result<WorkerGroup> {
         let (flush_sender, flush_receiver) = watch::channel(());
         let write_buffer_manager = write_buffer_manager.unwrap_or_else(|| {
@@ -363,6 +367,7 @@ impl WorkerGroup {
                     plugins: Plugins::new(),
                     schema_metadata_manager: schema_metadata_manager.clone(),
                     file_ref_manager: file_ref_manager.clone(),
+                    partition_expr_fetcher: partition_expr_fetcher.clone(),
                 }
                 .start()
             })
@@ -443,6 +448,7 @@ struct WorkerStarter<S> {
     plugins: Plugins,
     schema_metadata_manager: SchemaMetadataManagerRef,
     file_ref_manager: FileReferenceManagerRef,
+    partition_expr_fetcher: PartitionExprFetcherRef,
 }
 
 impl<S: LogStore> WorkerStarter<S> {
@@ -496,6 +502,7 @@ impl<S: LogStore> WorkerStarter<S> {
             region_edit_queues: RegionEditQueues::default(),
             schema_metadata_manager: self.schema_metadata_manager,
             file_ref_manager: self.file_ref_manager.clone(),
+            partition_expr_fetcher: self.partition_expr_fetcher,
         };
         let handle = common_runtime::spawn_global(async move {
             worker_thread.run().await;
@@ -747,6 +754,8 @@ struct RegionWorkerLoop<S> {
     schema_metadata_manager: SchemaMetadataManagerRef,
     /// Datanode level file references manager.
     file_ref_manager: FileReferenceManagerRef,
+    /// Partition expr fetcher used to backfill partition expr on open for compatibility.
+    partition_expr_fetcher: PartitionExprFetcherRef,
 }
 
 impl<S: LogStore> RegionWorkerLoop<S> {

src/mito2/src/worker/handle_catchup.rs

@@ -20,7 +20,7 @@ use common_telemetry::tracing::warn;
 use common_telemetry::{debug, info};
 use snafu::ensure;
 use store_api::logstore::LogStore;
-use store_api::region_engine::RegionRole;
+use store_api::region_engine::{RegionRole, SettableRegionRoleState};
 use store_api::region_request::{AffectedRows, RegionCatchupRequest};
 use store_api::storage::RegionId;
 use tokio::time::Instant;
@@ -136,6 +136,10 @@ impl<S: LogStore> RegionWorkerLoop<S> {
 
         if request.set_writable {
             region.set_role(RegionRole::Leader);
+            // Finalize leadership: persist backfilled metadata.
+            region
+                .set_role_state_gracefully(SettableRegionRoleState::Leader)
+                .await?;
         }
 
         Ok(0)
@@ -164,6 +168,7 @@ impl<S: LogStore> RegionWorkerLoop<S> {
                 self.intermediate_manager.clone(),
                 self.time_provider.clone(),
                 self.file_ref_manager.clone(),
+                self.partition_expr_fetcher.clone(),
             )
             .cache(Some(self.cache_manager.clone()))
             .options(region.version().options.clone())?

src/mito2/src/worker/handle_create.rs

@@ -67,6 +67,7 @@ impl<S: LogStore> RegionWorkerLoop<S> {
             self.intermediate_manager.clone(),
             self.time_provider.clone(),
             self.file_ref_manager.clone(),
+            self.partition_expr_fetcher.clone(),
         )
         .metadata_builder(builder)
         .parse_options(request.options)?

src/mito2/src/worker/handle_open.rs

@@ -104,6 +104,7 @@ impl<S: LogStore> RegionWorkerLoop<S> {
             self.intermediate_manager.clone(),
             self.time_provider.clone(),
             self.file_ref_manager.clone(),
+            self.partition_expr_fetcher.clone(),
         )
         .skip_wal_replay(request.skip_wal_replay)
         .cache(Some(self.cache_manager.clone()))