krishiv_sql/
scalar_udf.rs

1//! Scalar SQL-expression user functions, expanded (inlined) into native SQL
2//! before planning.
3//!
4//! A scalar SQL function such as `CREATE FUNCTION tax(x DOUBLE) RETURNS DOUBLE
5//! RETURN x * 1.1` is stored as a parsed body expression plus its parameter
6//! names. When a query references it (`SELECT tax(amount) FROM sales`), every
7//! call is replaced with the body — arguments substituted for parameters — so
8//! the query becomes pure native SQL (`SELECT (amount * 1.1) FROM sales`).
9//!
10//! Because the result is ordinary SQL with no UDF reference, it plans and runs
11//! anywhere the engine runs, INCLUDING distributed execution on the Rust
12//! executors — no Python interpreter and no per-executor function registration
13//! required. This is the portable, distributable counterpart to Python-callable
14//! UDFs, which can only run in the embedded (in-process) engine.
15
16use std::collections::HashMap;
17use std::ops::ControlFlow;
18
19use datafusion::sql::sqlparser::ast::{
20    Expr, FunctionArg, FunctionArgExpr, FunctionArguments, ObjectName, visit_expressions_mut,
21};
22use datafusion::sql::sqlparser::dialect::GenericDialect;
23use datafusion::sql::sqlparser::parser::Parser;
24
25/// A scalar SQL function definition: its (lower-cased) name, ordered parameter
26/// names, and the parsed body expression.
27#[derive(Clone, Debug)]
28pub struct ScalarSqlFunction {
29    pub name: String,
30    pub params: Vec<String>,
31    body: Expr,
32}
33
34impl ScalarSqlFunction {
35    /// Build from a name, parameter names, and a body SQL expression (e.g.
36    /// `"x * 1.1"`). Parameter/function names are matched case-insensitively.
37    pub fn new(name: &str, params: &[String], body_sql: &str) -> Result<Self, String> {
38        let dialect = GenericDialect {};
39        let body = Parser::new(&dialect)
40            .try_with_sql(body_sql)
41            .and_then(|mut p| p.parse_expr())
42            .map_err(|e| format!("invalid scalar function body '{body_sql}': {e}"))?;
43        Ok(Self {
44            name: name.trim().to_lowercase(),
45            params: params.iter().map(|p| p.trim().to_lowercase()).collect(),
46            body,
47        })
48    }
49}
50
51fn object_name_lower(name: &ObjectName) -> String {
52    name.to_string().to_lowercase()
53}
54
55fn unnamed_args(args: &FunctionArguments) -> Option<Vec<Expr>> {
56    match args {
57        FunctionArguments::List(list) => {
58            let mut out = Vec::with_capacity(list.args.len());
59            for arg in &list.args {
60                match arg {
61                    FunctionArg::Unnamed(FunctionArgExpr::Expr(expr)) => out.push(expr.clone()),
62                    _ => return None,
63                }
64            }
65            Some(out)
66        }
67        _ => None,
68    }
69}
70
71/// Replace every parameter identifier in `body` with the matching argument.
72fn substitute(body: &mut Expr, params: &[String], args: &[Expr]) {
73    let _: ControlFlow<()> = visit_expressions_mut(body, |expr| {
74        if let Expr::Identifier(ident) = expr {
75            let name = ident.value.to_lowercase();
76            if let Some(pos) = params.iter().position(|p| *p == name)
77                && let Some(arg) = args.get(pos)
78            {
79                *expr = arg.clone();
80            }
81        }
82        ControlFlow::Continue(())
83    });
84}
85
86/// Inline every call to a registered scalar SQL function in `sql`, returning
87/// pure native SQL. No-op (returns `sql` unchanged) when `funcs` is empty or the
88/// query references none of them. Errors only if the query cannot be parsed.
89pub fn expand_scalar_sql_functions(
90    sql: &str,
91    funcs: &HashMap<String, ScalarSqlFunction>,
92) -> Result<String, String> {
93    if funcs.is_empty() {
94        return Ok(sql.to_string());
95    }
96    let dialect = GenericDialect {};
97    let mut statements = Parser::parse_sql(&dialect, sql)
98        .map_err(|e| format!("cannot parse query for scalar-UDF expansion: {e}"))?;
99
100    // Bounded fixpoint so a function body that itself calls another registered
101    // function is also expanded; the bound guards against a mutually-recursive
102    // definition looping forever.
103    let mut any = false;
104    for _ in 0..32 {
105        let mut changed = false;
106        let _: ControlFlow<()> = visit_expressions_mut(&mut statements, |expr| {
107            if let Expr::Function(func) = expr {
108                let fname = object_name_lower(&func.name);
109                if let Some(def) = funcs.get(&fname)
110                    && let Some(args) = unnamed_args(&func.args)
111                    && args.len() == def.params.len()
112                {
113                    let mut body = def.body.clone();
114                    substitute(&mut body, &def.params, &args);
115                    // Parenthesize to preserve precedence at the call site.
116                    *expr = Expr::Nested(Box::new(body));
117                    changed = true;
118                }
119            }
120            ControlFlow::Continue(())
121        });
122        any |= changed;
123        if !changed {
124            break;
125        }
126    }
127
128    if !any {
129        return Ok(sql.to_string());
130    }
131    Ok(statements
132        .iter()
133        .map(ToString::to_string)
134        .collect::<Vec<_>>()
135        .join("; "))
136}
137
138#[cfg(test)]
139mod tests {
140    use super::*;
141
142    fn reg(defs: &[(&str, &[&str], &str)]) -> HashMap<String, ScalarSqlFunction> {
143        defs.iter()
144            .map(|(n, p, b)| {
145                let params: Vec<String> = p.iter().map(|s| s.to_string()).collect();
146                (
147                    n.to_lowercase(),
148                    ScalarSqlFunction::new(n, &params, b).unwrap(),
149                )
150            })
151            .collect()
152    }
153
154    fn norm(s: &str) -> String {
155        s.chars()
156            .filter(|c| !c.is_whitespace())
157            .collect::<String>()
158            .to_lowercase()
159    }
160
161    #[test]
162    fn inlines_single_call() {
163        let funcs = reg(&[("tax", &["x"], "x * 1.1")]);
164        let out = expand_scalar_sql_functions("SELECT tax(amount) FROM sales", &funcs).unwrap();
165        assert_eq!(norm(&out), norm("SELECT (amount * 1.1) FROM sales"));
166    }
167
168    #[test]
169    fn inlines_multiple_params_and_calls() {
170        let funcs = reg(&[("disc", &["p", "d"], "p * (1 - d)")]);
171        let out = expand_scalar_sql_functions(
172            "SELECT disc(price, rate), disc(unit_price, 0.2) FROM t WHERE disc(price, rate) > 10",
173            &funcs,
174        )
175        .unwrap();
176        assert!(norm(&out).contains(&norm("(price * (1 - rate))")));
177        assert!(norm(&out).contains(&norm("(unit_price * (1 - 0.2))")));
178    }
179
180    #[test]
181    fn inlines_nested_functions() {
182        let funcs = reg(&[("a", &["x"], "x + 1"), ("b", &["y"], "a(y) * 2")]);
183        let out = expand_scalar_sql_functions("SELECT b(v) FROM t", &funcs).unwrap();
184        // b(v) -> (a(v) * 2) -> ((v + 1) * 2)
185        assert!(norm(&out).contains(&norm("((v + 1) * 2)")));
186    }
187
188    #[test]
189    fn leaves_unrelated_and_builtin_calls_untouched() {
190        let funcs = reg(&[("tax", &["x"], "x * 1.1")]);
191        let out =
192            expand_scalar_sql_functions("SELECT SUM(amount), UPPER(region) FROM sales", &funcs)
193                .unwrap();
194        assert_eq!(
195            norm(&out),
196            norm("SELECT SUM(amount), UPPER(region) FROM sales")
197        );
198    }
199
200    #[test]
201    fn empty_registry_is_noop() {
202        let funcs = HashMap::new();
203        let sql = "SELECT tax(amount) FROM sales";
204        assert_eq!(expand_scalar_sql_functions(sql, &funcs).unwrap(), sql);
205    }
206}
krishiv_sql/scalar_udf.rs

krishiv_sql/
scalar_udf.rs
