1use crate::ast::*;
4use crate::cancel::CancelCheck;
5use crate::plan::*;
6use crate::planner::{
7 extract_single_bound, range_scan_for_target, try_extract_eq_index_key, RangeBound, RangeTarget,
8};
9use crate::result::{QueryError, QueryResult};
10use powdb_storage::btree::IndexStats;
11use powdb_storage::catalog::{Catalog, ExpressionIndexMeta, IndexOrderDirection};
12use powdb_storage::row::{decode_column, decode_row, patch_var_column_in_place, RowLayout};
13use powdb_storage::stored_json_path::StoredJsonPathV1;
14use powdb_storage::types::*;
15use std::cmp::Reverse;
16use std::collections::{BinaryHeap, HashSet};
17use std::ops::ControlFlow;
18
19use super::compiled::*;
20use super::eval::*;
21use super::row_body_base;
22use super::{check_join_limit, mem_budget, Engine, MAX_SORT_ROWS};
23use powdb_storage::view::ViewDef;
24
25const CANCELLABLE_SORT_RUN: usize = 2_048;
29
30#[cfg(test)]
34thread_local! {
35 static GENERIC_RID_MATCH_CALLS: std::cell::Cell<u64> = const { std::cell::Cell::new(0) };
36}
37
38#[cfg(test)]
39pub(super) fn reset_generic_rid_match_calls() {
40 GENERIC_RID_MATCH_CALLS.with(|calls| calls.set(0));
41}
42
43#[cfg(test)]
44pub(super) fn generic_rid_match_calls() -> u64 {
45 GENERIC_RID_MATCH_CALLS.with(std::cell::Cell::get)
46}
47
48pub(super) fn compare_order_values(
54 left: &Value,
55 right: &Value,
56 descending: bool,
57) -> std::cmp::Ordering {
58 use std::cmp::Ordering;
59
60 match (left, right) {
61 (Value::Empty, Value::Empty) => Ordering::Equal,
62 (Value::Empty, _) => Ordering::Greater,
63 (_, Value::Empty) => Ordering::Less,
64 _ if descending => left.cmp(right).reverse(),
65 _ => left.cmp(right),
66 }
67}
68
69pub(super) fn cooperative_stable_sort_by<T, F>(
78 values: &mut [T],
79 memory_limit: usize,
80 compare: F,
81) -> Result<(), QueryError>
82where
83 F: Fn(&T, &T) -> std::cmp::Ordering,
84{
85 crate::cancel::check()?;
86 let len = values.len();
87 if len < 2 {
88 return Ok(());
89 }
90
91 let scratch_bytes = len
92 .saturating_mul(std::mem::size_of::<usize>())
93 .saturating_mul(2);
94 mem_budget::charge(scratch_bytes, memory_limit)?;
95
96 let mut order: Vec<usize> = (0..len).collect();
97 let mut scratch = vec![0usize; len];
98
99 for run in order.chunks_mut(CANCELLABLE_SORT_RUN) {
100 crate::cancel::check()?;
101 run.sort_by(|&a, &b| compare(&values[a], &values[b]));
102 crate::cancel::check()?;
103 }
104
105 let mut cancel = CancelCheck::new();
106 let mut width = CANCELLABLE_SORT_RUN;
107 while width < len {
108 let step = width.saturating_mul(2);
109 let mut start = 0usize;
110 while start < len {
111 let mid = start.saturating_add(width).min(len);
112 let end = start.saturating_add(step).min(len);
113 let (mut left, mut right, mut out) = (start, mid, start);
114
115 while left < mid && right < end {
116 cancel.tick()?;
117 if compare(&values[order[left]], &values[order[right]])
118 != std::cmp::Ordering::Greater
119 {
120 scratch[out] = order[left];
121 left += 1;
122 } else {
123 scratch[out] = order[right];
124 right += 1;
125 }
126 out += 1;
127 }
128 while left < mid {
129 cancel.tick()?;
130 scratch[out] = order[left];
131 left += 1;
132 out += 1;
133 }
134 while right < end {
135 cancel.tick()?;
136 scratch[out] = order[right];
137 right += 1;
138 out += 1;
139 }
140 start = start.saturating_add(step);
141 }
142 std::mem::swap(&mut order, &mut scratch);
143 width = step;
144 }
145
146 for (new_position, &old_position) in order.iter().enumerate() {
149 cancel.tick()?;
150 scratch[old_position] = new_position;
151 }
152 drop(order);
153 for position in 0..len {
154 while scratch[position] != position {
155 cancel.tick()?;
156 let destination = scratch[position];
157 values.swap(position, destination);
158 scratch.swap(position, destination);
159 }
160 }
161 Ok(())
162}
163
164pub(super) fn for_each_row_raw_cancellable(
167 catalog: &Catalog,
168 table: &str,
169 mut f: impl FnMut(RowId, &[u8]),
170) -> Result<(), QueryError> {
171 if !crate::cancel::has_active_install() {
177 return catalog
178 .for_each_row_raw(table, f)
179 .map_err(|err| QueryError::StorageError(err.to_string()));
180 }
181
182 let mut cancel = CancelCheck::new();
183 let mut cancel_err: Option<QueryError> = None;
184 catalog
185 .try_for_each_row_raw(table, |rid, data| {
186 if let Err(err) = cancel.tick() {
187 cancel_err = Some(err);
188 return ControlFlow::Break(());
189 }
190 f(rid, data);
191 ControlFlow::Continue(())
192 })
193 .map_err(|err| QueryError::StorageError(err.to_string()))?;
194 match cancel_err {
195 Some(err) => Err(err),
196 None => Ok(()),
197 }
198}
199
200fn resolve_expression_index(
201 catalog: &Catalog,
202 table: &str,
203 path: &StoredJsonPathV1,
204) -> Option<ExpressionIndexMeta> {
205 catalog
206 .expression_index_metadata(table)?
207 .into_iter()
208 .find(|metadata| metadata.canonical_version == 1 && metadata.json_path == *path)
209}
210
211fn expression_index_fallback(plan: &PlanNode) -> Option<PlanNode> {
212 match plan {
213 PlanNode::ExprIndexScan { table, path, key } => Some(PlanNode::Filter {
214 input: Box::new(PlanNode::SeqScan {
215 table: table.clone(),
216 }),
217 predicate: Expr::BinaryOp(
218 Box::new(stored_json_path_expr(path)),
219 BinOp::Eq,
220 Box::new(key.clone()),
221 ),
222 }),
223 PlanNode::ExprRangeScan {
224 table,
225 path,
226 start,
227 end,
228 } => Some(PlanNode::Filter {
229 input: Box::new(PlanNode::SeqScan {
230 table: table.clone(),
231 }),
232 predicate: synthesize_expr_range_predicate(path, start, end),
233 }),
234 PlanNode::OrderedExprIndexScan {
235 table,
236 path,
237 descending,
238 limit,
239 offset,
240 } => {
241 let sorted = PlanNode::Sort {
242 input: Box::new(PlanNode::SeqScan {
243 table: table.clone(),
244 }),
245 keys: vec![SortKey {
246 expr: stored_json_path_expr(path),
247 descending: *descending,
248 }],
249 };
250 let sliced = match offset {
251 Some(count) => PlanNode::Offset {
252 input: Box::new(sorted),
253 count: count.clone(),
254 },
255 None => sorted,
256 };
257 Some(PlanNode::Limit {
258 input: Box::new(sliced),
259 count: limit.clone(),
260 })
261 }
262 _ => None,
263 }
264}
265
266#[derive(Debug)]
267pub(super) struct ProvenanceRows {
268 pub(super) columns: Vec<String>,
269 pub(super) rows: Vec<Vec<Value>>,
270 source_aliases: Vec<String>,
271 provenance: Vec<Vec<Option<RowId>>>,
272}
273
274impl ProvenanceRows {
275 fn source_index(&self, alias: &str) -> Option<usize> {
276 self.source_aliases
277 .iter()
278 .position(|source| source == alias)
279 }
280}
281
282impl Engine {
283 pub(super) fn execute_expression_index_plan(
284 &self,
285 plan: &PlanNode,
286 projected_fields: Option<&[ProjectField]>,
287 ) -> Result<Option<QueryResult>, QueryError> {
288 let (table, path) = match plan {
289 PlanNode::ExprIndexScan { table, path, .. }
290 | PlanNode::ExprRangeScan { table, path, .. }
291 | PlanNode::OrderedExprIndexScan { table, path, .. } => (table, path),
292 _ => return Ok(None),
293 };
294 let Some(index) = resolve_expression_index(&self.catalog, table, path) else {
295 return Ok(None);
296 };
297 let schema = self
298 .catalog
299 .schema(table)
300 .ok_or_else(|| QueryError::TableNotFound(table.clone()))?
301 .clone();
302 let all_columns: Vec<String> = schema
303 .columns
304 .iter()
305 .map(|column| column.name.clone())
306 .collect();
307
308 let projection = match projected_fields {
309 Some(fields) => {
310 if !fields
311 .iter()
312 .all(|field| matches!(field.expr, Expr::Field(_)))
313 {
314 return Ok(None);
315 }
316 let mut indices = Vec::with_capacity(fields.len());
317 let mut columns = Vec::with_capacity(fields.len());
318 for field in fields {
319 let Expr::Field(name) = &field.expr else {
320 unreachable!("plain-field projection checked above")
321 };
322 let index =
323 schema
324 .column_index(name)
325 .ok_or_else(|| QueryError::ColumnNotFound {
326 table: table.clone(),
327 column: name.clone(),
328 })?;
329 indices.push(index);
330 columns.push(field.alias.clone().unwrap_or_else(|| name.clone()));
331 }
332 Some((indices, columns))
333 }
334 None => None,
335 };
336
337 let (rids, range) = match plan {
338 PlanNode::ExprIndexScan { key, .. } => {
339 let key = literal_to_value(key)?;
340 let rids = if key.is_empty() {
341 self.catalog
342 .expression_index_btree(table, index.index_id)
343 .ok_or_else(|| {
344 QueryError::Execution("expression index disappeared".to_string())
345 })?
346 .empty_rids()
347 .to_vec()
348 } else {
349 self.catalog
350 .expression_index_lookup_all(table, index.index_id, &key)
351 .map_err(|error| QueryError::StorageError(error.to_string()))?
352 };
353 (rids, None)
354 }
355 PlanNode::ExprRangeScan { start, end, .. } => {
356 let start_value = start
357 .as_ref()
358 .map(|(expr, _)| literal_to_value(expr))
359 .transpose()?;
360 let end_value = end
361 .as_ref()
362 .map(|(expr, _)| literal_to_value(expr))
363 .transpose()?;
364 let rids = self
365 .catalog
366 .expression_index_range_rids(
367 table,
368 index.index_id,
369 start_value.as_ref(),
370 end_value.as_ref(),
371 )
372 .map_err(|error| QueryError::StorageError(error.to_string()))?;
373 (
374 rids,
375 Some((
376 start_value,
377 start.as_ref().is_none_or(|(_, inclusive)| *inclusive),
378 end_value,
379 end.as_ref().is_none_or(|(_, inclusive)| *inclusive),
380 )),
381 )
382 }
383 PlanNode::OrderedExprIndexScan {
384 descending,
385 limit,
386 offset,
387 ..
388 } => {
389 let Expr::Literal(Literal::Int(limit)) = limit else {
390 return Err(QueryError::Execution(
391 "expression-index limit must be a non-negative integer".to_string(),
392 ));
393 };
394 let offset = match offset {
395 Some(Expr::Literal(Literal::Int(offset))) if *offset >= 0 => *offset as usize,
396 None => 0,
397 _ => {
398 return Err(QueryError::Execution(
399 "expression-index offset must be a non-negative integer".to_string(),
400 ));
401 }
402 };
403 if *limit < 0 {
404 return Err(QueryError::Execution(
405 "expression-index limit must be a non-negative integer".to_string(),
406 ));
407 }
408 let rids = self
409 .catalog
410 .expression_index_ordered_rids_bounded(
411 table,
412 index.index_id,
413 if *descending {
414 IndexOrderDirection::Desc
415 } else {
416 IndexOrderDirection::Asc
417 },
418 offset,
419 *limit as usize,
420 )
421 .map_err(|error| QueryError::StorageError(error.to_string()))?;
422 (rids, None)
423 }
424 _ => unreachable!("expression-index plan checked above"),
425 };
426
427 let root_index =
428 schema
429 .column_index(&path.column)
430 .ok_or_else(|| QueryError::ColumnNotFound {
431 table: table.clone(),
432 column: path.column.clone(),
433 })?;
434 let path_expr = stored_json_path_expr(path);
435 let mut rows = Vec::with_capacity(rids.len());
436 let mut cancel = CancelCheck::new();
437 for rid in rids {
438 cancel.tick()?;
439 match &projection {
440 Some((projected_indices, _)) => {
441 let mut fetch_indices = projected_indices.clone();
442 let root_position = fetch_indices.iter().position(|index| *index == root_index);
443 let root_position = match root_position {
444 Some(position) => position,
445 None => {
446 fetch_indices.push(root_index);
447 fetch_indices.len() - 1
448 }
449 };
450 let Some(mut fetched) = self
451 .catalog
452 .get_projected(table, rid, &fetch_indices)
453 .map_err(|error| QueryError::StorageError(error.to_string()))?
454 else {
455 continue;
456 };
457 if let Some((start, start_inclusive, end, end_inclusive)) = &range {
458 let value = eval_expr(
459 &path_expr,
460 std::slice::from_ref(&fetched[root_position]),
461 std::slice::from_ref(&path.column),
462 );
463 if value.is_empty()
464 || !range_matches(&value, start, *start_inclusive, end, *end_inclusive)
465 {
466 continue;
467 }
468 }
469 fetched.truncate(projected_indices.len());
470 rows.push(fetched);
471 }
472 None => {
473 let Some(row) = self.catalog.get(table, rid) else {
474 continue;
475 };
476 if let Some((start, start_inclusive, end, end_inclusive)) = &range {
477 let value = eval_expr(&path_expr, &row, &all_columns);
478 if value.is_empty()
479 || !range_matches(&value, start, *start_inclusive, end, *end_inclusive)
480 {
481 continue;
482 }
483 }
484 rows.push(row);
485 }
486 }
487 }
488
489 let columns = projection
490 .map(|(_, columns)| columns)
491 .unwrap_or(all_columns);
492 Ok(Some(QueryResult::Rows { columns, rows }))
493 }
494
495 pub(super) fn try_filter_index_residual_fast(
509 &self,
510 input: &PlanNode,
511 predicate: &Expr,
512 ) -> Result<Option<QueryResult>, QueryError> {
513 if contains_subquery(predicate) {
516 return Ok(None);
517 }
518 let (table, rids) = match input {
519 PlanNode::IndexScan { table, column, key } => {
520 let Some(tbl) = self.catalog.get_table(table) else {
521 return Ok(None);
522 };
523 if !tbl.has_index(column) {
524 return Ok(None);
525 }
526 let key_value = literal_to_value(key)?;
527 (table.as_str(), tbl.index_lookup_all(column, &key_value))
528 }
529 PlanNode::ExprIndexScan { table, path, key } => {
530 let Some(index) = resolve_expression_index(&self.catalog, table, path) else {
531 return Ok(None);
532 };
533 let key_value = literal_to_value(key)?;
534 let rids = if key_value.is_empty() {
535 self.catalog
536 .expression_index_btree(table, index.index_id)
537 .ok_or_else(|| {
538 QueryError::Execution("expression index disappeared".to_string())
539 })?
540 .empty_rids()
541 .to_vec()
542 } else {
543 self.catalog
544 .expression_index_lookup_all(table, index.index_id, &key_value)
545 .map_err(|error| QueryError::StorageError(error.to_string()))?
546 };
547 (table.as_str(), rids)
548 }
549 _ => return Ok(None),
550 };
551
552 let schema = self
553 .catalog
554 .schema(table)
555 .ok_or_else(|| QueryError::TableNotFound(table.to_string()))?
556 .clone();
557 let all_columns: Vec<String> = schema.columns.iter().map(|c| c.name.clone()).collect();
558 let residual_indices = predicate_column_indices_json(predicate, &all_columns);
561 let residual_names: Vec<String> = residual_indices
562 .iter()
563 .map(|&index| all_columns[index].clone())
564 .collect();
565
566 let mut rows: Vec<Vec<Value>> = Vec::new();
567 let mut cancel = CancelCheck::new();
570 for rid in rids {
571 cancel.tick()?;
572 let Some(sparse) = self
573 .catalog
574 .get_projected(table, rid, &residual_indices)
575 .map_err(|error| QueryError::StorageError(error.to_string()))?
576 else {
577 continue;
578 };
579 if eval_predicate(predicate, &sparse, &residual_names) {
580 if let Some(full) = self.catalog.get(table, rid) {
581 rows.push(full);
582 }
583 }
584 }
585 Ok(Some(QueryResult::Rows {
586 columns: all_columns,
587 rows,
588 }))
589 }
590
591 fn charge_provenance(&self, rows: &ProvenanceRows) -> Result<(), QueryError> {
592 let aliases =
593 rows.source_aliases
594 .iter()
595 .fold(std::mem::size_of::<Vec<String>>(), |total, alias| {
596 total
597 .saturating_add(std::mem::size_of::<String>())
598 .saturating_add(alias.capacity())
599 });
600 let per_row = std::mem::size_of::<Vec<Option<RowId>>>().saturating_add(
601 rows.source_aliases
602 .len()
603 .saturating_mul(std::mem::size_of::<Option<RowId>>()),
604 );
605 mem_budget::charge(
606 aliases.saturating_add(rows.provenance.len().saturating_mul(per_row)),
607 self.query_memory_limit(),
608 )
609 }
610
611 fn provenance_scan(
612 &self,
613 table: &str,
614 alias: &str,
615 qualify_columns: bool,
616 ) -> Result<ProvenanceRows, QueryError> {
617 let schema = self
618 .catalog
619 .schema(table)
620 .ok_or_else(|| QueryError::TableNotFound(table.to_string()))?
621 .clone();
622 let columns = schema
623 .columns
624 .iter()
625 .map(|column| {
626 if qualify_columns {
627 format!("{alias}.{}", column.name)
628 } else {
629 column.name.clone()
630 }
631 })
632 .collect();
633 let mut rows = Vec::new();
634 let mut provenance = Vec::new();
635 let mut cancel = CancelCheck::new();
636 for (rid, row) in self
637 .catalog
638 .scan(table)
639 .map_err(|error| QueryError::StorageError(error.to_string()))?
640 {
641 cancel.tick()?;
642 rows.push(row);
643 provenance.push(vec![Some(rid)]);
644 }
645 let result = ProvenanceRows {
646 columns,
647 rows,
648 source_aliases: vec![alias.to_string()],
649 provenance,
650 };
651 Ok(result)
652 }
653
654 pub(super) fn materialize_rows_with_provenance(
655 &self,
656 plan: &PlanNode,
657 ) -> Result<ProvenanceRows, QueryError> {
658 let result = match plan {
659 PlanNode::SeqScan { table } => self.provenance_scan(table, table, false)?,
660 PlanNode::AliasScan { table, alias } => self.provenance_scan(table, alias, true)?,
661 PlanNode::IndexScan { table, column, key } => {
662 let fallback = PlanNode::Filter {
663 input: Box::new(PlanNode::SeqScan {
664 table: table.clone(),
665 }),
666 predicate: Expr::BinaryOp(
667 Box::new(Expr::Field(column.clone())),
668 BinOp::Eq,
669 Box::new(key.clone()),
670 ),
671 };
672 self.materialize_rows_with_provenance(&fallback)?
673 }
674 PlanNode::RangeScan {
675 table,
676 column,
677 start,
678 end,
679 } => {
680 let fallback = PlanNode::Filter {
681 input: Box::new(PlanNode::SeqScan {
682 table: table.clone(),
683 }),
684 predicate: synthesize_range_predicate(column, start, end),
685 };
686 self.materialize_rows_with_provenance(&fallback)?
687 }
688 PlanNode::ExprIndexScan { .. }
689 | PlanNode::ExprRangeScan { .. }
690 | PlanNode::OrderedExprIndexScan { .. } => {
691 let fallback = expression_index_fallback(plan)
692 .expect("expression-index branch always has a fallback");
693 self.materialize_rows_with_provenance(&fallback)?
694 }
695 PlanNode::Filter { input, predicate } => {
696 if contains_subquery(predicate) {
697 return Err(QueryError::Execution(
698 "symmetric aggregation over a subquery filter is not supported; use raw"
699 .to_string(),
700 ));
701 }
702 let input = self.materialize_rows_with_provenance(input)?;
703 let mut rows = Vec::new();
704 let mut provenance = Vec::new();
705 let mut cancel = CancelCheck::new();
706 for (row, row_provenance) in input.rows.into_iter().zip(input.provenance) {
707 cancel.tick()?;
708 if eval_predicate(predicate, &row, &input.columns) {
709 rows.push(row);
710 provenance.push(row_provenance);
711 }
712 }
713 ProvenanceRows {
714 columns: input.columns,
715 rows,
716 source_aliases: input.source_aliases,
717 provenance,
718 }
719 }
720 PlanNode::Project { input, fields } => {
721 let input = self.materialize_rows_with_provenance(input)?;
722 let columns = fields
723 .iter()
724 .map(|field| {
725 field.alias.clone().unwrap_or_else(|| match &field.expr {
726 Expr::Field(name) => name.clone(),
727 Expr::QualifiedField { qualifier, field } => {
728 format!("{qualifier}.{field}")
729 }
730 _ => expression_output_name(&field.expr),
731 })
732 })
733 .collect();
734 let mut rows = Vec::with_capacity(input.rows.len());
735 let mut cancel = CancelCheck::new();
736 for row in &input.rows {
737 cancel.tick()?;
738 rows.push(
739 fields
740 .iter()
741 .map(|field| eval_expr(&field.expr, row, &input.columns))
742 .collect(),
743 );
744 }
745 ProvenanceRows {
746 columns,
747 rows,
748 source_aliases: input.source_aliases,
749 provenance: input.provenance,
750 }
751 }
752 PlanNode::Sort { input, keys } => {
753 let input = self.materialize_rows_with_provenance(input)?;
754 if input.rows.len() > MAX_SORT_ROWS {
755 return Err(QueryError::SortLimitExceeded);
756 }
757 self.charge_rows(&input.rows)?;
758 let mut paired: Vec<_> = input.rows.into_iter().zip(input.provenance).collect();
759 cooperative_stable_sort_by(
760 &mut paired,
761 self.query_memory_limit(),
762 |(left, _), (right, _)| {
763 for key in keys {
764 let left_value = eval_expr(&key.expr, left, &input.columns);
765 let right_value = eval_expr(&key.expr, right, &input.columns);
766 let comparison =
767 compare_order_values(&left_value, &right_value, key.descending);
768 if comparison != std::cmp::Ordering::Equal {
769 return comparison;
770 }
771 }
772 std::cmp::Ordering::Equal
773 },
774 )?;
775 let (rows, provenance) = paired.into_iter().unzip();
776 ProvenanceRows {
777 columns: input.columns,
778 rows,
779 source_aliases: input.source_aliases,
780 provenance,
781 }
782 }
783 PlanNode::Limit { input, count } | PlanNode::Offset { input, count } => {
784 let input_rows = self.materialize_rows_with_provenance(input)?;
785 let Expr::Literal(Literal::Int(count)) = count else {
786 return Err(QueryError::Execution(
787 "limit/offset must be an integer literal".to_string(),
788 ));
789 };
790 let count = *count as usize;
791 let is_limit = matches!(plan, PlanNode::Limit { .. });
792 let iterator = input_rows.rows.into_iter().zip(input_rows.provenance);
793 let (rows, provenance) = if is_limit {
794 iterator.take(count).unzip()
795 } else {
796 iterator.skip(count).unzip()
797 };
798 ProvenanceRows {
799 columns: input_rows.columns,
800 rows,
801 source_aliases: input_rows.source_aliases,
802 provenance,
803 }
804 }
805 PlanNode::Distinct { input } => {
806 let input = self.materialize_rows_with_provenance(input)?;
807 let mut seen = HashSet::new();
808 let mut rows = Vec::new();
809 let mut provenance = Vec::new();
810 let mut cancel = CancelCheck::new();
811 for (row, row_provenance) in input.rows.into_iter().zip(input.provenance) {
812 cancel.tick()?;
813 if seen.insert(row.clone()) {
814 rows.push(row);
815 provenance.push(row_provenance);
816 }
817 }
818 ProvenanceRows {
819 columns: input.columns,
820 rows,
821 source_aliases: input.source_aliases,
822 provenance,
823 }
824 }
825 PlanNode::Union { left, right, all } => {
826 let mut left_rows = self.materialize_rows_with_provenance(left)?;
827 let right_rows = self.materialize_rows_with_provenance(right)?;
828 if left_rows.columns.len() != right_rows.columns.len() {
829 return Err(QueryError::Execution(
830 "union sides must have the same number of columns".to_string(),
831 ));
832 }
833 if left_rows.source_aliases != right_rows.source_aliases {
834 return Err(QueryError::Execution(
835 "symmetric aggregation over union requires matching source aliases; use raw"
836 .to_string(),
837 ));
838 }
839 left_rows.rows.extend(right_rows.rows);
840 left_rows.provenance.extend(right_rows.provenance);
841 if !all {
842 let mut seen = HashSet::new();
843 let mut rows = Vec::new();
844 let mut provenance = Vec::new();
845 for (row, row_provenance) in
846 left_rows.rows.into_iter().zip(left_rows.provenance)
847 {
848 if seen.insert(row.clone()) {
849 rows.push(row);
850 provenance.push(row_provenance);
851 }
852 }
853 left_rows.rows = rows;
854 left_rows.provenance = provenance;
855 }
856 left_rows
857 }
858 PlanNode::NestedLoopJoin {
859 left,
860 right,
861 on,
862 kind,
863 } => {
864 let left = self.materialize_rows_with_provenance(left)?;
865 let right = self.materialize_rows_with_provenance(right)?;
866 execute_provenance_join(
867 left,
868 right,
869 on.as_ref(),
870 *kind,
871 self.nested_loop_pair_limit,
872 )?
873 }
874 _ => {
875 return Err(QueryError::Execution(
876 "symmetric aggregation input shape is not supported; use raw".to_string(),
877 ));
878 }
879 };
880 self.charge_provenance(&result)?;
881 Ok(result)
882 }
883
884 pub(super) fn introspect_list_types(&self) -> Result<QueryResult, QueryError> {
887 let rows: Vec<Vec<Value>> = self
888 .catalog
889 .list_tables()
890 .iter()
891 .map(|name| {
892 let cols = self
893 .catalog
894 .schema(name)
895 .map(|s| s.columns.len())
896 .unwrap_or(0) as i64;
897 vec![Value::Str((*name).to_string()), Value::Int(cols)]
898 })
899 .collect();
900 Ok(QueryResult::Rows {
901 columns: vec!["name".to_string(), "columns".to_string()],
902 rows,
903 })
904 }
905
906 pub(super) fn introspect_describe(&self, table: &str) -> Result<QueryResult, QueryError> {
909 let schema = self
910 .catalog
911 .schema(table)
912 .ok_or_else(|| QueryError::TableNotFound(table.to_string()))?;
913 let rows: Vec<Vec<Value>> = schema
914 .columns
915 .iter()
916 .map(|c| {
917 let index = if self.catalog.has_index(table, &c.name) {
918 match self.catalog.is_index_unique(table, &c.name) {
919 Some(true) => "unique",
920 _ => "index",
921 }
922 } else {
923 ""
924 };
925 vec![
926 Value::Str(c.name.clone()),
927 Value::Str(type_id_to_name(c.type_id).to_string()),
928 Value::Bool(!c.required),
929 Value::Str(index.to_string()),
930 ]
931 })
932 .collect();
933 Ok(QueryResult::Rows {
934 columns: vec![
935 "column".to_string(),
936 "type".to_string(),
937 "nullable".to_string(),
938 "index".to_string(),
939 ],
940 rows,
941 })
942 }
943
944 pub fn execute_plan(&mut self, plan: &PlanNode) -> Result<QueryResult, QueryError> {
945 validate_no_stray_aggregates(plan)?;
951 validate_json_path_types(&self.catalog, plan)?;
952 match plan {
953 PlanNode::ExprIndexScan { .. }
954 | PlanNode::ExprRangeScan { .. }
955 | PlanNode::OrderedExprIndexScan { .. } => {
956 if let Some(result) = self.execute_expression_index_plan(plan, None)? {
957 return Ok(result);
958 }
959 let fallback = expression_index_fallback(plan)
960 .expect("expression-index branch always has a fallback");
961 self.execute_plan(&fallback)
962 }
963 PlanNode::SeqScan { table } => {
964 if self.view_registry.is_dirty(table) {
966 self.refresh_view(table)?;
967 }
968 let schema = self
969 .catalog
970 .schema(table)
971 .ok_or_else(|| QueryError::TableNotFound(table.to_string()))?
972 .clone();
973 let columns: Vec<String> = schema.columns.iter().map(|c| c.name.clone()).collect();
974 let mut cancel = CancelCheck::new();
977 let mut rows: Vec<Vec<Value>> = Vec::new();
978 for (_, row) in self
979 .catalog
980 .scan(table)
981 .map_err(|e| QueryError::StorageError(e.to_string()))?
982 {
983 cancel.tick()?;
984 rows.push(row);
985 }
986 Ok(QueryResult::Rows { columns, rows })
987 }
988
989 PlanNode::Filter { input, predicate } => {
990 let materialized;
994 let predicate = if contains_subquery(predicate) {
995 materialized = self.materialize_subqueries(predicate)?;
996 &materialized
997 } else {
998 predicate
999 };
1000
1001 if contains_subquery(predicate) {
1003 let result = self.execute_plan(input)?;
1004 return match result {
1005 QueryResult::Rows { columns, rows } => {
1006 let mut filtered = Vec::new();
1007 let mut cancel = CancelCheck::new();
1010 for row in rows {
1011 cancel.tick()?;
1012 let row_pred =
1013 self.materialize_correlated_for_row(predicate, &row, &columns)?;
1014 if eval_predicate(&row_pred, &row, &columns) {
1015 filtered.push(row);
1016 }
1017 }
1018 Ok(QueryResult::Rows {
1019 columns,
1020 rows: filtered,
1021 })
1022 }
1023 _ => Err("filter requires row input".into()),
1024 };
1025 }
1026
1027 if matches!(
1031 input.as_ref(),
1032 PlanNode::IndexScan { .. } | PlanNode::ExprIndexScan { .. }
1033 ) {
1034 if let Some(result) = self.try_filter_index_residual_fast(input, predicate)? {
1035 return Ok(result);
1036 }
1037 }
1038
1039 if let PlanNode::SeqScan { table } = input.as_ref() {
1047 if !self.catalog.table_has_overflow(table) {
1048 if self.view_registry.is_dirty(table) {
1050 self.refresh_view(table)?;
1051 }
1052 let schema = self
1053 .catalog
1054 .schema(table)
1055 .ok_or_else(|| QueryError::TableNotFound(table.to_string()))?
1056 .clone();
1057 let columns: Vec<String> =
1058 schema.columns.iter().map(|c| c.name.clone()).collect();
1059 let fast = FastLayout::new(&schema);
1060 let row_layout = RowLayout::new(&schema);
1061 let mut rows: Vec<Vec<Value>> = Vec::with_capacity(64);
1065
1066 let mut cancel = CancelCheck::new();
1073 let mut cancel_err: Option<QueryError> = None;
1074 if let Some(compiled) =
1075 compile_predicate(predicate, &columns, &fast, &schema)
1076 {
1077 self.catalog
1078 .try_for_each_row_raw(table, |_rid, data| {
1079 if let Err(e) = cancel.tick() {
1080 cancel_err = Some(e);
1081 return ControlFlow::Break(());
1082 }
1083 if compiled(data) {
1084 rows.push(decode_row(&schema, data));
1085 }
1086 ControlFlow::Continue(())
1087 })
1088 .map_err(|e| QueryError::StorageError(e.to_string()))?;
1089 } else {
1090 let pred_cols = predicate_column_indices_json(predicate, &columns);
1091 self.catalog
1092 .try_for_each_row_raw(table, |_rid, data| {
1093 if let Err(e) = cancel.tick() {
1094 cancel_err = Some(e);
1095 return ControlFlow::Break(());
1096 }
1097 let pred_row =
1098 decode_selective(&schema, &row_layout, data, &pred_cols);
1099 if eval_predicate(predicate, &pred_row, &columns) {
1100 rows.push(decode_row(&schema, data));
1101 }
1102 ControlFlow::Continue(())
1103 })
1104 .map_err(|e| QueryError::StorageError(e.to_string()))?;
1105 }
1106 if let Some(e) = cancel_err {
1107 return Err(e);
1108 }
1109
1110 return Ok(QueryResult::Rows { columns, rows });
1111 }
1112 }
1113
1114 let result = self.execute_plan(input)?;
1116 match result {
1117 QueryResult::Rows { columns, rows } => {
1118 let mut cancel = CancelCheck::new();
1119 let mut filtered: Vec<Vec<Value>> = Vec::new();
1120 for row in rows {
1121 cancel.tick()?;
1122 if eval_predicate(predicate, &row, &columns) {
1123 filtered.push(row);
1124 }
1125 }
1126 Ok(QueryResult::Rows {
1127 columns,
1128 rows: filtered,
1129 })
1130 }
1131 _ => Err("filter requires row input".into()),
1132 }
1133 }
1134
1135 PlanNode::Project { input, fields } => {
1136 if matches!(
1137 input.as_ref(),
1138 PlanNode::ExprIndexScan { .. }
1139 | PlanNode::ExprRangeScan { .. }
1140 | PlanNode::OrderedExprIndexScan { .. }
1141 ) {
1142 if let Some(result) = self.execute_expression_index_plan(input, Some(fields))? {
1143 return Ok(result);
1144 }
1145 }
1146 if let PlanNode::IndexScan { table, column, key } = input.as_ref() {
1149 let schema = self
1150 .catalog
1151 .schema(table)
1152 .ok_or_else(|| QueryError::TableNotFound(table.to_string()))?
1153 .clone();
1154 let all_columns: Vec<String> =
1155 schema.columns.iter().map(|c| c.name.clone()).collect();
1156 let key_value = literal_to_value(key)?;
1157 let tbl = self
1158 .catalog
1159 .get_table(table)
1160 .ok_or_else(|| QueryError::TableNotFound(table.to_string()))?;
1161
1162 let proj_columns: Vec<String> = fields
1163 .iter()
1164 .map(|f| {
1165 f.alias.clone().unwrap_or_else(|| match &f.expr {
1166 Expr::Field(name) => name.clone(),
1167 _ => "?".into(),
1168 })
1169 })
1170 .collect();
1171
1172 let proj_indices: Vec<usize> = fields
1174 .iter()
1175 .filter_map(|f| {
1176 if let Expr::Field(name) = &f.expr {
1177 all_columns.iter().position(|c| c == name)
1178 } else {
1179 None
1180 }
1181 })
1182 .collect();
1183
1184 let all_plain_fields = fields.iter().all(|f| matches!(f.expr, Expr::Field(_)));
1189 if tbl.has_index(column) && all_plain_fields {
1190 let rids = tbl.index_lookup_all(column, &key_value);
1191 let mut rows: Vec<Vec<Value>> = Vec::with_capacity(rids.len());
1192 let mut cancel = CancelCheck::new();
1193 for rid in rids {
1194 cancel.tick()?;
1195 if let Some(full) = tbl.get(rid) {
1201 let row: Vec<Value> =
1202 proj_indices.iter().map(|&ci| full[ci].clone()).collect();
1203 rows.push(row);
1204 }
1205 }
1206 return Ok(QueryResult::Rows {
1207 columns: proj_columns,
1208 rows,
1209 });
1210 }
1211 }
1212
1213 if let PlanNode::Limit {
1218 input: inner,
1219 count: limit_expr,
1220 } = input.as_ref()
1221 {
1222 if let PlanNode::Sort {
1223 input: sort_input,
1224 keys,
1225 } = inner.as_ref()
1226 {
1227 if keys.len() == 1 {
1229 if let Expr::Field(sort_field) = &keys[0].expr {
1230 let descending = keys[0].descending;
1231 let limit = match limit_expr {
1232 Expr::Literal(Literal::Int(v)) if *v >= 0 => *v as usize,
1233 _ => usize::MAX,
1234 };
1235 let (table_opt, pred_opt): (Option<&str>, Option<&Expr>) =
1236 match sort_input.as_ref() {
1237 PlanNode::SeqScan { table } => (Some(table.as_str()), None),
1238 PlanNode::Filter {
1239 input: fi,
1240 predicate,
1241 } => {
1242 if let PlanNode::SeqScan { table } = fi.as_ref() {
1243 (Some(table.as_str()), Some(predicate))
1244 } else {
1245 (None, None)
1246 }
1247 }
1248 _ => (None, None),
1249 };
1250 if let Some(table) = table_opt {
1251 if let Some(result) = self.project_filter_sort_limit_fast(
1252 table, fields, sort_field, descending, limit, pred_opt,
1253 )? {
1254 return Ok(result);
1255 }
1256 }
1257 }
1258 }
1259 }
1260 if let PlanNode::Filter {
1263 input: fi,
1264 predicate,
1265 } = inner.as_ref()
1266 {
1267 if let PlanNode::SeqScan { table } = fi.as_ref() {
1268 let limit = match limit_expr {
1269 Expr::Literal(Literal::Int(v)) if *v >= 0 => *v as usize,
1270 _ => usize::MAX,
1271 };
1272 if let Some(result) = self.project_filter_limit_fast(
1273 table,
1274 fields,
1275 limit,
1276 Some(predicate),
1277 )? {
1278 return Ok(result);
1279 }
1280 }
1281 }
1282 if let PlanNode::SeqScan { table } = inner.as_ref() {
1284 let limit = match limit_expr {
1285 Expr::Literal(Literal::Int(v)) if *v >= 0 => *v as usize,
1286 _ => usize::MAX,
1287 };
1288 if let Some(result) =
1289 self.project_filter_limit_fast(table, fields, limit, None)?
1290 {
1291 return Ok(result);
1292 }
1293 }
1294 }
1295
1296 if let PlanNode::Filter {
1307 input: fi,
1308 predicate,
1309 } = input.as_ref()
1310 {
1311 if let PlanNode::SeqScan { table } = fi.as_ref() {
1312 if let Some(result) = self.project_filter_limit_fast(
1313 table,
1314 fields,
1315 usize::MAX,
1316 Some(predicate),
1317 )? {
1318 return Ok(result);
1319 }
1320 }
1321 }
1322
1323 if let PlanNode::SeqScan { table } = input.as_ref() {
1327 if let Some(result) =
1328 self.project_filter_limit_fast(table, fields, usize::MAX, None)?
1329 {
1330 return Ok(result);
1331 }
1332 }
1333
1334 let result = self.execute_plan(input)?;
1335 match result {
1336 QueryResult::Rows { columns, rows } => {
1337 let proj_columns: Vec<String> = fields
1338 .iter()
1339 .map(|f| {
1340 f.alias.clone().unwrap_or_else(|| match &f.expr {
1341 Expr::Field(name) => name.clone(),
1342 Expr::QualifiedField { qualifier, field } => {
1346 format!("{qualifier}.{field}")
1347 }
1348 _ => "?".into(),
1349 })
1350 })
1351 .collect();
1352 let mut cancel = CancelCheck::new();
1353 let mut proj_rows: Vec<Vec<Value>> = Vec::with_capacity(rows.len());
1354 for row in &rows {
1355 cancel.tick()?;
1356 proj_rows.push(
1357 fields
1358 .iter()
1359 .map(|f| eval_expr(&f.expr, row, &columns))
1360 .collect(),
1361 );
1362 }
1363 Ok(QueryResult::Rows {
1364 columns: proj_columns,
1365 rows: proj_rows,
1366 })
1367 }
1368 _ => Err("project requires row input".into()),
1369 }
1370 }
1371
1372 PlanNode::Sort { input, keys } => {
1373 let result = self.execute_plan(input)?;
1374 match result {
1375 QueryResult::Rows { columns, mut rows } => {
1376 if rows.len() > MAX_SORT_ROWS {
1380 return Err(QueryError::SortLimitExceeded);
1381 }
1382 self.charge_rows(&rows)?;
1383 let key_specs: Vec<(Option<usize>, &Expr, bool)> = keys
1384 .iter()
1385 .map(|k| {
1386 let stored_name = match &k.expr {
1387 Expr::Field(name) => Some(name.clone()),
1388 Expr::QualifiedField { qualifier, field } => {
1389 Some(format!("{qualifier}.{field}"))
1390 }
1391 _ => None,
1392 };
1393 let index = stored_name
1394 .as_ref()
1395 .and_then(|name| columns.iter().position(|c| c == name));
1396 if let Some(name) = stored_name {
1397 if index.is_none() {
1398 return Err(QueryError::ColumnNotFound {
1399 table: String::new(),
1400 column: name,
1401 });
1402 }
1403 }
1404 Ok((index, &k.expr, k.descending))
1405 })
1406 .collect::<Result<_, QueryError>>()?;
1407 cooperative_stable_sort_by(&mut rows, self.query_memory_limit, |a, b| {
1408 for &(col_idx, expr, descending) in &key_specs {
1409 let (left_value, right_value) = match col_idx {
1410 Some(index) => (&a[index], &b[index]),
1411 None => {
1412 let left = eval_expr(expr, a, &columns);
1413 let right = eval_expr(expr, b, &columns);
1414 let cmp = compare_order_values(&left, &right, descending);
1415 if cmp != std::cmp::Ordering::Equal {
1416 return cmp;
1417 }
1418 continue;
1419 }
1420 };
1421 let cmp = compare_order_values(left_value, right_value, descending);
1422 if cmp != std::cmp::Ordering::Equal {
1423 return cmp;
1424 }
1425 }
1426 std::cmp::Ordering::Equal
1427 })?;
1428 Ok(QueryResult::Rows { columns, rows })
1429 }
1430 _ => Err("sort requires row input".into()),
1431 }
1432 }
1433
1434 PlanNode::Limit { input, count } => {
1435 let result = self.execute_plan(input)?;
1436 let n = match count {
1437 Expr::Literal(Literal::Int(v)) => *v as usize,
1438 _ => return Err("limit must be integer literal".into()),
1439 };
1440 match result {
1441 QueryResult::Rows { columns, rows } => {
1442 let mut cancel = CancelCheck::new();
1443 let mut limited = Vec::with_capacity(n.min(rows.len()));
1444 for row in rows.into_iter().take(n) {
1445 cancel.tick()?;
1446 limited.push(row);
1447 }
1448 Ok(QueryResult::Rows {
1449 columns,
1450 rows: limited,
1451 })
1452 }
1453 _ => Err("limit requires row input".into()),
1454 }
1455 }
1456
1457 PlanNode::Offset { input, count } => {
1458 let result = self.execute_plan(input)?;
1459 let n = match count {
1460 Expr::Literal(Literal::Int(v)) => *v as usize,
1461 _ => return Err("offset must be integer literal".into()),
1462 };
1463 match result {
1464 QueryResult::Rows { columns, rows } => {
1465 let mut cancel = CancelCheck::new();
1466 let mut offset = Vec::with_capacity(rows.len().saturating_sub(n));
1467 for (index, row) in rows.into_iter().enumerate() {
1468 cancel.tick()?;
1469 if index >= n {
1470 offset.push(row);
1471 }
1472 }
1473 Ok(QueryResult::Rows {
1474 columns,
1475 rows: offset,
1476 })
1477 }
1478 _ => Err("offset requires row input".into()),
1479 }
1480 }
1481
1482 PlanNode::Aggregate {
1483 input,
1484 function,
1485 argument,
1486 mode: _,
1487 provenance_alias,
1488 } => {
1489 if let Some(provenance_alias) = provenance_alias {
1490 let input = self.materialize_rows_with_provenance(input)?;
1491 self.charge_rows(&input.rows)?;
1492 return aggregate_rows_with_provenance(
1493 *function,
1494 argument.as_ref(),
1495 &input,
1496 provenance_alias,
1497 self.query_memory_limit(),
1498 );
1499 }
1500 if *function == AggFunc::Count {
1502 if let PlanNode::SeqScan { table } = input.as_ref() {
1506 if !self.catalog.table_has_overflow(table) {
1507 if self.view_registry.is_dirty(table) {
1511 self.refresh_view(table)?;
1512 }
1513 let mut count: i64 = 0;
1514 for_each_row_raw_cancellable(&self.catalog, table, |_rid, _data| {
1515 count += 1;
1516 })?;
1517 return Ok(QueryResult::Scalar(Value::Int(count)));
1518 }
1519 }
1520 if let PlanNode::Filter {
1528 input: inner,
1529 predicate,
1530 } = input.as_ref()
1531 {
1532 if let PlanNode::SeqScan { table } = inner.as_ref() {
1533 if self.view_registry.is_dirty(table) {
1534 self.refresh_view(table)?;
1535 }
1536 }
1537 if let (PlanNode::SeqScan { table }, false) =
1538 (inner.as_ref(), contains_subquery(predicate))
1539 {
1540 if !self.catalog.table_has_overflow(table) {
1541 let schema = self
1542 .catalog
1543 .schema(table)
1544 .ok_or_else(|| QueryError::TableNotFound(table.to_string()))?
1545 .clone();
1546 let columns: Vec<String> =
1547 schema.columns.iter().map(|c| c.name.clone()).collect();
1548 let fast = FastLayout::new(&schema);
1549 let row_layout = RowLayout::new(&schema);
1550
1551 if let Some(compiled) =
1554 compile_predicate(predicate, &columns, &fast, &schema)
1555 {
1556 let mut count: i64 = 0;
1557 for_each_row_raw_cancellable(
1558 &self.catalog,
1559 table,
1560 |_rid, data| {
1561 if compiled(data) {
1562 count += 1;
1563 }
1564 },
1565 )?;
1566 return Ok(QueryResult::Scalar(Value::Int(count)));
1567 }
1568
1569 let pred_cols = predicate_column_indices_json(predicate, &columns);
1571 let mut count: i64 = 0;
1572 for_each_row_raw_cancellable(
1573 &self.catalog,
1574 table,
1575 |_rid, data| {
1576 let pred_row = decode_selective(
1577 &schema,
1578 &row_layout,
1579 data,
1580 &pred_cols,
1581 );
1582 if eval_predicate(predicate, &pred_row, &columns) {
1583 count += 1;
1584 }
1585 },
1586 )?;
1587
1588 return Ok(QueryResult::Scalar(Value::Int(count)));
1589 }
1590 }
1591 }
1592 }
1593
1594 if matches!(
1598 function,
1599 AggFunc::Sum
1600 | AggFunc::Avg
1601 | AggFunc::Min
1602 | AggFunc::Max
1603 | AggFunc::CountDistinct
1604 ) {
1605 if let Some(Expr::Field(col)) = argument.as_ref() {
1606 let (table_opt, pred_opt): (Option<&str>, Option<&Expr>) =
1608 match input.as_ref() {
1609 PlanNode::SeqScan { table } => (Some(table.as_str()), None),
1610 PlanNode::Filter {
1611 input: inner,
1612 predicate,
1613 } => {
1614 if let PlanNode::SeqScan { table } = inner.as_ref() {
1615 (Some(table.as_str()), Some(predicate))
1616 } else {
1617 (None, None)
1618 }
1619 }
1620 _ => (None, None),
1621 };
1622 if let Some(table) = table_opt {
1623 if let Some(result) =
1624 self.agg_single_col_fast(table, col, *function, pred_opt)?
1625 {
1626 return Ok(result);
1627 }
1628 }
1629 }
1630 }
1631
1632 let result = self.execute_plan(input)?;
1637 match result {
1638 QueryResult::Rows { columns, rows } => {
1639 aggregate_rows(*function, argument.as_ref(), &columns, &rows)
1640 }
1641 _ => Err("aggregate requires row input".into()),
1642 }
1643 }
1644
1645 PlanNode::Insert {
1646 table,
1647 rows,
1648 returning,
1649 } => {
1650 let mut returning_columns: Vec<String> = Vec::new();
1655 let all_values: Vec<Vec<Value>> = {
1656 let schema = self
1657 .catalog
1658 .schema(table)
1659 .ok_or_else(|| QueryError::TableNotFound(table.to_string()))?;
1660 if *returning {
1661 returning_columns = schema.columns.iter().map(|c| c.name.clone()).collect();
1662 }
1663 let defaults = self.catalog.column_defaults(table).unwrap_or(&[]);
1664 let auto = self.catalog.auto_columns(table).unwrap_or(&[]);
1665 let mut all = Vec::with_capacity(rows.len());
1666 for assignments in rows {
1667 let mut values = vec![Value::Empty; schema.columns.len()];
1668 for a in assignments {
1669 let idx = schema.column_index(&a.field).ok_or_else(|| {
1670 QueryError::ColumnNotFound {
1671 table: String::new(),
1672 column: a.field.clone(),
1673 }
1674 })?;
1675 let raw = literal_to_value(&a.value)?;
1676 values[idx] = coerce_value(raw, &schema.columns[idx])?;
1677 }
1678 for (i, slot) in values.iter_mut().enumerate() {
1682 if slot.is_empty() {
1683 if let Some(Some(d)) = defaults.get(i) {
1684 *slot = d.clone();
1685 }
1686 }
1687 }
1688 for col in &schema.columns {
1689 let pos = col.position as usize;
1690 let is_auto = auto.get(pos).copied().unwrap_or(false);
1693 if col.required && !is_auto && matches!(values[pos], Value::Empty) {
1694 return Err(QueryError::Execution(format!(
1695 "column '{}' is required but no value was provided",
1696 col.name
1697 )));
1698 }
1699 }
1700 all.push(values);
1701 }
1702 all
1703 };
1704 let mut all_values = all_values;
1709 for values in all_values.iter_mut() {
1710 self.catalog.assign_auto_columns(table, values);
1711 }
1712 self.charge_rows(&all_values)?;
1718 let n = all_values.len() as u64;
1719 for values in &all_values {
1720 self.catalog
1721 .insert(table, values)
1722 .map_err(|e| QueryError::StorageError(e.to_string()))?;
1723 }
1724 self.view_registry.mark_dependents_dirty(table);
1725 if *returning {
1726 Ok(QueryResult::Rows {
1727 columns: returning_columns,
1728 rows: all_values,
1729 })
1730 } else {
1731 Ok(QueryResult::Modified(n))
1732 }
1733 }
1734
1735 PlanNode::Upsert {
1736 table,
1737 key_column,
1738 assignments,
1739 on_conflict,
1740 } => {
1741 let (values, key_idx) = {
1742 let schema = self
1743 .catalog
1744 .schema(table)
1745 .ok_or_else(|| QueryError::TableNotFound(table.to_string()))?;
1746 let mut values = vec![Value::Empty; schema.columns.len()];
1747 for a in assignments {
1748 let idx = schema.column_index(&a.field).ok_or_else(|| {
1749 QueryError::ColumnNotFound {
1750 table: String::new(),
1751 column: a.field.clone(),
1752 }
1753 })?;
1754 let raw = literal_to_value(&a.value)?;
1755 values[idx] = coerce_value(raw, &schema.columns[idx])?;
1756 }
1757 let defaults = self.catalog.column_defaults(table).unwrap_or(&[]);
1760 for (i, slot) in values.iter_mut().enumerate() {
1761 if slot.is_empty() {
1762 if let Some(Some(d)) = defaults.get(i) {
1763 *slot = d.clone();
1764 }
1765 }
1766 }
1767 for col in &schema.columns {
1768 if col.required && matches!(values[col.position as usize], Value::Empty) {
1769 return Err(QueryError::Execution(format!(
1770 "column '{}' is required but no value was provided",
1771 col.name
1772 )));
1773 }
1774 }
1775 let key_idx = schema
1776 .column_index(key_column)
1777 .ok_or_else(|| format!("key column '{key_column}' not found"))?;
1778 (values, key_idx)
1779 };
1780
1781 if self.catalog.is_index_unique(table, key_column) != Some(true) {
1785 return Err(QueryError::Execution(format!(
1786 "upsert on .{key_column} requires a unique column (declare it with \
1787 `unique {key_column}: <type>` or `alter {table} add unique .{key_column}`)"
1788 )));
1789 }
1790
1791 let key_value = values[key_idx].clone();
1792
1793 let existing = {
1795 let tbl = self
1796 .catalog
1797 .get_table(table)
1798 .ok_or_else(|| QueryError::TableNotFound(table.to_string()))?;
1799 let rids = tbl.index_lookup_all(key_column, &key_value);
1802 rids.into_iter()
1805 .next()
1806 .and_then(|rid| tbl.get(rid).map(|row| (rid, row)))
1807 };
1808
1809 if let Some((rid, mut existing_row)) = existing {
1810 let update_assignments = if on_conflict.is_empty() {
1812 assignments
1813 } else {
1814 on_conflict
1815 };
1816 let changed_cols: Vec<usize> = {
1817 let schema = self
1818 .catalog
1819 .schema(table)
1820 .ok_or_else(|| QueryError::TableNotFound(table.to_string()))?;
1821 let mut indices = Vec::new();
1822 for a in update_assignments {
1823 let idx = schema.column_index(&a.field).ok_or_else(|| {
1824 QueryError::ColumnNotFound {
1825 table: String::new(),
1826 column: a.field.clone(),
1827 }
1828 })?;
1829 if idx != key_idx {
1830 existing_row[idx] =
1835 coerce_value(literal_to_value(&a.value)?, &schema.columns[idx])
1836 .map_err(QueryError::TypeError)?;
1837 indices.push(idx);
1838 }
1839 }
1840 indices
1841 };
1842 self.catalog
1843 .update_hinted(table, rid, &existing_row, Some(&changed_cols))
1844 .map_err(|e| QueryError::StorageError(e.to_string()))?;
1845 self.view_registry.mark_dependents_dirty(table);
1846 Ok(QueryResult::Modified(1))
1847 } else {
1848 self.catalog
1850 .insert(table, &values)
1851 .map_err(|e| QueryError::StorageError(e.to_string()))?;
1852 self.view_registry.mark_dependents_dirty(table);
1853 Ok(QueryResult::Modified(1))
1854 }
1855 }
1856
1857 PlanNode::Update {
1858 input,
1859 table,
1860 assignments,
1861 returning,
1862 } => {
1863 let (col_indices, literal_vals, target_cols): (
1869 Vec<usize>,
1870 Option<Vec<Value>>,
1871 Vec<ColumnDef>,
1872 ) = {
1873 let schema_ref = self
1874 .catalog
1875 .schema(table)
1876 .ok_or_else(|| QueryError::TableNotFound(table.to_string()))?;
1877 let indices: Vec<usize> = assignments
1878 .iter()
1879 .map(|a| {
1880 schema_ref.column_index(&a.field).ok_or_else(|| {
1881 QueryError::ColumnNotFound {
1882 table: String::new(),
1883 column: a.field.clone(),
1884 }
1885 })
1886 })
1887 .collect::<Result<_, _>>()?;
1888 let target_cols: Vec<ColumnDef> = indices
1892 .iter()
1893 .map(|&idx| schema_ref.columns[idx].clone())
1894 .collect();
1895 let raw_vals: Result<Vec<Value>, _> = assignments
1899 .iter()
1900 .map(|a| literal_to_value(&a.value))
1901 .collect();
1902 let coerced = match raw_vals {
1911 Ok(raws) => {
1912 let mut out = Vec::with_capacity(raws.len());
1913 for (raw, &idx) in raws.into_iter().zip(indices.iter()) {
1914 out.push(
1915 coerce_value(raw, &schema_ref.columns[idx])
1916 .map_err(QueryError::TypeError)?,
1917 );
1918 }
1919 Some(out)
1920 }
1921 Err(_) => None,
1922 };
1923 (indices, coerced, target_cols)
1924 };
1925 let resolved_assignments: Option<Vec<(usize, Value)>> =
1926 literal_vals.map(|vals| col_indices.iter().copied().zip(vals).collect());
1927
1928 let changed_cols: Vec<usize> = col_indices.clone();
1931
1932 if *returning {
1939 let columns: Vec<String> = {
1940 let schema_ref = self
1941 .catalog
1942 .schema(table)
1943 .ok_or_else(|| QueryError::TableNotFound(table.to_string()))?;
1944 schema_ref.columns.iter().map(|c| c.name.clone()).collect()
1945 };
1946 let matching_rids = self.collect_rids_for_mutation(input, table)?;
1947 let mut out_rows: Vec<Vec<Value>> = Vec::with_capacity(matching_rids.len());
1948 crate::cancel::check()?;
1955 for rid in matching_rids {
1956 let mut row = match self.catalog.get(table, rid) {
1957 Some(r) => r,
1958 None => continue,
1959 };
1960 match &resolved_assignments {
1961 Some(resolved) => {
1963 for (idx, val) in resolved.iter() {
1964 row[*idx] = val.clone();
1965 }
1966 }
1967 None => {
1973 for (i, asgn) in assignments.iter().enumerate() {
1974 let val = eval_expr(&asgn.value, &row, &columns);
1975 row[col_indices[i]] = coerce_value(val, &target_cols[i])
1976 .map_err(QueryError::TypeError)?;
1977 }
1978 }
1979 }
1980 self.catalog
1981 .update_hinted(table, rid, &row, Some(&changed_cols))
1982 .map_err(|e| QueryError::StorageError(e.to_string()))?;
1983 out_rows.push(row);
1984 }
1985 self.view_registry.mark_dependents_dirty(table);
1986 return Ok(QueryResult::Rows {
1987 columns,
1988 rows: out_rows,
1989 });
1990 }
1991
1992 if let Some(ref resolved_assignments) = resolved_assignments {
1999 if let PlanNode::Filter {
2000 input: inner,
2001 predicate,
2002 } = input.as_ref()
2003 {
2004 if let PlanNode::SeqScan { table: t } = inner.as_ref() {
2005 if t == table {
2006 crate::cancel::check()?;
2012 let fused_result = self.try_fused_scan_update(
2013 table,
2014 predicate,
2015 resolved_assignments,
2016 &changed_cols,
2017 );
2018 if let Some(result) = fused_result {
2019 return result;
2020 }
2021 }
2022 }
2023 }
2024 }
2025
2026 let matching_rids = self.collect_rids_for_mutation(input, table)?;
2028 crate::cancel::check()?;
2031
2032 if let Some(ref resolved_assignments) = resolved_assignments {
2034 let fast_patch: Option<Vec<FastPatch>> = {
2040 let tbl = self
2041 .catalog
2042 .get_table(table)
2043 .ok_or_else(|| QueryError::TableNotFound(table.to_string()))?;
2044 let schema = tbl.schema();
2045 let all_fixed_nonnull = !tbl.has_overflow_rows()
2049 && resolved_assignments.iter().all(|(idx, val)| {
2050 is_fixed_size(schema.columns[*idx].type_id) && !val.is_empty()
2051 });
2052 let no_indexed = !resolved_assignments
2053 .iter()
2054 .any(|(idx, _)| tbl.has_indexed_col(*idx));
2055
2056 if all_fixed_nonnull && no_indexed {
2057 let layout = RowLayout::new(schema);
2058 let bitmap_size = layout.bitmap_size();
2059 let patches: Vec<FastPatch> = resolved_assignments
2060 .iter()
2061 .map(|(idx, val)| {
2062 let fixed_off = layout
2063 .fixed_offset(*idx)
2064 .expect("is_fixed_size already checked");
2065 let field_off = 2 + bitmap_size + fixed_off;
2066 let bytes: FixedBytes = match val {
2067 Value::Int(v) => FixedBytes::I64(v.to_le_bytes()),
2068 Value::Float(v) => FixedBytes::F64(v.to_le_bytes()),
2069 Value::Bool(v) => FixedBytes::Bool(if *v { 1 } else { 0 }),
2070 Value::DateTime(v) => FixedBytes::I64(v.to_le_bytes()),
2071 Value::Uuid(v) => FixedBytes::Uuid(*v),
2072 _ => unreachable!("all_fixed_nonnull guard lied"),
2073 };
2074 FastPatch {
2075 field_off,
2076 bitmap_byte_off: 2 + idx / 8,
2077 bit_mask: 1u8 << (idx % 8),
2078 bytes,
2079 }
2080 })
2081 .collect();
2082 Some(patches)
2083 } else {
2084 None
2085 }
2086 };
2087
2088 if let Some(patches) = fast_patch {
2089 let mut count = 0u64;
2090 let mut fallback_rids: Vec<RowId> = Vec::new();
2091 for rid in &matching_rids {
2092 let ok = self
2107 .catalog
2108 .update_row_bytes_logged(table, *rid, |row| {
2109 let base = row_body_base(row);
2110 for p in &patches {
2111 row[base + p.bitmap_byte_off] &= !p.bit_mask;
2112 let field_bytes = p.bytes.as_slice();
2113 row[base + p.field_off
2114 ..base + p.field_off + field_bytes.len()]
2115 .copy_from_slice(field_bytes);
2116 }
2117 })
2118 .map_err(|e| QueryError::StorageError(e.to_string()))?;
2119 if ok {
2120 count += 1;
2121 } else {
2122 fallback_rids.push(*rid);
2123 }
2124 }
2125 for rid in fallback_rids {
2126 let mut row = match self.catalog.get(table, rid) {
2127 Some(r) => r,
2128 None => continue,
2129 };
2130 for (idx, val) in resolved_assignments.iter() {
2131 row[*idx] = val.clone();
2132 }
2133 self.catalog
2134 .update_hinted(table, rid, &row, Some(&changed_cols))
2135 .map_err(|e| QueryError::StorageError(e.to_string()))?;
2136 count += 1;
2137 }
2138 self.view_registry.mark_dependents_dirty(table);
2139 return Ok(QueryResult::Modified(count));
2140 }
2141
2142 let var_fast: Option<(usize, Option<Vec<u8>>)> = {
2144 let tbl = self
2145 .catalog
2146 .get_table(table)
2147 .ok_or_else(|| QueryError::TableNotFound(table.to_string()))?;
2148 let schema = tbl.schema();
2149 let is_single = resolved_assignments.len() == 1 && !tbl.has_overflow_rows();
2153 let is_var_col = is_single
2154 && !is_fixed_size(schema.columns[resolved_assignments[0].0].type_id);
2155 let no_indexed = !resolved_assignments
2156 .iter()
2157 .any(|(idx, _)| tbl.has_indexed_col(*idx));
2158
2159 if is_single && is_var_col && no_indexed {
2160 let (idx, val) = &resolved_assignments[0];
2161 let bytes_opt: Option<Vec<u8>> = match val {
2162 Value::Str(s) => Some(s.as_bytes().to_vec()),
2163 Value::Bytes(b) => Some(b.clone()),
2164 Value::Json(b) => Some(b.to_vec()),
2168 Value::Empty => None,
2169 _ => {
2170 return Err(QueryError::TypeError(format!(
2171 "cannot assign non-var value to var column '{}'",
2172 schema.columns[*idx].name
2173 )))
2174 }
2175 };
2176 Some((*idx, bytes_opt))
2177 } else {
2178 None
2179 }
2180 };
2181
2182 if let Some((col_idx, new_bytes_opt)) = var_fast {
2183 let new_bytes_ref: Option<&[u8]> = new_bytes_opt.as_deref();
2184 let mut count = 0u64;
2185 let mut fallback_rids: Vec<RowId> = Vec::new();
2186 for rid in &matching_rids {
2187 let ok = self
2193 .catalog
2194 .patch_var_col_logged(table, *rid, col_idx, new_bytes_ref)
2195 .map_err(|e| QueryError::StorageError(e.to_string()))?;
2196 if ok {
2197 count += 1;
2198 } else {
2199 fallback_rids.push(*rid);
2200 }
2201 }
2202 for rid in fallback_rids {
2203 let mut row = match self.catalog.get(table, rid) {
2204 Some(r) => r,
2205 None => continue,
2206 };
2207 for (idx, val) in resolved_assignments.iter() {
2208 row[*idx] = val.clone();
2209 }
2210 self.catalog
2211 .update_hinted(table, rid, &row, Some(&changed_cols))
2212 .map_err(|e| QueryError::StorageError(e.to_string()))?;
2213 count += 1;
2214 }
2215 self.view_registry.mark_dependents_dirty(table);
2216 return Ok(QueryResult::Modified(count));
2217 }
2218
2219 let mut count = 0u64;
2221 for rid in matching_rids {
2222 let mut row = match self.catalog.get(table, rid) {
2223 Some(r) => r,
2224 None => continue,
2225 };
2226 for (idx, val) in resolved_assignments.iter() {
2227 row[*idx] = val.clone();
2228 }
2229 self.catalog
2230 .update_hinted(table, rid, &row, Some(&changed_cols))
2231 .map_err(|e| QueryError::StorageError(e.to_string()))?;
2232 count += 1;
2233 }
2234 self.view_registry.mark_dependents_dirty(table);
2235 return Ok(QueryResult::Modified(count));
2236 } let col_names: Vec<String> = {
2242 let schema_ref = self
2243 .catalog
2244 .schema(table)
2245 .ok_or_else(|| QueryError::TableNotFound(table.to_string()))?;
2246 schema_ref.columns.iter().map(|c| c.name.clone()).collect()
2247 };
2248 let mut count = 0u64;
2249 for rid in matching_rids {
2250 let mut row = match self.catalog.get(table, rid) {
2251 Some(r) => r,
2252 None => continue,
2253 };
2254 for (i, asgn) in assignments.iter().enumerate() {
2255 let val = eval_expr(&asgn.value, &row, &col_names);
2256 row[col_indices[i]] =
2262 coerce_value(val, &target_cols[i]).map_err(QueryError::TypeError)?;
2263 }
2264 self.catalog
2265 .update_hinted(table, rid, &row, Some(&changed_cols))
2266 .map_err(|e| QueryError::StorageError(e.to_string()))?;
2267 count += 1;
2268 }
2269 self.view_registry.mark_dependents_dirty(table);
2270 Ok(QueryResult::Modified(count))
2271 }
2272
2273 PlanNode::Delete {
2274 input,
2275 table,
2276 returning,
2277 } => {
2278 if *returning {
2285 let columns: Vec<String> = {
2286 let schema_ref = self
2287 .catalog
2288 .schema(table)
2289 .ok_or_else(|| QueryError::TableNotFound(table.to_string()))?;
2290 schema_ref.columns.iter().map(|c| c.name.clone()).collect()
2291 };
2292 let matching_rids = self.collect_rids_for_mutation(input, table)?;
2293 let mut out_rows: Vec<Vec<Value>> = Vec::with_capacity(matching_rids.len());
2294 let mut cancel = CancelCheck::new();
2298 for rid in &matching_rids {
2299 cancel.tick()?;
2300 if let Some(row) = self.catalog.get(table, *rid) {
2301 out_rows.push(row);
2302 }
2303 }
2304 crate::cancel::check()?;
2305 self.catalog
2306 .delete_many(table, &matching_rids)
2307 .map_err(|e| QueryError::StorageError(e.to_string()))?;
2308 self.view_registry.mark_dependents_dirty(table);
2309 return Ok(QueryResult::Rows {
2310 columns,
2311 rows: out_rows,
2312 });
2313 }
2314
2315 let delete_overflow = self.catalog.table_has_overflow(table);
2339 if let PlanNode::Filter {
2340 input: inner,
2341 predicate,
2342 } = input.as_ref()
2343 {
2344 if let PlanNode::SeqScan { table: t } = inner.as_ref() {
2345 if t == table && !delete_overflow {
2346 let schema = self
2347 .catalog
2348 .schema(table)
2349 .ok_or_else(|| QueryError::TableNotFound(table.to_string()))?;
2350 let columns: Vec<String> =
2351 schema.columns.iter().map(|c| c.name.clone()).collect();
2352 let fast = FastLayout::new(schema);
2353 if let Some(compiled) =
2354 compile_predicate(predicate, &columns, &fast, schema)
2355 {
2356 crate::cancel::check()?;
2362 let count = self
2363 .catalog
2364 .scan_delete_matching_logged(table, |data| compiled(data))
2365 .map_err(|e| QueryError::StorageError(e.to_string()))?;
2366 self.view_registry.mark_dependents_dirty(table);
2367 return Ok(QueryResult::Modified(count));
2368 }
2369 }
2370 }
2371 } else if let PlanNode::SeqScan { table: t } = input.as_ref() {
2372 if t == table && !delete_overflow {
2373 crate::cancel::check()?;
2377 let count = self
2378 .catalog
2379 .scan_delete_matching_logged(table, |_| true)
2380 .map_err(|e| QueryError::StorageError(e.to_string()))?;
2381 self.view_registry.mark_dependents_dirty(table);
2382 return Ok(QueryResult::Modified(count));
2383 }
2384 }
2385
2386 let matching_rids = self.collect_rids_for_mutation(input, table)?;
2387 crate::cancel::check()?;
2388 let count = self
2389 .catalog
2390 .delete_many(table, &matching_rids)
2391 .map_err(|e| QueryError::StorageError(e.to_string()))?;
2392 self.view_registry.mark_dependents_dirty(table);
2393 Ok(QueryResult::Modified(count))
2394 }
2395
2396 PlanNode::AliasScan { table, alias } => {
2397 let schema = self
2407 .catalog
2408 .schema(table)
2409 .ok_or_else(|| QueryError::TableNotFound(table.to_string()))?
2410 .clone();
2411 let columns: Vec<String> = schema
2412 .columns
2413 .iter()
2414 .map(|c| format!("{alias}.{}", c.name))
2415 .collect();
2416 let mut cancel = CancelCheck::new();
2417 let mut rows: Vec<Vec<Value>> = Vec::new();
2418 for (_, row) in self
2419 .catalog
2420 .scan(table)
2421 .map_err(|e| QueryError::StorageError(e.to_string()))?
2422 {
2423 cancel.tick()?;
2424 rows.push(row);
2425 }
2426 Ok(QueryResult::Rows { columns, rows })
2427 }
2428
2429 PlanNode::NestedLoopJoin {
2430 left,
2431 right,
2432 on,
2433 kind,
2434 } => {
2435 let left_result = self.execute_plan(left)?;
2446 let right_result = self.execute_plan(right)?;
2447 let (left_columns, left_rows) = match left_result {
2448 QueryResult::Rows { columns, rows } => (columns, rows),
2449 _ => return Err("join left side must produce rows".into()),
2450 };
2451 let (right_columns, right_rows) = match right_result {
2452 QueryResult::Rows { columns, rows } => (columns, rows),
2453 _ => return Err("join right side must produce rows".into()),
2454 };
2455
2456 self.charge_rows(&left_rows)?;
2460 self.charge_rows(&right_rows)?;
2461
2462 execute_materialized_join(
2463 left_columns,
2464 left_rows,
2465 right_columns,
2466 right_rows,
2467 on.as_ref(),
2468 *kind,
2469 self.nested_loop_pair_limit,
2470 )
2471 }
2472
2473 PlanNode::Distinct { input } => {
2474 let result = self.execute_plan(input)?;
2475 match result {
2476 QueryResult::Rows { columns, rows } => {
2477 let mut seen = std::collections::HashSet::new();
2478 let mut unique_rows = Vec::new();
2479 let mut cancel = CancelCheck::new();
2480 for row in rows {
2481 cancel.tick()?;
2482 if seen.insert(row.clone()) {
2483 unique_rows.push(row);
2484 }
2485 }
2486 Ok(QueryResult::Rows {
2487 columns,
2488 rows: unique_rows,
2489 })
2490 }
2491 other => Ok(other),
2492 }
2493 }
2494
2495 PlanNode::GroupBy {
2496 input,
2497 keys,
2498 aggregates,
2499 having,
2500 } => {
2501 if aggregates
2502 .iter()
2503 .any(|aggregate| aggregate.provenance_alias.is_some())
2504 {
2505 let input = self.materialize_rows_with_provenance(input)?;
2506 self.charge_rows(&input.rows)?;
2507 return exec_group_by_with_provenance(
2508 input,
2509 keys,
2510 aggregates,
2511 having,
2512 self.query_memory_limit(),
2513 );
2514 }
2515 let result = self.execute_plan(input)?;
2516 match result {
2517 QueryResult::Rows { columns, rows } => {
2518 self.charge_rows(&rows)?;
2521 exec_group_by(columns, rows, keys, aggregates, having)
2522 }
2523 _ => Err("group by requires row input".into()),
2524 }
2525 }
2526
2527 PlanNode::CreateTable {
2528 name,
2529 fields,
2530 if_not_exists,
2531 } => {
2532 if self.catalog.schema(name).is_some() {
2536 if *if_not_exists {
2537 return Ok(QueryResult::Executed {
2538 message: format!("type '{name}' already exists (skipped)"),
2539 });
2540 }
2541 return Err(QueryError::Execution(format!(
2544 "cannot create type '{name}': it already exists"
2545 )));
2546 }
2547 let columns: Vec<ColumnDef> = fields
2548 .iter()
2549 .enumerate()
2550 .map(|(i, f)| -> Result<ColumnDef, QueryError> {
2551 Ok(ColumnDef {
2552 name: f.name.clone(),
2553 type_id: type_name_to_id(&f.type_name)
2554 .map_err(QueryError::TypeError)?,
2555 required: f.required,
2556 position: i as u16,
2557 })
2558 })
2559 .collect::<Result<Vec<_>, _>>()?;
2560 let mut defaults: Vec<Option<Value>> = vec![None; columns.len()];
2564 let mut auto_cols: Vec<bool> = vec![false; columns.len()];
2565 for (i, f) in fields.iter().enumerate() {
2566 if let Some(lit) = &f.default {
2567 let raw = literal_value_from(lit);
2568 defaults[i] = Some(coerce_value(raw, &columns[i])?);
2569 }
2570 if f.auto {
2571 if columns[i].type_id != TypeId::Int {
2575 return Err(QueryError::TypeError(format!(
2576 "auto column '{}' must be of type int",
2577 f.name
2578 )));
2579 }
2580 if f.default.is_some() {
2581 return Err(QueryError::TypeError(format!(
2582 "auto column '{}' cannot also declare a default",
2583 f.name
2584 )));
2585 }
2586 auto_cols[i] = true;
2587 }
2588 }
2589 let schema = Schema {
2590 table_name: name.clone(),
2591 columns,
2592 };
2593 self.catalog
2594 .create_table_full(schema, defaults, auto_cols)
2595 .map_err(|e| QueryError::StorageError(e.to_string()))?;
2596 for f in fields.iter().filter(|f| f.unique) {
2599 self.catalog
2600 .create_index_unique(name, &f.name, true)
2601 .map_err(|e| QueryError::StorageError(e.to_string()))?;
2602 }
2603 Ok(QueryResult::Created(name.clone()))
2604 }
2605
2606 PlanNode::AlterTable { table, action } => match action {
2607 AlterAction::AddColumn {
2608 name,
2609 type_name,
2610 required,
2611 } => {
2612 let position = self
2613 .catalog
2614 .schema(table)
2615 .ok_or_else(|| QueryError::TableNotFound(table.to_string()))?
2616 .columns
2617 .len() as u16;
2618 let col = ColumnDef {
2619 name: name.clone(),
2620 type_id: type_name_to_id(type_name).map_err(QueryError::TypeError)?,
2621 required: *required,
2622 position,
2623 };
2624 self.catalog
2625 .alter_table_add_column(table, col)
2626 .map_err(|e| QueryError::StorageError(e.to_string()))?;
2627 Ok(QueryResult::Executed {
2628 message: format!("column '{name}' added to '{table}'"),
2629 })
2630 }
2631 AlterAction::DropColumn { name, if_exists } => {
2632 if *if_exists {
2635 let present = self
2636 .catalog
2637 .schema(table)
2638 .map(|s| s.column_index(name).is_some())
2639 .unwrap_or(false);
2640 if !present {
2641 return Ok(QueryResult::Executed {
2642 message: format!(
2643 "column '{name}' does not exist on '{table}' (skipped)"
2644 ),
2645 });
2646 }
2647 }
2648 self.catalog
2649 .alter_table_drop_column(table, name)
2650 .map_err(|e| QueryError::StorageError(e.to_string()))?;
2651 Ok(QueryResult::Executed {
2652 message: format!("column '{name}' dropped from '{table}'"),
2653 })
2654 }
2655 AlterAction::AddIndex {
2656 target,
2657 if_not_exists: _,
2658 } => {
2659 let IndexTarget::Column(column) = target else {
2660 let IndexTarget::JsonPath(path) = target else {
2661 unreachable!("index target variants are exhaustive")
2662 };
2663 if let Some(existing) = resolve_expression_index(&self.catalog, table, path)
2664 {
2665 return Ok(QueryResult::Executed {
2666 message: format!(
2667 "expression index {} on '{}' already exists (skipped)",
2668 existing.index_id, table
2669 ),
2670 });
2671 }
2672 crate::cancel::check()?;
2673 let index_id = self
2674 .catalog
2675 .create_expression_index_metadata(
2676 table,
2677 1,
2678 path.canonical_text(),
2679 path.clone(),
2680 false,
2681 )
2682 .map_err(|error| QueryError::StorageError(error.to_string()))?;
2683 return Ok(QueryResult::Executed {
2684 message: format!("expression index {index_id} on '{}' created", table),
2685 });
2686 };
2687 crate::cancel::check()?;
2691 self.catalog
2692 .create_index(table, column)
2693 .map_err(|e| QueryError::StorageError(e.to_string()))?;
2694 Ok(QueryResult::Executed {
2695 message: format!("index on '{table}.{column}' created"),
2696 })
2697 }
2698 AlterAction::AddUnique {
2699 target,
2700 if_not_exists,
2701 } => {
2702 let IndexTarget::Column(column) = target else {
2703 let IndexTarget::JsonPath(path) = target else {
2704 unreachable!("index target variants are exhaustive")
2705 };
2706 if let Some(existing) = resolve_expression_index(&self.catalog, table, path)
2707 {
2708 if *if_not_exists {
2709 return Ok(QueryResult::Executed {
2710 message: format!(
2711 "expression index {} on '{}' already exists (skipped)",
2712 existing.index_id, table
2713 ),
2714 });
2715 }
2716 return Err(QueryError::Execution(format!(
2717 "cannot add unique expression index on {}: path already indexed",
2718 table
2719 )));
2720 }
2721 crate::cancel::check()?;
2722 let index_id = self
2723 .catalog
2724 .create_expression_index_metadata(
2725 table,
2726 1,
2727 path.canonical_text(),
2728 path.clone(),
2729 true,
2730 )
2731 .map_err(|error| QueryError::StorageError(error.to_string()))?;
2732 return Ok(QueryResult::Executed {
2733 message: format!(
2734 "unique expression index {index_id} on '{}' created",
2735 table
2736 ),
2737 });
2738 };
2739 if self.catalog.has_index(table, column) {
2742 if *if_not_exists {
2743 return Ok(QueryResult::Executed {
2744 message: format!(
2745 "index on '{table}.{column}' already exists (skipped)"
2746 ),
2747 });
2748 }
2749 return Err(QueryError::Execution(format!(
2752 "cannot add unique on {table}.{column}: column already indexed"
2753 )));
2754 }
2755 {
2758 let tbl = self
2759 .catalog
2760 .get_table(table)
2761 .ok_or_else(|| QueryError::TableNotFound(table.to_string()))?;
2762 let col_idx = tbl.schema().column_index(column).ok_or_else(|| {
2763 QueryError::ColumnNotFound {
2764 table: table.to_string(),
2765 column: column.clone(),
2766 }
2767 })?;
2768 let mut seen = std::collections::HashSet::new();
2769 let mut cancel = CancelCheck::new();
2770 for (_, row) in tbl.scan() {
2771 cancel.tick()?;
2772 let v = &row[col_idx];
2773 if v.is_empty() {
2774 continue;
2775 }
2776 if !seen.insert(v.clone()) {
2777 return Err(QueryError::Execution(format!(
2778 "cannot add unique on {table}.{column}: \
2779 duplicate value {v:?} exists"
2780 )));
2781 }
2782 }
2783 }
2784 crate::cancel::check()?;
2785 self.catalog
2786 .create_index_unique(table, column, true)
2787 .map_err(|e| QueryError::StorageError(e.to_string()))?;
2788 Ok(QueryResult::Executed {
2789 message: format!("unique index on '{table}.{column}' created"),
2790 })
2791 }
2792 AlterAction::DropIndex { target, if_exists } => {
2793 let IndexTarget::JsonPath(path) = target else {
2794 return Err(QueryError::Execution(
2795 "dropping stored-column indexes is not supported".to_string(),
2796 ));
2797 };
2798 let Some(existing) = resolve_expression_index(&self.catalog, table, path)
2799 else {
2800 if *if_exists {
2801 return Ok(QueryResult::Executed {
2802 message: format!(
2803 "expression index on '{}' does not exist (skipped)",
2804 table
2805 ),
2806 });
2807 }
2808 return Err(QueryError::Execution(format!(
2809 "expression index on '{}' does not exist",
2810 table
2811 )));
2812 };
2813 crate::cancel::check()?;
2814 self.catalog
2815 .drop_expression_index(table, existing.index_id)
2816 .map_err(|error| QueryError::StorageError(error.to_string()))?;
2817 Ok(QueryResult::Executed {
2818 message: format!(
2819 "expression index {} on '{}' dropped",
2820 existing.index_id, table
2821 ),
2822 })
2823 }
2824 },
2825
2826 PlanNode::DropTable { name, if_exists } => {
2827 if *if_exists && self.catalog.schema(name).is_none() {
2828 return Ok(QueryResult::Executed {
2829 message: format!("type '{name}' does not exist (skipped)"),
2830 });
2831 }
2832 self.catalog
2833 .drop_table(name)
2834 .map_err(|e| QueryError::StorageError(e.to_string()))?;
2835 Ok(QueryResult::Executed {
2836 message: format!("table '{name}' dropped"),
2837 })
2838 }
2839
2840 PlanNode::ListTypes => self.introspect_list_types(),
2841
2842 PlanNode::Describe { table } => self.introspect_describe(table),
2843
2844 PlanNode::CreateView { name, query_text } => {
2845 self.create_view(name, query_text)?;
2846 Ok(QueryResult::Executed {
2847 message: format!("materialized view '{name}' created"),
2848 })
2849 }
2850
2851 PlanNode::RefreshView { name } => {
2852 self.refresh_view(name)?;
2853 Ok(QueryResult::Executed {
2854 message: format!("materialized view '{name}' refreshed"),
2855 })
2856 }
2857
2858 PlanNode::DropView { name, if_exists } => {
2859 if *if_exists && !self.view_registry.is_view(name) {
2860 return Ok(QueryResult::Executed {
2861 message: format!("view '{name}' does not exist (skipped)"),
2862 });
2863 }
2864 self.drop_view(name)?;
2865 Ok(QueryResult::Executed {
2866 message: format!("materialized view '{name}' dropped"),
2867 })
2868 }
2869
2870 PlanNode::Window { input, windows } => {
2871 let result = self.execute_plan(input)?;
2872 execute_window(result, windows, self.query_memory_limit)
2873 }
2874
2875 PlanNode::Union { left, right, all } => {
2876 let left_result = self.execute_plan(left)?;
2877 let right_result = self.execute_plan(right)?;
2878 let (left_cols, left_rows) = match left_result {
2879 QueryResult::Rows { columns, rows } => (columns, rows),
2880 _ => return Err("UNION requires query results on left side".into()),
2881 };
2882 let (_, right_rows) = match right_result {
2883 QueryResult::Rows { columns, rows } => (columns, rows),
2884 _ => return Err("UNION requires query results on right side".into()),
2885 };
2886 let mut combined = left_rows;
2887 let mut cancel = CancelCheck::new();
2888 if *all {
2889 for row in right_rows {
2891 cancel.tick()?;
2892 combined.push(row);
2893 }
2894 } else {
2895 let mut seen = std::collections::HashSet::new();
2898 for row in &combined {
2899 cancel.tick()?;
2900 seen.insert(row.clone());
2901 }
2902 for row in right_rows {
2903 cancel.tick()?;
2904 if seen.insert(row.clone()) {
2905 combined.push(row);
2906 }
2907 }
2908 }
2909 Ok(QueryResult::Rows {
2910 columns: left_cols,
2911 rows: combined,
2912 })
2913 }
2914
2915 PlanNode::Explain { input } => {
2916 let text = format_plan_tree(&self.catalog, input, 0);
2920 Ok(QueryResult::Rows {
2921 columns: vec!["plan".to_string()],
2922 rows: text
2923 .lines()
2924 .map(|line| vec![Value::Str(line.to_string())])
2925 .collect(),
2926 })
2927 }
2928
2929 PlanNode::Begin => {
2930 if self.in_transaction {
2931 return Err(QueryError::Execution(
2932 "already in a transaction (nested transactions not supported)".into(),
2933 ));
2934 }
2935 self.catalog
2936 .begin_transaction()
2937 .map_err(|e| QueryError::StorageError(e.to_string()))?;
2938 self.in_transaction = true;
2939 Ok(QueryResult::Executed {
2940 message: "transaction started".to_string(),
2941 })
2942 }
2943
2944 PlanNode::Commit => {
2945 if !self.in_transaction {
2946 return Err(QueryError::Execution(
2947 "no active transaction to commit".into(),
2948 ));
2949 }
2950 self.catalog
2951 .commit_transaction()
2952 .map_err(|e| QueryError::StorageError(e.to_string()))?;
2953 self.in_transaction = false;
2954 Ok(QueryResult::Executed {
2955 message: "transaction committed".to_string(),
2956 })
2957 }
2958
2959 PlanNode::Rollback => {
2960 if !self.in_transaction {
2961 return Err(QueryError::Execution(
2962 "no active transaction to roll back".into(),
2963 ));
2964 }
2965 self.rollback_transaction_preserving_wal_archive()
2966 }
2967
2968 PlanNode::IndexScan { table, column, key } => {
2969 let key_value = literal_to_value(key)?;
2970 let tbl = self
2971 .catalog
2972 .get_table(table)
2973 .ok_or_else(|| QueryError::TableNotFound(table.to_string()))?;
2974 let columns: Vec<String> = tbl
2975 .schema()
2976 .columns
2977 .iter()
2978 .map(|c| c.name.clone())
2979 .collect();
2980
2981 if tbl.has_index(column) {
2985 let rids = tbl.index_lookup_all(column, &key_value);
2986 let mut rows: Vec<Vec<Value>> = Vec::with_capacity(rids.len());
2987 let mut cancel = CancelCheck::new();
2988 for rid in rids {
2989 cancel.tick()?;
2990 if let Some(row) = tbl.get(rid) {
2994 rows.push(row);
2995 }
2996 }
2997 return Ok(QueryResult::Rows { columns, rows });
2998 }
2999
3000 let schema = tbl.schema();
3008 let fast = FastLayout::new(schema);
3009 let synth_pred = Expr::BinaryOp(
3010 Box::new(Expr::Field(column.clone())),
3011 BinOp::Eq,
3012 Box::new(key.clone()),
3013 );
3014 if !tbl.has_overflow_rows() {
3017 if let Some(compiled) = compile_predicate(&synth_pred, &columns, &fast, schema)
3018 {
3019 let mut rows: Vec<Vec<Value>> = Vec::with_capacity(64);
3021 for_each_row_raw_cancellable(&self.catalog, table, |_rid, data| {
3022 if compiled(data) {
3023 rows.push(decode_row(schema, data));
3024 }
3025 })?;
3026 return Ok(QueryResult::Rows { columns, rows });
3027 }
3028 }
3029
3030 let col_idx =
3032 schema
3033 .column_index(column)
3034 .ok_or_else(|| QueryError::ColumnNotFound {
3035 table: String::new(),
3036 column: column.clone(),
3037 })?;
3038 let mut cancel = CancelCheck::new();
3039 let mut rows: Vec<Vec<Value>> = Vec::new();
3040 for (_, row) in tbl.scan() {
3041 cancel.tick()?;
3042 if row[col_idx] == key_value {
3043 rows.push(row);
3044 }
3045 }
3046 Ok(QueryResult::Rows { columns, rows })
3047 }
3048
3049 PlanNode::RangeScan {
3050 table,
3051 column,
3052 start,
3053 end,
3054 } => {
3055 let tbl = self
3056 .catalog
3057 .get_table(table)
3058 .ok_or_else(|| QueryError::TableNotFound(table.to_string()))?;
3059 let columns: Vec<String> = tbl
3060 .schema()
3061 .columns
3062 .iter()
3063 .map(|c| c.name.clone())
3064 .collect();
3065 let schema = tbl.schema();
3066
3067 let start_val = match start {
3068 Some((expr, _)) => Some(literal_to_value(expr)?),
3069 None => None,
3070 };
3071 let end_val = match end {
3072 Some((expr, _)) => Some(literal_to_value(expr)?),
3073 None => None,
3074 };
3075 let start_inclusive = start.as_ref().map(|(_, inc)| *inc).unwrap_or(true);
3076 let end_inclusive = end.as_ref().map(|(_, inc)| *inc).unwrap_or(true);
3077
3078 if tbl.is_index_unique(column) == Some(false) {
3084 if let Some(btree) = tbl.index(column) {
3085 if start_val.is_some() || end_val.is_some() {
3086 let col_idx = schema.column_index(column).ok_or_else(|| {
3087 QueryError::ColumnNotFound {
3088 table: String::new(),
3089 column: column.clone(),
3090 }
3091 })?;
3092 let rids = btree.range_rids(start_val.as_ref(), end_val.as_ref());
3093 let mut rows: Vec<Vec<Value>> = Vec::with_capacity(rids.len());
3094 let mut cancel = CancelCheck::new();
3095 for rid in rids {
3096 cancel.tick()?;
3097 if let Some(row) = tbl.get(rid) {
3099 if !row[col_idx].is_empty()
3100 && range_matches(
3101 &row[col_idx],
3102 &start_val,
3103 start_inclusive,
3104 &end_val,
3105 end_inclusive,
3106 )
3107 {
3108 rows.push(row);
3109 }
3110 }
3111 }
3112 return Ok(QueryResult::Rows { columns, rows });
3113 }
3114 }
3115 }
3116
3117 if tbl.is_index_unique(column) == Some(true) {
3120 if let Some(btree) = tbl.index(column) {
3121 let hits: Vec<(Value, RowId)> = match (&start_val, &end_val) {
3122 (Some(s), Some(e)) => btree.range(s, e).collect(),
3123 (Some(s), None) => btree.range_from(s),
3124 (None, Some(e)) => btree.range_to(e),
3125 (None, None) => {
3126 let mut cancel = CancelCheck::new();
3127 let mut rows: Vec<Vec<Value>> = Vec::new();
3128 for (_, row) in tbl.scan() {
3129 cancel.tick()?;
3130 rows.push(row);
3131 }
3132 return Ok(QueryResult::Rows { columns, rows });
3133 }
3134 };
3135 let mut rows: Vec<Vec<Value>> = Vec::with_capacity(hits.len());
3136 let mut cancel = CancelCheck::new();
3137 for (key, rid) in hits {
3138 cancel.tick()?;
3139 if !start_inclusive {
3140 if let Some(ref s) = start_val {
3141 if &key == s {
3142 continue;
3143 }
3144 }
3145 }
3146 if !end_inclusive {
3147 if let Some(ref e) = end_val {
3148 if &key == e {
3149 continue;
3150 }
3151 }
3152 }
3153 if let Some(row) = tbl.get(rid) {
3155 rows.push(row);
3156 }
3157 }
3158 return Ok(QueryResult::Rows { columns, rows });
3159 }
3160 }
3161
3162 let fast = FastLayout::new(schema);
3166 let synth = synthesize_range_predicate(column, start, end);
3167 if !tbl.has_overflow_rows() {
3168 if let Some(compiled) = compile_predicate(&synth, &columns, &fast, schema) {
3169 let mut rows: Vec<Vec<Value>> = Vec::with_capacity(64);
3170 for_each_row_raw_cancellable(&self.catalog, table, |_rid, data| {
3171 if compiled(data) {
3172 rows.push(decode_row(schema, data));
3173 }
3174 })?;
3175 return Ok(QueryResult::Rows { columns, rows });
3176 }
3177 }
3178
3179 let col_idx =
3180 schema
3181 .column_index(column)
3182 .ok_or_else(|| QueryError::ColumnNotFound {
3183 table: String::new(),
3184 column: column.clone(),
3185 })?;
3186 let mut cancel = CancelCheck::new();
3187 let mut rows: Vec<Vec<Value>> = Vec::new();
3188 for (_, row) in tbl.scan() {
3189 cancel.tick()?;
3190 if range_matches(
3191 &row[col_idx],
3192 &start_val,
3193 start_inclusive,
3194 &end_val,
3195 end_inclusive,
3196 ) {
3197 rows.push(row);
3198 }
3199 }
3200 Ok(QueryResult::Rows { columns, rows })
3201 }
3202 }
3203 }
3204
3205 fn create_view(&mut self, name: &str, query_text: &str) -> Result<(), QueryError> {
3210 if self.view_registry.is_view(name) {
3211 return Err(QueryError::ViewError(format!(
3212 "materialized view '{name}' already exists"
3213 )));
3214 }
3215 let result = self.execute_powql(query_text)?;
3217 let (columns, rows) = match result {
3218 QueryResult::Rows { columns, rows } => (columns, rows),
3219 _ => return Err("view source query must be a SELECT".into()),
3220 };
3221 let schema = self.derive_view_schema(name, &columns, &rows);
3223 crate::cancel::check()?;
3225 self.catalog
3226 .create_table(schema)
3227 .map_err(|e| QueryError::StorageError(e.to_string()))?;
3228 for row in &rows {
3229 self.catalog
3230 .insert(name, row)
3231 .map_err(|e| QueryError::StorageError(e.to_string()))?;
3232 }
3233 let depends_on = self.extract_view_deps(query_text);
3235 self.view_registry
3236 .register(ViewDef {
3237 name: name.to_string(),
3238 query: query_text.to_string(),
3239 depends_on,
3240 dirty: false,
3241 })
3242 .map_err(|e| QueryError::StorageError(e.to_string()))?;
3243 Ok(())
3244 }
3245
3246 fn refresh_view(&mut self, name: &str) -> Result<(), QueryError> {
3249 let def = self
3250 .view_registry
3251 .get(name)
3252 .ok_or_else(|| format!("materialized view '{name}' not found"))?;
3253 let query_text = def.query.clone();
3254 let result = self.execute_powql(&query_text)?;
3256 let (_columns, rows) = match result {
3257 QueryResult::Rows { columns, rows } => (columns, rows),
3258 _ => return Err("view source query must be a SELECT".into()),
3259 };
3260 crate::cancel::check()?;
3264 self.catalog
3265 .scan_delete_matching_logged(name, |_| true)
3266 .map_err(|e| QueryError::StorageError(e.to_string()))?;
3267 for row in &rows {
3268 self.catalog
3269 .insert(name, row)
3270 .map_err(|e| QueryError::StorageError(e.to_string()))?;
3271 }
3272 self.view_registry.mark_clean(name);
3273 Ok(())
3274 }
3275
3276 fn drop_view(&mut self, name: &str) -> Result<(), QueryError> {
3278 if !self.view_registry.is_view(name) {
3279 return Err(QueryError::ViewError(format!(
3280 "materialized view '{name}' not found"
3281 )));
3282 }
3283 self.view_registry
3284 .unregister(name)
3285 .map_err(|e| QueryError::StorageError(e.to_string()))?;
3286 self.catalog
3287 .drop_table(name)
3288 .map_err(|e| QueryError::StorageError(e.to_string()))?;
3289 Ok(())
3290 }
3291
3292 fn derive_view_schema(&self, name: &str, columns: &[String], rows: &[Vec<Value>]) -> Schema {
3295 use powdb_storage::types::{ColumnDef, TypeId};
3296 let cols: Vec<ColumnDef> = columns
3297 .iter()
3298 .enumerate()
3299 .map(|(i, col_name)| {
3300 let type_id = rows
3301 .first()
3302 .and_then(|row| row.get(i))
3303 .map(|v| v.type_id())
3304 .unwrap_or(TypeId::Str);
3305 ColumnDef {
3306 name: col_name.clone(),
3307 type_id,
3308 required: false,
3309 position: i as u16,
3310 }
3311 })
3312 .collect();
3313 Schema {
3314 table_name: name.to_string(),
3315 columns: cols,
3316 }
3317 }
3318
3319 fn extract_view_deps(&self, query_text: &str) -> Vec<String> {
3322 use crate::parser::parse;
3323 match parse(query_text) {
3324 Ok(Statement::Query(q)) => {
3325 let mut deps = vec![q.source.clone()];
3326 for j in &q.joins {
3327 deps.push(j.source.clone());
3328 }
3329 deps
3330 }
3331 _ => Vec::new(),
3332 }
3333 }
3334
3335 pub(super) fn agg_single_col_fast(
3345 &self,
3346 table: &str,
3347 col: &str,
3348 function: AggFunc,
3349 predicate: Option<&Expr>,
3350 ) -> Result<Option<QueryResult>, QueryError> {
3351 if self.catalog.table_has_overflow(table) {
3355 return Ok(None);
3356 }
3357 let schema = self
3358 .catalog
3359 .schema(table)
3360 .ok_or_else(|| QueryError::TableNotFound(table.to_string()))?
3361 .clone();
3362 let columns: Vec<String> = schema.columns.iter().map(|c| c.name.clone()).collect();
3363 let col_idx = match schema.column_index(col) {
3364 Some(i) => i,
3365 None => return Ok(None),
3366 };
3367 let col_type = schema.columns[col_idx].type_id;
3374 if col_type != TypeId::Int && col_type != TypeId::Float {
3375 return Ok(None);
3376 }
3377
3378 let fast = FastLayout::new(&schema);
3379 let byte_offset = match fast.fixed_offsets[col_idx] {
3384 Some(o) => o,
3385 None => return Ok(None),
3386 };
3387 let bitmap_byte = col_idx / 8;
3388 let bitmap_bit = (col_idx % 8) as u32;
3389 let body_data_offset = 2 + fast.bitmap_size + byte_offset;
3390
3391 let compiled_pred: Option<CompiledPredicate> = match predicate {
3393 Some(pred) => match compile_predicate(pred, &columns, &fast, &schema) {
3394 Some(c) => Some(c),
3395 None => return Ok(None), },
3397 None => None,
3398 };
3399
3400 let result = match col_type {
3427 TypeId::Int => match function {
3428 AggFunc::Sum | AggFunc::Avg => {
3429 let mut sum_i128: i128 = 0;
3430 let mut count: i64 = 0;
3431 agg_int_loop!(
3432 self,
3433 table,
3434 compiled_pred,
3435 bitmap_byte,
3436 bitmap_bit,
3437 body_data_offset,
3438 |v: i64| {
3439 count += 1;
3440 sum_i128 += v as i128;
3441 }
3442 );
3443 if matches!(function, AggFunc::Sum) {
3444 let clamped = sum_i128.clamp(i64::MIN as i128, i64::MAX as i128) as i64;
3445 QueryResult::Scalar(Value::Int(clamped))
3446 } else if count == 0 {
3447 QueryResult::Scalar(Value::Empty)
3448 } else {
3449 let avg = (sum_i128 as f64) / (count as f64);
3450 QueryResult::Scalar(Value::Float(avg))
3451 }
3452 }
3453 AggFunc::Min => {
3454 let mut min_v: Option<i64> = None;
3455 agg_int_loop!(
3456 self,
3457 table,
3458 compiled_pred,
3459 bitmap_byte,
3460 bitmap_bit,
3461 body_data_offset,
3462 |v: i64| {
3463 min_v = Some(match min_v {
3464 Some(m) => m.min(v),
3465 None => v,
3466 });
3467 }
3468 );
3469 QueryResult::Scalar(min_v.map(Value::Int).unwrap_or(Value::Empty))
3470 }
3471 AggFunc::Max => {
3472 let mut max_v: Option<i64> = None;
3473 agg_int_loop!(
3474 self,
3475 table,
3476 compiled_pred,
3477 bitmap_byte,
3478 bitmap_bit,
3479 body_data_offset,
3480 |v: i64| {
3481 max_v = Some(match max_v {
3482 Some(m) => m.max(v),
3483 None => v,
3484 });
3485 }
3486 );
3487 QueryResult::Scalar(max_v.map(Value::Int).unwrap_or(Value::Empty))
3488 }
3489 AggFunc::Count => {
3490 let mut count: i64 = 0;
3491 agg_int_loop!(
3492 self,
3493 table,
3494 compiled_pred,
3495 bitmap_byte,
3496 bitmap_bit,
3497 body_data_offset,
3498 |_v: i64| {
3499 count += 1;
3500 }
3501 );
3502 QueryResult::Scalar(Value::Int(count))
3503 }
3504 AggFunc::CountDistinct => {
3505 let mut seen = rustc_hash::FxHashSet::default();
3506 agg_int_loop!(
3507 self,
3508 table,
3509 compiled_pred,
3510 bitmap_byte,
3511 bitmap_bit,
3512 body_data_offset,
3513 |v: i64| {
3514 seen.insert(v);
3515 }
3516 );
3517 QueryResult::Scalar(Value::Int(seen.len() as i64))
3518 }
3519 },
3520 TypeId::Float => match function {
3521 AggFunc::Sum => {
3522 let mut sum: f64 = 0.0;
3527 agg_float_loop!(
3528 self,
3529 table,
3530 compiled_pred,
3531 bitmap_byte,
3532 bitmap_bit,
3533 body_data_offset,
3534 |v: f64| {
3535 sum += v;
3536 }
3537 );
3538 QueryResult::Scalar(Value::Float(sum))
3539 }
3540 AggFunc::Avg => {
3541 let mut sum: f64 = 0.0;
3542 let mut count: i64 = 0;
3543 agg_float_loop!(
3544 self,
3545 table,
3546 compiled_pred,
3547 bitmap_byte,
3548 bitmap_bit,
3549 body_data_offset,
3550 |v: f64| {
3551 sum += v;
3552 count += 1;
3553 }
3554 );
3555 if count == 0 {
3556 QueryResult::Scalar(Value::Empty)
3557 } else {
3558 QueryResult::Scalar(Value::Float(sum / count as f64))
3559 }
3560 }
3561 AggFunc::Min => {
3562 let mut min_v: Option<f64> = None;
3566 agg_float_loop!(
3567 self,
3568 table,
3569 compiled_pred,
3570 bitmap_byte,
3571 bitmap_bit,
3572 body_data_offset,
3573 |v: f64| {
3574 min_v = Some(match min_v {
3575 Some(m) => {
3576 if v.total_cmp(&m).is_lt() {
3577 v
3578 } else {
3579 m
3580 }
3581 }
3582 None => v,
3583 });
3584 }
3585 );
3586 QueryResult::Scalar(min_v.map(Value::Float).unwrap_or(Value::Empty))
3587 }
3588 AggFunc::Max => {
3589 let mut max_v: Option<f64> = None;
3590 agg_float_loop!(
3591 self,
3592 table,
3593 compiled_pred,
3594 bitmap_byte,
3595 bitmap_bit,
3596 body_data_offset,
3597 |v: f64| {
3598 max_v = Some(match max_v {
3599 Some(m) => {
3600 if v.total_cmp(&m).is_gt() {
3601 v
3602 } else {
3603 m
3604 }
3605 }
3606 None => v,
3607 });
3608 }
3609 );
3610 QueryResult::Scalar(max_v.map(Value::Float).unwrap_or(Value::Empty))
3611 }
3612 AggFunc::Count => {
3613 let mut count: i64 = 0;
3614 agg_float_loop!(
3615 self,
3616 table,
3617 compiled_pred,
3618 bitmap_byte,
3619 bitmap_bit,
3620 body_data_offset,
3621 |_v: f64| {
3622 count += 1;
3623 }
3624 );
3625 QueryResult::Scalar(Value::Int(count))
3626 }
3627 AggFunc::CountDistinct => {
3628 let mut seen = rustc_hash::FxHashSet::default();
3634 agg_float_loop!(
3635 self,
3636 table,
3637 compiled_pred,
3638 bitmap_byte,
3639 bitmap_bit,
3640 body_data_offset,
3641 |v: f64| {
3642 seen.insert(v.to_bits());
3643 }
3644 );
3645 QueryResult::Scalar(Value::Int(seen.len() as i64))
3646 }
3647 },
3648 _ => unreachable!("type guard above restricts to Int/Float"),
3649 };
3650 Ok(Some(result))
3651 }
3652
3653 pub(super) fn project_filter_limit_fast(
3656 &self,
3657 table: &str,
3658 fields: &[ProjectField],
3659 limit: usize,
3660 predicate: Option<&Expr>,
3661 ) -> Result<Option<QueryResult>, QueryError> {
3662 if self.catalog.table_has_overflow(table) {
3666 return Ok(None);
3667 }
3668 let schema = self
3669 .catalog
3670 .schema(table)
3671 .ok_or_else(|| QueryError::TableNotFound(table.to_string()))?
3672 .clone();
3673 let all_columns: Vec<String> = schema.columns.iter().map(|c| c.name.clone()).collect();
3674
3675 let mut proj_indices: Vec<usize> = Vec::with_capacity(fields.len());
3678 let mut proj_columns: Vec<String> = Vec::with_capacity(fields.len());
3679 for f in fields {
3680 let name = match &f.expr {
3681 Expr::Field(n) => n.clone(),
3682 _ => return Ok(None),
3683 };
3684 let idx = match all_columns.iter().position(|c| c == &name) {
3685 Some(i) => i,
3686 None => return Ok(None),
3687 };
3688 proj_indices.push(idx);
3689 proj_columns.push(f.alias.clone().unwrap_or(name));
3690 }
3691
3692 let fast = FastLayout::new(&schema);
3693 let row_layout = RowLayout::new(&schema);
3694
3695 let compiled_pred: Option<CompiledPredicate> = match predicate {
3696 Some(pred) => match compile_predicate(pred, &all_columns, &fast, &schema) {
3697 Some(c) => Some(c),
3698 None => return Ok(None),
3699 },
3700 None => None,
3701 };
3702
3703 let mut out: Vec<Vec<Value>> = Vec::with_capacity(limit.min(1024));
3704 let mut cancel = CancelCheck::new();
3711 let mut cancel_err: Option<QueryError> = None;
3712 self.catalog
3713 .try_for_each_row_raw(table, |_rid, data| {
3714 if let Err(e) = cancel.tick() {
3715 cancel_err = Some(e);
3716 return ControlFlow::Break(());
3717 }
3718 if let Some(ref pred) = compiled_pred {
3719 if !pred(data) {
3720 return ControlFlow::Continue(());
3721 }
3722 }
3723 let row: Vec<Value> = proj_indices
3724 .iter()
3725 .map(|&ci| decode_column(&schema, &row_layout, data, ci))
3726 .collect();
3727 out.push(row);
3728 if out.len() >= limit {
3729 ControlFlow::Break(())
3730 } else {
3731 ControlFlow::Continue(())
3732 }
3733 })
3734 .map_err(|e| QueryError::StorageError(e.to_string()))?;
3735 if let Some(e) = cancel_err {
3736 return Err(e);
3737 }
3738
3739 Ok(Some(QueryResult::Rows {
3740 columns: proj_columns,
3741 rows: out,
3742 }))
3743 }
3744
3745 pub(super) fn project_filter_sort_limit_fast(
3750 &self,
3751 table: &str,
3752 fields: &[ProjectField],
3753 sort_field: &str,
3754 descending: bool,
3755 limit: usize,
3756 predicate: Option<&Expr>,
3757 ) -> Result<Option<QueryResult>, QueryError> {
3758 if self.catalog.table_has_overflow(table) {
3761 return Ok(None);
3762 }
3763 if limit == 0 {
3764 return Ok(None);
3767 }
3768 const TOPN_PREALLOC_CAP: usize = 4096;
3774 let prealloc = limit.min(TOPN_PREALLOC_CAP);
3775 let schema = self
3776 .catalog
3777 .schema(table)
3778 .ok_or_else(|| QueryError::TableNotFound(table.to_string()))?
3779 .clone();
3780 let all_columns: Vec<String> = schema.columns.iter().map(|c| c.name.clone()).collect();
3781
3782 let sort_idx = match schema.column_index(sort_field) {
3789 Some(i) => i,
3790 None => return Ok(None),
3791 };
3792 let sort_col_type = schema.columns[sort_idx].type_id;
3793 if sort_col_type != TypeId::Int && sort_col_type != TypeId::Float {
3794 return Ok(None);
3795 }
3796
3797 let mut proj_indices: Vec<usize> = Vec::with_capacity(fields.len());
3799 let mut proj_columns: Vec<String> = Vec::with_capacity(fields.len());
3800 for f in fields {
3801 let name = match &f.expr {
3802 Expr::Field(n) => n.clone(),
3803 _ => return Ok(None),
3804 };
3805 let idx = match all_columns.iter().position(|c| c == &name) {
3806 Some(i) => i,
3807 None => return Ok(None),
3808 };
3809 proj_indices.push(idx);
3810 proj_columns.push(f.alias.clone().unwrap_or(name));
3811 }
3812
3813 let fast = FastLayout::new(&schema);
3814 let row_layout = RowLayout::new(&schema);
3815 let sort_byte_offset = match fast.fixed_offsets[sort_idx] {
3817 Some(o) => o,
3818 None => return Ok(None),
3819 };
3820 let sort_bitmap_byte = sort_idx / 8;
3821 let sort_bitmap_bit = (sort_idx % 8) as u32;
3822 let sort_body_data_offset = 2 + fast.bitmap_size + sort_byte_offset;
3823
3824 let compiled_pred: Option<CompiledPredicate> = match predicate {
3825 Some(pred) => match compile_predicate(pred, &all_columns, &fast, &schema) {
3826 Some(c) => Some(c),
3827 None => return Ok(None),
3828 },
3829 None => None,
3830 };
3831
3832 let drained: Vec<Vec<u8>> = match sort_col_type {
3841 TypeId::Int => {
3842 let mut seq: u64 = 0;
3843 let mut heap_desc: BinaryHeap<Reverse<(i64, u64, Vec<u8>)>> =
3844 BinaryHeap::with_capacity(prealloc);
3845 let mut heap_asc: BinaryHeap<(i64, u64, Vec<u8>)> =
3846 BinaryHeap::with_capacity(prealloc);
3847 let mut null_rows: Vec<Vec<u8>> = Vec::with_capacity(prealloc);
3848
3849 for_each_row_raw_cancellable(&self.catalog, table, |_rid, data| {
3850 if let Some(ref pred) = compiled_pred {
3851 if !pred(data) {
3852 return;
3853 }
3854 }
3855 let base = row_body_base(data);
3857 let sort_data_offset = base + sort_body_data_offset;
3858 if data.len() < sort_data_offset + 8
3859 || data.len() <= base + 2 + sort_bitmap_byte
3860 {
3861 return;
3862 }
3863 let is_null = (data[base + 2 + sort_bitmap_byte] >> sort_bitmap_bit) & 1 == 1;
3864 let id = seq;
3865 seq += 1;
3866 if is_null {
3867 if null_rows.len() < limit {
3868 null_rows.push(data.to_vec());
3869 }
3870 return;
3871 }
3872 let key = i64::from_le_bytes(
3873 data[sort_data_offset..sort_data_offset + 8]
3874 .try_into()
3875 .unwrap_or_else(|_| unreachable!()),
3876 );
3877 if descending {
3878 if heap_desc.len() < limit {
3879 heap_desc.push(Reverse((key, id, data.to_vec())));
3880 } else if let Some(Reverse((top_key, _, _))) = heap_desc.peek() {
3881 if key > *top_key {
3882 heap_desc.pop();
3883 heap_desc.push(Reverse((key, id, data.to_vec())));
3884 }
3885 }
3886 } else if heap_asc.len() < limit {
3887 heap_asc.push((key, id, data.to_vec()));
3888 } else if let Some((top_key, _, _)) = heap_asc.peek() {
3889 if key < *top_key {
3890 heap_asc.pop();
3891 heap_asc.push((key, id, data.to_vec()));
3892 }
3893 }
3894 })?;
3895
3896 let mut drained: Vec<(i64, u64, Vec<u8>)> = if descending {
3897 heap_desc.into_iter().map(|Reverse(t)| t).collect()
3898 } else {
3899 heap_asc.into_iter().collect()
3900 };
3901 if descending {
3902 cooperative_stable_sort_by(&mut drained, self.query_memory_limit, |a, b| {
3903 b.0.cmp(&a.0).then(a.1.cmp(&b.1))
3904 })?;
3905 } else {
3906 cooperative_stable_sort_by(&mut drained, self.query_memory_limit, |a, b| {
3907 a.0.cmp(&b.0).then(a.1.cmp(&b.1))
3908 })?;
3909 }
3910 let mut rows: Vec<Vec<u8>> = drained.into_iter().map(|(_, _, d)| d).collect();
3911 rows.extend(null_rows.into_iter().take(limit.saturating_sub(rows.len())));
3912 rows
3913 }
3914 TypeId::Float => {
3915 let mut seq: u64 = 0;
3924 let mut heap_desc: BinaryHeap<Reverse<(u64, u64, Vec<u8>)>> =
3925 BinaryHeap::with_capacity(prealloc);
3926 let mut heap_asc: BinaryHeap<(u64, u64, Vec<u8>)> =
3927 BinaryHeap::with_capacity(prealloc);
3928 let mut null_rows: Vec<Vec<u8>> = Vec::with_capacity(prealloc);
3929
3930 for_each_row_raw_cancellable(&self.catalog, table, |_rid, data| {
3931 if let Some(ref pred) = compiled_pred {
3932 if !pred(data) {
3933 return;
3934 }
3935 }
3936 let base = row_body_base(data);
3937 let sort_data_offset = base + sort_body_data_offset;
3938 if data.len() < sort_data_offset + 8
3939 || data.len() <= base + 2 + sort_bitmap_byte
3940 {
3941 return;
3942 }
3943 let is_null = (data[base + 2 + sort_bitmap_byte] >> sort_bitmap_bit) & 1 == 1;
3944 let id = seq;
3945 seq += 1;
3946 if is_null {
3947 if null_rows.len() < limit {
3948 null_rows.push(data.to_vec());
3949 }
3950 return;
3951 }
3952 let bits = u64::from_le_bytes(
3953 data[sort_data_offset..sort_data_offset + 8]
3954 .try_into()
3955 .unwrap_or_else(|_| unreachable!()),
3956 );
3957 let key = f64_bits_to_sortable_u64(bits);
3958 if descending {
3959 if heap_desc.len() < limit {
3960 heap_desc.push(Reverse((key, id, data.to_vec())));
3961 } else if let Some(Reverse((top_key, _, _))) = heap_desc.peek() {
3962 if key > *top_key {
3963 heap_desc.pop();
3964 heap_desc.push(Reverse((key, id, data.to_vec())));
3965 }
3966 }
3967 } else if heap_asc.len() < limit {
3968 heap_asc.push((key, id, data.to_vec()));
3969 } else if let Some((top_key, _, _)) = heap_asc.peek() {
3970 if key < *top_key {
3971 heap_asc.pop();
3972 heap_asc.push((key, id, data.to_vec()));
3973 }
3974 }
3975 })?;
3976
3977 let mut drained: Vec<(u64, u64, Vec<u8>)> = if descending {
3978 heap_desc.into_iter().map(|Reverse(t)| t).collect()
3979 } else {
3980 heap_asc.into_iter().collect()
3981 };
3982 if descending {
3983 cooperative_stable_sort_by(&mut drained, self.query_memory_limit, |a, b| {
3984 b.0.cmp(&a.0).then(a.1.cmp(&b.1))
3985 })?;
3986 } else {
3987 cooperative_stable_sort_by(&mut drained, self.query_memory_limit, |a, b| {
3988 a.0.cmp(&b.0).then(a.1.cmp(&b.1))
3989 })?;
3990 }
3991 let mut rows: Vec<Vec<u8>> = drained.into_iter().map(|(_, _, d)| d).collect();
3992 rows.extend(null_rows.into_iter().take(limit.saturating_sub(rows.len())));
3993 rows
3994 }
3995 _ => unreachable!("type guard above restricts to Int/Float"),
3996 };
3997
3998 let mut cancel = CancelCheck::new();
3999 let mut rows: Vec<Vec<Value>> = Vec::with_capacity(drained.len());
4000 for data in drained {
4001 cancel.tick()?;
4002 rows.push(
4003 proj_indices
4004 .iter()
4005 .map(|&ci| decode_column(&schema, &row_layout, &data, ci))
4006 .collect(),
4007 );
4008 }
4009
4010 Ok(Some(QueryResult::Rows {
4011 columns: proj_columns,
4012 rows,
4013 }))
4014 }
4015
4016 fn try_fused_scan_update(
4033 &mut self,
4034 table: &str,
4035 predicate: &Expr,
4036 resolved: &[(usize, Value)],
4037 changed_cols: &[usize],
4038 ) -> Option<Result<QueryResult, QueryError>> {
4039 if self.catalog.table_has_overflow(table) {
4045 return None;
4046 }
4047 let compiled = {
4050 let schema = self.catalog.schema(table)?;
4051 let columns: Vec<String> = schema.columns.iter().map(|c| c.name.clone()).collect();
4052 let fast = FastLayout::new(schema);
4053 compile_predicate(predicate, &columns, &fast, schema)?
4054 };
4055
4056 let fixed_patches: Option<Vec<FastPatch>> = {
4058 let tbl = self.catalog.get_table(table)?;
4059 let schema = tbl.schema();
4060 let all_fixed_nonnull = resolved
4061 .iter()
4062 .all(|(idx, val)| is_fixed_size(schema.columns[*idx].type_id) && !val.is_empty());
4063 let no_indexed = !resolved.iter().any(|(idx, _)| tbl.has_indexed_col(*idx));
4064 if all_fixed_nonnull && no_indexed {
4065 let layout = RowLayout::new(schema);
4066 let bitmap_size = layout.bitmap_size();
4067 Some(
4068 resolved
4069 .iter()
4070 .map(|(idx, val)| {
4071 let fixed_off = layout
4072 .fixed_offset(*idx)
4073 .expect("is_fixed_size already checked");
4074 let field_off = 2 + bitmap_size + fixed_off;
4075 let bytes: FixedBytes = match val {
4076 Value::Int(v) => FixedBytes::I64(v.to_le_bytes()),
4077 Value::Float(v) => FixedBytes::F64(v.to_le_bytes()),
4078 Value::Bool(v) => FixedBytes::Bool(if *v { 1 } else { 0 }),
4079 Value::DateTime(v) => FixedBytes::I64(v.to_le_bytes()),
4080 Value::Uuid(v) => FixedBytes::Uuid(*v),
4081 _ => unreachable!("all_fixed_nonnull guard"),
4082 };
4083 FastPatch {
4084 field_off,
4085 bitmap_byte_off: 2 + idx / 8,
4086 bit_mask: 1u8 << (idx % 8),
4087 bytes,
4088 }
4089 })
4090 .collect(),
4091 )
4092 } else {
4093 None
4094 }
4095 };
4096 if let Some(patches) = fixed_patches {
4097 let result = self
4098 .catalog
4099 .scan_patch_matching_logged(table, compiled, |row| {
4100 let base = row_body_base(row);
4101 for p in &patches {
4102 row[base + p.bitmap_byte_off] &= !p.bit_mask;
4103 let field_bytes = p.bytes.as_slice();
4104 row[base + p.field_off..base + p.field_off + field_bytes.len()]
4105 .copy_from_slice(field_bytes);
4106 }
4107 Some(row.len() as u16)
4108 })
4109 .map_err(|e| e.to_string());
4110 match result {
4111 Ok((count, _)) => {
4112 self.view_registry.mark_dependents_dirty(table);
4113 return Some(Ok(QueryResult::Modified(count)));
4114 }
4115 Err(e) => return Some(Err(QueryError::Execution(e))),
4116 }
4117 }
4118
4119 let var_patch: Option<(usize, Option<Vec<u8>>)> = {
4121 let tbl = self.catalog.get_table(table)?;
4122 let schema = tbl.schema();
4123 let is_single = resolved.len() == 1;
4124 let is_var = is_single && !is_fixed_size(schema.columns[resolved[0].0].type_id);
4125 let no_indexed = !resolved.iter().any(|(idx, _)| tbl.has_indexed_col(*idx));
4126 if is_single && is_var && no_indexed {
4127 let (idx, val) = &resolved[0];
4128 let bytes_opt = match val {
4129 Value::Str(s) => Some(s.as_bytes().to_vec()),
4130 Value::Bytes(b) => Some(b.clone()),
4131 Value::Empty => None,
4132 _ => return None, };
4134 Some((*idx, bytes_opt))
4135 } else {
4136 None
4137 }
4138 };
4139 if let Some((col_idx, ref new_bytes_opt)) = var_patch {
4140 let layout = {
4142 let schema = self.catalog.schema(table)?;
4143 RowLayout::new(schema)
4144 };
4145 let new_bytes_ref: Option<&[u8]> = new_bytes_opt.as_deref();
4146 let result = self
4147 .catalog
4148 .scan_patch_matching_logged(table, compiled, |row| {
4149 patch_var_column_in_place(row, &layout, col_idx, new_bytes_ref)
4150 })
4151 .map_err(|e| e.to_string());
4152 match result {
4153 Ok((mut count, fallback_rids)) => {
4154 for rid in fallback_rids {
4156 let mut row = match self.catalog.get(table, rid) {
4157 Some(r) => r,
4158 None => continue,
4159 };
4160 for (idx, val) in resolved.iter() {
4161 row[*idx] = val.clone();
4162 }
4163 if let Err(e) =
4164 self.catalog
4165 .update_hinted(table, rid, &row, Some(changed_cols))
4166 {
4167 return Some(Err(QueryError::StorageError(e.to_string())));
4168 }
4169 count += 1;
4170 }
4171 self.view_registry.mark_dependents_dirty(table);
4172 return Some(Ok(QueryResult::Modified(count)));
4173 }
4174 Err(e) => return Some(Err(QueryError::Execution(e))),
4175 }
4176 }
4177
4178 None }
4180
4181 fn index_scan_rids(&self, scan: &PlanNode) -> Result<Option<Vec<RowId>>, QueryError> {
4189 match scan {
4190 PlanNode::IndexScan { table, column, key } => {
4191 let Some(tbl) = self.catalog.get_table(table) else {
4192 return Ok(None);
4193 };
4194 if !tbl.has_index(column) {
4195 return Ok(None);
4196 }
4197 let key_value = literal_to_value(key)?;
4198 Ok(Some(tbl.index_lookup_all(column, &key_value)))
4199 }
4200 PlanNode::ExprIndexScan { table, path, key } => {
4201 let Some(index) = resolve_expression_index(&self.catalog, table, path) else {
4202 return Ok(None);
4203 };
4204 let key_value = literal_to_value(key)?;
4205 let rids = if key_value.is_empty() {
4206 self.catalog
4207 .expression_index_btree(table, index.index_id)
4208 .ok_or_else(|| {
4209 QueryError::Execution("expression index disappeared".to_string())
4210 })?
4211 .empty_rids()
4212 .to_vec()
4213 } else {
4214 self.catalog
4215 .expression_index_lookup_all(table, index.index_id, &key_value)
4216 .map_err(|error| QueryError::StorageError(error.to_string()))?
4217 };
4218 Ok(Some(rids))
4219 }
4220 PlanNode::RangeScan {
4221 table,
4222 column,
4223 start,
4224 end,
4225 } => {
4226 let Some(tbl) = self.catalog.get_table(table) else {
4227 return Ok(None);
4228 };
4229 let start_val = start
4230 .as_ref()
4231 .map(|(expr, _)| literal_to_value(expr))
4232 .transpose()?;
4233 let end_val = end
4234 .as_ref()
4235 .map(|(expr, _)| literal_to_value(expr))
4236 .transpose()?;
4237 let start_inclusive = start.as_ref().map(|(_, inc)| *inc).unwrap_or(true);
4238 let end_inclusive = end.as_ref().map(|(_, inc)| *inc).unwrap_or(true);
4239 match tbl.is_index_unique(column) {
4242 Some(false) => {
4243 let col_idx = tbl.schema().column_index(column).ok_or_else(|| {
4244 QueryError::ColumnNotFound {
4245 table: String::new(),
4246 column: column.clone(),
4247 }
4248 })?;
4249 let Some(btree) = tbl.index(column) else {
4250 return Ok(None);
4251 };
4252 let candidates = btree.range_rids(start_val.as_ref(), end_val.as_ref());
4256 let mut rids = Vec::with_capacity(candidates.len());
4257 let mut cancel = CancelCheck::new();
4258 for rid in candidates {
4259 cancel.tick()?;
4260 if let Some(row) = tbl.get(rid) {
4261 if !row[col_idx].is_empty()
4262 && range_matches(
4263 &row[col_idx],
4264 &start_val,
4265 start_inclusive,
4266 &end_val,
4267 end_inclusive,
4268 )
4269 {
4270 rids.push(rid);
4271 }
4272 }
4273 }
4274 Ok(Some(rids))
4275 }
4276 Some(true) => {
4277 let Some(btree) = tbl.index(column) else {
4278 return Ok(None);
4279 };
4280 let hits: Vec<(Value, RowId)> = match (&start_val, &end_val) {
4284 (Some(s), Some(e)) => btree.range(s, e).collect(),
4285 (Some(s), None) => btree.range_from(s),
4286 (None, Some(e)) => btree.range_to(e),
4287 (None, None) => return Ok(None),
4288 };
4289 let mut rids = Vec::with_capacity(hits.len());
4290 let mut cancel = CancelCheck::new();
4291 for (key, rid) in hits {
4292 cancel.tick()?;
4293 if !start_inclusive {
4294 if let Some(ref s) = start_val {
4295 if &key == s {
4296 continue;
4297 }
4298 }
4299 }
4300 if !end_inclusive {
4301 if let Some(ref e) = end_val {
4302 if &key == e {
4303 continue;
4304 }
4305 }
4306 }
4307 rids.push(rid);
4308 }
4309 Ok(Some(rids))
4310 }
4311 None => Ok(None),
4312 }
4313 }
4314 PlanNode::ExprRangeScan {
4315 table,
4316 path,
4317 start,
4318 end,
4319 } => {
4320 let Some(index) = resolve_expression_index(&self.catalog, table, path) else {
4321 return Ok(None);
4322 };
4323 let start_val = start
4324 .as_ref()
4325 .map(|(expr, _)| literal_to_value(expr))
4326 .transpose()?;
4327 let end_val = end
4328 .as_ref()
4329 .map(|(expr, _)| literal_to_value(expr))
4330 .transpose()?;
4331 let start_inclusive = start.as_ref().map(|(_, inc)| *inc).unwrap_or(true);
4332 let end_inclusive = end.as_ref().map(|(_, inc)| *inc).unwrap_or(true);
4333 let candidates = self
4334 .catalog
4335 .expression_index_range_rids(
4336 table,
4337 index.index_id,
4338 start_val.as_ref(),
4339 end_val.as_ref(),
4340 )
4341 .map_err(|error| QueryError::StorageError(error.to_string()))?;
4342 let schema = self
4343 .catalog
4344 .schema(table)
4345 .ok_or_else(|| QueryError::TableNotFound(table.to_string()))?;
4346 let all_columns: Vec<String> =
4347 schema.columns.iter().map(|c| c.name.clone()).collect();
4348 let path_expr = stored_json_path_expr(path);
4349 let mut rids = Vec::with_capacity(candidates.len());
4350 let mut cancel = CancelCheck::new();
4351 for rid in candidates {
4352 cancel.tick()?;
4353 let Some(row) = self.catalog.get(table, rid) else {
4354 continue;
4355 };
4356 let value = eval_expr(&path_expr, &row, &all_columns);
4357 if value.is_empty()
4358 || !range_matches(
4359 &value,
4360 &start_val,
4361 start_inclusive,
4362 &end_val,
4363 end_inclusive,
4364 )
4365 {
4366 continue;
4367 }
4368 rids.push(rid);
4369 }
4370 Ok(Some(rids))
4371 }
4372 _ => Ok(None),
4373 }
4374 }
4375
4376 fn collect_rids_via_index_residual(
4384 &self,
4385 inner: &PlanNode,
4386 predicate: &Expr,
4387 table: &str,
4388 ) -> Result<Option<Vec<RowId>>, QueryError> {
4389 if contains_subquery(predicate) || scan_table(inner) != Some(table) {
4390 return Ok(None);
4391 }
4392 let Some(candidates) = self.index_scan_rids(inner)? else {
4393 return Ok(None);
4394 };
4395 let schema = self
4396 .catalog
4397 .schema(table)
4398 .ok_or_else(|| QueryError::TableNotFound(table.to_string()))?;
4399 let all_columns: Vec<String> = schema.columns.iter().map(|c| c.name.clone()).collect();
4400 let residual_indices = predicate_column_indices_json(predicate, &all_columns);
4401 let residual_names: Vec<String> = residual_indices
4402 .iter()
4403 .map(|&index| all_columns[index].clone())
4404 .collect();
4405 let mut rids = Vec::new();
4406 let mut cancel = CancelCheck::new();
4407 for rid in candidates {
4408 cancel.tick()?;
4409 let Some(sparse) = self
4410 .catalog
4411 .get_projected(table, rid, &residual_indices)
4412 .map_err(|error| QueryError::StorageError(error.to_string()))?
4413 else {
4414 continue;
4415 };
4416 if eval_predicate(predicate, &sparse, &residual_names) {
4417 rids.push(rid);
4418 }
4419 }
4420 Ok(Some(rids))
4421 }
4422
4423 fn collect_rids_for_mutation(
4429 &mut self,
4430 input: &PlanNode,
4431 table: &str,
4432 ) -> Result<Vec<RowId>, QueryError> {
4433 if self.catalog.table_has_overflow(table) {
4441 if let Some(rids) = self.collect_rids_decoded(input, table)? {
4442 return Ok(rids);
4443 }
4444 }
4445 match input {
4446 PlanNode::SeqScan { table: t } if t == table => {
4447 let mut cancel = CancelCheck::new();
4449 let mut rids: Vec<RowId> = Vec::new();
4450 for (rid, _) in self
4451 .catalog
4452 .scan(table)
4453 .map_err(|e| QueryError::StorageError(e.to_string()))?
4454 {
4455 cancel.tick()?;
4456 rids.push(rid);
4457 }
4458 Ok(rids)
4459 }
4460 PlanNode::IndexScan {
4461 table: t,
4462 column,
4463 key,
4464 } if t == table => {
4465 let key_value = literal_to_value(key)?;
4466
4467 {
4476 let tbl = self
4477 .catalog
4478 .get_table(table)
4479 .ok_or_else(|| QueryError::TableNotFound(table.to_string()))?;
4480 if tbl.has_index(column) {
4481 let rids = tbl.index_lookup_all(column, &key_value);
4482 return Ok(rids);
4483 }
4484 }
4485
4486 let schema = self
4491 .catalog
4492 .schema(table)
4493 .ok_or_else(|| QueryError::TableNotFound(table.to_string()))?;
4494 let columns: Vec<String> = schema.columns.iter().map(|c| c.name.clone()).collect();
4495 let fast = FastLayout::new(schema);
4496 let synth = Expr::BinaryOp(
4497 Box::new(Expr::Field(column.clone())),
4498 BinOp::Eq,
4499 Box::new(key.clone()),
4500 );
4501 if let Some(compiled) = compile_predicate(&synth, &columns, &fast, schema) {
4502 let mut rids: Vec<RowId> = Vec::with_capacity(64);
4504 let mut cancel = CancelCheck::new();
4505 let mut cancel_err: Option<QueryError> = None;
4506 self.catalog
4507 .try_for_each_row_raw(table, |rid, data| {
4508 if let Err(e) = cancel.tick() {
4509 cancel_err = Some(e);
4510 return ControlFlow::Break(());
4511 }
4512 if compiled(data) {
4513 rids.push(rid);
4514 }
4515 ControlFlow::Continue(())
4516 })
4517 .map_err(|e| QueryError::StorageError(e.to_string()))?;
4518 if let Some(e) = cancel_err {
4519 return Err(e);
4520 }
4521 return Ok(rids);
4522 }
4523
4524 let col_idx =
4526 schema
4527 .column_index(column)
4528 .ok_or_else(|| QueryError::ColumnNotFound {
4529 table: String::new(),
4530 column: column.clone(),
4531 })?;
4532 let mut cancel = CancelCheck::new();
4533 let mut rids: Vec<RowId> = Vec::new();
4534 for (rid, row) in self
4535 .catalog
4536 .scan(table)
4537 .map_err(|e| QueryError::StorageError(e.to_string()))?
4538 {
4539 cancel.tick()?;
4540 if row[col_idx] == key_value {
4541 rids.push(rid);
4542 }
4543 }
4544 Ok(rids)
4545 }
4546 PlanNode::RangeScan { table: t, .. }
4547 | PlanNode::ExprIndexScan { table: t, .. }
4548 | PlanNode::ExprRangeScan { table: t, .. }
4549 if t == table =>
4550 {
4551 match self.index_scan_rids(input)? {
4555 Some(rids) => Ok(rids),
4556 None => self.generic_rid_match(input, table),
4557 }
4558 }
4559 PlanNode::Filter {
4560 input: inner,
4561 predicate,
4562 } => {
4563 if let PlanNode::SeqScan { table: t } = inner.as_ref() {
4564 if t != table {
4565 return self.generic_rid_match(input, table);
4566 }
4567 let schema = self
4568 .catalog
4569 .schema(table)
4570 .ok_or_else(|| QueryError::TableNotFound(table.to_string()))?;
4571 let columns: Vec<String> =
4572 schema.columns.iter().map(|c| c.name.clone()).collect();
4573 let fast = FastLayout::new(schema);
4574 let row_layout = RowLayout::new(schema);
4575
4576 let mut cancel = CancelCheck::new();
4580 let mut cancel_err: Option<QueryError> = None;
4581 if let Some(compiled) = compile_predicate(predicate, &columns, &fast, schema) {
4583 let mut rids: Vec<RowId> = Vec::with_capacity(64);
4585 self.catalog
4586 .try_for_each_row_raw(table, |rid, data| {
4587 if let Err(e) = cancel.tick() {
4588 cancel_err = Some(e);
4589 return ControlFlow::Break(());
4590 }
4591 if compiled(data) {
4592 rids.push(rid);
4593 }
4594 ControlFlow::Continue(())
4595 })
4596 .map_err(|e| QueryError::StorageError(e.to_string()))?;
4597 if let Some(e) = cancel_err {
4598 return Err(e);
4599 }
4600 return Ok(rids);
4601 }
4602
4603 let pred_cols = predicate_column_indices_json(predicate, &columns);
4605 let mut rids: Vec<RowId> = Vec::with_capacity(64);
4606 self.catalog
4607 .try_for_each_row_raw(table, |rid, data| {
4608 if let Err(e) = cancel.tick() {
4609 cancel_err = Some(e);
4610 return ControlFlow::Break(());
4611 }
4612 let pred_row = decode_selective(schema, &row_layout, data, &pred_cols);
4613 if eval_predicate(predicate, &pred_row, &columns) {
4614 rids.push(rid);
4615 }
4616 ControlFlow::Continue(())
4617 })
4618 .map_err(|e| QueryError::StorageError(e.to_string()))?;
4619 if let Some(e) = cancel_err {
4620 return Err(e);
4621 }
4622 return Ok(rids);
4623 }
4624 if let Some(rids) = self.collect_rids_via_index_residual(inner, predicate, table)? {
4629 return Ok(rids);
4630 }
4631 self.generic_rid_match(input, table)
4632 }
4633 _ => self.generic_rid_match(input, table),
4634 }
4635 }
4636
4637 fn collect_rids_decoded(
4644 &mut self,
4645 input: &PlanNode,
4646 table: &str,
4647 ) -> Result<Option<Vec<RowId>>, QueryError> {
4648 let pred: Option<Expr> = match input {
4650 PlanNode::SeqScan { table: t } if t == table => None,
4651 PlanNode::Filter {
4652 input: inner,
4653 predicate,
4654 } => match inner.as_ref() {
4655 PlanNode::SeqScan { table: t } if t == table => Some(predicate.clone()),
4656 _ => return Ok(None),
4657 },
4658 PlanNode::IndexScan {
4659 table: t,
4660 column,
4661 key,
4662 } if t == table => {
4663 let indexed = self
4666 .catalog
4667 .get_table(table)
4668 .map(|tb| tb.has_index(column))
4669 .unwrap_or(false);
4670 if indexed {
4671 return Ok(None);
4672 }
4673 Some(Expr::BinaryOp(
4674 Box::new(Expr::Field(column.clone())),
4675 BinOp::Eq,
4676 Box::new(key.clone()),
4677 ))
4678 }
4679 _ => return Ok(None),
4680 };
4681
4682 let columns: Vec<String> = {
4683 let schema = self
4684 .catalog
4685 .schema(table)
4686 .ok_or_else(|| QueryError::TableNotFound(table.to_string()))?;
4687 schema.columns.iter().map(|c| c.name.clone()).collect()
4688 };
4689 let mut rids: Vec<RowId> = Vec::new();
4690 let mut cancel = CancelCheck::new();
4691 for (rid, row) in self
4692 .catalog
4693 .scan(table)
4694 .map_err(|e| QueryError::StorageError(e.to_string()))?
4695 {
4696 cancel.tick()?;
4697 let keep = match &pred {
4698 None => true,
4699 Some(p) => eval_predicate(p, &row, &columns),
4700 };
4701 if keep {
4702 rids.push(rid);
4703 }
4704 }
4705 Ok(Some(rids))
4706 }
4707
4708 fn generic_rid_match(
4712 &mut self,
4713 input: &PlanNode,
4714 table: &str,
4715 ) -> Result<Vec<RowId>, QueryError> {
4716 #[cfg(test)]
4717 GENERIC_RID_MATCH_CALLS.with(|calls| calls.set(calls.get() + 1));
4718 let result = self.execute_plan(input)?;
4719 let rows = match result {
4720 QueryResult::Rows { rows, .. } => rows,
4721 _ => return Err("mutation source must be rows".into()),
4722 };
4723 let mut matching: Vec<RowId> = Vec::new();
4724 let mut cancel = CancelCheck::new();
4725 for (rid, row) in self
4726 .catalog
4727 .scan(table)
4728 .map_err(|e| QueryError::StorageError(e.to_string()))?
4729 {
4730 cancel.tick()?;
4731 let mut matched = false;
4732 for candidate in &rows {
4733 cancel.tick()?;
4734 if candidate == &row {
4735 matched = true;
4736 break;
4737 }
4738 }
4739 if matched {
4740 matching.push(rid);
4741 }
4742 }
4743 Ok(matching)
4744 }
4745}
4746
4747pub(super) fn execute_window(
4748 result: QueryResult,
4749 windows: &[WindowDef],
4750 memory_limit: usize,
4751) -> Result<QueryResult, QueryError> {
4752 let (mut columns, mut rows) = match result {
4753 QueryResult::Rows { columns, rows } => (columns, rows),
4754 _ => return Err("window function requires row input".into()),
4755 };
4756
4757 let mut cancel = CancelCheck::new();
4758 for wdef in windows {
4759 cancel.tick()?;
4760 let part_indices: Vec<Option<usize>> = wdef
4763 .partition_by
4764 .iter()
4765 .map(|expr| resolve_direct_group_expr(expr, &columns))
4766 .collect::<Result<Vec<_>, _>>()?;
4767
4768 let ord_indices: Vec<(Option<usize>, &Expr, bool)> = wdef
4769 .order_by
4770 .iter()
4771 .map(|sk| {
4772 resolve_direct_group_expr(&sk.expr, &columns)
4773 .map(|index| (index, &sk.expr, sk.descending))
4774 })
4775 .collect::<Result<Vec<_>, _>>()?;
4776
4777 let arg_expr = wdef.args.first();
4778 let arg_col_idx = arg_expr
4779 .map(|expr| resolve_direct_group_expr(expr, &columns))
4780 .transpose()?
4781 .flatten();
4782
4783 let n = rows.len();
4787 let mut indices: Vec<usize> = (0..n).collect();
4788 cooperative_stable_sort_by(&mut indices, memory_limit, |&a, &b| {
4789 for (expr, index) in wdef.partition_by.iter().zip(&part_indices) {
4791 let av = index
4792 .map(|i| rows[a][i].clone())
4793 .unwrap_or_else(|| eval_expr(expr, &rows[a], &columns));
4794 let bv = index
4795 .map(|i| rows[b][i].clone())
4796 .unwrap_or_else(|| eval_expr(expr, &rows[b], &columns));
4797 let cmp = av.cmp(&bv);
4798 if cmp != std::cmp::Ordering::Equal {
4799 return cmp;
4800 }
4801 }
4802 for &(index, expr, desc) in &ord_indices {
4804 let av = index
4805 .map(|i| rows[a][i].clone())
4806 .unwrap_or_else(|| eval_expr(expr, &rows[a], &columns));
4807 let bv = index
4808 .map(|i| rows[b][i].clone())
4809 .unwrap_or_else(|| eval_expr(expr, &rows[b], &columns));
4810 let cmp = compare_order_values(&av, &bv, desc);
4811 if cmp != std::cmp::Ordering::Equal {
4812 return cmp;
4813 }
4814 }
4815 std::cmp::Ordering::Equal
4816 })?;
4817
4818 let whole_partition_frame = wdef.order_by.is_empty()
4826 && matches!(
4827 wdef.function,
4828 WindowFunc::Sum
4829 | WindowFunc::Avg
4830 | WindowFunc::Count
4831 | WindowFunc::Min
4832 | WindowFunc::Max
4833 );
4834 let mut partition_row_indices: Vec<usize> = Vec::new();
4837
4838 let mut win_values: Vec<Value> = vec![Value::Empty; n];
4840 let mut partition_start = 0usize;
4841 let mut running_count: i64 = 0;
4843 let mut running_int_sum: i64 = 0;
4844 let mut running_float_sum: f64 = 0.0;
4845 let mut running_saw_float = false;
4846 let mut running_min: Option<Value> = None;
4847 let mut running_max: Option<Value> = None;
4848 let mut rank_counter: i64 = 0;
4849 let mut dense_rank_counter: i64 = 0;
4850 let mut prev_order_key: Option<Vec<Value>> = None;
4851 let mut same_rank_count: i64 = 0;
4852
4853 for sorted_pos in 0..n {
4854 cancel.tick()?;
4855 let row_idx = indices[sorted_pos];
4856
4857 let new_partition = if sorted_pos == 0 {
4859 true
4860 } else {
4861 let prev_row_idx = indices[sorted_pos - 1];
4862 wdef.partition_by
4863 .iter()
4864 .zip(&part_indices)
4865 .any(|(expr, index)| {
4866 let current = index
4867 .map(|i| rows[row_idx][i].clone())
4868 .unwrap_or_else(|| eval_expr(expr, &rows[row_idx], &columns));
4869 let previous = index
4870 .map(|i| rows[prev_row_idx][i].clone())
4871 .unwrap_or_else(|| eval_expr(expr, &rows[prev_row_idx], &columns));
4872 current != previous
4873 })
4874 };
4875
4876 if new_partition {
4877 if whole_partition_frame && sorted_pos > 0 {
4881 let final_v = win_values[indices[sorted_pos - 1]].clone();
4882 for ri in partition_row_indices.drain(..) {
4883 cancel.tick()?;
4884 win_values[ri] = final_v.clone();
4885 }
4886 }
4887 partition_start = sorted_pos;
4888 running_count = 0;
4889 running_int_sum = 0;
4890 running_float_sum = 0.0;
4891 running_saw_float = false;
4892 running_min = None;
4893 running_max = None;
4894 rank_counter = 0;
4895 dense_rank_counter = 0;
4896 prev_order_key = None;
4897 same_rank_count = 0;
4898 }
4899
4900 let current_order_key: Vec<Value> = ord_indices
4902 .iter()
4903 .map(|&(index, expr, _)| {
4904 index
4905 .map(|i| rows[row_idx][i].clone())
4906 .unwrap_or_else(|| eval_expr(expr, &rows[row_idx], &columns))
4907 })
4908 .collect();
4909 let same_as_prev = prev_order_key.as_ref() == Some(¤t_order_key);
4910 let current_arg = || {
4911 arg_expr.map(|expr| {
4912 arg_col_idx
4913 .map(|index| rows[row_idx][index].clone())
4914 .unwrap_or_else(|| eval_expr(expr, &rows[row_idx], &columns))
4915 })
4916 };
4917 let count_all =
4918 arg_expr.is_none() || matches!(arg_expr, Some(Expr::Field(name)) if name == "*");
4919
4920 let value = match wdef.function {
4921 WindowFunc::RowNumber => Value::Int((sorted_pos - partition_start + 1) as i64),
4922 WindowFunc::Rank => {
4923 if same_as_prev {
4924 same_rank_count += 1;
4925 } else {
4926 rank_counter += same_rank_count + 1;
4927 same_rank_count = 0;
4928 if rank_counter == 0 {
4929 rank_counter = 1;
4930 }
4931 }
4932 Value::Int(rank_counter)
4933 }
4934 WindowFunc::DenseRank => {
4935 if !same_as_prev {
4936 dense_rank_counter += 1;
4937 }
4938 Value::Int(dense_rank_counter)
4939 }
4940 WindowFunc::Sum => {
4941 if let Some(value) = current_arg() {
4942 match value {
4943 Value::Int(v) => running_int_sum += v,
4944 Value::Float(v) => {
4945 running_float_sum += v;
4946 running_saw_float = true;
4947 }
4948 _ => {}
4949 }
4950 }
4951 if running_saw_float {
4952 Value::Float(running_float_sum + running_int_sum as f64)
4953 } else {
4954 Value::Int(running_int_sum)
4955 }
4956 }
4957 WindowFunc::Avg => {
4958 if let Some(value) = current_arg() {
4959 match value {
4960 Value::Int(v) => {
4961 running_float_sum += v as f64;
4962 running_count += 1;
4963 }
4964 Value::Float(v) => {
4965 running_float_sum += v;
4966 running_count += 1;
4967 }
4968 _ => {}
4969 }
4970 }
4971 if running_count == 0 {
4972 Value::Empty
4973 } else {
4974 Value::Float(running_float_sum / running_count as f64)
4975 }
4976 }
4977 WindowFunc::Count => {
4978 if count_all {
4979 running_count += 1;
4980 } else if let Some(value) = current_arg() {
4981 if !value.is_empty() {
4982 running_count += 1;
4983 }
4984 }
4985 Value::Int(running_count)
4986 }
4987 WindowFunc::Min => {
4988 if let Some(v) = current_arg() {
4989 if !v.is_empty() {
4990 running_min = Some(match &running_min {
4991 None => v,
4992 Some(cur) => {
4993 if v < *cur {
4994 v
4995 } else {
4996 cur.clone()
4997 }
4998 }
4999 });
5000 }
5001 }
5002 running_min.clone().unwrap_or(Value::Empty)
5003 }
5004 WindowFunc::Max => {
5005 if let Some(v) = current_arg() {
5006 if !v.is_empty() {
5007 running_max = Some(match &running_max {
5008 None => v,
5009 Some(cur) => {
5010 if v > *cur {
5011 v
5012 } else {
5013 cur.clone()
5014 }
5015 }
5016 });
5017 }
5018 }
5019 running_max.clone().unwrap_or(Value::Empty)
5020 }
5021 };
5022
5023 prev_order_key = Some(current_order_key);
5024 win_values[row_idx] = value;
5025 if whole_partition_frame {
5026 partition_row_indices.push(row_idx);
5027 }
5028 }
5029
5030 if whole_partition_frame && n > 0 {
5032 let final_v = win_values[indices[n - 1]].clone();
5033 for ri in partition_row_indices.drain(..) {
5034 cancel.tick()?;
5035 win_values[ri] = final_v.clone();
5036 }
5037 }
5038
5039 for (ri, row) in rows.iter_mut().enumerate() {
5041 cancel.tick()?;
5042 row.push(win_values[ri].clone());
5043 }
5044 columns.push(wdef.output_name.clone());
5045 }
5046
5047 Ok(QueryResult::Rows { columns, rows })
5048}
5049
5050pub(super) fn resolve_group_column(name: &str, columns: &[String]) -> Result<usize, QueryError> {
5063 if let Some(i) = columns.iter().position(|c| c == name) {
5064 return Ok(i);
5065 }
5066 if name.contains('.') {
5067 return Err(QueryError::ColumnNotFound {
5068 table: String::new(),
5069 column: name.to_string(),
5070 });
5071 }
5072 let suffix = format!(".{name}");
5073 let mut matches = columns
5074 .iter()
5075 .enumerate()
5076 .filter(|(_, c)| c.ends_with(&suffix));
5077 match matches.next() {
5078 None => Err(QueryError::ColumnNotFound {
5079 table: String::new(),
5080 column: name.to_string(),
5081 }),
5082 Some((first_idx, _)) => {
5083 let rest: Vec<&str> = matches.map(|(_, c)| c.as_str()).collect();
5084 if rest.is_empty() {
5085 Ok(first_idx)
5086 } else {
5087 let candidates: Vec<&str> = columns
5090 .iter()
5091 .filter(|c| c.ends_with(&suffix))
5092 .map(|c| c.as_str())
5093 .collect();
5094 Err(QueryError::Execution(format!(
5095 "cannot group by ambiguous column '{name}'; candidates: {}",
5096 candidates.join(", ")
5097 )))
5098 }
5099 }
5100 }
5101}
5102
5103pub(super) fn exec_group_by(
5108 columns: Vec<String>,
5109 rows: Vec<Vec<Value>>,
5110 keys: &[GroupKey],
5111 aggregates: &[GroupAgg],
5112 having: &Option<Expr>,
5113) -> Result<QueryResult, QueryError> {
5114 exec_group_by_internal(columns, rows, None, keys, aggregates, having)
5115}
5116
5117pub(super) fn exec_group_by_with_provenance(
5118 input: ProvenanceRows,
5119 keys: &[GroupKey],
5120 aggregates: &[GroupAgg],
5121 having: &Option<Expr>,
5122 memory_limit: usize,
5123) -> Result<QueryResult, QueryError> {
5124 let ProvenanceRows {
5125 columns,
5126 rows,
5127 source_aliases,
5128 provenance,
5129 } = input;
5130 exec_group_by_internal(
5131 columns,
5132 rows,
5133 Some(GroupProvenance {
5134 source_aliases,
5135 rows: provenance,
5136 memory_limit,
5137 }),
5138 keys,
5139 aggregates,
5140 having,
5141 )
5142}
5143
5144struct GroupProvenance {
5145 source_aliases: Vec<String>,
5146 rows: Vec<Vec<Option<RowId>>>,
5147 memory_limit: usize,
5148}
5149
5150fn exec_group_by_internal(
5151 columns: Vec<String>,
5152 rows: Vec<Vec<Value>>,
5153 provenance: Option<GroupProvenance>,
5154 keys: &[GroupKey],
5155 aggregates: &[GroupAgg],
5156 having: &Option<Expr>,
5157) -> Result<QueryResult, QueryError> {
5158 let key_indices: Vec<Option<usize>> = keys
5161 .iter()
5162 .map(|k| resolve_direct_group_expr(&k.expr, &columns))
5163 .collect::<Result<Vec<_>, _>>()?;
5164
5165 let agg_field_indices: Vec<Option<usize>> = aggregates
5166 .iter()
5167 .map(|a| resolve_direct_group_expr(&a.argument, &columns))
5168 .collect::<Result<Vec<_>, _>>()?;
5169 let agg_source_indices: Vec<Option<usize>> = aggregates
5170 .iter()
5171 .map(|aggregate| {
5172 aggregate
5173 .provenance_alias
5174 .as_ref()
5175 .map(|alias| {
5176 provenance
5177 .as_ref()
5178 .and_then(|provenance| {
5179 provenance
5180 .source_aliases
5181 .iter()
5182 .position(|source| source == alias)
5183 })
5184 .ok_or_else(|| {
5185 QueryError::Execution(format!(
5186 "symmetric aggregate source alias '{alias}' is not present in its input"
5187 ))
5188 })
5189 })
5190 .transpose()
5191 })
5192 .collect::<Result<Vec<_>, _>>()?;
5193
5194 let mut group_map: rustc_hash::FxHashMap<Vec<Value>, usize> = rustc_hash::FxHashMap::default();
5196 let mut groups: Vec<(Vec<Value>, Vec<usize>)> = Vec::new();
5197 let mut cancel = CancelCheck::new();
5198 for (ri, row) in rows.iter().enumerate() {
5199 cancel.tick()?;
5200 let key: Vec<Value> = keys
5201 .iter()
5202 .zip(&key_indices)
5203 .map(|(key, index)| match index {
5204 Some(index) => row[*index].clone(),
5205 None => eval_expr(&key.expr, row, &columns),
5206 })
5207 .collect();
5208 match group_map.get(&key) {
5209 Some(&idx) => groups[idx].1.push(ri),
5210 None => {
5211 let idx = groups.len();
5212 group_map.insert(key.clone(), idx);
5213 groups.push((key, vec![ri]));
5214 }
5215 }
5216 }
5217
5218 let mut out_columns: Vec<String> = keys.iter().map(|k| k.output_name()).collect();
5222 for agg in aggregates.iter() {
5223 out_columns.push(agg.output_name.clone());
5224 }
5225
5226 let mut out_rows: Vec<Vec<Value>> = Vec::with_capacity(groups.len());
5228 for (key_vals, row_indices) in &groups {
5229 cancel.tick()?;
5230 let mut row = key_vals.clone();
5231 for (ai, agg) in aggregates.iter().enumerate() {
5232 let val = compute_group_aggregate(
5233 agg.function,
5234 &agg.argument,
5235 agg_field_indices[ai],
5236 GroupAggregateContext {
5237 columns: &columns,
5238 all_rows: &rows,
5239 row_indices,
5240 source_index: agg_source_indices[ai],
5241 provenance: provenance
5242 .as_ref()
5243 .map(|provenance| (provenance.rows.as_slice(), provenance.memory_limit)),
5244 },
5245 )?;
5246 row.push(val);
5247 }
5248 out_rows.push(row);
5249 }
5250
5251 if let Some(having_expr) = having {
5253 let mut filtered = Vec::with_capacity(out_rows.len());
5254 for row in out_rows {
5255 cancel.tick()?;
5256 if eval_predicate(having_expr, &row, &out_columns) {
5257 filtered.push(row);
5258 }
5259 }
5260 out_rows = filtered;
5261 }
5262
5263 Ok(QueryResult::Rows {
5264 columns: out_columns,
5265 rows: out_rows,
5266 })
5267}
5268
5269fn resolve_direct_group_expr(expr: &Expr, columns: &[String]) -> Result<Option<usize>, QueryError> {
5270 match expr {
5271 Expr::Field(name) if name == "*" => Ok(None),
5272 Expr::Field(name) => resolve_group_column(name, columns).map(Some),
5273 Expr::QualifiedField { qualifier, field } => {
5274 resolve_group_column(&format!("{qualifier}.{field}"), columns).map(Some)
5275 }
5276 _ => Ok(None),
5277 }
5278}
5279
5280pub(super) fn predicate_column_indices_json(expr: &Expr, columns: &[String]) -> Vec<usize> {
5295 let mut indices = predicate_column_indices(expr, columns);
5296 collect_json_path_base_indices(expr, columns, &mut indices);
5297 indices.sort_unstable();
5298 indices.dedup();
5299 indices
5300}
5301
5302fn collect_json_path_base_indices(expr: &Expr, columns: &[String], out: &mut Vec<usize>) {
5304 match expr {
5305 Expr::JsonPath { base, .. } => {
5306 let name = match base.as_ref() {
5307 Expr::Field(n) => n.clone(),
5308 Expr::QualifiedField { qualifier, field } => format!("{qualifier}.{field}"),
5309 other => {
5310 collect_json_path_base_indices(other, columns, out);
5311 return;
5312 }
5313 };
5314 if let Some(idx) = columns.iter().position(|c| *c == name) {
5315 out.push(idx);
5316 }
5317 }
5318 Expr::BinaryOp(l, _, r) | Expr::Coalesce(l, r) => {
5319 collect_json_path_base_indices(l, columns, out);
5320 collect_json_path_base_indices(r, columns, out);
5321 }
5322 Expr::UnaryOp(_, i) | Expr::FunctionCall(_, i, _) | Expr::Cast(i, _) => {
5323 collect_json_path_base_indices(i, columns, out);
5324 }
5325 Expr::ScalarFunc(_, args) => {
5326 for a in args {
5327 collect_json_path_base_indices(a, columns, out);
5328 }
5329 }
5330 Expr::InList { expr, list, .. } => {
5331 collect_json_path_base_indices(expr, columns, out);
5332 for item in list {
5333 collect_json_path_base_indices(item, columns, out);
5334 }
5335 }
5336 Expr::InSubquery { expr, .. } => collect_json_path_base_indices(expr, columns, out),
5337 Expr::Case { whens, else_expr } => {
5338 for (c, r) in whens {
5339 collect_json_path_base_indices(c, columns, out);
5340 collect_json_path_base_indices(r, columns, out);
5341 }
5342 if let Some(e) = else_expr {
5343 collect_json_path_base_indices(e, columns, out);
5344 }
5345 }
5346 _ => {}
5347 }
5348}
5349
5350pub(super) fn validate_json_path_types(
5361 catalog: &Catalog,
5362 plan: &PlanNode,
5363) -> Result<(), QueryError> {
5364 let mut scope: Vec<(String, TypeId)> = Vec::new();
5365 collect_scan_columns(catalog, plan, &mut scope);
5366 let mut shadowed: std::collections::HashSet<String> = std::collections::HashSet::new();
5367 collect_projected_names(plan, &mut shadowed);
5368 check_plan_json_paths(plan, &scope, &shadowed)
5369}
5370
5371fn collect_scan_columns(catalog: &Catalog, plan: &PlanNode, out: &mut Vec<(String, TypeId)>) {
5375 match plan {
5376 PlanNode::SeqScan { table }
5377 | PlanNode::IndexScan { table, .. }
5378 | PlanNode::RangeScan { table, .. } => {
5379 if let Some(schema) = catalog.schema(table) {
5380 for c in &schema.columns {
5381 out.push((c.name.clone(), c.type_id));
5382 }
5383 }
5384 }
5385 PlanNode::AliasScan { table, alias } => {
5386 if let Some(schema) = catalog.schema(table) {
5387 for c in &schema.columns {
5388 out.push((format!("{alias}.{}", c.name), c.type_id));
5389 }
5390 }
5391 }
5392 PlanNode::Filter { input, .. }
5393 | PlanNode::Project { input, .. }
5394 | PlanNode::Sort { input, .. }
5395 | PlanNode::Limit { input, .. }
5396 | PlanNode::Offset { input, .. }
5397 | PlanNode::Aggregate { input, .. }
5398 | PlanNode::Distinct { input }
5399 | PlanNode::GroupBy { input, .. }
5400 | PlanNode::Window { input, .. }
5401 | PlanNode::Update { input, .. }
5402 | PlanNode::Delete { input, .. }
5403 | PlanNode::Explain { input } => collect_scan_columns(catalog, input, out),
5404 PlanNode::NestedLoopJoin { left, right, .. } | PlanNode::Union { left, right, .. } => {
5405 collect_scan_columns(catalog, left, out);
5406 collect_scan_columns(catalog, right, out);
5407 }
5408 _ => {}
5409 }
5410}
5411
5412fn collect_projected_names(plan: &PlanNode, out: &mut std::collections::HashSet<String>) {
5415 if let PlanNode::Project { fields, .. } = plan {
5416 for f in fields {
5417 if let Some(a) = &f.alias {
5418 out.insert(a.clone());
5419 } else {
5420 match &f.expr {
5421 Expr::Field(n) => {
5422 out.insert(n.clone());
5423 }
5424 Expr::QualifiedField { qualifier, field } => {
5425 out.insert(format!("{qualifier}.{field}"));
5426 }
5427 _ => {}
5428 }
5429 }
5430 }
5431 }
5432 match plan {
5433 PlanNode::Filter { input, .. }
5434 | PlanNode::Project { input, .. }
5435 | PlanNode::Sort { input, .. }
5436 | PlanNode::Limit { input, .. }
5437 | PlanNode::Offset { input, .. }
5438 | PlanNode::Aggregate { input, .. }
5439 | PlanNode::Distinct { input }
5440 | PlanNode::GroupBy { input, .. }
5441 | PlanNode::Window { input, .. }
5442 | PlanNode::Update { input, .. }
5443 | PlanNode::Delete { input, .. }
5444 | PlanNode::Explain { input } => collect_projected_names(input, out),
5445 PlanNode::NestedLoopJoin { left, right, .. } | PlanNode::Union { left, right, .. } => {
5446 collect_projected_names(left, out);
5447 collect_projected_names(right, out);
5448 }
5449 _ => {}
5450 }
5451}
5452
5453fn resolve_scan_type(name: &str, scope: &[(String, TypeId)]) -> Option<TypeId> {
5457 let mut found: Option<TypeId> = None;
5458 for (n, t) in scope {
5459 if n == name {
5460 match found {
5461 None => found = Some(*t),
5462 Some(prev) if prev == *t => {}
5463 Some(_) => return None, }
5465 }
5466 }
5467 found
5468}
5469
5470fn json_path_base_error(
5473 base: &Expr,
5474 scope: &[(String, TypeId)],
5475 shadowed: &std::collections::HashSet<String>,
5476) -> Option<String> {
5477 let name = match base {
5478 Expr::Field(n) => n.clone(),
5479 Expr::QualifiedField { qualifier, field } => format!("{qualifier}.{field}"),
5480 _ => return None,
5483 };
5484 if shadowed.contains(&name) {
5485 return None;
5486 }
5487 match resolve_scan_type(&name, scope) {
5488 Some(TypeId::Json) | None => None,
5489 Some(other) => Some(format!(
5490 "'{}' is a {} column, not json: the '->' path operator requires a json column",
5491 name,
5492 type_id_to_name(other)
5493 )),
5494 }
5495}
5496
5497fn check_expr_json_paths(
5499 expr: &Expr,
5500 scope: &[(String, TypeId)],
5501 shadowed: &std::collections::HashSet<String>,
5502) -> Result<(), QueryError> {
5503 match expr {
5504 Expr::JsonPath { base, .. } => {
5505 if let Some(msg) = json_path_base_error(base, scope, shadowed) {
5506 return Err(QueryError::TypeError(msg));
5507 }
5508 check_expr_json_paths(base, scope, shadowed)
5509 }
5510 Expr::BinaryOp(l, _, r) | Expr::Coalesce(l, r) => {
5511 check_expr_json_paths(l, scope, shadowed)?;
5512 check_expr_json_paths(r, scope, shadowed)
5513 }
5514 Expr::UnaryOp(_, inner) | Expr::FunctionCall(_, inner, _) | Expr::Cast(inner, _) => {
5515 check_expr_json_paths(inner, scope, shadowed)
5516 }
5517 Expr::ScalarFunc(_, args) => {
5518 for a in args {
5519 check_expr_json_paths(a, scope, shadowed)?;
5520 }
5521 Ok(())
5522 }
5523 Expr::Window {
5524 args,
5525 partition_by,
5526 order_by,
5527 ..
5528 } => {
5529 for expr in args.iter().chain(partition_by) {
5530 check_expr_json_paths(expr, scope, shadowed)?;
5531 }
5532 for key in order_by {
5533 check_expr_json_paths(&key.expr, scope, shadowed)?;
5534 }
5535 Ok(())
5536 }
5537 Expr::InList { expr, list, .. } => {
5538 check_expr_json_paths(expr, scope, shadowed)?;
5539 for item in list {
5540 check_expr_json_paths(item, scope, shadowed)?;
5541 }
5542 Ok(())
5543 }
5544 Expr::Case { whens, else_expr } => {
5545 for (c, r) in whens {
5546 check_expr_json_paths(c, scope, shadowed)?;
5547 check_expr_json_paths(r, scope, shadowed)?;
5548 }
5549 if let Some(e) = else_expr {
5550 check_expr_json_paths(e, scope, shadowed)?;
5551 }
5552 Ok(())
5553 }
5554 Expr::InSubquery { expr, .. } => check_expr_json_paths(expr, scope, shadowed),
5557 _ => Ok(()),
5558 }
5559}
5560
5561fn check_plan_json_paths(
5563 plan: &PlanNode,
5564 scope: &[(String, TypeId)],
5565 shadowed: &std::collections::HashSet<String>,
5566) -> Result<(), QueryError> {
5567 match plan {
5568 PlanNode::Filter { input, predicate } => {
5569 check_expr_json_paths(predicate, scope, shadowed)?;
5570 check_plan_json_paths(input, scope, shadowed)
5571 }
5572 PlanNode::Project { input, fields } => {
5573 for f in fields {
5574 check_expr_json_paths(&f.expr, scope, shadowed)?;
5575 }
5576 check_plan_json_paths(input, scope, shadowed)
5577 }
5578 PlanNode::GroupBy {
5579 input,
5580 keys,
5581 aggregates,
5582 having,
5583 } => {
5584 for key in keys {
5585 check_expr_json_paths(&key.expr, scope, shadowed)?;
5586 }
5587 for aggregate in aggregates {
5588 check_expr_json_paths(&aggregate.argument, scope, shadowed)?;
5589 }
5590 if let Some(h) = having {
5591 check_expr_json_paths(h, scope, shadowed)?;
5592 }
5593 check_plan_json_paths(input, scope, shadowed)
5594 }
5595 PlanNode::NestedLoopJoin {
5596 left, right, on, ..
5597 } => {
5598 if let Some(on) = on {
5599 check_expr_json_paths(on, scope, shadowed)?;
5600 }
5601 check_plan_json_paths(left, scope, shadowed)?;
5602 check_plan_json_paths(right, scope, shadowed)
5603 }
5604 PlanNode::Union { left, right, .. } => {
5605 check_plan_json_paths(left, scope, shadowed)?;
5606 check_plan_json_paths(right, scope, shadowed)
5607 }
5608 PlanNode::Sort { input, keys } => {
5609 for key in keys {
5610 check_expr_json_paths(&key.expr, scope, shadowed)?;
5611 }
5612 check_plan_json_paths(input, scope, shadowed)
5613 }
5614 PlanNode::Aggregate {
5615 input, argument, ..
5616 } => {
5617 if let Some(argument) = argument {
5618 check_expr_json_paths(argument, scope, shadowed)?;
5619 }
5620 check_plan_json_paths(input, scope, shadowed)
5621 }
5622 PlanNode::Window { input, windows } => {
5623 for window in windows {
5624 for expr in window.args.iter().chain(&window.partition_by) {
5625 check_expr_json_paths(expr, scope, shadowed)?;
5626 }
5627 for key in &window.order_by {
5628 check_expr_json_paths(&key.expr, scope, shadowed)?;
5629 }
5630 }
5631 check_plan_json_paths(input, scope, shadowed)
5632 }
5633 PlanNode::Limit { input, .. }
5634 | PlanNode::Offset { input, .. }
5635 | PlanNode::Distinct { input }
5636 | PlanNode::Update { input, .. }
5637 | PlanNode::Delete { input, .. }
5638 | PlanNode::Explain { input } => check_plan_json_paths(input, scope, shadowed),
5639 _ => Ok(()),
5640 }
5641}
5642
5643pub(super) fn validate_no_stray_aggregates(plan: &PlanNode) -> Result<(), QueryError> {
5644 match plan {
5645 PlanNode::Project { input, fields } => {
5646 for f in fields {
5647 check_expr_no_aggregate(&f.expr)?;
5648 }
5649 validate_no_stray_aggregates(input)?;
5650 }
5651 PlanNode::Filter { input, predicate } => {
5652 check_expr_no_aggregate(predicate)?;
5653 validate_no_stray_aggregates(input)?;
5654 }
5655 PlanNode::GroupBy {
5656 input,
5657 keys,
5658 aggregates,
5659 having,
5660 } => {
5661 for key in keys {
5662 check_expr_no_aggregate(&key.expr)?;
5663 }
5664 for aggregate in aggregates {
5665 check_expr_no_aggregate(&aggregate.argument)?;
5666 }
5667 if let Some(h) = having {
5668 check_expr_no_aggregate(h)?;
5669 }
5670 validate_no_stray_aggregates(input)?;
5671 }
5672 PlanNode::NestedLoopJoin {
5673 left, right, on, ..
5674 } => {
5675 if let Some(on) = on {
5676 check_expr_no_aggregate(on)?;
5677 }
5678 validate_no_stray_aggregates(left)?;
5679 validate_no_stray_aggregates(right)?;
5680 }
5681 PlanNode::Union { left, right, .. } => {
5682 validate_no_stray_aggregates(left)?;
5683 validate_no_stray_aggregates(right)?;
5684 }
5685 PlanNode::Sort { input, keys } => {
5686 for key in keys {
5687 check_expr_no_aggregate(&key.expr)?;
5688 }
5689 validate_no_stray_aggregates(input)?;
5690 }
5691 PlanNode::Aggregate {
5692 input, argument, ..
5693 } => {
5694 if let Some(argument) = argument {
5695 check_expr_no_aggregate(argument)?;
5696 }
5697 validate_no_stray_aggregates(input)?;
5698 }
5699 PlanNode::Window { input, windows } => {
5700 for window in windows {
5701 for expr in window.args.iter().chain(&window.partition_by) {
5702 check_expr_no_aggregate(expr)?;
5703 }
5704 for key in &window.order_by {
5705 check_expr_no_aggregate(&key.expr)?;
5706 }
5707 }
5708 validate_no_stray_aggregates(input)?;
5709 }
5710 PlanNode::Limit { input, .. }
5711 | PlanNode::Offset { input, .. }
5712 | PlanNode::Distinct { input }
5713 | PlanNode::Update { input, .. }
5714 | PlanNode::Delete { input, .. }
5715 | PlanNode::Explain { input } => {
5716 validate_no_stray_aggregates(input)?;
5717 }
5718 _ => {}
5719 }
5720 Ok(())
5721}
5722
5723fn check_expr_no_aggregate(expr: &Expr) -> Result<(), QueryError> {
5727 match expr {
5728 Expr::FunctionCall(..) => Err(QueryError::Execution(
5729 "invalid query: aggregate function in an unsupported position".to_string(),
5730 )),
5731 Expr::BinaryOp(l, _, r) | Expr::Coalesce(l, r) => {
5732 check_expr_no_aggregate(l)?;
5733 check_expr_no_aggregate(r)
5734 }
5735 Expr::UnaryOp(_, inner) | Expr::Cast(inner, _) | Expr::JsonPath { base: inner, .. } => {
5736 check_expr_no_aggregate(inner)
5737 }
5738 Expr::ScalarFunc(_, args) => {
5739 for a in args {
5740 check_expr_no_aggregate(a)?;
5741 }
5742 Ok(())
5743 }
5744 Expr::InList { expr: e, list, .. } => {
5745 check_expr_no_aggregate(e)?;
5746 for item in list {
5747 check_expr_no_aggregate(item)?;
5748 }
5749 Ok(())
5750 }
5751 Expr::InSubquery { expr: e, .. } => check_expr_no_aggregate(e),
5752 Expr::Case { whens, else_expr } => {
5753 for (c, r) in whens {
5754 check_expr_no_aggregate(c)?;
5755 check_expr_no_aggregate(r)?;
5756 }
5757 if let Some(e) = else_expr {
5758 check_expr_no_aggregate(e)?;
5759 }
5760 Ok(())
5761 }
5762 Expr::Window {
5763 args,
5764 partition_by,
5765 order_by,
5766 ..
5767 } => {
5768 for expr in args.iter().chain(partition_by) {
5769 check_expr_no_aggregate(expr)?;
5770 }
5771 for key in order_by {
5772 check_expr_no_aggregate(&key.expr)?;
5773 }
5774 Ok(())
5775 }
5776 _ => Ok(()),
5777 }
5778}
5779
5780pub(super) fn aggregate_rows(
5784 func: AggFunc,
5785 argument: Option<&Expr>,
5786 columns: &[String],
5787 rows: &[Vec<Value>],
5788) -> Result<QueryResult, QueryError> {
5789 let mut cancel = CancelCheck::new();
5790 if func == AggFunc::Count && argument.is_none() {
5791 return Ok(QueryResult::Scalar(Value::Int(rows.len() as i64)));
5792 }
5793 let argument = argument.ok_or_else(|| {
5794 QueryError::Execution(format!(
5795 "{} requires an argument",
5796 format!("{func:?}").to_lowercase()
5797 ))
5798 })?;
5799
5800 let mut values = Vec::with_capacity(rows.len());
5801 for row in rows {
5802 cancel.tick()?;
5803 values.push(eval_expr(argument, row, columns));
5804 }
5805
5806 let value = match func {
5807 AggFunc::Count => Value::Int(values.iter().filter(|v| !v.is_empty()).count() as i64),
5808 AggFunc::CountDistinct => {
5809 let seen: std::collections::HashSet<Value> =
5810 values.into_iter().filter(|v| !v.is_empty()).collect();
5811 Value::Int(seen.len() as i64)
5812 }
5813 AggFunc::Avg => {
5814 let mut sum = 0.0;
5815 let mut count = 0_u64;
5816 for value in values {
5817 match value {
5818 Value::Int(v) => {
5819 sum += v as f64;
5820 count += 1;
5821 }
5822 Value::Float(v) => {
5823 sum += v;
5824 count += 1;
5825 }
5826 _ => {}
5827 }
5828 }
5829 if count == 0 {
5830 Value::Empty
5831 } else {
5832 Value::Float(sum / count as f64)
5833 }
5834 }
5835 AggFunc::Sum => {
5836 let mut int_sum = 0_i64;
5837 let mut float_sum = 0.0;
5838 let mut saw_float = false;
5839 for value in values {
5840 match value {
5841 Value::Int(v) => int_sum += v,
5842 Value::Float(v) => {
5843 float_sum += v;
5844 saw_float = true;
5845 }
5846 _ => {}
5847 }
5848 }
5849 if saw_float {
5850 Value::Float(float_sum + int_sum as f64)
5851 } else {
5852 Value::Int(int_sum)
5853 }
5854 }
5855 AggFunc::Min | AggFunc::Max => {
5856 let mut result: Option<Value> = None;
5857 for value in values.into_iter().filter(|v| !v.is_empty()) {
5858 let replace = match &result {
5859 None => true,
5860 Some(current) if func == AggFunc::Min => value < *current,
5861 Some(current) => value > *current,
5862 };
5863 if replace {
5864 result = Some(value);
5865 }
5866 }
5867 result.unwrap_or(Value::Empty)
5868 }
5869 };
5870 Ok(QueryResult::Scalar(value))
5871}
5872
5873const SYMMETRIC_RID_SET_ENTRY_BYTES: usize =
5874 std::mem::size_of::<RowId>() + 2 * std::mem::size_of::<usize>();
5875
5876pub(super) fn aggregate_rows_with_provenance(
5877 func: AggFunc,
5878 argument: Option<&Expr>,
5879 input: &ProvenanceRows,
5880 provenance_alias: &str,
5881 memory_limit: usize,
5882) -> Result<QueryResult, QueryError> {
5883 if matches!(func, AggFunc::Min | AggFunc::Max | AggFunc::CountDistinct) {
5884 return aggregate_rows(func, argument, &input.columns, &input.rows);
5885 }
5886 let argument = argument.ok_or_else(|| {
5887 QueryError::Execution(
5888 "symmetric aggregate requires a source-valued argument; use raw".to_string(),
5889 )
5890 })?;
5891 let source_index = input.source_index(provenance_alias).ok_or_else(|| {
5892 QueryError::Execution(format!(
5893 "symmetric aggregate source alias '{provenance_alias}' is not present in its input"
5894 ))
5895 })?;
5896 let mut seen = HashSet::new();
5897 let mut int_sum = 0_i64;
5898 let mut float_sum = 0.0_f64;
5899 let mut saw_float = false;
5900 let mut count = 0_u64;
5901 let mut cancel = CancelCheck::new();
5902 for (row, row_provenance) in input.rows.iter().zip(&input.provenance) {
5903 cancel.tick()?;
5904 let value = eval_expr(argument, row, &input.columns);
5905 if value.is_empty() {
5906 continue;
5907 }
5908 let Some(rid) = row_provenance[source_index] else {
5909 continue;
5910 };
5911 if !seen.insert(rid) {
5912 continue;
5913 }
5914 mem_budget::charge(SYMMETRIC_RID_SET_ENTRY_BYTES, memory_limit)?;
5915 match func {
5916 AggFunc::Count => count += 1,
5917 AggFunc::Sum | AggFunc::Avg => match value {
5918 Value::Int(value) => {
5919 int_sum += value;
5920 count += 1;
5921 }
5922 Value::Float(value) => {
5923 float_sum += value;
5924 saw_float = true;
5925 count += 1;
5926 }
5927 _ => {}
5928 },
5929 AggFunc::CountDistinct | AggFunc::Min | AggFunc::Max => unreachable!(),
5930 }
5931 }
5932 let value = match func {
5933 AggFunc::Count => Value::Int(count as i64),
5934 AggFunc::Sum if saw_float => Value::Float(float_sum + int_sum as f64),
5935 AggFunc::Sum => Value::Int(int_sum),
5936 AggFunc::Avg if count == 0 => Value::Empty,
5937 AggFunc::Avg => Value::Float((float_sum + int_sum as f64) / count as f64),
5938 AggFunc::CountDistinct | AggFunc::Min | AggFunc::Max => unreachable!(),
5939 };
5940 Ok(QueryResult::Scalar(value))
5941}
5942
5943pub(super) struct GroupAggregateContext<'a> {
5945 pub(super) columns: &'a [String],
5946 pub(super) all_rows: &'a [Vec<Value>],
5947 pub(super) row_indices: &'a [usize],
5948 pub(super) source_index: Option<usize>,
5949 pub(super) provenance: Option<(&'a [Vec<Option<RowId>>], usize)>,
5950}
5951
5952pub(super) fn compute_group_aggregate(
5953 func: AggFunc,
5954 argument: &Expr,
5955 direct_index: Option<usize>,
5956 context: GroupAggregateContext<'_>,
5957) -> Result<Value, QueryError> {
5958 let GroupAggregateContext {
5959 columns,
5960 all_rows,
5961 row_indices,
5962 source_index,
5963 provenance,
5964 } = context;
5965 let count_all = matches!(argument, Expr::Field(name) if name == "*");
5966 let value_at = |ri: usize| match direct_index {
5967 Some(index) => all_rows[ri][index].clone(),
5968 None => eval_expr(argument, &all_rows[ri], columns),
5969 };
5970 let mut cancel = CancelCheck::new();
5971 let mut seen_rids = HashSet::new();
5972 match func {
5973 AggFunc::Count => {
5974 if count_all {
5975 return Ok(Value::Int(row_indices.len() as i64));
5977 }
5978 let mut count = 0usize;
5979 for &ri in row_indices {
5980 cancel.tick()?;
5981 let value = value_at(ri);
5982 if !value.is_empty()
5983 && accept_symmetric_contribution(ri, source_index, provenance, &mut seen_rids)?
5984 {
5985 count += 1;
5986 }
5987 }
5988 Ok(Value::Int(count as i64))
5989 }
5990 AggFunc::CountDistinct => {
5991 let mut seen = std::collections::HashSet::new();
5992 for &ri in row_indices {
5993 cancel.tick()?;
5994 let v = value_at(ri);
5995 if !v.is_empty() {
5996 seen.insert(v);
5997 }
5998 }
5999 Ok(Value::Int(seen.len() as i64))
6000 }
6001 AggFunc::Sum => {
6002 let mut int_sum: i64 = 0;
6007 let mut float_sum: f64 = 0.0;
6008 let mut saw_float = false;
6009 for &ri in row_indices {
6010 cancel.tick()?;
6011 let value = value_at(ri);
6012 if value.is_empty()
6013 || !accept_symmetric_contribution(ri, source_index, provenance, &mut seen_rids)?
6014 {
6015 continue;
6016 }
6017 match value {
6018 Value::Int(v) => int_sum += v,
6019 Value::Float(v) => {
6020 float_sum += v;
6021 saw_float = true;
6022 }
6023 _ => {}
6024 }
6025 }
6026 if saw_float {
6027 Ok(Value::Float(float_sum + int_sum as f64))
6028 } else {
6029 Ok(Value::Int(int_sum))
6030 }
6031 }
6032 AggFunc::Avg => {
6033 let mut sum = 0.0f64;
6034 let mut count = 0usize;
6035 for &ri in row_indices {
6036 cancel.tick()?;
6037 let value = value_at(ri);
6038 if value.is_empty()
6039 || !accept_symmetric_contribution(ri, source_index, provenance, &mut seen_rids)?
6040 {
6041 continue;
6042 }
6043 match value {
6044 Value::Int(v) => {
6045 sum += v as f64;
6046 count += 1;
6047 }
6048 Value::Float(v) => {
6049 sum += v;
6050 count += 1;
6051 }
6052 _ => {}
6053 }
6054 }
6055 if count == 0 {
6056 Ok(Value::Empty)
6057 } else {
6058 Ok(Value::Float(sum / count as f64))
6059 }
6060 }
6061 AggFunc::Min | AggFunc::Max => {
6062 let mut result: Option<Value> = None;
6063 for &ri in row_indices {
6064 cancel.tick()?;
6065 let value = value_at(ri);
6066 if value.is_empty() {
6067 continue;
6068 }
6069 let replace = match &result {
6070 None => true,
6071 Some(current) if func == AggFunc::Min => value < *current,
6072 Some(current) => value > *current,
6073 };
6074 if replace {
6075 result = Some(value);
6076 }
6077 }
6078 Ok(result.unwrap_or(Value::Empty))
6079 }
6080 }
6081}
6082
6083fn accept_symmetric_contribution(
6084 row_index: usize,
6085 source_index: Option<usize>,
6086 provenance: Option<(&[Vec<Option<RowId>>], usize)>,
6087 seen: &mut HashSet<RowId>,
6088) -> Result<bool, QueryError> {
6089 let Some(source_index) = source_index else {
6090 return Ok(true);
6091 };
6092 let Some((provenance, memory_limit)) = provenance else {
6093 return Err(QueryError::Execution(
6094 "symmetric aggregate provenance is unavailable; use raw".to_string(),
6095 ));
6096 };
6097 let Some(rid) = provenance[row_index][source_index] else {
6098 return Ok(false);
6099 };
6100 if !seen.insert(rid) {
6101 return Ok(false);
6102 }
6103 mem_budget::charge(SYMMETRIC_RID_SET_ENTRY_BYTES, memory_limit)?;
6104 Ok(true)
6105}
6106
6107struct HashJoinSpec<'a> {
6108 left_key_idx: usize,
6109 right_key_idx: usize,
6110 residuals: Vec<&'a Expr>,
6111}
6112
6113struct MaterializedJoinInputs {
6114 left_columns: Vec<String>,
6115 left_rows: Vec<Vec<Value>>,
6116 right_columns: Vec<String>,
6117 right_rows: Vec<Vec<Value>>,
6118}
6119
6120fn flatten_conjunctions<'a>(expr: &'a Expr, out: &mut Vec<&'a Expr>) {
6121 match expr {
6122 Expr::BinaryOp(left, BinOp::And, right) => {
6123 flatten_conjunctions(left, out);
6124 flatten_conjunctions(right, out);
6125 }
6126 _ => out.push(expr),
6127 }
6128}
6129
6130fn try_extract_hash_join<'a>(
6134 pred: &'a Expr,
6135 left_columns: &[String],
6136 right_columns: &[String],
6137) -> Option<HashJoinSpec<'a>> {
6138 let mut conjuncts = Vec::new();
6139 flatten_conjunctions(pred, &mut conjuncts);
6140 for (key_position, conjunct) in conjuncts.iter().enumerate() {
6141 let Some((left_key_idx, right_key_idx)) =
6142 try_extract_equi_join_keys(conjunct, left_columns, right_columns)
6143 else {
6144 continue;
6145 };
6146 let residuals = conjuncts
6147 .iter()
6148 .enumerate()
6149 .filter_map(|(position, residual)| (position != key_position).then_some(*residual))
6150 .collect();
6151 return Some(HashJoinSpec {
6152 left_key_idx,
6153 right_key_idx,
6154 residuals,
6155 });
6156 }
6157 None
6158}
6159
6160pub(super) fn try_extract_equi_join_keys(
6162 pred: &Expr,
6163 left_columns: &[String],
6164 right_columns: &[String],
6165) -> Option<(usize, usize)> {
6166 let (lhs, op, rhs) = match pred {
6167 Expr::BinaryOp(l, op, r) => (l.as_ref(), *op, r.as_ref()),
6168 _ => return None,
6169 };
6170 if op != BinOp::Eq {
6171 return None;
6172 }
6173 if let (Some(li), Some(ri)) = (
6175 resolve_side_column(lhs, left_columns),
6176 resolve_side_column(rhs, right_columns),
6177 ) {
6178 return Some((li, ri));
6179 }
6180 if let (Some(li), Some(ri)) = (
6183 resolve_side_column(rhs, left_columns),
6184 resolve_side_column(lhs, right_columns),
6185 ) {
6186 return Some((li, ri));
6187 }
6188 None
6189}
6190
6191fn resolve_side_column(expr: &Expr, columns: &[String]) -> Option<usize> {
6192 match expr {
6193 Expr::QualifiedField { qualifier, field } => {
6194 let q = qualifier.as_bytes();
6199 let f = field.as_bytes();
6200 columns.iter().position(|c| {
6201 let b = c.as_bytes();
6202 b.len() == q.len() + 1 + f.len()
6203 && b[..q.len()] == *q
6204 && b[q.len()] == b'.'
6205 && b[q.len() + 1..] == *f
6206 })
6207 }
6208 Expr::Field(name) => columns.iter().position(|c| c == name),
6209 _ => None,
6210 }
6211}
6212
6213fn hash_join(
6224 inputs: MaterializedJoinInputs,
6225 left_key_idx: usize,
6226 right_key_idx: usize,
6227 kind: JoinKind,
6228 residuals: &[&Expr],
6229) -> Result<QueryResult, QueryError> {
6230 use rustc_hash::FxHashMap;
6231
6232 let MaterializedJoinInputs {
6233 left_columns,
6234 left_rows,
6235 right_columns,
6236 right_rows,
6237 } = inputs;
6238
6239 let n_left = left_columns.len();
6240 let n_right = right_columns.len();
6241 let mut columns = Vec::with_capacity(n_left + n_right);
6242 columns.extend(left_columns);
6243 columns.extend(right_columns);
6244
6245 let mut cancel = CancelCheck::new();
6248
6249 let mut build: FxHashMap<Value, Vec<usize>> =
6252 FxHashMap::with_capacity_and_hasher(right_rows.len(), Default::default());
6253 for (i, row) in right_rows.iter().enumerate() {
6254 cancel.tick()?;
6255 build.entry(row[right_key_idx].clone()).or_default().push(i);
6259 }
6260
6261 let mut rows: Vec<Vec<Value>> = Vec::with_capacity(left_rows.len());
6264
6265 crate::cancel::check()?;
6266 for left_row in &left_rows {
6267 cancel.tick()?;
6268 let key = &left_row[left_key_idx];
6269 let candidates = build.get(key);
6270 let mut matched = false;
6271 match candidates {
6272 Some(matches) if !matches.is_empty() => {
6273 for &ri in matches {
6274 cancel.tick()?;
6275 let right_row = &right_rows[ri];
6276 let mut combined = Vec::with_capacity(n_left + n_right);
6277 combined.extend_from_slice(left_row);
6278 combined.extend_from_slice(right_row);
6279 if residuals
6280 .iter()
6281 .all(|residual| eval_predicate(residual, &combined, &columns))
6282 {
6283 rows.push(combined);
6284 check_join_limit(rows.len())?;
6285 matched = true;
6286 }
6287 }
6288 }
6289 _ => {}
6290 }
6291 if !matched && matches!(kind, JoinKind::LeftOuter) {
6292 let mut row = Vec::with_capacity(n_left + n_right);
6293 row.extend_from_slice(left_row);
6294 row.resize(n_left + n_right, Value::Empty);
6295 rows.push(row);
6296 check_join_limit(rows.len())?;
6297 }
6298 }
6299
6300 Ok(QueryResult::Rows { columns, rows })
6301}
6302
6303#[inline]
6304pub(super) fn check_nested_loop_pair_limit(
6305 left_rows: usize,
6306 right_rows: usize,
6307 pair_limit: usize,
6308) -> Result<usize, QueryError> {
6309 let candidate_pairs =
6310 left_rows
6311 .checked_mul(right_rows)
6312 .ok_or(QueryError::NestedLoopPairLimitExceeded {
6313 left_rows,
6314 right_rows,
6315 limit: pair_limit,
6316 })?;
6317 if candidate_pairs > pair_limit {
6318 return Err(QueryError::NestedLoopPairLimitExceeded {
6319 left_rows,
6320 right_rows,
6321 limit: pair_limit,
6322 });
6323 }
6324 Ok(candidate_pairs)
6325}
6326
6327pub(super) fn execute_materialized_join(
6331 left_columns: Vec<String>,
6332 left_rows: Vec<Vec<Value>>,
6333 right_columns: Vec<String>,
6334 right_rows: Vec<Vec<Value>>,
6335 on: Option<&Expr>,
6336 kind: JoinKind,
6337 pair_limit: usize,
6338) -> Result<QueryResult, QueryError> {
6339 crate::cancel::check()?;
6340 if !matches!(kind, JoinKind::Cross) {
6341 if let Some(pred) = on {
6342 if let Some(spec) = try_extract_hash_join(pred, &left_columns, &right_columns) {
6343 return hash_join(
6344 MaterializedJoinInputs {
6345 left_columns,
6346 left_rows,
6347 right_columns,
6348 right_rows,
6349 },
6350 spec.left_key_idx,
6351 spec.right_key_idx,
6352 kind,
6353 &spec.residuals,
6354 );
6355 }
6356 }
6357 }
6358
6359 check_nested_loop_pair_limit(left_rows.len(), right_rows.len(), pair_limit)?;
6360 let n_left = left_columns.len();
6361 let n_right = right_columns.len();
6362 let mut columns = Vec::with_capacity(n_left + n_right);
6363 columns.extend(left_columns);
6364 columns.extend(right_columns);
6365
6366 let mut rows: Vec<Vec<Value>> = Vec::with_capacity(left_rows.len());
6367 let mut combined: Vec<Value> = Vec::with_capacity(n_left + n_right);
6368 let mut cancel = CancelCheck::new();
6369 for left_row in &left_rows {
6370 let mut matched = false;
6371 for right_row in &right_rows {
6372 cancel.tick()?;
6373 combined.clear();
6374 combined.extend_from_slice(left_row);
6375 combined.extend_from_slice(right_row);
6376 let keep = match kind {
6377 JoinKind::Cross => true,
6378 JoinKind::Inner | JoinKind::LeftOuter => {
6379 on.is_none_or(|pred| eval_predicate(pred, &combined, &columns))
6380 }
6381 JoinKind::RightOuter => {
6382 unreachable!("planner rewrites RightOuter to LeftOuter")
6383 }
6384 };
6385 if keep {
6386 rows.push(combined.clone());
6387 check_join_limit(rows.len())?;
6388 matched = true;
6389 }
6390 }
6391 if !matched && matches!(kind, JoinKind::LeftOuter) {
6392 let mut row = Vec::with_capacity(n_left + n_right);
6393 row.extend_from_slice(left_row);
6394 row.resize(n_left + n_right, Value::Empty);
6395 rows.push(row);
6396 check_join_limit(rows.len())?;
6397 }
6398 }
6399 Ok(QueryResult::Rows { columns, rows })
6400}
6401
6402fn execute_provenance_join(
6403 left: ProvenanceRows,
6404 right: ProvenanceRows,
6405 on: Option<&Expr>,
6406 kind: JoinKind,
6407 pair_limit: usize,
6408) -> Result<ProvenanceRows, QueryError> {
6409 let left_width = left.columns.len();
6410 let right_width = right.columns.len();
6411 let right_source_count = right.source_aliases.len();
6412 let mut columns = left.columns.clone();
6413 columns.extend(right.columns.clone());
6414 let mut source_aliases = left.source_aliases.clone();
6415 source_aliases.extend(right.source_aliases.clone());
6416 let mut rows = Vec::new();
6417 let mut provenance = Vec::new();
6418 let mut cancel = CancelCheck::new();
6419
6420 if !matches!(kind, JoinKind::Cross) {
6421 if let Some(predicate) = on {
6422 if let Some(spec) = try_extract_hash_join(predicate, &left.columns, &right.columns) {
6423 let mut build: rustc_hash::FxHashMap<Value, Vec<usize>> =
6424 rustc_hash::FxHashMap::default();
6425 for (index, row) in right.rows.iter().enumerate() {
6426 cancel.tick()?;
6427 let key = &row[spec.right_key_idx];
6428 build.entry(key.clone()).or_default().push(index);
6429 }
6430 for (left_index, left_row) in left.rows.iter().enumerate() {
6431 cancel.tick()?;
6432 let key = &left_row[spec.left_key_idx];
6433 let candidates = build.get(key);
6434 let mut matched = false;
6435 if let Some(candidates) = candidates {
6436 for &right_index in candidates {
6437 cancel.tick()?;
6438 let mut row = Vec::with_capacity(left_width + right_width);
6439 row.extend_from_slice(left_row);
6440 row.extend_from_slice(&right.rows[right_index]);
6441 if spec
6442 .residuals
6443 .iter()
6444 .all(|residual| eval_predicate(residual, &row, &columns))
6445 {
6446 let mut row_provenance = left.provenance[left_index].clone();
6447 row_provenance.extend_from_slice(&right.provenance[right_index]);
6448 rows.push(row);
6449 provenance.push(row_provenance);
6450 check_join_limit(rows.len())?;
6451 matched = true;
6452 }
6453 }
6454 }
6455 if !matched && matches!(kind, JoinKind::LeftOuter) {
6456 let mut row = left_row.clone();
6457 row.resize(left_width + right_width, Value::Empty);
6458 let mut row_provenance = left.provenance[left_index].clone();
6459 row_provenance.extend(std::iter::repeat_n(None, right_source_count));
6460 rows.push(row);
6461 provenance.push(row_provenance);
6462 check_join_limit(rows.len())?;
6463 }
6464 }
6465 return Ok(ProvenanceRows {
6466 columns,
6467 rows,
6468 source_aliases,
6469 provenance,
6470 });
6471 }
6472 }
6473 }
6474
6475 check_nested_loop_pair_limit(left.rows.len(), right.rows.len(), pair_limit)?;
6476 for (left_index, left_row) in left.rows.iter().enumerate() {
6477 let mut matched = false;
6478 for (right_index, right_row) in right.rows.iter().enumerate() {
6479 cancel.tick()?;
6480 let mut row = Vec::with_capacity(left_width + right_width);
6481 row.extend_from_slice(left_row);
6482 row.extend_from_slice(right_row);
6483 let keep = match kind {
6484 JoinKind::Cross => true,
6485 JoinKind::Inner | JoinKind::LeftOuter => {
6486 on.is_none_or(|predicate| eval_predicate(predicate, &row, &columns))
6487 }
6488 JoinKind::RightOuter => {
6489 unreachable!("planner rewrites RightOuter to LeftOuter")
6490 }
6491 };
6492 if keep {
6493 let mut row_provenance = left.provenance[left_index].clone();
6494 row_provenance.extend_from_slice(&right.provenance[right_index]);
6495 rows.push(row);
6496 provenance.push(row_provenance);
6497 check_join_limit(rows.len())?;
6498 matched = true;
6499 }
6500 }
6501 if !matched && matches!(kind, JoinKind::LeftOuter) {
6502 let mut row = left_row.clone();
6503 row.resize(left_width + right_width, Value::Empty);
6504 let mut row_provenance = left.provenance[left_index].clone();
6505 row_provenance.extend(std::iter::repeat_n(None, right_source_count));
6506 rows.push(row);
6507 provenance.push(row_provenance);
6508 check_join_limit(rows.len())?;
6509 }
6510 }
6511 Ok(ProvenanceRows {
6512 columns,
6513 rows,
6514 source_aliases,
6515 provenance,
6516 })
6517}
6518
6519fn flatten_and<'a>(expr: &'a Expr, out: &mut Vec<&'a Expr>) {
6534 match expr {
6535 Expr::BinaryOp(lhs, BinOp::And, rhs) => {
6536 flatten_and(lhs, out);
6537 flatten_and(rhs, out);
6538 }
6539 other => out.push(other),
6540 }
6541}
6542
6543fn eq_candidate_tier(catalog: &Catalog, scan: &PlanNode) -> Option<u8> {
6547 match scan {
6548 PlanNode::IndexScan { table, column, .. } => match catalog.is_index_unique(table, column) {
6549 Some(true) => Some(0),
6550 Some(false) => Some(1),
6551 None => None,
6552 },
6553 PlanNode::ExprIndexScan { table, path, .. } => {
6554 resolve_expression_index(catalog, table, path).map(|meta| u8::from(!meta.unique))
6555 }
6556 _ => None,
6557 }
6558}
6559
6560fn range_candidate_resolves(catalog: &Catalog, scan: &PlanNode) -> bool {
6562 match scan {
6563 PlanNode::RangeScan { table, column, .. } => catalog.has_index(table, column),
6564 PlanNode::ExprRangeScan { table, path, .. } => {
6565 resolve_expression_index(catalog, table, path).is_some()
6566 }
6567 _ => false,
6568 }
6569}
6570
6571const UNKNOWN_EST: u64 = u64::MAX;
6575
6576fn probes_empty_sentinel(key: &Expr) -> bool {
6579 matches!(literal_to_value(key), Ok(Value::Empty))
6580}
6581
6582fn eq_est_rows(stats: &IndexStats, unique: bool, empty_probe: bool) -> u64 {
6589 if unique {
6590 1
6591 } else if empty_probe {
6592 stats.empty_count
6593 } else {
6594 stats.total_entries / stats.distinct_keys.max(1)
6595 }
6596}
6597
6598fn eq_candidate_est(catalog: &Catalog, scan: &PlanNode, tier: u8) -> u64 {
6602 let (stats, key) = match scan {
6603 PlanNode::IndexScan { table, column, key } => (catalog.index_stats(table, column), key),
6604 PlanNode::ExprIndexScan { table, path, key } => (
6605 resolve_expression_index(catalog, table, path)
6606 .and_then(|meta| catalog.expression_index_stats(table, meta.index_id)),
6607 key,
6608 ),
6609 _ => return UNKNOWN_EST,
6610 };
6611 stats.map_or(UNKNOWN_EST, |stats| {
6612 eq_est_rows(&stats, tier == 0, probes_empty_sentinel(key))
6613 })
6614}
6615
6616fn range_candidate_est(catalog: &Catalog, scan: &PlanNode) -> u64 {
6621 let stats = match scan {
6622 PlanNode::RangeScan { table, column, .. } => catalog.index_stats(table, column),
6623 PlanNode::ExprRangeScan { table, path, .. } => {
6624 resolve_expression_index(catalog, table, path)
6625 .and_then(|meta| catalog.expression_index_stats(table, meta.index_id))
6626 }
6627 _ => None,
6628 };
6629 stats.map_or(UNKNOWN_EST, |stats| stats.total_entries)
6630}
6631
6632fn column_type(catalog: &Catalog, table: &str, column: &str) -> Option<TypeId> {
6634 catalog
6635 .schema(table)?
6636 .find_column(column)
6637 .map(|col| col.type_id)
6638}
6639
6640fn coerce_column_index_key(col_type: TypeId, key: &Expr) -> Option<Expr> {
6654 match (key, col_type) {
6655 (Expr::Literal(Literal::Int(_)), TypeId::Int | TypeId::DateTime) => Some(key.clone()),
6659 (Expr::Literal(Literal::Float(_)), TypeId::Float) => Some(key.clone()),
6660 (Expr::Literal(Literal::String(_)), TypeId::Str) => Some(key.clone()),
6661 (Expr::Literal(Literal::Bool(_)), TypeId::Bool) => Some(key.clone()),
6662 (Expr::Literal(Literal::Int(v)), TypeId::Float) => {
6665 Some(Expr::Literal(Literal::Float(*v as f64)))
6666 }
6667 _ => None,
6670 }
6671}
6672
6673fn coerce_column_index_bound(
6677 col_type: TypeId,
6678 bound: Option<(Expr, bool)>,
6679) -> Option<Option<(Expr, bool)>> {
6680 match bound {
6681 None => Some(None),
6682 Some((expr, inclusive)) => {
6683 coerce_column_index_key(col_type, &expr).map(|expr| Some((expr, inclusive)))
6684 }
6685 }
6686}
6687
6688fn coerce_candidate_keys(catalog: &Catalog, scan: PlanNode) -> Option<PlanNode> {
6695 match scan {
6696 PlanNode::IndexScan { table, column, key } => {
6697 let col_type = column_type(catalog, &table, &column)?;
6698 let key = coerce_column_index_key(col_type, &key)?;
6699 Some(PlanNode::IndexScan { table, column, key })
6700 }
6701 PlanNode::RangeScan {
6702 table,
6703 column,
6704 start,
6705 end,
6706 } => {
6707 let col_type = column_type(catalog, &table, &column)?;
6708 let start = coerce_column_index_bound(col_type, start)?;
6709 let end = coerce_column_index_bound(col_type, end)?;
6710 Some(PlanNode::RangeScan {
6711 table,
6712 column,
6713 start,
6714 end,
6715 })
6716 }
6717 other => Some(other),
6718 }
6719}
6720
6721struct ConjunctionCandidate {
6724 plan: PlanNode,
6725 consumed: Vec<usize>,
6726 est: u64,
6728 tier: u8,
6729}
6730
6731fn lower_conjunction_scan(catalog: &Catalog, table: &str, predicate: &Expr) -> Option<PlanNode> {
6745 let mut conjuncts: Vec<&Expr> = Vec::new();
6746 flatten_and(predicate, &mut conjuncts);
6747 if conjuncts.len() < 2 {
6748 return None;
6749 }
6750
6751 let mut candidates: Vec<ConjunctionCandidate> = Vec::new();
6752
6753 for (i, conjunct) in conjuncts.iter().enumerate() {
6755 if let Some(scan) = try_extract_eq_index_key(table, conjunct) {
6756 if let Some(scan) = coerce_candidate_keys(catalog, scan) {
6760 if let Some(tier) = eq_candidate_tier(catalog, &scan) {
6761 let est = eq_candidate_est(catalog, &scan, tier);
6762 candidates.push(ConjunctionCandidate {
6763 plan: scan,
6764 consumed: vec![i],
6765 est,
6766 tier,
6767 });
6768 }
6769 }
6770 }
6771 }
6772
6773 let bounds: Vec<(usize, RangeBound)> = conjuncts
6778 .iter()
6779 .enumerate()
6780 .filter_map(|(i, conjunct)| extract_single_bound(conjunct).map(|bound| (i, bound)))
6781 .collect();
6782 let mut seen_targets: Vec<RangeTarget> = Vec::new();
6783 for (_, (target, _, _)) in &bounds {
6784 if !seen_targets.contains(target) {
6785 seen_targets.push(target.clone());
6786 }
6787 }
6788 for target in seen_targets {
6789 let mut lower: Option<(Expr, bool)> = None;
6790 let mut lower_idx: Option<usize> = None;
6791 let mut upper: Option<(Expr, bool)> = None;
6792 let mut upper_idx: Option<usize> = None;
6793 for (i, (candidate_target, start, end)) in &bounds {
6794 if *candidate_target != target {
6795 continue;
6796 }
6797 if lower.is_none() {
6798 if let Some(bound) = start.clone() {
6799 lower = Some(bound);
6800 lower_idx = Some(*i);
6801 }
6802 }
6803 if upper.is_none() {
6804 if let Some(bound) = end.clone() {
6805 upper = Some(bound);
6806 upper_idx = Some(*i);
6807 }
6808 }
6809 }
6810 if lower.is_none() && upper.is_none() {
6811 continue;
6812 }
6813 let scan = range_scan_for_target(table, target, lower, upper);
6814 let Some(scan) = coerce_candidate_keys(catalog, scan) else {
6818 continue;
6819 };
6820 if !range_candidate_resolves(catalog, &scan) {
6821 continue;
6822 }
6823 let mut consumed: Vec<usize> = Vec::new();
6824 if let Some(i) = lower_idx {
6825 consumed.push(i);
6826 }
6827 if let Some(i) = upper_idx {
6828 if !consumed.contains(&i) {
6829 consumed.push(i);
6830 }
6831 }
6832 let est = range_candidate_est(catalog, &scan);
6833 candidates.push(ConjunctionCandidate {
6834 plan: scan,
6835 consumed,
6836 est,
6837 tier: 2,
6838 });
6839 }
6840
6841 let winner = candidates
6845 .into_iter()
6846 .enumerate()
6847 .min_by_key(|(build_order, candidate)| (candidate.est, candidate.tier, *build_order))?
6848 .1;
6849
6850 let mut residual: Vec<Expr> = Vec::new();
6851 for (i, conjunct) in conjuncts.iter().enumerate() {
6852 if !winner.consumed.contains(&i) {
6853 residual.push((*conjunct).clone());
6854 }
6855 }
6856 if residual.is_empty() {
6857 return Some(winner.plan);
6858 }
6859 let residual_expr = residual
6860 .into_iter()
6861 .reduce(|acc, next| Expr::BinaryOp(Box::new(acc), BinOp::And, Box::new(next)))
6862 .expect("residual is non-empty");
6863 Some(PlanNode::Filter {
6864 input: Box::new(winner.plan),
6865 predicate: residual_expr,
6866 })
6867}
6868
6869pub(super) fn lower_unindexed_scans(catalog: &Catalog, plan: &PlanNode) -> PlanNode {
6871 match plan {
6872 PlanNode::ExprIndexScan { table, path, .. }
6873 | PlanNode::ExprRangeScan { table, path, .. }
6874 | PlanNode::OrderedExprIndexScan { table, path, .. } => {
6875 if resolve_expression_index(catalog, table, path).is_some() {
6876 plan.clone()
6877 } else {
6878 expression_index_fallback(plan)
6879 .expect("expression-index branch always has a fallback")
6880 }
6881 }
6882 PlanNode::RangeScan {
6883 table,
6884 column,
6885 start,
6886 end,
6887 } => {
6888 if let Some(tbl) = catalog.get_table(table) {
6889 if tbl.has_index(column) {
6895 return plan.clone();
6896 }
6897 }
6898 let pred = synthesize_range_predicate(column, start, end);
6899 PlanNode::Filter {
6900 input: Box::new(PlanNode::SeqScan {
6901 table: table.clone(),
6902 }),
6903 predicate: pred,
6904 }
6905 }
6906 PlanNode::Filter { input, predicate } => {
6907 if let PlanNode::SeqScan { table } = input.as_ref() {
6912 if let Some(lowered) = lower_conjunction_scan(catalog, table, predicate) {
6913 return lowered;
6914 }
6915 }
6916 PlanNode::Filter {
6917 input: Box::new(lower_unindexed_scans(catalog, input)),
6918 predicate: predicate.clone(),
6919 }
6920 }
6921 PlanNode::Project { input, fields } => PlanNode::Project {
6922 input: Box::new(lower_unindexed_scans(catalog, input)),
6923 fields: fields.clone(),
6924 },
6925 PlanNode::Sort { input, keys } => PlanNode::Sort {
6926 input: Box::new(lower_unindexed_scans(catalog, input)),
6927 keys: keys.clone(),
6928 },
6929 PlanNode::Limit { input, count } => PlanNode::Limit {
6930 input: Box::new(lower_unindexed_scans(catalog, input)),
6931 count: count.clone(),
6932 },
6933 PlanNode::Offset { input, count } => PlanNode::Offset {
6934 input: Box::new(lower_unindexed_scans(catalog, input)),
6935 count: count.clone(),
6936 },
6937 PlanNode::Aggregate {
6938 input,
6939 function,
6940 argument,
6941 mode,
6942 provenance_alias,
6943 } => PlanNode::Aggregate {
6944 input: Box::new(lower_unindexed_scans(catalog, input)),
6945 function: *function,
6946 argument: argument.clone(),
6947 mode: *mode,
6948 provenance_alias: provenance_alias.clone(),
6949 },
6950 PlanNode::Distinct { input } => PlanNode::Distinct {
6951 input: Box::new(lower_unindexed_scans(catalog, input)),
6952 },
6953 PlanNode::GroupBy {
6954 input,
6955 keys,
6956 aggregates,
6957 having,
6958 } => PlanNode::GroupBy {
6959 input: Box::new(lower_unindexed_scans(catalog, input)),
6960 keys: keys.clone(),
6961 aggregates: aggregates.clone(),
6962 having: having.clone(),
6963 },
6964 PlanNode::Update {
6965 input,
6966 table,
6967 assignments,
6968 returning,
6969 } => PlanNode::Update {
6970 input: Box::new(lower_unindexed_scans(catalog, input)),
6971 table: table.clone(),
6972 assignments: assignments.clone(),
6973 returning: *returning,
6974 },
6975 PlanNode::Delete {
6976 input,
6977 table,
6978 returning,
6979 } => PlanNode::Delete {
6980 input: Box::new(lower_unindexed_scans(catalog, input)),
6981 table: table.clone(),
6982 returning: *returning,
6983 },
6984 PlanNode::Window { input, windows } => PlanNode::Window {
6985 input: Box::new(lower_unindexed_scans(catalog, input)),
6986 windows: windows.clone(),
6987 },
6988 PlanNode::Union { left, right, all } => PlanNode::Union {
6989 left: Box::new(lower_unindexed_scans(catalog, left)),
6990 right: Box::new(lower_unindexed_scans(catalog, right)),
6991 all: *all,
6992 },
6993 PlanNode::Explain { input } => PlanNode::Explain {
6994 input: Box::new(lower_unindexed_scans(catalog, input)),
6995 },
6996 PlanNode::NestedLoopJoin {
6997 left,
6998 right,
6999 on,
7000 kind,
7001 } => PlanNode::NestedLoopJoin {
7002 left: Box::new(lower_unindexed_scans(catalog, left)),
7003 right: Box::new(lower_unindexed_scans(catalog, right)),
7004 on: on.clone(),
7005 kind: *kind,
7006 },
7007 PlanNode::IndexScan { table, column, key } => {
7008 if let Some(tbl) = catalog.get_table(table) {
7009 if tbl.has_index(column) {
7010 return plan.clone();
7011 }
7012 }
7013 PlanNode::Filter {
7014 input: Box::new(PlanNode::SeqScan {
7015 table: table.clone(),
7016 }),
7017 predicate: Expr::BinaryOp(
7018 Box::new(Expr::Field(column.clone())),
7019 BinOp::Eq,
7020 Box::new(key.clone()),
7021 ),
7022 }
7023 }
7024 _ => plan.clone(),
7026 }
7027}
7028
7029fn stored_json_path_expr(path: &powdb_storage::stored_json_path::StoredJsonPathV1) -> Expr {
7030 use powdb_storage::stored_json_path::StoredJsonPathSegmentV1;
7031
7032 Expr::JsonPath {
7033 base: Box::new(Expr::Field(path.column.clone())),
7034 segments: path
7035 .segments
7036 .iter()
7037 .map(|segment| match segment {
7038 StoredJsonPathSegmentV1::Key(key) => PathSeg::Key(key.clone()),
7039 StoredJsonPathSegmentV1::Index(index) => PathSeg::Index(*index),
7040 })
7041 .collect(),
7042 }
7043}
7044
7045fn synthesize_expr_range_predicate(
7046 path: &powdb_storage::stored_json_path::StoredJsonPathV1,
7047 start: &Option<(Expr, bool)>,
7048 end: &Option<(Expr, bool)>,
7049) -> Expr {
7050 let lower = start.as_ref().map(|(expr, inclusive)| {
7051 Expr::BinaryOp(
7052 Box::new(stored_json_path_expr(path)),
7053 if *inclusive { BinOp::Gte } else { BinOp::Gt },
7054 Box::new(expr.clone()),
7055 )
7056 });
7057 let upper = end.as_ref().map(|(expr, inclusive)| {
7058 Expr::BinaryOp(
7059 Box::new(stored_json_path_expr(path)),
7060 if *inclusive { BinOp::Lte } else { BinOp::Lt },
7061 Box::new(expr.clone()),
7062 )
7063 });
7064 match (lower, upper) {
7065 (Some(lower), Some(upper)) => Expr::BinaryOp(Box::new(lower), BinOp::And, Box::new(upper)),
7066 (Some(lower), None) => lower,
7067 (None, Some(upper)) => upper,
7068 (None, None) => Expr::Literal(Literal::Bool(true)),
7069 }
7070}
7071
7072pub(super) fn synthesize_range_predicate(
7074 column: &str,
7075 start: &Option<(Expr, bool)>,
7076 end: &Option<(Expr, bool)>,
7077) -> Expr {
7078 let lower = start.as_ref().map(|(expr, inclusive)| {
7079 let op = if *inclusive { BinOp::Gte } else { BinOp::Gt };
7080 Expr::BinaryOp(
7081 Box::new(Expr::Field(column.to_string())),
7082 op,
7083 Box::new(expr.clone()),
7084 )
7085 });
7086 let upper = end.as_ref().map(|(expr, inclusive)| {
7087 let op = if *inclusive { BinOp::Lte } else { BinOp::Lt };
7088 Expr::BinaryOp(
7089 Box::new(Expr::Field(column.to_string())),
7090 op,
7091 Box::new(expr.clone()),
7092 )
7093 });
7094 match (lower, upper) {
7095 (Some(l), Some(u)) => Expr::BinaryOp(Box::new(l), BinOp::And, Box::new(u)),
7096 (Some(l), None) => l,
7097 (None, Some(u)) => u,
7098 (None, None) => Expr::Literal(Literal::Bool(true)),
7099 }
7100}
7101
7102fn scan_table(scan: &PlanNode) -> Option<&str> {
7107 match scan {
7108 PlanNode::IndexScan { table, .. }
7109 | PlanNode::RangeScan { table, .. }
7110 | PlanNode::ExprIndexScan { table, .. }
7111 | PlanNode::ExprRangeScan { table, .. } => Some(table),
7112 _ => None,
7113 }
7114}
7115
7116pub(super) fn range_matches(
7117 val: &Value,
7118 start: &Option<Value>,
7119 start_inc: bool,
7120 end: &Option<Value>,
7121 end_inc: bool,
7122) -> bool {
7123 if let Some(ref s) = start {
7124 if start_inc {
7125 if val < s {
7126 return false;
7127 }
7128 } else if val <= s {
7129 return false;
7130 }
7131 }
7132 if let Some(ref e) = end {
7133 if end_inc {
7134 if val > e {
7135 return false;
7136 }
7137 } else if val >= e {
7138 return false;
7139 }
7140 }
7141 true
7142}
7143
7144fn collect_plan_qualifiers(plan: &PlanNode, qualifiers: &mut HashSet<String>) {
7145 match plan {
7146 PlanNode::SeqScan { table }
7147 | PlanNode::IndexScan { table, .. }
7148 | PlanNode::RangeScan { table, .. }
7149 | PlanNode::ExprIndexScan { table, .. }
7150 | PlanNode::ExprRangeScan { table, .. }
7151 | PlanNode::OrderedExprIndexScan { table, .. } => {
7152 qualifiers.insert(table.clone());
7153 }
7154 PlanNode::AliasScan { alias, .. } => {
7155 qualifiers.insert(alias.clone());
7156 }
7157 PlanNode::Filter { input, .. }
7158 | PlanNode::Project { input, .. }
7159 | PlanNode::Sort { input, .. }
7160 | PlanNode::Limit { input, .. }
7161 | PlanNode::Offset { input, .. }
7162 | PlanNode::Aggregate { input, .. }
7163 | PlanNode::Distinct { input }
7164 | PlanNode::GroupBy { input, .. }
7165 | PlanNode::Update { input, .. }
7166 | PlanNode::Delete { input, .. }
7167 | PlanNode::Window { input, .. }
7168 | PlanNode::Explain { input } => collect_plan_qualifiers(input, qualifiers),
7169 PlanNode::NestedLoopJoin { left, right, .. } | PlanNode::Union { left, right, .. } => {
7170 collect_plan_qualifiers(left, qualifiers);
7171 collect_plan_qualifiers(right, qualifiers);
7172 }
7173 _ => {}
7174 }
7175}
7176
7177fn qualified_ref(expr: &Expr) -> Option<&str> {
7178 match expr {
7179 Expr::QualifiedField { qualifier, .. } => Some(qualifier),
7180 _ => None,
7181 }
7182}
7183
7184fn explain_join_strategy(
7185 left: &PlanNode,
7186 right: &PlanNode,
7187 on: Option<&Expr>,
7188 kind: JoinKind,
7189) -> &'static str {
7190 if matches!(kind, JoinKind::Cross) {
7191 return "nested-loop-bounded";
7192 }
7193 let Some(predicate) = on else {
7194 return "nested-loop-bounded";
7195 };
7196 let mut conjunctions = Vec::new();
7197 flatten_conjunctions(predicate, &mut conjunctions);
7198 let mut left_qualifiers = HashSet::new();
7199 let mut right_qualifiers = HashSet::new();
7200 collect_plan_qualifiers(left, &mut left_qualifiers);
7201 collect_plan_qualifiers(right, &mut right_qualifiers);
7202
7203 let has_cross_side_equi = conjunctions.iter().any(|expr| {
7204 let Expr::BinaryOp(lhs, BinOp::Eq, rhs) = expr else {
7205 return false;
7206 };
7207 let (Some(lhs_q), Some(rhs_q)) = (qualified_ref(lhs), qualified_ref(rhs)) else {
7208 return false;
7209 };
7210 (left_qualifiers.contains(lhs_q) && right_qualifiers.contains(rhs_q))
7211 || (left_qualifiers.contains(rhs_q) && right_qualifiers.contains(lhs_q))
7212 });
7213 if has_cross_side_equi {
7214 if conjunctions.len() > 1 {
7215 "hash+residual"
7216 } else {
7217 "hash"
7218 }
7219 } else {
7220 "nested-loop-bounded"
7221 }
7222}
7223
7224pub(super) fn format_plan_tree(catalog: &Catalog, plan: &PlanNode, depth: usize) -> String {
7227 let indent = " ".repeat(depth);
7228 match plan {
7229 PlanNode::SeqScan { table } => format!("{indent}SeqScan table={table}"),
7230 PlanNode::AliasScan { table, alias } => {
7231 format!("{indent}AliasScan table={table} alias={alias}")
7232 }
7233 PlanNode::IndexScan { table, column, key } => {
7234 let base = format!("{indent}IndexScan table={table} column={column} key={key:?}");
7235 match catalog.index_stats(table, column) {
7236 Some(stats) => {
7237 let unique = catalog.is_index_unique(table, column) == Some(true);
7238 let est = eq_est_rows(&stats, unique, probes_empty_sentinel(key));
7239 format!(
7240 "{base} est_rows={est} entries={} distinct={}",
7241 stats.total_entries, stats.distinct_keys
7242 )
7243 }
7244 None => base,
7245 }
7246 }
7247 PlanNode::RangeScan {
7248 table,
7249 column,
7250 start,
7251 end,
7252 } => {
7253 let s = match start {
7254 Some((expr, inc)) => {
7255 let op = if *inc { ">=" } else { ">" };
7256 format!("{op}{expr:?}")
7257 }
7258 None => "unbounded".to_string(),
7259 };
7260 let e = match end {
7261 Some((expr, inc)) => {
7262 let op = if *inc { "<=" } else { "<" };
7263 format!("{op}{expr:?}")
7264 }
7265 None => "unbounded".to_string(),
7266 };
7267 format!("{indent}RangeScan table={table} column={column} [{s}, {e}]")
7268 }
7269 PlanNode::ExprIndexScan { table, path, key } => {
7270 let meta = resolve_expression_index(catalog, table, path);
7271 let index_id = meta
7272 .as_ref()
7273 .map(|metadata| metadata.index_id.to_string())
7274 .unwrap_or_else(|| "unresolved".to_string());
7275 let base = format!(
7276 "{indent}ExprIndexScan table={table} path={} index_id={index_id} key={key:?}",
7277 path.canonical_text()
7278 );
7279 match meta.and_then(|m| {
7280 catalog
7281 .expression_index_stats(table, m.index_id)
7282 .map(|stats| (m.unique, stats))
7283 }) {
7284 Some((unique, stats)) => {
7285 let est = eq_est_rows(&stats, unique, probes_empty_sentinel(key));
7286 format!(
7287 "{base} est_rows={est} entries={} distinct={}",
7288 stats.total_entries, stats.distinct_keys
7289 )
7290 }
7291 None => base,
7292 }
7293 }
7294 PlanNode::ExprRangeScan {
7295 table,
7296 path,
7297 start,
7298 end,
7299 } => {
7300 let index_id = resolve_expression_index(catalog, table, path)
7301 .map(|metadata| metadata.index_id.to_string())
7302 .unwrap_or_else(|| "unresolved".to_string());
7303 format!(
7304 "{indent}ExprRangeScan table={table} path={} index_id={index_id} start={start:?} end={end:?}",
7305 path.canonical_text()
7306 )
7307 }
7308 PlanNode::OrderedExprIndexScan {
7309 table,
7310 path,
7311 descending,
7312 limit,
7313 offset,
7314 } => {
7315 let index_id = resolve_expression_index(catalog, table, path)
7316 .map(|metadata| metadata.index_id.to_string())
7317 .unwrap_or_else(|| "unresolved".to_string());
7318 format!(
7319 "{indent}OrderedExprIndexScan table={table} path={} index_id={index_id} descending={descending} limit={limit:?} offset={offset:?}",
7320 path.canonical_text()
7321 )
7322 }
7323 PlanNode::Filter { input, predicate } => {
7324 let child = format_plan_tree(catalog, input, depth + 1);
7325 format!("{indent}Filter predicate={predicate:?}\n{child}")
7326 }
7327 PlanNode::Project { input, fields } => {
7328 let names: Vec<String> = fields
7329 .iter()
7330 .map(|f| match &f.alias {
7331 Some(a) => format!("{a}: {:?}", f.expr),
7332 None => format!("{:?}", f.expr),
7333 })
7334 .collect();
7335 let child = format_plan_tree(catalog, input, depth + 1);
7336 format!("{indent}Project fields=[{}]\n{child}", names.join(", "))
7337 }
7338 PlanNode::Sort { input, keys } => {
7339 let ks: Vec<String> = keys
7340 .iter()
7341 .map(|k| {
7342 let expr = expression_output_name(&k.expr);
7343 if k.descending {
7344 format!("{expr} desc")
7345 } else {
7346 expr
7347 }
7348 })
7349 .collect();
7350 let child = format_plan_tree(catalog, input, depth + 1);
7351 format!("{indent}Sort keys=[{}]\n{child}", ks.join(", "))
7352 }
7353 PlanNode::Limit { input, count } => {
7354 let child = format_plan_tree(catalog, input, depth + 1);
7355 format!("{indent}Limit count={count:?}\n{child}")
7356 }
7357 PlanNode::Offset { input, count } => {
7358 let child = format_plan_tree(catalog, input, depth + 1);
7359 format!("{indent}Offset count={count:?}\n{child}")
7360 }
7361 PlanNode::Aggregate {
7362 input,
7363 function,
7364 argument,
7365 mode,
7366 provenance_alias: _,
7367 } => {
7368 let argument = argument
7369 .as_ref()
7370 .map(expression_output_name)
7371 .unwrap_or_else(|| "*".to_string());
7372 let child = format_plan_tree(catalog, input, depth + 1);
7373 format!("{indent}Aggregate fn={function:?} mode={mode:?} argument={argument}\n{child}")
7374 }
7375 PlanNode::NestedLoopJoin {
7376 left,
7377 right,
7378 on,
7379 kind,
7380 } => {
7381 let left_child = format_plan_tree(catalog, left, depth + 1);
7382 let right_child = format_plan_tree(catalog, right, depth + 1);
7383 let on_str = match on {
7384 Some(pred) => format!("{pred:?}"),
7385 None => "none".to_string(),
7386 };
7387 let strategy = explain_join_strategy(left, right, on.as_ref(), *kind);
7388 format!(
7389 "{indent}NestedLoopJoin kind={kind:?} strategy={strategy} on={on_str}\n{left_child}\n{right_child}"
7390 )
7391 }
7392 PlanNode::Distinct { input } => {
7393 let child = format_plan_tree(catalog, input, depth + 1);
7394 format!("{indent}Distinct\n{child}")
7395 }
7396 PlanNode::GroupBy {
7397 input,
7398 keys,
7399 aggregates,
7400 having,
7401 } => {
7402 let agg_strs: Vec<String> = aggregates
7403 .iter()
7404 .map(|a| {
7405 format!(
7406 "{:?}({}) mode={:?} as {}",
7407 a.function,
7408 expression_output_name(&a.argument),
7409 a.mode,
7410 a.output_name
7411 )
7412 })
7413 .collect();
7414 let having_str = match having {
7415 Some(h) => format!(" having={h:?}"),
7416 None => String::new(),
7417 };
7418 let key_strs: Vec<String> = keys.iter().map(|k| k.output_name()).collect();
7419 let child = format_plan_tree(catalog, input, depth + 1);
7420 format!(
7421 "{indent}GroupBy keys=[{}] aggs=[{}]{having_str}\n{child}",
7422 key_strs.join(", "),
7423 agg_strs.join(", "),
7424 )
7425 }
7426 PlanNode::Insert { table, rows, .. } => {
7427 let cols: Vec<&str> = rows
7428 .first()
7429 .map(|r| r.iter().map(|a| a.field.as_str()).collect())
7430 .unwrap_or_default();
7431 format!(
7432 "{indent}Insert table={table} rows={} cols=[{}]",
7433 rows.len(),
7434 cols.join(", ")
7435 )
7436 }
7437 PlanNode::Upsert {
7438 table,
7439 key_column,
7440 assignments,
7441 on_conflict,
7442 } => {
7443 let cols: Vec<&str> = assignments.iter().map(|a| a.field.as_str()).collect();
7444 let conflict_cols: Vec<&str> = on_conflict.iter().map(|a| a.field.as_str()).collect();
7445 if conflict_cols.is_empty() {
7446 format!(
7447 "{indent}Upsert table={table} key={key_column} cols=[{}]",
7448 cols.join(", ")
7449 )
7450 } else {
7451 format!(
7452 "{indent}Upsert table={table} key={key_column} cols=[{}] on_conflict=[{}]",
7453 cols.join(", "),
7454 conflict_cols.join(", ")
7455 )
7456 }
7457 }
7458 PlanNode::Update {
7459 input,
7460 table,
7461 assignments,
7462 returning,
7463 } => {
7464 let cols: Vec<&str> = assignments.iter().map(|a| a.field.as_str()).collect();
7465 let child = format_plan_tree(catalog, input, depth + 1);
7466 let ret = if *returning { " returning" } else { "" };
7467 format!(
7468 "{indent}Update table={table} set=[{}]{ret}\n{child}",
7469 cols.join(", ")
7470 )
7471 }
7472 PlanNode::Delete {
7473 input,
7474 table,
7475 returning,
7476 } => {
7477 let child = format_plan_tree(catalog, input, depth + 1);
7478 let ret = if *returning { " returning" } else { "" };
7479 format!("{indent}Delete table={table}{ret}\n{child}")
7480 }
7481 PlanNode::CreateTable { name, fields, .. } => {
7482 let fs: Vec<String> = fields
7483 .iter()
7484 .map(|f| {
7485 let mut mods = String::new();
7486 if f.required {
7487 mods.push_str(" required");
7488 }
7489 if f.unique {
7490 mods.push_str(" unique");
7491 }
7492 format!("{}: {}{mods}", f.name, f.type_name)
7493 })
7494 .collect();
7495 format!("{indent}CreateTable name={name} fields=[{}]", fs.join(", "))
7496 }
7497 PlanNode::AlterTable { table, action } => {
7498 format!("{indent}AlterTable table={table} action={action:?}")
7499 }
7500 PlanNode::DropTable { name, .. } => format!("{indent}DropTable name={name}"),
7501 PlanNode::CreateView { name, .. } => format!("{indent}CreateView name={name}"),
7502 PlanNode::RefreshView { name } => format!("{indent}RefreshView name={name}"),
7503 PlanNode::DropView { name, .. } => format!("{indent}DropView name={name}"),
7504 PlanNode::ListTypes => format!("{indent}ListTypes"),
7505 PlanNode::Describe { table } => format!("{indent}Describe table={table}"),
7506 PlanNode::Window { input, windows } => {
7507 let ws: Vec<String> = windows
7508 .iter()
7509 .map(|w| format!("{:?} as {}", w.function, w.output_name))
7510 .collect();
7511 let child = format_plan_tree(catalog, input, depth + 1);
7512 format!("{indent}Window fns=[{}]\n{child}", ws.join(", "))
7513 }
7514 PlanNode::Union { left, right, all } => {
7515 let kind = if *all { "UNION ALL" } else { "UNION" };
7516 let left_child = format_plan_tree(catalog, left, depth + 1);
7517 let right_child = format_plan_tree(catalog, right, depth + 1);
7518 format!("{indent}{kind}\n{left_child}\n{right_child}")
7519 }
7520 PlanNode::Explain { input } => {
7521 let child = format_plan_tree(catalog, input, depth + 1);
7522 format!("{indent}Explain\n{child}")
7523 }
7524 PlanNode::Begin => format!("{indent}Begin"),
7525 PlanNode::Commit => format!("{indent}Commit"),
7526 PlanNode::Rollback => format!("{indent}Rollback"),
7527 }
7528}