uni_plugin_extism/
adapter_algorithm.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright 2024-2026 Dragonscale Team
3
4//! Algorithm adapter — an Extism guest algorithm driving GraphCompute kernels.
5//!
6//! ## Wire contract (per qname `q`)
7//! - `algo_<q>_invoke` — input is a JSON object `{session, graph, args}`. The
8//!   guest drives the coarse kernels via the imported `uni_graph_call` host fn
9//!   (referencing `session` on every call) and publishes its per-vertex result
10//!   with `emit`. The output bytes are ignored; the host reads the emitted
11//!   columns from the session registry (proposal §4.6).
12//!
13//! The host projects the graph (capability-gated on `GraphCompute` + `HostQuery`
14//! via the bridge), opens a session in the shared registry, invokes the guest,
15//! then closes the session and assembles the declared `(nodeId, …)` batch.
16//
17// Rust guideline compliant
18
19use std::sync::Arc;
20
21use arrow_array::{ArrayRef, Float64Array, Int64Array, RecordBatch};
22use arrow_schema::{DataType, Schema, SchemaRef};
23use datafusion::error::DataFusionError;
24use datafusion::execution::SendableRecordBatchStream;
25use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
26use uni_plugin::QName;
27use uni_plugin::errors::FnError;
28use uni_plugin::traits::algorithm::{
29    AlgorithmContext, AlgorithmProvider, AlgorithmSignature, GraphProjectionSpec,
30};
31use uni_plugin_builtin::algorithms::bridge::AlgorithmHostBridge;
32use uni_plugin_builtin::algorithms::graph_compute::handle::Handle;
33use uni_plugin_builtin::algorithms::graph_compute::{
34    AlgoSession, Arena, DEFAULT_ARENA_MAX_HANDLES, SharedRegistry, WorkBudget, next_session_epoch,
35};
36
37use crate::adapter_common::{acquire, sanitize_qname};
38use crate::pool::ExtismInstancePool;
39
40/// Plugin-side algorithm-invoke export name from a qname.
41#[must_use]
42pub(crate) fn algo_invoke_export_name(qname: &QName) -> String {
43    format!("algo_{}_invoke", sanitize_qname(qname))
44}
45
46/// `AlgorithmProvider` adapter wrapping an Extism plugin pool + session registry.
47pub struct ExtismAlgorithm {
48    pool: Arc<ExtismInstancePool<extism::Plugin>>,
49    registry: SharedRegistry,
50    qname: QName,
51    invoke_export: String,
52    signature: AlgorithmSignature,
53}
54
55impl std::fmt::Debug for ExtismAlgorithm {
56    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
57        f.debug_struct("ExtismAlgorithm")
58            .field("qname", &self.qname)
59            .finish_non_exhaustive()
60    }
61}
62
63impl ExtismAlgorithm {
64    /// Constructs an algorithm adapter against `pool` and the shared `registry`.
65    #[must_use]
66    pub fn new(
67        pool: Arc<ExtismInstancePool<extism::Plugin>>,
68        registry: SharedRegistry,
69        qname: QName,
70        signature: AlgorithmSignature,
71    ) -> Self {
72        let invoke_export = algo_invoke_export_name(&qname);
73        Self {
74            pool,
75            registry,
76            qname,
77            invoke_export,
78            signature,
79        }
80    }
81}
82
83fn to_i64(h: Handle) -> i64 {
84    #[expect(
85        clippy::cast_possible_wrap,
86        reason = "opaque handle round-trips bit-exact"
87    )]
88    let v = h.as_u64() as i64;
89    v
90}
91
92impl AlgorithmProvider for ExtismAlgorithm {
93    fn signature(&self) -> &AlgorithmSignature {
94        &self.signature
95    }
96
97    fn run(&self, ctx: AlgorithmContext<'_>) -> Result<SendableRecordBatchStream, FnError> {
98        let host = ctx
99            .host
100            .ok_or_else(|| FnError::new(0x800, "extism algorithm: host unbound"))?;
101        let bridge = host
102            .as_any()
103            .downcast_ref::<AlgorithmHostBridge>()
104            .ok_or_else(|| FnError::new(0x801, "extism algorithm: host is not the bridge"))?;
105
106        let json_args: Vec<serde_json::Value> = serde_json::from_str(ctx.config_json)
107            .map_err(|e| FnError::new(0x802, format!("extism algorithm: bad config json: {e}")))?;
108
109        let spec = GraphProjectionSpec {
110            include_reverse: true, // enable In-direction kernels (WCC/k-core/HITS)
111            ..GraphProjectionSpec::default()
112        };
113        let projection = bridge.project_for_graph_compute(&spec);
114        let (work_cap, arena_bytes) = bridge.graph_compute_caps();
115
116        let out_schema: SchemaRef = Arc::new(Schema::new(self.signature.output_fields.clone()));
117        let schema_for_batch = Arc::clone(&out_schema);
118        let registry = Arc::clone(&self.registry);
119        let pool = Arc::clone(&self.pool);
120        let invoke_export = self.invoke_export.clone();
121        let qname_str = self.qname.to_string();
122        let expected_cols = uni_plugin_builtin::algorithms::graph_compute::guest_emit_columns(
123            &self.signature.output_fields,
124        );
125
126        let stream = futures::stream::once(async move {
127            let graph = projection
128                .await
129                .map_err(|e| DataFusionError::Execution(format!("extism algorithm: {e}")))?;
130
131            // An explicit grant is authoritative and may raise the ceiling;
132            // otherwise the size-derived default (proposal §9).
133            let budget = WorkBudget::resolve(
134                work_cap,
135                graph.vertex_count() as u64,
136                graph.edge_count() as u64,
137            );
138            let mut session = AlgoSession::new(
139                next_session_epoch(),
140                budget,
141                Arena::new(arena_bytes, DEFAULT_ARENA_MAX_HANDLES),
142            )
143            .with_expected_columns(expected_cols);
144            let g = to_i64(session.bind_graph(Arc::clone(&graph)));
145            let sid = registry.open(session);
146
147            // Everything that can fail between open and close is wrapped so the
148            // session is ALWAYS closed (no leak of the projected graph into the
149            // process-global registry), even if input-build or `acquire` fails.
150            let started = std::time::Instant::now();
151            let call_result: Result<(), DataFusionError> = (|| {
152                let input = serde_json::to_vec(&serde_json::json!({
153                    "session": sid, "graph": g, "args": json_args,
154                }))
155                .map_err(|e| DataFusionError::Execution(format!("extism algorithm input: {e}")))?;
156                let mut leased = acquire(&pool)
157                    .map_err(|e| DataFusionError::Execution(format!("extism acquire: {e}")))?;
158                leased
159                    .get_mut()
160                    .call::<&[u8], &[u8]>(&invoke_export, &input)
161                    .map(|_| ())
162                    .map_err(|e| {
163                        DataFusionError::Execution(format!("extism call `{invoke_export}`: {e}"))
164                    })
165            })();
166
167            // Always close the session (freeing handles) even on guest error.
168            let closed = registry.close(sid);
169            if let Err(orig) = call_result {
170                // Classify from the closed session: a drained budget is a typed
171                // Exhausted outcome (§5.2). Extism sets no host wall-clock
172                // deadline, so a Timeout is never inferred. Other faults verbatim.
173                let (spent, budget) = closed
174                    .as_ref()
175                    .map_or((0, 0), |s| (s.work_spent(), s.work_budget()));
176                return Err(
177                    uni_plugin_builtin::algorithms::graph_compute::error::incomplete_tag_after_guest(
178                        &qname_str,
179                        false,
180                        spent,
181                        budget,
182                        started.elapsed().as_millis() as u64,
183                    )
184                    .map_or(orig, DataFusionError::Execution),
185                );
186            }
187            let mut closed =
188                closed.ok_or_else(|| DataFusionError::Execution("session vanished".into()))?;
189            let emitted = closed.take_emitted();
190
191            build_batch(&schema_for_batch, &graph, &emitted)
192                .map_err(|e| DataFusionError::Execution(format!("extism algorithm emit: {e}")))
193        });
194
195        Ok(Box::pin(RecordBatchStreamAdapter::new(out_schema, stream)))
196    }
197}
198
199/// Assembles the output batch from emitted columns + a `nodeId` column.
200fn build_batch(
201    schema: &SchemaRef,
202    graph: &uni_algo::algo::GraphProjection,
203    emitted: &[(String, Vec<f64>)],
204) -> Result<RecordBatch, FnError> {
205    let n = graph.vertex_count();
206    let mut columns: Vec<ArrayRef> = Vec::with_capacity(schema.fields().len());
207    for field in schema.fields() {
208        if field.name() == "nodeId" {
209            #[expect(
210                clippy::cast_possible_wrap,
211                reason = "vids fit i64 in practice; Cypher integers are i64"
212            )]
213            let ids: Vec<i64> = (0..n as u32)
214                .map(|slot| graph.to_vid(slot).as_u64() as i64)
215                .collect();
216            columns.push(Arc::new(Int64Array::from(ids)));
217            continue;
218        }
219        let (_, values) = emitted
220            .iter()
221            .find(|(name, _)| name == field.name())
222            .ok_or_else(|| {
223                FnError::new(
224                    0x869,
225                    format!("guest did not emit declared column `{}`", field.name()),
226                )
227            })?;
228        match field.data_type() {
229            DataType::Float64 => columns.push(Arc::new(Float64Array::from(values.clone()))),
230            #[expect(
231                clippy::cast_possible_truncation,
232                reason = "int columns hold whole f64 values"
233            )]
234            DataType::Int64 => {
235                let ints: Vec<i64> = values.iter().map(|&v| v as i64).collect();
236                columns.push(Arc::new(Int64Array::from(ints)));
237            }
238            other => {
239                return Err(FnError::new(
240                    0x862,
241                    format!("unsupported emit column type {other:?}"),
242                ));
243            }
244        }
245    }
246    RecordBatch::try_new(Arc::clone(schema), columns)
247        .map_err(|e| FnError::new(0x15, format!("extism algorithm batch: {e}")))
248}
uni_plugin_extism/adapter_algorithm.rs

uni_plugin_extism/
adapter_algorithm.rs
