influxdb3_client/
query.rs

1use std::collections::HashMap;
2use std::fmt;
3use std::ops::Index;
4use std::sync::Arc;
5
6use arrow_array::array::{
7    Array, BinaryArray, BooleanArray, Decimal128Array, Decimal256Array, DictionaryArray,
8    Float32Array, Float64Array, Int16Array, Int32Array, Int64Array, Int8Array, LargeBinaryArray,
9    LargeStringArray, StringArray, TimestampMicrosecondArray, TimestampMillisecondArray,
10    TimestampNanosecondArray, TimestampSecondArray, UInt16Array, UInt32Array, UInt64Array,
11    UInt8Array,
12};
13use arrow_array::types::{
14    Int16Type, Int32Type, Int64Type, Int8Type, UInt16Type, UInt32Type, UInt64Type, UInt8Type,
15};
16use arrow_array::RecordBatch;
17use arrow_schema::SchemaRef;
18
19use crate::error::Error;
20
21/// Selects the query language used for a query operation.
22#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Hash)]
23pub enum QueryType {
24    /// Standard SQL (default)
25    #[default]
26    Sql,
27    /// InfluxQL, the InfluxDB 1.x query language
28    InfluxQL,
29}
30
31impl QueryType {
32    pub fn as_str(self) -> &'static str {
33        match self {
34            QueryType::Sql => "sql",
35            QueryType::InfluxQL => "influxql",
36        }
37    }
38}
39
40/// Named query parameters for parameterised SQL / InfluxQL statements.
41///
42/// Prefer chaining `.param("k", v)` on [`crate::QueryRequest`]; use this type
43/// directly when you need to assemble parameters dynamically.
44pub type QueryParameters = HashMap<String, serde_json::Value>;
45
46/// Options controlling a single query operation.
47#[derive(Debug, Clone, Default)]
48pub struct QueryOptions {
49    pub(crate) query_type: QueryType,
50    /// Extra gRPC metadata headers sent with the Flight DoGet request.
51    pub headers: HashMap<String, String>,
52}
53
54/// A dynamically typed value extracted from a query result row.
55#[derive(Debug, Clone, PartialEq)]
56pub enum Value {
57    Bool(bool),
58    I8(i8),
59    I16(i16),
60    I32(i32),
61    I64(i64),
62    U8(u8),
63    U16(u16),
64    U32(u32),
65    U64(u64),
66    F32(f32),
67    F64(f64),
68    String(String),
69    Binary(Vec<u8>),
70    /// Nanosecond-epoch timestamp
71    Timestamp(i64),
72    Null,
73}
74
75impl Value {
76    pub fn as_f64(&self) -> Option<f64> {
77        match self {
78            Value::F64(v) => Some(*v),
79            Value::F32(v) => Some(*v as f64),
80            Value::I64(v) => Some(*v as f64),
81            Value::I32(v) => Some(*v as f64),
82            Value::U64(v) => Some(*v as f64),
83            Value::U32(v) => Some(*v as f64),
84            _ => None,
85        }
86    }
87
88    pub fn as_i64(&self) -> Option<i64> {
89        match self {
90            Value::I64(v) => Some(*v),
91            Value::I32(v) => Some(*v as i64),
92            Value::I16(v) => Some(*v as i64),
93            Value::I8(v) => Some(*v as i64),
94            Value::Timestamp(v) => Some(*v),
95            _ => None,
96        }
97    }
98
99    pub fn as_str(&self) -> Option<&str> {
100        match self {
101            Value::String(s) => Some(s.as_str()),
102            _ => None,
103        }
104    }
105
106    pub fn as_bool(&self) -> Option<bool> {
107        match self {
108            Value::Bool(b) => Some(*b),
109            _ => None,
110        }
111    }
112
113    pub fn is_null(&self) -> bool {
114        matches!(self, Value::Null)
115    }
116}
117
118impl fmt::Display for Value {
119    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
120        match self {
121            Value::Bool(v) => write!(f, "{v}"),
122            Value::I8(v) => write!(f, "{v}"),
123            Value::I16(v) => write!(f, "{v}"),
124            Value::I32(v) => write!(f, "{v}"),
125            Value::I64(v) => write!(f, "{v}"),
126            Value::U8(v) => write!(f, "{v}"),
127            Value::U16(v) => write!(f, "{v}"),
128            Value::U32(v) => write!(f, "{v}"),
129            Value::U64(v) => write!(f, "{v}"),
130            Value::F32(v) => write!(f, "{v}"),
131            Value::F64(v) => write!(f, "{v}"),
132            Value::String(v) => f.write_str(v),
133            Value::Binary(v) => write!(f, "{}b", v.len()),
134            Value::Timestamp(v) => write!(f, "{v}"),
135            Value::Null => f.write_str("null"),
136        }
137    }
138}
139
140/// A single row from a query result.
141///
142/// Holds the raw `Vec<Value>` (one slot per column) and a shared index mapping
143/// column names to slot positions.  Lookup by name is O(1) via the shared
144/// `Arc<HashMap>`, so iteration allocates no per-row map.
145#[derive(Debug, Clone)]
146pub struct Row {
147    values: Vec<Value>,
148    columns: Arc<Vec<String>>,
149    index: Arc<HashMap<String, usize>>,
150}
151
152impl Row {
153    /// Look up a value by column name.
154    pub fn get(&self, name: &str) -> Option<&Value> {
155        self.index.get(name).and_then(|&i| self.values.get(i))
156    }
157
158    /// Look up a value by column position.
159    pub fn at(&self, idx: usize) -> Option<&Value> {
160        self.values.get(idx)
161    }
162
163    /// All column names, in schema order.
164    pub fn columns(&self) -> &[String] {
165        &self.columns
166    }
167
168    /// All values, in schema order.
169    pub fn values(&self) -> &[Value] {
170        &self.values
171    }
172
173    /// Number of columns in this row.
174    pub fn len(&self) -> usize {
175        self.values.len()
176    }
177
178    pub fn is_empty(&self) -> bool {
179        self.values.is_empty()
180    }
181
182    /// Convert to a `HashMap<String, Value>` for callers that prefer map-shaped
183    /// rows.  Allocates one HashMap and clones every column name.
184    pub fn into_map(self) -> HashMap<String, Value> {
185        self.columns.iter().cloned().zip(self.values).collect()
186    }
187}
188
189impl Index<&str> for Row {
190    type Output = Value;
191    fn index(&self, name: &str) -> &Value {
192        self.get(name)
193            .unwrap_or_else(|| panic!("no column named '{name}'"))
194    }
195}
196
197impl Index<usize> for Row {
198    type Output = Value;
199    fn index(&self, idx: usize) -> &Value {
200        &self.values[idx]
201    }
202}
203
204/// The complete result of a query: a collection of Arrow [`RecordBatch`]es.
205///
206/// Use `for row in result` (yields [`Row`]) for row-oriented access, or
207/// [`QueryResult::record_batches()`] for direct Arrow access.
208pub struct QueryResult {
209    pub(crate) schema: SchemaRef,
210    pub(crate) batches: Vec<RecordBatch>,
211}
212
213impl QueryResult {
214    pub fn new(schema: SchemaRef, batches: Vec<RecordBatch>) -> Self {
215        QueryResult { schema, batches }
216    }
217
218    pub fn schema(&self) -> &SchemaRef {
219        &self.schema
220    }
221
222    /// The underlying Arrow record batches (zero-copy).
223    pub fn record_batches(&self) -> &[RecordBatch] {
224        &self.batches
225    }
226
227    /// Total number of rows across all batches.
228    pub fn num_rows(&self) -> usize {
229        self.batches.iter().map(|b| b.num_rows()).sum()
230    }
231
232    /// Column names in schema order.
233    pub fn column_names(&self) -> Vec<&str> {
234        self.schema
235            .fields()
236            .iter()
237            .map(|f| f.name().as_str())
238            .collect()
239    }
240
241    /// Collect all rows into a `Vec<Row>`.
242    pub fn rows(self) -> Result<Vec<Row>, Error> {
243        self.into_iter().collect()
244    }
245
246    /// Convert the query result to a polars [`DataFrame`].
247    ///
248    /// Requires the `polars` Cargo feature.
249    ///
250    /// Note: this serialises the batches to Arrow IPC and reads them back
251    /// through polars, so it transiently holds roughly twice the result in
252    /// memory. For very large results, prefer streaming the
253    /// [`RecordBatch`]es via [`crate::Client::sql`]`(..).stream()` and
254    /// converting incrementally.
255    #[cfg(feature = "polars")]
256    pub fn to_polars(self) -> crate::Result<polars::prelude::DataFrame> {
257        use arrow::ipc::writer::FileWriter;
258        use polars::io::SerReader;
259        use polars::prelude::IpcReader;
260        use std::io::Cursor;
261
262        let mut buf: Vec<u8> = Vec::new();
263        {
264            let mut writer = FileWriter::try_new(&mut buf, &self.schema)?;
265            for batch in &self.batches {
266                writer.write(batch)?;
267            }
268            writer.finish()?;
269        }
270
271        let cursor = Cursor::new(buf);
272        IpcReader::new(cursor)
273            .finish()
274            .map_err(|e| crate::error::Error::Config(format!("polars conversion error: {e}")))
275    }
276}
277
278impl IntoIterator for QueryResult {
279    type Item = Result<Row, Error>;
280    type IntoIter = QueryIterator;
281
282    fn into_iter(self) -> Self::IntoIter {
283        QueryIterator::new(self.schema, self.batches)
284    }
285}
286
287/// Row-by-row iterator over a [`QueryResult`].
288///
289/// Holds the column-name index in an `Arc` so each yielded [`Row`] can share
290/// the same name-to-position map, so there is no per-row HashMap allocation.
291pub struct QueryIterator {
292    batches: Vec<RecordBatch>,
293    batch_idx: usize,
294    row_idx: usize,
295    columns: Arc<Vec<String>>,
296    index: Arc<HashMap<String, usize>>,
297}
298
299impl QueryIterator {
300    pub(crate) fn new(schema: SchemaRef, batches: Vec<RecordBatch>) -> Self {
301        let columns: Vec<String> = schema.fields().iter().map(|f| f.name().clone()).collect();
302        let index: HashMap<String, usize> = columns
303            .iter()
304            .enumerate()
305            .map(|(i, n)| (n.clone(), i))
306            .collect();
307        QueryIterator {
308            batches,
309            batch_idx: 0,
310            row_idx: 0,
311            columns: Arc::new(columns),
312            index: Arc::new(index),
313        }
314    }
315
316    /// The column names, in schema order.
317    pub fn column_names(&self) -> &[String] {
318        &self.columns
319    }
320
321    /// Total number of rows across all batches.
322    pub fn num_rows(&self) -> usize {
323        self.batches.iter().map(|b| b.num_rows()).sum()
324    }
325}
326
327impl Iterator for QueryIterator {
328    type Item = Result<Row, Error>;
329
330    fn next(&mut self) -> Option<Self::Item> {
331        while self.batch_idx < self.batches.len()
332            && self.row_idx >= self.batches[self.batch_idx].num_rows()
333        {
334            self.batch_idx += 1;
335            self.row_idx = 0;
336        }
337
338        if self.batch_idx >= self.batches.len() {
339            return None;
340        }
341
342        let batch = &self.batches[self.batch_idx];
343        let row = self.row_idx;
344        self.row_idx += 1;
345
346        let values = (0..batch.num_columns())
347            .map(|col_idx| extract_value(batch.column(col_idx).as_ref(), row))
348            .collect::<Result<Vec<_>, _>>();
349
350        Some(values.map(|values| Row {
351            values,
352            columns: Arc::clone(&self.columns),
353            index: Arc::clone(&self.index),
354        }))
355    }
356}
357
358/// Extract a single row value from an Arrow array column.
359fn extract_value(array: &dyn Array, row: usize) -> Result<Value, Error> {
360    use arrow_schema::DataType::*;
361
362    if array.is_null(row) {
363        return Ok(Value::Null);
364    }
365
366    match array.data_type() {
367        Boolean => Ok(Value::Bool(
368            array
369                .as_any()
370                .downcast_ref::<BooleanArray>()
371                .unwrap()
372                .value(row),
373        )),
374        Int8 => Ok(Value::I8(
375            array
376                .as_any()
377                .downcast_ref::<Int8Array>()
378                .unwrap()
379                .value(row),
380        )),
381        Int16 => Ok(Value::I16(
382            array
383                .as_any()
384                .downcast_ref::<Int16Array>()
385                .unwrap()
386                .value(row),
387        )),
388        Int32 => Ok(Value::I32(
389            array
390                .as_any()
391                .downcast_ref::<Int32Array>()
392                .unwrap()
393                .value(row),
394        )),
395        Int64 => Ok(Value::I64(
396            array
397                .as_any()
398                .downcast_ref::<Int64Array>()
399                .unwrap()
400                .value(row),
401        )),
402        UInt8 => Ok(Value::U8(
403            array
404                .as_any()
405                .downcast_ref::<UInt8Array>()
406                .unwrap()
407                .value(row),
408        )),
409        UInt16 => Ok(Value::U16(
410            array
411                .as_any()
412                .downcast_ref::<UInt16Array>()
413                .unwrap()
414                .value(row),
415        )),
416        UInt32 => Ok(Value::U32(
417            array
418                .as_any()
419                .downcast_ref::<UInt32Array>()
420                .unwrap()
421                .value(row),
422        )),
423        UInt64 => Ok(Value::U64(
424            array
425                .as_any()
426                .downcast_ref::<UInt64Array>()
427                .unwrap()
428                .value(row),
429        )),
430        Float32 => Ok(Value::F32(
431            array
432                .as_any()
433                .downcast_ref::<Float32Array>()
434                .unwrap()
435                .value(row),
436        )),
437        Float64 => Ok(Value::F64(
438            array
439                .as_any()
440                .downcast_ref::<Float64Array>()
441                .unwrap()
442                .value(row),
443        )),
444        Utf8 => Ok(Value::String(
445            array
446                .as_any()
447                .downcast_ref::<StringArray>()
448                .unwrap()
449                .value(row)
450                .to_owned(),
451        )),
452        LargeUtf8 => Ok(Value::String(
453            array
454                .as_any()
455                .downcast_ref::<LargeStringArray>()
456                .unwrap()
457                .value(row)
458                .to_owned(),
459        )),
460        Binary => Ok(Value::Binary(
461            array
462                .as_any()
463                .downcast_ref::<BinaryArray>()
464                .unwrap()
465                .value(row)
466                .to_owned(),
467        )),
468        LargeBinary => Ok(Value::Binary(
469            array
470                .as_any()
471                .downcast_ref::<LargeBinaryArray>()
472                .unwrap()
473                .value(row)
474                .to_owned(),
475        )),
476        Timestamp(arrow_schema::TimeUnit::Nanosecond, _) => Ok(Value::Timestamp(
477            array
478                .as_any()
479                .downcast_ref::<TimestampNanosecondArray>()
480                .unwrap()
481                .value(row),
482        )),
483        Timestamp(arrow_schema::TimeUnit::Microsecond, _) => Ok(Value::Timestamp(
484            array
485                .as_any()
486                .downcast_ref::<TimestampMicrosecondArray>()
487                .unwrap()
488                .value(row)
489                * 1_000,
490        )),
491        Timestamp(arrow_schema::TimeUnit::Millisecond, _) => Ok(Value::Timestamp(
492            array
493                .as_any()
494                .downcast_ref::<TimestampMillisecondArray>()
495                .unwrap()
496                .value(row)
497                * 1_000_000,
498        )),
499        Timestamp(arrow_schema::TimeUnit::Second, _) => Ok(Value::Timestamp(
500            array
501                .as_any()
502                .downcast_ref::<TimestampSecondArray>()
503                .unwrap()
504                .value(row)
505                * 1_000_000_000,
506        )),
507        // Dictionary-encoded columns: InfluxDB 3 returns tag columns as
508        // Dictionary(Int32, Utf8).  Resolve the key for this row and recurse
509        // into the values array, so the actual tag value is returned rather
510        // than a debug dump of the column.
511        Dictionary(key_type, _) => {
512            macro_rules! resolve {
513                ($t:ty) => {{
514                    let dict = array
515                        .as_any()
516                        .downcast_ref::<DictionaryArray<$t>>()
517                        .unwrap();
518                    let key = dict.keys().value(row) as usize;
519                    extract_value(dict.values().as_ref(), key)
520                }};
521            }
522            match key_type.as_ref() {
523                Int8 => resolve!(Int8Type),
524                Int16 => resolve!(Int16Type),
525                Int32 => resolve!(Int32Type),
526                Int64 => resolve!(Int64Type),
527                UInt8 => resolve!(UInt8Type),
528                UInt16 => resolve!(UInt16Type),
529                UInt32 => resolve!(UInt32Type),
530                UInt64 => resolve!(UInt64Type),
531                _ => Err(Error::UnsupportedArrowType {
532                    data_type: array.data_type().to_string(),
533                }),
534            }
535        }
536        // Decimals carry a scale that doesn't map onto an f64/i64 cleanly;
537        // render them as their exact decimal string.
538        Decimal128(_, _) => Ok(Value::String(
539            array
540                .as_any()
541                .downcast_ref::<Decimal128Array>()
542                .unwrap()
543                .value_as_string(row),
544        )),
545        Decimal256(_, _) => Ok(Value::String(
546            array
547                .as_any()
548                .downcast_ref::<Decimal256Array>()
549                .unwrap()
550                .value_as_string(row),
551        )),
552        _other => Err(Error::UnsupportedArrowType {
553            data_type: array.data_type().to_string(),
554        }),
555    }
556}
influxdb3_client/query.rs

influxdb3_client/
query.rs
