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

common_query/
lib.rs
