1use alopex_core::kv::KVStore;
2use alopex_core::sql::stream::ByteSized;
3use std::collections::{HashMap, HashSet};
4
5use crate::ast::LITERAL_TABLE;
6use crate::catalog::{Catalog, StorageType};
7use crate::executor::evaluator::EvalContext;
8use crate::executor::memory::{MemoryPolicy, MemoryTracker, map_core_memory_error};
9use crate::executor::{ExecutionResult, ExecutorError, QueryResult, QueryRowIterator, Result};
10use crate::planner::logical_plan::{LogicalPlan, RecursiveCteLimits, SetOperator};
11use crate::planner::typed_expr::{Projection, SortExpr};
12use crate::storage::{SqlTxn, SqlValue};
13
14use super::{ColumnInfo, Row};
15
16pub mod aggregate;
17pub mod columnar_scan;
18pub mod iterator;
19pub mod join;
20mod knn;
21mod project;
22mod scan;
23pub mod subquery;
24pub mod window;
25
26pub use columnar_scan::{ColumnarScanIterator, create_columnar_scan_iterator};
27pub use iterator::{
28 DistinctOnIterator, FilterIterator, LimitIterator, RowIterator, ScanIterator, SortIterator,
29};
30pub use project::{project_row_values, projected_columns};
31pub use scan::{
32 create_fenced_range_scan_iterator, create_scan_iterator, execute_fenced_range_scan,
33};
34
35#[derive(Clone)]
36struct RecursiveWorkingTable {
37 rows: Vec<Vec<SqlValue>>,
38 schema: Vec<crate::catalog::ColumnMetadata>,
39}
40
41#[derive(Clone, Default)]
44struct QueryExecutionContext {
45 recursive_tables: HashMap<String, RecursiveWorkingTable>,
46}
47
48struct RecursiveCteExecution {
49 name: String,
50 anchor: LogicalPlan,
51 recursive_term: LogicalPlan,
52 union_all: bool,
53 schema: Vec<crate::catalog::ColumnMetadata>,
54 limits: RecursiveCteLimits,
55}
56
57impl QueryExecutionContext {
58 fn with_recursive_table(
59 &self,
60 name: String,
61 table: RecursiveWorkingTable,
62 ) -> QueryExecutionContext {
63 let mut next = self.clone();
64 next.recursive_tables.insert(name, table);
65 next
66 }
67}
68
69fn plan_contains_recursive_cte(plan: &LogicalPlan) -> bool {
70 match plan {
71 LogicalPlan::RecursiveCte { .. } | LogicalPlan::RecursiveReference { .. } => true,
72 LogicalPlan::Filter { input, .. }
73 | LogicalPlan::Project { input, .. }
74 | LogicalPlan::Aggregate { input, .. }
75 | LogicalPlan::Window { input, .. }
76 | LogicalPlan::Sort { input, .. }
77 | LogicalPlan::DistinctOn { input, .. }
78 | LogicalPlan::Limit { input, .. } => plan_contains_recursive_cte(input),
79 LogicalPlan::Join { left, right, .. }
80 | LogicalPlan::LateralJoin { left, right, .. }
81 | LogicalPlan::SetOperation { left, right, .. } => {
82 plan_contains_recursive_cte(left) || plan_contains_recursive_cte(right)
83 }
84 _ => false,
85 }
86}
87
88fn plan_contains_lateral(plan: &LogicalPlan) -> bool {
95 match plan {
96 LogicalPlan::LateralJoin { .. } | LogicalPlan::TableFunction { .. } => true,
97 LogicalPlan::Filter { input, .. }
98 | LogicalPlan::Project { input, .. }
99 | LogicalPlan::Aggregate { input, .. }
100 | LogicalPlan::Window { input, .. }
101 | LogicalPlan::Sort { input, .. }
102 | LogicalPlan::DistinctOn { input, .. }
103 | LogicalPlan::Limit { input, .. } => plan_contains_lateral(input),
104 LogicalPlan::Join { left, right, .. } | LogicalPlan::SetOperation { left, right, .. } => {
105 plan_contains_lateral(left) || plan_contains_lateral(right)
106 }
107 LogicalPlan::RecursiveCte {
108 anchor,
109 recursive_term,
110 ..
111 } => plan_contains_lateral(anchor) || plan_contains_lateral(recursive_term),
112 _ => false,
113 }
114}
115
116pub fn execute_query<'txn, S: KVStore + 'txn, C: Catalog + ?Sized, T: SqlTxn<'txn, S>>(
128 txn: &mut T,
129 catalog: &C,
130 plan: LogicalPlan,
131) -> Result<ExecutionResult> {
132 execute_query_with_policy(txn, catalog, plan, None)
133}
134
135pub fn execute_query_with_policy<
136 'txn,
137 S: KVStore + 'txn,
138 C: Catalog + ?Sized,
139 T: SqlTxn<'txn, S>,
140>(
141 txn: &mut T,
142 catalog: &C,
143 plan: LogicalPlan,
144 memory: Option<&MemoryPolicy>,
145) -> Result<ExecutionResult> {
146 if let Some((pattern, projection, filter)) = knn::extract_knn_context(&plan) {
147 return knn::execute_knn_query(txn, catalog, &pattern, &projection, filter.as_ref());
148 }
149
150 let result = execute_query_result_with_outer_and_policy(txn, catalog, plan, None, memory)?;
151 Ok(ExecutionResult::Query(result))
152}
153
154pub(crate) fn execute_query_result_with_outer<
155 'txn,
156 S: KVStore + 'txn,
157 C: Catalog + ?Sized,
158 T: SqlTxn<'txn, S>,
159>(
160 txn: &mut T,
161 catalog: &C,
162 plan: LogicalPlan,
163 outer: Option<&Row>,
164) -> Result<QueryResult> {
165 execute_query_result_with_outer_and_policy(txn, catalog, plan, outer, None)
166}
167
168fn execute_query_result_with_outer_and_policy<
169 'txn,
170 S: KVStore + 'txn,
171 C: Catalog + ?Sized,
172 T: SqlTxn<'txn, S>,
173>(
174 txn: &mut T,
175 catalog: &C,
176 plan: LogicalPlan,
177 outer: Option<&Row>,
178 memory: Option<&MemoryPolicy>,
179) -> Result<QueryResult> {
180 execute_query_result_with_context(
181 txn,
182 catalog,
183 plan,
184 outer,
185 memory,
186 &QueryExecutionContext::default(),
187 )
188}
189
190fn execute_query_result_with_context<
191 'txn,
192 S: KVStore + 'txn,
193 C: Catalog + ?Sized,
194 T: SqlTxn<'txn, S>,
195>(
196 txn: &mut T,
197 catalog: &C,
198 plan: LogicalPlan,
199 outer: Option<&Row>,
200 memory: Option<&MemoryPolicy>,
201 context: &QueryExecutionContext,
202) -> Result<QueryResult> {
203 let (mut iter, projection, schema) =
204 build_iterator_pipeline_with_outer(txn, catalog, plan, memory, outer, context)?;
205 let mut rows = Vec::new();
206 while let Some(result) = iter.next_row() {
207 rows.push(result?);
208 }
209 execute_project_with_subqueries(txn, catalog, rows, &projection, &schema, outer)
210}
211
212pub fn execute_query_streaming<'txn, S: KVStore + 'txn, C: Catalog + ?Sized, T: SqlTxn<'txn, S>>(
228 txn: &mut T,
229 catalog: &C,
230 plan: LogicalPlan,
231) -> Result<QueryRowIterator<'static>> {
232 execute_query_streaming_with_policy(txn, catalog, plan, None)
233}
234
235pub fn execute_query_streaming_with_policy<
236 'txn,
237 S: KVStore + 'txn,
238 C: Catalog + ?Sized,
239 T: SqlTxn<'txn, S>,
240>(
241 txn: &mut T,
242 catalog: &C,
243 plan: LogicalPlan,
244 memory: Option<&MemoryPolicy>,
245) -> Result<QueryRowIterator<'static>> {
246 if knn::extract_knn_context(&plan).is_some() {
248 let result = execute_query_with_policy(txn, catalog, plan, memory)?;
250 if let ExecutionResult::Query(qr) = result {
251 let (iter, projection, schema) = materialize_query_result(qr);
252 return Ok(QueryRowIterator::new(iter, projection, schema));
253 }
254 return Err(ExecutorError::InvalidOperation {
255 operation: "execute_query_streaming".into(),
256 reason: "KNN query did not return Query result".into(),
257 });
258 }
259
260 if subquery::plan_contains_subquery(&plan)
265 || plan_contains_recursive_cte(&plan)
266 || plan_contains_lateral(&plan)
267 {
268 let result = execute_query_result_with_outer_and_policy(txn, catalog, plan, None, memory)?;
269 let (iter, projection, schema) = materialize_query_result(result);
270 return Ok(QueryRowIterator::new(iter, projection, schema));
271 }
272
273 let (iter, projection, schema) = build_iterator_pipeline(txn, catalog, plan, memory)?;
274
275 Ok(QueryRowIterator::new(iter, projection, schema))
276}
277
278fn materialize_query_result(
283 result: QueryResult,
284) -> (
285 Box<dyn RowIterator>,
286 Projection,
287 Vec<crate::catalog::ColumnMetadata>,
288) {
289 let column_names: Vec<String> = result.columns.iter().map(|c| c.name.clone()).collect();
290 let schema: Vec<crate::catalog::ColumnMetadata> = result
291 .columns
292 .iter()
293 .map(|c| crate::catalog::ColumnMetadata::new(&c.name, c.data_type.clone()))
294 .collect();
295 let rows: Vec<Row> = result
296 .rows
297 .into_iter()
298 .enumerate()
299 .map(|(i, values)| Row::new(i as u64, values))
300 .collect();
301 let iter = iterator::VecIterator::new(rows, schema.clone());
302 (Box::new(iter), Projection::All(column_names), schema)
303}
304
305fn execute_recursive_cte_result<
306 'txn,
307 S: KVStore + 'txn,
308 C: Catalog + ?Sized,
309 T: SqlTxn<'txn, S>,
310>(
311 txn: &mut T,
312 catalog: &C,
313 context: &QueryExecutionContext,
314 outer: Option<&Row>,
315 memory: Option<&MemoryPolicy>,
316 execution: RecursiveCteExecution,
317) -> Result<QueryResult> {
318 let RecursiveCteExecution {
319 name,
320 anchor,
321 recursive_term,
322 union_all,
323 schema,
324 limits,
325 } = execution;
326 let anchor_result =
327 execute_query_result_with_context(txn, catalog, anchor, outer, memory, context)?;
328 let mut accumulated = Vec::new();
329 let mut seen = HashSet::new();
330 let mut accumulated_bytes = 0u64;
331 let mut seen_bytes = 0u64;
332 for row in anchor_result.rows {
333 let should_accumulate = if union_all {
334 true
335 } else {
336 let key = aggregate::encode_group_key(&row)?;
337 let key_bytes = key.len() as u64;
338 if seen.insert(key) {
339 seen_bytes = seen_bytes.saturating_add(key_bytes);
340 true
341 } else {
342 false
343 }
344 };
345 if should_accumulate {
346 accumulated_bytes = accumulated_bytes.saturating_add(estimated_row_bytes(&row));
347 accumulated.push(row);
348 }
349 }
350 ensure_recursive_row_limit(&name, accumulated.len(), limits.max_rows)?;
351
352 let mut working = accumulated.clone();
353 let mut working_bytes = accumulated_bytes;
354 enforce_recursive_memory(
355 memory,
356 accumulated_bytes
357 .saturating_add(working_bytes)
358 .saturating_add(seen_bytes),
359 )?;
360 let mut iterations = 0usize;
361 while !working.is_empty() {
362 if iterations >= limits.max_iterations {
363 return Err(ExecutorError::ResourceExhausted {
364 message: format!(
365 "recursive CTE '{name}' reached iteration limit {}",
366 limits.max_iterations
367 ),
368 });
369 }
370 iterations += 1;
371
372 enforce_recursive_memory(
376 memory,
377 accumulated_bytes
378 .saturating_add(working_bytes.saturating_mul(2))
379 .saturating_add(seen_bytes),
380 )?;
381
382 let iteration_context = context.with_recursive_table(
383 name.clone(),
384 RecursiveWorkingTable {
385 rows: working,
386 schema: schema.clone(),
387 },
388 );
389 let recursive_result = execute_query_result_with_context(
390 txn,
391 catalog,
392 recursive_term.clone(),
393 outer,
394 memory,
395 &iteration_context,
396 )?;
397 let recursive_result_bytes = recursive_result
398 .rows
399 .iter()
400 .map(|row| estimated_row_bytes(row))
401 .sum::<u64>();
402 enforce_recursive_memory(
408 memory,
409 accumulated_bytes
410 .saturating_add(working_bytes)
411 .saturating_add(recursive_result_bytes)
412 .saturating_add(seen_bytes),
413 )?;
414 drop(iteration_context);
415 let mut next = Vec::new();
416 let mut next_bytes = 0u64;
417 for row in recursive_result.rows {
418 let should_accumulate = if union_all {
419 true
420 } else {
421 let key = aggregate::encode_group_key(&row)?;
422 let key_bytes = key.len() as u64;
423 if seen.insert(key) {
424 seen_bytes = seen_bytes.saturating_add(key_bytes);
425 true
426 } else {
427 false
428 }
429 };
430 if should_accumulate {
431 next_bytes = next_bytes.saturating_add(estimated_row_bytes(&row));
432 next.push(row);
433 }
434 }
435 ensure_recursive_row_limit(
436 &name,
437 accumulated.len().saturating_add(next.len()),
438 limits.max_rows,
439 )?;
440 accumulated.extend(next.iter().cloned());
441 accumulated_bytes = accumulated_bytes.saturating_add(next_bytes);
442 working_bytes = next_bytes;
443 enforce_recursive_memory(
444 memory,
445 accumulated_bytes
446 .saturating_add(working_bytes)
447 .saturating_add(seen_bytes),
448 )?;
449 working = next;
450 }
451
452 let columns = schema
453 .into_iter()
454 .map(|column| ColumnInfo::new(column.name, column.data_type))
455 .collect();
456 Ok(QueryResult::new(columns, accumulated))
457}
458
459fn estimated_row_bytes(row: &[SqlValue]) -> u64 {
460 row.iter().map(ByteSized::estimated_bytes).sum()
461}
462
463fn enforce_recursive_memory(memory: Option<&MemoryPolicy>, bytes: u64) -> Result<()> {
464 let Some(policy) = memory else {
465 return Ok(());
466 };
467 let mut tracker = MemoryTracker::new(policy.clone());
468 tracker.add_bytes(bytes).map_err(map_core_memory_error)?;
469 if tracker.over_limit() {
470 return Err(ExecutorError::ResourceExhausted {
471 message: format!(
472 "query memory limit exceeded by recursive CTE materialization at {} bytes",
473 tracker.used_bytes()
474 ),
475 });
476 }
477 Ok(())
478}
479
480fn ensure_recursive_row_limit(name: &str, rows: usize, max_rows: usize) -> Result<()> {
481 if rows <= max_rows {
482 return Ok(());
483 }
484 Err(ExecutorError::ResourceExhausted {
485 message: format!("recursive CTE '{name}' reached row limit {max_rows}"),
486 })
487}
488
489fn build_iterator_pipeline<'txn, S: KVStore + 'txn, C: Catalog + ?Sized, T: SqlTxn<'txn, S>>(
496 txn: &mut T,
497 catalog: &C,
498 plan: LogicalPlan,
499 memory: Option<&MemoryPolicy>,
500) -> Result<(
501 Box<dyn RowIterator>,
502 Projection,
503 Vec<crate::catalog::ColumnMetadata>,
504)> {
505 build_iterator_pipeline_with_outer(
506 txn,
507 catalog,
508 plan,
509 memory,
510 None,
511 &QueryExecutionContext::default(),
512 )
513}
514
515fn build_iterator_pipeline_with_outer<
516 'txn,
517 S: KVStore + 'txn,
518 C: Catalog + ?Sized,
519 T: SqlTxn<'txn, S>,
520>(
521 txn: &mut T,
522 catalog: &C,
523 plan: LogicalPlan,
524 memory: Option<&MemoryPolicy>,
525 outer: Option<&Row>,
526 context: &QueryExecutionContext,
527) -> Result<(
528 Box<dyn RowIterator>,
529 Projection,
530 Vec<crate::catalog::ColumnMetadata>,
531)> {
532 match plan {
533 LogicalPlan::RecursiveReference { name, schema } => {
534 let table = context.recursive_tables.get(&name).ok_or_else(|| {
535 ExecutorError::InvalidOperation {
536 operation: "recursive common table expression".into(),
537 reason: format!("working table '{name}' is not active"),
538 }
539 })?;
540 if table.schema.len() != schema.len()
541 || table.schema.iter().zip(&schema).any(|(left, right)| {
542 left.name != right.name || left.data_type != right.data_type
543 })
544 {
545 return Err(ExecutorError::InvalidOperation {
546 operation: "recursive common table expression".into(),
547 reason: format!("working table '{name}' schema changed during evaluation"),
548 });
549 }
550 let rows = table
551 .rows
552 .iter()
553 .cloned()
554 .enumerate()
555 .map(|(index, values)| Row::new(index as u64, values))
556 .collect();
557 let projection =
558 Projection::All(schema.iter().map(|column| column.name.clone()).collect());
559 Ok((
560 Box::new(iterator::VecIterator::new(rows, schema.clone())),
561 projection,
562 schema,
563 ))
564 }
565 LogicalPlan::RecursiveCte {
566 name,
567 anchor,
568 recursive_term,
569 union_all,
570 schema,
571 limits,
572 } => {
573 let result = execute_recursive_cte_result(
574 txn,
575 catalog,
576 context,
577 outer,
578 memory,
579 RecursiveCteExecution {
580 name,
581 anchor: *anchor,
582 recursive_term: *recursive_term,
583 union_all,
584 schema,
585 limits,
586 },
587 )?;
588 Ok(materialize_query_result(result))
589 }
590 LogicalPlan::Scan { table, projection } => {
591 if table == LITERAL_TABLE {
592 let schema = Vec::new();
593 let rows = vec![Row::new(0, Vec::new())];
594 let iter = iterator::VecIterator::new(rows, schema.clone());
595 return Ok((Box::new(iter), projection, schema));
596 }
597 let table_meta = catalog
598 .get_table(&table)
599 .cloned()
600 .ok_or_else(|| ExecutorError::TableNotFound(table.clone()))?;
601
602 if table_meta.storage_options.storage_type == StorageType::Columnar {
603 let columnar_scan = columnar_scan::build_columnar_scan(&table_meta, &projection);
604 let rows = columnar_scan::execute_columnar_scan(txn, &table_meta, &columnar_scan)?;
605 let schema = table_meta.columns.clone();
606 let iter = iterator::VecIterator::new(rows, schema.clone());
607 return Ok((Box::new(iter), projection, schema));
608 }
609
610 let rows = scan::execute_scan(txn, &table_meta)?;
614 let schema = table_meta.columns.clone();
615
616 let iter = iterator::VecIterator::new(rows, schema.clone());
618 Ok((Box::new(iter), projection, schema))
619 }
620 LogicalPlan::Values { rows, schema } => {
621 let projection =
622 Projection::All(schema.iter().map(|column| column.name.clone()).collect());
623 let iterator = iterator::ValuesIterator::new(rows, schema.clone(), outer);
624 Ok((Box::new(iterator), projection, schema))
625 }
626 LogicalPlan::Filter { input, predicate } => {
627 if let LogicalPlan::Scan { table, projection } = input.as_ref()
628 && let Some(table_meta) = catalog.get_table(table)
629 && table_meta.storage_options.storage_type == StorageType::Columnar
630 {
631 let evaluated_here = outer.is_some() || subquery::contains_subquery(&predicate);
639 let projection = projection.clone();
640 let schema = table_meta.columns.clone();
641 let columnar_scan = if evaluated_here {
642 columnar_scan::build_columnar_scan_for_external_filter(
643 table_meta,
644 projection.clone(),
645 &predicate,
646 )
647 } else {
648 columnar_scan::build_columnar_scan_for_filter(
649 table_meta,
650 projection.clone(),
651 &predicate,
652 )
653 };
654 let rows = columnar_scan::execute_columnar_scan(txn, table_meta, &columnar_scan)?;
655 if !evaluated_here {
656 let iter = iterator::VecIterator::new(rows, schema.clone());
657 return Ok((Box::new(iter), projection, schema));
658 }
659 let mut kept = Vec::new();
660 for row in rows {
661 let eval_row = combine_outer_for_eval(&row, outer);
662 if let SqlValue::Boolean(true) = subquery::evaluate_expr_with_subqueries(
663 txn, catalog, &predicate, &eval_row,
664 )? {
665 kept.push(row);
666 }
667 }
668 let iter = iterator::VecIterator::new(kept, schema.clone());
669 return Ok((Box::new(iter), projection, schema));
670 }
671 let (mut input_iter, projection, schema) =
672 build_iterator_pipeline_with_outer(txn, catalog, *input, memory, outer, context)?;
673 if outer.is_some() || subquery::contains_subquery(&predicate) {
674 let mut rows = Vec::new();
675 while let Some(result) = input_iter.next_row() {
676 let row = result?;
677 let eval_row = combine_outer_for_eval(&row, outer);
678 if let SqlValue::Boolean(true) = subquery::evaluate_expr_with_subqueries(
679 txn, catalog, &predicate, &eval_row,
680 )? {
681 rows.push(row);
682 }
683 }
684 let iter = iterator::VecIterator::new(rows, schema.clone());
685 return Ok((Box::new(iter), projection, schema));
686 }
687 let filter_iter = FilterIterator::new(input_iter, predicate);
688 Ok((Box::new(filter_iter), projection, schema))
689 }
690 LogicalPlan::Project { input, projection } => {
691 let input_result =
692 execute_query_result_with_context(txn, catalog, *input, outer, memory, context)?;
693 let (mut input_iter, _input_projection, schema) =
694 materialize_query_result(input_result);
695 let mut rows = Vec::new();
696 while let Some(result) = input_iter.next_row() {
697 rows.push(result?);
698 }
699 let projected =
700 execute_project_with_subqueries(txn, catalog, rows, &projection, &schema, outer)?;
701 let output_schema = projected
702 .columns
703 .iter()
704 .map(|col| crate::catalog::ColumnMetadata::new(&col.name, col.data_type.clone()))
705 .collect::<Vec<_>>();
706 let rows = projected
707 .rows
708 .into_iter()
709 .enumerate()
710 .map(|(idx, values)| Row::new(idx as u64, values))
711 .collect::<Vec<_>>();
712 let output_projection =
713 Projection::All(output_schema.iter().map(|col| col.name.clone()).collect());
714 let iter = iterator::VecIterator::new(rows, output_schema.clone());
715 Ok((Box::new(iter), output_projection, output_schema))
716 }
717 LogicalPlan::Join {
718 left,
719 right,
720 join_type,
721 condition,
722 using: _,
723 } => {
724 let (mut left_iter, _left_projection, left_schema) =
725 build_iterator_pipeline_with_outer(txn, catalog, *left, memory, outer, context)?;
726 let (mut right_iter, _right_projection, right_schema) =
727 build_iterator_pipeline_with_outer(txn, catalog, *right, memory, outer, context)?;
728 let mut left_rows = Vec::new();
729 while let Some(result) = left_iter.next_row() {
730 left_rows.push(result?);
731 }
732 let mut right_rows = Vec::new();
733 while let Some(result) = right_iter.next_row() {
734 right_rows.push(result?);
735 }
736 let left_width = left_schema.len();
737 let right_width = right_schema.len();
738 let rows = join::execute_join_with_widths(
739 left_rows,
740 right_rows,
741 join_type,
742 condition.as_ref(),
743 left_width,
744 right_width,
745 )?;
746 let mut schema = left_schema;
747 schema.extend(right_schema);
748 let projection = Projection::All(schema.iter().map(|col| col.name.clone()).collect());
749 let iter = iterator::VecIterator::new(rows, schema.clone());
750 Ok((Box::new(iter), projection, schema))
751 }
752 LogicalPlan::LateralJoin {
753 left,
754 right,
755 join_type,
756 condition,
757 right_schema,
758 } => {
759 let (mut left_iter, _left_projection, left_schema) =
760 build_iterator_pipeline_with_outer(txn, catalog, *left, memory, outer, context)?;
761 let mut left_rows = Vec::new();
762 while let Some(result) = left_iter.next_row() {
763 left_rows.push(result?);
764 }
765 drop(left_iter);
766 let rows = execute_lateral_join(
767 txn,
768 catalog,
769 left_rows,
770 &right,
771 join_type,
772 condition.as_ref(),
773 right_schema.len(),
774 outer,
775 memory,
776 context,
777 )?;
778 let mut schema = left_schema;
779 schema.extend(right_schema);
780 let projection = Projection::All(schema.iter().map(|col| col.name.clone()).collect());
781 let iter = iterator::VecIterator::new(rows, schema.clone());
782 Ok((Box::new(iter), projection, schema))
783 }
784 LogicalPlan::TableFunction {
785 function,
786 args,
787 schema,
788 } => {
789 let rows = execute_table_function(txn, catalog, function, &args, outer)?;
790 let projection =
791 Projection::All(schema.iter().map(|column| column.name.clone()).collect());
792 let iter = iterator::VecIterator::new(rows, schema.clone());
793 Ok((Box::new(iter), projection, schema))
794 }
795 LogicalPlan::Aggregate {
796 input,
797 group_keys,
798 aggregates,
799 having,
800 projection,
801 grouping_sets,
802 } => {
803 let (input_iter, _projection, input_schema) =
804 build_iterator_pipeline_with_outer(txn, catalog, *input, memory, outer, context)?;
805 let input_iter: Box<dyn RowIterator> = match outer {
811 Some(outer_row) => {
812 let rows = widen_rows_with_outer(input_iter, outer_row)?;
813 Box::new(iterator::VecIterator::new(rows, input_schema))
814 }
815 None => input_iter,
816 };
817 let mut schema = aggregate::build_aggregate_schema(&group_keys, &aggregates);
818 if grouping_sets.is_some() {
819 schema.push(crate::catalog::ColumnMetadata::new(
820 crate::planner::GROUPING_ID_COLUMN,
821 crate::planner::ResolvedType::BigInt,
822 ));
823 }
824 if grouping_sets.is_none()
828 && let Some(policy) = memory
829 && policy.spill_directory().is_some()
830 {
831 if group_keys.is_empty() {
832 let iter = aggregate::StreamingAggregateIterator::new(
833 input_iter,
834 group_keys,
835 aggregates,
836 having,
837 schema.clone(),
838 );
839 return Ok((Box::new(iter), projection, schema));
840 }
841 let order_by = group_keys
842 .iter()
843 .cloned()
844 .map(|expr| SortExpr {
845 expr,
846 asc: true,
847 nulls_first: false,
848 })
849 .collect::<Vec<_>>();
850 let sort_iter =
851 SortIterator::new_with_policy(input_iter, &order_by, Some(policy.clone()))?;
852 let iter = aggregate::StreamingAggregateIterator::new(
853 Box::new(sort_iter),
854 group_keys,
855 aggregates,
856 having,
857 schema.clone(),
858 );
859 return Ok((Box::new(iter), projection, schema));
860 }
861
862 let parallelism = std::thread::available_parallelism()
863 .map(usize::from)
864 .unwrap_or(1);
865 if grouping_sets.is_none()
866 && !aggregate::should_use_single_for_parallel(parallelism, &aggregates)
867 {
868 let rows = aggregate::execute_parallel_aggregate_rows_with_policy(
869 input_iter,
870 group_keys,
871 aggregates,
872 having,
873 schema.clone(),
874 parallelism,
875 memory.cloned(),
876 1_000_000,
877 )?;
878 let iter = iterator::VecIterator::new(rows, schema.clone());
879 return Ok((Box::new(iter), projection, schema));
880 }
881
882 let mut iter = aggregate::AggregateIterator::new(
883 input_iter,
884 group_keys,
885 aggregates,
886 having,
887 schema.clone(),
888 )
889 .with_grouping_sets(grouping_sets);
890 if let Some(policy) = memory {
891 iter = iter.with_memory_policy(Some(policy.clone()));
892 }
893 Ok((Box::new(iter), projection, schema))
894 }
895 LogicalPlan::Window { input, windows } => {
896 let (input_iter, _projection, input_schema) =
897 build_iterator_pipeline_with_outer(txn, catalog, *input, memory, outer, context)?;
898 if let Some(outer_row) = outer {
904 let inner_width = input_schema.len();
905 let outer_width = outer_row.len();
906 let mut wide_schema = input_schema.clone();
907 wide_schema.extend((0..outer_width).map(|index| {
908 crate::catalog::ColumnMetadata::new(
909 format!("__outer_{index}"),
910 crate::planner::ResolvedType::Text,
911 )
912 }));
913 let rows = widen_rows_with_outer(input_iter, outer_row)?;
914 let wide_iter = iterator::VecIterator::new(rows, wide_schema.clone());
915 let window_iter = window::WindowIterator::new(wide_iter, windows, memory)?;
916 let mut schema = input_schema;
917 schema.extend_from_slice(&window_iter.schema()[wide_schema.len()..]);
918 let rows = drain_rows(Box::new(window_iter))?
919 .into_iter()
920 .map(|mut row| {
921 let end = (inner_width + outer_width).min(row.values.len());
922 row.values.drain(inner_width.min(end)..end);
923 row
924 })
925 .collect::<Vec<_>>();
926 let projection =
927 Projection::All(schema.iter().map(|column| column.name.clone()).collect());
928 let iter = iterator::VecIterator::new(rows, schema.clone());
929 return Ok((Box::new(iter), projection, schema));
930 }
931 let iter = window::WindowIterator::new(input_iter, windows, memory)?;
932 let schema = iter.schema().to_vec();
933 let projection =
934 Projection::All(schema.iter().map(|column| column.name.clone()).collect());
935 Ok((Box::new(iter), projection, schema))
936 }
937 LogicalPlan::SetOperation {
938 left,
939 right,
940 operator,
941 all,
942 } => {
943 let left_result =
944 execute_query_result_with_context(txn, catalog, *left, outer, memory, context)?;
945 let right_result =
946 execute_query_result_with_context(txn, catalog, *right, outer, memory, context)?;
947 let rows = execute_set_operation(operator, all, left_result.rows, right_result.rows)?;
948 let schema = left_result
949 .columns
950 .iter()
951 .map(|column| {
952 crate::catalog::ColumnMetadata::new(&column.name, column.data_type.clone())
953 })
954 .collect::<Vec<_>>();
955 let projection = Projection::All(
956 left_result
957 .columns
958 .iter()
959 .map(|column| column.name.clone())
960 .collect(),
961 );
962 let rows = rows
963 .into_iter()
964 .enumerate()
965 .map(|(index, values)| Row::new(index as u64, values))
966 .collect();
967 Ok((
968 Box::new(iterator::VecIterator::new(rows, schema.clone())),
969 projection,
970 schema,
971 ))
972 }
973 LogicalPlan::Sort { input, order_by } => {
974 let (input_iter, projection, schema) =
975 build_iterator_pipeline_with_outer(txn, catalog, *input, memory, outer, context)?;
976 if let Some(outer_row) = outer {
980 let rows = widen_rows_with_outer(input_iter, outer_row)?;
981 let wide_iter = iterator::VecIterator::new(rows, schema.clone());
982 let sort_iter = if let Some(policy) = memory {
983 SortIterator::new_with_policy(wide_iter, &order_by, Some(policy.clone()))?
984 } else {
985 SortIterator::new(wide_iter, &order_by)?
986 };
987 let rows =
988 narrow_rows_from_outer(drain_rows(Box::new(sort_iter))?, outer_row.len());
989 let iter = iterator::VecIterator::new(rows, schema.clone());
990 return Ok((Box::new(iter), projection, schema));
991 }
992 let sort_iter = if let Some(policy) = memory {
993 SortIterator::new_with_policy(input_iter, &order_by, Some(policy.clone()))?
994 } else {
995 SortIterator::new(input_iter, &order_by)?
996 };
997 Ok((Box::new(sort_iter), projection, schema))
998 }
999 LogicalPlan::DistinctOn {
1000 input,
1001 key_count,
1002 order_by,
1003 } => {
1004 let (input_iter, projection, schema) =
1005 build_iterator_pipeline_with_outer(txn, catalog, *input, memory, outer, context)?;
1006 if let Some(outer_row) = outer {
1007 let rows = widen_rows_with_outer(input_iter, outer_row)?;
1008 let wide_iter = iterator::VecIterator::new(rows, schema.clone());
1009 let distinct_iter =
1010 DistinctOnIterator::new(wide_iter, &order_by, key_count, memory.cloned())?;
1011 let rows =
1012 narrow_rows_from_outer(drain_rows(Box::new(distinct_iter))?, outer_row.len());
1013 let iter = iterator::VecIterator::new(rows, schema.clone());
1014 return Ok((Box::new(iter), projection, schema));
1015 }
1016 let distinct_iter =
1017 DistinctOnIterator::new(input_iter, &order_by, key_count, memory.cloned())?;
1018 Ok((Box::new(distinct_iter), projection, schema))
1019 }
1020 LogicalPlan::Limit {
1021 input,
1022 limit,
1023 offset,
1024 ties,
1025 } => {
1026 let (input_iter, projection, schema) =
1027 build_iterator_pipeline_with_outer(txn, catalog, *input, memory, outer, context)?;
1028 if let (Some(outer_row), Some(tie_keys)) = (outer, ties.as_ref()) {
1031 let rows = widen_rows_with_outer(input_iter, outer_row)?;
1032 let wide_iter = iterator::VecIterator::new(rows, schema.clone());
1033 let limit_iter =
1034 LimitIterator::with_ties(wide_iter, limit, offset, tie_keys.clone());
1035 let rows =
1036 narrow_rows_from_outer(drain_rows(Box::new(limit_iter))?, outer_row.len());
1037 let iter = iterator::VecIterator::new(rows, schema.clone());
1038 return Ok((Box::new(iter), projection, schema));
1039 }
1040 let limit_iter = match ties {
1041 Some(tie_keys) => LimitIterator::with_ties(input_iter, limit, offset, tie_keys),
1042 None => LimitIterator::new(input_iter, limit, offset),
1043 };
1044 Ok((Box::new(limit_iter), projection, schema))
1045 }
1046 other => Err(ExecutorError::UnsupportedOperation(format!(
1047 "unsupported query plan: {other:?}"
1048 ))),
1049 }
1050}
1051
1052pub fn build_streaming_pipeline<
1064 'a,
1065 'txn: 'a,
1066 S: KVStore + 'txn,
1067 C: Catalog + ?Sized,
1068 T: SqlTxn<'txn, S>,
1069>(
1070 txn: &'a mut T,
1071 catalog: &C,
1072 plan: LogicalPlan,
1073) -> Result<(
1074 Box<dyn RowIterator + 'a>,
1075 Projection,
1076 Vec<crate::catalog::ColumnMetadata>,
1077)> {
1078 build_streaming_pipeline_with_policy(txn, catalog, plan, None)
1079}
1080
1081pub fn build_streaming_pipeline_with_policy<
1082 'a,
1083 'txn: 'a,
1084 S: KVStore + 'txn,
1085 C: Catalog + ?Sized,
1086 T: SqlTxn<'txn, S>,
1087>(
1088 txn: &'a mut T,
1089 catalog: &C,
1090 plan: LogicalPlan,
1091 memory: Option<&MemoryPolicy>,
1092) -> Result<(
1093 Box<dyn RowIterator + 'a>,
1094 Projection,
1095 Vec<crate::catalog::ColumnMetadata>,
1096)> {
1097 if subquery::plan_contains_subquery(&plan)
1103 || plan_contains_recursive_cte(&plan)
1104 || plan_contains_lateral(&plan)
1105 {
1106 let result = execute_query_result_with_outer_and_policy(txn, catalog, plan, None, memory)?;
1107 return Ok(materialize_query_result(result));
1108 }
1109
1110 match plan {
1111 LogicalPlan::SetOperation {
1112 left,
1113 right,
1114 operator,
1115 all,
1116 } => build_streaming_set_operation(txn, catalog, *left, *right, operator, all, memory),
1117 other => build_streaming_pipeline_inner(txn, catalog, other, memory),
1118 }
1119}
1120
1121fn build_streaming_set_operation<
1122 'a,
1123 'txn: 'a,
1124 S: KVStore + 'txn,
1125 C: Catalog + ?Sized,
1126 T: SqlTxn<'txn, S>,
1127>(
1128 txn: &'a mut T,
1129 catalog: &C,
1130 left: LogicalPlan,
1131 right: LogicalPlan,
1132 operator: SetOperator,
1133 all: bool,
1134 memory: Option<&MemoryPolicy>,
1135) -> Result<(
1136 Box<dyn RowIterator + 'a>,
1137 Projection,
1138 Vec<crate::catalog::ColumnMetadata>,
1139)> {
1140 let (mut left_iter, left_projection, left_schema) =
1141 build_streaming_pipeline_with_policy(txn, catalog, left, memory)?;
1142 let mut left_rows = Vec::new();
1143 while let Some(result) = left_iter.next_row() {
1144 left_rows.push(result?);
1145 }
1146 drop(left_iter);
1147 let left_result = project::execute_project(left_rows, &left_projection, &left_schema)?;
1148
1149 let (mut right_iter, right_projection, right_schema) =
1150 build_streaming_pipeline_with_policy(txn, catalog, right, memory)?;
1151 let mut right_rows = Vec::new();
1152 while let Some(result) = right_iter.next_row() {
1153 right_rows.push(result?);
1154 }
1155 let right_result = project::execute_project(right_rows, &right_projection, &right_schema)?;
1156
1157 let rows = execute_set_operation(operator, all, left_result.rows, right_result.rows)?;
1158 let schema = left_result
1159 .columns
1160 .iter()
1161 .map(|column| crate::catalog::ColumnMetadata::new(&column.name, column.data_type.clone()))
1162 .collect::<Vec<_>>();
1163 let projection = Projection::All(
1164 left_result
1165 .columns
1166 .iter()
1167 .map(|column| column.name.clone())
1168 .collect(),
1169 );
1170 let rows = rows
1171 .into_iter()
1172 .enumerate()
1173 .map(|(index, values)| Row::new(index as u64, values))
1174 .collect();
1175 Ok((
1176 Box::new(iterator::VecIterator::new(rows, schema.clone())),
1177 projection,
1178 schema,
1179 ))
1180}
1181
1182fn execute_set_operation(
1183 operator: SetOperator,
1184 all: bool,
1185 mut left: Vec<Vec<SqlValue>>,
1186 right: Vec<Vec<SqlValue>>,
1187) -> Result<Vec<Vec<SqlValue>>> {
1188 if all {
1189 match operator {
1190 SetOperator::Union => {
1191 left.extend(right);
1192 return Ok(left);
1193 }
1194 SetOperator::Intersect | SetOperator::Except => {
1195 let mut right_counts = HashMap::<Vec<u8>, usize>::new();
1196 for row in right {
1197 *right_counts
1198 .entry(aggregate::encode_group_key(&row)?)
1199 .or_default() += 1;
1200 }
1201 let mut output = Vec::new();
1202 for row in left {
1203 let key = aggregate::encode_group_key(&row)?;
1204 let remaining = right_counts.entry(key).or_default();
1205 if *remaining > 0 {
1206 *remaining -= 1;
1207 if operator == SetOperator::Intersect {
1208 output.push(row);
1209 }
1210 } else if operator == SetOperator::Except {
1211 output.push(row);
1212 }
1213 }
1214 return Ok(output);
1215 }
1216 }
1217 }
1218
1219 let right_keys = right
1220 .iter()
1221 .map(|row| aggregate::encode_group_key(row))
1222 .collect::<Result<HashSet<_>>>()?;
1223 let mut seen = HashSet::new();
1224 let mut output = Vec::new();
1225 match operator {
1226 SetOperator::Union => {
1227 for row in left.into_iter().chain(right) {
1228 if seen.insert(aggregate::encode_group_key(&row)?) {
1229 output.push(row);
1230 }
1231 }
1232 }
1233 SetOperator::Intersect => {
1234 for row in left {
1235 let key = aggregate::encode_group_key(&row)?;
1236 if right_keys.contains(&key) && seen.insert(key) {
1237 output.push(row);
1238 }
1239 }
1240 }
1241 SetOperator::Except => {
1242 for row in left {
1243 let key = aggregate::encode_group_key(&row)?;
1244 if !right_keys.contains(&key) && seen.insert(key) {
1245 output.push(row);
1246 }
1247 }
1248 }
1249 }
1250 Ok(output)
1251}
1252
1253fn build_streaming_pipeline_inner<
1255 'a,
1256 'txn: 'a,
1257 S: KVStore + 'txn,
1258 C: Catalog + ?Sized,
1259 T: SqlTxn<'txn, S>,
1260>(
1261 txn: &'a mut T,
1262 catalog: &C,
1263 plan: LogicalPlan,
1264 memory: Option<&MemoryPolicy>,
1265) -> Result<(
1266 Box<dyn RowIterator + 'a>,
1267 Projection,
1268 Vec<crate::catalog::ColumnMetadata>,
1269)> {
1270 match plan {
1271 LogicalPlan::Scan { table, projection } => {
1272 if table == LITERAL_TABLE {
1273 let schema = Vec::new();
1274 let rows = vec![Row::new(0, Vec::new())];
1275 let iter = iterator::VecIterator::new(rows, schema.clone());
1276 return Ok((Box::new(iter), projection, schema));
1277 }
1278 let table_meta = catalog
1279 .get_table(&table)
1280 .cloned()
1281 .ok_or_else(|| ExecutorError::TableNotFound(table.clone()))?;
1282
1283 if table_meta.storage_options.storage_type == StorageType::Columnar {
1284 let columnar_scan = columnar_scan::build_columnar_scan(&table_meta, &projection);
1286 let schema = table_meta.columns.clone();
1287 let iter =
1288 columnar_scan::create_columnar_scan_iterator(txn, &table_meta, &columnar_scan)?;
1289 return Ok((Box::new(iter), projection, schema));
1290 }
1291
1292 let schema = table_meta.columns.clone();
1294 let scan_iter = scan::create_scan_iterator(txn, &table_meta)?;
1295 Ok((Box::new(scan_iter), projection, schema))
1296 }
1297 LogicalPlan::Values { rows, schema } => {
1298 let projection =
1299 Projection::All(schema.iter().map(|column| column.name.clone()).collect());
1300 let iterator = iterator::ValuesIterator::new(rows, schema.clone(), None);
1301 Ok((Box::new(iterator), projection, schema))
1302 }
1303 LogicalPlan::Filter { input, predicate } => {
1304 if let LogicalPlan::Scan { table, projection } = input.as_ref()
1305 && let Some(table_meta) = catalog.get_table(table)
1306 && table_meta.storage_options.storage_type == StorageType::Columnar
1307 {
1308 let columnar_scan = columnar_scan::build_columnar_scan_for_filter(
1310 table_meta,
1311 projection.clone(),
1312 &predicate,
1313 );
1314 let schema = table_meta.columns.clone();
1315 let iter =
1316 columnar_scan::create_columnar_scan_iterator(txn, table_meta, &columnar_scan)?;
1317 return Ok((Box::new(iter), projection.clone(), schema));
1318 }
1319 let (input_iter, projection, schema) =
1320 build_streaming_pipeline_with_policy(txn, catalog, *input, memory)?;
1321 let filter_iter = FilterIterator::new(input_iter, predicate);
1322 Ok((Box::new(filter_iter), projection, schema))
1323 }
1324 LogicalPlan::Project { input, projection } => {
1325 let (mut input_iter, _input_projection, schema) =
1326 build_streaming_pipeline_with_policy(txn, catalog, *input, memory)?;
1327 let mut rows = Vec::new();
1328 while let Some(result) = input_iter.next_row() {
1329 rows.push(result?);
1330 }
1331 let projected = project::execute_project(rows, &projection, &schema)?;
1332 let output_schema = projected
1333 .columns
1334 .iter()
1335 .map(|col| crate::catalog::ColumnMetadata::new(&col.name, col.data_type.clone()))
1336 .collect::<Vec<_>>();
1337 let rows = projected
1338 .rows
1339 .into_iter()
1340 .enumerate()
1341 .map(|(idx, values)| Row::new(idx as u64, values))
1342 .collect::<Vec<_>>();
1343 let output_projection =
1344 Projection::All(output_schema.iter().map(|col| col.name.clone()).collect());
1345 let iter = iterator::VecIterator::new(rows, output_schema.clone());
1346 Ok((Box::new(iter), output_projection, output_schema))
1347 }
1348 LogicalPlan::Join {
1349 left,
1350 right,
1351 join_type,
1352 condition,
1353 using: _,
1354 } => {
1355 let (mut left_iter, _left_projection, left_schema) =
1356 build_streaming_pipeline_with_policy(txn, catalog, *left, memory)?;
1357 let mut left_rows = Vec::new();
1358 while let Some(result) = left_iter.next_row() {
1359 left_rows.push(result?);
1360 }
1361 drop(left_iter);
1362 let (mut right_iter, _right_projection, right_schema) =
1363 build_streaming_pipeline_with_policy(txn, catalog, *right, memory)?;
1364 let mut right_rows = Vec::new();
1365 while let Some(result) = right_iter.next_row() {
1366 right_rows.push(result?);
1367 }
1368 let rows = join::execute_join_with_widths(
1369 left_rows,
1370 right_rows,
1371 join_type,
1372 condition.as_ref(),
1373 left_schema.len(),
1374 right_schema.len(),
1375 )?;
1376 let mut schema = left_schema;
1377 schema.extend(right_schema);
1378 let projection = Projection::All(schema.iter().map(|col| col.name.clone()).collect());
1379 let iter = iterator::VecIterator::new(rows, schema.clone());
1380 Ok((Box::new(iter), projection, schema))
1381 }
1382 LogicalPlan::Aggregate {
1383 input,
1384 group_keys,
1385 aggregates,
1386 having,
1387 projection,
1388 grouping_sets,
1389 } => {
1390 let (input_iter, _projection, _schema) =
1391 build_streaming_pipeline_with_policy(txn, catalog, *input, memory)?;
1392 let mut schema = aggregate::build_aggregate_schema(&group_keys, &aggregates);
1393 if grouping_sets.is_some() {
1394 schema.push(crate::catalog::ColumnMetadata::new(
1395 crate::planner::GROUPING_ID_COLUMN,
1396 crate::planner::ResolvedType::BigInt,
1397 ));
1398 }
1399 if grouping_sets.is_none()
1401 && let Some(policy) = memory
1402 && policy.spill_directory().is_some()
1403 {
1404 if group_keys.is_empty() {
1405 let iter = aggregate::StreamingAggregateIterator::new(
1406 input_iter,
1407 group_keys,
1408 aggregates,
1409 having,
1410 schema.clone(),
1411 );
1412 return Ok((Box::new(iter), projection, schema));
1413 }
1414 let order_by = group_keys
1415 .iter()
1416 .cloned()
1417 .map(|expr| SortExpr {
1418 expr,
1419 asc: true,
1420 nulls_first: false,
1421 })
1422 .collect::<Vec<_>>();
1423 let sort_iter =
1424 SortIterator::new_with_policy(input_iter, &order_by, Some(policy.clone()))?;
1425 let iter = aggregate::StreamingAggregateIterator::new(
1426 Box::new(sort_iter),
1427 group_keys,
1428 aggregates,
1429 having,
1430 schema.clone(),
1431 );
1432 return Ok((Box::new(iter), projection, schema));
1433 }
1434
1435 let parallelism = std::thread::available_parallelism()
1436 .map(usize::from)
1437 .unwrap_or(1);
1438 if grouping_sets.is_none()
1439 && !aggregate::should_use_single_for_parallel(parallelism, &aggregates)
1440 {
1441 let rows = aggregate::execute_parallel_aggregate_rows_with_policy(
1442 input_iter,
1443 group_keys,
1444 aggregates,
1445 having,
1446 schema.clone(),
1447 parallelism,
1448 memory.cloned(),
1449 1_000_000,
1450 )?;
1451 let iter = iterator::VecIterator::new(rows, schema.clone());
1452 return Ok((Box::new(iter), projection, schema));
1453 }
1454
1455 let mut iter = aggregate::AggregateIterator::new(
1456 input_iter,
1457 group_keys,
1458 aggregates,
1459 having,
1460 schema.clone(),
1461 )
1462 .with_grouping_sets(grouping_sets);
1463 if let Some(policy) = memory {
1464 iter = iter.with_memory_policy(Some(policy.clone()));
1465 }
1466 Ok((Box::new(iter), projection, schema))
1467 }
1468 LogicalPlan::Window { input, windows } => {
1469 let (input_iter, _projection, _schema) =
1470 build_streaming_pipeline_inner(txn, catalog, *input, memory)?;
1471 let iter = window::WindowIterator::new(input_iter, windows, memory)?;
1472 let schema = iter.schema().to_vec();
1473 let projection =
1474 Projection::All(schema.iter().map(|column| column.name.clone()).collect());
1475 Ok((Box::new(iter), projection, schema))
1476 }
1477 LogicalPlan::Sort { input, order_by } => {
1478 let (input_iter, projection, schema) =
1479 build_streaming_pipeline_with_policy(txn, catalog, *input, memory)?;
1480 let sort_iter = if let Some(policy) = memory {
1481 SortIterator::new_with_policy(input_iter, &order_by, Some(policy.clone()))?
1482 } else {
1483 SortIterator::new(input_iter, &order_by)?
1484 };
1485 Ok((Box::new(sort_iter), projection, schema))
1486 }
1487 LogicalPlan::DistinctOn {
1488 input,
1489 key_count,
1490 order_by,
1491 } => {
1492 let (input_iter, projection, schema) =
1493 build_streaming_pipeline_with_policy(txn, catalog, *input, memory)?;
1494 let distinct_iter =
1495 DistinctOnIterator::new(input_iter, &order_by, key_count, memory.cloned())?;
1496 Ok((Box::new(distinct_iter), projection, schema))
1497 }
1498 LogicalPlan::Limit {
1499 input,
1500 limit,
1501 offset,
1502 ties,
1503 } => {
1504 let (input_iter, projection, schema) =
1505 build_streaming_pipeline_with_policy(txn, catalog, *input, memory)?;
1506 let limit_iter = match ties {
1507 Some(tie_keys) => LimitIterator::with_ties(input_iter, limit, offset, tie_keys),
1508 None => LimitIterator::new(input_iter, limit, offset),
1509 };
1510 Ok((Box::new(limit_iter), projection, schema))
1511 }
1512 other => Err(ExecutorError::UnsupportedOperation(format!(
1513 "unsupported query plan: {other:?}"
1514 ))),
1515 }
1516}
1517
1518#[allow(clippy::too_many_arguments)]
1525fn execute_lateral_join<'txn, S: KVStore + 'txn, C: Catalog + ?Sized, T: SqlTxn<'txn, S>>(
1526 txn: &mut T,
1527 catalog: &C,
1528 left_rows: Vec<Row>,
1529 right: &LogicalPlan,
1530 join_type: crate::planner::JoinType,
1531 condition: Option<&crate::planner::typed_expr::TypedExpr>,
1532 right_width: usize,
1533 outer: Option<&Row>,
1534 memory: Option<&MemoryPolicy>,
1535 context: &QueryExecutionContext,
1536) -> Result<Vec<Row>> {
1537 use crate::planner::JoinType;
1538
1539 let mut rows = Vec::new();
1540 let mut row_id = 0u64;
1541 for left_row in left_rows {
1542 let eval_outer = combine_outer_for_eval(&left_row, outer);
1545 let right_result = execute_query_result_with_context(
1546 txn,
1547 catalog,
1548 right.clone(),
1549 Some(&eval_outer),
1550 memory,
1551 context,
1552 )?;
1553 let mut matched = false;
1554 for right_values in right_result.rows {
1555 let mut values = left_row.values.clone();
1556 values.extend(right_values);
1557 let combined = Row::new(row_id, values);
1558 if let Some(condition) = condition {
1559 let eval_row = combine_outer_for_eval(&combined, outer);
1560 let keep =
1561 subquery::evaluate_expr_with_subqueries(txn, catalog, condition, &eval_row)?;
1562 if !matches!(keep, SqlValue::Boolean(true)) {
1563 continue;
1564 }
1565 }
1566 matched = true;
1567 rows.push(combined);
1568 row_id += 1;
1569 }
1570 if !matched && matches!(join_type, JoinType::Left) {
1571 let mut values = left_row.values.clone();
1572 values.extend(std::iter::repeat_n(SqlValue::Null, right_width));
1573 rows.push(Row::new(row_id, values));
1574 row_id += 1;
1575 }
1576 }
1577 Ok(rows)
1578}
1579
1580fn execute_table_function<'txn, S: KVStore + 'txn, C: Catalog + ?Sized, T: SqlTxn<'txn, S>>(
1585 txn: &mut T,
1586 catalog: &C,
1587 function: crate::planner::TableFunctionKind,
1588 args: &[crate::planner::typed_expr::TypedExpr],
1589 outer: Option<&Row>,
1590) -> Result<Vec<Row>> {
1591 use crate::planner::TableFunctionKind;
1592
1593 let empty = Row::new(0, Vec::new());
1594 let eval_row = combine_outer_for_eval(&empty, outer);
1595 let mut values = Vec::with_capacity(args.len());
1596 for arg in args {
1597 values.push(subquery::evaluate_expr_with_subqueries(
1598 txn, catalog, arg, &eval_row,
1599 )?);
1600 }
1601
1602 match function {
1603 TableFunctionKind::Unnest | TableFunctionKind::UnnestWithOrdinality => {
1604 let Some(argument) = values.into_iter().next() else {
1605 return Err(ExecutorError::InvalidOperation {
1606 operation: "UNNEST".into(),
1607 reason: "UNNEST requires exactly one argument".into(),
1608 });
1609 };
1610 match argument {
1611 SqlValue::Null => Ok(Vec::new()),
1613 SqlValue::Vector(elements) => Ok(elements
1614 .into_iter()
1615 .enumerate()
1616 .map(|(index, element)| {
1617 let mut values = vec![SqlValue::Float(element)];
1618 if function == TableFunctionKind::UnnestWithOrdinality {
1619 values.push(SqlValue::BigInt(index as i64 + 1));
1620 }
1621 Row::new(index as u64, values)
1622 })
1623 .collect()),
1624 SqlValue::Array(elements) => Ok(elements
1625 .into_iter()
1626 .enumerate()
1627 .map(|(index, element)| {
1628 let mut values = vec![element];
1629 if function == TableFunctionKind::UnnestWithOrdinality {
1630 values.push(SqlValue::BigInt(index as i64 + 1));
1631 }
1632 Row::new(index as u64, values)
1633 })
1634 .collect()),
1635 other => Err(ExecutorError::InvalidOperation {
1636 operation: "UNNEST".into(),
1637 reason: format!(
1638 "UNNEST requires an ARRAY or VECTOR argument, found {}",
1639 other.type_name()
1640 ),
1641 }),
1642 }
1643 }
1644 TableFunctionKind::GenerateSeries => {
1645 if matches!(values.first(), Some(SqlValue::Timestamp(_)))
1646 || matches!(values.get(1), Some(SqlValue::Timestamp(_)))
1647 || matches!(values.get(2), Some(SqlValue::Interval { .. }))
1648 {
1649 generate_timestamp_series(&values)
1650 } else {
1651 generate_integer_series(&values)
1652 }
1653 }
1654 TableFunctionKind::JsonEach | TableFunctionKind::JsonTree => {
1655 crate::executor::evaluator::json::table_rows(function.name(), &values).map(|rows| {
1656 rows.into_iter()
1657 .enumerate()
1658 .map(|(index, values)| Row::new(index as u64, values))
1659 .collect()
1660 })
1661 }
1662 TableFunctionKind::FtsSearch => execute_fts_search(txn, catalog, &values),
1663 }
1664}
1665
1666fn execute_fts_search<'txn, S: KVStore + 'txn, C: Catalog + ?Sized, T: SqlTxn<'txn, S>>(
1667 txn: &mut T,
1668 catalog: &C,
1669 values: &[SqlValue],
1670) -> Result<Vec<Row>> {
1671 let argument = |index: usize| match values.get(index) {
1672 Some(SqlValue::Text(value)) => Ok(value.as_str()),
1673 Some(value) => Err(ExecutorError::InvalidOperation {
1674 operation: "FTS_SEARCH".into(),
1675 reason: format!("expected TEXT, found {}", value.type_name()),
1676 }),
1677 None => Err(ExecutorError::InvalidOperation {
1678 operation: "FTS_SEARCH".into(),
1679 reason: "missing argument".into(),
1680 }),
1681 };
1682 let table_name = argument(0)?;
1683 let column_name = argument(1)?;
1684 let query_text = argument(2)?;
1685 let config = if values.len() == 4 {
1686 argument(3)?
1687 } else {
1688 "simple"
1689 };
1690 let table = catalog
1691 .get_table(table_name)
1692 .cloned()
1693 .ok_or_else(|| ExecutorError::TableNotFound(table_name.into()))?;
1694 let column = table
1695 .get_column_index(column_name)
1696 .ok_or_else(|| ExecutorError::ColumnNotFound(column_name.into()))?;
1697 if table.columns[column].data_type != crate::planner::ResolvedType::Text {
1698 return Err(ExecutorError::InvalidOperation {
1699 operation: "FTS_SEARCH".into(),
1700 reason: "the searched column must be TEXT".into(),
1701 });
1702 }
1703 let index = catalog
1704 .get_indexes_for_table(table_name)
1705 .into_iter()
1706 .find(|index| {
1707 matches!(index.method, Some(crate::ast::ddl::IndexMethod::Fts))
1708 && index.column_indices == [column]
1709 && crate::executor::fts_bridge::config(index).eq_ignore_ascii_case(config)
1710 });
1711 if index.is_some_and(|index| {
1712 index.get_option("fts_format_version") != Some(crate::fts::INDEX_FORMAT_VERSION)
1713 }) {
1714 return Err(ExecutorError::InvalidOperation {
1715 operation: "FTS_SEARCH".into(),
1716 reason: "unsupported or missing FTS index format version".into(),
1717 });
1718 }
1719 let query = crate::fts::parse_tsquery(config, query_text).map_err(|reason| {
1720 ExecutorError::InvalidOperation {
1721 operation: "FTS_SEARCH".into(),
1722 reason,
1723 }
1724 })?;
1725
1726 let candidates = if let Some(index) = index {
1727 let mut terms = std::collections::BTreeSet::new();
1728 if crate::fts::index_terms(&query, &mut terms) {
1729 let mut ids = std::collections::BTreeSet::new();
1730 let mut storage = txn.index_storage(index.index_id, false, vec![column]);
1731 for term in terms {
1732 ids.extend(storage.lookup(&SqlValue::Text(term))?);
1733 }
1734 Some(ids)
1735 } else {
1736 None
1737 }
1738 } else {
1739 None
1740 };
1741
1742 let mut output = Vec::new();
1743 let mut start_row_id = 0u64;
1744 loop {
1745 let rows = txn.with_table(&table, |storage| {
1746 storage
1747 .range_scan(start_row_id, u64::MAX)?
1748 .take(2048)
1749 .collect::<std::result::Result<Vec<_>, _>>()
1750 })?;
1751 if rows.is_empty() {
1752 break;
1753 }
1754 for (row_id, row) in rows {
1755 start_row_id = row_id.saturating_add(1);
1756 if candidates
1757 .as_ref()
1758 .is_some_and(|ids| !ids.contains(&row_id))
1759 {
1760 continue;
1761 }
1762 let SqlValue::Text(document) = &row[column] else {
1763 continue;
1764 };
1765 let tokens = crate::fts::tokenize(config, document).map_err(|reason| {
1766 ExecutorError::InvalidOperation {
1767 operation: "FTS_SEARCH".into(),
1768 reason,
1769 }
1770 })?;
1771 if !crate::fts::matches_query(&tokens, &query) {
1772 continue;
1773 }
1774 let row_id = i64::try_from(row_id).map_err(|_| ExecutorError::InvalidOperation {
1775 operation: "FTS_SEARCH".into(),
1776 reason: "row id exceeds BIGINT".into(),
1777 })?;
1778 let headline = crate::fts::headline(config, document, &query).map_err(|reason| {
1779 ExecutorError::InvalidOperation {
1780 operation: "FTS_SEARCH".into(),
1781 reason,
1782 }
1783 })?;
1784 if output.len() == 100_000 {
1785 return Err(ExecutorError::InvalidOperation {
1786 operation: "FTS_SEARCH".into(),
1787 reason: "result exceeds 100000 rows".into(),
1788 });
1789 }
1790 output.push(Row::new(
1791 row_id as u64,
1792 vec![
1793 SqlValue::BigInt(row_id),
1794 SqlValue::Text(document.clone()),
1795 SqlValue::Double(crate::fts::rank(&tokens, &query)),
1796 SqlValue::Text(headline),
1797 ],
1798 ));
1799 }
1800 }
1801 Ok(output)
1802}
1803
1804const MAX_GENERATED_SERIES_ROWS: usize = 100_000;
1805
1806fn generate_integer_series(values: &[SqlValue]) -> Result<Vec<Row>> {
1807 if values.iter().any(SqlValue::is_null) {
1808 return Ok(Vec::new());
1809 }
1810 let integer = |index: usize| match values.get(index) {
1811 Some(SqlValue::Integer(value)) => Ok(i64::from(*value)),
1812 Some(SqlValue::BigInt(value)) => Ok(*value),
1813 Some(value) => Err(ExecutorError::InvalidOperation {
1814 operation: "GENERATE_SERIES".into(),
1815 reason: format!("expected INTEGER, found {}", value.type_name()),
1816 }),
1817 None => Err(ExecutorError::InvalidOperation {
1818 operation: "GENERATE_SERIES".into(),
1819 reason: "missing argument".into(),
1820 }),
1821 };
1822 let start = integer(0)?;
1823 let stop = integer(1)?;
1824 let step = if values.len() == 3 { integer(2)? } else { 1 };
1825 if step == 0 {
1826 return Err(ExecutorError::InvalidOperation {
1827 operation: "GENERATE_SERIES".into(),
1828 reason: "step must not be zero".into(),
1829 });
1830 }
1831 if (step > 0 && start > stop) || (step < 0 && start < stop) {
1832 return Ok(Vec::new());
1833 }
1834
1835 let use_bigint = values
1836 .iter()
1837 .any(|value| matches!(value, SqlValue::BigInt(_)));
1838 let mut rows = Vec::new();
1839 let mut current = start;
1840 loop {
1841 if rows.len() == MAX_GENERATED_SERIES_ROWS {
1842 return Err(ExecutorError::ResourceExhausted {
1843 message: format!("GENERATE_SERIES is limited to {MAX_GENERATED_SERIES_ROWS} rows"),
1844 });
1845 }
1846 let value = if use_bigint {
1847 SqlValue::BigInt(current)
1848 } else {
1849 SqlValue::Integer(i32::try_from(current).map_err(|_| {
1850 ExecutorError::Evaluation(crate::executor::EvaluationError::Overflow)
1851 })?)
1852 };
1853 rows.push(Row::new(rows.len() as u64, vec![value]));
1854 if current == stop {
1855 break;
1856 }
1857 current = current.checked_add(step).ok_or(ExecutorError::Evaluation(
1858 crate::executor::EvaluationError::Overflow,
1859 ))?;
1860 if (step > 0 && current > stop) || (step < 0 && current < stop) {
1861 break;
1862 }
1863 }
1864 Ok(rows)
1865}
1866
1867fn generate_timestamp_series(values: &[SqlValue]) -> Result<Vec<Row>> {
1868 if values.iter().any(SqlValue::is_null) {
1869 return Ok(Vec::new());
1870 }
1871 let [
1872 SqlValue::Timestamp(start),
1873 SqlValue::Timestamp(stop),
1874 SqlValue::Interval {
1875 months,
1876 days,
1877 micros,
1878 },
1879 ] = values
1880 else {
1881 return Err(ExecutorError::InvalidOperation {
1882 operation: "GENERATE_SERIES".into(),
1883 reason: "expected (TIMESTAMP, TIMESTAMP, INTERVAL)".into(),
1884 });
1885 };
1886 if *months == 0 && *days == 0 && *micros == 0 {
1887 return Err(ExecutorError::InvalidOperation {
1888 operation: "GENERATE_SERIES".into(),
1889 reason: "step must not be zero".into(),
1890 });
1891 }
1892 if start == stop {
1893 return Ok(vec![Row::new(0, vec![SqlValue::Timestamp(*start)])]);
1894 }
1895
1896 let advance = |value| -> Result<i64> {
1897 match crate::executor::evaluator::binary_op::add_timestamp_interval(
1898 value, *months, *days, *micros,
1899 )? {
1900 SqlValue::Timestamp(value) => Ok(value),
1901 _ => unreachable!("timestamp plus interval returns timestamp"),
1902 }
1903 };
1904 let first_next = advance(*start)?;
1905 let ascending = first_next > *start;
1906 if first_next == *start {
1907 return Err(ExecutorError::InvalidOperation {
1908 operation: "GENERATE_SERIES".into(),
1909 reason: "step must advance the timestamp".into(),
1910 });
1911 }
1912 if (ascending && start > stop) || (!ascending && start < stop) {
1913 return Ok(Vec::new());
1914 }
1915
1916 let mut rows = Vec::new();
1917 let mut current = *start;
1918 loop {
1919 if rows.len() == MAX_GENERATED_SERIES_ROWS {
1920 return Err(ExecutorError::ResourceExhausted {
1921 message: format!("GENERATE_SERIES is limited to {MAX_GENERATED_SERIES_ROWS} rows"),
1922 });
1923 }
1924 rows.push(Row::new(
1925 rows.len() as u64,
1926 vec![SqlValue::Timestamp(current)],
1927 ));
1928 if current == *stop {
1929 break;
1930 }
1931 let next = advance(current)?;
1932 if (ascending && next <= current) || (!ascending && next >= current) {
1933 return Err(ExecutorError::InvalidOperation {
1934 operation: "GENERATE_SERIES".into(),
1935 reason: "step must advance consistently".into(),
1936 });
1937 }
1938 if (ascending && next > *stop) || (!ascending && next < *stop) {
1939 break;
1940 }
1941 current = next;
1942 }
1943 Ok(rows)
1944}
1945
1946fn eval_expr(expr: &crate::planner::typed_expr::TypedExpr, row: &Row) -> Result<SqlValue> {
1948 let ctx = EvalContext::new(&row.values);
1949 crate::executor::evaluator::evaluate(expr, &ctx)
1950}
1951
1952fn combine_outer_for_eval(row: &Row, outer: Option<&Row>) -> Row {
1953 let Some(outer) = outer else {
1954 return row.clone();
1955 };
1956 let mut values = Vec::with_capacity(row.len() + outer.len());
1957 values.extend(row.values.clone());
1958 values.extend(outer.values.clone());
1959 Row::new(row.row_id, values)
1960}
1961
1962fn widen_rows_with_outer(mut input: Box<dyn RowIterator>, outer: &Row) -> Result<Vec<Row>> {
1974 let mut rows = Vec::new();
1975 while let Some(result) = input.next_row() {
1976 rows.push(combine_outer_for_eval(&result?, Some(outer)));
1977 }
1978 Ok(rows)
1979}
1980
1981fn narrow_rows_from_outer(rows: Vec<Row>, outer_width: usize) -> Vec<Row> {
1987 rows.into_iter()
1988 .map(|mut row| {
1989 let keep = row.values.len().saturating_sub(outer_width);
1990 row.values.truncate(keep);
1991 row
1992 })
1993 .collect()
1994}
1995
1996fn drain_rows(mut input: Box<dyn RowIterator>) -> Result<Vec<Row>> {
1997 let mut rows = Vec::new();
1998 while let Some(result) = input.next_row() {
1999 rows.push(result?);
2000 }
2001 Ok(rows)
2002}
2003
2004fn execute_project_with_subqueries<
2005 'txn,
2006 S: KVStore + 'txn,
2007 C: Catalog + ?Sized,
2008 T: SqlTxn<'txn, S>,
2009>(
2010 txn: &mut T,
2011 catalog: &C,
2012 rows: Vec<Row>,
2013 projection: &Projection,
2014 schema: &[crate::catalog::ColumnMetadata],
2015 outer: Option<&Row>,
2016) -> Result<QueryResult> {
2017 match projection {
2018 Projection::All(_) => project::execute_project(rows, projection, schema),
2019 Projection::Columns(cols)
2020 if outer.is_some() || cols.iter().any(|c| subquery::contains_subquery(&c.expr)) =>
2021 {
2022 let columns: Vec<_> = cols
2023 .iter()
2024 .enumerate()
2025 .map(|(i, c)| column_info_from_projection(c, i))
2026 .collect();
2027 let mut projected_rows = Vec::with_capacity(rows.len());
2028 for row in rows {
2029 let eval_row = combine_outer_for_eval(&row, outer);
2030 let mut values = Vec::with_capacity(cols.len());
2031 for col in cols {
2032 values.push(subquery::evaluate_expr_with_subqueries(
2033 txn, catalog, &col.expr, &eval_row,
2034 )?);
2035 }
2036 projected_rows.push(values);
2037 }
2038 Ok(QueryResult::new(columns, projected_rows))
2039 }
2040 Projection::Columns(_) => project::execute_project(rows, projection, schema),
2041 }
2042}
2043
2044fn column_name_from_projection(
2046 projected: &crate::planner::typed_expr::ProjectedColumn,
2047 idx: usize,
2048) -> String {
2049 use crate::planner::typed_expr::TypedExprKind;
2050
2051 projected
2052 .alias
2053 .clone()
2054 .or_else(|| match &projected.expr.kind {
2055 TypedExprKind::ColumnRef { column, .. } => Some(column.clone()),
2056 TypedExprKind::FunctionCall { name, args, .. }
2059 if name == "coalesce" && !args.is_empty() =>
2060 {
2061 let first_column = match &args[0].kind {
2062 TypedExprKind::ColumnRef { column, .. } => Some(column),
2063 _ => None,
2064 };
2065 first_column
2066 .filter(|column| {
2067 args.iter().all(|arg| {
2068 matches!(
2069 &arg.kind,
2070 TypedExprKind::ColumnRef { column: other, .. } if other == *column
2071 )
2072 })
2073 })
2074 .cloned()
2075 }
2076 _ => None,
2077 })
2078 .unwrap_or_else(|| format!("col_{idx}"))
2079}
2080
2081fn column_info_from_projection(
2083 projected: &crate::planner::typed_expr::ProjectedColumn,
2084 idx: usize,
2085) -> ColumnInfo {
2086 ColumnInfo::new(
2087 column_name_from_projection(projected, idx),
2088 projected.expr.resolved_type.clone(),
2089 )
2090}
2091
2092fn column_infos_from_all(
2094 schema: &[crate::catalog::ColumnMetadata],
2095 names: &[String],
2096) -> Result<Vec<ColumnInfo>> {
2097 if names.len() == schema.len() {
2098 return Ok(names
2099 .iter()
2100 .zip(schema)
2101 .map(|(name, column)| ColumnInfo::new(name.clone(), column.data_type.clone()))
2102 .collect());
2103 }
2104 names
2105 .iter()
2106 .map(|name| {
2107 let col = schema
2108 .iter()
2109 .find(|c| &c.name == name)
2110 .ok_or_else(|| ExecutorError::ColumnNotFound(name.clone()))?;
2111 Ok(ColumnInfo::new(name.clone(), col.data_type.clone()))
2112 })
2113 .collect()
2114}
2115
2116#[cfg(test)]
2117mod tests {
2118 use super::*;
2119 use crate::catalog::{ColumnMetadata, MemoryCatalog, TableMetadata};
2120 use crate::executor::SpillPolicy;
2121 use crate::executor::ddl::create_table::execute_create_table;
2122 use crate::planner::typed_expr::{ProjectedColumn, TypedExpr};
2123 use crate::planner::types::ResolvedType;
2124 use crate::storage::TxnBridge;
2125 use alopex_core::kv::memory::MemoryKV;
2126 use std::sync::Arc;
2127
2128 fn text_literal_plan(value: &str) -> LogicalPlan {
2129 LogicalPlan::Project {
2130 input: Box::new(LogicalPlan::Scan {
2131 table: LITERAL_TABLE.into(),
2132 projection: Projection::All(Vec::new()),
2133 }),
2134 projection: Projection::Columns(vec![ProjectedColumn::new(TypedExpr::literal(
2135 crate::ast::expr::Literal::String(value.into()),
2136 ResolvedType::Text,
2137 crate::Span::default(),
2138 ))]),
2139 }
2140 }
2141
2142 #[test]
2143 fn execute_query_scan_only_returns_rows() {
2144 let bridge = TxnBridge::new(Arc::new(MemoryKV::new()));
2145 let mut catalog = MemoryCatalog::new();
2146 let table = TableMetadata::new(
2147 "users",
2148 vec![
2149 ColumnMetadata::new("id", ResolvedType::Integer),
2150 ColumnMetadata::new("name", ResolvedType::Text),
2151 ],
2152 );
2153 let mut ddl_txn = bridge.begin_write().unwrap();
2154 execute_create_table(&mut ddl_txn, &mut catalog, table.clone(), vec![], false).unwrap();
2155 ddl_txn.commit().unwrap();
2156
2157 let mut txn = bridge.begin_write().unwrap();
2158 crate::executor::dml::execute_insert(
2159 &mut txn,
2160 &catalog,
2161 "users",
2162 vec!["id".into(), "name".into()],
2163 vec![vec![
2164 TypedExpr::literal(
2165 crate::ast::expr::Literal::Number("1".into()),
2166 ResolvedType::Integer,
2167 crate::Span::default(),
2168 ),
2169 TypedExpr::literal(
2170 crate::ast::expr::Literal::String("alice".into()),
2171 ResolvedType::Text,
2172 crate::Span::default(),
2173 ),
2174 ]],
2175 )
2176 .unwrap();
2177
2178 let result = execute_query(
2179 &mut txn,
2180 &catalog,
2181 LogicalPlan::scan(
2182 "users".into(),
2183 Projection::All(vec!["id".into(), "name".into()]),
2184 ),
2185 )
2186 .unwrap();
2187
2188 match result {
2189 ExecutionResult::Query(q) => {
2190 assert_eq!(q.rows.len(), 1);
2191 assert_eq!(q.columns.len(), 2);
2192 assert_eq!(
2193 q.rows[0],
2194 vec![SqlValue::Integer(1), SqlValue::Text("alice".into())]
2195 );
2196 }
2197 other => panic!("unexpected result {other:?}"),
2198 }
2199 }
2200
2201 #[test]
2202 fn recursive_cte_accounts_for_reference_clone_high_water() {
2203 let bridge = TxnBridge::new(Arc::new(MemoryKV::new()));
2204 let catalog = MemoryCatalog::new();
2205 let schema = vec![ColumnMetadata::new("value", ResolvedType::Text)];
2206 let anchor = text_literal_plan(&"x".repeat(32));
2207 let recursive_term = LogicalPlan::RecursiveReference {
2208 name: "memory_cycle".into(),
2209 schema: schema.clone(),
2210 };
2211 let plan = LogicalPlan::RecursiveCte {
2212 name: "memory_cycle".into(),
2213 anchor: Box::new(anchor),
2214 recursive_term: Box::new(recursive_term),
2215 union_all: false,
2216 schema,
2217 limits: RecursiveCteLimits::default(),
2218 };
2219 let policy = MemoryPolicy::new(Some(115), SpillPolicy::FailFast);
2220 let mut txn = bridge.begin_write().unwrap();
2221
2222 let error = execute_query_with_policy(&mut txn, &catalog, plan, Some(&policy))
2223 .expect_err("the working table and its iterator clone must both be accounted");
2224
2225 assert!(
2226 matches!(&error, ExecutorError::ResourceExhausted { message }
2227 if message.contains("query memory limit exceeded")),
2228 "expected recursive materialization to honor the query memory limit, got: {error}"
2229 );
2230 }
2231
2232 #[test]
2233 fn recursive_cte_enforces_accumulated_row_limit() {
2234 let bridge = TxnBridge::new(Arc::new(MemoryKV::new()));
2235 let catalog = MemoryCatalog::new();
2236 let schema = vec![ColumnMetadata::new("value", ResolvedType::Text)];
2237 let anchor = LogicalPlan::SetOperation {
2238 left: Box::new(text_literal_plan("first")),
2239 right: Box::new(text_literal_plan("second")),
2240 operator: SetOperator::Union,
2241 all: true,
2242 };
2243 let plan = LogicalPlan::RecursiveCte {
2244 name: "bounded".into(),
2245 anchor: Box::new(anchor),
2246 recursive_term: Box::new(LogicalPlan::RecursiveReference {
2247 name: "bounded".into(),
2248 schema: schema.clone(),
2249 }),
2250 union_all: true,
2251 schema,
2252 limits: RecursiveCteLimits {
2253 max_iterations: 10,
2254 max_rows: 1,
2255 },
2256 };
2257 let mut txn = bridge.begin_write().unwrap();
2258
2259 let error = execute_query(&mut txn, &catalog, plan)
2260 .expect_err("the accumulated row limit must be checked before iteration");
2261
2262 assert!(
2263 matches!(&error, ExecutorError::ResourceExhausted { message }
2264 if message.contains("recursive CTE 'bounded' reached row limit 1")),
2265 "expected a named recursive row-limit error, got: {error}"
2266 );
2267 }
2268
2269 #[test]
2270 fn recursive_cte_accounts_for_result_before_releasing_working_table() {
2271 let bridge = TxnBridge::new(Arc::new(MemoryKV::new()));
2272 let catalog = MemoryCatalog::new();
2273 let value = "y".repeat(16);
2274 let schema = vec![ColumnMetadata::new("value", ResolvedType::Text)];
2275 let reference = LogicalPlan::RecursiveReference {
2276 name: "result_overlap".into(),
2277 schema: schema.clone(),
2278 };
2279 let two_literals = LogicalPlan::SetOperation {
2280 left: Box::new(text_literal_plan(&value)),
2281 right: Box::new(text_literal_plan(&value)),
2282 operator: SetOperator::Union,
2283 all: true,
2284 };
2285 let recursive_term = LogicalPlan::SetOperation {
2286 left: Box::new(reference),
2287 right: Box::new(two_literals),
2288 operator: SetOperator::Union,
2289 all: true,
2290 };
2291 let plan = LogicalPlan::RecursiveCte {
2292 name: "result_overlap".into(),
2293 anchor: Box::new(text_literal_plan(&value)),
2294 recursive_term: Box::new(recursive_term),
2295 union_all: false,
2296 schema,
2297 limits: RecursiveCteLimits::default(),
2298 };
2299 let policy = MemoryPolicy::new(Some(92), SpillPolicy::FailFast);
2300 let mut txn = bridge.begin_write().unwrap();
2301
2302 let error = execute_query_with_policy(&mut txn, &catalog, plan, Some(&policy))
2303 .expect_err("recursive result must be counted before the working table is released");
2304
2305 assert!(
2306 matches!(&error, ExecutorError::ResourceExhausted { message }
2307 if message.contains("query memory limit exceeded")),
2308 "expected result/working-table overlap to honor the query memory limit, got: {error}"
2309 );
2310 }
2311}