fv_plan/
convert.rs

1//! Row ⇄ Arrow conversion for the DataFusion execution path (reshapes). Column types are inferred
2//! from the values (real pipelines will pass declared schemas later); the value dialect stays dynamic.
3
4use std::sync::Arc;
5
6use datafusion::arrow::array::{
7    Array, BooleanArray, Float32Array, Float64Array, Int16Array, Int32Array, Int64Array, Int8Array, LargeListArray,
8    LargeStringArray, ListArray, StringArray, StringViewArray, StructArray, UInt16Array, UInt32Array, UInt64Array,
9    UInt8Array,
10};
11use datafusion::arrow::datatypes::{DataType, Field, Schema};
12use datafusion::arrow::record_batch::RecordBatch;
13use fv_value::Value;
14
15use crate::row::Row;
16
17/// Infer the Arrow type for a column across rows: numbers → Float64, booleans → Boolean, anything
18/// with strings or mixed/complex values → Utf8 (JSON-encoded), all-null → Float64 (round-trips Null).
19pub fn infer_type(rows: &[Row], col: &str) -> DataType {
20    infer(rows, col)
21}
22
23fn infer(rows: &[Row], col: &str) -> DataType {
24    let (mut num, mut boolean, mut string, mut other) = (false, false, false, false);
25    for r in rows {
26        match r.get(col) {
27            Value::Null => {}
28            Value::Num(_) => num = true,
29            Value::Bool(_) => boolean = true,
30            Value::Str(_) => string = true,
31            _ => other = true,
32        }
33    }
34    if other || string {
35        DataType::Utf8
36    } else if boolean && !num {
37        DataType::Boolean
38    } else {
39        DataType::Float64
40    }
41}
42
43fn scalar_to_string(v: &Value) -> Option<String> {
44    match v {
45        Value::Null => None,
46        Value::Str(s) => Some(s.clone()),
47        Value::Bool(b) => Some(if *b { "true" } else { "false" }.to_string()),
48        Value::Num(n) => Some(if n.fract() == 0.0 && n.is_finite() {
49            format!("{}", *n as i64)
50        } else {
51            format!("{n}")
52        }),
53        // lists / objects → a stable JSON-ish encoding so grouping/joining on them is well-defined
54        other => Some(format!("{other:?}")),
55    }
56}
57
58/// Build a RecordBatch from rows in the given column order (types inferred).
59pub fn rows_to_batch(rows: &[Row], columns: &[String]) -> Result<RecordBatch, String> {
60    let mut fields = Vec::with_capacity(columns.len());
61    let mut arrays: Vec<Arc<dyn Array>> = Vec::with_capacity(columns.len());
62    for col in columns {
63        let dt = infer(rows, col);
64        fields.push(Field::new(col, dt.clone(), true));
65        let arr: Arc<dyn Array> = match dt {
66            DataType::Float64 => Arc::new(Float64Array::from(
67                rows.iter().map(|r| r.get(col).as_num()).collect::<Vec<_>>(),
68            )),
69            DataType::Boolean => Arc::new(BooleanArray::from(
70                rows.iter()
71                    .map(|r| match r.get(col) {
72                        Value::Bool(b) => Some(b),
73                        _ => None,
74                    })
75                    .collect::<Vec<_>>(),
76            )),
77            _ => Arc::new(StringArray::from(
78                rows.iter().map(|r| scalar_to_string(&r.get(col))).collect::<Vec<_>>(),
79            )),
80        };
81        arrays.push(arr);
82    }
83    let schema = Arc::new(Schema::new(fields));
84    RecordBatch::try_new(schema, arrays).map_err(|e| e.to_string())
85}
86
87/// Read a RecordBatch back into ordered rows (schema field order preserved).
88pub fn batch_to_rows(batch: &RecordBatch) -> Vec<Row> {
89    let schema = batch.schema();
90    let mut rows = vec![Row::new(); batch.num_rows()];
91    for (c, field) in schema.fields().iter().enumerate() {
92        let col = batch.column(c);
93        let name = field.name();
94        for (i, row) in rows.iter_mut().enumerate() {
95            let v = read_cell(col, i);
96            row.0.push((name.clone(), v));
97        }
98    }
99    rows
100}
101
102/// Decode a single Arrow cell to a dialect Value (used by the value-dialect ScalarUDF, per row).
103pub fn array_value(col: &Arc<dyn Array>, i: usize) -> Value {
104    read_cell(col, i)
105}
106
107fn read_cell(col: &Arc<dyn Array>, i: usize) -> Value {
108    if col.is_null(i) {
109        return Value::Null;
110    }
111    // DataFusion produces various numeric widths (COUNT → Int64, ROW_NUMBER/RANK → UInt64, …); all
112    // collapse to the dialect's single numeric type (f64).
113    match col.data_type() {
114        DataType::Float64 => Value::Num(col.as_any().downcast_ref::<Float64Array>().unwrap().value(i)),
115        DataType::Float32 => Value::Num(col.as_any().downcast_ref::<Float32Array>().unwrap().value(i) as f64),
116        DataType::Int64 => Value::Num(col.as_any().downcast_ref::<Int64Array>().unwrap().value(i) as f64),
117        DataType::Int32 => Value::Num(col.as_any().downcast_ref::<Int32Array>().unwrap().value(i) as f64),
118        DataType::Int16 => Value::Num(col.as_any().downcast_ref::<Int16Array>().unwrap().value(i) as f64),
119        DataType::Int8 => Value::Num(col.as_any().downcast_ref::<Int8Array>().unwrap().value(i) as f64),
120        DataType::UInt64 => Value::Num(col.as_any().downcast_ref::<UInt64Array>().unwrap().value(i) as f64),
121        DataType::UInt32 => Value::Num(col.as_any().downcast_ref::<UInt32Array>().unwrap().value(i) as f64),
122        DataType::UInt16 => Value::Num(col.as_any().downcast_ref::<UInt16Array>().unwrap().value(i) as f64),
123        DataType::UInt8 => Value::Num(col.as_any().downcast_ref::<UInt8Array>().unwrap().value(i) as f64),
124        DataType::Boolean => Value::Bool(col.as_any().downcast_ref::<BooleanArray>().unwrap().value(i)),
125        DataType::Utf8 => Value::Str(col.as_any().downcast_ref::<StringArray>().unwrap().value(i).to_string()),
126        DataType::LargeUtf8 => Value::Str(
127            col.as_any()
128                .downcast_ref::<LargeStringArray>()
129                .unwrap()
130                .value(i)
131                .to_string(),
132        ),
133        // DataFusion 54 emits Utf8View (StringViewArray) from VARCHAR casts and string kernels —
134        // without this arm every cast string silently collapsed to Null (live-hit: the union
135        // reshape's type-aligning CASTs nulled rid/position and the whole pipeline emptied).
136        DataType::Utf8View => Value::Str(
137            col.as_any()
138                .downcast_ref::<StringViewArray>()
139                .unwrap()
140                .value(i)
141                .to_string(),
142        ),
143        DataType::BinaryView => Value::Null, // no dialect binary type — explicit, not fall-through
144        // ── Nested types (e.g. an Iceberg GEO_SHAPE column read as a native Arrow struct/list) ──
145        // Decoded recursively so a nested GeoJSON object round-trips to a `Value::Obj`/`Value::List`
146        // rather than collapsing to Null.
147        DataType::Struct(fields) => {
148            let arr = col.as_any().downcast_ref::<StructArray>().unwrap();
149            Value::Obj(
150                fields
151                    .iter()
152                    .enumerate()
153                    .map(|(c, f)| (f.name().clone(), read_cell(arr.column(c), i)))
154                    .collect(),
155            )
156        }
157        DataType::List(_) => read_list(&col.as_any().downcast_ref::<ListArray>().unwrap().value(i)),
158        DataType::LargeList(_) => read_list(&col.as_any().downcast_ref::<LargeListArray>().unwrap().value(i)),
159        _ => Value::Null,
160    }
161}
162
163/// Decode one row's list sub-array (each element recursively) → a `Value::List`.
164fn read_list(sub: &Arc<dyn Array>) -> Value {
165    Value::List((0..sub.len()).map(|j| read_cell(sub, j)).collect())
166}
167
168#[cfg(test)]
169mod tests {
170    use super::*;
171    use datafusion::arrow::array::ArrayRef;
172    use datafusion::arrow::datatypes::{Field, Float64Type};
173
174    fn num(n: f64) -> Value {
175        Value::Num(n)
176    }
177
178    #[test]
179    fn decodes_a_list_column_to_value_list() {
180        let list: ArrayRef = Arc::new(ListArray::from_iter_primitive::<Float64Type, _, _>(vec![
181            Some(vec![Some(1.0), Some(2.0)]),
182            Some(vec![Some(3.0)]),
183        ]));
184        assert_eq!(array_value(&list, 0), Value::List(vec![num(1.0), num(2.0)]));
185        assert_eq!(array_value(&list, 1), Value::List(vec![num(3.0)]));
186    }
187
188    #[test]
189    fn decodes_a_struct_column_to_value_obj() {
190        let ty: ArrayRef = Arc::new(StringArray::from(vec!["Polygon"]));
191        let coords: ArrayRef = Arc::new(ListArray::from_iter_primitive::<Float64Type, _, _>(vec![Some(vec![
192            Some(1.0),
193            Some(2.0),
194        ])]));
195        let s: ArrayRef = Arc::new(StructArray::from(vec![
196            (Arc::new(Field::new("type", DataType::Utf8, true)), ty),
197            (
198                Arc::new(Field::new("coordinates", coords.data_type().clone(), true)),
199                coords,
200            ),
201        ]));
202        let got = array_value(&s, 0);
203        let want = Value::Obj(
204            [
205                ("type".to_string(), Value::Str("Polygon".to_string())),
206                ("coordinates".to_string(), Value::List(vec![num(1.0), num(2.0)])),
207            ]
208            .into_iter()
209            .collect(),
210        );
211        assert_eq!(got, want);
212    }
213
214    #[test]
215    fn decodes_nested_lists_of_lists() {
216        // A GeoJSON-shaped List<List<Float64>> (e.g. a ring of positions).
217        let inner = ListArray::from_iter_primitive::<Float64Type, _, _>(vec![
218            Some(vec![Some(1.0), Some(2.0)]),
219            Some(vec![Some(3.0), Some(4.0)]),
220        ]);
221        // Wrap the two inner lists into one outer list element.
222        use datafusion::arrow::buffer::OffsetBuffer;
223        let field = Arc::new(Field::new("item", inner.data_type().clone(), true));
224        let outer: ArrayRef = Arc::new(ListArray::new(
225            field,
226            OffsetBuffer::new(vec![0, 2].into()),
227            Arc::new(inner),
228            None,
229        ));
230        assert_eq!(
231            array_value(&outer, 0),
232            Value::List(vec![
233                Value::List(vec![num(1.0), num(2.0)]),
234                Value::List(vec![num(3.0), num(4.0)]),
235            ])
236        );
237    }
238
239    #[test]
240    fn a_null_nested_cell_reads_null() {
241        let list: ArrayRef = Arc::new(ListArray::from_iter_primitive::<Float64Type, _, _>(vec![
242            Some(vec![Some(1.0)]),
243            None,
244        ]));
245        assert_eq!(array_value(&list, 1), Value::Null);
246    }
247}
fv_plan/convert.rs

fv_plan/
convert.rs
