perf(catalog): optimize single-table lookups in information_schema.partitions (#9473)

* 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>

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<KvBackendCatalogManager> {
+        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<KvBackendCatalogManager> {
+        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<dyn CatalogManager>,
+        filters: Vec<Expr>,
+    ) -> (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::<Vec<_>>().await.unwrap();
+        let rows = batches.iter().map(|batch| batch.num_rows()).sum();
+        let output = batches
+            .iter()
+            .map(|batch| batch.pretty_print())
+            .collect::<Vec<_>>()
+            .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<dyn CatalogManager> = 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<dyn CatalogManager> = 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}");
+    }
+}

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<Predicate>) -> 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 = [