From 8485c9af3394a74665b1010cc2b7ef1a966907e9 Mon Sep 17 00:00:00 2001
From: Ruihang Xia <waynestxia@gmail.com>
Date: Tue, 23 Jan 2024 17:10:27 +0800
Subject: [PATCH] feat: read column and region info from state cache (#3222)

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>
---
 src/metric-engine/src/engine/put.rs   | 15 ++++-----------
 src/metric-engine/src/engine/state.rs | 21 +++++++++++++++++++++
 2 files changed, 25 insertions(+), 11 deletions(-)

diff --git a/src/metric-engine/src/engine/put.rs b/src/metric-engine/src/engine/put.rs
index 10ffd39fee..3d4bc1485b 100644
--- a/src/metric-engine/src/engine/put.rs
+++ b/src/metric-engine/src/engine/put.rs
@@ -98,11 +98,9 @@ impl MetricEngineInner {
     ) -> Result<()> {
         // check if the region exists
         let metadata_region_id = to_metadata_region_id(physical_region_id);
-        if !self
-            .metadata_region
-            .is_logical_region_exists(metadata_region_id, logical_region_id)
-            .await?
-        {
+        let data_region_id = to_data_region_id(physical_region_id);
+        let state = self.state.read().unwrap();
+        if !state.is_logical_region_exist(logical_region_id) {
             error!("Trying to write to an nonexistent region {logical_region_id}");
             return LogicalRegionNotFoundSnafu {
                 region_id: logical_region_id,
@@ -112,12 +110,7 @@ impl MetricEngineInner {
 
         // check if the columns exist
         for col in &request.rows.schema {
-            if self
-                .metadata_region
-                .column_semantic_type(metadata_region_id, logical_region_id, &col.column_name)
-                .await?
-                .is_none()
-            {
+            if !state.is_physical_column_exist(data_region_id, &col.column_name)? {
                 return ColumnNotFoundSnafu {
                     name: col.column_name.clone(),
                     region_id: logical_region_id,
diff --git a/src/metric-engine/src/engine/state.rs b/src/metric-engine/src/engine/state.rs
index 154115138d..cc36fe27df 100644
--- a/src/metric-engine/src/engine/state.rs
+++ b/src/metric-engine/src/engine/state.rs
@@ -131,4 +131,25 @@ impl MetricEngineState {
 
         Ok(())
     }
+
+    /// Check if a physical column exists.
+    pub fn is_physical_column_exist(
+        &self,
+        physical_region_id: RegionId,
+        column_name: &str,
+    ) -> Result<bool> {
+        let data_region_id = to_data_region_id(physical_region_id);
+        let exist = self
+            .physical_columns()
+            .get(&data_region_id)
+            .context(PhysicalRegionNotFoundSnafu {
+                region_id: data_region_id,
+            })?
+            .contains(column_name);
+        Ok(exist)
+    }
+
+    pub fn is_logical_region_exist(&self, logical_region_id: RegionId) -> bool {
+        self.logical_regions().contains_key(&logical_region_id)
+    }
 }