catalog/system_schema/information_schema/
flow_statistics.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, Weak};
16
17use common_catalog::consts::INFORMATION_SCHEMA_FLOW_STATISTICS_TABLE_ID;
18use common_error::ext::BoxedError;
19use common_meta::key::FlowId;
20use common_meta::key::flow::FlowMetadataManager;
21use common_meta::key::flow::flow_state::FlowStat;
22use common_recordbatch::adapter::RecordBatchStreamAdapter;
23use common_recordbatch::{DfSendableRecordBatchStream, RecordBatch, SendableRecordBatchStream};
24use common_time::util::current_time_millis;
25use datafusion::execution::TaskContext;
26use datafusion::physical_plan::stream::RecordBatchStreamAdapter as DfRecordBatchStreamAdapter;
27use datafusion::physical_plan::streaming::PartitionStream as DfPartitionStream;
28use datatypes::prelude::ConcreteDataType as CDT;
29use datatypes::scalars::ScalarVectorBuilder;
30use datatypes::schema::{ColumnSchema, Schema, SchemaRef};
31use datatypes::timestamp::TimestampMillisecond;
32use datatypes::value::Value;
33use datatypes::vectors::{
34    Int64VectorBuilder, StringVectorBuilder, TimestampMillisecondVectorBuilder,
35    UInt32VectorBuilder, UInt64VectorBuilder, VectorRef,
36};
37use futures::TryStreamExt;
38use snafu::ResultExt;
39use store_api::storage::{ScanRequest, TableId};
40
41use crate::CatalogManager;
42use crate::error::{CreateRecordBatchSnafu, InternalSnafu, ListFlowsSnafu, Result};
43use crate::information_schema::{FLOW_STATISTICS, Predicates};
44use crate::system_schema::information_schema::InformationTable;
45use crate::system_schema::utils;
46
47const INIT_CAPACITY: usize = 42;
48
49// rows of information_schema.flow_statistics
50pub const FLOW_ID: &str = "flow_id";
51pub const FLOW_NAME: &str = "flow_name";
52pub const START_TIME: &str = "start_time";
53pub const LAST_EXECUTION_TIME: &str = "last_execution_time";
54pub const UPTIME_SECONDS: &str = "uptime_seconds";
55pub const STATE_SIZE: &str = "state_size";
56
57/// The `information_schema.flow_statistics` provides runtime statistics about flows.
58#[derive(Debug)]
59pub(super) struct InformationSchemaFlowStatistics {
60    schema: SchemaRef,
61    catalog_name: String,
62    catalog_manager: Weak<dyn CatalogManager>,
63    flow_metadata_manager: Arc<FlowMetadataManager>,
64}
65
66impl InformationSchemaFlowStatistics {
67    pub(super) fn new(
68        catalog_name: String,
69        catalog_manager: Weak<dyn CatalogManager>,
70        flow_metadata_manager: Arc<FlowMetadataManager>,
71    ) -> Self {
72        Self {
73            schema: Self::schema(),
74            catalog_name,
75            catalog_manager,
76            flow_metadata_manager,
77        }
78    }
79
80    pub(crate) fn schema() -> SchemaRef {
81        Arc::new(Schema::new(
82            vec![
83                (FLOW_ID, CDT::uint32_datatype(), false),
84                (FLOW_NAME, CDT::string_datatype(), false),
85                (START_TIME, CDT::timestamp_millisecond_datatype(), true),
86                (
87                    LAST_EXECUTION_TIME,
88                    CDT::timestamp_millisecond_datatype(),
89                    true,
90                ),
91                (UPTIME_SECONDS, CDT::int64_datatype(), true),
92                (STATE_SIZE, CDT::uint64_datatype(), true),
93            ]
94            .into_iter()
95            .map(|(name, ty, nullable)| ColumnSchema::new(name, ty, nullable))
96            .collect(),
97        ))
98    }
99
100    fn builder(&self) -> InformationSchemaFlowStatisticsBuilder {
101        InformationSchemaFlowStatisticsBuilder::new(
102            self.schema.clone(),
103            self.catalog_name.clone(),
104            self.catalog_manager.clone(),
105            &self.flow_metadata_manager,
106        )
107    }
108}
109
110impl InformationTable for InformationSchemaFlowStatistics {
111    fn table_id(&self) -> TableId {
112        INFORMATION_SCHEMA_FLOW_STATISTICS_TABLE_ID
113    }
114
115    fn table_name(&self) -> &'static str {
116        FLOW_STATISTICS
117    }
118
119    fn schema(&self) -> SchemaRef {
120        self.schema.clone()
121    }
122
123    fn to_stream(&self, request: ScanRequest) -> Result<SendableRecordBatchStream> {
124        let schema = self.schema.arrow_schema().clone();
125        let mut builder = self.builder();
126        let stream = Box::pin(DfRecordBatchStreamAdapter::new(
127            schema,
128            futures::stream::once(async move {
129                builder
130                    .make_flow_statistics(Some(request))
131                    .await
132                    .map(|x| x.into_df_record_batch())
133                    .map_err(|err| datafusion::error::DataFusionError::External(Box::new(err)))
134            }),
135        ));
136        Ok(Box::pin(
137            RecordBatchStreamAdapter::try_new(stream)
138                .map_err(BoxedError::new)
139                .context(InternalSnafu)?,
140        ))
141    }
142}
143
144/// Builds the `information_schema.flow_statistics` table row by row.
145struct InformationSchemaFlowStatisticsBuilder {
146    schema: SchemaRef,
147    catalog_name: String,
148    catalog_manager: Weak<dyn CatalogManager>,
149    flow_metadata_manager: Arc<FlowMetadataManager>,
150
151    flow_ids: UInt32VectorBuilder,
152    flow_names: StringVectorBuilder,
153    start_times: TimestampMillisecondVectorBuilder,
154    last_execution_times: TimestampMillisecondVectorBuilder,
155    uptime_seconds: Int64VectorBuilder,
156    state_sizes: UInt64VectorBuilder,
157}
158
159impl InformationSchemaFlowStatisticsBuilder {
160    fn new(
161        schema: SchemaRef,
162        catalog_name: String,
163        catalog_manager: Weak<dyn CatalogManager>,
164        flow_metadata_manager: &Arc<FlowMetadataManager>,
165    ) -> Self {
166        Self {
167            schema,
168            catalog_name,
169            catalog_manager,
170            flow_metadata_manager: flow_metadata_manager.clone(),
171
172            flow_ids: UInt32VectorBuilder::with_capacity(INIT_CAPACITY),
173            flow_names: StringVectorBuilder::with_capacity(INIT_CAPACITY),
174            start_times: TimestampMillisecondVectorBuilder::with_capacity(INIT_CAPACITY),
175            last_execution_times: TimestampMillisecondVectorBuilder::with_capacity(INIT_CAPACITY),
176            uptime_seconds: Int64VectorBuilder::with_capacity(INIT_CAPACITY),
177            state_sizes: UInt64VectorBuilder::with_capacity(INIT_CAPACITY),
178        }
179    }
180
181    /// Construct the `information_schema.flow_statistics` virtual table.
182    async fn make_flow_statistics(&mut self, request: Option<ScanRequest>) -> Result<RecordBatch> {
183        let catalog_name = self.catalog_name.clone();
184        let predicates = Predicates::from_scan_request(&request);
185
186        let flow_info_manager = self.flow_metadata_manager.clone();
187
188        let mut stream = flow_info_manager
189            .flow_name_manager()
190            .flow_names(&catalog_name)
191            .await;
192
193        let flow_stat = {
194            let information_extension = utils::information_extension(&self.catalog_manager)?;
195            information_extension.flow_stats().await?
196        };
197
198        let now = current_time_millis();
199
200        while let Some((flow_name, flow_id)) = stream
201            .try_next()
202            .await
203            .map_err(BoxedError::new)
204            .context(ListFlowsSnafu {
205                catalog: &catalog_name,
206            })?
207        {
208            self.add_flow_statistic(&predicates, flow_id.flow_id(), &flow_name, &flow_stat, now);
209        }
210
211        self.finish()
212    }
213
214    fn add_flow_statistic(
215        &mut self,
216        predicates: &Predicates,
217        flow_id: FlowId,
218        flow_name: &str,
219        flow_stat: &Option<FlowStat>,
220        now: i64,
221    ) {
222        let row = [
223            (FLOW_ID, &Value::from(flow_id)),
224            (FLOW_NAME, &Value::from(flow_name.to_string())),
225        ];
226        if !predicates.eval(&row) {
227            return;
228        }
229
230        let start_time = flow_stat
231            .as_ref()
232            .and_then(|stat| stat.start_time_map.get(&flow_id).copied());
233
234        self.flow_ids.push(Some(flow_id));
235        self.flow_names.push(Some(flow_name));
236        self.start_times
237            .push(start_time.map(TimestampMillisecond::new));
238        self.last_execution_times
239            .push(flow_stat.as_ref().and_then(|stat| {
240                stat.last_exec_time_map
241                    .get(&flow_id)
242                    .map(|v| TimestampMillisecond::new(*v))
243            }));
244        self.uptime_seconds
245            .push(start_time.map(|start| ((now - start) / 1000).max(0)));
246        self.state_sizes.push(
247            flow_stat
248                .as_ref()
249                .and_then(|stat| stat.state_size.get(&flow_id).map(|v| *v as u64)),
250        );
251    }
252
253    fn finish(&mut self) -> Result<RecordBatch> {
254        let columns: Vec<VectorRef> = vec![
255            Arc::new(self.flow_ids.finish()),
256            Arc::new(self.flow_names.finish()),
257            Arc::new(self.start_times.finish()),
258            Arc::new(self.last_execution_times.finish()),
259            Arc::new(self.uptime_seconds.finish()),
260            Arc::new(self.state_sizes.finish()),
261        ];
262        RecordBatch::new(self.schema.clone(), columns).context(CreateRecordBatchSnafu)
263    }
264}
265
266impl DfPartitionStream for InformationSchemaFlowStatistics {
267    fn schema(&self) -> &arrow_schema::SchemaRef {
268        self.schema.arrow_schema()
269    }
270
271    fn execute(&self, _: Arc<TaskContext>) -> DfSendableRecordBatchStream {
272        let schema: Arc<arrow_schema::Schema> = self.schema.arrow_schema().clone();
273        let mut builder = self.builder();
274        Box::pin(DfRecordBatchStreamAdapter::new(
275            schema,
276            futures::stream::once(async move {
277                builder
278                    .make_flow_statistics(None)
279                    .await
280                    .map(|x| x.into_df_record_batch())
281                    .map_err(Into::into)
282            }),
283        ))
284    }
285}
catalog/system_schema/information_schema/flow_statistics.rs

catalog/system_schema/information_schema/
flow_statistics.rs
