linkrs_query/executor/expression/functions/
custom.rs

1//! Custom (non-built-in) function definitions.
2//!
3//! Covers Rust closures, C ABI callbacks and dynamically loaded UDF plugins.
4
5use crate::executor::expression::{ExpressionError, ExpressionErrorType};
6use linkrs_core::value_conversion::core_value_to_linkrs;
7use linkrs_core::Value;
8
9use std::ffi::c_void;
10
11use super::udf::{execute_isolated, SharedPlugin};
12use super::ExpressionFunction;
13
14/// C Function Context Structure (Opaque Pointers)
15pub struct CFunctionContext {
16    /// Result value
17    pub result: Option<Value>,
18    /// Error message
19    pub error: Option<String>,
20    /// Aggregation status (used for aggregate functions)
21    pub aggregate_state: Option<Box<dyn std::any::Any + Send>>,
22    /// User data pointer
23    pub user_data: usize,
24    /// Number of parameters
25    pub argc: usize,
26    /// Parameter array (converted to C API format)
27    pub argv: Vec<linkrs_core::types::c_api::linkrs_value_t>,
28}
29
30impl Default for CFunctionContext {
31    fn default() -> Self {
32        Self::new()
33    }
34}
35
36impl CFunctionContext {
37    pub fn new() -> Self {
38        Self {
39            result: None,
40            error: None,
41            aggregate_state: None,
42            user_data: 0,
43            argc: 0,
44            argv: Vec::new(),
45        }
46    }
47
48    pub fn with_user_data(user_data: usize) -> Self {
49        Self {
50            result: None,
51            error: None,
52            aggregate_state: None,
53            user_data,
54            argc: 0,
55            argv: Vec::new(),
56        }
57    }
58
59    pub fn set_result(&mut self, value: Value) {
60        self.result = Some(value);
61    }
62
63    pub fn set_error(&mut self, error: String) {
64        self.error = Some(error);
65    }
66
67    /// Set the aggregation status
68    pub fn set_aggregate_state<T: std::any::Any + Send + 'static>(&mut self, state: T) {
69        self.aggregate_state = Some(Box::new(state));
70    }
71
72    /// Obtain the aggregated status.
73    pub fn get_aggregate_state<T: std::any::Any + Send + 'static>(&self) -> Option<&T> {
74        self.aggregate_state.as_ref()?.downcast_ref::<T>()
75    }
76
77    /// Obtain a variable reference to the aggregated status
78    pub fn get_aggregate_state_mut<T: std::any::Any + Send + 'static>(&mut self) -> Option<&mut T> {
79        self.aggregate_state.as_mut()?.downcast_mut::<T>()
80    }
81}
82
83/// Scalar function callback type
84pub type ScalarFunctionCallback =
85    extern "C" fn(*mut CFunctionContext, i32, *mut linkrs_core::types::c_api::linkrs_value_t);
86
87/// Aggregation step callback type
88pub type AggregateStepCallback =
89    extern "C" fn(*mut CFunctionContext, i32, *mut linkrs_core::types::c_api::linkrs_value_t);
90
91/// Aggregate final callback type
92pub type AggregateFinalCallback = extern "C" fn(*mut CFunctionContext);
93
94/// Implementation of custom functions and their types
95#[derive(Clone)]
96pub enum CustomFunctionImpl {
97    /// Custom functions implemented in Rust
98    Rust(fn(&[Value]) -> Result<Value, ExpressionError>),
99    /// A scalar function implemented using a C callback
100    C {
101        /// Scalar function callback (stores the address of the function pointer)
102        scalar_callback: usize,
103        /// User data (storage pointer addresses)
104        user_data: usize,
105    },
106    /// Aggregate functions implemented using C callbacks
107    Aggregate {
108        /// Aggregation step callback (pointer address to the storage function)
109        step_callback: usize,
110        /// Aggregated final callback (address of the stored function pointer)
111        final_callback: usize,
112        /// User data (storage pointer address)
113        user_data: usize,
114    },
115    /// A function backed by a dynamically loaded UDF plugin.
116    Dynamic {
117        /// Shared plugin handle; keeps the dynamic library alive.
118        plugin: SharedPlugin,
119    },
120}
121
122impl std::fmt::Debug for CustomFunctionImpl {
123    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
124        match self {
125            CustomFunctionImpl::Rust(_) => write!(f, "Rust closure"),
126            CustomFunctionImpl::C { .. } => write!(f, "C scalar callback"),
127            CustomFunctionImpl::Aggregate { .. } => write!(f, "C aggregate callback"),
128            CustomFunctionImpl::Dynamic { plugin } => {
129                write!(f, "Dynamic UDF '{}'", plugin.name())
130            }
131        }
132    }
133}
134
135/// Custom function definition
136#[derive(Debug, Clone)]
137pub struct CustomFunction {
138    /// Function name
139    pub name: String,
140    /// Number of parameters
141    pub arity: usize,
142    /// "Do you accept variable parameters?"
143    pub is_variadic: bool,
144    /// Function description
145    pub description: String,
146    /// Function implementation
147    pub implementation: CustomFunctionImpl,
148}
149
150impl CustomFunction {
151    /// Create a new custom Rust function.
152    pub fn new_rust(
153        name: impl Into<String>,
154        arity: usize,
155        is_variadic: bool,
156        description: impl Into<String>,
157        implementation: fn(&[Value]) -> Result<Value, ExpressionError>,
158    ) -> Self {
159        Self {
160            name: name.into(),
161            arity,
162            is_variadic,
163            description: description.into(),
164            implementation: CustomFunctionImpl::Rust(implementation),
165        }
166    }
167
168    /// Create a new custom C callback function.
169    pub fn new_c(
170        name: impl Into<String>,
171        arity: usize,
172        is_variadic: bool,
173        description: impl Into<String>,
174        scalar_callback: ScalarFunctionCallback,
175        user_data: *mut c_void,
176    ) -> Self {
177        Self {
178            name: name.into(),
179            arity,
180            is_variadic,
181            description: description.into(),
182            implementation: CustomFunctionImpl::C {
183                scalar_callback: scalar_callback as usize,
184                user_data: user_data as usize,
185            },
186        }
187    }
188
189    /// Create a new C callback aggregation function
190    pub fn new_c_aggregate(
191        name: impl Into<String>,
192        arity: usize,
193        is_variadic: bool,
194        description: impl Into<String>,
195        step_callback: AggregateStepCallback,
196        final_callback: AggregateFinalCallback,
197        user_data: *mut c_void,
198    ) -> Self {
199        Self {
200            name: name.into(),
201            arity,
202            is_variadic,
203            description: description.into(),
204            implementation: CustomFunctionImpl::Aggregate {
205                step_callback: step_callback as usize,
206                final_callback: final_callback as usize,
207                user_data: user_data as usize,
208            },
209        }
210    }
211
212    /// Create a custom function backed by a dynamically loaded UDF plugin.
213    ///
214    /// Arity metadata is read from the plugin so registration stays in sync
215    /// with the library implementation.
216    pub fn new_dynamic(plugin: SharedPlugin) -> Self {
217        let (min_arity, max_arity) = (plugin.min_arity(), plugin.max_arity());
218        Self {
219            name: plugin.name().to_string(),
220            arity: min_arity,
221            is_variadic: min_arity != max_arity,
222            description: plugin.description().to_string(),
223            implementation: CustomFunctionImpl::Dynamic { plugin },
224        }
225    }
226
227    /// Check whether the function is backed by a dynamic UDF plugin.
228    pub fn is_dynamic(&self) -> bool {
229        matches!(self.implementation, CustomFunctionImpl::Dynamic { .. })
230    }
231
232    /// Check whether it is an aggregate function.
233    pub fn is_aggregate(&self) -> bool {
234        matches!(self.implementation, CustomFunctionImpl::Aggregate { .. })
235    }
236
237    /// Execute the function
238    pub fn execute(&self, args: &[Value]) -> Result<Value, ExpressionError> {
239        match &self.implementation {
240            CustomFunctionImpl::Rust(func) => func(args),
241            CustomFunctionImpl::C {
242                scalar_callback,
243                user_data: _,
244            } => {
245                // Creating a C function context
246                let mut ctx = CFunctionContext::new();
247                ctx.argc = args.len();
248                ctx.argv = args.iter().map(core_value_to_linkrs).collect();
249                let ctx_ptr = &mut ctx as *mut CFunctionContext;
250
251                // Convert a `usize` value back to a function pointer
252                let callback: ScalarFunctionCallback =
253                    unsafe { std::mem::transmute(*scalar_callback) };
254
255                // Calling a C callback
256                let argv_ptr = if ctx.argv.is_empty() {
257                    std::ptr::null_mut()
258                } else {
259                    ctx.argv.as_mut_ptr()
260                };
261                callback(ctx_ptr, args.len() as i32, argv_ptr);
262
263                // Check for errors.
264                if let Some(error) = ctx.error {
265                    return Err(ExpressionError::new(
266                        ExpressionErrorType::FunctionExecutionError,
267                        error,
268                    ));
269                }
270
271                // Return the value set by the callback.
272                ctx.result.ok_or_else(|| {
273                    ExpressionError::new(
274                        ExpressionErrorType::FunctionExecutionError,
275                        format!("The function '{}' does not set a return value", self.name),
276                    )
277                })
278            }
279            CustomFunctionImpl::Aggregate { .. } => Err(ExpressionError::new(
280                ExpressionErrorType::InvalidOperation,
281                "Aggregation functions need to be executed within the aggregation context"
282                    .to_string(),
283            )),
284            CustomFunctionImpl::Dynamic { plugin } => execute_isolated(plugin.as_ref(), args),
285        }
286    }
287}
288
289impl ExpressionFunction for CustomFunction {
290    fn name(&self) -> &str {
291        &self.name
292    }
293
294    fn arity(&self) -> usize {
295        self.arity
296    }
297
298    fn is_variadic(&self) -> bool {
299        self.is_variadic
300    }
301
302    fn execute(&self, args: &[Value]) -> Result<Value, ExpressionError> {
303        self.execute(args)
304    }
305
306    fn description(&self) -> &str {
307        &self.description
308    }
309}
linkrs_query/executor/expression/functions/custom.rs

linkrs_query/executor/expression/functions/
custom.rs
