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 snafu::IntoError;
42use sqlparser_derive::{Visit, VisitMut};
43
44/// new Output struct with output data(previously Output) and output meta
45#[derive(Debug)]
46pub struct Output {
47    pub data: OutputData,
48    pub meta: OutputMeta,
49}
50
51/// Original Output struct
52/// carrying result data to response/client/user interface
53pub enum OutputData {
54    AffectedRows(OutputRows),
55    RecordBatches(RecordBatches),
56    Stream(SendableRecordBatchStream),
57}
58
59impl OutputData {
60    /// Consume the data to pretty printed string.
61    pub async fn pretty_print(self) -> String {
62        match self {
63            OutputData::AffectedRows(x) => {
64                format!("Affected Rows: {x}")
65            }
66            OutputData::RecordBatches(x) => x.pretty_print().unwrap_or_else(|e| e.to_string()),
67            OutputData::Stream(x) => common_recordbatch::util::collect_batches(x)
68                .await
69                .and_then(|x| x.pretty_print())
70                .unwrap_or_else(|e| e.to_string()),
71        }
72    }
73}
74
75/// OutputMeta stores meta information produced/generated during the execution
76#[derive(Debug, Default)]
77pub struct OutputMeta {
78    /// May exist for query output. One can retrieve execution metrics from this plan.
79    pub plan: Option<Arc<dyn ExecutionPlan>>,
80    pub cost: OutputCost,
81    /// Process-local storage results retained for asynchronous batching or confirmed writes.
82    /// Callers needing confirmed writes must await these. Dropping them does not cancel writes.
83    pub write_completions: Vec<WriteCompletion>,
84}
85
86/// Storage outcome of an already accepted write, whether synchronous or asynchronous.
87pub struct WriteCompletion(Pin<Box<dyn Future<Output = WriteProgress> + Send>>);
88
89/// Acknowledged rows for one original input, even when another part failed.
90/// Failed or lost acknowledgements do not establish how many rows were persisted.
91#[derive(Debug, Default, Clone)]
92pub struct WriteProgress {
93    pub affected_rows: usize,
94    pub error: Option<Arc<BoxedError>>,
95}
96
97impl WriteProgress {
98    pub fn failed(error: BoxedError) -> Self {
99        Self {
100            affected_rows: 0,
101            error: Some(Arc::new(error)),
102        }
103    }
104
105    pub fn merge(&mut self, other: Self) {
106        self.affected_rows += other.affected_rows;
107        if self.error.is_none() {
108            self.error = other.error;
109        }
110    }
111
112    pub fn into_result(self) -> Result<usize, BoxedError> {
113        match self.error {
114            Some(error) => Err(BoxedError::new(error::StorageWriteSnafu.into_error(error))),
115            None => Ok(self.affected_rows),
116        }
117    }
118}
119
120impl WriteCompletion {
121    /// Wraps the storage result without changing how the write is scheduled.
122    /// Success reports zero affected rows; use [`Self::with_progress`] for counted writes.
123    pub fn new(completion: impl Future<Output = Result<(), BoxedError>> + Send + 'static) -> Self {
124        Self::with_progress(async move {
125            match completion.await {
126                Ok(()) => WriteProgress::default(),
127                Err(error) => WriteProgress::failed(error),
128            }
129        })
130    }
131
132    pub fn with_progress(completion: impl Future<Output = WriteProgress> + Send + 'static) -> Self {
133        Self(Box::pin(completion))
134    }
135
136    pub async fn progress(self) -> WriteProgress {
137        self.0.await
138    }
139
140    /// Waits for storage success or failure, independently of queue admission.
141    pub async fn wait(self) -> Result<(), BoxedError> {
142        self.0.await.into_result().map(|_| ())
143    }
144}
145
146impl Debug for WriteCompletion {
147    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
148        f.debug_struct("WriteCompletion").finish_non_exhaustive()
149    }
150}
151
152impl Output {
153    /// Observes every dispatched completion before returning any failure.
154    pub async fn write_progress(self) -> WriteProgress {
155        if self.meta.write_completions.is_empty() {
156            return WriteProgress {
157                affected_rows: self.extract_rows_and_cost().0,
158                error: None,
159            };
160        }
161        let mut progress = WriteProgress::default();
162        for completion in self.meta.write_completions {
163            progress.merge(completion.progress().await);
164        }
165        progress
166    }
167    pub fn new_with_affected_rows(affected_rows: OutputRows) -> Self {
168        Self {
169            data: OutputData::AffectedRows(affected_rows),
170            meta: Default::default(),
171        }
172    }
173
174    pub fn new_with_record_batches(recordbatches: RecordBatches) -> Self {
175        Self {
176            data: OutputData::RecordBatches(recordbatches),
177            meta: Default::default(),
178        }
179    }
180
181    pub fn new_with_stream(stream: SendableRecordBatchStream) -> Self {
182        Self {
183            data: OutputData::Stream(stream),
184            meta: Default::default(),
185        }
186    }
187
188    pub fn new(data: OutputData, meta: OutputMeta) -> Self {
189        Self { data, meta }
190    }
191
192    /// Expands dictionary arrays before exposing a query result to a client.
193    pub fn map_dictionary_to_values(self) -> common_recordbatch::error::Result<Self> {
194        let Self { data, meta } = self;
195        let data = match data {
196            OutputData::AffectedRows(rows) => OutputData::AffectedRows(rows),
197            OutputData::RecordBatches(record_batches) => {
198                let original_schema = record_batches.schema();
199                let (mapped_schema, apply_mapper) =
200                    map_dictionary_to_values_schema(original_schema.clone());
201                if !apply_mapper {
202                    OutputData::RecordBatches(record_batches)
203                } else {
204                    let batches = record_batches
205                        .into_iter()
206                        .map(|batch| {
207                            map_dictionary_to_values(batch, &original_schema, &mapped_schema)
208                        })
209                        .collect::<common_recordbatch::error::Result<Vec<_>>>()?;
210                    OutputData::RecordBatches(RecordBatches::try_new(mapped_schema, batches)?)
211                }
212            }
213            OutputData::Stream(stream) => {
214                let (_, apply_mapper) = map_dictionary_to_values_schema(stream.schema());
215                if apply_mapper {
216                    OutputData::Stream(Box::pin(SendableRecordBatchMapper::new(
217                        stream,
218                        map_dictionary_to_values,
219                        map_dictionary_to_values_schema,
220                    )))
221                } else {
222                    OutputData::Stream(stream)
223                }
224            }
225        };
226        Ok(Self { data, meta })
227    }
228
229    pub fn extract_rows_and_cost(&self) -> (OutputRows, OutputCost) {
230        match self.data {
231            OutputData::AffectedRows(rows) => (rows, self.meta.cost),
232            _ => (0, self.meta.cost),
233        }
234    }
235}
236
237impl Debug for OutputData {
238    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
239        match self {
240            OutputData::AffectedRows(rows) => write!(f, "OutputData::AffectedRows({rows})"),
241            OutputData::RecordBatches(recordbatches) => {
242                write!(f, "OutputData::RecordBatches({recordbatches:?})")
243            }
244            OutputData::Stream(s) => {
245                write!(f, "OutputData::Stream(<{}>)", s.name())
246            }
247        }
248    }
249}
250
251impl OutputMeta {
252    pub fn new(plan: Option<Arc<dyn ExecutionPlan>>, cost: usize) -> Self {
253        Self {
254            plan,
255            cost,
256            ..Default::default()
257        }
258    }
259
260    pub fn new_with_plan(plan: Arc<dyn ExecutionPlan>) -> Self {
261        Self {
262            plan: Some(plan),
263            cost: 0,
264            ..Default::default()
265        }
266    }
267
268    pub fn new_with_cost(cost: usize) -> Self {
269        Self {
270            cost,
271            ..Default::default()
272        }
273    }
274}
275
276#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Visit, VisitMut)]
277pub enum AddColumnLocation {
278    First,
279    After { column_name: String },
280}
281
282impl Display for AddColumnLocation {
283    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
284        match self {
285            AddColumnLocation::First => write!(f, r#"FIRST"#),
286            AddColumnLocation::After { column_name } => {
287                write!(f, r#"AFTER {column_name}"#)
288            }
289        }
290    }
291}
292
293impl From<&AddColumnLocation> for Location {
294    fn from(value: &AddColumnLocation) -> Self {
295        match value {
296            AddColumnLocation::First => Location {
297                location_type: LocationType::First.into(),
298                after_column_name: String::default(),
299            },
300            AddColumnLocation::After { column_name } => Location {
301                location_type: LocationType::After.into(),
302                after_column_name: column_name.clone(),
303            },
304        }
305    }
306}
307
308pub type OutputRows = usize;
309pub type OutputCost = usize;
common_query/lib.rs

common_query/
lib.rs
