tecton_compute/
query.rs

1//! Query request/response types and Arrow → JSON conversion.
2
3use datafusion::arrow::array::Array;
4use datafusion::arrow::record_batch::RecordBatch;
5use datafusion::arrow::util::display::array_value_to_string;
6use serde::{Deserialize, Serialize};
7use serde_json::{json, Map, Value};
8use tecton_core::{Result, TectonError};
9
10/// Incoming SQL query from CLI or UI.
11#[derive(Debug, Clone, Serialize, Deserialize)]
12pub struct QueryRequest {
13    pub sql: String,
14}
15
16impl QueryRequest {
17    pub fn new(sql: impl Into<String>) -> Self {
18        Self { sql: sql.into() }
19    }
20}
21
22/// Rich result returned by [`crate::execute_query`] for large-dataset workloads.
23#[derive(Debug, Clone, Serialize, Deserialize)]
24pub struct QueryResult {
25    /// Wall-clock execution time in milliseconds.
26    pub execution_time_ms: u64,
27    /// Rows matched / estimated (may exceed `data.len()` when a safety LIMIT applied).
28    pub total_rows_affected: u64,
29    /// Limited result set as JSON objects (one per row).
30    pub data: Vec<Value>,
31    /// Human-readable note about automatic optimizations (Parquet conversion, LIMIT, …).
32    pub optimization_applied: Option<String>,
33}
34
35impl QueryResult {
36    /// Serialize as pretty JSON for CLI / Tauri string bridges.
37    pub fn to_pretty_json(&self) -> std::result::Result<String, serde_json::Error> {
38        serde_json::to_string_pretty(self)
39    }
40}
41
42/// Tabular result returned to callers (legacy engine path).
43#[derive(Debug, Clone, Serialize, Deserialize)]
44pub struct QueryResponse {
45    pub columns: Vec<String>,
46    pub rows: Vec<Vec<Option<String>>>,
47    pub row_count: usize,
48    pub truncated: bool,
49    pub elapsed_ms: u64,
50}
51
52/// Intermediate conversion helper.
53pub struct QueryResultSet {
54    pub columns: Vec<String>,
55    pub rows: Vec<Vec<Option<String>>>,
56    pub row_count: usize,
57    pub truncated: bool,
58}
59
60impl QueryResultSet {
61    pub fn from_batches(batches: &[RecordBatch], max_rows: usize) -> Result<Self> {
62        if batches.is_empty() {
63            return Ok(Self {
64                columns: vec![],
65                rows: vec![],
66                row_count: 0,
67                truncated: false,
68            });
69        }
70
71        let columns: Vec<String> = batches[0]
72            .schema()
73            .fields()
74            .iter()
75            .map(|f| f.name().clone())
76            .collect();
77
78        let mut rows = Vec::new();
79        let mut truncated = false;
80
81        'outer: for batch in batches {
82            let n = batch.num_rows();
83            for row_idx in 0..n {
84                if rows.len() >= max_rows {
85                    truncated = true;
86                    break 'outer;
87                }
88                let mut row = Vec::with_capacity(batch.num_columns());
89                for col_idx in 0..batch.num_columns() {
90                    let array = batch.column(col_idx);
91                    if array.is_null(row_idx) {
92                        row.push(None);
93                    } else {
94                        let value = array_value_to_string(array.as_ref(), row_idx)
95                            .map_err(|e| TectonError::compute(e.to_string()))?;
96                        row.push(Some(value));
97                    }
98                }
99                rows.push(row);
100            }
101        }
102
103        let row_count = rows.len();
104        Ok(Self {
105            columns,
106            rows,
107            row_count,
108            truncated,
109        })
110    }
111}
112
113/// Convert a record batch into JSON objects (column name → value).
114pub fn batch_to_json_values(batch: &RecordBatch) -> std::result::Result<Vec<Value>, TectonError> {
115    let columns: Vec<String> = batch
116        .schema()
117        .fields()
118        .iter()
119        .map(|f| f.name().clone())
120        .collect();
121
122    let mut rows = Vec::with_capacity(batch.num_rows());
123    for row_idx in 0..batch.num_rows() {
124        let mut map = Map::with_capacity(batch.num_columns());
125        for (col_idx, name) in columns.iter().enumerate() {
126            let array = batch.column(col_idx);
127            if array.is_null(row_idx) {
128                map.insert(name.clone(), Value::Null);
129                continue;
130            }
131            let text = array_value_to_string(array.as_ref(), row_idx)
132                .map_err(|e| TectonError::compute(e.to_string()))?;
133            map.insert(name.clone(), json_primitive_from_text(&text));
134        }
135        rows.push(Value::Object(map));
136    }
137    Ok(rows)
138}
139
140fn json_primitive_from_text(text: &str) -> Value {
141    if let Ok(n) = text.parse::<i64>() {
142        return json!(n);
143    }
144    if let Ok(n) = text.parse::<f64>() {
145        if n.is_finite() {
146            return json!(n);
147        }
148    }
149    if text.eq_ignore_ascii_case("true") {
150        return Value::Bool(true);
151    }
152    if text.eq_ignore_ascii_case("false") {
153        return Value::Bool(false);
154    }
155    Value::String(text.to_owned())
156}
tecton_compute/query.rs

tecton_compute/
query.rs
