From 732784d3f88788af9bc31920e7f97215ec19df08 Mon Sep 17 00:00:00 2001
From: Weny Xu <wenymedia@gmail.com>
Date: Fri, 26 May 2023 12:30:46 +0900
Subject: [PATCH] feat: support to load missing region (#1651)

* feat: support to load missing region

* Update src/mito/src/table.rs

Co-authored-by: Lei, HUANG <6406592+v0y4g3r@users.noreply.github.com>

---------

Co-authored-by: Lei, HUANG <6406592+v0y4g3r@users.noreply.github.com>
---
 src/mito/src/engine.rs |  1 -
 src/mito/src/table.rs  | 20 +++++++++++++++-----
 2 files changed, 15 insertions(+), 6 deletions(-)

diff --git a/src/mito/src/engine.rs b/src/mito/src/engine.rs
index f559c90d20..b6ccfa30f4 100644
--- a/src/mito/src/engine.rs
+++ b/src/mito/src/engine.rs
@@ -585,7 +585,6 @@ impl<S: StorageEngine> MitoEngineInner<S> {
                     Some(table)
                 } else {
                     // Loads missing regions
-                    // TODO(weny): Supports to load regions
                     self.load_missing_regions(ctx, table.clone(), &request.region_numbers)
                         .await?;
 
diff --git a/src/mito/src/table.rs b/src/mito/src/table.rs
index a4c2b502e5..ed3301a1e2 100644
--- a/src/mito/src/table.rs
+++ b/src/mito/src/table.rs
@@ -28,7 +28,7 @@ use common_query::logical_plan::Expr;
 use common_query::physical_plan::PhysicalPlanRef;
 use common_recordbatch::error::{ExternalSnafu, Result as RecordBatchResult};
 use common_recordbatch::{RecordBatch, RecordBatchStream};
-use common_telemetry::{logging, warn};
+use common_telemetry::{info, logging};
 use datatypes::schema::Schema;
 use futures::task::{Context, Poll};
 use futures::Stream;
@@ -615,12 +615,22 @@ impl<R: Region> MitoTable<R> {
         Ok(())
     }
 
-    pub async fn load_region(&self, region_number: RegionNumber, _region: R) -> TableResult<()> {
+    // Loads a region if the slot of the corresponding region number was not occupied.
+    // Assuming the regions with the same region_number are the same.
+    pub async fn load_region(&self, region_number: RegionNumber, region: R) -> TableResult<()> {
         let info = self.table_info.load();
 
-        // TODO(weny): Supports to load the region
-        warn!(
-            "MitoTable try to load region: {} in table: {}",
+        self.regions.rcu(|regions| {
+            let mut regions = HashMap::clone(regions);
+            regions
+                .entry(region_number)
+                .or_insert_with(|| region.clone());
+
+            Arc::new(regions)
+        });
+
+        info!(
+            "MitoTable loads new region: {} in table: {}",
             region_number,
             format!("{}.{}.{}", info.catalog_name, info.schema_name, info.name)
         );