From 4ade13d9b664670fb20951bdb45a8298d43eb854 Mon Sep 17 00:00:00 2001 From: fengyue-xve <1759519137@qq.com> Date: Fri, 9 Oct 2026 07:39:03 +0000 Subject: [PATCH] perf(catalog): optimize single-table lookups in information_schema.partitions (#9473) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * perf(catalog): optimize single-table lookups in information_schema.partitions Signed-off-by: 风月 <1759519137@qq.com> * style: apply rustfmt formatting to partitions tests Signed-off-by: 风月 <1759519137@qq.com> --------- Signed-off-by: 风月 <1759519137@qq.com> --- .../information_schema/partitions.rs | 295 ++++++++++++++++++ src/catalog/src/system_schema/predicate.rs | 100 ++++++ 2 files changed, 395 insertions(+) diff --git a/src/catalog/src/system_schema/information_schema/partitions.rs b/src/catalog/src/system_schema/information_schema/partitions.rs index d2cce93150a..1bfe3f634fd 100644 --- a/src/catalog/src/system_schema/information_schema/partitions.rs +++ b/src/catalog/src/system_schema/information_schema/partitions.rs @@ -254,6 +254,51 @@ impl InformationSchemaPartitionsBuilder { let predicates = Predicates::from_scan_request(&request); + // Fast path: when the predicates pin both the schema and the table name + // with equality, resolve the table directly instead of enumerating all + // tables in all schemas and fetching partition metadata for each of them. + if let (Some(schema_name), Some(table_name)) = ( + predicates + .find_eq_value(TABLE_SCHEMA) + .and_then(Value::as_string), + predicates + .find_eq_value(TABLE_NAME) + .and_then(Value::as_string), + ) { + if let Some(table) = catalog_manager + .table(&catalog_name, &schema_name, &table_name, None) + .await? + { + let table_info = table.table_info(); + + // The enumeration path below skips temporary tables, keep the + // behavior identical. + if table_info.table_type != TableType::Temporary { + let mut table_partitions = partition_manager + .batch_find_table_partitions(&[table_info.table_id()]) + .await + .context(FindPartitionsSnafu)?; + + let partitions = table_partitions + .remove(&table_info.table_id()) + .unwrap_or_default(); + + self.add_partitions( + &predicates, + &table_info, + &catalog_name, + &schema_name, + &table_info.name, + &partitions, + ); + } + } + + // The pinned table fully determines the result: the table above was + // found (or not) unambiguously, so no other table can contribute rows. + return self.finish(); + } + for schema_name in catalog_manager.schema_names(&catalog_name, None).await? { let table_info_stream = catalog_manager .tables(&catalog_name, &schema_name, None) @@ -421,3 +466,253 @@ impl DfPartitionStream for InformationSchemaPartitions { )) } } + +#[cfg(test)] +mod tests { + use api::v1::meta::Peer; + use cache::{build_fundamental_cache_registry, with_default_composite_cache_registry}; + use common_catalog::consts::{DEFAULT_CATALOG_NAME, INFORMATION_SCHEMA_NAME}; + use common_meta::cache::{CacheRegistryBuilder, LayeredCacheRegistryBuilder}; + use common_meta::key::schema_name::SchemaNameKey; + use common_meta::key::table_route::TableRouteValue; + use common_meta::kv_backend::memory::MemoryKvBackend; + use common_meta::rpc::router::{Region, RegionRoute}; + use common_meta::wal_provider::RegionWalOptions; + use datafusion::logical_expr::{BinaryExpr, Expr, Operator, col, lit}; + use futures::TryStreamExt; + use store_api::storage::RegionId; + + use super::*; + use crate::information_schema::NoopInformationExtension; + use crate::kvbackend::KvBackendCatalogManagerBuilder; + + const SCHEMA_1: &str = "partitions_schema_1"; + const SCHEMA_2: &str = "partitions_schema_2"; + const PARTITIONED_TABLE: &str = "partitioned_table"; + const MISSING_TABLE: &str = "missing_table"; + const TABLE_ID_1: TableId = 1; + const TABLE_ID_2: TableId = 2; + const REGION_NUMBER_1: u32 = 1; + const REGION_NUMBER_2: u32 = 2; + const REGION_NUMBER_7: u32 = 7; + + fn kv_catalog_manager() -> Arc { + let backend = Arc::new(MemoryKvBackend::default()); + let layered_cache_builder = LayeredCacheRegistryBuilder::default() + .add_cache_registry(CacheRegistryBuilder::default().build()); + let fundamental_cache_registry = build_fundamental_cache_registry(backend.clone()); + let layered_cache_registry = Arc::new( + with_default_composite_cache_registry( + layered_cache_builder.add_cache_registry(fundamental_cache_registry), + ) + .unwrap() + .build(), + ); + + KvBackendCatalogManagerBuilder::new( + Arc::new(NoopInformationExtension), + backend, + layered_cache_registry, + ) + .build() + } + + fn region_route(table_id: TableId, region_number: u32) -> RegionRoute { + RegionRoute { + region: Region { + id: RegionId::new(table_id, region_number), + ..Default::default() + }, + leader_peer: Some(Peer { + id: 1, + addr: "127.0.0.1:3001".to_string(), + }), + ..Default::default() + } + } + + fn table_info(schema: &str, table: &str, table_id: TableId) -> TableInfo { + // The catalog defaults to the default catalog name. + let mut table_info = + common_meta::key::test_utils::new_test_table_info_with_name(table_id, table); + table_info.schema_name = schema.to_string(); + table_info + } + + /// Creates a catalog with two partitioned tables named `PARTITIONED_TABLE` + /// in two different schemas: + /// + /// - `SCHEMA_1.PARTITIONED_TABLE` has two regions. + /// - `SCHEMA_2.PARTITIONED_TABLE` has one region. + async fn setup_catalog() -> Arc { + let manager = kv_catalog_manager(); + let metadata = manager.table_metadata_manager_ref().clone(); + + for schema in [SCHEMA_1, SCHEMA_2] { + metadata + .schema_manager() + .create(SchemaNameKey::new(DEFAULT_CATALOG_NAME, schema), None, true) + .await + .unwrap(); + } + + for (table_id, schema, region_numbers) in [ + (TABLE_ID_1, SCHEMA_1, vec![REGION_NUMBER_1, REGION_NUMBER_2]), + (TABLE_ID_2, SCHEMA_2, vec![REGION_NUMBER_7]), + ] { + let routes = region_numbers + .into_iter() + .map(|region_number| region_route(table_id, region_number)) + .collect(); + metadata + .create_table_metadata( + table_info(schema, PARTITIONED_TABLE, table_id), + TableRouteValue::physical(routes), + RegionWalOptions::default(), + ) + .await + .unwrap(); + } + + manager + } + + /// Scans `information_schema.partitions` and returns the total row count + /// and the pretty printed batches. + async fn scan_partitions( + manager: &Arc, + filters: Vec, + ) -> (usize, String) { + let table = manager + .table( + DEFAULT_CATALOG_NAME, + INFORMATION_SCHEMA_NAME, + PARTITIONS, + None, + ) + .await + .unwrap() + .unwrap(); + let stream = table + .scan_to_stream(ScanRequest { + filters, + ..Default::default() + }) + .await + .unwrap(); + let batches = stream.try_collect::>().await.unwrap(); + let rows = batches.iter().map(|batch| batch.num_rows()).sum(); + let output = batches + .iter() + .map(|batch| batch.pretty_print()) + .collect::>() + .join("\n"); + (rows, output) + } + + fn eq_filter(column: &str, value: &str) -> Expr { + Expr::BinaryExpr(BinaryExpr::new( + Box::new(col(column)), + Operator::Eq, + Box::new(lit(value)), + )) + } + + fn not_eq_filter(column: &str, value: &str) -> Expr { + Expr::BinaryExpr(BinaryExpr::new( + Box::new(col(column)), + Operator::NotEq, + Box::new(lit(value)), + )) + } + + #[tokio::test] + async fn test_scan_pinned_table_skips_catalog_enumeration() { + let manager: Arc = setup_catalog().await; + + // No predicates: every partitioned table is listed. + let (rows, output) = scan_partitions(&manager, vec![]).await; + assert_eq!(3, rows, "{output}"); + + // Pins both the schema and the table name, so the scan resolves the + // table directly instead of enumerating all catalog entries. + let (rows, output) = scan_partitions( + &manager, + vec![ + eq_filter(TABLE_SCHEMA, SCHEMA_1), + eq_filter(TABLE_NAME, PARTITIONED_TABLE), + ], + ) + .await; + assert_eq!(2, rows, "{output}"); + assert!(output.contains(SCHEMA_1), "{output}"); + assert!(!output.contains(SCHEMA_2), "{output}"); + assert!( + output.contains( + &RegionId::new(TABLE_ID_1, REGION_NUMBER_1) + .as_u64() + .to_string() + ), + "{output}" + ); + assert!( + output.contains( + &RegionId::new(TABLE_ID_1, REGION_NUMBER_2) + .as_u64() + .to_string() + ), + "{output}" + ); + + // The fallback enumeration path must return the identical result. The + // inequality on `table_name` doesn't pin a single name, so no direct + // lookup happens here. + let (fallback_rows, fallback_output) = scan_partitions( + &manager, + vec![ + eq_filter(TABLE_SCHEMA, SCHEMA_1), + not_eq_filter(TABLE_NAME, MISSING_TABLE), + ], + ) + .await; + assert_eq!(rows, fallback_rows, "{fallback_output}"); + assert_eq!(output, fallback_output); + + // Only the table name is pinned: the schema can't be resolved and both + // tables of that name are kept. + let (rows, output) = + scan_partitions(&manager, vec![eq_filter(TABLE_NAME, PARTITIONED_TABLE)]).await; + assert_eq!(3, rows, "{output}"); + assert!(output.contains(SCHEMA_1), "{output}"); + assert!(output.contains(SCHEMA_2), "{output}"); + + // An unknown table in a pinned schema yields no rows. + let (rows, output) = scan_partitions( + &manager, + vec![ + eq_filter(TABLE_SCHEMA, SCHEMA_1), + eq_filter(TABLE_NAME, MISSING_TABLE), + ], + ) + .await; + assert_eq!(0, rows, "{output}"); + } + + #[tokio::test] + async fn test_scan_equality_under_disjunction_falls_back_to_enumeration() { + let manager: Arc = setup_catalog().await; + + // `table_schema = SCHEMA_1 OR table_name = PARTITIONED_TABLE` matches + // the tables in both schemas: equality nested under a disjunction must + // not pin a single table. + let filter = Expr::BinaryExpr(BinaryExpr::new( + Box::new(eq_filter(TABLE_SCHEMA, SCHEMA_1)), + Operator::Or, + Box::new(eq_filter(TABLE_NAME, PARTITIONED_TABLE)), + )); + let (rows, output) = scan_partitions(&manager, vec![filter]).await; + assert_eq!(3, rows, "{output}"); + assert!(output.contains(SCHEMA_1), "{output}"); + assert!(output.contains(SCHEMA_2), "{output}"); + } +} diff --git a/src/catalog/src/system_schema/predicate.rs b/src/catalog/src/system_schema/predicate.rs index 7957ba454d4..cb71e51ed0e 100644 --- a/src/catalog/src/system_schema/predicate.rs +++ b/src/catalog/src/system_schema/predicate.rs @@ -214,6 +214,22 @@ impl Predicate { _ => None, } } + + /// Collects the values that constrain `column` to be equal to them at this + /// predicate's top-level conjunctions. + /// + /// Constraints under `Or`/`Not` are ignored because they don't guarantee + /// that `column` equals a single value. + fn collect_eq_values<'a>(&'a self, column: &str, values: &mut Vec<&'a Value>) { + match self { + Predicate::Eq(c, v) if c == column => values.push(v), + Predicate::And(left, right) => { + left.collect_eq_values(column, values); + right.collect_eq_values(column, values); + } + _ => {} + } + } } /// Perform SQL left LIKE right, return `None` if fail to evaluate. @@ -269,6 +285,25 @@ impl Predicates { } } + /// Returns the single equality value that constrains `column` in the + /// top-level conjunctions of the predicates. + /// + /// Returns `None` if there is no equality constraint on `column`, or if + /// the predicates constrain it with multiple distinct values. + pub fn find_eq_value(&self, column: &str) -> Option<&Value> { + let mut values = Vec::new(); + for predicate in &self.predicates { + predicate.collect_eq_values(column, &mut values); + } + + let (first, rest) = values.split_first()?; + if rest.iter().all(|v| v == first) { + Some(*first) + } else { + None + } + } + /// Evaluate the predicates with the row. /// returns true when all the predicates are satisfied or can't be evaluated. pub fn eval(&self, row: &[(&str, &Value)]) -> bool { @@ -556,6 +591,71 @@ mod tests { ); } + #[test] + fn test_find_eq_value() { + fn eq(name: &str, value: &str) -> Predicate { + Predicate::Eq(name.to_string(), Value::from(value)) + } + + fn and(left: Predicate, right: Predicate) -> Predicate { + Predicate::And(Box::new(left), Box::new(right)) + } + + fn or(left: Predicate, right: Predicate) -> Predicate { + Predicate::Or(Box::new(left), Box::new(right)) + } + + fn predicates(predicates: Vec) -> Predicates { + Predicates { predicates } + } + + // No predicates, no equality value. + assert!(predicates(vec![]).find_eq_value("a").is_none()); + + // No equality on the column. + let p = predicates(vec![Predicate::NotEq("a".to_string(), Value::from("x"))]); + assert!(p.find_eq_value("a").is_none()); + assert!(p.find_eq_value("b").is_none()); + + // Simple equality. + let p = predicates(vec![eq("a", "x")]); + let v = p.find_eq_value("a").unwrap(); + assert!(match_string_value(v, "x")); + assert!(p.find_eq_value("b").is_none()); + + // Duplicated equality on the same value. + let p = predicates(vec![and(eq("a", "x"), eq("a", "x"))]); + let v = p.find_eq_value("a").unwrap(); + assert!(match_string_value(v, "x")); + + // Conflicting equalities reject the lookup. + let p = predicates(vec![and(eq("a", "x"), eq("a", "y"))]); + assert!(p.find_eq_value("a").is_none()); + + // Conflicting equalities in separate predicates reject the lookup too. + let p = predicates(vec![eq("a", "x"), eq("a", "y")]); + assert!(p.find_eq_value("a").is_none()); + + // Equality nested in deeper conjunctions. + let p = predicates(vec![and(and(eq("b", "y"), eq("a", "x")), eq("a", "x"))]); + let v = p.find_eq_value("a").unwrap(); + assert!(match_string_value(v, "x")); + + // Equality under a disjunction or a negation doesn't constrain the column. + let p = predicates(vec![or(eq("a", "x"), eq("a", "x"))]); + assert!(p.find_eq_value("a").is_none()); + let p = predicates(vec![Predicate::Not(Box::new(eq("a", "x")))]); + assert!(p.find_eq_value("a").is_none()); + + // Equality value in a top-level conjunction composed of separate predicates. + let p = predicates(vec![ + eq("a", "x"), + Predicate::Like("b".to_string(), "%y".to_string(), false), + ]); + let v = p.find_eq_value("a").unwrap(); + assert!(match_string_value(v, "x")); + } + #[test] fn test_predicates_eval_row() { let wrong_row = [