uni_plugin_extism/
adapter.rs1use std::sync::Arc;
20
21use arrow::array::RecordBatch;
22use arrow_schema::{Field, Schema, SchemaRef};
23use datafusion::logical_expr::ColumnarValue;
24use uni_plugin::QName;
25use uni_plugin::errors::FnError;
26use uni_plugin::traits::scalar::{FnSignature, ScalarPluginFn};
27
28use crate::adapter_common::{acquire, extism_err_to_fn_err, sanitize_qname};
29use crate::ipc::{decode_batch, encode_batch};
30use crate::pool::ExtismInstancePool;
31
32pub(crate) fn scalar_export_name(qname: &QName) -> String {
41 format!("invoke_{}", sanitize_qname(qname))
42}
43
44pub struct ExtismScalarFn {
52 pool: Arc<ExtismInstancePool<extism::Plugin>>,
53 qname: QName,
54 export_name: String,
55 sig: FnSignature,
56}
57
58impl std::fmt::Debug for ExtismScalarFn {
59 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
60 f.debug_struct("ExtismScalarFn")
61 .field("qname", &self.qname)
62 .field("export_name", &self.export_name)
63 .field("signature", &self.sig)
64 .finish_non_exhaustive()
65 }
66}
67
68impl ExtismScalarFn {
69 #[must_use]
71 pub fn new(
72 pool: Arc<ExtismInstancePool<extism::Plugin>>,
73 qname: QName,
74 sig: FnSignature,
75 ) -> Self {
76 let export_name = scalar_export_name(&qname);
77 Self {
78 pool,
79 qname,
80 export_name,
81 sig,
82 }
83 }
84
85 fn args_to_batch(&self, args: &[ColumnarValue], rows: usize) -> Result<RecordBatch, FnError> {
90 let arrays: Vec<arrow::array::ArrayRef> = args
91 .iter()
92 .map(|c| {
93 c.clone().into_array(rows).map_err(|e| {
94 FnError::new(
95 FnError::CODE_TYPE_COERCION,
96 format!("ColumnarValue::into_array: {e}"),
97 )
98 })
99 })
100 .collect::<Result<_, _>>()?;
101 let fields: Vec<Field> = arrays
102 .iter()
103 .enumerate()
104 .map(|(i, a)| Field::new(format!("arg{i}"), a.data_type().clone(), true))
105 .collect();
106 let schema: SchemaRef = Arc::new(Schema::new(fields));
107 RecordBatch::try_new(schema, arrays).map_err(|e| {
108 FnError::new(
109 FnError::CODE_TYPE_COERCION,
110 format!("RecordBatch assembly: {e}"),
111 )
112 })
113 }
114}
115
116impl ScalarPluginFn for ExtismScalarFn {
117 fn signature(&self) -> &FnSignature {
118 &self.sig
119 }
120
121 fn invoke(&self, args: &[ColumnarValue], rows: usize) -> Result<ColumnarValue, FnError> {
122 let batch = self.args_to_batch(args, rows)?;
123 let bytes = encode_batch(&batch).map_err(extism_err_to_fn_err)?;
124
125 let mut leased = acquire(&self.pool)?;
126 let out_bytes: Vec<u8> = leased
128 .get_mut()
129 .call::<&[u8], &[u8]>(&self.export_name, bytes.as_slice())
130 .map_err(|e| {
131 FnError::new(
132 FnError::CODE_UNEXPECTED_NULL,
133 format!("extism call `{}` failed: {e}", self.export_name),
134 )
135 })?
136 .to_vec();
137 drop(leased);
138
139 let out_batch = decode_batch(&out_bytes)
140 .map_err(extism_err_to_fn_err)?
141 .ok_or_else(|| {
142 FnError::new(
143 FnError::CODE_UNEXPECTED_NULL,
144 format!("plugin `{}` returned an empty IPC stream", self.export_name),
145 )
146 })?;
147
148 if out_batch.num_columns() != 1 {
149 return Err(FnError::new(
150 FnError::CODE_TYPE_COERCION,
151 format!(
152 "plugin `{}` returned {} columns; scalar fns must return exactly 1",
153 self.export_name,
154 out_batch.num_columns()
155 ),
156 ));
157 }
158 Ok(ColumnarValue::Array(out_batch.column(0).clone()))
159 }
160}
161
162#[cfg(test)]
163mod tests {
164 use super::*;
165
166 #[test]
167 fn scalar_export_name_format() {
168 let q = QName::parse("geo.haversine").expect("valid");
169 assert_eq!(scalar_export_name(&q), "invoke_geo_haversine");
170 }
171}