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 sub_plan: Arc<PhysicalPlan>,
39 bindings: Box<QueryBindings>,
42 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#[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 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 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 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 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 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 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 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}