catalog/system_schema/
utils.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15pub mod tables;
16
17use std::sync::{Arc, Weak};
18
19use common_meta::key::TableMetadataManagerRef;
20use common_query::logical_plan::is_lossless_substrait_type;
21use datafusion::common::DFSchema;
22use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion};
23use datafusion::logical_expr::ExprSchemable;
24use datafusion::physical_plan::ExecutionPlan;
25use datatypes::schema::SchemaRef;
26use snafu::{OptionExt, ResultExt};
27use store_api::storage::ScanRequest;
28
29use crate::CatalogManager;
30use crate::error::{
31    GetInformationExtensionSnafu, ProjectSchemaSnafu, Result, UpgradeWeakCatalogManagerRefSnafu,
32};
33use crate::information_schema::{
34    DatanodeInspectKind, DatanodeInspectRequest, InformationExtensionRef,
35};
36use crate::kvbackend::KvBackendCatalogManager;
37use crate::system_schema::semantic_graph::EntityGraphProviderRef;
38
39/// Projects an inspection table schema according to the scan request.
40pub(crate) fn projected_schema(schema: &SchemaRef, request: &ScanRequest) -> Result<SchemaRef> {
41    if let Some(p) = request.projection.as_deref() {
42        Ok(Arc::new(schema.try_project(p).context(ProjectSchemaSnafu)?))
43    } else {
44        Ok(schema.clone())
45    }
46}
47
48/// Builds a datanode inspection plan with the requested output schema.
49pub(crate) fn inspect_plan(
50    catalog_manager: &Weak<dyn CatalogManager>,
51    kind: DatanodeInspectKind,
52    schema: &SchemaRef,
53    request: ScanRequest,
54) -> Result<Option<Arc<dyn ExecutionPlan>>> {
55    let projected_schema = projected_schema(schema, &request)?;
56    let info_ext = information_extension(catalog_manager)?;
57    let req = inspect_request(kind, schema, request)?;
58
59    info_ext.inspect_datanode_plan(req, projected_schema)
60}
61
62/// System table predicates are inexact: the frontend keeps a residual filter. Dropping
63/// an unsafe remote filter preserves correctness only if no remote limit truncates it.
64pub(crate) fn inspect_request(
65    kind: DatanodeInspectKind,
66    schema: &SchemaRef,
67    mut scan: ScanRequest,
68) -> Result<DatanodeInspectRequest> {
69    let schema =
70        DFSchema::try_from(schema.arrow_schema().clone()).context(crate::error::DatafusionSnafu)?;
71    let filter_count = scan.filters.len();
72    scan.filters.retain(|filter| {
73        let mut supported = true;
74        let visited = filter.apply(|expr| {
75            if !expr
76                .get_type(&schema)
77                .is_ok_and(|data_type| is_lossless_substrait_type(&data_type))
78            {
79                supported = false;
80                return Ok(TreeNodeRecursion::Stop);
81            }
82            Ok(TreeNodeRecursion::Continue)
83        });
84        visited.is_ok() && supported
85    });
86    if scan.filters.len() != filter_count {
87        scan.limit = None;
88    }
89    Ok(DatanodeInspectRequest { kind, scan })
90}
91
92/// Try to get the entity-graph provider from a `[CatalogManager]` weak reference.
93/// Returns `None` when the manager is not the kv-backed one or the provider has
94/// not been injected yet (the computed graph tables then stream empty).
95pub fn entity_graph_provider(
96    catalog_manager: &Weak<dyn CatalogManager>,
97) -> Result<Option<EntityGraphProviderRef>> {
98    let catalog_manager = catalog_manager
99        .upgrade()
100        .context(UpgradeWeakCatalogManagerRefSnafu)?;
101
102    Ok(catalog_manager
103        .as_any()
104        .downcast_ref::<KvBackendCatalogManager>()
105        .and_then(|manager| manager.entity_graph_provider()))
106}
107
108/// Try to get the `[InformationExtension]` from `[CatalogManager]` weak reference.
109pub fn information_extension(
110    catalog_manager: &Weak<dyn CatalogManager>,
111) -> Result<InformationExtensionRef> {
112    let catalog_manager = catalog_manager
113        .upgrade()
114        .context(UpgradeWeakCatalogManagerRefSnafu)?;
115
116    let information_extension = catalog_manager
117        .as_any()
118        .downcast_ref::<KvBackendCatalogManager>()
119        .map(|manager| manager.information_extension())
120        .context(GetInformationExtensionSnafu)?;
121
122    Ok(information_extension)
123}
124
125/// Try to get the `[TableMetadataManagerRef]` from `[CatalogManager]` weak reference.
126pub fn table_meta_manager(
127    catalog_manager: &Weak<dyn CatalogManager>,
128) -> Result<Option<TableMetadataManagerRef>> {
129    let catalog_manager = catalog_manager
130        .upgrade()
131        .context(UpgradeWeakCatalogManagerRefSnafu)?;
132
133    Ok(catalog_manager
134        .as_any()
135        .downcast_ref::<KvBackendCatalogManager>()
136        .map(|manager| manager.table_metadata_manager_ref().clone()))
137}
catalog/system_schema/utils.rs

catalog/system_schema/
utils.rs
