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            let size_budget =
132                WorkBudget::from_graph_size(graph.vertex_count() as u64, graph.edge_count() as u64)
133                    .total();
134            let total = work_cap.map_or(size_budget, |w| w.min(size_budget));
135            let mut session = AlgoSession::new(
136                next_session_epoch(),
137                WorkBudget::new(total.max(1)),
138                Arena::new(arena_bytes, DEFAULT_ARENA_MAX_HANDLES),
139            )
140            .with_expected_columns(expected_cols);
141            let g = to_i64(session.bind_graph(Arc::clone(&graph)));
142            let sid = registry.open(session);
143
144            // Everything that can fail between open and close is wrapped so the
145            // session is ALWAYS closed (no leak of the projected graph into the
146            // process-global registry), even if input-build or `acquire` fails.
147            let started = std::time::Instant::now();
148            let call_result: Result<(), DataFusionError> = (|| {
149                let input = serde_json::to_vec(&serde_json::json!({
150                    "session": sid, "graph": g, "args": json_args,
151                }))
152                .map_err(|e| DataFusionError::Execution(format!("extism algorithm input: {e}")))?;
153                let mut leased = acquire(&pool)
154                    .map_err(|e| DataFusionError::Execution(format!("extism acquire: {e}")))?;
155                leased
156                    .get_mut()
157                    .call::<&[u8], &[u8]>(&invoke_export, &input)
158                    .map(|_| ())
159                    .map_err(|e| {
160                        DataFusionError::Execution(format!("extism call `{invoke_export}`: {e}"))
161                    })
162            })();
163
164            // Always close the session (freeing handles) even on guest error.
165            let closed = registry.close(sid);
166            if let Err(orig) = call_result {
167                // Classify from the closed session: a drained budget is a typed
168                // Exhausted outcome (§5.2). Extism sets no host wall-clock
169                // deadline, so a Timeout is never inferred. Other faults verbatim.
170                let (spent, budget) = closed
171                    .as_ref()
172                    .map_or((0, 0), |s| (s.work_spent(), s.work_budget()));
173                return Err(
174                    uni_plugin_builtin::algorithms::graph_compute::error::incomplete_tag_after_guest(
175                        &qname_str,
176                        false,
177                        spent,
178                        budget,
179                        started.elapsed().as_millis() as u64,
180                    )
181                    .map_or(orig, DataFusionError::Execution),
182                );
183            }
184            let mut closed =
185                closed.ok_or_else(|| DataFusionError::Execution("session vanished".into()))?;
186            let emitted = closed.take_emitted();
187
188            build_batch(&schema_for_batch, &graph, &emitted)
189                .map_err(|e| DataFusionError::Execution(format!("extism algorithm emit: {e}")))
190        });
191
192        Ok(Box::pin(RecordBatchStreamAdapter::new(out_schema, stream)))
193    }
194}
195
196/// Assembles the output batch from emitted columns + a `nodeId` column.
197fn build_batch(
198    schema: &SchemaRef,
199    graph: &uni_algo::algo::GraphProjection,
200    emitted: &[(String, Vec<f64>)],
201) -> Result<RecordBatch, FnError> {
202    let n = graph.vertex_count();
203    let mut columns: Vec<ArrayRef> = Vec::with_capacity(schema.fields().len());
204    for field in schema.fields() {
205        if field.name() == "nodeId" {
206            #[expect(
207                clippy::cast_possible_wrap,
208                reason = "vids fit i64 in practice; Cypher integers are i64"
209            )]
210            let ids: Vec<i64> = (0..n as u32)
211                .map(|slot| graph.to_vid(slot).as_u64() as i64)
212                .collect();
213            columns.push(Arc::new(Int64Array::from(ids)));
214            continue;
215        }
216        let (_, values) = emitted
217            .iter()
218            .find(|(name, _)| name == field.name())
219            .ok_or_else(|| {
220                FnError::new(
221                    0x869,
222                    format!("guest did not emit declared column `{}`", field.name()),
223                )
224            })?;
225        match field.data_type() {
226            DataType::Float64 => columns.push(Arc::new(Float64Array::from(values.clone()))),
227            #[expect(
228                clippy::cast_possible_truncation,
229                reason = "int columns hold whole f64 values"
230            )]
231            DataType::Int64 => {
232                let ints: Vec<i64> = values.iter().map(|&v| v as i64).collect();
233                columns.push(Arc::new(Int64Array::from(ints)));
234            }
235            other => {
236                return Err(FnError::new(
237                    0x862,
238                    format!("unsupported emit column type {other:?}"),
239                ));
240            }
241        }
242    }
243    RecordBatch::try_new(Arc::clone(schema), columns)
244        .map_err(|e| FnError::new(0x15, format!("extism algorithm batch: {e}")))
245}
uni_plugin_extism/adapter_algorithm.rs

uni_plugin_extism/
adapter_algorithm.rs
