1use alopex_core::kv::KVStore;
2use std::collections::{HashMap, HashSet};
3
4use crate::ast::LITERAL_TABLE;
5use crate::catalog::{Catalog, StorageType};
6use crate::executor::evaluator::EvalContext;
7use crate::executor::memory::MemoryPolicy;
8use crate::executor::{ExecutionResult, ExecutorError, QueryResult, QueryRowIterator, Result};
9use crate::planner::logical_plan::{LogicalPlan, SetOperator};
10use crate::planner::typed_expr::{Projection, SortExpr};
11use crate::storage::{SqlTxn, SqlValue};
12
13use super::{ColumnInfo, Row};
14
15pub mod aggregate;
16pub mod columnar_scan;
17pub mod iterator;
18pub mod join;
19mod knn;
20mod project;
21mod scan;
22pub mod subquery;
23pub mod window;
24
25pub use columnar_scan::{ColumnarScanIterator, create_columnar_scan_iterator};
26pub use iterator::{FilterIterator, LimitIterator, RowIterator, ScanIterator, SortIterator};
27pub use project::{project_row_values, projected_columns};
28pub use scan::{
29 create_fenced_range_scan_iterator, create_scan_iterator, execute_fenced_range_scan,
30};
31
32pub fn execute_query<'txn, S: KVStore + 'txn, C: Catalog + ?Sized, T: SqlTxn<'txn, S>>(
44 txn: &mut T,
45 catalog: &C,
46 plan: LogicalPlan,
47) -> Result<ExecutionResult> {
48 execute_query_with_policy(txn, catalog, plan, None)
49}
50
51pub fn execute_query_with_policy<
52 'txn,
53 S: KVStore + 'txn,
54 C: Catalog + ?Sized,
55 T: SqlTxn<'txn, S>,
56>(
57 txn: &mut T,
58 catalog: &C,
59 plan: LogicalPlan,
60 memory: Option<&MemoryPolicy>,
61) -> Result<ExecutionResult> {
62 if let Some((pattern, projection, filter)) = knn::extract_knn_context(&plan) {
63 return knn::execute_knn_query(txn, catalog, &pattern, &projection, filter.as_ref());
64 }
65
66 let result = execute_query_result_with_outer_and_policy(txn, catalog, plan, None, memory)?;
67 Ok(ExecutionResult::Query(result))
68}
69
70pub(crate) fn execute_query_result_with_outer<
71 'txn,
72 S: KVStore + 'txn,
73 C: Catalog + ?Sized,
74 T: SqlTxn<'txn, S>,
75>(
76 txn: &mut T,
77 catalog: &C,
78 plan: LogicalPlan,
79 outer: Option<&Row>,
80) -> Result<QueryResult> {
81 execute_query_result_with_outer_and_policy(txn, catalog, plan, outer, None)
82}
83
84fn execute_query_result_with_outer_and_policy<
85 'txn,
86 S: KVStore + 'txn,
87 C: Catalog + ?Sized,
88 T: SqlTxn<'txn, S>,
89>(
90 txn: &mut T,
91 catalog: &C,
92 plan: LogicalPlan,
93 outer: Option<&Row>,
94 memory: Option<&MemoryPolicy>,
95) -> Result<QueryResult> {
96 let (mut iter, projection, schema) =
97 build_iterator_pipeline_with_outer(txn, catalog, plan, memory, outer)?;
98 let mut rows = Vec::new();
99 while let Some(result) = iter.next_row() {
100 rows.push(result?);
101 }
102 execute_project_with_subqueries(txn, catalog, rows, &projection, &schema, outer)
103}
104
105pub fn execute_query_streaming<'txn, S: KVStore + 'txn, C: Catalog + ?Sized, T: SqlTxn<'txn, S>>(
121 txn: &mut T,
122 catalog: &C,
123 plan: LogicalPlan,
124) -> Result<QueryRowIterator<'static>> {
125 execute_query_streaming_with_policy(txn, catalog, plan, None)
126}
127
128pub fn execute_query_streaming_with_policy<
129 'txn,
130 S: KVStore + 'txn,
131 C: Catalog + ?Sized,
132 T: SqlTxn<'txn, S>,
133>(
134 txn: &mut T,
135 catalog: &C,
136 plan: LogicalPlan,
137 memory: Option<&MemoryPolicy>,
138) -> Result<QueryRowIterator<'static>> {
139 if knn::extract_knn_context(&plan).is_some() {
141 let result = execute_query_with_policy(txn, catalog, plan, memory)?;
143 if let ExecutionResult::Query(qr) = result {
144 let (iter, projection, schema) = materialize_query_result(qr);
145 return Ok(QueryRowIterator::new(iter, projection, schema));
146 }
147 return Err(ExecutorError::InvalidOperation {
148 operation: "execute_query_streaming".into(),
149 reason: "KNN query did not return Query result".into(),
150 });
151 }
152
153 if subquery::plan_contains_subquery(&plan) {
158 let result = execute_query_result_with_outer_and_policy(txn, catalog, plan, None, memory)?;
159 let (iter, projection, schema) = materialize_query_result(result);
160 return Ok(QueryRowIterator::new(iter, projection, schema));
161 }
162
163 let (iter, projection, schema) = build_iterator_pipeline(txn, catalog, plan, memory)?;
164
165 Ok(QueryRowIterator::new(iter, projection, schema))
166}
167
168fn materialize_query_result(
173 result: QueryResult,
174) -> (
175 Box<dyn RowIterator>,
176 Projection,
177 Vec<crate::catalog::ColumnMetadata>,
178) {
179 let column_names: Vec<String> = result.columns.iter().map(|c| c.name.clone()).collect();
180 let schema: Vec<crate::catalog::ColumnMetadata> = result
181 .columns
182 .iter()
183 .map(|c| crate::catalog::ColumnMetadata::new(&c.name, c.data_type.clone()))
184 .collect();
185 let rows: Vec<Row> = result
186 .rows
187 .into_iter()
188 .enumerate()
189 .map(|(i, values)| Row::new(i as u64, values))
190 .collect();
191 let iter = iterator::VecIterator::new(rows, schema.clone());
192 (Box::new(iter), Projection::All(column_names), schema)
193}
194
195fn build_iterator_pipeline<'txn, S: KVStore + 'txn, C: Catalog + ?Sized, T: SqlTxn<'txn, S>>(
202 txn: &mut T,
203 catalog: &C,
204 plan: LogicalPlan,
205 memory: Option<&MemoryPolicy>,
206) -> Result<(
207 Box<dyn RowIterator>,
208 Projection,
209 Vec<crate::catalog::ColumnMetadata>,
210)> {
211 build_iterator_pipeline_with_outer(txn, catalog, plan, memory, None)
212}
213
214fn build_iterator_pipeline_with_outer<
215 'txn,
216 S: KVStore + 'txn,
217 C: Catalog + ?Sized,
218 T: SqlTxn<'txn, S>,
219>(
220 txn: &mut T,
221 catalog: &C,
222 plan: LogicalPlan,
223 memory: Option<&MemoryPolicy>,
224 outer: Option<&Row>,
225) -> Result<(
226 Box<dyn RowIterator>,
227 Projection,
228 Vec<crate::catalog::ColumnMetadata>,
229)> {
230 match plan {
231 LogicalPlan::Scan { table, projection } => {
232 if table == LITERAL_TABLE {
233 let schema = Vec::new();
234 let rows = vec![Row::new(0, Vec::new())];
235 let iter = iterator::VecIterator::new(rows, schema.clone());
236 return Ok((Box::new(iter), projection, schema));
237 }
238 let table_meta = catalog
239 .get_table(&table)
240 .cloned()
241 .ok_or_else(|| ExecutorError::TableNotFound(table.clone()))?;
242
243 if table_meta.storage_options.storage_type == StorageType::Columnar {
244 let columnar_scan = columnar_scan::build_columnar_scan(&table_meta, &projection);
245 let rows = columnar_scan::execute_columnar_scan(txn, &table_meta, &columnar_scan)?;
246 let schema = table_meta.columns.clone();
247 let iter = iterator::VecIterator::new(rows, schema.clone());
248 return Ok((Box::new(iter), projection, schema));
249 }
250
251 let rows = scan::execute_scan(txn, &table_meta)?;
255 let schema = table_meta.columns.clone();
256
257 let iter = iterator::VecIterator::new(rows, schema.clone());
259 Ok((Box::new(iter), projection, schema))
260 }
261 LogicalPlan::Filter { input, predicate } => {
262 if let LogicalPlan::Scan { table, projection } = input.as_ref()
263 && let Some(table_meta) = catalog.get_table(table)
264 && table_meta.storage_options.storage_type == StorageType::Columnar
265 {
266 let columnar_scan = columnar_scan::build_columnar_scan_for_filter(
267 table_meta,
268 projection.clone(),
269 &predicate,
270 );
271 let rows = columnar_scan::execute_columnar_scan(txn, table_meta, &columnar_scan)?;
272 let schema = table_meta.columns.clone();
273 let iter = iterator::VecIterator::new(rows, schema.clone());
274 return Ok((Box::new(iter), projection.clone(), schema));
275 }
276 let (mut input_iter, projection, schema) =
277 build_iterator_pipeline_with_outer(txn, catalog, *input, memory, outer)?;
278 if outer.is_some() || subquery::contains_subquery(&predicate) {
279 let mut rows = Vec::new();
280 while let Some(result) = input_iter.next_row() {
281 let row = result?;
282 let eval_row = combine_outer_for_eval(&row, outer);
283 if let SqlValue::Boolean(true) = subquery::evaluate_expr_with_subqueries(
284 txn, catalog, &predicate, &eval_row,
285 )? {
286 rows.push(row);
287 }
288 }
289 let iter = iterator::VecIterator::new(rows, schema.clone());
290 return Ok((Box::new(iter), projection, schema));
291 }
292 let filter_iter = FilterIterator::new(input_iter, predicate);
293 Ok((Box::new(filter_iter), projection, schema))
294 }
295 LogicalPlan::Project { input, projection } => {
296 let input_result =
297 execute_query_result_with_outer_and_policy(txn, catalog, *input, outer, memory)?;
298 let (mut input_iter, _input_projection, schema) =
299 materialize_query_result(input_result);
300 let mut rows = Vec::new();
301 while let Some(result) = input_iter.next_row() {
302 rows.push(result?);
303 }
304 let projected =
305 execute_project_with_subqueries(txn, catalog, rows, &projection, &schema, outer)?;
306 let output_schema = projected
307 .columns
308 .iter()
309 .map(|col| crate::catalog::ColumnMetadata::new(&col.name, col.data_type.clone()))
310 .collect::<Vec<_>>();
311 let rows = projected
312 .rows
313 .into_iter()
314 .enumerate()
315 .map(|(idx, values)| Row::new(idx as u64, values))
316 .collect::<Vec<_>>();
317 let output_projection =
318 Projection::All(output_schema.iter().map(|col| col.name.clone()).collect());
319 let iter = iterator::VecIterator::new(rows, output_schema.clone());
320 Ok((Box::new(iter), output_projection, output_schema))
321 }
322 LogicalPlan::Join {
323 left,
324 right,
325 join_type,
326 condition,
327 using: _,
328 } => {
329 let (mut left_iter, _left_projection, left_schema) =
330 build_iterator_pipeline_with_outer(txn, catalog, *left, memory, outer)?;
331 let (mut right_iter, _right_projection, right_schema) =
332 build_iterator_pipeline_with_outer(txn, catalog, *right, memory, outer)?;
333 let mut left_rows = Vec::new();
334 while let Some(result) = left_iter.next_row() {
335 left_rows.push(result?);
336 }
337 let mut right_rows = Vec::new();
338 while let Some(result) = right_iter.next_row() {
339 right_rows.push(result?);
340 }
341 let left_width = left_schema.len();
342 let right_width = right_schema.len();
343 let rows = join::execute_join_with_widths(
344 left_rows,
345 right_rows,
346 join_type,
347 condition.as_ref(),
348 left_width,
349 right_width,
350 )?;
351 let mut schema = left_schema;
352 schema.extend(right_schema);
353 let projection = Projection::All(schema.iter().map(|col| col.name.clone()).collect());
354 let iter = iterator::VecIterator::new(rows, schema.clone());
355 Ok((Box::new(iter), projection, schema))
356 }
357 LogicalPlan::Aggregate {
358 input,
359 group_keys,
360 aggregates,
361 having,
362 projection,
363 } => {
364 let (input_iter, _projection, _schema) =
365 build_iterator_pipeline_with_outer(txn, catalog, *input, memory, outer)?;
366 let schema = aggregate::build_aggregate_schema(&group_keys, &aggregates);
367 if let Some(policy) = memory
368 && policy.spill_directory().is_some()
369 {
370 if group_keys.is_empty() {
371 let iter = aggregate::StreamingAggregateIterator::new(
372 input_iter,
373 group_keys,
374 aggregates,
375 having,
376 schema.clone(),
377 );
378 return Ok((Box::new(iter), projection, schema));
379 }
380 let order_by = group_keys
381 .iter()
382 .cloned()
383 .map(|expr| SortExpr {
384 expr,
385 asc: true,
386 nulls_first: false,
387 })
388 .collect::<Vec<_>>();
389 let sort_iter =
390 SortIterator::new_with_policy(input_iter, &order_by, Some(policy.clone()))?;
391 let iter = aggregate::StreamingAggregateIterator::new(
392 Box::new(sort_iter),
393 group_keys,
394 aggregates,
395 having,
396 schema.clone(),
397 );
398 return Ok((Box::new(iter), projection, schema));
399 }
400
401 let parallelism = std::thread::available_parallelism()
402 .map(usize::from)
403 .unwrap_or(1);
404 if !aggregate::should_use_single_for_parallel(parallelism, &aggregates) {
405 let rows = aggregate::execute_parallel_aggregate_rows_with_policy(
406 input_iter,
407 group_keys,
408 aggregates,
409 having,
410 schema.clone(),
411 parallelism,
412 memory.cloned(),
413 1_000_000,
414 )?;
415 let iter = iterator::VecIterator::new(rows, schema.clone());
416 return Ok((Box::new(iter), projection, schema));
417 }
418
419 let mut iter = aggregate::AggregateIterator::new(
420 input_iter,
421 group_keys,
422 aggregates,
423 having,
424 schema.clone(),
425 );
426 if let Some(policy) = memory {
427 iter = iter.with_memory_policy(Some(policy.clone()));
428 }
429 Ok((Box::new(iter), projection, schema))
430 }
431 LogicalPlan::Window { input, windows } => {
432 let (input_iter, _projection, _schema) =
433 build_iterator_pipeline_with_outer(txn, catalog, *input, memory, outer)?;
434 let iter = window::WindowIterator::new(input_iter, windows, memory)?;
435 let schema = iter.schema().to_vec();
436 let projection =
437 Projection::All(schema.iter().map(|column| column.name.clone()).collect());
438 Ok((Box::new(iter), projection, schema))
439 }
440 LogicalPlan::SetOperation {
441 left,
442 right,
443 operator,
444 all,
445 } => {
446 let left_result =
447 execute_query_result_with_outer_and_policy(txn, catalog, *left, outer, memory)?;
448 let right_result =
449 execute_query_result_with_outer_and_policy(txn, catalog, *right, outer, memory)?;
450 let rows = execute_set_operation(operator, all, left_result.rows, right_result.rows)?;
451 let schema = left_result
452 .columns
453 .iter()
454 .map(|column| {
455 crate::catalog::ColumnMetadata::new(&column.name, column.data_type.clone())
456 })
457 .collect::<Vec<_>>();
458 let projection = Projection::All(
459 left_result
460 .columns
461 .iter()
462 .map(|column| column.name.clone())
463 .collect(),
464 );
465 let rows = rows
466 .into_iter()
467 .enumerate()
468 .map(|(index, values)| Row::new(index as u64, values))
469 .collect();
470 Ok((
471 Box::new(iterator::VecIterator::new(rows, schema.clone())),
472 projection,
473 schema,
474 ))
475 }
476 LogicalPlan::Sort { input, order_by } => {
477 let (input_iter, projection, schema) =
478 build_iterator_pipeline_with_outer(txn, catalog, *input, memory, outer)?;
479 let sort_iter = if let Some(policy) = memory {
480 SortIterator::new_with_policy(input_iter, &order_by, Some(policy.clone()))?
481 } else {
482 SortIterator::new(input_iter, &order_by)?
483 };
484 Ok((Box::new(sort_iter), projection, schema))
485 }
486 LogicalPlan::Limit {
487 input,
488 limit,
489 offset,
490 } => {
491 let (input_iter, projection, schema) =
492 build_iterator_pipeline_with_outer(txn, catalog, *input, memory, outer)?;
493 let limit_iter = LimitIterator::new(input_iter, limit, offset);
494 Ok((Box::new(limit_iter), projection, schema))
495 }
496 other => Err(ExecutorError::UnsupportedOperation(format!(
497 "unsupported query plan: {other:?}"
498 ))),
499 }
500}
501
502pub fn build_streaming_pipeline<
514 'a,
515 'txn: 'a,
516 S: KVStore + 'txn,
517 C: Catalog + ?Sized,
518 T: SqlTxn<'txn, S>,
519>(
520 txn: &'a mut T,
521 catalog: &C,
522 plan: LogicalPlan,
523) -> Result<(
524 Box<dyn RowIterator + 'a>,
525 Projection,
526 Vec<crate::catalog::ColumnMetadata>,
527)> {
528 build_streaming_pipeline_with_policy(txn, catalog, plan, None)
529}
530
531pub fn build_streaming_pipeline_with_policy<
532 'a,
533 'txn: 'a,
534 S: KVStore + 'txn,
535 C: Catalog + ?Sized,
536 T: SqlTxn<'txn, S>,
537>(
538 txn: &'a mut T,
539 catalog: &C,
540 plan: LogicalPlan,
541 memory: Option<&MemoryPolicy>,
542) -> Result<(
543 Box<dyn RowIterator + 'a>,
544 Projection,
545 Vec<crate::catalog::ColumnMetadata>,
546)> {
547 if subquery::plan_contains_subquery(&plan) {
553 let result = execute_query_result_with_outer_and_policy(txn, catalog, plan, None, memory)?;
554 return Ok(materialize_query_result(result));
555 }
556
557 match plan {
558 LogicalPlan::SetOperation {
559 left,
560 right,
561 operator,
562 all,
563 } => build_streaming_set_operation(txn, catalog, *left, *right, operator, all, memory),
564 other => build_streaming_pipeline_inner(txn, catalog, other, memory),
565 }
566}
567
568fn build_streaming_set_operation<
569 'a,
570 'txn: 'a,
571 S: KVStore + 'txn,
572 C: Catalog + ?Sized,
573 T: SqlTxn<'txn, S>,
574>(
575 txn: &'a mut T,
576 catalog: &C,
577 left: LogicalPlan,
578 right: LogicalPlan,
579 operator: SetOperator,
580 all: bool,
581 memory: Option<&MemoryPolicy>,
582) -> Result<(
583 Box<dyn RowIterator + 'a>,
584 Projection,
585 Vec<crate::catalog::ColumnMetadata>,
586)> {
587 let (mut left_iter, left_projection, left_schema) =
588 build_streaming_pipeline_with_policy(txn, catalog, left, memory)?;
589 let mut left_rows = Vec::new();
590 while let Some(result) = left_iter.next_row() {
591 left_rows.push(result?);
592 }
593 drop(left_iter);
594 let left_result = project::execute_project(left_rows, &left_projection, &left_schema)?;
595
596 let (mut right_iter, right_projection, right_schema) =
597 build_streaming_pipeline_with_policy(txn, catalog, right, memory)?;
598 let mut right_rows = Vec::new();
599 while let Some(result) = right_iter.next_row() {
600 right_rows.push(result?);
601 }
602 let right_result = project::execute_project(right_rows, &right_projection, &right_schema)?;
603
604 let rows = execute_set_operation(operator, all, left_result.rows, right_result.rows)?;
605 let schema = left_result
606 .columns
607 .iter()
608 .map(|column| crate::catalog::ColumnMetadata::new(&column.name, column.data_type.clone()))
609 .collect::<Vec<_>>();
610 let projection = Projection::All(
611 left_result
612 .columns
613 .iter()
614 .map(|column| column.name.clone())
615 .collect(),
616 );
617 let rows = rows
618 .into_iter()
619 .enumerate()
620 .map(|(index, values)| Row::new(index as u64, values))
621 .collect();
622 Ok((
623 Box::new(iterator::VecIterator::new(rows, schema.clone())),
624 projection,
625 schema,
626 ))
627}
628
629fn execute_set_operation(
630 operator: SetOperator,
631 all: bool,
632 mut left: Vec<Vec<SqlValue>>,
633 right: Vec<Vec<SqlValue>>,
634) -> Result<Vec<Vec<SqlValue>>> {
635 if all {
636 match operator {
637 SetOperator::Union => {
638 left.extend(right);
639 return Ok(left);
640 }
641 SetOperator::Intersect | SetOperator::Except => {
642 let mut right_counts = HashMap::<Vec<u8>, usize>::new();
643 for row in right {
644 *right_counts
645 .entry(aggregate::encode_group_key(&row)?)
646 .or_default() += 1;
647 }
648 let mut output = Vec::new();
649 for row in left {
650 let key = aggregate::encode_group_key(&row)?;
651 let remaining = right_counts.entry(key).or_default();
652 if *remaining > 0 {
653 *remaining -= 1;
654 if operator == SetOperator::Intersect {
655 output.push(row);
656 }
657 } else if operator == SetOperator::Except {
658 output.push(row);
659 }
660 }
661 return Ok(output);
662 }
663 }
664 }
665
666 let right_keys = right
667 .iter()
668 .map(|row| aggregate::encode_group_key(row))
669 .collect::<Result<HashSet<_>>>()?;
670 let mut seen = HashSet::new();
671 let mut output = Vec::new();
672 match operator {
673 SetOperator::Union => {
674 for row in left.into_iter().chain(right) {
675 if seen.insert(aggregate::encode_group_key(&row)?) {
676 output.push(row);
677 }
678 }
679 }
680 SetOperator::Intersect => {
681 for row in left {
682 let key = aggregate::encode_group_key(&row)?;
683 if right_keys.contains(&key) && seen.insert(key) {
684 output.push(row);
685 }
686 }
687 }
688 SetOperator::Except => {
689 for row in left {
690 let key = aggregate::encode_group_key(&row)?;
691 if !right_keys.contains(&key) && seen.insert(key) {
692 output.push(row);
693 }
694 }
695 }
696 }
697 Ok(output)
698}
699
700fn build_streaming_pipeline_inner<
702 'a,
703 'txn: 'a,
704 S: KVStore + 'txn,
705 C: Catalog + ?Sized,
706 T: SqlTxn<'txn, S>,
707>(
708 txn: &'a mut T,
709 catalog: &C,
710 plan: LogicalPlan,
711 memory: Option<&MemoryPolicy>,
712) -> Result<(
713 Box<dyn RowIterator + 'a>,
714 Projection,
715 Vec<crate::catalog::ColumnMetadata>,
716)> {
717 match plan {
718 LogicalPlan::Scan { table, projection } => {
719 if table == LITERAL_TABLE {
720 let schema = Vec::new();
721 let rows = vec![Row::new(0, Vec::new())];
722 let iter = iterator::VecIterator::new(rows, schema.clone());
723 return Ok((Box::new(iter), projection, schema));
724 }
725 let table_meta = catalog
726 .get_table(&table)
727 .cloned()
728 .ok_or_else(|| ExecutorError::TableNotFound(table.clone()))?;
729
730 if table_meta.storage_options.storage_type == StorageType::Columnar {
731 let columnar_scan = columnar_scan::build_columnar_scan(&table_meta, &projection);
733 let schema = table_meta.columns.clone();
734 let iter =
735 columnar_scan::create_columnar_scan_iterator(txn, &table_meta, &columnar_scan)?;
736 return Ok((Box::new(iter), projection, schema));
737 }
738
739 let schema = table_meta.columns.clone();
741 let scan_iter = scan::create_scan_iterator(txn, &table_meta)?;
742 Ok((Box::new(scan_iter), projection, schema))
743 }
744 LogicalPlan::Filter { input, predicate } => {
745 if let LogicalPlan::Scan { table, projection } = input.as_ref()
746 && let Some(table_meta) = catalog.get_table(table)
747 && table_meta.storage_options.storage_type == StorageType::Columnar
748 {
749 let columnar_scan = columnar_scan::build_columnar_scan_for_filter(
751 table_meta,
752 projection.clone(),
753 &predicate,
754 );
755 let schema = table_meta.columns.clone();
756 let iter =
757 columnar_scan::create_columnar_scan_iterator(txn, table_meta, &columnar_scan)?;
758 return Ok((Box::new(iter), projection.clone(), schema));
759 }
760 let (input_iter, projection, schema) =
761 build_streaming_pipeline_with_policy(txn, catalog, *input, memory)?;
762 let filter_iter = FilterIterator::new(input_iter, predicate);
763 Ok((Box::new(filter_iter), projection, schema))
764 }
765 LogicalPlan::Project { input, projection } => {
766 let (mut input_iter, _input_projection, schema) =
767 build_streaming_pipeline_with_policy(txn, catalog, *input, memory)?;
768 let mut rows = Vec::new();
769 while let Some(result) = input_iter.next_row() {
770 rows.push(result?);
771 }
772 let projected = project::execute_project(rows, &projection, &schema)?;
773 let output_schema = projected
774 .columns
775 .iter()
776 .map(|col| crate::catalog::ColumnMetadata::new(&col.name, col.data_type.clone()))
777 .collect::<Vec<_>>();
778 let rows = projected
779 .rows
780 .into_iter()
781 .enumerate()
782 .map(|(idx, values)| Row::new(idx as u64, values))
783 .collect::<Vec<_>>();
784 let output_projection =
785 Projection::All(output_schema.iter().map(|col| col.name.clone()).collect());
786 let iter = iterator::VecIterator::new(rows, output_schema.clone());
787 Ok((Box::new(iter), output_projection, output_schema))
788 }
789 LogicalPlan::Join {
790 left,
791 right,
792 join_type,
793 condition,
794 using: _,
795 } => {
796 let (mut left_iter, _left_projection, left_schema) =
797 build_streaming_pipeline_with_policy(txn, catalog, *left, memory)?;
798 let mut left_rows = Vec::new();
799 while let Some(result) = left_iter.next_row() {
800 left_rows.push(result?);
801 }
802 drop(left_iter);
803 let (mut right_iter, _right_projection, right_schema) =
804 build_streaming_pipeline_with_policy(txn, catalog, *right, memory)?;
805 let mut right_rows = Vec::new();
806 while let Some(result) = right_iter.next_row() {
807 right_rows.push(result?);
808 }
809 let rows = join::execute_join_with_widths(
810 left_rows,
811 right_rows,
812 join_type,
813 condition.as_ref(),
814 left_schema.len(),
815 right_schema.len(),
816 )?;
817 let mut schema = left_schema;
818 schema.extend(right_schema);
819 let projection = Projection::All(schema.iter().map(|col| col.name.clone()).collect());
820 let iter = iterator::VecIterator::new(rows, schema.clone());
821 Ok((Box::new(iter), projection, schema))
822 }
823 LogicalPlan::Aggregate {
824 input,
825 group_keys,
826 aggregates,
827 having,
828 projection,
829 } => {
830 let (input_iter, _projection, _schema) =
831 build_streaming_pipeline_with_policy(txn, catalog, *input, memory)?;
832 let schema = aggregate::build_aggregate_schema(&group_keys, &aggregates);
833 if let Some(policy) = memory
834 && policy.spill_directory().is_some()
835 {
836 if group_keys.is_empty() {
837 let iter = aggregate::StreamingAggregateIterator::new(
838 input_iter,
839 group_keys,
840 aggregates,
841 having,
842 schema.clone(),
843 );
844 return Ok((Box::new(iter), projection, schema));
845 }
846 let order_by = group_keys
847 .iter()
848 .cloned()
849 .map(|expr| SortExpr {
850 expr,
851 asc: true,
852 nulls_first: false,
853 })
854 .collect::<Vec<_>>();
855 let sort_iter =
856 SortIterator::new_with_policy(input_iter, &order_by, Some(policy.clone()))?;
857 let iter = aggregate::StreamingAggregateIterator::new(
858 Box::new(sort_iter),
859 group_keys,
860 aggregates,
861 having,
862 schema.clone(),
863 );
864 return Ok((Box::new(iter), projection, schema));
865 }
866
867 let parallelism = std::thread::available_parallelism()
868 .map(usize::from)
869 .unwrap_or(1);
870 if !aggregate::should_use_single_for_parallel(parallelism, &aggregates) {
871 let rows = aggregate::execute_parallel_aggregate_rows_with_policy(
872 input_iter,
873 group_keys,
874 aggregates,
875 having,
876 schema.clone(),
877 parallelism,
878 memory.cloned(),
879 1_000_000,
880 )?;
881 let iter = iterator::VecIterator::new(rows, schema.clone());
882 return Ok((Box::new(iter), projection, schema));
883 }
884
885 let mut iter = aggregate::AggregateIterator::new(
886 input_iter,
887 group_keys,
888 aggregates,
889 having,
890 schema.clone(),
891 );
892 if let Some(policy) = memory {
893 iter = iter.with_memory_policy(Some(policy.clone()));
894 }
895 Ok((Box::new(iter), projection, schema))
896 }
897 LogicalPlan::Window { input, windows } => {
898 let (input_iter, _projection, _schema) =
899 build_streaming_pipeline_inner(txn, catalog, *input, memory)?;
900 let iter = window::WindowIterator::new(input_iter, windows, memory)?;
901 let schema = iter.schema().to_vec();
902 let projection =
903 Projection::All(schema.iter().map(|column| column.name.clone()).collect());
904 Ok((Box::new(iter), projection, schema))
905 }
906 LogicalPlan::Sort { input, order_by } => {
907 let (input_iter, projection, schema) =
908 build_streaming_pipeline_with_policy(txn, catalog, *input, memory)?;
909 let sort_iter = if let Some(policy) = memory {
910 SortIterator::new_with_policy(input_iter, &order_by, Some(policy.clone()))?
911 } else {
912 SortIterator::new(input_iter, &order_by)?
913 };
914 Ok((Box::new(sort_iter), projection, schema))
915 }
916 LogicalPlan::Limit {
917 input,
918 limit,
919 offset,
920 } => {
921 let (input_iter, projection, schema) =
922 build_streaming_pipeline_with_policy(txn, catalog, *input, memory)?;
923 let limit_iter = LimitIterator::new(input_iter, limit, offset);
924 Ok((Box::new(limit_iter), projection, schema))
925 }
926 other => Err(ExecutorError::UnsupportedOperation(format!(
927 "unsupported query plan: {other:?}"
928 ))),
929 }
930}
931
932fn eval_expr(expr: &crate::planner::typed_expr::TypedExpr, row: &Row) -> Result<SqlValue> {
934 let ctx = EvalContext::new(&row.values);
935 crate::executor::evaluator::evaluate(expr, &ctx)
936}
937
938fn combine_outer_for_eval(row: &Row, outer: Option<&Row>) -> Row {
939 let Some(outer) = outer else {
940 return row.clone();
941 };
942 let mut values = Vec::with_capacity(row.len() + outer.len());
943 values.extend(row.values.clone());
944 values.extend(outer.values.clone());
945 Row::new(row.row_id, values)
946}
947
948fn execute_project_with_subqueries<
949 'txn,
950 S: KVStore + 'txn,
951 C: Catalog + ?Sized,
952 T: SqlTxn<'txn, S>,
953>(
954 txn: &mut T,
955 catalog: &C,
956 rows: Vec<Row>,
957 projection: &Projection,
958 schema: &[crate::catalog::ColumnMetadata],
959 outer: Option<&Row>,
960) -> Result<QueryResult> {
961 match projection {
962 Projection::All(_) => project::execute_project(rows, projection, schema),
963 Projection::Columns(cols)
964 if outer.is_some() || cols.iter().any(|c| subquery::contains_subquery(&c.expr)) =>
965 {
966 let columns: Vec<_> = cols
967 .iter()
968 .enumerate()
969 .map(|(i, c)| column_info_from_projection(c, i))
970 .collect();
971 let mut projected_rows = Vec::with_capacity(rows.len());
972 for row in rows {
973 let eval_row = combine_outer_for_eval(&row, outer);
974 let mut values = Vec::with_capacity(cols.len());
975 for col in cols {
976 values.push(subquery::evaluate_expr_with_subqueries(
977 txn, catalog, &col.expr, &eval_row,
978 )?);
979 }
980 projected_rows.push(values);
981 }
982 Ok(QueryResult::new(columns, projected_rows))
983 }
984 Projection::Columns(_) => project::execute_project(rows, projection, schema),
985 }
986}
987
988fn column_name_from_projection(
990 projected: &crate::planner::typed_expr::ProjectedColumn,
991 idx: usize,
992) -> String {
993 use crate::planner::typed_expr::TypedExprKind;
994
995 projected
996 .alias
997 .clone()
998 .or_else(|| match &projected.expr.kind {
999 TypedExprKind::ColumnRef { column, .. } => Some(column.clone()),
1000 TypedExprKind::FunctionCall { name, args, .. }
1003 if name == "coalesce" && !args.is_empty() =>
1004 {
1005 let first_column = match &args[0].kind {
1006 TypedExprKind::ColumnRef { column, .. } => Some(column),
1007 _ => None,
1008 };
1009 first_column
1010 .filter(|column| {
1011 args.iter().all(|arg| {
1012 matches!(
1013 &arg.kind,
1014 TypedExprKind::ColumnRef { column: other, .. } if other == *column
1015 )
1016 })
1017 })
1018 .cloned()
1019 }
1020 _ => None,
1021 })
1022 .unwrap_or_else(|| format!("col_{idx}"))
1023}
1024
1025fn column_info_from_projection(
1027 projected: &crate::planner::typed_expr::ProjectedColumn,
1028 idx: usize,
1029) -> ColumnInfo {
1030 ColumnInfo::new(
1031 column_name_from_projection(projected, idx),
1032 projected.expr.resolved_type.clone(),
1033 )
1034}
1035
1036fn column_infos_from_all(
1038 schema: &[crate::catalog::ColumnMetadata],
1039 names: &[String],
1040) -> Result<Vec<ColumnInfo>> {
1041 if names.len() == schema.len() {
1042 return Ok(names
1043 .iter()
1044 .zip(schema)
1045 .map(|(name, column)| ColumnInfo::new(name.clone(), column.data_type.clone()))
1046 .collect());
1047 }
1048 names
1049 .iter()
1050 .map(|name| {
1051 let col = schema
1052 .iter()
1053 .find(|c| &c.name == name)
1054 .ok_or_else(|| ExecutorError::ColumnNotFound(name.clone()))?;
1055 Ok(ColumnInfo::new(name.clone(), col.data_type.clone()))
1056 })
1057 .collect()
1058}
1059
1060#[cfg(test)]
1061mod tests {
1062 use super::*;
1063 use crate::catalog::{ColumnMetadata, MemoryCatalog, TableMetadata};
1064 use crate::executor::ddl::create_table::execute_create_table;
1065 use crate::planner::typed_expr::TypedExpr;
1066 use crate::planner::types::ResolvedType;
1067 use crate::storage::TxnBridge;
1068 use alopex_core::kv::memory::MemoryKV;
1069 use std::sync::Arc;
1070
1071 #[test]
1072 fn execute_query_scan_only_returns_rows() {
1073 let bridge = TxnBridge::new(Arc::new(MemoryKV::new()));
1074 let mut catalog = MemoryCatalog::new();
1075 let table = TableMetadata::new(
1076 "users",
1077 vec![
1078 ColumnMetadata::new("id", ResolvedType::Integer),
1079 ColumnMetadata::new("name", ResolvedType::Text),
1080 ],
1081 );
1082 let mut ddl_txn = bridge.begin_write().unwrap();
1083 execute_create_table(&mut ddl_txn, &mut catalog, table.clone(), vec![], false).unwrap();
1084 ddl_txn.commit().unwrap();
1085
1086 let mut txn = bridge.begin_write().unwrap();
1087 crate::executor::dml::execute_insert(
1088 &mut txn,
1089 &catalog,
1090 "users",
1091 vec!["id".into(), "name".into()],
1092 vec![vec![
1093 TypedExpr::literal(
1094 crate::ast::expr::Literal::Number("1".into()),
1095 ResolvedType::Integer,
1096 crate::Span::default(),
1097 ),
1098 TypedExpr::literal(
1099 crate::ast::expr::Literal::String("alice".into()),
1100 ResolvedType::Text,
1101 crate::Span::default(),
1102 ),
1103 ]],
1104 )
1105 .unwrap();
1106
1107 let result = execute_query(
1108 &mut txn,
1109 &catalog,
1110 LogicalPlan::scan(
1111 "users".into(),
1112 Projection::All(vec!["id".into(), "name".into()]),
1113 ),
1114 )
1115 .unwrap();
1116
1117 match result {
1118 ExecutionResult::Query(q) => {
1119 assert_eq!(q.rows.len(), 1);
1120 assert_eq!(q.columns.len(), 2);
1121 assert_eq!(
1122 q.rows[0],
1123 vec![SqlValue::Integer(1), SqlValue::Text("alice".into())]
1124 );
1125 }
1126 other => panic!("unexpected result {other:?}"),
1127 }
1128 }
1129}