vortex_tui/
datafusion_helper.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Shared DataFusion query execution utilities for both CLI and TUI.
5
6use std::sync::Arc;
7
8use arrow_array::Array as ArrowArray;
9use arrow_array::RecordBatch;
10use datafusion::datasource::listing::ListingOptions;
11use datafusion::datasource::listing::ListingTable;
12use datafusion::datasource::listing::ListingTableConfig;
13use datafusion::datasource::listing::ListingTableUrl;
14use datafusion::prelude::SessionContext;
15use vortex::session::VortexSession;
16use vortex_datafusion::VortexFormat;
17
18/// Execute a SQL query against a Vortex file.
19///
20/// The file is registered as a table named "data".
21/// Returns the result as a vector of RecordBatches.
22///
23/// # Errors
24///
25/// Returns an error if the query fails to parse or execute.
26pub async fn execute_vortex_query(
27    session: &VortexSession,
28    file_path: &str,
29    sql: &str,
30) -> Result<Vec<RecordBatch>, String> {
31    let ctx = create_vortex_context(session, file_path).await?;
32
33    let df = ctx.sql(sql).await.map_err(|e| format!("SQL error: {e}"))?;
34
35    df.collect()
36        .await
37        .map_err(|e| format!("Query execution error: {e}"))
38}
39
40/// Create a DataFusion SessionContext with a Vortex file registered as "data".
41///
42/// # Errors
43///
44/// Returns an error if the context cannot be created.
45pub async fn create_vortex_context(
46    session: &VortexSession,
47    file_path: &str,
48) -> Result<SessionContext, String> {
49    let ctx = SessionContext::new();
50    let format = Arc::new(VortexFormat::new(session.clone()));
51
52    // Resolve to an absolute, canonical path before handing it to `ListingTableUrl`.
53    // `ListingTable` is collection-oriented: if the object `head` lookup for the URL
54    // fails (which happens for relative or non-canonical paths), DataFusion silently
55    // retries the path as a directory prefix and recursively lists it, crawling
56    // sibling and child directories. Canonicalizing to the existing file keeps us on
57    // the single-file fast path.
58    let canonical = std::fs::canonicalize(file_path)
59        .map_err(|e| format!("Failed to resolve file path '{file_path}': {e}"))?;
60    let canonical = canonical
61        .to_str()
62        .ok_or_else(|| "Resolved path is not valid UTF-8".to_string())?;
63
64    let table_url =
65        ListingTableUrl::parse(canonical).map_err(|e| format!("Failed to parse file path: {e}"))?;
66
67    // Clear the file extension filter (it defaults to `format.get_ext()`, i.e.
68    // "vortex"). The CLI/TUI can open a Vortex file under any name, so registering
69    // it for query must not require a `.vortex` extension. This is safe because the
70    // canonical path above keeps us on the single-file path, so no directory listing
71    // happens that an empty extension could over-match.
72    let listing_options = ListingOptions::new(format)
73        .with_file_extension("")
74        .with_session_config_options(ctx.state().config());
75
76    let config = ListingTableConfig::new(table_url)
77        .with_listing_options(listing_options)
78        .infer_schema(&ctx.state())
79        .await
80        .map_err(|e| format!("Failed to infer schema: {e}"))?;
81
82    let listing_table = Arc::new(
83        ListingTable::try_new(config).map_err(|e| format!("Failed to create table: {e}"))?,
84    );
85
86    ctx.register_table("data", listing_table)
87        .map_err(|e| format!("Failed to register table: {e}"))?;
88
89    Ok(ctx)
90}
91
92/// Convert an Arrow array value at a given index to a JSON value.
93///
94/// # Panics
95///
96/// Panics if the array type doesn't match the expected Arrow array type during downcast.
97/// This should not happen for well-formed Arrow arrays.
98#[expect(clippy::unwrap_used)]
99pub fn arrow_value_to_json(array: &dyn ArrowArray, idx: usize) -> serde_json::Value {
100    use arrow_array::*;
101    use arrow_schema::DataType;
102
103    if array.is_null(idx) {
104        return serde_json::Value::Null;
105    }
106
107    match array.data_type() {
108        DataType::Null => serde_json::Value::Null,
109        DataType::Boolean => {
110            let arr = array.as_any().downcast_ref::<BooleanArray>().unwrap();
111            serde_json::Value::Bool(arr.value(idx))
112        }
113        DataType::Int8 => {
114            let arr = array.as_any().downcast_ref::<Int8Array>().unwrap();
115            serde_json::json!(arr.value(idx))
116        }
117        DataType::Int16 => {
118            let arr = array.as_any().downcast_ref::<Int16Array>().unwrap();
119            serde_json::json!(arr.value(idx))
120        }
121        DataType::Int32 => {
122            let arr = array.as_any().downcast_ref::<Int32Array>().unwrap();
123            serde_json::json!(arr.value(idx))
124        }
125        DataType::Int64 => {
126            let arr = array.as_any().downcast_ref::<Int64Array>().unwrap();
127            serde_json::json!(arr.value(idx))
128        }
129        DataType::UInt8 => {
130            let arr = array.as_any().downcast_ref::<UInt8Array>().unwrap();
131            serde_json::json!(arr.value(idx))
132        }
133        DataType::UInt16 => {
134            let arr = array.as_any().downcast_ref::<UInt16Array>().unwrap();
135            serde_json::json!(arr.value(idx))
136        }
137        DataType::UInt32 => {
138            let arr = array.as_any().downcast_ref::<UInt32Array>().unwrap();
139            serde_json::json!(arr.value(idx))
140        }
141        DataType::UInt64 => {
142            let arr = array.as_any().downcast_ref::<UInt64Array>().unwrap();
143            serde_json::json!(arr.value(idx))
144        }
145        DataType::Float16 => {
146            let arr = array.as_any().downcast_ref::<Float16Array>().unwrap();
147            serde_json::json!(arr.value(idx).to_f32())
148        }
149        DataType::Float32 => {
150            let arr = array.as_any().downcast_ref::<Float32Array>().unwrap();
151            serde_json::json!(arr.value(idx))
152        }
153        DataType::Float64 => {
154            let arr = array.as_any().downcast_ref::<Float64Array>().unwrap();
155            serde_json::json!(arr.value(idx))
156        }
157        DataType::Utf8 => {
158            let arr = array.as_any().downcast_ref::<StringArray>().unwrap();
159            serde_json::Value::String(arr.value(idx).to_string())
160        }
161        DataType::LargeUtf8 => {
162            let arr = array.as_any().downcast_ref::<LargeStringArray>().unwrap();
163            serde_json::Value::String(arr.value(idx).to_string())
164        }
165        DataType::Utf8View => {
166            let arr = array.as_any().downcast_ref::<StringViewArray>().unwrap();
167            serde_json::Value::String(arr.value(idx).to_string())
168        }
169        DataType::Binary => {
170            let arr = array.as_any().downcast_ref::<BinaryArray>().unwrap();
171            let hex: String = arr.value(idx).iter().map(|b| format!("{b:02x}")).collect();
172            serde_json::Value::String(hex)
173        }
174        DataType::LargeBinary => {
175            let arr = array.as_any().downcast_ref::<LargeBinaryArray>().unwrap();
176            let hex: String = arr.value(idx).iter().map(|b| format!("{b:02x}")).collect();
177            serde_json::Value::String(hex)
178        }
179        DataType::BinaryView => {
180            let arr = array.as_any().downcast_ref::<BinaryViewArray>().unwrap();
181            let hex: String = arr.value(idx).iter().map(|b| format!("{b:02x}")).collect();
182            serde_json::Value::String(hex)
183        }
184        DataType::Date32 => {
185            let arr = array.as_any().downcast_ref::<Date32Array>().unwrap();
186            serde_json::json!(arr.value(idx))
187        }
188        DataType::Date64 => {
189            let arr = array.as_any().downcast_ref::<Date64Array>().unwrap();
190            serde_json::json!(arr.value(idx))
191        }
192        DataType::Timestamp(..) => {
193            if let Some(arr) = array.as_any().downcast_ref::<TimestampMicrosecondArray>() {
194                serde_json::json!(arr.value(idx))
195            } else if let Some(arr) = array.as_any().downcast_ref::<TimestampMillisecondArray>() {
196                serde_json::json!(arr.value(idx))
197            } else if let Some(arr) = array.as_any().downcast_ref::<TimestampSecondArray>() {
198                serde_json::json!(arr.value(idx))
199            } else if let Some(arr) = array.as_any().downcast_ref::<TimestampNanosecondArray>() {
200                serde_json::json!(arr.value(idx))
201            } else {
202                serde_json::Value::String("<timestamp>".to_string())
203            }
204        }
205        DataType::Decimal128(..) => {
206            let arr = array.as_any().downcast_ref::<Decimal128Array>().unwrap();
207            serde_json::Value::String(arr.value_as_string(idx))
208        }
209        DataType::Decimal256(..) => {
210            let arr = array.as_any().downcast_ref::<Decimal256Array>().unwrap();
211            serde_json::Value::String(arr.value_as_string(idx))
212        }
213        DataType::List(_) => {
214            let arr = array.as_any().downcast_ref::<ListArray>().unwrap();
215            let value_arr = arr.value(idx);
216            let elements: Vec<serde_json::Value> = (0..value_arr.len())
217                .map(|i| arrow_value_to_json(value_arr.as_ref(), i))
218                .collect();
219            serde_json::Value::Array(elements)
220        }
221        DataType::LargeList(_) => {
222            let arr = array.as_any().downcast_ref::<LargeListArray>().unwrap();
223            let value_arr = arr.value(idx);
224            let elements: Vec<serde_json::Value> = (0..value_arr.len())
225                .map(|i| arrow_value_to_json(value_arr.as_ref(), i))
226                .collect();
227            serde_json::Value::Array(elements)
228        }
229        DataType::Struct(_) => {
230            let arr = array.as_any().downcast_ref::<StructArray>().unwrap();
231            let mut obj = serde_json::Map::new();
232            for (i, field) in arr.fields().iter().enumerate() {
233                let col = arr.column(i);
234                obj.insert(field.name().clone(), arrow_value_to_json(col.as_ref(), idx));
235            }
236            serde_json::Value::Object(obj)
237        }
238        _ => {
239            // Fallback for unsupported types
240            serde_json::Value::String(format!("<{}>", array.data_type()))
241        }
242    }
243}
244
245/// Format a JSON value for display in the TUI.
246///
247/// - Null becomes "NULL"
248/// - Strings are displayed without quotes
249/// - Other values use their JSON string representation
250pub fn json_value_to_display(value: serde_json::Value) -> String {
251    match value {
252        serde_json::Value::Null => "NULL".to_string(),
253        serde_json::Value::String(s) => s,
254        other => other.to_string(),
255    }
256}
vortex_tui/datafusion_helper.rs

vortex_tui/
datafusion_helper.rs
