fix: metrics without physical partition columns query push down (#6694)

* fix: metrics no part cols

Signed-off-by: discord9 <discord9@163.com>

* chore: typos

Signed-off-by: discord9 <discord9@163.com>

* chore: clippy

Signed-off-by: discord9 <discord9@163.com>

* chore: rename stuff

Signed-off-by: discord9 <discord9@163.com>

* refactor: put partition rules in table

Signed-off-by: discord9 <discord9@163.com>

* more tests

Signed-off-by: discord9 <discord9@163.com>

* test: redact more

Signed-off-by: discord9 <discord9@163.com>

---------

Signed-off-by: discord9 <discord9@163.com>

src/catalog/src/kvbackend/manager.rs

@@ -44,6 +44,7 @@ use store_api::metric_engine_consts::METRIC_ENGINE_NAME;
 use table::dist_table::DistTable;
 use table::metadata::{TableId, TableInfoRef};
 use table::table::numbers::{NumbersTable, NUMBERS_TABLE_NAME};
+use table::table::PartitionRules;
 use table::table_name::TableName;
 use table::TableRef;
 use tokio::sync::Semaphore;
@@ -132,6 +133,8 @@ impl KvBackendCatalogManager {
         {
             let mut new_table_info = (*table.table_info()).clone();
 
+            let mut phy_part_cols_not_in_logical_table = vec![];
+
             // Remap partition key indices from physical table to logical table
             new_table_info.meta.partition_key_indices = physical_table_info_value
                 .table_info
@@ -148,15 +151,30 @@ impl KvBackendCatalogManager {
                         .get(physical_index)
                         .and_then(|physical_column| {
                             // Find the corresponding index in the logical table schema
-                            new_table_info
+                            let idx = new_table_info
                                 .meta
                                 .schema
-                                .column_index_by_name(physical_column.name.as_str())
+                                .column_index_by_name(physical_column.name.as_str());
+                            if idx.is_none() {
+                                // not all part columns in physical table that are also in logical table
+                                phy_part_cols_not_in_logical_table
+                                    .push(physical_column.name.clone());
+                            }
+
+                            idx
                         })
                 })
                 .collect();
 
-            let new_table = DistTable::table(Arc::new(new_table_info));
+            let partition_rules = if !phy_part_cols_not_in_logical_table.is_empty() {
+                Some(PartitionRules {
+                    extra_phy_cols_not_in_logical_table: phy_part_cols_not_in_logical_table,
+                })
+            } else {
+                None
+            };
+
+            let new_table = DistTable::table_partitioned(Arc::new(new_table_info), partition_rules);
 
             return Ok(new_table);
         }

src/query/src/dist_plan/analyzer.rs

@@ -482,14 +482,32 @@ impl PlanRewriter {
                     .as_any()
                     .downcast_ref::<DfTableProviderAdapter>()
                 {
-                    if provider.table().table_type() == TableType::Base {
-                        let info = provider.table().table_info();
+                    let table = provider.table();
+                    if table.table_type() == TableType::Base {
+                        let info = table.table_info();
                         let partition_key_indices = info.meta.partition_key_indices.clone();
                         let schema = info.meta.schema.clone();
-                        let partition_cols = partition_key_indices
+                        let mut partition_cols = partition_key_indices
                             .into_iter()
                             .map(|index| schema.column_name_by_index(index).to_string())
                             .collect::<Vec<String>>();
+
+                        let partition_rules = table.partition_rules();
+                        let exist_phy_part_cols_not_in_logical_table = partition_rules
+                            .map(|r| !r.extra_phy_cols_not_in_logical_table.is_empty())
+                            .unwrap_or(false);
+
+                        if exist_phy_part_cols_not_in_logical_table && partition_cols.is_empty() {
+                            // there are other physical partition columns that are not in logical table and part cols are empty
+                            // so we need to add a placeholder for it to prevent certain optimization
+                            // this is used to make sure the final partition columns(that optimizer see) are not empty
+                            // notice if originally partition_cols is not empty, then there is no need to add this place holder,
+                            // as subset of phy part cols can still be used for certain optimization, and it works as if
+                            // those columns are always null
+                            // This helps with distinguishing between non-partitioned table and partitioned table with all phy part cols not in logical table
+                            partition_cols
+                                .push("__OTHER_PHYSICAL_PART_COLS_PLACEHOLDER__".to_string());
+                        }
                         self.partition_cols = Some(partition_cols);
                     }
                 }

src/table/src/dist_table.rs

@@ -21,6 +21,7 @@ use store_api::storage::ScanRequest;
 
 use crate::error::UnsupportedSnafu;
 use crate::metadata::{FilterPushDownType, TableInfoRef};
+use crate::table::PartitionRules;
 use crate::{Table, TableRef};
 
 #[derive(Clone)]
@@ -32,6 +33,20 @@ impl DistTable {
         let table = Table::new(table_info, FilterPushDownType::Inexact, data_source);
         Arc::new(table)
     }
+
+    pub fn table_partitioned(
+        table_info: TableInfoRef,
+        partition_rule: Option<PartitionRules>,
+    ) -> TableRef {
+        let data_source = Arc::new(DummyDataSource);
+        let table = Table::new_partitioned(
+            table_info,
+            FilterPushDownType::Inexact,
+            data_source,
+            partition_rule,
+        );
+        Arc::new(table)
+    }
 }
 
 pub struct DummyDataSource;

src/table/src/table.rs

@@ -52,6 +52,16 @@ lazy_static! {
     };
 }
 
+/// Defines partition rules for a table.
+/// TODO(discord9): add the entire partition exprs rules here later
+pub struct PartitionRules {
+    /// The physical partition columns that are not in the logical table.
+    /// only used in kvbackend manager to store the physical partition columns that are not in the logical table.
+    /// This is used to avoid the partition columns in the physical table that are not in the logical table
+    /// to prevent certain optimizations, if table is not a logical table, this should be empty
+    pub extra_phy_cols_not_in_logical_table: Vec<String>,
+}
+
 pub type TableRef = Arc<Table>;
 
 /// Table handle.
@@ -61,6 +71,7 @@ pub struct Table {
     data_source: DataSourceRef,
     /// Columns default [`Expr`]
     column_defaults: HashMap<String, Expr>,
+    partition_rules: Option<PartitionRules>,
 }
 
 impl Table {
@@ -74,6 +85,22 @@ impl Table {
             table_info,
             filter_pushdown,
             data_source,
+            partition_rules: None,
+        }
+    }
+
+    pub fn new_partitioned(
+        table_info: TableInfoRef,
+        filter_pushdown: FilterPushDownType,
+        data_source: DataSourceRef,
+        partition_rules: Option<PartitionRules>,
+    ) -> Self {
+        Self {
+            column_defaults: collect_column_defaults(table_info.meta.schema.column_schemas()),
+            table_info,
+            filter_pushdown,
+            data_source,
+            partition_rules,
         }
     }
 
@@ -101,6 +128,10 @@ impl Table {
         self.table_info.table_type
     }
 
+    pub fn partition_rules(&self) -> Option<&PartitionRules> {
+        self.partition_rules.as_ref()
+    }
+
     pub async fn scan_to_stream(&self, request: ScanRequest) -> Result<SendableRecordBatchStream> {
         self.data_source
             .get_stream(request)