databend_driver_core/
rows.rs

1// Copyright 2021 Datafuse Labs
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
15use std::pin::Pin;
16use std::task::Context;
17use std::task::Poll;
18
19use serde::Deserialize;
20use tokio_stream::{Stream, StreamExt};
21
22use crate::error::{Error, Result};
23use crate::schema::SchemaRef;
24use crate::value::Value;
25use arrow::record_batch::RecordBatch;
26use chrono_tz::Tz;
27
28#[derive(Clone, Debug)]
29pub enum RowWithStats {
30    Row(Row),
31    Stats(ServerStats),
32}
33
34#[derive(Deserialize, Clone, Debug, Default)]
35pub struct ServerStats {
36    #[serde(default)]
37    pub total_rows: usize,
38    #[serde(default)]
39    pub total_bytes: usize,
40
41    #[serde(default)]
42    pub read_rows: usize,
43    #[serde(default)]
44    pub read_bytes: usize,
45
46    #[serde(default)]
47    pub write_rows: usize,
48    #[serde(default)]
49    pub write_bytes: usize,
50
51    #[serde(default)]
52    pub running_time_ms: f64,
53
54    #[serde(default)]
55    pub spill_file_nums: usize,
56
57    #[serde(default)]
58    pub spill_bytes: usize,
59}
60
61impl ServerStats {
62    pub fn normalize(&mut self) {
63        if self.total_rows == 0 {
64            self.total_rows = self.read_rows;
65        }
66        if self.total_bytes == 0 {
67            self.total_bytes = self.read_bytes;
68        }
69    }
70
71    pub fn merge(&mut self, other: &ServerStats) {
72        self.total_rows += other.total_rows;
73        self.total_bytes += other.total_bytes;
74        self.read_rows += other.read_rows;
75        self.read_bytes += other.read_bytes;
76        self.write_rows += other.write_rows;
77        self.write_bytes += other.write_bytes;
78        self.running_time_ms += other.running_time_ms;
79        self.spill_file_nums += other.spill_file_nums;
80        self.spill_bytes += other.spill_bytes;
81    }
82}
83
84impl From<databend_client::QueryStats> for ServerStats {
85    fn from(stats: databend_client::QueryStats) -> Self {
86        let mut p = Self {
87            total_rows: 0,
88            total_bytes: 0,
89            read_rows: stats.progresses.scan_progress.rows,
90            read_bytes: stats.progresses.scan_progress.bytes,
91            write_rows: stats.progresses.write_progress.rows,
92            write_bytes: stats.progresses.write_progress.bytes,
93            spill_file_nums: stats.progresses.spill_progress.file_nums,
94            spill_bytes: stats.progresses.spill_progress.bytes,
95            running_time_ms: stats.running_time_ms,
96        };
97        if let Some(total) = stats.progresses.total_scan {
98            p.total_rows = total.rows;
99            p.total_bytes = total.bytes;
100        }
101        p
102    }
103}
104
105#[derive(Clone, Debug, Default)]
106pub struct Row {
107    schema: SchemaRef,
108    values: Vec<Value>,
109}
110
111impl Row {
112    pub fn new(schema: SchemaRef, values: Vec<Value>) -> Self {
113        Self { schema, values }
114    }
115
116    pub fn len(&self) -> usize {
117        self.values.len()
118    }
119
120    pub fn is_empty(&self) -> bool {
121        self.values.is_empty()
122    }
123
124    pub fn values(&self) -> &[Value] {
125        &self.values
126    }
127
128    pub fn schema(&self) -> SchemaRef {
129        self.schema.clone()
130    }
131
132    pub fn from_vec(schema: SchemaRef, values: Vec<Value>) -> Self {
133        Self { schema, values }
134    }
135}
136
137impl TryFrom<(SchemaRef, Vec<Option<String>>, Tz)> for Row {
138    type Error = Error;
139
140    fn try_from((schema, data, tz): (SchemaRef, Vec<Option<String>>, Tz)) -> Result<Self> {
141        let mut values: Vec<Value> = Vec::with_capacity(data.len());
142        for (field, val) in schema.fields().iter().zip(data.into_iter()) {
143            values.push(Value::try_from((&field.data_type, val, tz))?);
144        }
145        Ok(Self::new(schema, values))
146    }
147}
148
149impl IntoIterator for Row {
150    type Item = Value;
151    type IntoIter = std::vec::IntoIter<Self::Item>;
152
153    fn into_iter(self) -> Self::IntoIter {
154        self.values.into_iter()
155    }
156}
157
158#[derive(Clone, Debug)]
159pub struct Rows {
160    rows: Vec<Row>,
161}
162
163impl Rows {
164    pub fn new(rows: Vec<Row>) -> Self {
165        Self { rows }
166    }
167
168    // pub fn schema(&self) -> SchemaRef {
169    //     self.schema.clone()
170    // }
171
172    pub fn rows(&self) -> &[Row] {
173        &self.rows
174    }
175
176    pub fn len(&self) -> usize {
177        self.rows.len()
178    }
179
180    pub fn is_empty(&self) -> bool {
181        self.rows.is_empty()
182    }
183}
184
185impl TryFrom<(RecordBatch, Tz)> for Rows {
186    type Error = Error;
187    fn try_from((batch, ltz): (RecordBatch, Tz)) -> Result<Self> {
188        let batch_schema = batch.schema();
189        let schema = SchemaRef::new(batch_schema.clone().try_into()?);
190        let mut rows: Vec<Row> = Vec::new();
191        for i in 0..batch.num_rows() {
192            let mut values: Vec<Value> = Vec::new();
193            for j in 0..batch_schema.fields().len() {
194                let v = batch.column(j);
195                let field = batch_schema.field(j);
196                let value = Value::try_from((field, v, i, ltz))?;
197                values.push(value);
198            }
199            rows.push(Row::new(schema.clone(), values));
200        }
201        Ok(Self::new(rows))
202    }
203}
204
205impl IntoIterator for Rows {
206    type Item = Row;
207    type IntoIter = std::vec::IntoIter<Self::Item>;
208
209    fn into_iter(self) -> Self::IntoIter {
210        self.rows.into_iter()
211    }
212}
213
214macro_rules! replace_expr {
215    ($_t:tt $sub:expr) => {
216        $sub
217    };
218}
219
220// This macro implements TryFrom for tuple of types
221macro_rules! impl_tuple_from_row {
222    ( $($Ti:tt),+ ) => {
223        impl<$($Ti),+> TryFrom<Row> for ($($Ti,)+)
224        where
225            $($Ti: TryFrom<Value>),+
226        {
227            type Error = String;
228            fn try_from(row: Row) -> Result<Self, String> {
229                // It is not possible yet to get the number of metavariable repetitions
230                // ref: https://github.com/rust-lang/lang-team/issues/28#issue-644523674
231                // This is a workaround
232                let expected_len = <[()]>::len(&[$(replace_expr!(($Ti) ())),*]);
233
234                if expected_len != row.len() {
235                    return Err(format!("row size mismatch: expected {} columns, got {}", expected_len, row.len()));
236                }
237                let mut vals_iter = row.into_iter().enumerate();
238
239                Ok((
240                    $(
241                        {
242                            let (col_ix, col_value) = vals_iter
243                                .next()
244                                .unwrap(); // vals_iter size is checked before this code is reached,
245                                           // so it is safe to unwrap
246                            let t = col_value.get_type();
247                            $Ti::try_from(col_value)
248                                .map_err(|_| format!("failed converting column {} from type({:?}) to type({})", col_ix, t, std::any::type_name::<$Ti>()))?
249                        }
250                    ,)+
251                ))
252            }
253        }
254    }
255}
256
257// Implement FromRow for tuples of size up to 16
258impl_tuple_from_row!(T1);
259impl_tuple_from_row!(T1, T2);
260impl_tuple_from_row!(T1, T2, T3);
261impl_tuple_from_row!(T1, T2, T3, T4);
262impl_tuple_from_row!(T1, T2, T3, T4, T5);
263impl_tuple_from_row!(T1, T2, T3, T4, T5, T6);
264impl_tuple_from_row!(T1, T2, T3, T4, T5, T6, T7);
265impl_tuple_from_row!(T1, T2, T3, T4, T5, T6, T7, T8);
266impl_tuple_from_row!(T1, T2, T3, T4, T5, T6, T7, T8, T9);
267impl_tuple_from_row!(T1, T2, T3, T4, T5, T6, T7, T8, T9, T10);
268impl_tuple_from_row!(T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11);
269impl_tuple_from_row!(T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12);
270impl_tuple_from_row!(T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13);
271impl_tuple_from_row!(T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13, T14);
272impl_tuple_from_row!(T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13, T14, T15);
273impl_tuple_from_row!(T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13, T14, T15, T16);
274
275pub struct RowIterator {
276    schema: SchemaRef,
277    it: Option<Pin<Box<dyn Stream<Item = Result<Row>> + Send>>>,
278}
279
280impl RowIterator {
281    pub fn new(schema: SchemaRef, it: Pin<Box<dyn Stream<Item = Result<Row>> + Send>>) -> Self {
282        Self {
283            schema,
284            it: Some(it),
285        }
286    }
287
288    pub fn schema(&self) -> SchemaRef {
289        self.schema.clone()
290    }
291
292    pub async fn try_collect<T>(mut self) -> Result<Vec<T>>
293    where
294        T: TryFrom<Row>,
295        T::Error: std::fmt::Display,
296    {
297        if let Some(it) = &mut self.it {
298            let mut ret = Vec::new();
299            while let Some(row) = it.next().await {
300                let v = T::try_from(row?).map_err(|e| Error::Parsing(e.to_string()))?;
301                ret.push(v)
302            }
303            Ok(ret)
304        } else {
305            Err(Error::BadArgument("RowIterator already closed".to_owned()))
306        }
307    }
308
309    pub fn close(&mut self) {
310        self.it = None;
311    }
312}
313
314impl Stream for RowIterator {
315    type Item = Result<Row>;
316
317    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
318        if let Some(it) = self.it.as_mut() {
319            Pin::new(it).poll_next(cx)
320        } else {
321            Poll::Ready(Some(Err(Error::BadArgument(
322                "RowIterator already closed".to_owned(),
323            ))))
324        }
325    }
326}
327
328pub struct RowStatsIterator {
329    schema: SchemaRef,
330    it: Option<Pin<Box<dyn Stream<Item = Result<RowWithStats>> + Send>>>,
331}
332
333impl RowStatsIterator {
334    pub fn new(
335        schema: SchemaRef,
336        it: Pin<Box<dyn Stream<Item = Result<RowWithStats>> + Send>>,
337    ) -> Self {
338        Self {
339            schema,
340            it: Some(it),
341        }
342    }
343
344    pub fn schema(&self) -> SchemaRef {
345        self.schema.clone()
346    }
347
348    pub async fn filter_rows(self) -> Result<RowIterator> {
349        if let Some(it) = self.it {
350            let it = it.filter_map(|r| match r {
351                Ok(RowWithStats::Row(r)) => Some(Ok(r)),
352                Ok(_) => None,
353                Err(err) => Some(Err(err)),
354            });
355            Ok(RowIterator::new(self.schema, Box::pin(it)))
356        } else {
357            Err(Error::BadArgument("RowIterator already closed".to_owned()))
358        }
359    }
360
361    pub fn close(&mut self) {
362        self.it = None;
363    }
364}
365
366impl Stream for RowStatsIterator {
367    type Item = Result<RowWithStats>;
368
369    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
370        if let Some(it) = self.it.as_mut() {
371            Pin::new(it).poll_next(cx)
372        } else {
373            Poll::Ready(Some(Err(Error::BadArgument(
374                "RowStatsIterator already closed".to_owned(),
375            ))))
376        }
377    }
378}
databend_driver_core/rows.rs

databend_driver_core/
rows.rs
