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

common_query/
lib.rs
