From 32be67c0242c19dc43f01fb5b4dffb51b28ef1a4 Mon Sep 17 00:00:00 2001 From: discord9 Date: Thu, 30 Jul 2026 15:05:38 +0800 Subject: [PATCH] fix(query): use physical partition types for metric route pruning (#8590) * fix(query): prune logical metric regions after repartition Signed-off-by: discord9 * fix(query): centralize conservative region pruning Signed-off-by: discord9 * refactor(query): extract pruning metadata helper Signed-off-by: discord9 * docs(query): explain pruning metadata helper Signed-off-by: discord9 * docs(query): clarify pruning metadata mismatch Signed-off-by: discord9 --------- Signed-off-by: discord9 --- Cargo.lock | 1 + src/partition/src/manager.rs | 20 +- src/query/Cargo.toml | 1 + src/query/src/dist_plan/planner.rs | 395 +++++++++++++++++++++-- src/query/src/dist_plan/region_pruner.rs | 57 +++- 5 files changed, 434 insertions(+), 40 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index aa1ac6687d..1faabf151f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11544,6 +11544,7 @@ dependencies = [ "log-query", "meter-core", "meter-macros", + "moka", "nalgebra", "num", "num-traits", diff --git a/src/partition/src/manager.rs b/src/partition/src/manager.rs index cc00f9caba..89bb61c417 100644 --- a/src/partition/src/manager.rs +++ b/src/partition/src/manager.rs @@ -74,6 +74,14 @@ impl PartitionRuleManager { &self, table_id: TableId, ) -> Result> { + Ok(self.find_physical_table_route_with_id(table_id).await?.1) + } + + /// Returns the resolved physical table id and its route. + pub async fn find_physical_table_route_with_id( + &self, + table_id: TableId, + ) -> Result<(TableId, Arc)> { match self .table_route_cache .get(table_id) @@ -82,7 +90,9 @@ impl PartitionRuleManager { .context(error::TableRouteNotFoundSnafu { table_id })? .as_ref() { - TableRoute::Physical(physical_table_route) => Ok(physical_table_route.clone()), + TableRoute::Physical(physical_table_route) => { + Ok((table_id, physical_table_route.clone())) + } TableRoute::Logical(logical_table_route) => { let physical_table_id = logical_table_route.physical_table_id(); let physical_table_route = self @@ -90,17 +100,19 @@ impl PartitionRuleManager { .get(physical_table_id) .await .context(error::TableRouteManagerSnafu)? - .context(error::TableRouteNotFoundSnafu { table_id })?; + .context(error::TableRouteNotFoundSnafu { + table_id: physical_table_id, + })?; let physical_table_route = physical_table_route .as_physical_table_route_ref() .context(error::UnexpectedSnafu{ err_msg: format!( - "Expected the physical table route, but got logical table route, table: {table_id}" + "Expected the physical table route, but got logical table route, table: {physical_table_id}" ), })?; - Ok(physical_table_route.clone()) + Ok((physical_table_id, physical_table_route.clone())) } } } diff --git a/src/query/Cargo.toml b/src/query/Cargo.toml index 9d9f8b24b0..3ce1b3df39 100644 --- a/src/query/Cargo.toml +++ b/src/query/Cargo.toml @@ -87,6 +87,7 @@ common-function.workspace = true common-macro.workspace = true common-query = { workspace = true, features = ["testing"] } fastrand = "2.0" +moka = { workspace = true, features = ["future"] } nalgebra.workspace = true num = "0.4" num-traits = "0.2" diff --git a/src/query/src/dist_plan/planner.rs b/src/query/src/dist_plan/planner.rs index f3b7851fde..999515df73 100644 --- a/src/query/src/dist_plan/planner.rs +++ b/src/query/src/dist_plan/planner.rs @@ -31,10 +31,13 @@ use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion, TreeNodeVisitor} use datafusion_common::{DataFusionError, TableReference}; use datafusion_expr::{LogicalPlan, UserDefinedLogicalNode}; use datafusion_physical_expr::{LexOrdering, PhysicalSortExpr}; +use datatypes::prelude::ConcreteDataType; +use partition::expr::PartitionExpr; use partition::manager::{PartitionRuleManagerRef, create_partitions_from_region_routes}; use session::context::QueryContext; use snafu::{OptionExt, ResultExt}; use store_api::storage::RegionId; +use table::metadata::TableInfo; pub use table::metadata::TableType; use table::table::adapter::DfTableProviderAdapter; use table::table_name::TableName; @@ -242,9 +245,9 @@ impl DistExtensionPlanner { })?; let table_info = table.table_info(); - let physical_table_route = self + let (physical_table_id, physical_table_route) = self .partition_rule_manager - .find_physical_table_route(table_info.table_id()) + .find_physical_table_route_with_id(table_info.table_id()) .await .context(PartitionRuleManagerSnafu)?; let all_regions = physical_table_route @@ -252,9 +255,11 @@ impl DistExtensionPlanner { .iter() .map(|r| RegionId::new(table_info.table_id(), r.region.id.region_number())) .collect::>(); - // Extract partition columns - let partition_columns: Vec = - table_info.meta.partition_column_names().cloned().collect(); + let logical_partition_columns = partition_column_types(&table_info); + let partition_columns = logical_partition_columns + .iter() + .map(|(name, _)| name.clone()) + .collect::>(); debug!( "DistExtensionPlanner: loaded table partition metadata, table: {}, table_id: {}, partition_key_indices: {:?}, partition_columns: {:?}, all_regions: {:?}", table_name, @@ -266,21 +271,6 @@ impl DistExtensionPlanner { if partition_columns.is_empty() { return Ok(all_regions); } - let partition_column_types = partition_columns - .iter() - .map(|col_name| { - let data_type = table_info - .meta - .schema - .column_schema_by_name(col_name) - // Safety: names are retrieved above from the same table - .unwrap() - .data_type - .clone(); - (col_name.clone(), data_type) - }) - .collect::>(); - // Extract predicates from logical plan let partition_expressions = match PredicateExtractor::extract_partition_expressions( logical_plan, @@ -302,6 +292,19 @@ impl DistExtensionPlanner { return Ok(all_regions); } + let Some(partition_column_types) = self + .partition_column_types_for_pruning( + table_name, + table_info.as_ref(), + physical_table_id, + &partition_expressions, + &all_regions, + ) + .await + else { + return Ok(all_regions); + }; + // Get partition information for the table if partition rule manager is available let partitions = match create_partitions_from_region_routes( table_info.table_id(), @@ -320,7 +323,6 @@ impl DistExtensionPlanner { if partitions.is_empty() { return Ok(all_regions); } - // Apply region pruning based on partition rules let pruned_regions = match ConstraintPruner::prune_regions( &partition_expressions, @@ -349,6 +351,72 @@ impl DistExtensionPlanner { Ok(pruned_regions) } + /// Resolves the partition-column types that are safe to use for region pruning. + /// + /// A logical metric table may not contain every physical partition column, either for backward + /// compatibility or because its physical table was repartitioned after the logical table was + /// created. Predicate extraction must remain bounded by the logical schema, while pruning + /// needs the physical datatypes to evaluate route expressions. Any lookup failure or + /// logical/physical datatype mismatch returns `None`, causing the caller to scan all regions. + async fn partition_column_types_for_pruning( + &self, + table_name: &TableName, + logical_table_info: &TableInfo, + physical_table_id: u32, + partition_expressions: &[PartitionExpr], + all_regions: &[RegionId], + ) -> Option> { + let physical_partition_columns = if physical_table_id == logical_table_info.table_id() { + partition_column_types(logical_table_info) + } else { + match self + .catalog_manager + .table_info_by_id(physical_table_id) + .await + { + Ok(Some(physical_table_info)) => { + partition_column_types(physical_table_info.as_ref()) + } + Ok(None) => { + debug!( + "DistExtensionPlanner: physical table info not found for table {} (id: {}), using all regions: {:?}", + table_name, physical_table_id, all_regions + ); + return None; + } + Err(err) => { + debug!( + "DistExtensionPlanner: failed to load physical table info for table {} (id: {}): {}, using all regions: {:?}", + table_name, physical_table_id, err, all_regions + ); + return None; + } + } + }; + let physical_column_types = physical_partition_columns + .into_iter() + .collect::>(); + let logical_column_types = partition_column_types(logical_table_info) + .into_iter() + .collect::>(); + let mut predicate_column_names = std::collections::HashSet::new(); + for expression in partition_expressions { + expression.collect_column_names(&mut predicate_column_names); + } + if predicate_column_names + .iter() + .any(|name| logical_column_types.get(name) != physical_column_types.get(name)) + { + debug!( + "DistExtensionPlanner: logical and physical partition metadata mismatch for table {} (physical id: {}), using all regions: {:?}", + table_name, physical_table_id, all_regions + ); + return None; + } + + Some(physical_column_types) + } + /// Input logical plan is analyzed. Thus only call logical optimizer to optimize it. fn optimize_input_logical_plan( &self, @@ -360,6 +428,14 @@ impl DistExtensionPlanner { } } +fn partition_column_types(table_info: &TableInfo) -> Vec<(String, ConcreteDataType)> { + table_info + .meta + .partition_columns() + .map(|column| (column.name.clone(), column.data_type.clone())) + .collect() +} + /// Visitor to extract table name from logical plan (TableScan node) #[derive(Default)] struct TableNameExtractor { @@ -424,3 +500,280 @@ impl TreeNodeVisitor<'_> for TableNameExtractor { } } } + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + use std::sync::Arc; + + use api::v1::region::{RemoteDynFilterUnregister, RemoteDynFilterUpdate}; + use async_trait::async_trait; + use catalog::memory::MemoryCatalogManager; + use catalog::{CatalogManagerRef, RegisterTableRequest}; + use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; + use common_meta::cache::new_table_route_cache; + use common_meta::key::TableMetadataManager; + use common_meta::key::table_route::TableRouteValue; + use common_meta::kv_backend::memory::MemoryKvBackend; + use common_meta::rpc::router::{Region, RegionRoute}; + use common_query::request::QueryRequest; + use common_recordbatch::SendableRecordBatchStream; + use datafusion::datasource::DefaultTableSource; + use datafusion_expr::{LogicalPlan, LogicalPlanBuilder, col as df_col, lit}; + use datatypes::prelude::ConcreteDataType; + use datatypes::schema::{ColumnSchema, Schema}; + use datatypes::value::Value; + use moka::future::CacheBuilder; + use partition::cache::new_partition_info_cache; + use partition::expr::{PartitionExpr, col as partition_col}; + use partition::manager::PartitionRuleManager; + use session::ReadPreference; + use store_api::storage::RegionId; + use table::metadata::{TableInfo, TableInfoBuilder, TableMeta, TableType}; + use table::table::adapter::DfTableProviderAdapter; + use table::table_name::TableName; + use table::test_util::EmptyTable; + + use super::DistExtensionPlanner; + use crate::region_query::RegionQueryHandler; + + const LOGICAL_TABLE_ID: u32 = 1024; + const PHYSICAL_TABLE_ID: u32 = 2048; + + struct UnusedRegionQueryHandler; + + #[async_trait] + impl RegionQueryHandler for UnusedRegionQueryHandler { + async fn do_get( + &self, + _read_preference: ReadPreference, + _request: QueryRequest, + ) -> crate::error::Result { + unreachable!("get_regions does not query regions") + } + + async fn handle_remote_dyn_filter_update( + &self, + _region_id: RegionId, + _query_id: String, + _update: RemoteDynFilterUpdate, + ) -> crate::error::Result<()> { + unreachable!("get_regions does not update dynamic filters") + } + + async fn handle_remote_dyn_filter_unregister( + &self, + _region_id: RegionId, + _query_id: String, + _unregister: RemoteDynFilterUnregister, + ) -> crate::error::Result<()> { + unreachable!("get_regions does not unregister dynamic filters") + } + } + + fn table_info( + table_id: u32, + name: &str, + columns: &[&str], + partition_keys: Vec, + ) -> TableInfo { + let schema = Arc::new(Schema::new( + columns + .iter() + .map(|name| ColumnSchema::new(*name, ConcreteDataType::string_datatype(), true)) + .collect(), + )); + let meta = TableMeta { + schema, + primary_key_indices: vec![], + value_indices: vec![], + engine: "metric".to_string(), + next_column_id: columns.len() as u32, + options: Default::default(), + created_on: Default::default(), + updated_on: Default::default(), + partition_key_indices: partition_keys, + column_ids: (0..columns.len() as u32).collect(), + }; + TableInfoBuilder::default() + .table_id(table_id) + .table_version(0) + .name(name.to_string()) + .catalog_name(DEFAULT_CATALOG_NAME.to_string()) + .schema_name(DEFAULT_SCHEMA_NAME.to_string()) + .desc(None) + .table_type(TableType::Base) + .meta(meta) + .build() + .unwrap() + } + + fn region_route(region_number: u32, expression: Option) -> RegionRoute { + RegionRoute { + region: Region { + id: RegionId::new(PHYSICAL_TABLE_ID, region_number), + partition_expr: expression + .map(|expression| expression.as_json_str().unwrap()) + .unwrap_or_default(), + ..Default::default() + }, + ..Default::default() + } + } + + async fn planner_and_plan( + physical_partition_keys: Vec, + expressions: Vec>, + ) -> (DistExtensionPlanner, LogicalPlan, TableName) { + let logical_info = table_info(LOGICAL_TABLE_ID, "logical", &["host"], vec![0]); + let physical_info = table_info( + PHYSICAL_TABLE_ID, + "physical", + &["host", "rack"], + physical_partition_keys, + ); + let logical_table = EmptyTable::from_table_info(&logical_info); + let physical_table = EmptyTable::from_table_info(&physical_info); + let catalog_manager = MemoryCatalogManager::with_default_setup(); + for table in [&logical_table, &physical_table] { + let info = table.table_info(); + catalog_manager + .register_table_sync(RegisterTableRequest { + catalog: info.catalog_name.clone(), + schema: info.schema_name.clone(), + table_name: info.name.clone(), + table_id: info.table_id(), + table: table.clone(), + }) + .unwrap(); + } + + let backend = Arc::new(MemoryKvBackend::default()); + let metadata_manager = TableMetadataManager::new(backend.clone()); + let routes = expressions + .into_iter() + .enumerate() + .map(|(index, expression)| region_route(index as u32 + 1, expression)) + .collect(); + metadata_manager + .create_table_metadata( + physical_info, + TableRouteValue::physical(routes), + HashMap::new(), + ) + .await + .unwrap(); + metadata_manager + .create_table_metadata( + logical_info, + TableRouteValue::logical(PHYSICAL_TABLE_ID), + HashMap::new(), + ) + .await + .unwrap(); + + let table_route_cache = Arc::new(new_table_route_cache( + "planner-test-routes".to_string(), + CacheBuilder::new(16).build(), + backend.clone(), + )); + let partition_info_cache = Arc::new(new_partition_info_cache( + "planner-test-partitions".to_string(), + CacheBuilder::new(16).build(), + table_route_cache.clone(), + )); + let partition_rule_manager = Arc::new(PartitionRuleManager::new( + backend, + table_route_cache, + partition_info_cache, + )); + let (resolved_physical_id, physical_route) = partition_rule_manager + .find_physical_table_route_with_id(PHYSICAL_TABLE_ID) + .await + .unwrap(); + assert_eq!(PHYSICAL_TABLE_ID, resolved_physical_id); + let (resolved_logical_id, logical_route) = partition_rule_manager + .find_physical_table_route_with_id(LOGICAL_TABLE_ID) + .await + .unwrap(); + assert_eq!(PHYSICAL_TABLE_ID, resolved_logical_id); + assert_eq!(physical_route.region_routes, logical_route.region_routes); + let catalog_manager: CatalogManagerRef = catalog_manager; + let planner = DistExtensionPlanner::new( + catalog_manager, + partition_rule_manager, + Arc::new(UnusedRegionQueryHandler), + false, + ); + let table_source = Arc::new(DefaultTableSource::new(Arc::new( + DfTableProviderAdapter::new(logical_table), + ))); + let plan = LogicalPlanBuilder::scan_with_filters("logical", table_source, None, vec![]) + .unwrap() + .filter(df_col("host").eq(lit("a"))) + .unwrap() + .build() + .unwrap(); + ( + planner, + plan, + TableName::new(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, "logical"), + ) + } + + fn physical_partition_expressions() -> Vec> { + vec![ + Some(partition_col("host").lt(Value::String("m".into()))), + Some( + partition_col("host") + .gt_eq(Value::String("m".into())) + .and(partition_col("rack").lt(Value::String("n".into()))), + ), + Some( + partition_col("host") + .gt_eq(Value::String("m".into())) + .and(partition_col("rack").gt_eq(Value::String("n".into()))), + ), + ] + } + + #[tokio::test] + async fn logical_table_pruning_uses_physical_partition_datatypes() { + let (planner, plan, table_name) = + planner_and_plan(vec![0, 1], physical_partition_expressions()).await; + + assert_eq!( + vec![RegionId::new(LOGICAL_TABLE_ID, 1)], + planner.get_regions(&table_name, &plan).await.unwrap() + ); + } + + #[tokio::test] + async fn missing_physical_partition_datatype_falls_back_to_all_logical_regions() { + let (planner, plan, table_name) = + planner_and_plan(vec![0], physical_partition_expressions()).await; + + assert_all_logical_regions(planner.get_regions(&table_name, &plan).await.unwrap()); + } + + #[tokio::test] + async fn missing_route_partition_expression_falls_back_to_all_logical_regions() { + let mut expressions = physical_partition_expressions(); + expressions[1] = None; + let (planner, plan, table_name) = planner_and_plan(vec![0, 1], expressions).await; + + assert_all_logical_regions(planner.get_regions(&table_name, &plan).await.unwrap()); + } + + fn assert_all_logical_regions(mut regions: Vec) { + regions.sort_unstable(); + assert_eq!( + vec![ + RegionId::new(LOGICAL_TABLE_ID, 1), + RegionId::new(LOGICAL_TABLE_ID, 2), + RegionId::new(LOGICAL_TABLE_ID, 3), + ], + regions + ); + } +} diff --git a/src/query/src/dist_plan/region_pruner.rs b/src/query/src/dist_plan/region_pruner.rs index 3304bf23df..eca6a60ebd 100644 --- a/src/query/src/dist_plan/region_pruner.rs +++ b/src/query/src/dist_plan/region_pruner.rs @@ -38,29 +38,32 @@ impl ConstraintPruner { column_datatypes: HashMap, ) -> Result> { let start = std::time::Instant::now(); + let all_regions = partitions + .iter() + .map(|partition| partition.id) + .collect::>(); if query_expressions.is_empty() || partitions.is_empty() { // No constraints, return all regions - return Ok(partitions.iter().map(|p| p.id).collect()); + return Ok(all_regions); } // Collect all partition expressions for unified normalization - let mut expression_to_partition = Vec::with_capacity(partitions.len()); - let mut all_partition_expressions = Vec::with_capacity(partitions.len()); - for partition in partitions { - if let Some(expr) = &partition.partition_expr { - expression_to_partition.push(partition.id); - all_partition_expressions.push(expr.clone()); - } - } - if all_partition_expressions.is_empty() { - return Ok(partitions.iter().map(|p| p.id).collect()); - } + let Some(all_partition_expressions) = partitions + .iter() + .map(|partition| partition.partition_expr.clone()) + .collect::>>() + else { + debug!( + "Partition metadata contains a missing partition expression, returning all regions conservatively" + ); + return Ok(all_regions); + }; // Create unified collider with both query and partition expressions for consistent normalization let mut all_expressions = query_expressions.to_vec(); all_expressions.extend(all_partition_expressions.iter().cloned()); if !Self::normalize_datatype(&mut all_expressions, &column_datatypes) { - return Ok(partitions.iter().map(|p| p.id).collect()); + return Ok(all_regions); } let collider = match Collider::new(&all_expressions) { @@ -70,7 +73,7 @@ impl ConstraintPruner { "Failed to create unified collider: {}, returning all regions conservatively", err ); - return Ok(partitions.iter().map(|p| p.id).collect()); + return Ok(all_regions); } }; @@ -91,7 +94,7 @@ impl ConstraintPruner { if Self::atomic_sets_overlap(&query_atomics, region_atomics) { let partition_expr_index = region_atomics.source_expr_index - query_expressions.len(); - candidate_regions.insert(expression_to_partition[partition_expr_index]); + candidate_regions.insert(all_regions[partition_expr_index]); } } @@ -384,6 +387,30 @@ mod tests { assert_eq!(pruned.len(), 2); } + #[test] + fn test_missing_partition_expression_returns_all_regions() { + let partitions = vec![ + create_test_partition_info(1, Some(col("user_id").lt(Value::Int64(100)))), + create_test_partition_info(2, None), + create_test_partition_info(3, Some(col("user_id").gt_eq(Value::Int64(200)))), + ]; + let query_exprs = vec![col("user_id").eq(Value::Int64(150))]; + let mut column_datatypes = HashMap::default(); + column_datatypes.insert("user_id".to_string(), ConcreteDataType::int64_datatype()); + + let pruned = + ConstraintPruner::prune_regions(&query_exprs, &partitions, column_datatypes).unwrap(); + + assert_eq!( + vec![ + RegionId::new(1, 1), + RegionId::new(1, 2), + RegionId::new(1, 3), + ], + pruned + ); + } + #[test] fn test_prune_regions_with_simple_equality() { let partitions = vec![