fv_plan/
reshape.rs

1//! Reshape steps on DataFusion — the executors that genuinely need a relational engine (they change
2//! cardinality/shape, unlike the row-wise inline ops). This crate lands `aggregate` (group-by) on
3//! DataFusion's SQL planner/executor via a MemTable; join/window/union follow. Value↔Arrow via
4//! `convert`. Semantics are SQL's (null semantics: sum/avg/min/max ignore
5//! nulls, all-null group → null; count(col) = non-null count, count(*) = rows).
6
7use std::sync::Arc;
8
9use datafusion::arrow::datatypes::DataType;
10use datafusion::datasource::MemTable;
11use datafusion::prelude::SessionContext;
12
13use crate::session::geo_context;
14use fv_value::Value;
15use serde_json::Value as J;
16
17use crate::convert::{batch_to_rows, infer_type, rows_to_batch};
18use crate::row::Row;
19
20/// One transform input: its declared/observed columns + rows.
21pub struct Input {
22    pub columns: Vec<String>,
23    pub rows: Vec<Row>,
24}
25
26/// The reshape (whole-transform) step ops this engine executes natively on DataFusion. Adding a
27/// reshape step = a new executor fn + one `run` arm + one entry here — all in THIS file.
28pub const NATIVE_OPS: &[&str] = &["join", "aggregate", "window", "union"];
29
30/// How a native op's generated SQL is executed — the injection seam that lets the SAME reshape SQL
31/// run single-node ([`LocalSqlRunner`]) or distributed (the runner's Ballista impl over
32/// `remote_with_state` + Parquet-staged inputs), without `fv-plan` depending on Ballista. The op
33/// builds the `(name, input)` table bindings + SQL; the runner registers + executes them.
34#[async_trait::async_trait]
35pub trait SqlRunner: Send + Sync {
36    async fn run(&self, tables: &[(&str, &Input)], sql: &str) -> Result<Vec<Row>, String>;
37}
38
39/// The default single-node runner: each input a `MemTable` on a `geo_context` (ST_* available).
40pub struct LocalSqlRunner;
41
42#[async_trait::async_trait]
43impl SqlRunner for LocalSqlRunner {
44    async fn run(&self, tables: &[(&str, &Input)], sql: &str) -> Result<Vec<Row>, String> {
45        let ctx = geo_context();
46        for (name, input) in tables {
47            register(&ctx, name, input)?;
48        }
49        run_sql(&ctx, sql).await
50    }
51}
52
53fn str_list(j: &J) -> Vec<String> {
54    j.as_array()
55        .map(|a| a.iter().filter_map(|v| v.as_str().map(String::from)).collect())
56        .unwrap_or_default()
57}
58
59/// Execute one reshape step (decoded from its JSON spec) over the transform's inputs, single-node.
60pub async fn run(step: &J, inputs: Vec<Input>) -> Result<Vec<Row>, String> {
61    run_with(step, inputs, &LocalSqlRunner).await
62}
63
64/// As [`run`], but the generated SQL is executed by `runner` — single-node ([`LocalSqlRunner`]) or
65/// distributed (the runner service's Ballista impl). The op logic + SQL are identical either way.
66pub async fn run_with(step: &J, inputs: Vec<Input>, runner: &dyn SqlRunner) -> Result<Vec<Row>, String> {
67    match step["op"].as_str().unwrap_or("") {
68        "aggregate" => {
69            let group_by = str_list(&step["groupBy"]);
70            let aggs: Vec<Aggregation> = step["aggregations"]
71                .as_array()
72                .map(|a| {
73                    a.iter()
74                        .map(|g| Aggregation {
75                            op: g["op"].as_str().unwrap_or("").to_string(),
76                            column: g["column"].as_str().map(String::from),
77                            alias: g["as"].as_str().unwrap_or("").to_string(),
78                        })
79                        .collect()
80                })
81                .unwrap_or_default();
82            aggregate(&inputs[0], &group_by, &aggs, runner).await
83        }
84        "join" => {
85            join(
86                &inputs[0],
87                &inputs[1],
88                step["how"].as_str().unwrap_or("inner"),
89                &str_list(&step["on"]),
90                runner,
91            )
92            .await
93        }
94        "union" => union(&inputs, runner).await,
95        "window" => {
96            let funcs: Vec<WindowFn> = step["functions"]
97                .as_array()
98                .map(|a| {
99                    a.iter()
100                        .map(|g| WindowFn {
101                            op: g["op"].as_str().unwrap_or("").to_string(),
102                            column: g["column"].as_str().map(String::from),
103                            alias: g["as"].as_str().unwrap_or("").to_string(),
104                        })
105                        .collect()
106                })
107                .unwrap_or_default();
108            window(
109                &inputs[0],
110                &str_list(&step["partitionBy"]),
111                &str_list(&step["orderBy"]),
112                &funcs,
113                runner,
114            )
115            .await
116        }
117        other => Err(format!("unknown native op '{other}'")),
118    }
119}
120
121/// Register an input as a named MemTable in `ctx`.
122fn register(ctx: &SessionContext, name: &str, input: &Input) -> Result<(), String> {
123    let batch = rows_to_batch(&input.rows, &input.columns)?;
124    let provider = MemTable::try_new(batch.schema(), vec![vec![batch]]).map_err(|e| e.to_string())?;
125    ctx.register_table(name, Arc::new(provider))
126        .map_err(|e| e.to_string())?;
127    Ok(())
128}
129
130/// Run a generated SQL query and read the result back into ordered rows.
131async fn run_sql(ctx: &SessionContext, sql: &str) -> Result<Vec<Row>, String> {
132    let df = ctx.sql(sql).await.map_err(|e| e.to_string())?;
133    let batches = df.collect().await.map_err(|e| e.to_string())?;
134    Ok(batches.iter().flat_map(batch_to_rows).collect())
135}
136
137/// The DataFusion SQL type name for a CAST (used by `union` to align branch types).
138fn sql_type(dt: &DataType) -> &'static str {
139    match dt {
140        DataType::Boolean => "BOOLEAN",
141        DataType::Utf8 => "VARCHAR",
142        _ => "DOUBLE",
143    }
144}
145
146#[derive(Debug, Clone)]
147pub struct Aggregation {
148    pub op: String, // sum | avg | min | max | count
149    pub column: Option<String>,
150    pub alias: String,
151}
152
153/// Aggregate (group-by) the single input on DataFusion, executed via `runner`.
154pub async fn aggregate(
155    input: &Input,
156    group_by: &[String],
157    aggregations: &[Aggregation],
158    runner: &dyn SqlRunner,
159) -> Result<Vec<Row>, String> {
160    let mut selects: Vec<String> = group_by.iter().map(|g| q(g)).collect();
161    for a in aggregations {
162        selects.push(format!("{} AS {}", agg_expr(a)?, q(&a.alias)));
163    }
164    let mut sql = format!("SELECT {} FROM input", selects.join(", "));
165    if !group_by.is_empty() {
166        sql.push_str(&format!(
167            " GROUP BY {}",
168            group_by.iter().map(|g| q(g)).collect::<Vec<_>>().join(", ")
169        ));
170    }
171    runner.run(&[("input", input)], &sql).await
172}
173
174fn agg_expr(a: &Aggregation) -> Result<String, String> {
175    let col = || {
176        a.column
177            .as_deref()
178            .map(q)
179            .ok_or_else(|| format!("{}: column required", a.op))
180    };
181    Ok(match a.op.as_str() {
182        "sum" => format!("SUM({})", col()?),
183        "avg" => format!("AVG({})", col()?),
184        "min" => format!("MIN({})", col()?),
185        "max" => format!("MAX({})", col()?),
186        "count" => match &a.column {
187            Some(c) => format!("COUNT({})", q(c)),
188            None => "COUNT(*)".to_string(),
189        },
190        other => return Err(format!("unknown aggregation op '{other}'")),
191    })
192}
193
194/// Quote a DataFusion SQL identifier (our column names are simple, but quote to be safe).
195fn q(ident: &str) -> String {
196    format!("\"{}\"", ident.replace('"', "\"\""))
197}
198
199/// Equi-join two inputs on shared key columns. Output = left columns + right's NON-KEY columns, with
200/// right winning on a non-key name collision (a left join keeps the left value where the right is
201/// absent). NB uses SQL equi-join semantics: NULL keys do not match (the standard, and what any real
202/// engine does) — this is a deliberate, documented refinement over the reference's dict semantics.
203pub async fn join(
204    left: &Input,
205    right: &Input,
206    how: &str,
207    on: &[String],
208    runner: &dyn SqlRunner,
209) -> Result<Vec<Row>, String> {
210    let is_left = how.eq_ignore_ascii_case("left");
211    let right_nonkey: Vec<&String> = right.columns.iter().filter(|c| !on.contains(c)).collect();
212
213    let mut selects: Vec<String> = Vec::new();
214    for c in &left.columns {
215        if right_nonkey.contains(&c) {
216            // collision: right wins on a match; a left join falls back to the left value.
217            if is_left {
218                selects.push(format!("COALESCE(r.{c}, l.{c}) AS {c}", c = q(c)));
219            } else {
220                selects.push(format!("r.{c} AS {c}", c = q(c)));
221            }
222        } else {
223            selects.push(format!("l.{c} AS {c}", c = q(c)));
224        }
225    }
226    for c in &right_nonkey {
227        if !left.columns.iter().any(|lc| &lc == c) {
228            selects.push(format!("r.{c} AS {c}", c = q(c)));
229        }
230    }
231
232    let join_kw = if is_left { "LEFT JOIN" } else { "JOIN" };
233    let cond = on
234        .iter()
235        .map(|k| format!("l.{k} = r.{k}", k = q(k)))
236        .collect::<Vec<_>>()
237        .join(" AND ");
238    let sql = format!("SELECT {} FROM l {join_kw} r ON {cond}", selects.join(", "));
239    runner.run(&[("l", left), ("r", right)], &sql).await
240}
241
242/// Stack all inputs (UNION ALL by name): the unified column set is the first-seen order across
243/// inputs; a column missing from an input lands null. Each branch casts to the unified column type.
244pub async fn union(inputs: &[Input], runner: &dyn SqlRunner) -> Result<Vec<Row>, String> {
245    // An EMPTY input contributes nothing — skip it rather than registering a zero-column batch
246    // (arrow refuses those). A regional feed going momentarily dark must not fail the whole union;
247    // the output guard belongs to the caller's expectations (e.g. rowCountMin).
248    let inputs: Vec<&Input> = inputs.iter().filter(|i| !i.rows.is_empty()).collect();
249    if inputs.is_empty() {
250        return Ok(Vec::new());
251    }
252    let mut columns: Vec<String> = Vec::new();
253    for inp in &inputs {
254        for c in &inp.columns {
255            if !columns.contains(c) {
256                columns.push(c.clone());
257            }
258        }
259    }
260    // Reconcile each column's type across the inputs that carry it (any string/mixed → Utf8).
261    let types: Vec<DataType> = columns
262        .iter()
263        .map(|c| {
264            let (mut s, mut b, mut num) = (false, false, false);
265            for inp in &inputs {
266                if inp.columns.contains(c) {
267                    match infer_type(&inp.rows, c) {
268                        DataType::Utf8 => s = true,
269                        DataType::Boolean => b = true,
270                        _ => num = true,
271                    }
272                }
273            }
274            if s || (b && num) {
275                DataType::Utf8
276            } else if b {
277                DataType::Boolean
278            } else {
279                DataType::Float64
280            }
281        })
282        .collect();
283
284    let names: Vec<String> = (0..inputs.len()).map(|i| format!("u{i}")).collect();
285    let mut branches = Vec::new();
286    for (i, inp) in inputs.iter().enumerate() {
287        let name = &names[i];
288        let sel: Vec<String> = columns
289            .iter()
290            .zip(&types)
291            .map(|(c, dt)| {
292                if inp.columns.contains(c) {
293                    format!("CAST({col} AS {ty}) AS {col}", col = q(c), ty = sql_type(dt))
294                } else {
295                    format!("CAST(NULL AS {ty}) AS {col}", col = q(c), ty = sql_type(dt))
296                }
297            })
298            .collect();
299        branches.push(format!("SELECT {} FROM {name}", sel.join(", ")));
300    }
301    let bindings: Vec<(&str, &Input)> = names.iter().map(String::as_str).zip(inputs.iter().copied()).collect();
302    runner.run(&bindings, &branches.join(" UNION ALL ")).await
303}
304
305#[derive(Debug, Clone)]
306pub struct WindowFn {
307    pub op: String, // rowNumber | rank | denseRank | cumSum | cumCount
308    pub column: Option<String>,
309    pub alias: String,
310}
311
312/// Keep every input row and add window-function columns over partitions. Output row order is the
313/// input order (a synthetic order column restores it and breaks ties deterministically, mirroring the
314/// reference's stable sort).
315pub async fn window(
316    input: &Input,
317    partition_by: &[String],
318    order_by: &[String],
319    funcs: &[WindowFn],
320    runner: &dyn SqlRunner,
321) -> Result<Vec<Row>, String> {
322    const ORD: &str = "__fv_ord";
323    let mut cols = input.columns.clone();
324    cols.push(ORD.to_string());
325    let mut rows = input.rows.clone();
326    for (i, r) in rows.iter_mut().enumerate() {
327        r.0.push((ORD.to_string(), Value::Num(i as f64)));
328    }
329    let augmented = Input { columns: cols, rows };
330
331    let part = if partition_by.is_empty() {
332        String::new()
333    } else {
334        format!(
335            "PARTITION BY {} ",
336            partition_by.iter().map(|c| q(c)).collect::<Vec<_>>().join(", ")
337        )
338    };
339    // Stable order = the step's orderBy plus the input-order tiebreaker (for rowNumber / running aggs).
340    let mut stable: Vec<String> = order_by.iter().map(|c| q(c)).collect();
341    stable.push(q(ORD));
342    let ord_stable = stable.join(", ");
343    // Peer order = the step's orderBy only (so RANK/DENSE_RANK give ties equal ranks).
344    let ord_peers = if order_by.is_empty() {
345        q(ORD)
346    } else {
347        order_by.iter().map(|c| q(c)).collect::<Vec<_>>().join(", ")
348    };
349
350    let running = format!("OVER ({part}ORDER BY {ord_stable} ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW)");
351    let numbered = format!("OVER ({part}ORDER BY {ord_stable})");
352    let ranked = format!("OVER ({part}ORDER BY {ord_peers})");
353
354    let mut selects: Vec<String> = input.columns.iter().map(|c| q(c)).collect();
355    for f in funcs {
356        let expr = match f.op.as_str() {
357            "rowNumber" => format!("ROW_NUMBER() {numbered}"),
358            "rank" => format!("RANK() {ranked}"),
359            "denseRank" => format!("DENSE_RANK() {ranked}"),
360            "cumSum" => format!(
361                "SUM({}) {running}",
362                q(f.column.as_deref().ok_or("cumSum: column required")?)
363            ),
364            "cumCount" => match &f.column {
365                Some(c) => format!("COUNT({}) {running}", q(c)),
366                None => format!("COUNT(*) {running}"),
367            },
368            other => return Err(format!("unknown window op '{other}'")),
369        };
370        selects.push(format!("{expr} AS {}", q(&f.alias)));
371    }
372    let sql = format!("SELECT {} FROM w ORDER BY {}", selects.join(", "), q(ORD));
373    runner.run(&[("w", &augmented)], &sql).await
374}
fv_plan/reshape.rs

fv_plan/
reshape.rs
