ddx_datafusion/
markers.rs

1// SPDX-FileCopyrightText: 2026 Alexander Merose <al@merose.com> & ddx Authors
2//
3// SPDX-License-Identifier: Apache-2.0
4
5//! The `grad`/`jvp` marker UDFs.
6//!
7//! These are **not** row functions. A scalar UDF only ever receives evaluated
8//! *values*, never the symbolic expression of its argument, but differentiation
9//! is a function of the symbolic form — so `grad` cannot be computed at
10//! runtime. (Empirically pinned on a live engine in
11//! `docs/spikes/datafusion_python_analyzer_rule_r2.py`, T4: a `grad` UDF given
12//! `grad(x*x, x)` over `x = [1,2,3]` receives `[1.0, 4.0, 9.0]`.)
13//!
14//! Registration exists for exactly one reason: to make the marker call *parse
15//! and plan*, so [`crate::analyzer::DdxAnalyzer`] can find it in the
16//! `LogicalPlan` and rewrite it away. Reaching execution is therefore always a
17//! bug, and these deliberately error there rather than returning a number.
18
19use datafusion::arrow::datatypes::DataType;
20use datafusion::error::{DataFusionError, Result};
21use datafusion::logical_expr::{
22    ColumnarValue, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, Volatility,
23};
24
25/// The name of the gradient marker, as written in SQL.
26pub const GRAD: &str = "grad";
27/// The name of the forward-mode (directional derivative) marker.
28pub const JVP: &str = "jvp";
29
30/// A marker UDF: parses and plans, never executes.
31#[derive(Debug, PartialEq, Eq, Hash)]
32struct Marker {
33    name: &'static str,
34    signature: Signature,
35}
36
37impl Marker {
38    fn new(name: &'static str, arg_count: usize) -> Self {
39        Marker {
40            name,
41            // `Signature::any` accepts the arguments at whatever types they
42            // arrive in. That tolerance is required here: `add_analyzer_rule`
43            // installs the rule to run AFTER `TypeCoercion`, so a
44            // stricter signature would make the planner inject casts into the
45            // marker's argument before ddx ever sees it. (ddx-core does have a
46            // `Cast` rule, so an injected cast is survivable — but not
47            // provoking one keeps the differentiated expression closer to what
48            // the user actually wrote.)
49            signature: Signature::any(arg_count, Volatility::Immutable),
50        }
51    }
52}
53
54impl ScalarUDFImpl for Marker {
55    fn name(&self) -> &str {
56        self.name
57    }
58
59    fn signature(&self) -> &Signature {
60        &self.signature
61    }
62
63    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
64        // Derivatives are always emitted DOUBLE-typed: differentiation runs
65        // pre-binding, so operand types are unknown, and SQL integer division
66        // truncates on some engines but not others.
67        Ok(DataType::Float64)
68    }
69
70    fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> Result<ColumnarValue> {
71        Err(DataFusionError::Execution(format!(
72            "ddx: `{name}` reached execution, which never happens in a correct \
73             rewrite — it is a compile-time marker, not a row function.\n\n\
74             The `{name}()` call was not rewritten away before planning finished. \
75             Either the ddx analyzer rule is not installed on this SessionContext \
76             (use `ddx_datafusion::install(&ctx)`), or the marker sits somewhere \
77             the rule does not reach — in which case rewrite the SQL text instead \
78             with `ddx_datafusion::ddx_sql(&ctx, sql)`.",
79            name = self.name
80        )))
81    }
82}
83
84/// The `grad(expr, column)` marker: `d(expr)/d(column)`.
85pub fn grad_udf() -> ScalarUDF {
86    ScalarUDF::new_from_impl(Marker::new(GRAD, 2))
87}
88
89/// The `jvp(expr, column, tangent)` marker: `d(expr)/d(column) · tangent`.
90pub fn jvp_udf() -> ScalarUDF {
91    ScalarUDF::new_from_impl(Marker::new(JVP, 3))
92}
93
94/// Is `name` one of ddx's marker functions? Case-folded, because SQL function
95/// names are case-insensitive and `GRAD(x, x)` must be caught too.
96pub(crate) fn marker_kind(name: &str) -> Option<&'static str> {
97    if name.eq_ignore_ascii_case(GRAD) {
98        Some(GRAD)
99    } else if name.eq_ignore_ascii_case(JVP) {
100        Some(JVP)
101    } else {
102        None
103    }
104}
105
106#[cfg(test)]
107mod tests {
108    use super::*;
109
110    #[test]
111    fn markers_are_named_and_arity_checked() {
112        assert_eq!(grad_udf().name(), "grad");
113        assert_eq!(jvp_udf().name(), "jvp");
114    }
115
116    #[test]
117    fn marker_names_are_matched_case_insensitively() {
118        assert_eq!(marker_kind("grad"), Some(GRAD));
119        assert_eq!(marker_kind("GRAD"), Some(GRAD));
120        assert_eq!(marker_kind("Jvp"), Some(JVP));
121        assert_eq!(marker_kind("gradient"), None);
122        assert_eq!(marker_kind("mygrad"), None);
123    }
124
125    #[test]
126    fn executing_a_marker_is_a_loud_error_not_a_number() {
127        // Markers deliberately error if one reaches execution, rather than
128        // silently producing a value.
129        let udf = grad_udf();
130        let err = udf
131            .invoke_with_args(ScalarFunctionArgs {
132                args: vec![],
133                arg_fields: vec![],
134                number_rows: 1,
135                return_field: std::sync::Arc::new(datafusion::arrow::datatypes::Field::new(
136                    "d",
137                    DataType::Float64,
138                    true,
139                )),
140                config_options: std::sync::Arc::new(datafusion::config::ConfigOptions::default()),
141            })
142            .expect_err("a marker must never execute successfully");
143        let msg = err.to_string();
144        assert!(msg.contains("reached execution"), "unexpected: {msg}");
145        // The error has to tell the user what to actually do about it.
146        assert!(msg.contains("install"), "no remedy in message: {msg}");
147        assert!(msg.contains("ddx_sql"), "no fallback in message: {msg}");
148    }
149}
ddx_datafusion/markers.rs

ddx_datafusion/
markers.rs
