1use crate::execution::TruncationReason;
8
9#[derive(Debug, Clone, Default)]
14pub enum StopCondition {
15 MaxRows(u64),
17 MaxBytes(u64),
27 SchemaStable {
30 consecutive_stable_rows: u64,
34 },
35 ConfidenceThreshold(f64),
41 MemoryPressure(f64),
45 Any(Vec<StopCondition>),
47 All(Vec<StopCondition>),
49 #[default]
51 Never,
52}
53
54impl StopCondition {
55 pub fn schema_inference() -> Self {
58 StopCondition::Any(vec![
59 StopCondition::MaxRows(10_000),
60 StopCondition::SchemaStable {
61 consecutive_stable_rows: 1_000,
62 },
63 ])
64 }
65
66 pub fn quality_sample() -> Self {
69 StopCondition::Any(vec![
70 StopCondition::MaxRows(50_000),
71 StopCondition::MaxBytes(50 * 1024 * 1024),
72 StopCondition::ConfidenceThreshold(0.95),
73 ])
74 }
75
76 pub fn max_rows(&self) -> Option<u64> {
91 match self {
92 StopCondition::MaxRows(n) => Some(*n),
93 StopCondition::Any(conditions) => {
95 conditions.iter().filter_map(StopCondition::max_rows).min()
96 }
97 StopCondition::All(conditions) => {
100 if conditions.is_empty() {
101 return None; }
103 conditions
104 .iter()
105 .map(StopCondition::max_rows)
106 .try_fold(0u64, |acc, cap| Some(acc.max(cap?)))
107 }
108 _ => None,
109 }
110 }
111
112 pub fn max_bytes(&self) -> Option<u64> {
120 match self {
121 StopCondition::MaxBytes(n) => Some(*n),
122 StopCondition::Any(conditions) => {
123 conditions.iter().filter_map(StopCondition::max_bytes).min()
124 }
125 StopCondition::All(conditions) => {
126 if conditions.is_empty() {
127 return None;
128 }
129 conditions
130 .iter()
131 .map(StopCondition::max_bytes)
132 .try_fold(0u64, |acc, cap| Some(acc.max(cap?)))
133 }
134 _ => None,
135 }
136 }
137
138 pub fn is_row_limit_only(&self) -> bool {
145 match self {
146 StopCondition::Never | StopCondition::MaxRows(_) => true,
147 StopCondition::Any(conditions) | StopCondition::All(conditions) => {
148 conditions.iter().all(StopCondition::is_row_limit_only)
149 }
150 _ => false,
151 }
152 }
153
154 pub fn needs_row_count(&self) -> bool {
156 match self {
157 StopCondition::ConfidenceThreshold(_) => true,
158 StopCondition::Any(conditions) | StopCondition::All(conditions) => {
159 conditions.iter().any(StopCondition::needs_row_count)
160 }
161 _ => false,
162 }
163 }
164}
165
166pub struct StopEvaluator {
171 condition: StopCondition,
172 rows_processed: u64,
173 bytes_consumed: u64,
174 estimated_total_rows: Option<u64>,
175 schema_stable_rows: Option<u64>,
176 triggered_reason: Option<TruncationReason>,
177}
178
179impl StopEvaluator {
180 pub fn new(condition: StopCondition) -> Self {
181 let condition = Self::clamp_thresholds(condition);
182 Self {
183 condition,
184 rows_processed: 0,
185 bytes_consumed: 0,
186 estimated_total_rows: None,
187 schema_stable_rows: None,
188 triggered_reason: None,
189 }
190 }
191
192 fn clamp_thresholds(condition: StopCondition) -> StopCondition {
195 match condition {
196 StopCondition::ConfidenceThreshold(t) => {
197 StopCondition::ConfidenceThreshold(t.clamp(0.0, 1.0))
198 }
199 StopCondition::MemoryPressure(t) => StopCondition::MemoryPressure(t.clamp(0.0, 1.0)),
200 StopCondition::Any(cs) => {
201 StopCondition::Any(cs.into_iter().map(Self::clamp_thresholds).collect())
202 }
203 StopCondition::All(cs) => {
204 StopCondition::All(cs.into_iter().map(Self::clamp_thresholds).collect())
205 }
206 other => other,
207 }
208 }
209
210 pub fn with_estimated_total(mut self, rows: u64) -> Self {
212 self.estimated_total_rows = Some(rows);
213 self
214 }
215
216 pub fn set_estimated_total(&mut self, rows: u64) {
218 self.estimated_total_rows = Some(rows);
219 }
220
221 pub fn set_schema_stable_rows(&mut self, rows: u64) {
224 self.schema_stable_rows = Some(rows);
225 }
226
227 pub fn update(&mut self, chunk_rows: u64, chunk_bytes: u64, memory_fraction: f64) -> bool {
235 self.rows_processed += chunk_rows;
236 self.bytes_consumed += chunk_bytes;
237
238 if self.triggered_reason.is_some() {
239 return true;
240 }
241
242 let reason = evaluate(
243 &self.condition,
244 self.rows_processed,
245 self.bytes_consumed,
246 memory_fraction,
247 self.estimated_total_rows,
248 self.schema_stable_rows,
249 );
250
251 if reason.is_some() {
252 self.triggered_reason = reason;
253 true
254 } else {
255 false
256 }
257 }
258
259 pub fn should_stop(&self) -> bool {
261 self.triggered_reason.is_some()
262 }
263
264 pub fn truncation_reason(&self) -> Option<TruncationReason> {
266 self.triggered_reason.clone()
267 }
268
269 pub fn rows_processed(&self) -> u64 {
271 self.rows_processed
272 }
273
274 pub fn bytes_consumed(&self) -> u64 {
276 self.bytes_consumed
277 }
278}
279
280fn evaluate(
283 condition: &StopCondition,
284 rows: u64,
285 bytes: u64,
286 memory_fraction: f64,
287 estimated_total: Option<u64>,
288 schema_stable_rows: Option<u64>,
289) -> Option<TruncationReason> {
290 match condition {
291 StopCondition::MaxRows(limit) => {
292 if rows >= *limit {
293 Some(TruncationReason::MaxRows(*limit))
294 } else {
295 None
296 }
297 }
298 StopCondition::MaxBytes(limit) => {
299 if bytes >= *limit {
300 Some(TruncationReason::MaxBytes(*limit))
301 } else {
302 None
303 }
304 }
305 StopCondition::SchemaStable {
306 consecutive_stable_rows,
307 } => {
308 if schema_stable_rows.is_some_and(|rows| rows >= *consecutive_stable_rows) {
309 Some(TruncationReason::StopCondition(format!(
310 "schema_stable({consecutive_stable_rows})"
311 )))
312 } else {
313 None
314 }
315 }
316 StopCondition::ConfidenceThreshold(threshold) => {
317 if let Some(total) = estimated_total
318 && total > 0
319 {
320 let confidence = rows as f64 / total as f64;
321 if confidence >= *threshold {
322 return Some(TruncationReason::StopCondition(format!(
323 "confidence_threshold({})",
324 threshold
325 )));
326 }
327 }
328 None
329 }
330 StopCondition::MemoryPressure(threshold) => {
331 if memory_fraction >= *threshold {
332 Some(TruncationReason::MemoryPressure)
333 } else {
334 None
335 }
336 }
337 StopCondition::Any(conditions) => {
338 for c in conditions {
339 if let Some(reason) = evaluate(
340 c,
341 rows,
342 bytes,
343 memory_fraction,
344 estimated_total,
345 schema_stable_rows,
346 ) {
347 return Some(reason);
348 }
349 }
350 None
351 }
352 StopCondition::All(conditions) => {
353 if conditions.is_empty() {
354 return None;
355 }
356 let mut first_reason = None;
358 for c in conditions {
359 let reason = evaluate(
360 c,
361 rows,
362 bytes,
363 memory_fraction,
364 estimated_total,
365 schema_stable_rows,
366 )?;
367 if first_reason.is_none() {
368 first_reason = Some(reason);
369 }
370 }
371 first_reason
372 }
373 StopCondition::Never => None,
374 }
375}
376
377pub fn schema_stable_threshold(condition: &StopCondition) -> Option<u64> {
381 match condition {
382 StopCondition::SchemaStable {
383 consecutive_stable_rows,
384 } => Some(*consecutive_stable_rows),
385 StopCondition::Any(conditions) | StopCondition::All(conditions) => {
386 conditions.iter().find_map(schema_stable_threshold)
387 }
388 _ => None,
389 }
390}
391
392pub struct SchemaStabilityTracker {
399 threshold: u64,
400 consecutive_stable: u64,
401 last_fingerprint: Option<u64>,
402}
403
404impl SchemaStabilityTracker {
405 pub fn from_condition(condition: &StopCondition) -> Option<Self> {
408 schema_stable_threshold(condition).map(|threshold| Self {
409 threshold,
410 consecutive_stable: 0,
411 last_fingerprint: None,
412 })
413 }
414
415 pub fn update(&mut self, fingerprint: u64, chunk_rows: u64) -> bool {
419 match self.last_fingerprint {
420 Some(prev) if prev == fingerprint => {
421 self.consecutive_stable += chunk_rows;
422 }
423 _ => {
424 self.consecutive_stable = 0;
428 self.last_fingerprint = Some(fingerprint);
429 }
430 }
431 self.consecutive_stable >= self.threshold
432 }
433
434 pub fn consecutive_stable_rows(&self) -> u64 {
436 self.consecutive_stable
437 }
438
439 pub fn truncation_reason(&self) -> TruncationReason {
441 TruncationReason::StopCondition(format!("schema_stable({})", self.threshold))
442 }
443}
444
445#[cfg(test)]
446mod tests {
447 use super::*;
448
449 #[test]
450 fn max_bytes_follows_the_any_and_all_rules() {
451 assert_eq!(StopCondition::MaxBytes(10).max_bytes(), Some(10));
452 assert_eq!(StopCondition::MaxRows(10).max_bytes(), None);
453 let any = StopCondition::Any(vec![
454 StopCondition::MaxBytes(20),
455 StopCondition::MaxRows(5),
456 StopCondition::MaxBytes(10),
457 ]);
458 assert_eq!(any.max_bytes(), Some(10));
459 let all = StopCondition::All(vec![
460 StopCondition::MaxBytes(10),
461 StopCondition::MaxBytes(20),
462 ]);
463 assert_eq!(all.max_bytes(), Some(20));
464 let mixed =
467 StopCondition::All(vec![StopCondition::MaxBytes(10), StopCondition::MaxRows(5)]);
468 assert_eq!(mixed.max_bytes(), None);
469 assert_eq!(StopCondition::All(vec![]).max_bytes(), None);
470 }
471
472 #[test]
473 fn test_max_rows_leaf_and_never() {
474 assert_eq!(StopCondition::MaxRows(10).max_rows(), Some(10));
475 assert_eq!(StopCondition::Never.max_rows(), None);
476 assert_eq!(StopCondition::MaxBytes(10).max_rows(), None);
477 }
478
479 #[test]
480 fn test_max_rows_any_takes_the_earliest_cap() {
481 let condition =
483 StopCondition::Any(vec![StopCondition::MaxRows(20), StopCondition::MaxRows(10)]);
484 assert_eq!(condition.max_rows(), Some(10));
485
486 let flipped =
488 StopCondition::Any(vec![StopCondition::MaxRows(10), StopCondition::MaxRows(20)]);
489 assert_eq!(flipped.max_rows(), Some(10));
490
491 let with_never = StopCondition::Any(vec![StopCondition::Never, StopCondition::MaxRows(10)]);
493 assert_eq!(with_never.max_rows(), Some(10));
494
495 let mixed = StopCondition::Any(vec![
497 StopCondition::MaxBytes(1024),
498 StopCondition::MaxRows(10),
499 ]);
500 assert_eq!(mixed.max_rows(), Some(10));
501 }
502
503 #[test]
504 fn test_max_rows_all_needs_every_child_row_triggerable() {
505 let condition =
507 StopCondition::All(vec![StopCondition::MaxRows(10), StopCondition::MaxRows(20)]);
508 assert_eq!(condition.max_rows(), Some(20));
509
510 let with_bytes = StopCondition::All(vec![
512 StopCondition::MaxRows(10),
513 StopCondition::MaxBytes(1024),
514 ]);
515 assert_eq!(with_bytes.max_rows(), None);
516
517 let with_never = StopCondition::All(vec![StopCondition::MaxRows(10), StopCondition::Never]);
518 assert_eq!(with_never.max_rows(), None);
519
520 assert_eq!(StopCondition::All(vec![]).max_rows(), None);
522 assert_eq!(StopCondition::Any(vec![]).max_rows(), None);
523 }
524
525 #[test]
526 fn test_is_row_limit_only_recurses() {
527 assert!(StopCondition::Never.is_row_limit_only());
528 assert!(StopCondition::MaxRows(10).is_row_limit_only());
529 assert!(!StopCondition::MaxBytes(10).is_row_limit_only());
530
531 assert!(
533 StopCondition::Any(vec![StopCondition::MaxRows(10), StopCondition::MaxRows(20)])
534 .is_row_limit_only()
535 );
536 assert!(
537 StopCondition::All(vec![StopCondition::Never, StopCondition::MaxRows(20)])
538 .is_row_limit_only()
539 );
540
541 assert!(
543 !StopCondition::Any(vec![
544 StopCondition::MaxRows(10),
545 StopCondition::MaxBytes(1024)
546 ])
547 .is_row_limit_only()
548 );
549
550 assert!(!StopCondition::schema_inference().is_row_limit_only());
552 assert_eq!(StopCondition::schema_inference().max_rows(), Some(10_000));
554 }
555
556 #[test]
557 fn test_max_rows_stops() {
558 let mut eval = StopEvaluator::new(StopCondition::MaxRows(100));
559 assert!(!eval.update(50, 0, 0.0));
560 assert!(!eval.should_stop());
561 assert!(eval.update(50, 0, 0.0));
562 assert!(eval.should_stop());
563 assert_eq!(
564 eval.truncation_reason(),
565 Some(TruncationReason::MaxRows(100))
566 );
567 }
568
569 #[test]
570 fn test_max_bytes_stops() {
571 let mut eval = StopEvaluator::new(StopCondition::MaxBytes(1000));
572 assert!(!eval.update(10, 500, 0.0));
573 assert!(eval.update(10, 600, 0.0));
574 assert_eq!(
575 eval.truncation_reason(),
576 Some(TruncationReason::MaxBytes(1000))
577 );
578 }
579
580 #[test]
581 fn test_memory_pressure_stops() {
582 let mut eval = StopEvaluator::new(StopCondition::MemoryPressure(0.9));
583 assert!(!eval.update(100, 0, 0.5));
584 assert!(eval.update(100, 0, 0.95));
585 assert_eq!(
586 eval.truncation_reason(),
587 Some(TruncationReason::MemoryPressure)
588 );
589 }
590
591 #[test]
592 fn test_confidence_threshold_without_estimate() {
593 let mut eval = StopEvaluator::new(StopCondition::ConfidenceThreshold(0.95));
595 assert!(!eval.update(100_000, 0, 0.0));
596 assert!(!eval.should_stop());
597 }
598
599 #[test]
600 fn test_confidence_threshold_with_estimate() {
601 let mut eval =
602 StopEvaluator::new(StopCondition::ConfidenceThreshold(0.95)).with_estimated_total(1000);
603 assert!(!eval.update(900, 0, 0.0));
604 assert!(eval.update(100, 0, 0.0)); }
606
607 #[test]
608 fn test_never_never_stops() {
609 let mut eval = StopEvaluator::new(StopCondition::Never);
610 for _ in 0..100 {
611 assert!(!eval.update(1_000_000, 1_000_000, 1.0));
612 }
613 assert!(!eval.should_stop());
614 }
615
616 #[test]
617 fn test_any_stops_on_first() {
618 let condition = StopCondition::Any(vec![
619 StopCondition::MaxRows(100),
620 StopCondition::MaxBytes(1_000_000),
621 ]);
622 let mut eval = StopEvaluator::new(condition);
623 assert!(eval.update(100, 500, 0.0));
625 assert_eq!(
626 eval.truncation_reason(),
627 Some(TruncationReason::MaxRows(100))
628 );
629 }
630
631 #[test]
632 fn test_all_requires_all() {
633 let condition = StopCondition::All(vec![
634 StopCondition::MaxRows(100),
635 StopCondition::MaxBytes(1000),
636 ]);
637 let mut eval = StopEvaluator::new(condition);
638 assert!(!eval.update(100, 500, 0.0));
640 assert!(eval.update(0, 600, 0.0));
642 assert_eq!(
643 eval.truncation_reason(),
644 Some(TruncationReason::MaxRows(100))
645 );
646 }
647
648 #[test]
649 fn test_all_empty_never_triggers() {
650 let mut eval = StopEvaluator::new(StopCondition::All(vec![]));
651 assert!(!eval.update(100, 100, 1.0));
652 }
653
654 #[test]
655 fn test_convenience_schema_inference() {
656 let condition = StopCondition::schema_inference();
657 match &condition {
658 StopCondition::Any(conditions) => {
659 assert_eq!(conditions.len(), 2);
660 assert!(matches!(conditions[0], StopCondition::MaxRows(10_000)));
661 assert!(matches!(
662 conditions[1],
663 StopCondition::SchemaStable {
664 consecutive_stable_rows: 1_000
665 }
666 ));
667 }
668 _ => panic!("Expected Any variant"),
669 }
670 }
671
672 #[test]
673 fn test_convenience_quality_sample() {
674 let condition = StopCondition::quality_sample();
675 match &condition {
676 StopCondition::Any(conditions) => {
677 assert_eq!(conditions.len(), 3);
678 assert!(matches!(conditions[0], StopCondition::MaxRows(50_000)));
679 assert!(matches!(conditions[1], StopCondition::MaxBytes(52_428_800)));
680 }
681 _ => panic!("Expected Any variant"),
682 }
683 }
684
685 #[test]
686 fn test_once_triggered_stays_triggered() {
687 let mut eval = StopEvaluator::new(StopCondition::MaxRows(10));
688 assert!(eval.update(10, 0, 0.0));
689 assert!(eval.update(5, 0, 0.0));
691 assert!(eval.should_stop());
692 }
693
694 #[test]
695 fn test_schema_stability_tracker() {
696 let condition = StopCondition::SchemaStable {
697 consecutive_stable_rows: 3,
698 };
699 let mut tracker = SchemaStabilityTracker::from_condition(&condition).unwrap();
700
701 let fp: u64 = 0xABCD;
702 assert!(!tracker.update(fp, 1)); assert!(!tracker.update(fp, 1)); assert!(!tracker.update(fp, 1)); assert!(tracker.update(fp, 1)); }
707
708 #[test]
709 fn test_schema_stability_tracker_resets_on_change() {
710 let condition = StopCondition::SchemaStable {
711 consecutive_stable_rows: 3,
712 };
713 let mut tracker = SchemaStabilityTracker::from_condition(&condition).unwrap();
714
715 let fp1: u64 = 0x1111;
716 let fp2: u64 = 0x2222;
717
718 assert!(!tracker.update(fp1, 1)); assert!(!tracker.update(fp1, 1)); assert!(!tracker.update(fp2, 1)); assert!(!tracker.update(fp2, 1)); assert!(!tracker.update(fp2, 1)); assert!(tracker.update(fp2, 1)); }
726
727 #[test]
728 fn test_schema_stable_threshold_extraction() {
729 assert_eq!(schema_stable_threshold(&StopCondition::Never), None);
730 assert_eq!(
731 schema_stable_threshold(&StopCondition::SchemaStable {
732 consecutive_stable_rows: 500
733 }),
734 Some(500)
735 );
736 let nested = StopCondition::Any(vec![
738 StopCondition::MaxRows(100),
739 StopCondition::SchemaStable {
740 consecutive_stable_rows: 200,
741 },
742 ]);
743 assert_eq!(schema_stable_threshold(&nested), Some(200));
744 }
745
746 #[test]
747 fn test_rows_and_bytes_accessors() {
748 let mut eval = StopEvaluator::new(StopCondition::Never);
749 eval.update(100, 500, 0.0);
750 eval.update(200, 1000, 0.0);
751 assert_eq!(eval.rows_processed(), 300);
752 assert_eq!(eval.bytes_consumed(), 1500);
753 }
754}