1use std::sync::Arc;
46
47use rudb_catalog::{Catalog, QualifiedName};
48use rudb_common::{Cancel, Memory, Result, Session, Value};
49use rudb_functions::TableFunction;
50use rudb_metrics::{Counters, Driver, Report};
51use rudb_parquet::{Bound, Op};
52use rudb_pipeline::{
53 BufferId, DynSink, DynStream, Pipeline, PipelineId, Source, Watched, root, root_in_order,
54};
55use rudb_plan::{
56 CompareOp, ConjunctionOp, Expr, ExprRef, Node, NodeRef, PipelineRef, Plan, ROOT, Shape, Slice,
57 seams_of,
58};
59use rudb_seam::Settings;
60
61use crate::enginenames::{
62 database_size, dialects, extensions, grammar_extensions, optimizers, platform, user_agent,
63 version,
64};
65use crate::entrynames::{columnnames, databasenames, schemanames, tablenames, viewnames};
66use crate::fetch::{Fetch, TableFetch};
67use crate::functionnames::functionnames;
68use crate::gather::{Gather, Keep};
69use crate::group::{Aggregate, Distinct};
70use crate::join::{CrossProduct, Gathered, Join};
71use crate::keywords::keywords;
72use crate::query::Query;
73use crate::register::registries;
74use crate::schema::Schema;
75use crate::setop::SetOp;
76use crate::settingnames::settingnames;
77use crate::sort::Sort;
78use crate::source::{Dummy, FileScan, Scan, Series, Values};
79use crate::strategies::strategies;
80use crate::stream::{Filter, Limit, Project};
81use crate::topn::TopN;
82use crate::typenames::typenames;
83
84pub fn build<'a>(plan: &'a Plan, catalog: &'a Catalog) -> Result<Query<'a>> {
94 build_with(plan, catalog, &Cancel::new(), &Memory::unlimited(), &Settings::new())
95}
96
97pub fn build_with<'a>(
125 plan: &'a Plan,
126 catalog: &'a Catalog,
127 cancel: &Cancel,
128 memory: &Memory,
129 seams: &Settings,
130) -> Result<Query<'a>> {
131 build_measured(plan, catalog, cancel, memory, seams, &Session::new(), &Report::new())
132}
133
134pub fn build_measured<'a>(
150 plan: &'a Plan,
151 catalog: &'a Catalog,
152 cancel: &Cancel,
153 memory: &Memory,
154 seams: &Settings,
155 session: &Session,
156 report: &Report,
157) -> Result<Query<'a>> {
158 build_measured_with_sink(
159 plan,
160 catalog,
161 BuildUnder { cancel, memory, seams, session, report },
162 None,
163 )
164}
165
166pub fn build_measured_into<'a>(
176 plan: &'a Plan,
177 catalog: &'a Catalog,
178 cancel: &Cancel,
179 memory: &Memory,
180 seams: &Settings,
181 session: &Session,
182 sink: Arc<dyn DynSink + 'a>,
183) -> Result<Query<'a>> {
184 let report = Report::new();
185 build_measured_with_sink(
186 plan,
187 catalog,
188 BuildUnder { cancel, memory, seams, session, report: &report },
189 Some(sink),
190 )
191}
192
193struct BuildUnder<'a> {
194 cancel: &'a Cancel,
195 memory: &'a Memory,
196 seams: &'a Settings,
197 session: &'a Session,
198 report: &'a Report,
199}
200
201#[derive(Clone, Copy, Default)]
202struct AggregateBound {
203 max_groups: Option<usize>,
204 top_counts: Option<usize>,
205 having_count: Option<(usize, i64)>,
206}
207
208fn build_measured_with_sink<'a>(
209 plan: &'a Plan,
210 catalog: &'a Catalog,
211 under: BuildUnder<'_>,
212 sink: Option<Arc<dyn DynSink + 'a>>,
213) -> Result<Query<'a>> {
214 let BuildUnder { cancel, memory, seams, session, report } = under;
215 let shape = Shape::of(plan);
216 for pipeline in shape.all() {
217 report.pipeline(pipeline);
218 for waits_for in shape.waits_for(pipeline) {
219 report.depends(pipeline, *waits_for);
220 }
221 }
222 let mut building = Building {
223 plan,
224 catalog,
225 cancel,
226 memory,
227 seams,
228 session,
229 report,
230 shape,
231 done: Vec::new(),
232 drivers: Vec::new(),
233 pruning: Vec::new(),
234 top_counts: Vec::new(),
235 };
236 let segment = building.node(plan.root())?;
237 let schema = segment.schema.clone();
238 let reader = if let Some(sink) = sink {
244 building.close(segment, ROOT, sink);
245 None
246 } else {
247 let (sink, reader) = if ordered(plan, plan.root()) {
248 root(BufferId(0), None)
249 } else {
250 root_in_order(BufferId(0), None)
251 };
252 building.close(segment, ROOT, Arc::new(sink));
253 Some(reader)
254 };
255 let Building { done, drivers, .. } = building;
256 Query::new(done, drivers, reader, schema)
257}
258
259fn ordered(plan: &Plan, node: NodeRef) -> bool {
265 match *plan.node(node) {
266 Node::Sort { .. } | Node::TopN { .. } => true,
267 Node::Project { input, .. }
268 | Node::Filter { input, .. }
269 | Node::Limit { input, .. }
270 | Node::Fetch { input, .. } => ordered(plan, input),
271 _ => false,
272 }
273}
274
275fn count_top_aggregate(plan: &Plan, input: NodeRef, keys: Slice) -> Option<NodeRef> {
282 let [key] = plan.sort_key_list(keys) else { return None };
283 if !key.descending {
284 return None;
285 }
286 let Expr::Column(ordered) = *plan.expr(key.expr) else { return None };
287 let (aggregate, output) = match *plan.node(input) {
288 Node::Project { input, index, exprs, .. } => {
289 if ordered.table != index {
290 return None;
291 }
292 let projected = *plan.expr_list(exprs).get(ordered.column as usize)?;
293 let Expr::Column(output) = *plan.expr(projected) else { return None };
294 (input, output)
295 }
296 Node::Aggregate { index, .. } if ordered.table == index => (input, ordered),
297 _ => return None,
298 };
299 let Node::Aggregate { index, groups, aggregates, .. } = *plan.node(aggregate) else {
300 return None;
301 };
302 if output.table != index || output.column as usize != plan.expr_list(groups).len() {
303 return None;
304 }
305 let first = *plan.expr_list(aggregates).first()?;
306 let Expr::Aggregate { name, args, distinct, filter } = *plan.expr(first) else {
307 return None;
308 };
309 let count_star = plan.string(name) == "count_star"
310 && plan.expr_list(args).is_empty()
311 && !distinct
312 && filter.is_none();
313 let distinct_count = plan.string(name) == "count"
314 && plan.expr_list(args).len() == 1
315 && distinct
316 && filter.is_none();
317 (count_star || distinct_count).then_some(aggregate)
318}
319
320fn count_having_aggregate(
325 plan: &Plan,
326 input: NodeRef,
327 predicate: ExprRef,
328) -> Option<(NodeRef, usize, i64)> {
329 let Node::Aggregate { index, groups, aggregates, .. } = *plan.node(input) else { return None };
330 let Expr::Compare { op, left, right } = *plan.expr(predicate) else { return None };
331 let Expr::Column(column) = *plan.expr(left) else { return None };
332 let Expr::Constant(value) = *plan.expr(right) else { return None };
333 let Value::BigInt(value) = *plan.value(value) else { return None };
334 if column.table != index {
335 return None;
336 }
337 let call = (column.column as usize).checked_sub(plan.expr_list(groups).len())?;
338 let aggregate = *plan.expr_list(aggregates).get(call)?;
339 let Expr::Aggregate { name, args, distinct, filter } = *plan.expr(aggregate) else {
340 return None;
341 };
342 if plan.string(name) != "count_star"
343 || !plan.expr_list(args).is_empty()
344 || distinct
345 || filter.is_some()
346 {
347 return None;
348 }
349 let minimum = match op {
350 CompareOp::Greater => value.checked_add(1)?,
351 CompareOp::GreaterOrEqual => value,
352 _ => return None,
353 };
354 Some((input, call, minimum))
355}
356
357fn bounds(plan: &Plan, input: NodeRef, predicate: ExprRef) -> Vec<(usize, Op, Bound)> {
372 let index = match *plan.node(input) {
373 Node::TableFunction { index, function, .. } => {
374 if TableFunction::lookup(plan.string(function)) != Some(TableFunction::ReadParquet) {
375 return Vec::new();
376 }
377 index
378 }
379 Node::Get { index, .. } => index,
380 _ => return Vec::new(),
381 };
382 let mut tests = Vec::new();
383 conjuncts(plan, predicate, index, &mut tests);
384 tests
385}
386
387fn conjuncts(plan: &Plan, predicate: ExprRef, index: u32, out: &mut Vec<(usize, Op, Bound)>) {
389 match *plan.expr(predicate) {
390 Expr::Conjunction { op: ConjunctionOp::And, children } => {
391 for child in plan.expr_list(children) {
392 conjuncts(plan, *child, index, out);
393 }
394 }
395 Expr::Compare { op, left, right } => {
396 if let Some(test) = comparison(plan, op, left, right, index) {
397 out.push(test);
398 }
399 }
400 _ => {}
401 }
402}
403
404fn comparison(
412 plan: &Plan,
413 op: CompareOp,
414 left: ExprRef,
415 right: ExprRef,
416 index: u32,
417) -> Option<(usize, Op, Bound)> {
418 let op = match op {
419 CompareOp::Equal => Op::Equal,
420 CompareOp::Less => Op::Less,
421 CompareOp::LessOrEqual => Op::LessOrEqual,
422 CompareOp::Greater => Op::Greater,
423 CompareOp::GreaterOrEqual => Op::GreaterOrEqual,
424 CompareOp::NotEqual | CompareOp::DistinctFrom | CompareOp::NotDistinctFrom => return None,
425 };
426 let (op, binding, value) = match (plan.expr(left), plan.expr(right)) {
427 (Expr::Column(binding), Expr::Constant(value)) => (op, *binding, *value),
428 (Expr::Constant(value), Expr::Column(binding)) => (op.flipped(), *binding, *value),
429 _ => return None,
430 };
431 if binding.table != index {
432 return None;
433 }
434 Some((binding.column as usize, op, Bound::of_value(plan.value(value))?))
435}
436
437struct Segment<'a> {
443 source: Arc<dyn Source + 'a>,
444 streams: Vec<Arc<dyn DynStream + 'a>>,
446 schema: Schema,
448 after: Vec<PipelineRef>,
450}
451
452impl<'a> Segment<'a> {
453 fn new(source: Arc<dyn Source + 'a>, schema: Schema) -> Self {
455 Self { source, streams: Vec::new(), schema, after: Vec::new() }
456 }
457
458 fn reading(source: Arc<dyn Source + 'a>, schema: Schema, after: PipelineRef) -> Self {
460 Self { source, streams: Vec::new(), schema, after: vec![after] }
461 }
462
463 fn then(mut self, stream: Arc<dyn DynStream + 'a>, schema: Schema) -> Self {
465 self.streams.push(stream);
466 self.schema = schema;
467 self
468 }
469}
470
471struct Building<'a, 'b> {
473 plan: &'a Plan,
474 catalog: &'a Catalog,
475 cancel: &'b Cancel,
476 memory: &'b Memory,
477 seams: &'b Settings,
478 session: &'b Session,
480 report: &'b Report,
481 shape: Shape,
482 done: Vec<Pipeline<'a>>,
484 drivers: Vec<Arc<Driver>>,
486 pruning: Vec<(usize, Op, Bound)>,
494 top_counts: Vec<(NodeRef, usize)>,
496}
497
498impl<'a> Building<'a, '_> {
499 fn gathered(&self, node: NodeRef) -> u32 {
505 self.shape.gathered(node).expect("a node with two inputs has a second operator")
506 }
507
508 fn close(&mut self, segment: Segment<'a>, id: PipelineRef, sink: Arc<dyn DynSink + 'a>) {
510 let mut pipeline = Pipeline::new(PipelineId(id), segment.source, sink);
511 for stream in segment.streams {
512 pipeline = pipeline.then(stream);
513 }
514 for after in segment.after {
515 pipeline = pipeline.after(PipelineId(after));
516 }
517 self.done.push(pipeline);
518 self.drivers.push(self.report.driving(id));
519 }
520
521 fn watch(
536 &self,
537 node: NodeRef,
538 id: u32,
539 pipeline: u32,
540 kind: &str,
541 detail: Option<&str>,
542 ) -> Arc<Counters> {
543 let mut counters = Counters::new(id, pipeline, kind);
544 if let Some(detail) = detail {
545 counters = counters.detailed(detail);
546 }
547 for seam in seams_of(self.plan.node(node)) {
548 if let Some(running) = registries().running(*seam, self.seams) {
549 counters = counters.chose(seam.name(), &running.name, running.is_reference);
550 }
551 }
552 self.report.watch(counters)
553 }
554
555 fn aggregate(
556 &mut self,
557 reference: NodeRef,
558 input: NodeRef,
559 index: u32,
560 groups: Slice,
561 aggregates: Slice,
562 bound: AggregateBound,
563 ) -> Result<Segment<'a>> {
564 let below = self.node(input)?;
565 let (aggregate, out) =
566 Aggregate::new(self.plan, &below.schema, index, groups, aggregates, self.memory)?;
567 let aggregate = aggregate.in_session(self.session);
568 let aggregate = match bound.max_groups {
569 Some(limit) => aggregate.limit_groups(limit),
570 None => aggregate,
571 };
572 let aggregate = match bound.top_counts {
573 Some(bound) => aggregate.top_counts(bound),
574 None => aggregate,
575 };
576 let aggregate = match bound.having_count {
577 Some((call, minimum)) => aggregate.having_count(call, minimum),
578 None => aggregate,
579 };
580 let schema = aggregate.schema().clone();
581 let id = self.shape.operator(reference);
582 let pipeline = self.shape.pipeline(reference);
583 let counters = self.watch(reference, id, pipeline, "Aggregate", None);
584 let reading = Arc::clone(&counters);
585 self.close(below, pipeline, Arc::new(Watched::new(aggregate, counters)));
586 Ok(Segment::reading(Arc::new(Watched::new(out, reading)), schema, pipeline))
587 }
588
589 fn node(&mut self, reference: NodeRef) -> Result<Segment<'a>> {
591 let plan = self.plan;
592 let memory = self.memory;
593 let id = self.shape.operator(reference);
594 let pipeline = self.shape.pipeline(reference);
595 let segment = match *plan.node(reference) {
596 Node::Get { catalog: database, schema, table, index, columns, .. } => {
597 let name = QualifiedName::new(
598 plan.string(database),
599 plan.string(schema),
600 plan.string(table),
601 );
602 let tests = std::mem::take(&mut self.pruning);
603 let scan = Scan::new(plan, self.catalog.table(&name)?, index, columns, tests)?;
604 let schema = scan.schema().clone();
605 let counters =
606 self.watch(reference, id, pipeline, "Scan", Some(plan.string(table)));
607 Segment::new(Arc::new(Watched::new(scan, counters)), schema)
608 }
609 Node::Dummy => {
610 let dummy = Dummy::new();
611 let schema = dummy.schema().clone();
612 let counters = self.watch(reference, id, pipeline, "Dummy", None);
613 Segment::new(Arc::new(Watched::new(dummy, counters)), schema)
614 }
615 Node::Values { index, columns, rows } => {
616 let values = Values::new(plan, index, columns, rows, self.session)?;
617 let schema = values.schema().clone();
618 let counters = self.watch(reference, id, pipeline, "Values", None);
619 Segment::new(Arc::new(Watched::new(values, counters)), schema)
620 }
621 Node::TableFunction { index, function, args, options, settings, columns } => {
622 let name = plan.string(function);
623 match TableFunction::lookup(name) {
624 Some(function @ (TableFunction::ReadParquet | TableFunction::ReadCsv)) => {
625 let counters = self.watch(reference, id, pipeline, "FileScan", Some(name));
626 let tests = std::mem::take(&mut self.pruning);
627 let scan = FileScan::new(
628 plan, index, function, args, options, settings, columns, tests,
629 )?
630 .watched(counters.clone());
631 let schema = scan.schema().clone();
632 Segment::new(Arc::new(Watched::new(scan, counters)), schema)
633 }
634 Some(
635 function @ (TableFunction::RudbStrategies
636 | TableFunction::DuckdbKeywords
637 | TableFunction::DuckdbTypes
638 | TableFunction::DuckdbFunctions
639 | TableFunction::DuckdbSettings
640 | TableFunction::DuckdbDatabases
641 | TableFunction::DuckdbSchemas
642 | TableFunction::DuckdbTables
643 | TableFunction::DuckdbViews
644 | TableFunction::DuckdbColumns
645 | TableFunction::DuckdbExtensions
646 | TableFunction::DuckdbOptimizers
647 | TableFunction::DuckdbDialects
648 | TableFunction::DuckdbGrammarExtensions
649 | TableFunction::PragmaVersion
650 | TableFunction::PragmaPlatform
651 | TableFunction::PragmaUserAgent
652 | TableFunction::PragmaDatabaseSize),
653 ) => {
654 let table = match function {
655 TableFunction::DuckdbKeywords => keywords(plan, index, columns)?,
656 TableFunction::DuckdbTypes => typenames(plan, index, columns)?,
657 TableFunction::DuckdbFunctions => functionnames(plan, index, columns)?,
658 TableFunction::DuckdbSettings => {
659 settingnames(self.session, plan, index, columns)?
660 }
661 TableFunction::DuckdbDatabases => {
662 databasenames(self.catalog, plan, index, columns)?
663 }
664 TableFunction::DuckdbSchemas => {
665 schemanames(self.catalog, plan, index, columns)?
666 }
667 TableFunction::DuckdbTables => {
668 tablenames(self.catalog, plan, index, columns)?
669 }
670 TableFunction::DuckdbViews => {
671 viewnames(self.catalog, plan, index, columns)?
672 }
673 TableFunction::DuckdbColumns => {
674 columnnames(self.catalog, plan, index, columns)?
675 }
676 TableFunction::DuckdbExtensions => extensions(plan, index, columns)?,
677 TableFunction::DuckdbOptimizers => optimizers(plan, index, columns)?,
678 TableFunction::DuckdbDialects => dialects(plan, index, columns)?,
679 TableFunction::DuckdbGrammarExtensions => {
680 grammar_extensions(plan, index, columns)?
681 }
682 TableFunction::PragmaVersion => version(plan, index, columns)?,
683 TableFunction::PragmaPlatform => platform(plan, index, columns)?,
684 TableFunction::PragmaUserAgent => user_agent(plan, index, columns)?,
685 TableFunction::PragmaDatabaseSize => {
686 database_size(self.catalog, self.memory, plan, index, columns)?
687 }
688 _ => strategies(plan, index, columns)?,
689 };
690 let schema = table.schema().clone();
691 let counters =
695 self.watch(reference, id, pipeline, "Metadata", Some(function.name()));
696 Segment::new(Arc::new(Watched::new(table, counters)), schema)
697 }
698 _ => {
699 let series = Series::new(plan, index, name, args)?;
700 let schema = series.schema().clone();
701 let counters = self.watch(reference, id, pipeline, "Series", Some(name));
702 Segment::new(Arc::new(Watched::new(series, counters)), schema)
703 }
704 }
705 }
706 Node::Fetch { input, index, args, columns, row } => {
707 let below = self.node(input)?;
708 let counters = self.watch(reference, id, pipeline, "Fetch", None);
709 let fetch = Fetch::new(plan, &below.schema, index, args, columns, row)?
710 .in_session(self.session)
711 .watched(counters.clone());
712 let schema = fetch.schema().clone();
713 below.then(Arc::new(Watched::new(fetch, counters)), schema)
714 }
715 Node::TableFetch { input, index, catalog, schema, table, columns, row } => {
716 let below = self.node(input)?;
717 let name = QualifiedName::new(
718 plan.string(catalog),
719 plan.string(schema),
720 plan.string(table),
721 );
722 let counters = self.watch(reference, id, pipeline, "TableFetch", None);
723 let fetch = TableFetch::new(
724 plan,
725 &below.schema,
726 index,
727 self.catalog.table(&name)?,
728 columns,
729 row,
730 )?
731 .in_session(self.session);
732 let schema = fetch.schema().clone();
733 below.then(Arc::new(Watched::new(fetch, counters)), schema)
734 }
735 Node::Filter { input, predicate } => {
736 self.pruning = bounds(plan, input, predicate);
737 let below = match count_having_aggregate(plan, input, predicate) {
738 Some((aggregate, call, minimum)) => {
739 let Node::Aggregate { input: under, index, groups, aggregates } =
740 *plan.node(aggregate)
741 else {
742 unreachable!("count_having_aggregate returned another node")
743 };
744 self.aggregate(
745 aggregate,
746 under,
747 index,
748 groups,
749 aggregates,
750 AggregateBound {
751 max_groups: None,
752 top_counts: None,
753 having_count: Some((call, minimum)),
754 },
755 )?
756 }
757 None => self.node(input)?,
758 };
759 self.pruning = Vec::new();
762 let schema = below.schema.clone();
763 let filter = Filter::new(plan, reference, predicate, &schema, self.seams)?
764 .in_session(self.session);
765 let counters = self.watch(reference, id, pipeline, "Filter", None);
766 below.then(Arc::new(Watched::new(filter, counters)), schema)
767 }
768 Node::Project { input, index, exprs, names } => {
769 let below = self.node(input)?;
770 let project = Project::new(plan, &below.schema, index, exprs, names)?
771 .in_session(self.session);
772 let schema = project.schema().clone();
773 let counters = self.watch(reference, id, pipeline, "Project", None);
774 below.then(Arc::new(Watched::new(project, counters)), schema)
775 }
776 Node::Aggregate { input, index, groups, aggregates } => {
777 let top_counts = self
778 .top_counts
779 .iter()
780 .find_map(|&(aggregate, bound)| (aggregate == reference).then_some(bound));
781 self.aggregate(
782 reference,
783 input,
784 index,
785 groups,
786 aggregates,
787 AggregateBound { max_groups: None, top_counts, having_count: None },
788 )?
789 }
790 Node::Sort { input, keys } => {
791 let below = self.node(input)?;
792 let schema = below.schema.clone();
793 let (sort, out) = Sort::new(plan, &schema, keys, memory)?;
794 let sort = sort.in_session(self.session);
795 let counters = self.watch(reference, id, pipeline, "Sort", None);
796 let reading = Arc::clone(&counters);
797 self.close(below, pipeline, Arc::new(Watched::new(sort, counters)));
798 Segment::reading(Arc::new(Watched::new(out, reading)), schema, pipeline)
799 }
800 Node::Limit { input, count, offset } => {
801 let max_groups = count
802 .and_then(|count| count.checked_add(offset))
803 .and_then(|count| usize::try_from(count).ok());
804 let below = match (plan.node(input).clone(), max_groups) {
805 (
806 Node::Aggregate { input: under, index, groups, aggregates },
807 Some(max_groups),
808 ) => self.aggregate(
809 input,
810 under,
811 index,
812 groups,
813 aggregates,
814 AggregateBound {
815 max_groups: Some(max_groups),
816 top_counts: None,
817 having_count: None,
818 },
819 )?,
820 _ => self.node(input)?,
821 };
822 let schema = below.schema.clone();
823 let limit = Limit::new(count, offset);
824 let counters = self.watch(reference, id, pipeline, "Limit", None);
825 below.then(Arc::new(Watched::new(limit, counters)), schema)
826 }
827 Node::TopN { input, keys, count, offset } => {
828 if let Some(aggregate) = count_top_aggregate(plan, input, keys) {
829 let bound = count.saturating_add(offset);
830 if let Ok(bound) = usize::try_from(bound) {
831 self.top_counts.push((aggregate, bound));
832 }
833 }
834 let below = self.node(input)?;
835 let schema = below.schema.clone();
836 let (top, out) = TopN::new(plan, &schema, keys, count, offset, memory)?;
837 let top = top.in_session(self.session);
838 let counters = self.watch(reference, id, pipeline, "TopN", None);
839 let reading = Arc::clone(&counters);
840 self.close(below, pipeline, Arc::new(Watched::new(top, counters)));
841 Segment::reading(Arc::new(Watched::new(out, reading)), schema, pipeline)
842 }
843 Node::Distinct { input, on } => {
844 let below = self.node(input)?;
845 let schema = below.schema.clone();
846 let (distinct, out) = Distinct::new(plan, &schema, on, memory)?;
847 let distinct = distinct.in_session(self.session);
848 let counters = self.watch(reference, id, pipeline, "Distinct", None);
849 let reading = Arc::clone(&counters);
850 self.close(below, pipeline, Arc::new(Watched::new(distinct, counters)));
851 Segment::reading(Arc::new(Watched::new(out, reading)), schema, pipeline)
852 }
853 Node::Join { left, right, kind, conditions } => {
854 let gather_id = self.gathered(reference);
860 let gathering = self.shape.pipeline(right);
861 let right = self.node(right)?;
862 let right_schema = right.schema.clone();
863 let (gather, gathered) = Gather::new(memory);
864 let kept = self.watch(reference, gather_id, gathering, "Gather", None);
865 self.close(right, gathering, Arc::new(Watched::new(gather, kept)));
866 let mut left = self.node(left)?;
867 let side = Gathered { schema: &right_schema, rows: gathered };
868 let (join, out) =
869 Join::new(plan, &left.schema, side, kind, conditions, self.cancel, memory);
870 let join = join.in_session(self.session);
871 let schema = join.schema().clone();
872 let counters = self.watch(reference, id, pipeline, "Join", None);
873 let reading = Arc::clone(&counters);
874 left.after.push(gathering);
875 self.close(left, pipeline, Arc::new(Watched::new(join, counters)));
876 Segment::reading(Arc::new(Watched::new(out, reading)), schema, pipeline)
877 }
878 Node::CrossProduct { left, right } => {
879 let keep_id = self.gathered(reference);
884 let aside = self.shape.pipeline(right);
885 let right = self.node(right)?;
886 let right_schema = right.schema.clone();
887 let (keep, kept) = Keep::new(memory);
888 let held = self.watch(reference, keep_id, aside, "Keep", None);
889 self.close(right, aside, Arc::new(Watched::new(keep, held)));
890 let mut left = self.node(left)?;
891 let cross = CrossProduct::new(&left.schema, &right_schema, kept);
892 let schema = cross.schema().clone();
893 let counters = self.watch(reference, id, pipeline, "CrossProduct", None);
894 left.after.push(aside);
895 left.then(Arc::new(Watched::new(cross, counters)), schema)
896 }
897 Node::SetOp { left, right, kind, all, index } => {
898 let gather_id = self.gathered(reference);
901 let counting = self.shape.pipeline(right);
902 let right = self.node(right)?;
903 let (gather, gathered) = Gather::new(memory);
904 let kept = self.watch(reference, gather_id, counting, "Gather", None);
905 self.close(right, counting, Arc::new(Watched::new(gather, kept)));
906 let mut left = self.node(left)?;
907 let (setop, out) = SetOp::new(&left.schema, gathered, kind, all, index, memory);
908 let schema = setop.schema().clone();
909 let counters = self.watch(reference, id, pipeline, "SetOp", None);
910 let reading = Arc::clone(&counters);
911 left.after.push(counting);
912 self.close(left, pipeline, Arc::new(Watched::new(setop, counters)));
913 Segment::reading(Arc::new(Watched::new(out, reading)), schema, pipeline)
914 }
915 };
916 Ok(segment)
917 }
918}
919
920#[cfg(test)]
921mod tests {
922 use rudb_common::LogicalType;
923 use rudb_plan::{CompareOp, Expr, Node, Plan};
924
925 use super::{count_having_aggregate, count_top_aggregate};
926
927 fn plan(direction: &str) -> Plan {
928 Plan::parse(&format!(
929 "TopN 10 offset 0 [#2.2::BIGINT {direction} NULLS LAST]\n \
930 Project #2 [#1.0::BIGINT AS WatchID, #1.1::INTEGER AS ClientIP, #1.2::BIGINT AS c]\n \
931 Aggregate #1 groups=[#0.0::BIGINT, #0.1::INTEGER] \
932 aggregates=[count_star()::BIGINT]\n \
933 Values #0 [WatchID::BIGINT, ClientIP::INTEGER] rows=[]"
934 ))
935 .expect("a grouped count plan")
936 }
937
938 #[test]
939 fn count_descending_topn_marks_its_aggregate() {
940 let plan = plan("DESC");
941 let Node::TopN { input, keys, .. } = *plan.node(plan.root()) else {
942 panic!("the root is a TopN")
943 };
944 let aggregate = count_top_aggregate(&plan, input, keys).expect("the grouped count");
945 assert!(matches!(plan.node(aggregate), Node::Aggregate { .. }));
946 }
947
948 #[test]
949 fn count_ascending_cannot_discard_large_counts() {
950 let plan = plan("ASC");
951 let Node::TopN { input, keys, .. } = *plan.node(plan.root()) else {
952 panic!("the root is a TopN")
953 };
954 assert!(count_top_aggregate(&plan, input, keys).is_none());
955 }
956
957 #[test]
958 fn distinct_count_descending_topn_marks_its_aggregate() {
959 let plan = Plan::parse(
960 "TopN 10 offset 0 [#1.1::BIGINT DESC NULLS LAST]\n \
961 Aggregate #1 groups=[#0.0::VARCHAR] \
962 aggregates=[count(DISTINCT #0.1::BIGINT)::BIGINT]\n \
963 Values #0 [SearchPhrase::VARCHAR, UserID::BIGINT] rows=[]",
964 )
965 .expect("a grouped distinct count plan");
966 let Node::TopN { input, keys, .. } = *plan.node(plan.root()) else {
967 panic!("the root is a TopN")
968 };
969 let aggregate = count_top_aggregate(&plan, input, keys).expect("the distinct count");
970 assert!(matches!(plan.node(aggregate), Node::Aggregate { .. }));
971 }
972
973 #[test]
974 fn a_count_having_lower_bound_marks_the_count_call() {
975 let plan = Plan::parse(
976 "Filter (#1.2::BIGINT > 100::BIGINT)::BOOLEAN\n \
977 Aggregate #1 groups=[#0.0::BIGINT] \
978 aggregates=[avg(#0.1::BIGINT)::DOUBLE, count_star()::BIGINT]\n \
979 Values #0 [key::BIGINT, value::BIGINT] rows=[]",
980 )
981 .expect("an aggregate with a HAVING filter");
982 let Node::Filter { input, predicate } = *plan.node(plan.root()) else {
983 panic!("the root is a Filter")
984 };
985 let (aggregate, call, minimum) =
986 count_having_aggregate(&plan, input, predicate).expect("the count bound");
987 assert_eq!(aggregate, input);
988 assert_eq!((call, minimum), (1, 101));
989 }
990
991 #[test]
992 fn an_upper_count_having_bound_cannot_drop_aggregate_output() {
993 let mut plan = Plan::parse(
994 "Filter (#1.1::BIGINT > 100::BIGINT)::BOOLEAN\n \
995 Aggregate #1 groups=[#0.0::BIGINT] aggregates=[count_star()::BIGINT]\n \
996 Values #0 [key::BIGINT] rows=[]",
997 )
998 .expect("an aggregate with a HAVING filter");
999 let Node::Filter { input, predicate } = *plan.node(plan.root()) else {
1000 panic!("the root is a Filter")
1001 };
1002 let Expr::Compare { left, right, .. } = *plan.expr(predicate) else {
1003 panic!("the predicate is a comparison")
1004 };
1005 let less =
1006 plan.add_expr(Expr::Compare { op: CompareOp::Less, left, right }, LogicalType::Boolean);
1007 assert!(count_having_aggregate(&plan, input, less).is_none());
1008 }
1009}