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

common_query/
lib.rs
