mirror of
https://github.com/GreptimeTeam/greptimedb.git
synced 2026-08-18 12:08:22 +00:00
fix(query): use physical partition types for metric route pruning (#8590)
* fix(query): prune logical metric regions after repartition Signed-off-by: discord9 <discord9@163.com> * fix(query): centralize conservative region pruning Signed-off-by: discord9 <discord9@163.com> * refactor(query): extract pruning metadata helper Signed-off-by: discord9 <discord9@163.com> * docs(query): explain pruning metadata helper Signed-off-by: discord9 <discord9@163.com> * docs(query): clarify pruning metadata mismatch Signed-off-by: discord9 <discord9@163.com> --------- Signed-off-by: discord9 <discord9@163.com>
This commit is contained in:
discord9
GitHub
Generated
+1
@@ -11544,6 +11544,7 @@ dependencies = [
|
||||
"log-query",
|
||||
"meter-core",
|
||||
"meter-macros",
|
||||
"moka",
|
||||
"nalgebra",
|
||||
"num",
|
||||
"num-traits",
|
||||
|
||||
@@ -74,6 +74,14 @@ impl PartitionRuleManager {
|
||||
&self,
|
||||
table_id: TableId,
|
||||
) -> Result<Arc<PhysicalTableRouteValue>> {
|
||||
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<PhysicalTableRouteValue>)> {
|
||||
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()))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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::<Vec<_>>();
|
||||
// Extract partition columns
|
||||
let partition_columns: Vec<String> =
|
||||
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::<Vec<_>>();
|
||||
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::<HashMap<_, _>>();
|
||||
|
||||
// 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<HashMap<String, ConcreteDataType>> {
|
||||
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::<HashMap<_, _>>();
|
||||
let logical_column_types = partition_column_types(logical_table_info)
|
||||
.into_iter()
|
||||
.collect::<HashMap<_, _>>();
|
||||
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<SendableRecordBatchStream> {
|
||||
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<usize>,
|
||||
) -> 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<PartitionExpr>) -> 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<usize>,
|
||||
expressions: Vec<Option<PartitionExpr>>,
|
||||
) -> (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<Option<PartitionExpr>> {
|
||||
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<RegionId>) {
|
||||
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
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -38,29 +38,32 @@ impl ConstraintPruner {
|
||||
column_datatypes: HashMap<String, ConcreteDataType>,
|
||||
) -> Result<Vec<RegionId>> {
|
||||
let start = std::time::Instant::now();
|
||||
let all_regions = partitions
|
||||
.iter()
|
||||
.map(|partition| partition.id)
|
||||
.collect::<Vec<_>>();
|
||||
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::<Option<Vec<_>>>()
|
||||
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![
|
||||
|
||||
Reference in New Issue
Block a user