Skip to main content

linkrs_query/executor/streaming/operators/
apply_operator.rs

1use std::sync::Arc;
2
3use crate::executor::base::{MemoryBudget, MemoryTracker};
4use crate::executor::streaming::chunk::DataChunk;
5use crate::executor::streaming::executor::StreamingExecutor;
6use crate::executor::streaming::instance::QueryBindings;
7use crate::executor::streaming::join_helpers::evaluate_join_key;
8use crate::executor::streaming::operators::source_operator::OperatorConfig;
9use crate::executor::streaming::operators::spec::{ApplyKind, ApplySpec};
10use crate::executor::streaming::plan::materializer::PhysicalPlanMaterializer;
11use crate::executor::streaming::plan::types::PhysicalPlan;
12use crate::executor::streaming::runtime::ExecutionRuntime;
13use crate::executor::streaming::slot::SlotLayout;
14use linkrs_core::error::QueryError;
15use linkrs_core::{NullType, Value};
16
17#[derive(Debug)]
18pub enum ApplyOperatorKind {
19    Apply {
20        kind: ApplyKind,
21        correlated_columns: Vec<String>,
22        right_rows: Option<Vec<Vec<Value>>>,
23        right_layout: Option<Arc<SlotLayout>>,
24        memory_tracker: MemoryTracker,
25    },
26    PatternApply {
27        hash_keys: Vec<linkrs_core::types::expr::Expression>,
28        probe_keys: Vec<linkrs_core::types::expr::Expression>,
29        anti: bool,
30        right_rows: Option<Vec<Vec<Value>>>,
31        right_layout: Option<Arc<SlotLayout>>,
32        memory_tracker: MemoryTracker,
33    },
34    CorrelatedApply {
35        /// Self-contained right subtree (rooted at an `Argument` source),
36        /// re-executed per outer row with the outer row bound as the
37        /// correlation frame.
38        sub_plan: Arc<PhysicalPlan>,
39        /// Bindings with the parameter maps stripped, cached for the nested
40        /// materialization.
41        bindings: Box<QueryBindings>,
42        /// Materialized once on the first outer row and reused via
43        /// `reset()` for every subsequent row.
44        sub_executor: Option<Box<StreamingExecutor>>,
45        anti: bool,
46        right_rows: Option<Vec<Vec<Value>>>,
47        right_layout: Option<Arc<SlotLayout>>,
48        memory_tracker: MemoryTracker,
49    },
50    RollUpApply {
51        compare_columns: Vec<String>,
52        collect_column: Option<String>,
53        right_rows: Option<Vec<Vec<Value>>>,
54        right_layout: Option<Arc<SlotLayout>>,
55        memory_tracker: MemoryTracker,
56    },
57}
58
59/// Apply operator.
60///
61/// Wraps [`ApplyOperatorKind`] with the runtime context injected at `open()`.
62/// Lifecycle state is owned exclusively by the executor; operators never
63/// write it.
64#[derive(Debug)]
65pub struct ApplyOperator {
66    pub kind: ApplyOperatorKind,
67    pub runtime: Option<Arc<ExecutionRuntime>>,
68    pub output_layout: Arc<SlotLayout>,
69    pub config: OperatorConfig,
70}
71
72impl ApplyOperator {
73    pub fn from_spec(
74        spec: &ApplySpec,
75        budget: &MemoryBudget,
76        output_layout: Arc<SlotLayout>,
77    ) -> Self {
78        let kind = match spec {
79            ApplySpec::Apply {
80                kind,
81                correlated_columns,
82            } => ApplyOperatorKind::Apply {
83                kind: *kind,
84                correlated_columns: correlated_columns.clone(),
85                right_rows: None,
86                right_layout: None,
87                memory_tracker: MemoryTracker::new(budget.clone()),
88            },
89            ApplySpec::PatternApply {
90                hash_keys,
91                probe_keys,
92                anti,
93            } => ApplyOperatorKind::PatternApply {
94                hash_keys: hash_keys.clone(),
95                probe_keys: probe_keys.clone(),
96                anti: *anti,
97                right_rows: None,
98                right_layout: None,
99                memory_tracker: MemoryTracker::new(budget.clone()),
100            },
101            ApplySpec::CorrelatedApply { .. } => {
102                // `CorrelatedApply` carries a nested physical plan and the
103                // stripped bindings, neither available from a spec alone; the
104                // materializer constructs it via the literal variant.
105                unreachable!("CorrelatedApply is constructed by the materializer")
106            }
107            ApplySpec::RollUpApply {
108                compare_columns,
109                collect_column,
110            } => ApplyOperatorKind::RollUpApply {
111                compare_columns: compare_columns.clone(),
112                collect_column: collect_column.clone(),
113                right_rows: None,
114                right_layout: None,
115                memory_tracker: MemoryTracker::new(budget.clone()),
116            },
117        };
118        Self::new(kind, output_layout)
119    }
120
121    pub fn new(kind: ApplyOperatorKind, output_layout: Arc<SlotLayout>) -> Self {
122        Self {
123            kind,
124            runtime: None,
125            output_layout,
126            config: OperatorConfig::default(),
127        }
128    }
129
130    /// Inject the runtime and execution config (called once by the executor
131    /// before this operator produces any data).
132    pub fn inject_context(
133        &mut self,
134        runtime: Option<&Arc<ExecutionRuntime>>,
135        config: OperatorConfig,
136    ) {
137        if let Some(rt) = runtime {
138            self.runtime = Some(rt.clone());
139        }
140        self.config = config;
141    }
142
143    pub fn memory_tracker(&self) -> &MemoryTracker {
144        match &self.kind {
145            ApplyOperatorKind::Apply { memory_tracker, .. }
146            | ApplyOperatorKind::PatternApply { memory_tracker, .. }
147            | ApplyOperatorKind::CorrelatedApply { memory_tracker, .. }
148            | ApplyOperatorKind::RollUpApply { memory_tracker, .. } => memory_tracker,
149        }
150    }
151
152    pub fn open(
153        &mut self,
154        left: &mut StreamingExecutor,
155        right: &mut StreamingExecutor,
156    ) -> Result<(), QueryError> {
157        left.open()?;
158        // CorrelatedApply keeps only the left input; the right placeholder
159        // (`SourceOperator::Start`) is never executed because the nested right
160        // subtree is re-materialized per outer row at runtime.
161        if !matches!(&self.kind, ApplyOperatorKind::CorrelatedApply { .. }) {
162            right.open()?;
163        }
164        Ok(())
165    }
166
167    pub fn next(
168        &mut self,
169        left: &mut StreamingExecutor,
170        right: &mut StreamingExecutor,
171    ) -> Result<Option<DataChunk>, QueryError> {
172        match &mut self.kind {
173            ApplyOperatorKind::Apply {
174                kind,
175                correlated_columns,
176                right_rows,
177                right_layout,
178                memory_tracker,
179            } => {
180                materialize_right(
181                    &self.runtime,
182                    right,
183                    right_rows,
184                    right_layout,
185                    memory_tracker,
186                )?;
187                let output_layout = Arc::clone(&self.output_layout);
188                let right_rows = right_rows.as_deref().unwrap_or_default();
189                let right_layout = right_layout
190                    .as_ref()
191                    .cloned()
192                    .unwrap_or_else(|| Arc::new(SlotLayout::new(Vec::new())));
193                loop {
194                    let Some(mut left_chunk) = left.advance()? else {
195                        return Ok(None);
196                    };
197                    left_chunk.normalize_for_opaque("Apply");
198                    let mut output = Vec::new();
199                    for left_row in left_chunk.rows {
200                        if let Some(rt) = &self.runtime {
201                            rt.ensure_not_cancelled()?;
202                        }
203                        let matches = matching_rows(
204                            &left_row,
205                            &left_chunk.layout,
206                            right_rows,
207                            &right_layout,
208                            correlated_columns,
209                        )?;
210                        match kind {
211                            ApplyKind::Standard => {
212                                for right_row in matches {
213                                    let mut row = left_row.clone();
214                                    row.extend_from_slice(right_row);
215                                    output.push(row);
216                                }
217                            }
218                            ApplyKind::Semi if !matches.is_empty() => output.push(left_row),
219                            ApplyKind::Anti if matches.is_empty() => output.push(left_row),
220                            ApplyKind::Single => match matches.as_slice() {
221                                [] => {
222                                    let mut row = left_row;
223                                    row.extend(std::iter::repeat_n(
224                                        Value::Null(NullType::Null),
225                                        right_layout.len(),
226                                    ));
227                                    output.push(row);
228                                }
229                                [right_row] => {
230                                    let mut row = left_row;
231                                    row.extend_from_slice(right_row);
232                                    output.push(row);
233                                }
234                                _ => {
235                                    return Err(QueryError::execution(
236                                        "Single Apply produced more than one matching row"
237                                            .to_string(),
238                                    ));
239                                }
240                            },
241                            ApplyKind::All if matches.len() == right_rows.len() => {
242                                output.push(left_row);
243                            }
244                            _ => {}
245                        }
246                    }
247                    if !output.is_empty() {
248                        return Ok(Some(DataChunk::new_with_layout(
249                            output,
250                            Arc::clone(&output_layout),
251                        )));
252                    }
253                }
254            }
255            ApplyOperatorKind::PatternApply {
256                hash_keys,
257                probe_keys,
258                anti,
259                right_rows,
260                right_layout,
261                memory_tracker,
262            } => {
263                materialize_right(
264                    &self.runtime,
265                    right,
266                    right_rows,
267                    right_layout,
268                    memory_tracker,
269                )?;
270                let output_layout = Arc::clone(&self.output_layout);
271                let right_rows = right_rows.as_deref().unwrap_or_default();
272                let right_layout = right_layout
273                    .as_ref()
274                    .cloned()
275                    .unwrap_or_else(|| Arc::new(SlotLayout::new(Vec::new())));
276                loop {
277                    let Some(mut left_chunk) = left.advance()? else {
278                        return Ok(None);
279                    };
280                    left_chunk.normalize_for_opaque("Apply");
281                    let mut output = Vec::new();
282                    for left_row in left_chunk.rows {
283                        if let Some(rt) = &self.runtime {
284                            rt.ensure_not_cancelled()?;
285                        }
286                        let left_key =
287                            evaluate_join_key(&left_row, left_chunk.layout.clone(), hash_keys)?;
288                        let mut exists = false;
289                        for right_row in right_rows {
290                            let right_key =
291                                evaluate_join_key(right_row, right_layout.clone(), probe_keys)?;
292                            if keys_match(&left_key, &right_key) {
293                                exists = true;
294                                break;
295                            }
296                        }
297                        if exists != *anti {
298                            output.push(left_row);
299                        }
300                    }
301                    if !output.is_empty() {
302                        return Ok(Some(DataChunk::new_with_layout(
303                            output,
304                            Arc::clone(&output_layout),
305                        )));
306                    }
307                }
308            }
309            ApplyOperatorKind::CorrelatedApply {
310                sub_plan,
311                bindings,
312                anti,
313                sub_executor,
314                ..
315            } => {
316                let rt = self.runtime.clone().ok_or_else(|| {
317                    QueryError::execution(
318                        "CorrelatedApply requires an execution runtime".to_string(),
319                    )
320                })?;
321                let output_layout = Arc::clone(&self.output_layout);
322                loop {
323                    let Some(mut left_chunk) = left.advance()? else {
324                        return Ok(None);
325                    };
326                    left_chunk.normalize_for_opaque("CorrelatedApply");
327                    let mut output = Vec::new();
328                    for left_row in left_chunk.rows {
329                        if let Some(rt) = &self.runtime {
330                            rt.ensure_not_cancelled()?;
331                        }
332                        // Materialize the self-contained right subtree once
333                        // and reuse it for every outer row via the reset
334                        // protocol. The frame is injected per row into the
335                        // Argument source's private slot.
336                        if sub_executor.is_none() {
337                            let (mut exec, _) =
338                                PhysicalPlanMaterializer::materialize(sub_plan, bindings)?;
339                            exec.set_chunk_size(bindings.chunk_size);
340                            exec.set_runtime(Some(rt.clone()));
341                            exec.open()?;
342                            *sub_executor = Some(Box::new(exec));
343                        }
344                        let exec = sub_executor.as_mut().ok_or_else(|| {
345                            QueryError::execution(
346                                "CorrelatedApply sub-executor failed to materialize".to_string(),
347                            )
348                        })?;
349                        exec.inject_correlation_frame(left_chunk.layout.clone(), left_row.clone());
350                        exec.reset()?;
351                        let mut exists = false;
352                        while let Some(mut sub_chunk) = exec.advance()? {
353                            sub_chunk.normalize_for_opaque("CorrelatedApply");
354                            if !sub_chunk.rows.is_empty() {
355                                exists = true;
356                                break;
357                            }
358                        }
359                        if exists != *anti {
360                            output.push(left_row);
361                        }
362                    }
363                    if !output.is_empty() {
364                        return Ok(Some(DataChunk::new_with_layout(
365                            output,
366                            Arc::clone(&output_layout),
367                        )));
368                    }
369                }
370            }
371            ApplyOperatorKind::RollUpApply {
372                compare_columns,
373                collect_column,
374                right_rows,
375                right_layout,
376                memory_tracker,
377            } => {
378                materialize_right(
379                    &self.runtime,
380                    right,
381                    right_rows,
382                    right_layout,
383                    memory_tracker,
384                )?;
385                let output_layout = Arc::clone(&self.output_layout);
386                let right_rows = right_rows.as_deref().unwrap_or_default();
387                let right_layout = right_layout
388                    .as_ref()
389                    .cloned()
390                    .unwrap_or_else(|| Arc::new(SlotLayout::new(Vec::new())));
391                let collect_slot = collect_column
392                    .as_deref()
393                    .map(|column| {
394                        right_layout.resolve(column).ok_or_else(|| {
395                            QueryError::execution(format!(
396                                "RollUpApply collect column not found: {column}"
397                            ))
398                        })
399                    })
400                    .transpose()?;
401                loop {
402                    let Some(mut left_chunk) = left.advance()? else {
403                        return Ok(None);
404                    };
405                    left_chunk.normalize_for_opaque("Apply");
406                    let mut output = Vec::with_capacity(left_chunk.rows.len());
407                    for left_row in left_chunk.rows {
408                        if let Some(rt) = &self.runtime {
409                            rt.ensure_not_cancelled()?;
410                        }
411                        let matches = matching_rows(
412                            &left_row,
413                            &left_chunk.layout,
414                            right_rows,
415                            &right_layout,
416                            compare_columns,
417                        )?;
418                        let values = matches
419                            .into_iter()
420                            .map(|row| {
421                                collect_slot
422                                    .and_then(|slot| row.get(slot).cloned())
423                                    .unwrap_or_else(|| {
424                                        Value::List(Box::new(linkrs_core::value::List {
425                                            values: row.clone(),
426                                        }))
427                                    })
428                            })
429                            .collect();
430                        let mut row = left_row;
431                        row.push(Value::List(Box::new(linkrs_core::value::List { values })));
432                        output.push(row);
433                    }
434                    if !output.is_empty() {
435                        return Ok(Some(DataChunk::new_with_layout(
436                            output,
437                            Arc::clone(&output_layout),
438                        )));
439                    }
440                }
441            }
442        }
443    }
444
445    pub fn stop(&mut self) -> Result<(), QueryError> {
446        Ok(())
447    }
448
449    /// Reset per-run materialized state and rewind both inputs.
450    ///
451    /// `Apply`/`PatternApply`/`RollUpApply` drop their materialized right
452    /// side so the next run re-pulls it; `CorrelatedApply` additionally
453    /// resets its reused sub-executor (nested subqueries reset recursively).
454    pub fn reset(
455        &mut self,
456        left: &mut StreamingExecutor,
457        right: &mut StreamingExecutor,
458    ) -> Result<bool, QueryError> {
459        match &mut self.kind {
460            ApplyOperatorKind::Apply {
461                right_rows,
462                right_layout,
463                ..
464            }
465            | ApplyOperatorKind::PatternApply {
466                right_rows,
467                right_layout,
468                ..
469            }
470            | ApplyOperatorKind::RollUpApply {
471                right_rows,
472                right_layout,
473                ..
474            } => {
475                right_rows.take();
476                right_layout.take();
477            }
478            ApplyOperatorKind::CorrelatedApply {
479                sub_executor,
480                right_rows,
481                right_layout,
482                ..
483            } => {
484                if let Some(exec) = sub_executor {
485                    exec.reset()?;
486                }
487                right_rows.take();
488                right_layout.take();
489            }
490        }
491        left.reset()?;
492        right.reset()?;
493        Ok(false)
494    }
495
496    pub fn close(&mut self) -> Result<(), QueryError> {
497        match &mut self.kind {
498            ApplyOperatorKind::Apply {
499                right_rows,
500                right_layout,
501                memory_tracker,
502                ..
503            }
504            | ApplyOperatorKind::PatternApply {
505                right_rows,
506                right_layout,
507                memory_tracker,
508                ..
509            }
510            | ApplyOperatorKind::CorrelatedApply {
511                right_rows,
512                right_layout,
513                memory_tracker,
514                ..
515            }
516            | ApplyOperatorKind::RollUpApply {
517                right_rows,
518                right_layout,
519                memory_tracker,
520                ..
521            } => {
522                right_rows.take();
523                right_layout.take();
524                memory_tracker.reset();
525            }
526        }
527        if let ApplyOperatorKind::CorrelatedApply { sub_executor, .. } = &mut self.kind {
528            if let Some(mut exec) = sub_executor.take() {
529                exec.close_tree()?;
530            }
531        }
532        Ok(())
533    }
534}
535
536fn materialize_right(
537    runtime: &Option<Arc<ExecutionRuntime>>,
538    right: &mut StreamingExecutor,
539    rows: &mut Option<Vec<Vec<Value>>>,
540    layout: &mut Option<Arc<SlotLayout>>,
541    memory_tracker: &mut MemoryTracker,
542) -> Result<(), QueryError> {
543    if rows.is_some() {
544        return Ok(());
545    }
546    let mut materialized = Vec::new();
547    while let Some(mut chunk) = right.advance()? {
548        chunk.normalize_for_opaque("Apply");
549        if let Some(rt) = runtime {
550            rt.ensure_not_cancelled()?;
551        }
552        if layout.is_none() {
553            *layout = Some(chunk.get_layout());
554        }
555        for row in &chunk.rows {
556            memory_tracker.try_reserve_row(row)?;
557        }
558        materialized.extend(chunk.rows);
559    }
560    *rows = Some(materialized);
561    Ok(())
562}
563
564fn matching_rows<'a>(
565    left_row: &[Value],
566    left_layout: &SlotLayout,
567    right_rows: &'a [Vec<Value>],
568    right_layout: &SlotLayout,
569    correlated_columns: &[String],
570) -> Result<Vec<&'a Vec<Value>>, QueryError> {
571    let mut slots = Vec::with_capacity(correlated_columns.len());
572    for column in correlated_columns {
573        let left_slot = left_layout.resolve(column).ok_or_else(|| {
574            QueryError::execution(format!("Apply left correlation column not found: {column}"))
575        })?;
576        let right_slot = right_layout.resolve(column).ok_or_else(|| {
577            QueryError::execution(format!(
578                "Apply right correlation column not found: {column}"
579            ))
580        })?;
581        slots.push((left_slot, right_slot));
582    }
583    Ok(right_rows
584        .iter()
585        .filter(|right_row| {
586            slots.iter().all(|(left_slot, right_slot)| {
587                match (left_row.get(*left_slot), right_row.get(*right_slot)) {
588                    (Some(Value::Null(_)), _) | (_, Some(Value::Null(_))) => false,
589                    (Some(left), Some(right)) => left == right,
590                    _ => false,
591                }
592            })
593        })
594        .collect())
595}
596
597fn keys_match(left: &[Value], right: &[Value]) -> bool {
598    left.len() == right.len()
599        && left.iter().zip(right).all(|(left, right)| {
600            !matches!(left, Value::Null(_)) && !matches!(right, Value::Null(_)) && left == right
601        })
602}
603
604#[cfg(test)]
605mod tests {
606    use super::*;
607    use crate::executor::streaming::operators::base::OperatorBase;
608    use crate::executor::streaming::operators::source_operator::{
609        SourceOperator, SourceOperatorKind,
610    };
611    use crate::executor::streaming::runtime::ExecutionRuntime;
612    use crate::planning::plan::core::nodes::base::plan_node_enum::PlanNodeEnum;
613    use crate::planning::plan::core::nodes::base::plan_node_traits::PlanNode;
614    use linkrs_core::types::ContextualExpression;
615
616    fn scan(rows: Vec<Vec<Value>>, col_names: Vec<String>) -> StreamingExecutor {
617        let layout = Arc::new(SlotLayout::from_names(&col_names));
618        StreamingExecutor::Source(
619            OperatorBase::new(0),
620            SourceOperator::new(
621                SourceOperatorKind::ScanVertices {
622                    buffer: rows,
623                    current_index: 0,
624                    col_names,
625                },
626                layout,
627            ),
628        )
629    }
630
631    fn execute_apply(
632        spec: ApplySpec,
633        left_rows: Vec<Vec<Value>>,
634        right_rows: Vec<Vec<Value>>,
635    ) -> Result<Vec<Vec<Value>>, QueryError> {
636        let left = scan(left_rows, vec!["id".to_string()]);
637        let right = scan(right_rows, vec!["id".to_string()]);
638        let budget = MemoryBudget::default_budget();
639        let operator =
640            ApplyOperator::from_spec(&spec, &budget, Arc::new(SlotLayout::new(Vec::new())));
641        let mut executor = StreamingExecutor::Apply(
642            OperatorBase::new(3),
643            Box::new(left),
644            Box::new(right),
645            operator,
646        );
647        executor.open()?;
648        let mut rows = Vec::new();
649        while let Some(mut chunk) = executor.advance()? {
650            chunk.normalize_for_opaque("Apply");
651            rows.extend(chunk.rows);
652        }
653        executor.close()?;
654        Ok(rows)
655    }
656
657    #[test]
658    fn semi_and_anti_apply_consume_the_right_input() {
659        let semi = execute_apply(
660            ApplySpec::Apply {
661                kind: ApplyKind::Semi,
662                correlated_columns: vec!["id".to_string()],
663            },
664            vec![vec![Value::Int(1)], vec![Value::Int(2)]],
665            vec![vec![Value::Int(2)]],
666        )
667        .expect("semi apply should execute");
668        assert_eq!(semi, vec![vec![Value::Int(2)]]);
669
670        let anti = execute_apply(
671            ApplySpec::Apply {
672                kind: ApplyKind::Anti,
673                correlated_columns: vec!["id".to_string()],
674            },
675            vec![vec![Value::Int(1)], vec![Value::Int(2)]],
676            vec![vec![Value::Int(2)]],
677        )
678        .expect("anti apply should execute");
679        assert_eq!(anti, vec![vec![Value::Int(1)]]);
680    }
681
682    #[test]
683    fn single_apply_rejects_multiple_matches() {
684        let result = execute_apply(
685            ApplySpec::Apply {
686                kind: ApplyKind::Single,
687                correlated_columns: vec!["id".to_string()],
688            },
689            vec![vec![Value::Int(1)]],
690            vec![vec![Value::Int(1)], vec![Value::Int(1)]],
691        );
692        assert!(result.is_err());
693    }
694
695    #[test]
696    fn apply_uses_the_planned_output_layout() {
697        let left = scan(vec![vec![Value::Int(7)]], vec!["left_input".to_string()]);
698        let right = scan(vec![vec![Value::Int(7)]], vec!["left_input".to_string()]);
699        let output_layout = Arc::new(SlotLayout::from_names(&[
700            "planned_left".to_string(),
701            "planned_right".to_string(),
702        ]));
703        let operator = ApplyOperator::from_spec(
704            &ApplySpec::Apply {
705                kind: ApplyKind::Standard,
706                correlated_columns: vec!["left_input".to_string()],
707            },
708            &MemoryBudget::default_budget(),
709            output_layout,
710        );
711        let mut executor = StreamingExecutor::Apply(
712            OperatorBase::new(3),
713            Box::new(left),
714            Box::new(right),
715            operator,
716        );
717
718        executor.open().expect("apply should open");
719        let chunk = executor
720            .advance()
721            .expect("apply should advance")
722            .expect("apply should produce a row");
723        executor.close().expect("apply should close");
724
725        assert_eq!(chunk.rows, vec![vec![Value::Int(7), Value::Int(7)]]);
726        assert_eq!(
727            chunk.layout.names(),
728            vec!["planned_left".to_string(), "planned_right".to_string()]
729        );
730    }
731
732    /// Build a self-contained right subtree rooted at an `Argument` source:
733    /// `Filter(condition) -> CrossJoin(Argument(col_names = ["id"]), Start)`.
734    /// When `condition` is `None` the filter is omitted, so the subtree is
735    /// non-empty for every correlation frame.
736    fn build_sub_plan(condition: Option<ContextualExpression>) -> Arc<PhysicalPlan> {
737        use crate::executor::base::ExecutionContext;
738        use crate::executor::streaming::plan::arena_builder::PhysicalPlanBuilder;
739        use crate::executor::streaming::plan::context::PhysicalPlanBuildContext;
740        use crate::executor::streaming::plan::validator::PhysicalPlanValidator;
741        use crate::planning::plan::core::nodes::control_flow::start_node::StartNode;
742        use crate::planning::plan::core::nodes::control_flow::ArgumentNode;
743        use crate::planning::plan::core::nodes::join::CrossJoinNode;
744        use crate::planning::plan::core::nodes::operation::filter_node::FilterNode;
745        use linkrs_core::types::expr::expression_context::ExpressionAnalysisContext;
746
747        let mut argument = ArgumentNode::new(-2, "_correlated_apply");
748        argument.set_col_names(vec!["id".to_string()]);
749        let start = PlanNodeEnum::Start(StartNode::new());
750        let cross = CrossJoinNode::new(PlanNodeEnum::Argument(argument), start)
751            .expect("cross join should build");
752        let node = match condition {
753            Some(condition) => PlanNodeEnum::Filter(
754                FilterNode::new(cross.into_enum(), condition).expect("filter should build"),
755            ),
756            None => cross.into_enum(),
757        };
758
759        let mut ctx = PhysicalPlanBuildContext::new();
760        let exec_ctx = ExecutionContext::new(Arc::new(ExpressionAnalysisContext::new()));
761        let plan =
762            PhysicalPlanBuilder::build(&node, &mut ctx, &exec_ctx).expect("sub-plan should build");
763        let plan = Arc::new(plan);
764        PhysicalPlanValidator::validate(&plan).expect("sub-plan should validate");
765        plan
766    }
767
768    fn correlated_bindings() -> QueryBindings {
769        QueryBindings {
770            parameters: Arc::new(std::collections::HashMap::new()),
771            session_variables: Arc::new(std::collections::HashMap::new()),
772            parameter_frame: None,
773            space_name: None,
774            storage: None,
775            bound_snapshot: None,
776            memory_budget: MemoryBudget::default_budget(),
777            max_workers: 1,
778            chunk_size: 2048,
779            max_buffered_chunks: 4,
780            query_id: 1,
781            cancel_token: None,
782            session_id: None,
783            user_name: None,
784            query_text: None,
785            transaction: crate::executor::streaming::transaction_scope::TransactionScope::None,
786            shared_scheduler: None,
787            partition_count: 0,
788            arena: None,
789            feedback_history: None,
790            columnar_policy: None,
791            macro_manager: None,
792            type_alias_manager: None,
793            sequence_manager: None,
794            search: crate::executor::base::SearchContext::default(),
795        }
796    }
797
798    fn execute_correlated_apply(
799        sub_plan: Arc<PhysicalPlan>,
800        anti: bool,
801        left_rows: Vec<Vec<Value>>,
802    ) -> Result<Vec<Vec<Value>>, QueryError> {
803        let left = scan(left_rows, vec!["id".to_string()]);
804        let right = StreamingExecutor::Source(
805            OperatorBase::new(0),
806            SourceOperator::new(
807                SourceOperatorKind::Start,
808                Arc::new(SlotLayout::new(Vec::new())),
809            ),
810        );
811        let budget = MemoryBudget::default_budget();
812        let operator = ApplyOperator::new(
813            ApplyOperatorKind::CorrelatedApply {
814                sub_plan,
815                bindings: Box::new(correlated_bindings()),
816                sub_executor: None,
817                anti,
818                right_rows: None,
819                right_layout: None,
820                memory_tracker: MemoryTracker::new(budget),
821            },
822            Arc::new(SlotLayout::new(Vec::new())),
823        );
824        let mut executor = StreamingExecutor::Apply(
825            OperatorBase::new(3),
826            Box::new(left),
827            Box::new(right),
828            operator,
829        );
830        executor.set_chunk_size(2048);
831        executor.set_runtime(Some(Arc::new(ExecutionRuntime::default_budget())));
832        executor.open()?;
833        let mut rows = Vec::new();
834        while let Some(mut chunk) = executor.advance()? {
835            chunk.normalize_for_opaque("CorrelatedApply");
836            rows.extend(chunk.rows);
837        }
838        executor.close()?;
839        Ok(rows)
840    }
841
842    #[test]
843    fn correlated_apply_semi_and_anti_use_right_subtree_existence() {
844        let sub_plan = build_sub_plan(None);
845
846        let semi = execute_correlated_apply(
847            sub_plan.clone(),
848            false,
849            vec![vec![Value::Int(1)], vec![Value::Int(2)]],
850        )
851        .expect("semi correlated apply should execute");
852        assert_eq!(
853            semi,
854            vec![vec![Value::Int(1)], vec![Value::Int(2)]],
855            "semi keeps rows whose right subtree is non-empty"
856        );
857
858        let anti = execute_correlated_apply(
859            sub_plan,
860            true,
861            vec![vec![Value::Int(1)], vec![Value::Int(2)]],
862        )
863        .expect("anti correlated apply should execute");
864        assert!(
865            anti.is_empty(),
866            "anti drops rows whose right subtree is non-empty"
867        );
868    }
869
870    #[test]
871    fn correlated_apply_binds_the_frame_per_row() {
872        // Filter that only passes the frame whose `id` equals 2. The left
873        // input carries two rows, so the per-row frame must differ.
874        use linkrs_core::types::expr::expression_context::ExpressionAnalysisContext;
875        use linkrs_core::types::operators::BinaryOperator;
876        use linkrs_core::types::{ContextualExpression, ExpressionMeta};
877        use linkrs_core::Expression;
878
879        let condition_expr = Expression::binary(
880            Expression::variable("id"),
881            BinaryOperator::Equal,
882            Expression::literal(linkrs_core::Value::Int(2)),
883        );
884        let expr_ctx = Arc::new(ExpressionAnalysisContext::new());
885        let expr_id = expr_ctx.register_expression(ExpressionMeta::new(condition_expr));
886        let condition = ContextualExpression::new(expr_id, expr_ctx);
887        let sub_plan = build_sub_plan(Some(condition));
888
889        let semi = execute_correlated_apply(
890            sub_plan.clone(),
891            false,
892            vec![vec![Value::Int(1)], vec![Value::Int(2)]],
893        )
894        .expect("semi correlated apply should execute");
895        assert_eq!(
896            semi,
897            vec![vec![Value::Int(2)]],
898            "only the row whose bound frame satisfies the filter is kept"
899        );
900
901        let anti = execute_correlated_apply(
902            sub_plan,
903            true,
904            vec![vec![Value::Int(1)], vec![Value::Int(2)]],
905        )
906        .expect("anti correlated apply should execute");
907        assert_eq!(
908            anti,
909            vec![vec![Value::Int(1)]],
910            "anti keeps only rows whose bound frame makes the subtree empty"
911        );
912    }
913}