linkrs_query/executor/streaming/operators/
wco_operator.rs

1//! Streaming N-way WCO intersect operator.
2//!
3//! One probe input plus N build inputs share a single intersect variable.
4//! Every build side is drained into a sorted
5//! [`IntersectBuild`](super::intersect_build::IntersectBuild) table keyed by
6//! its bound value; each probe row then emits one output row per intersect
7//! value present on ALL sides, crossed with the matching payload rows.
8//!
9//! Column positions resolve per input chunk from the spec's variable names,
10//! so upstream layout changes cannot silently misalign keys. Build sides
11//! are memory-tracked per row (the hash-join discipline); spilling the
12//! sorted adjacency tables is future work, shared with the other blocking
13//! join builds.
14
15use std::sync::Arc;
16
17use crate::executor::base::MemoryTracker;
18use crate::executor::streaming::chunk::DataChunk;
19use crate::executor::streaming::executor::StreamingExecutor;
20use crate::executor::streaming::operators::intersect_build::IntersectBuild;
21use crate::executor::streaming::operators::source_operator::OperatorConfig;
22use crate::executor::streaming::runtime::ExecutionRuntime;
23use crate::executor::streaming::slot::SlotLayout;
24use linkrs_core::error::QueryError;
25use linkrs_core::value::NullType;
26use linkrs_core::Value;
27
28use super::intersect::WcoIntersectExecutor;
29use super::spec::WcoSpec;
30
31/// Column positions of one build side inside its own chunks.
32#[derive(Debug, Clone, Default)]
33struct BuildLayout {
34    bound_col: usize,
35    intersect_col: usize,
36    col_names: Vec<String>,
37}
38
39/// N-way WCO intersect operator.
40///
41/// Wraps the value-level [`WcoIntersectExecutor`] with chunk draining,
42/// name-based column resolution, memory tracking, and output assembly in
43/// the plan's column order. Lifecycle state is owned exclusively by the
44/// executor; operators never write it.
45#[derive(Debug)]
46pub struct WcoIntersectOperator {
47    spec: WcoSpec,
48    builds: Vec<IntersectBuild>,
49    build_layouts: Vec<BuildLayout>,
50    build_done: bool,
51    memory_tracker: MemoryTracker,
52    runtime: Option<Arc<ExecutionRuntime>>,
53    output_layout: Arc<SlotLayout>,
54    config: OperatorConfig,
55}
56
57impl WcoIntersectOperator {
58    pub fn new(spec: WcoSpec, output_layout: Arc<SlotLayout>) -> Self {
59        Self {
60            spec,
61            builds: Vec::new(),
62            build_layouts: Vec::new(),
63            build_done: false,
64            memory_tracker: MemoryTracker::new(crate::executor::base::MemoryBudget::new(
65                crate::executor::base::MemoryBudget::DEFAULT_MAX,
66            )),
67            runtime: None,
68            output_layout,
69            config: OperatorConfig::default(),
70        }
71    }
72
73    pub fn from_spec(
74        spec: &WcoSpec,
75        memory_budget: &crate::executor::base::MemoryBudget,
76        output_layout: Arc<SlotLayout>,
77    ) -> Self {
78        let mut op = Self::new(spec.clone(), output_layout);
79        op.memory_tracker = MemoryTracker::new(memory_budget.clone());
80        op
81    }
82
83    /// Inject the runtime and execution config (called once by the executor
84    /// before this operator produces any data).
85    pub fn inject_context(
86        &mut self,
87        runtime: Option<&Arc<ExecutionRuntime>>,
88        config: OperatorConfig,
89    ) {
90        if let Some(rt) = runtime {
91            self.runtime = Some(rt.clone());
92        }
93        self.config = config;
94    }
95
96    pub fn memory_tracker(&self) -> &MemoryTracker {
97        &self.memory_tracker
98    }
99
100    pub fn open(
101        &mut self,
102        probe: &mut StreamingExecutor,
103        builds: &mut [StreamingExecutor],
104    ) -> Result<(), QueryError> {
105        probe.open()?;
106        for build in builds.iter_mut() {
107            build.open()?;
108        }
109        Ok(())
110    }
111
112    pub fn next(
113        &mut self,
114        probe: &mut StreamingExecutor,
115        builds: &mut [StreamingExecutor],
116    ) -> Result<Option<DataChunk>, QueryError> {
117        if builds.len() != self.spec.num_builds() {
118            return Err(QueryError::execution(format!(
119                "WcoIntersect planned {} build sides but received {} inputs",
120                self.spec.num_builds(),
121                builds.len()
122            )));
123        }
124        if !self.build_done {
125            self.drain_builds(builds)?;
126        }
127        let output_layout = Arc::clone(&self.output_layout);
128        while let Some(probe_chunk) = probe.advance()? {
129            if let Some(rt) = self.runtime.as_ref() {
130                rt.ensure_not_cancelled()?;
131            }
132            let chunk_names = probe_chunk.col_names();
133            let mut bound_cols = Vec::with_capacity(self.spec.bound_names.len());
134            for name in &self.spec.bound_names {
135                bound_cols.push(chunk_names.iter().position(|c| c == name).ok_or_else(|| {
136                    QueryError::execution(format!(
137                        "WcoIntersect probe side is missing bound column `{name}`"
138                    ))
139                })?);
140            }
141            // The executor borrows the sealed tables; move them out and
142            // restore afterwards so consecutive chunks reuse the same state.
143            let builds = std::mem::take(&mut self.builds);
144            let executor = WcoIntersectExecutor::new(builds, bound_cols);
145            let mut result_rows = Vec::new();
146            for row_idx in probe_chunk.visible_indices() {
147                for wide in executor.probe_row(&probe_chunk.rows[row_idx]) {
148                    result_rows.push(self.assemble_output_row(&wide, &chunk_names));
149                }
150            }
151            self.builds = executor.into_builds();
152            if !result_rows.is_empty() {
153                return Ok(Some(DataChunk::new_with_layout(result_rows, output_layout)));
154            }
155        }
156        Ok(None)
157    }
158
159    /// Reset build tables and rewind every input so the operator
160    /// re-produces the same result set.
161    pub fn reset(
162        &mut self,
163        probe: &mut StreamingExecutor,
164        builds: &mut [StreamingExecutor],
165    ) -> Result<bool, QueryError> {
166        self.builds.clear();
167        self.build_layouts.clear();
168        self.build_done = false;
169        self.memory_tracker.reset();
170        probe.reset()?;
171        for build in builds.iter_mut() {
172            build.reset()?;
173        }
174        Ok(false)
175    }
176
177    pub fn stop(&mut self) -> Result<(), QueryError> {
178        Ok(())
179    }
180
181    pub fn close(&mut self) -> Result<(), QueryError> {
182        self.memory_tracker.reset();
183        for build in &mut self.builds {
184            build.clear();
185        }
186        Ok(())
187    }
188
189    pub fn spill_with_manager(
190        &mut self,
191        _sm: &crate::executor::streaming::spill::SpillManager,
192    ) -> Result<(), QueryError> {
193        Ok(())
194    }
195
196    pub fn spilled_bytes(&self) -> u64 {
197        0
198    }
199
200    pub fn spilled_rows(&self) -> u64 {
201        0
202    }
203
204    /// Drain every build input into its sorted adjacency table.
205    fn drain_builds(&mut self, builds: &mut [StreamingExecutor]) -> Result<(), QueryError> {
206        self.builds = Vec::with_capacity(builds.len());
207        self.build_layouts = Vec::with_capacity(builds.len());
208        for (side, build) in builds.iter_mut().enumerate() {
209            let mut table: Option<IntersectBuild> = None;
210            let mut layout = BuildLayout::default();
211            let mut col_names = Vec::new();
212            while let Some(mut chunk) = build.advance()? {
213                if let Some(rt) = self.runtime.as_ref() {
214                    rt.ensure_not_cancelled()?;
215                }
216                chunk.normalize_for_opaque("WcoBuild");
217                if col_names.is_empty() {
218                    col_names = chunk.col_names();
219                    layout = self.resolve_build_layout(side, &col_names)?;
220                }
221                for row in chunk.rows.iter() {
222                    self.memory_tracker.try_reserve_row(row)?;
223                }
224                let table = table.get_or_insert_with(|| {
225                    IntersectBuild::new(layout.bound_col, layout.intersect_col)
226                });
227                for row in chunk.rows.iter().cloned() {
228                    table.append_row(row);
229                }
230            }
231            // An empty build side yields an empty (sealed) table: every
232            // probe misses, which is the correct inner-join semantics.
233            let mut table = table.unwrap_or_else(|| IntersectBuild::new(0, 0));
234            table.finish();
235            layout.col_names = col_names;
236            self.builds.push(table);
237            self.build_layouts.push(layout);
238        }
239        self.build_done = true;
240        Ok(())
241    }
242
243    /// Resolve one build side's column positions from its chunk names.
244    fn resolve_build_layout(
245        &self,
246        side: usize,
247        col_names: &[String],
248    ) -> Result<BuildLayout, QueryError> {
249        let bound_name = &self.spec.bound_names[side];
250        let bound_col = col_names
251            .iter()
252            .position(|c| c == bound_name)
253            .ok_or_else(|| {
254                QueryError::execution(format!(
255                    "WcoIntersect build side {side} is missing bound column `{bound_name}`"
256                ))
257            })?;
258        let intersect_name = &self.spec.intersect_name;
259        let intersect_col = col_names
260            .iter()
261            .position(|c| c == intersect_name)
262            .ok_or_else(|| {
263                QueryError::execution(format!(
264                    "WcoIntersect build side {side} is missing intersect column `{intersect_name}`"
265                ))
266            })?;
267        Ok(BuildLayout {
268            bound_col,
269            intersect_col,
270            col_names: Vec::new(),
271        })
272    }
273
274    /// Map one wide library row (`probe ++ [intersect] ++ builds`) to the
275    /// plan's output column order.
276    fn assemble_output_row(&self, wide: &[Value], probe_names: &[String]) -> Vec<Value> {
277        // Library layout: probe row, one intersect value, then each full
278        // build row in side order.
279        let probe_len = probe_names.len();
280        let mut build_offsets = Vec::with_capacity(self.build_layouts.len());
281        let mut offset = probe_len + 1;
282        for layout in &self.build_layouts {
283            build_offsets.push(offset);
284            offset += layout.col_names.len();
285        }
286        let rep = wide
287            .get(probe_len)
288            .cloned()
289            .unwrap_or(Value::Null(NullType::Null));
290        let mut out = Vec::with_capacity(self.spec.output_col_names.len());
291        for name in &self.spec.output_col_names {
292            if name == &self.spec.intersect_name {
293                out.push(rep.clone());
294                continue;
295            }
296            if let Some(idx) = probe_names.iter().position(|c| c == name) {
297                out.push(
298                    wide.get(idx)
299                        .cloned()
300                        .unwrap_or(Value::Null(NullType::Null)),
301                );
302                continue;
303            }
304            let mut pushed = false;
305            for (side, layout) in self.build_layouts.iter().enumerate() {
306                if let Some(idx) = layout.col_names.iter().position(|c| c == name) {
307                    let pos = build_offsets[side] + idx;
308                    out.push(
309                        wide.get(pos)
310                            .cloned()
311                            .unwrap_or(Value::Null(NullType::Null)),
312                    );
313                    pushed = true;
314                    break;
315                }
316            }
317            if !pushed {
318                out.push(Value::Null(NullType::Null));
319            }
320        }
321        out
322    }
323}
324
325#[cfg(test)]
326mod tests {
327    use super::*;
328    use crate::executor::streaming::operators::source_operator::{
329        SourceOperator, SourceOperatorKind,
330    };
331    use crate::executor::streaming::slot::SlotLayout;
332    use linkrs_core::Value;
333
334    fn source(rows: Vec<Vec<Value>>, col_names: Vec<String>) -> StreamingExecutor {
335        let layout = Arc::new(SlotLayout::from_names(&col_names));
336        let base = crate::executor::streaming::operators::base::OperatorBase::new(1)
337            .with_output_layout(layout.clone());
338        let op = SourceOperator::new(
339            SourceOperatorKind::ScanVertices {
340                buffer: rows,
341                current_index: 0,
342                col_names,
343            },
344            layout,
345        );
346        StreamingExecutor::Source(base, op)
347    }
348
349    fn triangle_spec() -> WcoSpec {
350        WcoSpec {
351            bound_names: vec!["a".to_string(), "b".to_string()],
352            intersect_name: "c".to_string(),
353            output_col_names: vec![
354                "a".to_string(),
355                "b".to_string(),
356                "e1".to_string(),
357                "c".to_string(),
358                "e2".to_string(),
359                "e3".to_string(),
360            ],
361        }
362    }
363
364    /// Triangle: probe (a,b,e1); build e2 keyed by a carries c; build e3
365    /// keyed by b carries c. Only c=20 is common to a=1 and b=2.
366    #[test]
367    fn triangle_probes_common_adjacency() {
368        let layout = Arc::new(SlotLayout::from_names(&triangle_spec().output_col_names));
369        let mut op = WcoIntersectOperator::from_spec(
370            &triangle_spec(),
371            &crate::executor::base::MemoryBudget::new(
372                crate::executor::base::MemoryBudget::DEFAULT_MAX,
373            ),
374            layout.clone(),
375        );
376        let mut probe = source(
377            vec![vec![Value::Int(1), Value::Int(2), Value::Int(7)]],
378            vec!["a".to_string(), "b".to_string(), "e1".to_string()],
379        );
380        let mut builds = vec![
381            source(
382                vec![
383                    vec![Value::Int(1), Value::Int(10), Value::Int(11)],
384                    vec![Value::Int(1), Value::Int(20), Value::Int(12)],
385                ],
386                vec!["a".to_string(), "c".to_string(), "e2".to_string()],
387            ),
388            source(
389                vec![
390                    vec![Value::Int(2), Value::Int(20), Value::Int(13)],
391                    vec![Value::Int(2), Value::Int(40), Value::Int(14)],
392                ],
393                vec!["b".to_string(), "c".to_string(), "e3".to_string()],
394            ),
395        ];
396        op.open(&mut probe, &mut builds).expect("open");
397        let chunk = op
398            .next(&mut probe, &mut builds)
399            .expect("next")
400            .expect("one chunk");
401        assert_eq!(chunk.rows.len(), 1);
402        assert_eq!(
403            chunk.rows[0],
404            vec![
405                Value::Int(1),
406                Value::Int(2),
407                Value::Int(7),
408                Value::Int(20),
409                Value::Int(12),
410                Value::Int(13),
411            ]
412        );
413        assert!(op.next(&mut probe, &mut builds).expect("drain").is_none());
414        op.close().expect("close");
415    }
416
417    #[test]
418    fn disjoint_adjacency_emits_nothing() {
419        let layout = Arc::new(SlotLayout::from_names(&triangle_spec().output_col_names));
420        let mut op = WcoIntersectOperator::from_spec(
421            &triangle_spec(),
422            &crate::executor::base::MemoryBudget::new(
423                crate::executor::base::MemoryBudget::DEFAULT_MAX,
424            ),
425            layout,
426        );
427        let mut probe = source(
428            vec![vec![Value::Int(2), Value::Int(2), Value::Int(7)]],
429            vec!["a".to_string(), "b".to_string(), "e1".to_string()],
430        );
431        let mut builds = vec![
432            source(
433                vec![vec![Value::Int(2), Value::Int(30), Value::Int(11)]],
434                vec!["a".to_string(), "c".to_string(), "e2".to_string()],
435            ),
436            source(
437                vec![vec![Value::Int(2), Value::Int(40), Value::Int(13)]],
438                vec!["b".to_string(), "c".to_string(), "e3".to_string()],
439            ),
440        ];
441        op.open(&mut probe, &mut builds).expect("open");
442        assert!(op.next(&mut probe, &mut builds).expect("next").is_none());
443        op.close().expect("close");
444    }
445
446    #[test]
447    fn missing_bound_column_is_loud_error() {
448        let layout = Arc::new(SlotLayout::from_names(&triangle_spec().output_col_names));
449        let mut op = WcoIntersectOperator::from_spec(
450            &triangle_spec(),
451            &crate::executor::base::MemoryBudget::new(
452                crate::executor::base::MemoryBudget::DEFAULT_MAX,
453            ),
454            layout,
455        );
456        let mut probe = source(
457            vec![vec![Value::Int(1), Value::Int(2), Value::Int(7)]],
458            vec!["a".to_string(), "b".to_string(), "e1".to_string()],
459        );
460        // e2 side lacks the intersect column `c`.
461        let mut builds = vec![
462            source(
463                vec![vec![Value::Int(1), Value::Int(11)]],
464                vec!["a".to_string(), "e2".to_string()],
465            ),
466            source(
467                vec![vec![Value::Int(2), Value::Int(20), Value::Int(13)]],
468                vec!["b".to_string(), "c".to_string(), "e3".to_string()],
469            ),
470        ];
471        op.open(&mut probe, &mut builds).expect("open");
472        let err = op.next(&mut probe, &mut builds).expect_err("must fail");
473        assert!(err.to_string().contains("intersect column"));
474    }
475
476    #[test]
477    fn reset_rebuilds_and_replays() {
478        let layout = Arc::new(SlotLayout::from_names(&triangle_spec().output_col_names));
479        let mut op = WcoIntersectOperator::from_spec(
480            &triangle_spec(),
481            &crate::executor::base::MemoryBudget::new(
482                crate::executor::base::MemoryBudget::DEFAULT_MAX,
483            ),
484            layout,
485        );
486        let mut probe = source(
487            vec![vec![Value::Int(1), Value::Int(2), Value::Int(7)]],
488            vec!["a".to_string(), "b".to_string(), "e1".to_string()],
489        );
490        let mut builds = vec![
491            source(
492                vec![vec![Value::Int(1), Value::Int(20), Value::Int(12)]],
493                vec!["a".to_string(), "c".to_string(), "e2".to_string()],
494            ),
495            source(
496                vec![vec![Value::Int(2), Value::Int(20), Value::Int(13)]],
497                vec!["b".to_string(), "c".to_string(), "e3".to_string()],
498            ),
499        ];
500        op.open(&mut probe, &mut builds).expect("open");
501        let first = op
502            .next(&mut probe, &mut builds)
503            .expect("next")
504            .expect("rows");
505        assert_eq!(first.rows.len(), 1);
506        op.reset(&mut probe, &mut builds).expect("reset");
507        let replayed = op
508            .next(&mut probe, &mut builds)
509            .expect("next")
510            .expect("rows");
511        assert_eq!(replayed.rows, first.rows);
512        op.close().expect("close");
513    }
514}
linkrs_query/executor/streaming/operators/wco_operator.rs

linkrs_query/executor/streaming/operators/
wco_operator.rs
