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::future::Future;
29use std::pin::Pin;
30use std::sync::Arc;
31
32use api::greptime_proto::v1::AddColumnLocation as Location;
33use api::greptime_proto::v1::add_column_location::LocationType;
34use common_error::ext::BoxedError;
35use common_recordbatch::{
36    RecordBatches, SendableRecordBatchMapper, SendableRecordBatchStream, map_dictionary_to_values,
37    map_dictionary_to_values_schema,
38};
39use datafusion::physical_plan::ExecutionPlan;
40use serde::{Deserialize, Serialize};
41use sqlparser_derive::{Visit, VisitMut};
42
43/// new Output struct with output data(previously Output) and output meta
44#[derive(Debug)]
45pub struct Output {
46    pub data: OutputData,
47    pub meta: OutputMeta,
48}
49
50/// Original Output struct
51/// carrying result data to response/client/user interface
52pub enum OutputData {
53    AffectedRows(OutputRows),
54    RecordBatches(RecordBatches),
55    Stream(SendableRecordBatchStream),
56}
57
58impl OutputData {
59    /// Consume the data to pretty printed string.
60    pub async fn pretty_print(self) -> String {
61        match self {
62            OutputData::AffectedRows(x) => {
63                format!("Affected Rows: {x}")
64            }
65            OutputData::RecordBatches(x) => x.pretty_print().unwrap_or_else(|e| e.to_string()),
66            OutputData::Stream(x) => common_recordbatch::util::collect_batches(x)
67                .await
68                .and_then(|x| x.pretty_print())
69                .unwrap_or_else(|e| e.to_string()),
70        }
71    }
72}
73
74/// OutputMeta stores meta information produced/generated during the execution
75#[derive(Debug, Default)]
76pub struct OutputMeta {
77    /// May exist for query output. One can retrieve execution metrics from this plan.
78    pub plan: Option<Arc<dyn ExecutionPlan>>,
79    pub cost: OutputCost,
80    /// Process-local storage results attached by asynchronous ordinary-table batching.
81    /// Callers needing confirmed writes must await these. Dropping them does not cancel writes.
82    pub write_completions: Vec<WriteCompletion>,
83}
84
85/// Actual storage completion of an already accepted asynchronous write.
86pub struct WriteCompletion(Pin<Box<dyn Future<Output = Result<(), BoxedError>> + Send>>);
87
88impl WriteCompletion {
89    /// Wraps the storage result without changing how the write is scheduled.
90    pub fn new(completion: impl Future<Output = Result<(), BoxedError>> + Send + 'static) -> Self {
91        Self(Box::pin(completion))
92    }
93
94    /// Waits for storage success or failure, independently of queue admission.
95    pub async fn wait(self) -> Result<(), BoxedError> {
96        self.0.await
97    }
98}
99
100impl Debug for WriteCompletion {
101    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
102        f.debug_struct("WriteCompletion").finish_non_exhaustive()
103    }
104}
105
106impl Output {
107    pub fn new_with_affected_rows(affected_rows: OutputRows) -> Self {
108        Self {
109            data: OutputData::AffectedRows(affected_rows),
110            meta: Default::default(),
111        }
112    }
113
114    pub fn new_with_record_batches(recordbatches: RecordBatches) -> Self {
115        Self {
116            data: OutputData::RecordBatches(recordbatches),
117            meta: Default::default(),
118        }
119    }
120
121    pub fn new_with_stream(stream: SendableRecordBatchStream) -> Self {
122        Self {
123            data: OutputData::Stream(stream),
124            meta: Default::default(),
125        }
126    }
127
128    pub fn new(data: OutputData, meta: OutputMeta) -> Self {
129        Self { data, meta }
130    }
131
132    /// Expands dictionary arrays before exposing a query result to a client.
133    pub fn map_dictionary_to_values(self) -> common_recordbatch::error::Result<Self> {
134        let Self { data, meta } = self;
135        let data = match data {
136            OutputData::AffectedRows(rows) => OutputData::AffectedRows(rows),
137            OutputData::RecordBatches(record_batches) => {
138                let original_schema = record_batches.schema();
139                let (mapped_schema, apply_mapper) =
140                    map_dictionary_to_values_schema(original_schema.clone());
141                if !apply_mapper {
142                    OutputData::RecordBatches(record_batches)
143                } else {
144                    let batches = record_batches
145                        .into_iter()
146                        .map(|batch| {
147                            map_dictionary_to_values(batch, &original_schema, &mapped_schema)
148                        })
149                        .collect::<common_recordbatch::error::Result<Vec<_>>>()?;
150                    OutputData::RecordBatches(RecordBatches::try_new(mapped_schema, batches)?)
151                }
152            }
153            OutputData::Stream(stream) => {
154                let (_, apply_mapper) = map_dictionary_to_values_schema(stream.schema());
155                if apply_mapper {
156                    OutputData::Stream(Box::pin(SendableRecordBatchMapper::new(
157                        stream,
158                        map_dictionary_to_values,
159                        map_dictionary_to_values_schema,
160                    )))
161                } else {
162                    OutputData::Stream(stream)
163                }
164            }
165        };
166        Ok(Self { data, meta })
167    }
168
169    pub fn extract_rows_and_cost(&self) -> (OutputRows, OutputCost) {
170        match self.data {
171            OutputData::AffectedRows(rows) => (rows, self.meta.cost),
172            _ => (0, self.meta.cost),
173        }
174    }
175}
176
177impl Debug for OutputData {
178    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
179        match self {
180            OutputData::AffectedRows(rows) => write!(f, "OutputData::AffectedRows({rows})"),
181            OutputData::RecordBatches(recordbatches) => {
182                write!(f, "OutputData::RecordBatches({recordbatches:?})")
183            }
184            OutputData::Stream(s) => {
185                write!(f, "OutputData::Stream(<{}>)", s.name())
186            }
187        }
188    }
189}
190
191impl OutputMeta {
192    pub fn new(plan: Option<Arc<dyn ExecutionPlan>>, cost: usize) -> Self {
193        Self {
194            plan,
195            cost,
196            ..Default::default()
197        }
198    }
199
200    pub fn new_with_plan(plan: Arc<dyn ExecutionPlan>) -> Self {
201        Self {
202            plan: Some(plan),
203            cost: 0,
204            ..Default::default()
205        }
206    }
207
208    pub fn new_with_cost(cost: usize) -> Self {
209        Self {
210            cost,
211            ..Default::default()
212        }
213    }
214}
215
216#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Visit, VisitMut)]
217pub enum AddColumnLocation {
218    First,
219    After { column_name: String },
220}
221
222impl Display for AddColumnLocation {
223    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
224        match self {
225            AddColumnLocation::First => write!(f, r#"FIRST"#),
226            AddColumnLocation::After { column_name } => {
227                write!(f, r#"AFTER {column_name}"#)
228            }
229        }
230    }
231}
232
233impl From<&AddColumnLocation> for Location {
234    fn from(value: &AddColumnLocation) -> Self {
235        match value {
236            AddColumnLocation::First => Location {
237                location_type: LocationType::First.into(),
238                after_column_name: String::default(),
239            },
240            AddColumnLocation::After { column_name } => Location {
241                location_type: LocationType::After.into(),
242                after_column_name: column_name.clone(),
243            },
244        }
245    }
246}
247
248pub type OutputRows = usize;
249pub type OutputCost = usize;
common_query/lib.rs

common_query/
lib.rs
