lora_executor/pull/
mutable.rs

1use std::collections::{BTreeMap, BTreeSet};
2use std::mem::ManuallyDrop;
3use std::sync::Arc;
4
5use lora_compiler::physical::{PhysicalNodeId, PhysicalPlan};
6use lora_compiler::CompiledQuery;
7use lora_store::{GraphStorage, GraphStorageMut};
8
9use crate::errors::{ExecResult, ExecutorError};
10use crate::executor::{GroupValueKey, MutableExecutionContext, MutableExecutor};
11use crate::value::{LoraValue, Row};
12
13use super::traits::write_op_input;
14use super::{build_streaming, subtree_is_fully_streaming, BufferedRowSource, RowSource};
15
16/// Pull-based read-write executor. Wraps the existing
17/// [`MutableExecutor`] under the same row-cursor API. Mutations are
18/// applied during `open_compiled`; the returned cursor yields the
19/// resulting rows lazily.
20pub struct MutablePullExecutor<'a, S: GraphStorageMut> {
21    storage: &'a mut S,
22    params: BTreeMap<String, LoraValue>,
23}
24
25impl<'a, S: GraphStorageMut + GraphStorage> MutablePullExecutor<'a, S> {
26    pub fn new(storage: &'a mut S, params: BTreeMap<String, LoraValue>) -> Self {
27        Self { storage, params }
28    }
29
30    /// Open a cursor for a compiled write query.
31    ///
32    /// Fast path: when a branch root is one of `Create` / `Set` /
33    /// `Delete` / `Remove` / `Merge` and its input subtree is fully
34    /// streamable, returns a [`StreamingWriteCursor`] that pulls input
35    /// row-by-row and applies the per-row write through
36    /// [`MutableExecutor::apply_write_op`]. `UNION ALL` plans stream
37    /// one branch at a time. Plain `UNION` drains branches first so
38    /// rows can be deduplicated by name.
39    ///
40    /// Fallback: a branch that is not streamable materializes through
41    /// [`MutableExecutor::execute_rows`] and wraps the result in a
42    /// [`BufferedRowSource`].
43    pub fn open_compiled(self, compiled: &'a CompiledQuery) -> ExecResult<Box<dyn RowSource + 'a>>
44    where
45        S: 'a,
46    {
47        if compiled.unions.is_empty() {
48            return open_mutable_plan_cursor(self.storage, &compiled.physical, self.params);
49        }
50
51        MutableUnionSource::open(self.storage, compiled, self.params)
52            .map(|source| Box::new(source) as Box<dyn RowSource + 'a>)
53    }
54}
55
56fn open_mutable_plan_cursor<'a, S: GraphStorageMut + GraphStorage + 'a>(
57    storage: &'a mut S,
58    plan: &'a PhysicalPlan,
59    params: BTreeMap<String, LoraValue>,
60) -> ExecResult<Box<dyn RowSource + 'a>> {
61    if let Some(input) = write_op_input(plan, plan.root) {
62        if subtree_is_fully_streaming(plan, input) {
63            return StreamingWriteCursor::open(storage, plan, plan.root, params)
64                .map(|c| Box::new(c) as Box<dyn RowSource + 'a>);
65        }
66    }
67
68    let mut executor = MutableExecutor::new(MutableExecutionContext { storage, params });
69    let rows = executor.execute_rows(plan)?;
70    Ok(Box::new(BufferedRowSource::new(rows)))
71}
72
73#[derive(Clone, Copy)]
74struct StoragePtr<S> {
75    ptr: *mut S,
76}
77
78impl<S> StoragePtr<S> {
79    fn from_mut(storage: &mut S) -> Self {
80        Self {
81            ptr: storage as *mut S,
82        }
83    }
84
85    unsafe fn as_ref<'a>(&self) -> &'a S {
86        unsafe { &*self.ptr }
87    }
88
89    unsafe fn as_mut<'a>(&self) -> &'a mut S {
90        unsafe { &mut *self.ptr }
91    }
92}
93
94/// Mutable UNION cursor. `UNION ALL` streams one branch at a time
95/// against the same staged graph. Plain `UNION` streams branch-by-branch
96/// while retaining only a seen-key set for deduplication.
97pub struct MutableUnionSource<'a, S: GraphStorageMut + GraphStorage + 'a> {
98    storage_ptr: StoragePtr<S>,
99    compiled: &'a CompiledQuery,
100    params: BTreeMap<String, LoraValue>,
101    branch_idx: usize,
102    current: Option<Box<dyn RowSource + 'a>>,
103    needs_dedup: bool,
104    seen: BTreeSet<Vec<(String, GroupValueKey)>>,
105    _phantom: std::marker::PhantomData<&'a mut S>,
106}
107
108impl<'a, S: GraphStorageMut + GraphStorage + 'a> MutableUnionSource<'a, S> {
109    fn open(
110        storage: &'a mut S,
111        compiled: &'a CompiledQuery,
112        params: BTreeMap<String, LoraValue>,
113    ) -> ExecResult<Self> {
114        let needs_dedup = compiled.unions.iter().any(|branch| !branch.all);
115        Ok(Self {
116            storage_ptr: StoragePtr::from_mut(storage),
117            compiled,
118            params,
119            branch_idx: 0,
120            current: None,
121            needs_dedup,
122            seen: BTreeSet::new(),
123            _phantom: std::marker::PhantomData,
124        })
125    }
126
127    fn branch_count(&self) -> usize {
128        self.compiled.unions.len() + 1
129    }
130
131    fn branch_plan(&self, idx: usize) -> &'a PhysicalPlan {
132        if idx == 0 {
133            &self.compiled.physical
134        } else {
135            &self.compiled.unions[idx - 1].physical
136        }
137    }
138
139    fn open_branch(&mut self, idx: usize) -> ExecResult<Box<dyn RowSource + 'a>> {
140        let plan = self.branch_plan(idx);
141        // SAFETY: MutableUnionSource keeps at most one branch cursor
142        // alive at a time. `current` is dropped before advancing to
143        // the next branch, so each mutable reborrow is temporally
144        // disjoint.
145        let storage = unsafe { self.storage_ptr.as_mut() };
146        open_mutable_plan_cursor(storage, plan, self.params.clone())
147    }
148}
149
150impl<'a, S: GraphStorageMut + GraphStorage + 'a> RowSource for MutableUnionSource<'a, S> {
151    fn next_row(&mut self) -> ExecResult<Option<Row>> {
152        loop {
153            if self.branch_idx >= self.branch_count() {
154                return Ok(None);
155            }
156
157            if self.current.is_none() {
158                self.current = Some(self.open_branch(self.branch_idx)?);
159            }
160
161            let Some(current) = self.current.as_mut() else {
162                return Err(ExecutorError::RuntimeError(
163                    "mutable UNION cursor lost its current branch".into(),
164                ));
165            };
166
167            match current.next_row()? {
168                Some(row) => {
169                    if self.needs_dedup {
170                        let key = row
171                            .iter_named()
172                            .map(|(_, name, val)| {
173                                (name.into_owned(), GroupValueKey::from_value(val))
174                            })
175                            .collect();
176                        if !self.seen.insert(key) {
177                            continue;
178                        }
179                    }
180                    return Ok(Some(row));
181                }
182                None => {
183                    self.current.take();
184                    self.branch_idx += 1;
185                }
186            }
187        }
188    }
189}
190
191/// Streaming write cursor for plans whose root is one of
192/// `Create` / `Set` / `Delete` / `Remove` / `Merge` and whose input
193/// subtree is fully streamable.
194///
195/// # Layout invariant
196///
197/// The cursor owns a raw alias of the original `&'a mut S`.
198/// Its `upstream` was constructed using a `&'a S` reborrow derived
199/// from `storage_ptr` via unsafe lifetime extension. This is sound
200/// because the existing read-side `RowSource` impls (see
201/// `NodeScanSource::cur_ids`, `ExpandSource::cur_edges`, etc.)
202/// materialize their iteration state into owned `Vec`s at
203/// construction or first call, so no live `&S` borrow into storage
204/// persists across `next_row` calls. Read-only access happens
205/// transiently inside each `upstream.next_row` call; mutable access
206/// happens between calls inside [`MutableExecutor::apply_write_op`].
207/// The borrows never overlap in time.
208///
209/// # Drop order
210///
211/// `upstream` must drop before any caller may regain `&mut S` access
212/// to the underlying storage. The explicit `Drop` impl enforces
213/// that order — `ManuallyDrop` lets us force the sequence.
214pub struct StreamingWriteCursor<'a, S: GraphStorageMut + GraphStorage + 'a> {
215    /// SAFETY: borrows from `*storage_ptr`. Must drop first.
216    upstream: ManuallyDrop<Box<dyn RowSource + 'a>>,
217    /// Raw alias of the `&'a mut S` handed in at construction. Used
218    /// as `&S` by `upstream` and as `&mut S` inside this cursor's `next_row`.
219    storage_ptr: StoragePtr<S>,
220    /// Physical plan — kept alive for the per-row op borrow.
221    plan: &'a PhysicalPlan,
222    /// Index into `plan.nodes` of the write operator.
223    /// We re-fetch the op per call so this struct doesn't need to
224    /// be parameterized by the specific op type.
225    write_op_node: PhysicalNodeId,
226    /// Parameters; cloned per row into a fresh `MutableExecutor`.
227    /// In typical bulk-write workloads this is empty or tiny.
228    params: BTreeMap<String, LoraValue>,
229    _phantom: std::marker::PhantomData<&'a mut S>,
230}
231
232impl<'a, S: GraphStorageMut + GraphStorage + 'a> StreamingWriteCursor<'a, S> {
233    /// Build a cursor. Caller must already have verified that
234    /// `plan.nodes[write_op_node]` is a streamable write op via
235    /// [`write_op_input`] and [`subtree_is_fully_streaming`].
236    pub(crate) fn open(
237        storage: &'a mut S,
238        plan: &'a PhysicalPlan,
239        write_op_node: PhysicalNodeId,
240        params: BTreeMap<String, LoraValue>,
241    ) -> ExecResult<Self> {
242        let input = match write_op_input(plan, write_op_node) {
243            Some(i) => i,
244            None => {
245                return Err(ExecutorError::RuntimeError(format!(
246                    "StreamingWriteCursor::open called with non-write node {write_op_node:?}"
247                )));
248            }
249        };
250        let storage_ptr = StoragePtr::from_mut(storage);
251
252        // SAFETY: see struct-level comment.
253        let storage_ref: &'a S = unsafe { storage_ptr.as_ref() };
254        let upstream = build_streaming(plan, input, storage_ref, Arc::new(params.clone()))?;
255
256        Ok(Self {
257            upstream: ManuallyDrop::new(upstream),
258            storage_ptr,
259            plan,
260            write_op_node,
261            params,
262            _phantom: std::marker::PhantomData,
263        })
264    }
265}
266
267impl<'a, S: GraphStorageMut + GraphStorage + 'a> RowSource for StreamingWriteCursor<'a, S> {
268    fn next_row(&mut self) -> ExecResult<Option<Row>> {
269        let mut row = match self.upstream.next_row()? {
270            Some(r) => r,
271            None => return Ok(None),
272        };
273
274        // SAFETY: upstream's `next_row` has returned, so its
275        // dormant `&S` borrow is not in active use right now. We
276        // reborrow `&mut S` for the per-row write and drop the
277        // borrow before the next pull.
278        let storage_mut: &mut S = unsafe { self.storage_ptr.as_mut() };
279        let mut exec = MutableExecutor::new(MutableExecutionContext {
280            storage: storage_mut,
281            params: self.params.clone(),
282        });
283        let op = &self.plan.nodes[self.write_op_node];
284        exec.apply_write_op(op, &mut row)?;
285        let row = exec.hydrate_row(row);
286        Ok(Some(row))
287    }
288}
289
290impl<'a, S: GraphStorageMut + GraphStorage + 'a> Drop for StreamingWriteCursor<'a, S> {
291    fn drop(&mut self) {
292        // SAFETY: drop `upstream` first to release its borrow into
293        // `*storage_ptr`. Subsequent fields drop via the normal
294        // field-drop sequence and don't touch storage.
295        unsafe {
296            ManuallyDrop::drop(&mut self.upstream);
297        }
298    }
299}
lora_executor/pull/mutable.rs

lora_executor/pull/
mutable.rs
