Skip to content

Commit 4c7bdf4

Browse files
committed
refactor: replace datafusion::arrow with arrow crate, delete dead UDFs
- Add arrow 58 as direct dependency, replace all datafusion::arrow re-exports in 10 non-UDF files (response_codec, cold storage, etc.) - Delete inline_rewrite.rs and user_function.rs (dead — planning no longer uses DataFusion UDF inlining) - Remove register_udfs_on from context.rs (no callers) - Replace validate.rs DataFusion-based validation with sqlparser-based syntax check - eval.rs: replace datafusion::arrow with arrow crate DataFusion remains only for: - Procedural executor expression evaluation (eval_sql_expr) - UDF trait implementations (still needed by eval.rs indirectly) - Temp table MemTable storage
1 parent 72b8a86 commit 4c7bdf4

19 files changed

Lines changed: 76 additions & 635 deletions

File tree

Cargo.toml

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -130,7 +130,10 @@ snap = "1"
130130
flate2 = "1"
131131
tonic = "0.14"
132132

133-
# Query engine
133+
# Arrow (columnar format)
134+
arrow = { version = "58", default-features = false, features = ["ipc"] }
135+
136+
# Query engine (retained for procedural executor + function body validation)
134137
datafusion = { version = "53", default-features = false, features = ["sql"] }
135138
datafusion-common = "53"
136139
datafusion-execution = "53"

nodedb/Cargo.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@ nodedb-spatial = { workspace = true }
3636
nodedb-graph = { workspace = true }
3737
nodedb-vector = { workspace = true, features = ["collection"] }
3838
nodedb-sql = { workspace = true }
39+
arrow = { workspace = true }
3940

4041
# Async runtime (Control Plane)
4142
tokio = { workspace = true }

nodedb/src/control/arrow_convert.rs

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -13,9 +13,9 @@
1313
1414
use std::sync::Arc;
1515

16-
use datafusion::arrow::array::{ArrayRef, Float64Array, Int64Array, StringArray};
17-
use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef};
18-
use datafusion::arrow::record_batch::RecordBatch;
16+
use arrow::array::{ArrayRef, Float64Array, Int64Array, StringArray};
17+
use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
18+
use arrow::record_batch::RecordBatch;
1919

2020
/// Convert a JSON array of document rows into an Arrow RecordBatch.
2121
///
@@ -163,9 +163,9 @@ pub fn arrow_sum(batch: &RecordBatch, column_name: &str) -> Option<f64> {
163163
let array = batch.column(idx);
164164

165165
if let Some(f64_arr) = array.as_any().downcast_ref::<Float64Array>() {
166-
Some(datafusion::arrow::compute::kernels::aggregate::sum(f64_arr).unwrap_or(0.0))
166+
Some(arrow::compute::kernels::aggregate::sum(f64_arr).unwrap_or(0.0))
167167
} else if let Some(i64_arr) = array.as_any().downcast_ref::<Int64Array>() {
168-
let sum = datafusion::arrow::compute::kernels::aggregate::sum(i64_arr).unwrap_or(0);
168+
let sum = arrow::compute::kernels::aggregate::sum(i64_arr).unwrap_or(0);
169169
Some(sum as f64)
170170
} else {
171171
None
@@ -178,9 +178,9 @@ pub fn arrow_min(batch: &RecordBatch, column_name: &str) -> Option<f64> {
178178
let array = batch.column(idx);
179179

180180
if let Some(f64_arr) = array.as_any().downcast_ref::<Float64Array>() {
181-
datafusion::arrow::compute::kernels::aggregate::min(f64_arr)
181+
arrow::compute::kernels::aggregate::min(f64_arr)
182182
} else if let Some(i64_arr) = array.as_any().downcast_ref::<Int64Array>() {
183-
datafusion::arrow::compute::kernels::aggregate::min(i64_arr).map(|v| v as f64)
183+
arrow::compute::kernels::aggregate::min(i64_arr).map(|v| v as f64)
184184
} else {
185185
None
186186
}
@@ -192,9 +192,9 @@ pub fn arrow_max(batch: &RecordBatch, column_name: &str) -> Option<f64> {
192192
let array = batch.column(idx);
193193

194194
if let Some(f64_arr) = array.as_any().downcast_ref::<Float64Array>() {
195-
datafusion::arrow::compute::kernels::aggregate::max(f64_arr)
195+
arrow::compute::kernels::aggregate::max(f64_arr)
196196
} else if let Some(i64_arr) = array.as_any().downcast_ref::<Int64Array>() {
197-
datafusion::arrow::compute::kernels::aggregate::max(i64_arr).map(|v| v as f64)
197+
arrow::compute::kernels::aggregate::max(i64_arr).map(|v| v as f64)
198198
} else {
199199
None
200200
}
@@ -223,7 +223,7 @@ pub fn arrow_avg(batch: &RecordBatch, column_name: &str) -> Option<f64> {
223223
/// Used by the Control Plane to receive Arrow data from the Data Plane
224224
/// across the SPSC bridge.
225225
pub fn decode_arrow_ipc(bytes: &[u8]) -> Option<RecordBatch> {
226-
use datafusion::arrow::ipc::reader::StreamReader;
226+
use arrow::ipc::reader::StreamReader;
227227
let cursor = std::io::Cursor::new(bytes);
228228
let mut reader = StreamReader::try_new(cursor, None).ok()?;
229229
reader.next()?.ok()

nodedb/src/control/planner/context.rs

Lines changed: 0 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -207,42 +207,3 @@ pub const SYSTEM_FUNCTION_NAMES: &[&str] = &[
207207
"setval",
208208
"next_preview",
209209
];
210-
211-
/// Register all system UDFs on a DataFusion `SessionContext`.
212-
///
213-
/// Used only by the DDL body validator (CREATE FUNCTION) which needs
214-
/// a temporary session with all system UDFs registered for validation.
215-
pub fn register_udfs_on(session: &datafusion::execution::context::SessionContext) {
216-
use super::udf::spatial::{
217-
GeoDistance, StContains, StDistance, StDwithin, StIntersects, StWithin,
218-
};
219-
use super::udf::{
220-
Allocate, Bm25Score, ChunkText, ConvertCurrency, Distribute, DocArrayContains, DocExists,
221-
DocGet, MultiVectorScore, MultiVectorSearch, RoundDecimal, RrfScore, SparseScore,
222-
TextMatch, VectorDistance, VectorMetadata,
223-
};
224-
use datafusion::logical_expr::ScalarUDF;
225-
session.register_udf(ScalarUDF::new_from_impl(ChunkText::new()));
226-
session.register_udf(ScalarUDF::new_from_impl(DocGet::new()));
227-
session.register_udf(ScalarUDF::new_from_impl(DocExists::new()));
228-
session.register_udf(ScalarUDF::new_from_impl(DocArrayContains::new()));
229-
session.register_udf(ScalarUDF::new_from_impl(VectorDistance::new()));
230-
session.register_udf(ScalarUDF::new_from_impl(VectorMetadata::new()));
231-
session.register_udf(ScalarUDF::new_from_impl(MultiVectorScore::new()));
232-
session.register_udf(ScalarUDF::new_from_impl(MultiVectorSearch::new()));
233-
session.register_udf(ScalarUDF::new_from_impl(RrfScore::new()));
234-
session.register_udf(ScalarUDF::new_from_impl(SparseScore::new()));
235-
session.register_udf(ScalarUDF::new_from_impl(Bm25Score::new()));
236-
session.register_udf(ScalarUDF::new_from_impl(TextMatch::new()));
237-
session.register_udf(ScalarUDF::new_from_impl(StDwithin::new()));
238-
session.register_udf(ScalarUDF::new_from_impl(StContains::new()));
239-
session.register_udf(ScalarUDF::new_from_impl(StIntersects::new()));
240-
session.register_udf(ScalarUDF::new_from_impl(StWithin::new()));
241-
session.register_udf(ScalarUDF::new_from_impl(StDistance::new()));
242-
session.register_udf(ScalarUDF::new_from_impl(GeoDistance::new()));
243-
session.register_udf(ScalarUDF::new_from_impl(RoundDecimal::new("HALF_EVEN")));
244-
session.register_udf(ScalarUDF::new_from_impl(Distribute::new()));
245-
session.register_udf(ScalarUDF::new_from_impl(Allocate::new()));
246-
session.register_udf(ScalarUDF::new_from_impl(ConvertCurrency::new()));
247-
nodedb_query::ts_udfs::register_timeseries_udfs(session);
248-
}

nodedb/src/control/planner/procedural/executor/arrow_conv.rs

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -6,8 +6,8 @@
66
77
use std::sync::Arc;
88

9-
use datafusion::arrow::array::*;
10-
use datafusion::arrow::datatypes::DataType;
9+
use arrow::array::*;
10+
use arrow::datatypes::DataType;
1111
use nodedb_types::Value;
1212

1313
/// Extract a single scalar value from an Arrow array at the given row index.
@@ -87,10 +87,9 @@ pub fn arrow_scalar_to_value(col: &Arc<dyn Array>, row: usize) -> Value {
8787
.map(|a| Value::String(a.value(row).to_string()))
8888
.unwrap_or(Value::Null),
8989
_ => {
90-
// Fallback: format as string via ScalarValue.
91-
let scalar = datafusion::common::ScalarValue::try_from_array(col, row);
92-
match scalar {
93-
Ok(s) => Value::String(s.to_string()),
90+
// Fallback: format as display string.
91+
match arrow::util::display::array_value_to_string(col, row) {
92+
Ok(s) => Value::String(s),
9493
Err(_) => Value::Null,
9594
}
9695
}

nodedb/src/control/planner/procedural/executor/eval.rs

Lines changed: 12 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,11 @@
11
//! Expression evaluation helpers for the statement executor.
22
//!
3-
//! Evaluates SQL expressions via DataFusion and converts results to
4-
//! Rust types (bool, i64, nodedb_types::Value). Used by ASSIGN, IF/WHILE
5-
//! conditions, FOR bounds, and OUT parameter capture.
3+
//! Evaluates SQL expressions and converts results to Rust types (bool, i64,
4+
//! nodedb_types::Value). Used by ASSIGN, IF/WHILE conditions, FOR bounds,
5+
//! and OUT parameter capture.
6+
//!
7+
//! Uses DataFusion for complex expression evaluation (arithmetic, CASE, etc.).
8+
//! Literal fast paths avoid DataFusion for simple constants.
69
710
use crate::control::state::SharedState;
811
use crate::types::TenantId;
@@ -29,13 +32,13 @@ pub async fn evaluate_condition(
2932
let col = batch.column(0);
3033
if let Some(bool_arr) = col
3134
.as_any()
32-
.downcast_ref::<datafusion::arrow::array::BooleanArray>()
35+
.downcast_ref::<arrow::array::BooleanArray>()
3336
{
3437
return Ok(bool_arr.value(0));
3538
}
3639
if let Some(int_arr) = col
3740
.as_any()
38-
.downcast_ref::<datafusion::arrow::array::Int32Array>()
41+
.downcast_ref::<arrow::array::Int32Array>()
3942
{
4043
return Ok(int_arr.value(0) != 0);
4144
}
@@ -64,13 +67,13 @@ pub async fn evaluate_int(
6467
let col = batch.column(0);
6568
if let Some(arr) = col
6669
.as_any()
67-
.downcast_ref::<datafusion::arrow::array::Int64Array>()
70+
.downcast_ref::<arrow::array::Int64Array>()
6871
{
6972
return Ok(arr.value(0));
7073
}
7174
if let Some(arr) = col
7275
.as_any()
73-
.downcast_ref::<datafusion::arrow::array::Int32Array>()
76+
.downcast_ref::<arrow::array::Int32Array>()
7477
{
7578
return Ok(arr.value(0) as i64);
7679
}
@@ -128,13 +131,13 @@ pub async fn evaluate_to_value(
128131
/// Execute a SQL expression via DataFusion and return the result batches.
129132
///
130133
/// Wraps the expression in `SELECT (<expr>) as __result` and collects.
131-
/// Uses a standalone DataFusion session (not QueryContext) for scalar evaluation.
134+
/// Uses a standalone DataFusion session for scalar evaluation only.
132135
pub async fn eval_sql_expr(
133136
_state: &SharedState,
134137
_tenant_id: TenantId,
135138
expr: &str,
136139
context: &str,
137-
) -> crate::Result<Vec<datafusion::arrow::record_batch::RecordBatch>> {
140+
) -> crate::Result<Vec<arrow::record_batch::RecordBatch>> {
138141
let session = datafusion::execution::context::SessionContext::new();
139142
let select_sql = format!("SELECT ({expr}) as __result");
140143
let df = session

0 commit comments

Comments
 (0)