catalog/
information_extension.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
15use std::sync::Arc;
16
17use api::v1::meta::ProcedureStatus;
18use common_error::ext::BoxedError;
19use common_meta::cluster::{ClusterInfo, NodeInfo, Role};
20use common_meta::datanode::RegionStat;
21use common_meta::key::flow::flow_state::FlowStat;
22use common_meta::node_manager::DatanodeManagerRef;
23use common_meta::procedure_executor::{ExecutorContext, ProcedureExecutor};
24use common_meta::rpc::procedure;
25use common_procedure::{ProcedureInfo, ProcedureState};
26use common_query::request::QueryRequest;
27use common_recordbatch::adapter::{AsyncRecordBatchStreamAdapter, DfRecordBatchStreamAdapter};
28use common_recordbatch::util::{ChainedRecordBatchStream, LimitedRecordBatchStream};
29use common_recordbatch::{DfSendableRecordBatchStream, SendableRecordBatchStream};
30use datafusion::common::Result;
31use datafusion::common::tree_node::TreeNodeRecursion;
32use datafusion::execution::TaskContext;
33use datafusion::physical_expr::{EquivalenceProperties, Partitioning, PhysicalExpr};
34use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType};
35use datafusion::physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties};
36use datatypes::arrow::datatypes::SchemaRef as ArrowSchemaRef;
37use datatypes::schema::SchemaRef;
38use futures_util::future::try_join_all;
39use meta_client::MetaClientRef;
40use snafu::ResultExt;
41use store_api::storage::RegionId;
42
43use crate::error;
44use crate::information_schema::{DatanodeInspectRequest, InformationExtension};
45
46pub struct DistributedInformationExtension {
47    meta_client: MetaClientRef,
48    datanode_manager: DatanodeManagerRef,
49}
50
51impl DistributedInformationExtension {
52    pub fn new(meta_client: MetaClientRef, datanode_manager: DatanodeManagerRef) -> Self {
53        Self {
54            meta_client,
55            datanode_manager,
56        }
57    }
58
59    async fn inspect_datanode_stream(
60        meta_client: MetaClientRef,
61        datanode_manager: DatanodeManagerRef,
62        request: DatanodeInspectRequest,
63    ) -> std::result::Result<SendableRecordBatchStream, crate::error::Error> {
64        let nodes = meta_client
65            .list_nodes(Some(Role::Datanode))
66            .await
67            .map_err(BoxedError::new)
68            .context(crate::error::ListNodesSnafu)?;
69
70        let limit = request.scan.limit;
71        let plan = request
72            .build_plan()
73            .context(crate::error::DatafusionSnafu)?;
74
75        let streams = try_join_all(nodes.into_iter().map(|node| {
76            let datanode_manager = datanode_manager.clone();
77            let plan = plan.clone();
78            async move {
79                let client = datanode_manager.datanode(&node.peer).await;
80                client
81                    .handle_query(QueryRequest {
82                        plan,
83                        region_id: RegionId::default(),
84                        header: None,
85                    })
86                    .await
87                    .context(crate::error::HandleQuerySnafu {
88                        addr: node.peer.addr,
89                    })
90            }
91        }))
92        .await?;
93
94        let chained =
95            ChainedRecordBatchStream::new(streams).context(crate::error::CreateRecordBatchSnafu)?;
96        match limit {
97            Some(limit) => Ok(Box::pin(LimitedRecordBatchStream::new(
98                Box::pin(chained),
99                limit,
100            ))),
101            None => Ok(Box::pin(chained)),
102        }
103    }
104}
105
106struct DistributedInspectExec {
107    request: DatanodeInspectRequest,
108    meta_client: MetaClientRef,
109    datanode_manager: DatanodeManagerRef,
110    schema: SchemaRef,
111    arrow_schema: ArrowSchemaRef,
112    properties: Arc<PlanProperties>,
113}
114
115impl std::fmt::Debug for DistributedInspectExec {
116    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
117        f.debug_struct("DistributedInspectExec")
118            .field("request", &self.request)
119            .field("schema", &self.schema)
120            .finish_non_exhaustive()
121    }
122}
123
124impl DistributedInspectExec {
125    fn new(
126        request: DatanodeInspectRequest,
127        schema: SchemaRef,
128        meta_client: MetaClientRef,
129        datanode_manager: DatanodeManagerRef,
130    ) -> Self {
131        let arrow_schema = schema.arrow_schema().clone();
132        let properties = Arc::new(PlanProperties::new(
133            EquivalenceProperties::new(arrow_schema.clone()),
134            Partitioning::UnknownPartitioning(1),
135            EmissionType::Incremental,
136            Boundedness::Bounded,
137        ));
138        Self {
139            request,
140            meta_client,
141            datanode_manager,
142            schema,
143            arrow_schema,
144            properties,
145        }
146    }
147}
148
149impl DisplayAs for DistributedInspectExec {
150    fn fmt_as(&self, _t: DisplayFormatType, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
151        write!(
152            f,
153            "DistributedInspectExec: kind={:?}, scan={}, schema={:?}",
154            self.request.kind, self.request.scan, self.arrow_schema
155        )
156    }
157}
158
159impl ExecutionPlan for DistributedInspectExec {
160    fn name(&self) -> &str {
161        "DistributedInspectExec"
162    }
163
164    fn schema(&self) -> ArrowSchemaRef {
165        self.arrow_schema.clone()
166    }
167
168    fn properties(&self) -> &Arc<PlanProperties> {
169        &self.properties
170    }
171
172    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
173        vec![]
174    }
175
176    fn apply_expressions(
177        &self,
178        _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
179    ) -> Result<TreeNodeRecursion> {
180        Ok(TreeNodeRecursion::Continue)
181    }
182
183    fn with_new_children(
184        self: Arc<Self>,
185        children: Vec<Arc<dyn ExecutionPlan>>,
186    ) -> datafusion::error::Result<Arc<dyn ExecutionPlan>> {
187        if children.is_empty() {
188            Ok(self)
189        } else {
190            Err(datafusion::error::DataFusionError::Internal(
191                "DistributedInspectExec is a leaf execution plan and cannot accept children"
192                    .to_string(),
193            ))
194        }
195    }
196
197    fn execute(
198        &self,
199        partition: usize,
200        _context: Arc<TaskContext>,
201    ) -> datafusion::error::Result<DfSendableRecordBatchStream> {
202        if partition != 0 {
203            return Err(datafusion::error::DataFusionError::Execution(format!(
204                "DistributedInspectExec invalid partition. Expected 0, got {partition}"
205            )));
206        }
207
208        let meta_client = self.meta_client.clone();
209        let datanode_manager = self.datanode_manager.clone();
210        let request = self.request.clone();
211        let schema = self.schema.clone();
212        let future = async move {
213            DistributedInformationExtension::inspect_datanode_stream(
214                meta_client,
215                datanode_manager,
216                request,
217            )
218            .await
219            .map_err(BoxedError::new)
220            .context(common_recordbatch::error::ExternalSnafu)
221        };
222        let stream = AsyncRecordBatchStreamAdapter::new(schema, Box::pin(future));
223        Ok(Box::pin(DfRecordBatchStreamAdapter::new(Box::pin(stream))))
224    }
225}
226
227#[async_trait::async_trait]
228impl InformationExtension for DistributedInformationExtension {
229    type Error = crate::error::Error;
230
231    async fn nodes(&self) -> std::result::Result<Vec<NodeInfo>, Self::Error> {
232        self.meta_client
233            .list_nodes(None)
234            .await
235            .map_err(BoxedError::new)
236            .context(error::ListNodesSnafu)
237    }
238
239    async fn procedures(&self) -> std::result::Result<Vec<(String, ProcedureInfo)>, Self::Error> {
240        let procedures = self
241            .meta_client
242            .list_procedures(&ExecutorContext::default())
243            .await
244            .map_err(BoxedError::new)
245            .context(error::ListProceduresSnafu)?
246            .procedures;
247        let mut result = Vec::with_capacity(procedures.len());
248        for procedure in procedures {
249            let pid = match procedure.id {
250                Some(pid) => pid,
251                None => return error::ProcedureIdNotFoundSnafu {}.fail(),
252            };
253            let pid = procedure::pb_pid_to_pid(&pid)
254                .map_err(BoxedError::new)
255                .context(error::ConvertProtoDataSnafu)?;
256            let status = ProcedureStatus::try_from(procedure.status)
257                .map(|v| v.as_str_name())
258                .unwrap_or("Unknown")
259                .to_string();
260            let procedure_info = ProcedureInfo {
261                id: pid,
262                type_name: procedure.type_name,
263                start_time_ms: procedure.start_time_ms,
264                end_time_ms: procedure.end_time_ms,
265                state: ProcedureState::Running,
266                lock_keys: procedure.lock_keys,
267            };
268            result.push((status, procedure_info));
269        }
270
271        Ok(result)
272    }
273
274    async fn region_stats(&self) -> std::result::Result<Vec<RegionStat>, Self::Error> {
275        self.meta_client
276            .list_region_stats()
277            .await
278            .map_err(BoxedError::new)
279            .context(error::ListRegionStatsSnafu)
280    }
281
282    async fn flow_stats(&self) -> std::result::Result<Option<FlowStat>, Self::Error> {
283        self.meta_client
284            .list_flow_stats()
285            .await
286            .map_err(BoxedError::new)
287            .context(crate::error::ListFlowStatsSnafu)
288    }
289
290    async fn inspect_datanode(
291        &self,
292        request: DatanodeInspectRequest,
293    ) -> std::result::Result<SendableRecordBatchStream, Self::Error> {
294        Self::inspect_datanode_stream(
295            self.meta_client.clone(),
296            self.datanode_manager.clone(),
297            request,
298        )
299        .await
300    }
301
302    fn inspect_datanode_plan(
303        &self,
304        request: DatanodeInspectRequest,
305        schema: SchemaRef,
306    ) -> std::result::Result<Option<Arc<dyn ExecutionPlan>>, Self::Error> {
307        Ok(Some(Arc::new(DistributedInspectExec::new(
308            request,
309            schema,
310            self.meta_client.clone(),
311            self.datanode_manager.clone(),
312        ))))
313    }
314}
315
316#[cfg(test)]
317mod tests {
318    use common_meta::node_manager::{DatanodeManager, DatanodeRef};
319    use common_meta::peer::Peer;
320    use datafusion::common::tree_node::TreeNode;
321    use datafusion::logical_expr::LogicalPlan;
322    use datafusion::physical_plan::sorts::sort::SortExec;
323    use meta_client::client::MetaClientBuilder;
324    use store_api::sst_entry::ManifestSstEntry;
325    use store_api::storage::ScanRequest;
326
327    use super::*;
328    use crate::information_schema::DatanodeInspectKind;
329
330    struct UnusedDatanodeManager;
331
332    #[async_trait::async_trait]
333    impl DatanodeManager for UnusedDatanodeManager {
334        async fn datanode(&self, _node: &Peer) -> DatanodeRef {
335            panic!("plan tests must not contact datanodes")
336        }
337    }
338
339    #[derive(Debug)]
340    struct ManifestScanProvider;
341
342    #[async_trait::async_trait]
343    impl datafusion::catalog::TableProvider for ManifestScanProvider {
344        fn schema(&self) -> ArrowSchemaRef {
345            ManifestSstEntry::schema().arrow_schema().clone()
346        }
347
348        fn table_type(&self) -> datafusion::logical_expr::TableType {
349            datafusion::logical_expr::TableType::View
350        }
351
352        fn supports_filters_pushdown(
353            &self,
354            filters: &[&datafusion::logical_expr::Expr],
355        ) -> Result<Vec<datafusion::logical_expr::TableProviderFilterPushDown>> {
356            Ok(vec![datafusion::logical_expr::TableProviderFilterPushDown::Inexact; filters.len()])
357        }
358
359        async fn scan(
360            &self,
361            _state: &dyn datafusion::catalog::Session,
362            projection: Option<&Vec<usize>>,
363            filters: &[datafusion::logical_expr::Expr],
364            limit: Option<usize>,
365        ) -> Result<Arc<dyn ExecutionPlan>> {
366            let schema = ManifestSstEntry::schema();
367            let schema = match projection {
368                Some(projection) => Arc::new(schema.try_project(projection).unwrap()),
369                None => schema,
370            };
371            Ok(Arc::new(DistributedInspectExec::new(
372                DatanodeInspectRequest {
373                    kind: DatanodeInspectKind::SstManifest,
374                    scan: ScanRequest {
375                        projection: projection.cloned(),
376                        filters: filters.to_vec(),
377                        limit,
378                        ..Default::default()
379                    },
380                },
381                schema,
382                Arc::new(MetaClientBuilder::new(0, api::v1::meta::Role::Frontend).build()),
383                Arc::new(UnusedDatanodeManager),
384            )))
385        }
386    }
387
388    #[tokio::test]
389    async fn test_sql_manifest_sort_stays_on_frontend() {
390        let ctx = datafusion::prelude::SessionContext::new();
391        ctx.register_table("ssts_manifest", Arc::new(ManifestScanProvider))
392            .unwrap();
393        for sql in [
394            "SELECT region_id, file_id, min_ts, max_ts, file_size FROM ssts_manifest WHERE region_id = 42 ORDER BY min_ts LIMIT 100",
395            "SELECT file_id FROM ssts_manifest WHERE region_id = 42 ORDER BY max_ts DESC LIMIT 10 OFFSET 3",
396            "SELECT min_ts AS ts, file_id FROM ssts_manifest WHERE region_id = 42 ORDER BY ts DESC NULLS FIRST LIMIT 10",
397            "SELECT file_id FROM ssts_manifest WHERE region_id = 42 ORDER BY min_ts, max_ts DESC LIMIT 10",
398        ] {
399            let plan = ctx
400                .sql(sql)
401                .await
402                .unwrap()
403                .create_physical_plan()
404                .await
405                .unwrap();
406            let mut found_scan = false;
407            let mut found_sort = false;
408            plan.apply(|plan| {
409                found_sort |= plan.is::<SortExec>();
410                if let Some(scan) = plan.downcast_ref::<DistributedInspectExec>() {
411                    found_scan = true;
412                    // All matching rows must reach the frontend global sort.
413                    assert_eq!(scan.request.scan.limit, None);
414                    let remote = scan.request.clone().build_plan().unwrap();
415                    remote.apply(|node| {
416                        assert!(!matches!(
417                            node,
418                            LogicalPlan::Sort(_) | LogicalPlan::Limit(_)
419                        ));
420                        Ok(TreeNodeRecursion::Continue)
421                    })?;
422                }
423                Ok(TreeNodeRecursion::Continue)
424            })
425            .unwrap();
426            assert!(
427                found_scan && found_sort,
428                "{}",
429                datafusion::physical_plan::displayable(plan.as_ref()).indent(true)
430            );
431        }
432    }
433}
catalog/information_extension.rs

catalog/
information_extension.rs
