laminar_sql/datafusion/
ai_udf.rs

1//! Marker UDFs for the `ai_*` SQL functions.
2//!
3//! These exist so that `ai_classify`, `ai_embed`, … resolve as known functions
4//! with a fixed return type, the same way `tumble()` and `watermark()` are
5//! markers the engine rewrites rather than evaluates. In the normal path the AI
6//! detector lifts an `ai_*` call out of the projection and replaces it with a
7//! computed column before DataFusion plans the residual query, so these markers
8//! are never invoked. If one survives to execution — used in a position the
9//! detector does not rewrite — `invoke` returns a clear error rather than
10//! producing a wrong value.
11//!
12//! The function name → return type map here must stay in step with the
13//! name → task map in `laminar-db`'s `sql_analysis` (the eight locked AI
14//! functions). The crate does not depend on `laminar-ai`, so the two lists are
15//! intentionally independent.
16
17use std::hash::{Hash, Hasher};
18use std::sync::Arc;
19
20use arrow::datatypes::{DataType, Field};
21use datafusion_common::{exec_err, Result};
22use datafusion_expr::{
23    ColumnarValue, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, Volatility,
24};
25
26/// A placeholder scalar UDF for an `ai_*` function. Holds the function's name
27/// and fixed return type; errors if evaluated directly.
28#[derive(Debug)]
29pub struct AiFunctionMarker {
30    name: &'static str,
31    signature: Signature,
32    return_type: DataType,
33}
34
35impl AiFunctionMarker {
36    /// Create a marker for `name` returning `return_type`.
37    #[must_use]
38    pub fn new(name: &'static str, return_type: DataType) -> Self {
39        Self {
40            name,
41            // Accept any arguments (input plus `model =>` / `labels =>` named
42            // args). Volatile so it is never const-folded away before the
43            // detector can see it.
44            signature: Signature::variadic_any(Volatility::Volatile),
45            return_type,
46        }
47    }
48}
49
50impl PartialEq for AiFunctionMarker {
51    fn eq(&self, other: &Self) -> bool {
52        self.name == other.name
53    }
54}
55
56impl Eq for AiFunctionMarker {}
57
58impl Hash for AiFunctionMarker {
59    fn hash<H: Hasher>(&self, state: &mut H) {
60        self.name.hash(state);
61    }
62}
63
64impl ScalarUDFImpl for AiFunctionMarker {
65    fn as_any(&self) -> &dyn std::any::Any {
66        self
67    }
68
69    fn name(&self) -> &'static str {
70        self.name
71    }
72
73    fn signature(&self) -> &Signature {
74        &self.signature
75    }
76
77    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
78        Ok(self.return_type.clone())
79    }
80
81    fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> Result<ColumnarValue> {
82        exec_err!(
83            "{} is an AI function and must be the top-level expression of a SELECT \
84             projection over a stream; it cannot be evaluated in this position",
85            self.name
86        )
87    }
88}
89
90/// Build the eight `ai_*` marker UDFs, ready to register on a session context.
91///
92/// Text-generating and discriminative-label tasks return `Utf8`; `ai_sentiment`
93/// returns a `Float64` score in `[-1, 1]`; `ai_embed` returns a `List<Float32>`
94/// embedding.
95#[must_use]
96pub fn ai_function_markers() -> Vec<ScalarUDF> {
97    let embedding = DataType::List(Arc::new(Field::new("item", DataType::Float32, true)));
98    let specs: [(&'static str, DataType); 8] = [
99        ("ai_classify", DataType::Utf8),
100        ("ai_sentiment", DataType::Float64),
101        ("ai_embed", embedding),
102        ("ai_extract", DataType::Utf8),
103        ("ai_complete", DataType::Utf8),
104        ("ai_summarize", DataType::Utf8),
105        ("ai_translate", DataType::Utf8),
106        ("ai_gen", DataType::Utf8),
107    ];
108    specs
109        .into_iter()
110        .map(|(name, return_type)| {
111            ScalarUDF::new_from_impl(AiFunctionMarker::new(name, return_type))
112        })
113        .collect()
114}
115
116#[cfg(test)]
117mod tests {
118    use super::*;
119    use datafusion_common::config::ConfigOptions;
120
121    #[test]
122    fn registers_eight_markers_with_expected_return_types() {
123        let markers = ai_function_markers();
124        assert_eq!(markers.len(), 8);
125
126        let classify = markers.iter().find(|m| m.name() == "ai_classify").unwrap();
127        assert_eq!(classify.return_type(&[]).unwrap(), DataType::Utf8);
128
129        let sentiment = markers.iter().find(|m| m.name() == "ai_sentiment").unwrap();
130        assert_eq!(sentiment.return_type(&[]).unwrap(), DataType::Float64);
131
132        let embed = markers.iter().find(|m| m.name() == "ai_embed").unwrap();
133        assert!(matches!(embed.return_type(&[]).unwrap(), DataType::List(_)));
134    }
135
136    #[test]
137    fn invoking_a_marker_is_an_error() {
138        let marker = AiFunctionMarker::new("ai_classify", DataType::Utf8);
139        let args = ScalarFunctionArgs {
140            args: vec![],
141            arg_fields: vec![],
142            number_rows: 1,
143            return_field: Arc::new(Field::new("out", DataType::Utf8, true)),
144            config_options: Arc::new(ConfigOptions::default()),
145        };
146        assert!(marker.invoke_with_args(args).is_err());
147    }
148}
laminar_sql/datafusion/ai_udf.rs

laminar_sql/datafusion/
ai_udf.rs
