common_query/
lib.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 columnar_value;
16pub mod error;
17pub mod logical_plan;
18pub mod native_histogram;
19pub mod prelude;
20pub mod prometheus;
21pub mod request;
22pub mod stream;
23#[cfg(any(test, feature = "testing"))]
24pub mod test_util;
25
26use std::fmt::{Debug, Display, Formatter};
27use std::sync::Arc;
28
29use api::greptime_proto::v1::AddColumnLocation as Location;
30use api::greptime_proto::v1::add_column_location::LocationType;
31use common_recordbatch::{RecordBatches, SendableRecordBatchStream};
32use datafusion::physical_plan::ExecutionPlan;
33use serde::{Deserialize, Serialize};
34use sqlparser_derive::{Visit, VisitMut};
35
36/// new Output struct with output data(previously Output) and output meta
37#[derive(Debug)]
38pub struct Output {
39    pub data: OutputData,
40    pub meta: OutputMeta,
41}
42
43/// Original Output struct
44/// carrying result data to response/client/user interface
45pub enum OutputData {
46    AffectedRows(OutputRows),
47    RecordBatches(RecordBatches),
48    Stream(SendableRecordBatchStream),
49}
50
51impl OutputData {
52    /// Consume the data to pretty printed string.
53    pub async fn pretty_print(self) -> String {
54        match self {
55            OutputData::AffectedRows(x) => {
56                format!("Affected Rows: {x}")
57            }
58            OutputData::RecordBatches(x) => x.pretty_print().unwrap_or_else(|e| e.to_string()),
59            OutputData::Stream(x) => common_recordbatch::util::collect_batches(x)
60                .await
61                .and_then(|x| x.pretty_print())
62                .unwrap_or_else(|e| e.to_string()),
63        }
64    }
65}
66
67/// OutputMeta stores meta information produced/generated during the execution
68#[derive(Debug, Default)]
69pub struct OutputMeta {
70    /// May exist for query output. One can retrieve execution metrics from this plan.
71    pub plan: Option<Arc<dyn ExecutionPlan>>,
72    pub cost: OutputCost,
73}
74
75impl Output {
76    pub fn new_with_affected_rows(affected_rows: OutputRows) -> Self {
77        Self {
78            data: OutputData::AffectedRows(affected_rows),
79            meta: Default::default(),
80        }
81    }
82
83    pub fn new_with_record_batches(recordbatches: RecordBatches) -> Self {
84        Self {
85            data: OutputData::RecordBatches(recordbatches),
86            meta: Default::default(),
87        }
88    }
89
90    pub fn new_with_stream(stream: SendableRecordBatchStream) -> Self {
91        Self {
92            data: OutputData::Stream(stream),
93            meta: Default::default(),
94        }
95    }
96
97    pub fn new(data: OutputData, meta: OutputMeta) -> Self {
98        Self { data, meta }
99    }
100
101    pub fn extract_rows_and_cost(&self) -> (OutputRows, OutputCost) {
102        match self.data {
103            OutputData::AffectedRows(rows) => (rows, self.meta.cost),
104            _ => (0, self.meta.cost),
105        }
106    }
107}
108
109impl Debug for OutputData {
110    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
111        match self {
112            OutputData::AffectedRows(rows) => write!(f, "OutputData::AffectedRows({rows})"),
113            OutputData::RecordBatches(recordbatches) => {
114                write!(f, "OutputData::RecordBatches({recordbatches:?})")
115            }
116            OutputData::Stream(s) => {
117                write!(f, "OutputData::Stream(<{}>)", s.name())
118            }
119        }
120    }
121}
122
123impl OutputMeta {
124    pub fn new(plan: Option<Arc<dyn ExecutionPlan>>, cost: usize) -> Self {
125        Self { plan, cost }
126    }
127
128    pub fn new_with_plan(plan: Arc<dyn ExecutionPlan>) -> Self {
129        Self {
130            plan: Some(plan),
131            cost: 0,
132        }
133    }
134
135    pub fn new_with_cost(cost: usize) -> Self {
136        Self { plan: None, cost }
137    }
138}
139
140#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Visit, VisitMut)]
141pub enum AddColumnLocation {
142    First,
143    After { column_name: String },
144}
145
146impl Display for AddColumnLocation {
147    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
148        match self {
149            AddColumnLocation::First => write!(f, r#"FIRST"#),
150            AddColumnLocation::After { column_name } => {
151                write!(f, r#"AFTER {column_name}"#)
152            }
153        }
154    }
155}
156
157impl From<&AddColumnLocation> for Location {
158    fn from(value: &AddColumnLocation) -> Self {
159        match value {
160            AddColumnLocation::First => Location {
161                location_type: LocationType::First.into(),
162                after_column_name: String::default(),
163            },
164            AddColumnLocation::After { column_name } => Location {
165                location_type: LocationType::After.into(),
166                after_column_name: column_name.clone(),
167            },
168        }
169    }
170}
171
172pub type OutputRows = usize;
173pub type OutputCost = usize;
common_query/lib.rs

common_query/
lib.rs
