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    schema: SchemaRef,
293    batches: Vec<RecordBatch>,
294    batch_idx: usize,
295    row_idx: usize,
296    columns: Arc<Vec<String>>,
297    index: Arc<HashMap<String, usize>>,
298}
299
300impl QueryIterator {
301    pub(crate) fn new(schema: SchemaRef, batches: Vec<RecordBatch>) -> Self {
302        let columns: Vec<String> = schema.fields().iter().map(|f| f.name().clone()).collect();
303        let index: HashMap<String, usize> = columns
304            .iter()
305            .enumerate()
306            .map(|(i, n)| (n.clone(), i))
307            .collect();
308        QueryIterator {
309            schema,
310            batches,
311            batch_idx: 0,
312            row_idx: 0,
313            columns: Arc::new(columns),
314            index: Arc::new(index),
315        }
316    }
317
318    /// The column names, in schema order.
319    pub fn column_names(&self) -> &[String] {
320        &self.columns
321    }
322
323    /// Total number of rows across all batches.
324    pub fn num_rows(&self) -> usize {
325        self.batches.iter().map(|b| b.num_rows()).sum()
326    }
327}
328
329impl Iterator for QueryIterator {
330    type Item = Result<Row, Error>;
331
332    fn next(&mut self) -> Option<Self::Item> {
333        while self.batch_idx < self.batches.len()
334            && self.row_idx >= self.batches[self.batch_idx].num_rows()
335        {
336            self.batch_idx += 1;
337            self.row_idx = 0;
338        }
339
340        if self.batch_idx >= self.batches.len() {
341            return None;
342        }
343
344        let batch = &self.batches[self.batch_idx];
345        let row = self.row_idx;
346        self.row_idx += 1;
347
348        let values = (0..batch.num_columns())
349            .map(|col_idx| extract_value(batch.column(col_idx).as_ref(), row))
350            .collect::<Result<Vec<_>, _>>();
351
352        Some(values.map(|values| Row {
353            values,
354            columns: Arc::clone(&self.columns),
355            index: Arc::clone(&self.index),
356        }))
357    }
358}
359
360/// Extract a single row value from an Arrow array column.
361fn extract_value(array: &dyn Array, row: usize) -> Result<Value, Error> {
362    use arrow_schema::DataType::*;
363
364    if array.is_null(row) {
365        return Ok(Value::Null);
366    }
367
368    match array.data_type() {
369        Boolean => Ok(Value::Bool(
370            array
371                .as_any()
372                .downcast_ref::<BooleanArray>()
373                .unwrap()
374                .value(row),
375        )),
376        Int8 => Ok(Value::I8(
377            array
378                .as_any()
379                .downcast_ref::<Int8Array>()
380                .unwrap()
381                .value(row),
382        )),
383        Int16 => Ok(Value::I16(
384            array
385                .as_any()
386                .downcast_ref::<Int16Array>()
387                .unwrap()
388                .value(row),
389        )),
390        Int32 => Ok(Value::I32(
391            array
392                .as_any()
393                .downcast_ref::<Int32Array>()
394                .unwrap()
395                .value(row),
396        )),
397        Int64 => Ok(Value::I64(
398            array
399                .as_any()
400                .downcast_ref::<Int64Array>()
401                .unwrap()
402                .value(row),
403        )),
404        UInt8 => Ok(Value::U8(
405            array
406                .as_any()
407                .downcast_ref::<UInt8Array>()
408                .unwrap()
409                .value(row),
410        )),
411        UInt16 => Ok(Value::U16(
412            array
413                .as_any()
414                .downcast_ref::<UInt16Array>()
415                .unwrap()
416                .value(row),
417        )),
418        UInt32 => Ok(Value::U32(
419            array
420                .as_any()
421                .downcast_ref::<UInt32Array>()
422                .unwrap()
423                .value(row),
424        )),
425        UInt64 => Ok(Value::U64(
426            array
427                .as_any()
428                .downcast_ref::<UInt64Array>()
429                .unwrap()
430                .value(row),
431        )),
432        Float32 => Ok(Value::F32(
433            array
434                .as_any()
435                .downcast_ref::<Float32Array>()
436                .unwrap()
437                .value(row),
438        )),
439        Float64 => Ok(Value::F64(
440            array
441                .as_any()
442                .downcast_ref::<Float64Array>()
443                .unwrap()
444                .value(row),
445        )),
446        Utf8 => Ok(Value::String(
447            array
448                .as_any()
449                .downcast_ref::<StringArray>()
450                .unwrap()
451                .value(row)
452                .to_owned(),
453        )),
454        LargeUtf8 => Ok(Value::String(
455            array
456                .as_any()
457                .downcast_ref::<LargeStringArray>()
458                .unwrap()
459                .value(row)
460                .to_owned(),
461        )),
462        Binary => Ok(Value::Binary(
463            array
464                .as_any()
465                .downcast_ref::<BinaryArray>()
466                .unwrap()
467                .value(row)
468                .to_owned(),
469        )),
470        LargeBinary => Ok(Value::Binary(
471            array
472                .as_any()
473                .downcast_ref::<LargeBinaryArray>()
474                .unwrap()
475                .value(row)
476                .to_owned(),
477        )),
478        Timestamp(arrow_schema::TimeUnit::Nanosecond, _) => Ok(Value::Timestamp(
479            array
480                .as_any()
481                .downcast_ref::<TimestampNanosecondArray>()
482                .unwrap()
483                .value(row),
484        )),
485        Timestamp(arrow_schema::TimeUnit::Microsecond, _) => Ok(Value::Timestamp(
486            array
487                .as_any()
488                .downcast_ref::<TimestampMicrosecondArray>()
489                .unwrap()
490                .value(row)
491                * 1_000,
492        )),
493        Timestamp(arrow_schema::TimeUnit::Millisecond, _) => Ok(Value::Timestamp(
494            array
495                .as_any()
496                .downcast_ref::<TimestampMillisecondArray>()
497                .unwrap()
498                .value(row)
499                * 1_000_000,
500        )),
501        Timestamp(arrow_schema::TimeUnit::Second, _) => Ok(Value::Timestamp(
502            array
503                .as_any()
504                .downcast_ref::<TimestampSecondArray>()
505                .unwrap()
506                .value(row)
507                * 1_000_000_000,
508        )),
509        // Dictionary-encoded columns: InfluxDB 3 returns tag columns as
510        // Dictionary(Int32, Utf8).  Resolve the key for this row and recurse
511        // into the values array, so the actual tag value is returned rather
512        // than a debug dump of the column.
513        Dictionary(key_type, _) => {
514            macro_rules! resolve {
515                ($t:ty) => {{
516                    let dict = array
517                        .as_any()
518                        .downcast_ref::<DictionaryArray<$t>>()
519                        .unwrap();
520                    let key = dict.keys().value(row) as usize;
521                    extract_value(dict.values().as_ref(), key)
522                }};
523            }
524            match key_type.as_ref() {
525                Int8 => resolve!(Int8Type),
526                Int16 => resolve!(Int16Type),
527                Int32 => resolve!(Int32Type),
528                Int64 => resolve!(Int64Type),
529                UInt8 => resolve!(UInt8Type),
530                UInt16 => resolve!(UInt16Type),
531                UInt32 => resolve!(UInt32Type),
532                UInt64 => resolve!(UInt64Type),
533                _ => Err(Error::UnsupportedArrowType {
534                    data_type: array.data_type().to_string(),
535                }),
536            }
537        }
538        // Decimals carry a scale that doesn't map onto an f64/i64 cleanly;
539        // render them as their exact decimal string.
540        Decimal128(_, _) => Ok(Value::String(
541            array
542                .as_any()
543                .downcast_ref::<Decimal128Array>()
544                .unwrap()
545                .value_as_string(row),
546        )),
547        Decimal256(_, _) => Ok(Value::String(
548            array
549                .as_any()
550                .downcast_ref::<Decimal256Array>()
551                .unwrap()
552                .value_as_string(row),
553        )),
554        _other => Err(Error::UnsupportedArrowType {
555            data_type: array.data_type().to_string(),
556        }),
557    }
558}
influxdb3_client/query.rs

influxdb3_client/
query.rs
