1use std::sync::Arc;
2
3use crate::executor::base::MemoryTracker;
4use crate::executor::streaming::operators::source_operator::OperatorConfig;
5use crate::executor::streaming::runtime::ExecutionRuntime;
6use crate::executor::streaming::slot::SlotLayout;
7use linkrs_core::error::QueryError;
8
9pub mod aggregate;
10pub mod aggregate_operator;
11pub mod helpers;
12pub mod materialize;
13pub mod materialize_operator;
14pub mod sort;
15pub mod sort_operator;
16pub mod window;
17pub mod window_operator;
18
19pub use aggregate::{AggregateState, FinalAggregateState, GroupByState, PartialAggregateState};
20pub use materialize::{DataCollectState, DistinctState, MaterializeState, RollUpApplyState};
21pub use sort::{MergeState, RunBuffer, SortState, TopNState};
22pub use window::{WindowFunctionState, WindowState};
23
24pub use crate::executor::streaming::chunk::ColumnarBatch;
25pub use crate::executor::streaming::executor::SortDirection;
26pub use linkrs_core::types::expr::Expression;
27pub use linkrs_core::Value;
28
29use crate::executor::streaming::chunk::DataChunk;
30use crate::executor::streaming::executor::StreamingExecutor;
31
32#[derive(Debug)]
33pub enum BlockingOperatorKind {
34 Sort {
35 sort_expressions: Vec<linkrs_core::types::expr::Expression>,
36 sort_directions: Vec<crate::executor::streaming::executor::SortDirection>,
37 memory_tracker: MemoryTracker,
38 state: Option<SortState>,
39 },
40 Aggregate {
41 group_by_expressions: Vec<linkrs_core::types::expr::Expression>,
42 aggregate_functions: Vec<(
43 linkrs_core::types::operators::AggregateFunction,
44 Vec<linkrs_core::types::expr::Expression>,
45 )>,
46 output_col_names: Vec<String>,
47 memory_tracker: MemoryTracker,
48 state: Option<AggregateState>,
49 },
50 GroupBy {
51 group_by_expressions: Vec<linkrs_core::types::expr::Expression>,
52 memory_tracker: MemoryTracker,
53 state: Option<GroupByState>,
54 },
55 WindowFunction {
56 window_exprs: Vec<linkrs_core::types::expr::Expression>,
57 partition_by_exprs: Vec<linkrs_core::types::expr::Expression>,
58 order_by_exprs: Vec<linkrs_core::types::expr::Expression>,
59 order_by_directions: Vec<crate::executor::streaming::executor::SortDirection>,
60 memory_tracker: MemoryTracker,
61 state: Option<WindowFunctionState>,
62 },
63 Window {
64 window_exprs: Vec<linkrs_core::types::expr::Expression>,
65 partition_by_exprs: Vec<linkrs_core::types::expr::Expression>,
66 order_by_exprs: Vec<linkrs_core::types::expr::Expression>,
67 order_by_directions: Vec<crate::executor::streaming::executor::SortDirection>,
68 memory_tracker: MemoryTracker,
69 state: Option<WindowState>,
70 },
71 TopN {
72 n: u32,
73 sort_expressions: Vec<linkrs_core::types::expr::Expression>,
74 sort_directions: Vec<crate::executor::streaming::executor::SortDirection>,
75 memory_tracker: MemoryTracker,
76 state: Option<TopNState>,
77 },
78 Distinct {
79 memory_tracker: MemoryTracker,
80 state: Option<DistinctState>,
81 },
82 Materialize {
83 memory_tracker: MemoryTracker,
84 state: Option<MaterializeState>,
85 },
86 DataCollect {
87 memory_tracker: MemoryTracker,
88 state: Option<DataCollectState>,
89 },
90 RollUpApply {
91 rollup_expressions: Vec<linkrs_core::types::expr::Expression>,
92 memory_tracker: MemoryTracker,
93 state: Option<RollUpApplyState>,
94 },
95 PartialAggregate {
96 group_by_expressions: Vec<linkrs_core::types::expr::Expression>,
97 aggregate_functions: Vec<(
98 linkrs_core::types::operators::AggregateFunction,
99 Vec<linkrs_core::types::expr::Expression>,
100 )>,
101 output_col_names: Vec<String>,
102 memory_tracker: MemoryTracker,
103 state: Option<PartialAggregateState>,
104 },
105 FinalAggregate {
106 group_by_expressions: Vec<linkrs_core::types::expr::Expression>,
107 aggregate_functions: Vec<(
108 linkrs_core::types::operators::AggregateFunction,
109 Vec<linkrs_core::types::expr::Expression>,
110 )>,
111 output_col_names: Vec<String>,
112 memory_tracker: MemoryTracker,
113 state: Option<FinalAggregateState>,
114 },
115}
116
117#[derive(Debug)]
118pub struct BlockingOperator {
119 pub kind: BlockingOperatorKind,
120 pub runtime: Option<Arc<ExecutionRuntime>>,
121 pub output_layout: Arc<SlotLayout>,
122 pub config: OperatorConfig,
123}
124
125impl BlockingOperator {
126 pub fn from_spec(
127 spec: &super::spec::BlockingSpec,
128 memory_budget: &crate::executor::base::MemoryBudget,
129 output_layout: Arc<SlotLayout>,
130 ) -> Self {
131 let kind = match spec {
132 super::spec::BlockingSpec::Sort {
133 sort_expressions,
134 sort_directions,
135 } => BlockingOperatorKind::Sort {
136 sort_expressions: sort_expressions.clone(),
137 sort_directions: sort_directions.clone(),
138 memory_tracker: crate::executor::base::MemoryTracker::new(memory_budget.clone()),
139 state: None,
140 },
141 super::spec::BlockingSpec::Aggregate {
142 group_by_expressions,
143 aggregate_functions,
144 output_col_names,
145 } => BlockingOperatorKind::Aggregate {
146 group_by_expressions: group_by_expressions.clone(),
147 aggregate_functions: aggregate_functions.clone(),
148 output_col_names: output_col_names.clone(),
149 memory_tracker: crate::executor::base::MemoryTracker::new(memory_budget.clone()),
150 state: None,
151 },
152 super::spec::BlockingSpec::GroupBy {
153 group_by_expressions,
154 } => BlockingOperatorKind::GroupBy {
155 group_by_expressions: group_by_expressions.clone(),
156 memory_tracker: crate::executor::base::MemoryTracker::new(memory_budget.clone()),
157 state: None,
158 },
159 super::spec::BlockingSpec::WindowFunction {
160 window_exprs,
161 partition_by_exprs,
162 order_by_exprs,
163 order_by_directions,
164 } => BlockingOperatorKind::WindowFunction {
165 window_exprs: window_exprs.clone(),
166 partition_by_exprs: partition_by_exprs.clone(),
167 order_by_exprs: order_by_exprs.clone(),
168 order_by_directions: order_by_directions.clone(),
169 memory_tracker: crate::executor::base::MemoryTracker::new(memory_budget.clone()),
170 state: None,
171 },
172 super::spec::BlockingSpec::Window {
173 window_exprs,
174 partition_by_exprs,
175 order_by_exprs,
176 order_by_directions,
177 } => BlockingOperatorKind::Window {
178 window_exprs: window_exprs.clone(),
179 partition_by_exprs: partition_by_exprs.clone(),
180 order_by_exprs: order_by_exprs.clone(),
181 order_by_directions: order_by_directions.clone(),
182 memory_tracker: crate::executor::base::MemoryTracker::new(memory_budget.clone()),
183 state: None,
184 },
185 super::spec::BlockingSpec::TopN {
186 n,
187 sort_expressions,
188 sort_directions,
189 } => BlockingOperatorKind::TopN {
190 n: *n,
191 sort_expressions: sort_expressions.clone(),
192 sort_directions: sort_directions.clone(),
193 memory_tracker: crate::executor::base::MemoryTracker::new(memory_budget.clone()),
194 state: None,
195 },
196 super::spec::BlockingSpec::Distinct => BlockingOperatorKind::Distinct {
197 memory_tracker: crate::executor::base::MemoryTracker::new(memory_budget.clone()),
198 state: None,
199 },
200 super::spec::BlockingSpec::Materialize => BlockingOperatorKind::Materialize {
201 memory_tracker: crate::executor::base::MemoryTracker::new(memory_budget.clone()),
202 state: None,
203 },
204 super::spec::BlockingSpec::DataCollect => BlockingOperatorKind::DataCollect {
205 memory_tracker: crate::executor::base::MemoryTracker::new(memory_budget.clone()),
206 state: None,
207 },
208 super::spec::BlockingSpec::RollUpApply { rollup_expressions } => {
209 BlockingOperatorKind::RollUpApply {
210 rollup_expressions: rollup_expressions.clone(),
211 memory_tracker: crate::executor::base::MemoryTracker::new(
212 memory_budget.clone(),
213 ),
214 state: None,
215 }
216 }
217 super::spec::BlockingSpec::PartialAggregate {
218 group_by_expressions,
219 aggregate_functions,
220 output_col_names,
221 } => BlockingOperatorKind::PartialAggregate {
222 group_by_expressions: group_by_expressions.clone(),
223 aggregate_functions: aggregate_functions.clone(),
224 output_col_names: output_col_names.clone(),
225 memory_tracker: crate::executor::base::MemoryTracker::new(memory_budget.clone()),
226 state: None,
227 },
228 super::spec::BlockingSpec::FinalAggregate {
229 group_by_expressions,
230 aggregate_functions,
231 output_col_names,
232 } => BlockingOperatorKind::FinalAggregate {
233 group_by_expressions: group_by_expressions.clone(),
234 aggregate_functions: aggregate_functions.clone(),
235 output_col_names: output_col_names.clone(),
236 memory_tracker: crate::executor::base::MemoryTracker::new(memory_budget.clone()),
237 state: None,
238 },
239 };
240 Self::new(kind, output_layout)
241 }
242
243 pub fn new(kind: BlockingOperatorKind, output_layout: Arc<SlotLayout>) -> Self {
244 Self {
245 kind,
246 runtime: None,
247 output_layout,
248 config: OperatorConfig::default(),
249 }
250 }
251
252 pub fn inject_context(
253 &mut self,
254 runtime: Option<&Arc<ExecutionRuntime>>,
255 config: OperatorConfig,
256 ) {
257 if let Some(rt) = runtime {
258 self.runtime = Some(rt.clone());
259 }
260 self.config = config;
261 }
262
263 pub fn memory_tracker(&self) -> &MemoryTracker {
264 match &self.kind {
265 BlockingOperatorKind::Sort { memory_tracker, .. }
266 | BlockingOperatorKind::Aggregate { memory_tracker, .. }
267 | BlockingOperatorKind::GroupBy { memory_tracker, .. }
268 | BlockingOperatorKind::WindowFunction { memory_tracker, .. }
269 | BlockingOperatorKind::Window { memory_tracker, .. }
270 | BlockingOperatorKind::TopN { memory_tracker, .. }
271 | BlockingOperatorKind::Distinct { memory_tracker, .. }
272 | BlockingOperatorKind::Materialize { memory_tracker, .. }
273 | BlockingOperatorKind::DataCollect { memory_tracker, .. }
274 | BlockingOperatorKind::RollUpApply { memory_tracker, .. }
275 | BlockingOperatorKind::PartialAggregate { memory_tracker, .. }
276 | BlockingOperatorKind::FinalAggregate { memory_tracker, .. } => memory_tracker,
277 }
278 }
279
280 pub fn open(&mut self, input: &mut StreamingExecutor) -> Result<(), QueryError> {
281 match &mut self.kind {
282 BlockingOperatorKind::Sort { state, .. } => {
283 sort_operator::open_sort(state);
284 }
285 BlockingOperatorKind::Aggregate {
286 state,
287 aggregate_functions,
288 ..
289 } => {
290 aggregate_operator::open_aggregate(state, aggregate_functions.len());
291 }
292 BlockingOperatorKind::GroupBy { state, .. } => {
293 aggregate_operator::open_groupby(state);
294 }
295 BlockingOperatorKind::WindowFunction { state, .. } => {
296 window_operator::open_window_function(state);
297 }
298 BlockingOperatorKind::Window { state, .. } => {
299 window_operator::open_window(state);
300 }
301 BlockingOperatorKind::TopN { state, .. } => {
302 sort_operator::open_topn(state);
303 }
304 BlockingOperatorKind::Distinct { state, .. } => {
305 materialize_operator::open_distinct(state);
306 }
307 BlockingOperatorKind::Materialize { state, .. } => {
308 materialize_operator::open_materialize(state);
309 }
310 BlockingOperatorKind::DataCollect { state, .. } => {
311 materialize_operator::open_data_collect(state);
312 }
313 BlockingOperatorKind::RollUpApply { state, .. } => {
314 materialize_operator::open_rollup_apply(state);
315 }
316 BlockingOperatorKind::PartialAggregate { state, .. } => {
317 aggregate_operator::open_partial_aggregate(state);
318 }
319 BlockingOperatorKind::FinalAggregate { state, .. } => {
320 aggregate_operator::open_final_aggregate(state);
321 }
322 }
323 input.open()?;
324 Ok(())
325 }
326
327 pub fn next(&mut self, input: &mut StreamingExecutor) -> Result<Option<DataChunk>, QueryError> {
328 let ctx = helpers::BlockingContext {
329 runtime: &self.runtime,
330 output_layout: &self.output_layout,
331 config: &self.config,
332 };
333
334 match &mut self.kind {
335 BlockingOperatorKind::Sort {
336 sort_expressions,
337 sort_directions,
338 memory_tracker,
339 state,
340 ..
341 } => sort_operator::next_sort(
342 sort_expressions,
343 sort_directions,
344 memory_tracker,
345 state.as_mut().unwrap(),
346 &ctx,
347 input,
348 ),
349 BlockingOperatorKind::Aggregate {
350 group_by_expressions,
351 aggregate_functions,
352 memory_tracker,
353 state,
354 ..
355 } => aggregate_operator::next_aggregate(
356 group_by_expressions,
357 aggregate_functions,
358 memory_tracker,
359 state.as_mut().unwrap(),
360 &ctx,
361 input,
362 ),
363 BlockingOperatorKind::GroupBy {
364 group_by_expressions,
365 memory_tracker,
366 state,
367 ..
368 } => aggregate_operator::next_groupby(
369 group_by_expressions,
370 memory_tracker,
371 state.as_mut().unwrap(),
372 &ctx,
373 input,
374 ),
375 BlockingOperatorKind::WindowFunction {
376 window_exprs,
377 partition_by_exprs,
378 order_by_exprs,
379 order_by_directions,
380 memory_tracker,
381 state,
382 ..
383 } => window_operator::next_window_function(
384 window_exprs,
385 partition_by_exprs,
386 order_by_exprs,
387 order_by_directions,
388 memory_tracker,
389 state.as_mut().unwrap(),
390 &ctx,
391 input,
392 ),
393 BlockingOperatorKind::Window {
394 window_exprs,
395 partition_by_exprs,
396 order_by_exprs,
397 order_by_directions,
398 memory_tracker,
399 state,
400 ..
401 } => window_operator::next_window(
402 window_exprs,
403 partition_by_exprs,
404 order_by_exprs,
405 order_by_directions,
406 memory_tracker,
407 state.as_mut().unwrap(),
408 &ctx,
409 input,
410 ),
411 BlockingOperatorKind::TopN {
412 n,
413 sort_expressions,
414 sort_directions,
415 memory_tracker,
416 state,
417 ..
418 } => sort_operator::next_topn(
419 *n,
420 sort_expressions,
421 sort_directions,
422 memory_tracker,
423 state.as_mut().unwrap(),
424 &ctx,
425 input,
426 ),
427 BlockingOperatorKind::Distinct {
428 memory_tracker,
429 state,
430 ..
431 } => materialize_operator::next_distinct(
432 memory_tracker,
433 state.as_mut().unwrap(),
434 &ctx,
435 input,
436 ),
437 BlockingOperatorKind::Materialize {
438 memory_tracker,
439 state,
440 ..
441 } => materialize_operator::next_materialize(
442 memory_tracker,
443 state.as_mut().unwrap(),
444 &ctx,
445 input,
446 ),
447 BlockingOperatorKind::DataCollect {
448 memory_tracker,
449 state,
450 ..
451 } => materialize_operator::next_data_collect(
452 memory_tracker,
453 state.as_mut().unwrap(),
454 &ctx,
455 input,
456 ),
457 BlockingOperatorKind::RollUpApply {
458 rollup_expressions,
459 memory_tracker,
460 state,
461 ..
462 } => materialize_operator::next_rollup_apply(
463 rollup_expressions,
464 memory_tracker,
465 state.as_mut().unwrap(),
466 &ctx,
467 input,
468 ),
469 BlockingOperatorKind::PartialAggregate {
470 group_by_expressions,
471 aggregate_functions,
472 memory_tracker,
473 state,
474 ..
475 } => aggregate_operator::next_partial_aggregate(
476 group_by_expressions,
477 aggregate_functions,
478 memory_tracker,
479 state.as_mut().unwrap(),
480 &ctx,
481 input,
482 ),
483 BlockingOperatorKind::FinalAggregate {
484 group_by_expressions,
485 aggregate_functions,
486 memory_tracker,
487 state,
488 ..
489 } => aggregate_operator::next_final_aggregate(
490 group_by_expressions,
491 aggregate_functions,
492 memory_tracker,
493 state.as_mut().unwrap(),
494 &ctx,
495 input,
496 ),
497 }
498 }
499
500 pub fn stop(&mut self) -> Result<(), QueryError> {
501 Ok(())
502 }
503
504 pub fn close(&mut self) -> Result<(), QueryError> {
505 match &mut self.kind {
506 BlockingOperatorKind::Sort {
507 state,
508 memory_tracker,
509 ..
510 } => {
511 sort_operator::close_sort(state);
512 memory_tracker.reset();
513 Ok(())
514 }
515 BlockingOperatorKind::Aggregate {
516 state,
517 memory_tracker,
518 ..
519 } => {
520 aggregate_operator::close_aggregate(state);
521 memory_tracker.reset();
522 Ok(())
523 }
524 BlockingOperatorKind::GroupBy {
525 state,
526 memory_tracker,
527 ..
528 } => {
529 aggregate_operator::close_groupby(state);
530 memory_tracker.reset();
531 Ok(())
532 }
533 BlockingOperatorKind::WindowFunction {
534 state,
535 memory_tracker,
536 ..
537 } => {
538 window_operator::close_window_function(state);
539 memory_tracker.reset();
540 Ok(())
541 }
542 BlockingOperatorKind::Window {
543 state,
544 memory_tracker,
545 ..
546 } => {
547 window_operator::close_window(state);
548 memory_tracker.reset();
549 Ok(())
550 }
551 BlockingOperatorKind::TopN {
552 state,
553 memory_tracker,
554 ..
555 } => {
556 sort_operator::close_topn(state);
557 memory_tracker.reset();
558 Ok(())
559 }
560 BlockingOperatorKind::RollUpApply {
561 state,
562 memory_tracker,
563 ..
564 } => {
565 materialize_operator::close_rollup_apply(state);
566 memory_tracker.reset();
567 Ok(())
568 }
569 BlockingOperatorKind::Distinct {
570 state,
571 memory_tracker,
572 ..
573 } => {
574 materialize_operator::close_distinct(state);
575 memory_tracker.reset();
576 Ok(())
577 }
578 BlockingOperatorKind::PartialAggregate {
579 state,
580 memory_tracker,
581 ..
582 } => {
583 aggregate_operator::close_partial_aggregate(state);
584 memory_tracker.reset();
585 Ok(())
586 }
587 BlockingOperatorKind::FinalAggregate {
588 state,
589 memory_tracker,
590 ..
591 } => {
592 aggregate_operator::close_final_aggregate(state);
593 memory_tracker.reset();
594 Ok(())
595 }
596 BlockingOperatorKind::Materialize {
597 state,
598 memory_tracker,
599 ..
600 } => {
601 materialize_operator::close_materialize(state);
602 memory_tracker.reset();
603 Ok(())
604 }
605 BlockingOperatorKind::DataCollect {
606 state,
607 memory_tracker,
608 ..
609 } => {
610 materialize_operator::close_data_collect(state);
611 memory_tracker.reset();
612 Ok(())
613 }
614 }
615 }
616
617 pub fn spill_with_manager(
618 &mut self,
619 sm: &crate::executor::streaming::spill::SpillManager,
620 ) -> Result<(), QueryError> {
621 match &mut self.kind {
622 BlockingOperatorKind::Sort {
623 state,
624 memory_tracker,
625 sort_expressions,
626 sort_directions,
627 ..
628 } => sort_operator::spill_sort(
629 state.as_mut().unwrap(),
630 memory_tracker,
631 sort_expressions,
632 sort_directions,
633 sm,
634 ),
635 BlockingOperatorKind::Aggregate {
636 state,
637 memory_tracker,
638 ..
639 } => aggregate_operator::spill_aggregate(state.as_mut().unwrap(), memory_tracker, sm),
640 BlockingOperatorKind::GroupBy {
641 state,
642 memory_tracker,
643 group_by_expressions,
644 ..
645 } => aggregate_operator::spill_groupby(
646 state.as_mut().unwrap(),
647 memory_tracker,
648 group_by_expressions,
649 sm,
650 ),
651 BlockingOperatorKind::WindowFunction {
652 state,
653 memory_tracker,
654 partition_by_exprs,
655 ..
656 } => window_operator::spill_window_function(
657 state.as_mut().unwrap(),
658 memory_tracker,
659 partition_by_exprs,
660 sm,
661 ),
662 BlockingOperatorKind::Window {
663 state,
664 memory_tracker,
665 partition_by_exprs,
666 ..
667 } => window_operator::spill_window(
668 state.as_mut().unwrap(),
669 memory_tracker,
670 partition_by_exprs,
671 sm,
672 ),
673 BlockingOperatorKind::TopN { .. } => Ok(()),
674 BlockingOperatorKind::Distinct {
675 state,
676 memory_tracker,
677 ..
678 } => materialize_operator::spill_distinct(state.as_mut().unwrap(), memory_tracker, sm),
679 BlockingOperatorKind::Materialize {
680 state,
681 memory_tracker,
682 ..
683 } => {
684 materialize_operator::spill_materialize(state.as_mut().unwrap(), sm, memory_tracker)
685 }
686 BlockingOperatorKind::DataCollect {
687 state,
688 memory_tracker,
689 ..
690 } => materialize_operator::spill_data_collect(
691 state.as_mut().unwrap(),
692 sm,
693 memory_tracker,
694 ),
695 BlockingOperatorKind::PartialAggregate { state, .. } => {
696 if state.as_ref().is_some_and(|s| !s.group_map.is_empty()) {
697 return Err(QueryError::execution(
698 "Query memory budget exceeded during partial aggregation. Reduce GROUP BY cardinality (add LIMIT or narrower filters) or raise the query memory budget".to_string(),
699 ));
700 }
701 Ok(())
702 }
703 BlockingOperatorKind::RollUpApply {
704 state,
705 memory_tracker,
706 ..
707 } => materialize_operator::spill_rollup_apply(
708 state.as_mut().unwrap(),
709 sm,
710 memory_tracker,
711 ),
712 BlockingOperatorKind::FinalAggregate { state, .. } => {
713 if state.as_ref().is_some_and(|s| !s.group_map.is_empty()) {
714 return Err(QueryError::execution(
715 "Query memory budget exceeded during final aggregation. Reduce GROUP BY cardinality (add LIMIT or narrower filters) or raise the query memory budget".to_string(),
716 ));
717 }
718 Ok(())
719 }
720 }
721 }
722
723 pub fn spill_count(&self) -> u64 {
724 match &self.kind {
725 BlockingOperatorKind::Sort { state, .. } => {
726 state.as_ref().map_or(0, |s| s.runs.len() as u64)
727 }
728 BlockingOperatorKind::Aggregate { state, .. } => state.as_ref().map_or(0, |s| {
729 s.spilled_runs.iter().filter_map(|r| r.as_ref()).count() as u64
730 }),
731 BlockingOperatorKind::GroupBy { state, .. } => state.as_ref().map_or(0, |s| {
732 s.spilled_runs.iter().filter_map(|r| r.as_ref()).count() as u64
733 }),
734 BlockingOperatorKind::WindowFunction { state, .. } => state.as_ref().map_or(0, |s| {
735 s.spilled_runs.iter().filter_map(|r| r.as_ref()).count() as u64
736 }),
737 BlockingOperatorKind::Window { state, .. } => state.as_ref().map_or(0, |s| {
738 s.spilled_runs.iter().filter_map(|r| r.as_ref()).count() as u64
739 }),
740 BlockingOperatorKind::Distinct { state, .. } => state.as_ref().map_or(0, |s| {
741 s.spilled_runs.iter().filter_map(|r| r.as_ref()).count() as u64
742 }),
743 _ => 0,
744 }
745 }
746
747 pub fn spilled_bytes(&self) -> u64 {
748 macro_rules! sum_runs {
749 ($state:expr) => {
750 $state.as_ref().map_or(0, |s| {
751 s.spilled_runs
752 .iter()
753 .filter_map(|r| r.as_ref())
754 .map(|r| r.byte_size)
755 .sum::<u64>()
756 })
757 };
758 }
759 match &self.kind {
760 BlockingOperatorKind::Sort { state, .. } => state
761 .as_ref()
762 .map_or(0, |s| s.runs.iter().map(|r| r.byte_size).sum::<u64>()),
763 BlockingOperatorKind::Aggregate { state, .. } => sum_runs!(state),
764 BlockingOperatorKind::GroupBy { state, .. } => sum_runs!(state),
765 BlockingOperatorKind::WindowFunction { state, .. } => sum_runs!(state),
766 BlockingOperatorKind::Window { state, .. } => sum_runs!(state),
767 BlockingOperatorKind::TopN { .. } => 0,
768 BlockingOperatorKind::Distinct { state, .. } => sum_runs!(state),
769 BlockingOperatorKind::Materialize { .. } => 0,
770 BlockingOperatorKind::DataCollect { .. } => 0,
771 BlockingOperatorKind::RollUpApply { .. } => 0,
772 BlockingOperatorKind::PartialAggregate { .. } => 0,
773 BlockingOperatorKind::FinalAggregate { .. } => 0,
774 }
775 }
776
777 pub fn spilled_rows(&self) -> u64 {
778 macro_rules! sum_rows {
779 ($state:expr) => {
780 $state.as_ref().map_or(0, |s| {
781 s.spilled_runs
782 .iter()
783 .filter_map(|r| r.as_ref())
784 .map(|r| r.row_count)
785 .sum::<u64>()
786 })
787 };
788 }
789 match &self.kind {
790 BlockingOperatorKind::Sort { state, .. } => state
791 .as_ref()
792 .map_or(0, |s| s.runs.iter().map(|r| r.row_count).sum::<u64>()),
793 BlockingOperatorKind::Aggregate { state, .. } => sum_rows!(state),
794 BlockingOperatorKind::GroupBy { state, .. } => sum_rows!(state),
795 BlockingOperatorKind::WindowFunction { state, .. } => sum_rows!(state),
796 BlockingOperatorKind::Window { state, .. } => sum_rows!(state),
797 BlockingOperatorKind::TopN { .. } => 0,
798 BlockingOperatorKind::Distinct { state, .. } => sum_rows!(state),
799 BlockingOperatorKind::Materialize { .. } => 0,
800 BlockingOperatorKind::DataCollect { .. } => 0,
801 BlockingOperatorKind::RollUpApply { .. } => 0,
802 BlockingOperatorKind::PartialAggregate { .. } => 0,
803 BlockingOperatorKind::FinalAggregate { .. } => 0,
804 }
805 }
806}
807
808#[cfg(test)]
809mod test;