1use std::any::Any;
104use std::cell::RefCell;
105use std::collections::{BTreeMap, BTreeSet};
106use std::convert::Infallible;
107use std::future::Future;
108use std::pin::Pin;
109use std::rc::{Rc, Weak};
110use std::sync::Arc;
111use std::task::Poll;
112
113use differential_dataflow::dynamic::pointstamp::PointStamp;
114use differential_dataflow::lattice::Lattice;
115use differential_dataflow::operators::arrange::Arranged;
116use differential_dataflow::operators::arrange::ShutdownButton;
117use differential_dataflow::operators::iterate::Variable;
118use differential_dataflow::trace::{BatchReader, TraceReader};
119use differential_dataflow::{AsCollection, Data, VecCollection};
120use futures::FutureExt;
121use futures::channel::oneshot;
122use itertools::Itertools;
123use mz_compute_types::dataflows::{DataflowDescription, IndexDesc};
124use mz_compute_types::dyncfgs::{
125 COMPUTE_APPLY_COLUMN_DEMANDS, COMPUTE_LOGICAL_BACKPRESSURE_INFLIGHT_SLACK,
126 COMPUTE_LOGICAL_BACKPRESSURE_MAX_RETAINED_CAPABILITIES, ENABLE_COMPUTE_LOGICAL_BACKPRESSURE,
127 ENABLE_COMPUTE_TEMPORAL_BUCKETING, SUBSCRIBE_SNAPSHOT_OPTIMIZATION, TEMPORAL_BUCKETING_SUMMARY,
128};
129use mz_compute_types::plan::render_plan::{
130 self, BindStage, LetBind, LetFreePlan, RecBind, RenderPlan,
131};
132use mz_compute_types::plan::scalar::LirScalarExpr;
133use mz_compute_types::plan::{ArrangementStrategy, LirId};
134use mz_expr::{EvalError, Id, LocalId, permutation_for_arrangement};
135use mz_persist_client::operators::shard_source::{ErrorHandler, SnapshotMode};
136use mz_repr::explain::DummyHumanizer;
137use mz_repr::fixed_length::ExtendDatums;
138use mz_repr::{Datum, DatumVec, Diff, GlobalId, ReprRelationType, Row, RowArena, SharedRow};
139use mz_storage_operators::persist_source;
140use mz_storage_types::controller::CollectionMetadata;
141use mz_timely_util::columnation::ColumnationChunker;
142use mz_timely_util::operator::{CollectionExt, StreamExt};
143use mz_timely_util::probe::{Handle as MzProbeHandle, ProbeNotify};
144use mz_timely_util::scope_label::ScopeExt;
145use timely::PartialOrder;
146use timely::container::CapacityContainerBuilder;
147use timely::dataflow::channels::pact::Pipeline;
148use timely::dataflow::operators::vec::ToStream;
149use timely::dataflow::operators::vec::{BranchWhen, Filter};
150use timely::dataflow::operators::{Capability, Operator, Probe, probe};
151use timely::dataflow::{Scope, Stream, StreamVec};
152use timely::order::{Product, TotalOrder};
153use timely::progress::timestamp::Refines;
154use timely::progress::{Antichain, Timestamp};
155use timely::scheduling::ActivateOnDrop;
156use timely::worker::Worker as TimelyWorker;
157
158use crate::arrangement::manager::TraceBundle;
159use crate::compute_state::ComputeState;
160use crate::extensions::arrange::{KeyCollection, MzArrange};
161use crate::extensions::reduce::MzReduce;
162use crate::extensions::temporal_bucket::TemporalBucketing;
163use crate::logging::compute::{
164 ComputeEvent, DataflowGlobal, LirMapping, LirMetadata, LogDataflowErrors, OperatorHydration,
165};
166use crate::render::context::{ArrangementFlavor, Context};
167use crate::render::errors::DataflowErrorSer;
168use crate::typedefs::{ErrBatcher, ErrBuilder, ErrSpine, KeyBatcher, MzTimestamp};
169use mz_row_spine::{DatumSeq, RowRowBatcher, RowRowBuilder};
170
171pub mod context;
172pub(crate) mod errors;
173mod flat_map;
174mod join;
175mod reduce;
176pub mod sinks;
177mod threshold;
178mod top_k;
179
180pub use context::CollectionBundle;
181pub use join::LinearJoinSpec;
182
183struct PressOnDrop<T>(ShutdownButton<T>);
187
188impl<T> Drop for PressOnDrop<T> {
189 fn drop(&mut self) {
190 self.0.press();
191 }
192}
193
194pub fn build_compute_dataflow(
200 timely_worker: &mut TimelyWorker,
201 compute_state: &mut ComputeState,
202 dataflow: DataflowDescription<RenderPlan, CollectionMetadata>,
203 start_signal: StartSignal,
204 until: Antichain<mz_repr::Timestamp>,
205 dataflow_expiration: Antichain<mz_repr::Timestamp>,
206) {
207 let recursive = dataflow
209 .objects_to_build
210 .iter()
211 .any(|object| object.plan.is_recursive());
212
213 let indexes = dataflow
215 .index_exports
216 .iter()
217 .map(|(idx_id, (idx, _typ))| (*idx_id, dataflow.depends_on(idx.on_id), idx.as_lir()))
218 .collect::<Vec<_>>();
219
220 let sinks = dataflow
222 .sink_exports
223 .iter()
224 .map(|(sink_id, sink)| (*sink_id, dataflow.depends_on(sink.from), sink.clone()))
225 .collect::<Vec<_>>();
226
227 let worker_logging = timely_worker.logger_for("timely").map(Into::into);
228 let apply_demands = COMPUTE_APPLY_COLUMN_DEMANDS.get(&compute_state.worker_config);
229 let subscribe_snapshot_optimization =
230 SUBSCRIBE_SNAPSHOT_OPTIMIZATION.get(&compute_state.worker_config);
231
232 let name = format!("Dataflow: {}", &dataflow.debug_name);
233 let input_name = format!("InputRegion: {}", &dataflow.debug_name);
234 let build_name = format!("BuildRegion: {}", &dataflow.debug_name);
235
236 timely_worker.dataflow_core(&name, worker_logging, Box::new(()), |_, scope| {
237 let scope = scope.with_label();
238
239 let mut imported_sources = Vec::new();
244 let mut tokens: BTreeMap<_, Rc<dyn Any>> = BTreeMap::new();
245 let output_probe = MzProbeHandle::default();
246
247 scope.clone().region_named(&input_name, |region| {
248 for (source_id, import) in dataflow.source_imports.iter() {
250 region.region_named(&format!("Source({:?})", source_id), |inner| {
251 let mut read_schema = None;
252 let mut mfp = import.desc.arguments.operators.clone().map(|mut ops| {
253 if apply_demands {
256 let demands = ops.demand();
257 let new_desc = import
258 .desc
259 .storage_metadata
260 .relation_desc
261 .apply_demand(&demands);
262 let new_arity = demands.len();
263 let remap: BTreeMap<_, _> = demands
264 .into_iter()
265 .enumerate()
266 .map(|(new, old)| (old, new))
267 .collect();
268 ops.permute_fn(|old_idx| remap[&old_idx], new_arity);
269 read_schema = Some(new_desc);
270 }
271
272 mz_expr::MfpPlan::create_from(ops)
273 .expect("Linear operators should always be valid")
274 });
275
276 let snapshot_mode = if import.with_snapshot || !subscribe_snapshot_optimization
277 {
278 SnapshotMode::Include
279 } else {
280 compute_state.metrics.inc_subscribe_snapshot_optimization();
281 SnapshotMode::Exclude
282 };
283 let suppress_early_progress_as_of = dataflow.as_of.clone();
284
285 let (mut ok_stream, err_stream, token) =
288 persist_source::persist_source::<DataflowErrorSer>(
289 inner,
290 *source_id,
291 Arc::clone(&compute_state.persist_clients),
292 &compute_state.txns_ctx,
293 import.desc.storage_metadata.clone(),
294 read_schema,
295 dataflow.as_of.clone(),
296 snapshot_mode,
297 until.clone(),
298 mfp.as_mut(),
299 compute_state.dataflow_max_inflight_bytes(),
300 start_signal.clone().into_send_future(),
301 ErrorHandler::Halt("compute_import"),
302 );
303
304 assert!(mfp.map(|x| x.is_identity()).unwrap_or(true));
307
308 if let Some(as_of) = suppress_early_progress_as_of {
312 ok_stream = suppress_early_progress(ok_stream, as_of);
313 }
314
315 if ENABLE_COMPUTE_LOGICAL_BACKPRESSURE.get(&compute_state.worker_config) {
316 let limit = COMPUTE_LOGICAL_BACKPRESSURE_MAX_RETAINED_CAPABILITIES
318 .get(&compute_state.worker_config);
319 let slack = COMPUTE_LOGICAL_BACKPRESSURE_INFLIGHT_SLACK
320 .get(&compute_state.worker_config)
321 .as_millis()
322 .try_into()
323 .expect("must fit");
324
325 let stream = ok_stream.limit_progress(
326 output_probe.clone(),
327 slack,
328 limit,
329 import.upper.clone(),
330 name.clone(),
331 );
332 ok_stream = stream;
333 }
334
335 let input_probe =
337 compute_state.input_probe_for(*source_id, dataflow.export_ids());
338 ok_stream = ok_stream.probe_with(&input_probe);
339
340 let (oks, errs) = (
341 ok_stream
342 .as_collection()
343 .leave_region(region)
344 .leave_region(scope),
345 err_stream
346 .as_collection()
347 .leave_region(region)
348 .leave_region(scope),
349 );
350
351 imported_sources.push((mz_expr::Id::Global(*source_id), (oks, errs)));
352
353 tokens.insert(*source_id, Rc::new(token));
355 });
356 }
357 });
358
359 if recursive {
362 scope.clone().iterative::<PointStamp<u64>, _, _>(|region| {
363 let mut context = Context::for_dataflow_in(
364 &dataflow,
365 region.clone(),
366 compute_state,
367 until,
368 dataflow_expiration,
369 );
370
371 for (id, (oks, errs)) in imported_sources.into_iter() {
372 let bundle = crate::render::CollectionBundle::from_collections(
373 oks.enter(region),
374 errs.enter(region),
375 );
376 context.insert_id(id, bundle);
378 }
379
380 for (idx_id, idx) in &dataflow.index_imports {
382 let input_probe = compute_state.input_probe_for(*idx_id, dataflow.export_ids());
383 let snapshot_mode = if idx.with_snapshot || !subscribe_snapshot_optimization {
384 SnapshotMode::Include
385 } else {
386 compute_state.metrics.inc_subscribe_snapshot_optimization();
387 SnapshotMode::Exclude
388 };
389 context.import_index(
390 scope,
391 compute_state,
392 &mut tokens,
393 input_probe,
394 *idx_id,
395 &idx.desc.as_lir(),
396 &idx.typ,
397 snapshot_mode,
398 start_signal.clone(),
399 );
400 }
401
402 for object in dataflow.objects_to_build {
404 let bundle = context.scope.clone().region_named(
405 &format!("BuildingObject({:?})", object.id),
406 |region| {
407 let depends = object.plan.depends();
408 let in_let = object.plan.is_recursive();
409 context
410 .enter_region(region, Some(&depends))
411 .render_recursive_plan(
412 object.id,
413 0,
414 object.plan,
415 BindingInfo::Body { in_let },
417 )
418 .leave_region(context.scope)
419 },
420 );
421 let global_id = object.id;
422
423 context.log_dataflow_global_id(
424 *bundle
425 .scope()
426 .addr()
427 .first()
428 .expect("Dataflow root id must exist"),
429 global_id,
430 );
431 context.insert_id(Id::Global(object.id), bundle);
432 }
433
434 for (idx_id, dependencies, idx) in indexes {
436 context.export_index_iterative(
437 scope,
438 compute_state,
439 &tokens,
440 dependencies,
441 idx_id,
442 &idx,
443 &output_probe,
444 );
445 }
446
447 for (sink_id, dependencies, sink) in sinks {
449 context.export_sink(
450 compute_state,
451 &tokens,
452 dependencies,
453 sink_id,
454 &sink,
455 start_signal.clone(),
456 &output_probe,
457 scope,
458 );
459 }
460 });
461 } else {
462 scope.clone().region_named(&build_name, |region| {
463 let mut context = Context::for_dataflow_in(
464 &dataflow,
465 region.clone(),
466 compute_state,
467 until,
468 dataflow_expiration,
469 );
470
471 for (id, (oks, errs)) in imported_sources.into_iter() {
472 let bundle = crate::render::CollectionBundle::from_collections(
473 oks.enter_region(region),
474 errs.enter_region(region),
475 );
476 context.insert_id(id, bundle);
478 }
479
480 for (idx_id, idx) in &dataflow.index_imports {
482 let input_probe = compute_state.input_probe_for(*idx_id, dataflow.export_ids());
483 let snapshot_mode = if idx.with_snapshot || !subscribe_snapshot_optimization {
484 SnapshotMode::Include
485 } else {
486 compute_state.metrics.inc_subscribe_snapshot_optimization();
487 SnapshotMode::Exclude
488 };
489 context.import_index(
490 scope,
491 compute_state,
492 &mut tokens,
493 input_probe,
494 *idx_id,
495 &idx.desc.as_lir(),
496 &idx.typ,
497 snapshot_mode,
498 start_signal.clone(),
499 );
500 }
501
502 for object in dataflow.objects_to_build {
504 let bundle = context.scope.clone().region_named(
505 &format!("BuildingObject({:?})", object.id),
506 |region| {
507 let depends = object.plan.depends();
508 context
509 .enter_region(region, Some(&depends))
510 .render_plan(object.id, object.plan)
511 .leave_region(context.scope)
512 },
513 );
514 let global_id = object.id;
515 context.log_dataflow_global_id(
516 *bundle
517 .scope()
518 .addr()
519 .first()
520 .expect("Dataflow root id must exist"),
521 global_id,
522 );
523 context.insert_id(Id::Global(object.id), bundle);
524 }
525
526 for (idx_id, dependencies, idx) in indexes {
528 context.export_index(
529 compute_state,
530 &tokens,
531 dependencies,
532 idx_id,
533 &idx,
534 &output_probe,
535 );
536 }
537
538 for (sink_id, dependencies, sink) in sinks {
540 context.export_sink(
541 compute_state,
542 &tokens,
543 dependencies,
544 sink_id,
545 &sink,
546 start_signal.clone(),
547 &output_probe,
548 scope,
549 );
550 }
551 });
552 }
553 });
554}
555
556impl<'g, T> Context<'g, T>
559where
560 T: Refines<mz_repr::Timestamp> + RenderTimestamp,
561{
562 fn import_filtered_index_collection<
566 'outer,
567 Tr: TraceReader<Time = mz_repr::Timestamp> + Clone,
568 V: Data,
569 >(
570 &self,
571 arranged: Arranged<'outer, Tr>,
572 start_signal: StartSignal,
573 mut logic: impl FnMut(Tr::Key<'_>, Tr::Val<'_>) -> V + 'static,
574 ) -> VecCollection<'g, T, V, Tr::Diff>
575 where
576 mz_repr::Timestamp: TotalOrder,
579 {
580 let oks = arranged.stream.with_start_signal(start_signal).filter({
581 let as_of = self.as_of_frontier.clone();
582 move |b| !<Antichain<mz_repr::Timestamp> as PartialOrder>::less_equal(b.upper(), &as_of)
583 });
584 Arranged::<'outer, Tr>::flat_map_batches(oks, move |a, b| [logic(a, b)]).enter(self.scope)
585 }
586
587 pub(crate) fn import_index<'outer>(
588 &mut self,
589 outer: Scope<'outer, mz_repr::Timestamp>,
590 compute_state: &mut ComputeState,
591 tokens: &mut BTreeMap<GlobalId, Rc<dyn std::any::Any>>,
592 input_probe: probe::Handle<mz_repr::Timestamp>,
593 idx_id: GlobalId,
594 idx: &IndexDesc<LirScalarExpr>,
595 typ: &ReprRelationType,
596 snapshot_mode: SnapshotMode,
597 start_signal: StartSignal,
598 ) {
599 if let Some(traces) = compute_state.traces.get_mut(&idx_id) {
600 assert!(
601 PartialOrder::less_equal(&traces.compaction_frontier(), &self.as_of_frontier),
602 "Index {idx_id} has been allowed to compact beyond the dataflow as_of"
603 );
604
605 let token = traces.to_drop().clone();
606
607 let (mut oks, ok_button) = traces.oks_mut().import_frontier_core(
608 outer,
609 &format!("Index({}, {:?})", idx.on_id, idx.key),
610 self.as_of_frontier.clone(),
611 self.until.clone(),
612 );
613
614 oks.stream = oks.stream.probe_with(&input_probe);
615
616 let (err_arranged, err_button) = traces.errs_mut().import_frontier_core(
617 outer,
618 &format!("ErrIndex({}, {:?})", idx.on_id, idx.key),
619 self.as_of_frontier.clone(),
620 self.until.clone(),
621 );
622
623 let bundle = match snapshot_mode {
624 SnapshotMode::Include => {
625 let ok_arranged = oks
626 .enter(self.scope)
627 .with_start_signal(start_signal.clone());
628 let err_arranged = err_arranged
629 .enter(self.scope)
630 .with_start_signal(start_signal);
631 CollectionBundle::from_expressions(
632 idx.key.clone(),
633 ArrangementFlavor::Trace(idx_id, ok_arranged, err_arranged),
634 )
635 }
636 SnapshotMode::Exclude => {
637 let oks = {
644 let mut datums = DatumVec::new();
645 let (permutation, _thinning) =
646 permutation_for_arrangement(&idx.key, typ.arity());
647 self.import_filtered_index_collection(
648 oks,
649 start_signal.clone(),
650 move |k: DatumSeq, v: DatumSeq| {
651 let temp_storage = RowArena::new();
652 let mut datums_borrow = datums.borrow();
653 k.extend_datums(&temp_storage, &mut datums_borrow, None);
654 v.extend_datums(&temp_storage, &mut datums_borrow, None);
655 SharedRow::pack(permutation.iter().map(|i| datums_borrow[*i]))
656 },
657 )
658 };
659 let errs = self.import_filtered_index_collection(
660 err_arranged,
661 start_signal,
662 |e, _| e.clone(),
663 );
664 CollectionBundle::from_collections(oks, errs)
665 }
666 };
667 self.update_id(Id::Global(idx.on_id), bundle);
668 tokens.insert(
669 idx_id,
670 Rc::new((PressOnDrop(ok_button), PressOnDrop(err_button), token)),
671 );
672 } else {
673 panic!(
674 "import of index {} failed while building dataflow {}",
675 idx_id, self.dataflow_id
676 );
677 }
678 }
679}
680
681impl<'g> Context<'g, mz_repr::Timestamp> {
684 pub(crate) fn export_index(
685 &self,
686 compute_state: &mut ComputeState,
687 tokens: &BTreeMap<GlobalId, Rc<dyn std::any::Any>>,
688 dependency_ids: BTreeSet<GlobalId>,
689 idx_id: GlobalId,
690 idx: &IndexDesc<LirScalarExpr>,
691 output_probe: &MzProbeHandle<mz_repr::Timestamp>,
692 ) {
693 let mut needed_tokens = Vec::new();
695 for dep_id in dependency_ids {
696 if let Some(token) = tokens.get(&dep_id) {
697 needed_tokens.push(Rc::clone(token));
698 }
699 }
700 let bundle = self.lookup_id(Id::Global(idx_id)).unwrap_or_else(|| {
701 panic!(
702 "Arrangement alarmingly absent! id: {:?}",
703 Id::Global(idx_id)
704 )
705 });
706
707 let key = &idx.key;
708 match bundle.arrangement(key) {
709 Some(ArrangementFlavor::Local(mut oks, mut errs)) => {
710 if let Some(&expiration) = self.dataflow_expiration.as_option() {
713 oks.stream = oks.stream.expire_stream_at(
714 &format!("{}_export_index_oks", self.debug_name),
715 expiration,
716 );
717 errs.stream = errs.stream.expire_stream_at(
718 &format!("{}_export_index_errs", self.debug_name),
719 expiration,
720 );
721 }
722
723 oks.stream = oks.stream.probe_notify_with(vec![output_probe.clone()]);
724
725 if let Some(logger) = compute_state.compute_logger.clone() {
727 errs.stream = errs.stream.log_dataflow_errors(logger, idx_id);
728 }
729
730 compute_state.traces.set(
731 idx_id,
732 TraceBundle::new(oks.trace, errs.trace).with_drop(needed_tokens),
733 );
734 }
735 Some(ArrangementFlavor::Trace(gid, _, _)) => {
736 let trace = compute_state.traces.get(&gid).unwrap().clone();
739 compute_state.traces.set(idx_id, trace);
740 }
741 None => {
742 println!("collection available: {:?}", bundle.collection.is_none());
743 println!(
744 "keys available: {:?}",
745 bundle.arranged.keys().collect::<Vec<_>>()
746 );
747 panic!(
748 "Arrangement alarmingly absent! id: {:?}, keys: {:?}",
749 Id::Global(idx_id),
750 &key
751 );
752 }
753 };
754 }
755}
756
757impl<'g, T> Context<'g, T>
760where
761 T: RenderTimestamp,
762{
763 pub(crate) fn export_index_iterative<'outer>(
764 &self,
765 outer: Scope<'outer, mz_repr::Timestamp>,
766 compute_state: &mut ComputeState,
767 tokens: &BTreeMap<GlobalId, Rc<dyn std::any::Any>>,
768 dependency_ids: BTreeSet<GlobalId>,
769 idx_id: GlobalId,
770 idx: &IndexDesc<LirScalarExpr>,
771 output_probe: &MzProbeHandle<mz_repr::Timestamp>,
772 ) {
773 let mut needed_tokens = Vec::new();
775 for dep_id in dependency_ids {
776 if let Some(token) = tokens.get(&dep_id) {
777 needed_tokens.push(Rc::clone(token));
778 }
779 }
780 let bundle = self.lookup_id(Id::Global(idx_id)).unwrap_or_else(|| {
781 panic!(
782 "Arrangement alarmingly absent! id: {:?}",
783 Id::Global(idx_id)
784 )
785 });
786
787 let key = &idx.key;
788 match bundle.arrangement(key) {
789 Some(ArrangementFlavor::Local(oks, errs)) => {
790 let mut oks = oks
794 .as_collection(|k, v| (k.to_row(), v.to_row()))
795 .leave(outer)
796 .mz_arrange::<
797 ColumnationChunker<_>,
798 RowRowBatcher<_, _>,
799 RowRowBuilder<_, _>,
800 _,
801 >(
802 "Arrange export iterative",
803 );
804
805 let mut errs = errs
806 .as_collection(|k, v| (k.clone(), v.clone()))
807 .leave(outer)
808 .mz_arrange::<ColumnationChunker<_>, ErrBatcher<_, _>, ErrBuilder<_, _>, _>(
809 "Arrange export iterative err",
810 );
811
812 if let Some(&expiration) = self.dataflow_expiration.as_option() {
815 oks.stream = oks.stream.expire_stream_at(
816 &format!("{}_export_index_iterative_oks", self.debug_name),
817 expiration,
818 );
819 errs.stream = errs.stream.expire_stream_at(
820 &format!("{}_export_index_iterative_err", self.debug_name),
821 expiration,
822 );
823 }
824
825 oks.stream = oks.stream.probe_notify_with(vec![output_probe.clone()]);
826
827 if let Some(logger) = compute_state.compute_logger.clone() {
829 errs.stream = errs.stream.log_dataflow_errors(logger, idx_id);
830 }
831
832 compute_state.traces.set(
833 idx_id,
834 TraceBundle::new(oks.trace, errs.trace).with_drop(needed_tokens),
835 );
836 }
837 Some(ArrangementFlavor::Trace(gid, _, _)) => {
838 let trace = compute_state.traces.get(&gid).unwrap().clone();
841 compute_state.traces.set(idx_id, trace);
842 }
843 None => {
844 println!("collection available: {:?}", bundle.collection.is_none());
845 println!(
846 "keys available: {:?}",
847 bundle.arranged.keys().collect::<Vec<_>>()
848 );
849 panic!(
850 "Arrangement alarmingly absent! id: {:?}, keys: {:?}",
851 Id::Global(idx_id),
852 &key,
853 );
854 }
855 };
856 }
857}
858
859enum BindingInfo {
865 Body { in_let: bool },
866 Let { id: LocalId, last: bool },
867 LetRec { id: LocalId, last: bool },
868}
869
870impl<'scope> Context<'scope, Product<mz_repr::Timestamp, PointStamp<u64>>> {
871 fn render_recursive_plan(
884 &mut self,
885 object_id: GlobalId,
886 level: usize,
887 plan: RenderPlan,
888 binding: BindingInfo,
889 ) -> CollectionBundle<'scope, Product<mz_repr::Timestamp, PointStamp<u64>>> {
890 for BindStage { lets, recs } in plan.binds {
891 let mut let_iter = lets.into_iter().peekable();
893 while let Some(LetBind { id, value }) = let_iter.next() {
894 let bundle =
895 self.scope
896 .clone()
897 .region_named(&format!("Binding({:?})", id), |region| {
898 let depends = value.depends();
899 let last = let_iter.peek().is_none();
900 let binding = BindingInfo::Let { id, last };
901 self.enter_region(region, Some(&depends))
902 .render_letfree_plan(object_id, value, binding)
903 .leave_region(self.scope)
904 });
905 self.insert_id(Id::Local(id), bundle);
906 }
907
908 let rec_ids: Vec<_> = recs.iter().map(|r| r.id).collect();
909
910 let mut variables = BTreeMap::new();
914 for id in rec_ids.iter() {
915 use differential_dataflow::dynamic::feedback_summary;
916 let inner = feedback_summary::<u64>(level + 1, 1);
917 let (oks_v, oks_collection) =
918 Variable::new(self.scope, Product::new(Default::default(), inner.clone()));
919 let (err_v, err_collection) =
920 Variable::new(self.scope, Product::new(Default::default(), inner));
921
922 self.insert_id(
923 Id::Local(*id),
924 CollectionBundle::from_collections(oks_collection, err_collection),
925 );
926 variables.insert(Id::Local(*id), (oks_v, err_v));
927 }
928 let mut rec_iter = recs.into_iter().peekable();
930 while let Some(RecBind { id, value, limit }) = rec_iter.next() {
931 let last = rec_iter.peek().is_none();
932 let binding = BindingInfo::LetRec { id, last };
933 let bundle = self.render_recursive_plan(object_id, level + 1, value, binding);
934 let (oks, mut err) = bundle.collection.clone().unwrap();
937 self.insert_id(Id::Local(id), bundle);
938 let (oks_v, err_v) = variables.remove(&Id::Local(id)).unwrap();
939
940 let mut oks = CollectionExt::consolidate_named::<KeyBatcher<_, _, _>>(
942 oks,
943 "LetRecConsolidation",
944 );
945
946 if let Some(limit) = limit {
947 let (in_limit, over_limit) =
950 oks.inner.branch_when(move |Product { inner: ps, .. }| {
951 let iteration_index = *ps.get(level).unwrap_or(&0);
953 iteration_index + 1 >= limit.max_iters.into()
955 });
956 oks = VecCollection::new(in_limit);
957 if !limit.return_at_limit {
958 err = err.concat(VecCollection::new(over_limit).map(move |_data| {
959 DataflowErrorSer::from(EvalError::LetRecLimitExceeded(
960 format!("{}", limit.max_iters.get()).into(),
961 ))
962 }));
963 }
964 }
965
966 let err: KeyCollection<_, _, _> = err.into();
972 let errs = err
973 .mz_arrange::<
974 ColumnationChunker<_>,
975 ErrBatcher<_, _>,
976 ErrBuilder<_, _>,
977 ErrSpine<_, _>,
978 >(
979 "Arrange recursive err",
980 )
981 .mz_reduce_abelian::<_, ErrBuilder<_, _>, ErrSpine<_, _>>(
982 "Distinct recursive err",
983 move |_k, _s, t| t.push(((), Diff::ONE)),
984 )
985 .as_collection(|k, _| k.clone());
986
987 oks_v.set(oks);
988 err_v.set(errs);
989 }
990 for id in rec_ids.into_iter() {
992 let bundle = self.remove_id(Id::Local(id)).unwrap();
993 let (oks, err) = bundle.collection.unwrap();
994 self.insert_id(
995 Id::Local(id),
996 CollectionBundle::from_collections(
997 oks.leave_dynamic(level + 1),
998 err.leave_dynamic(level + 1),
999 ),
1000 );
1001 }
1002 }
1003
1004 self.render_letfree_plan(object_id, plan.body, binding)
1005 }
1006}
1007
1008impl<'scope, T: RenderTimestamp + MaybeBucketByTime> Context<'scope, T> {
1009 fn render_plan(
1020 &mut self,
1021 object_id: GlobalId,
1022 plan: RenderPlan,
1023 ) -> CollectionBundle<'scope, T> {
1024 let mut in_let = false;
1025 for BindStage { lets, recs } in plan.binds {
1026 assert!(recs.is_empty());
1027
1028 let mut let_iter = lets.into_iter().peekable();
1029 while let Some(LetBind { id, value }) = let_iter.next() {
1030 in_let = true;
1032 let bundle =
1033 self.scope
1034 .clone()
1035 .region_named(&format!("Binding({:?})", id), |region| {
1036 let depends = value.depends();
1037 let last = let_iter.peek().is_none();
1038 let binding = BindingInfo::Let { id, last };
1039 self.enter_region(region, Some(&depends))
1040 .render_letfree_plan(object_id, value, binding)
1041 .leave_region(self.scope)
1042 });
1043 self.insert_id(Id::Local(id), bundle);
1044 }
1045 }
1046
1047 self.scope.clone().region_named("Main Body", |region| {
1048 let depends = plan.body.depends();
1049 self.enter_region(region, Some(&depends))
1050 .render_letfree_plan(object_id, plan.body, BindingInfo::Body { in_let })
1051 .leave_region(self.scope)
1052 })
1053 }
1054
1055 fn render_letfree_plan(
1057 &self,
1058 object_id: GlobalId,
1059 plan: LetFreePlan,
1060 binding: BindingInfo,
1061 ) -> CollectionBundle<'scope, T> {
1062 let (mut nodes, root_id, topological_order) = plan.destruct();
1063
1064 let mut collections = BTreeMap::new();
1066
1067 let should_compute_lir_metadata = self.compute_logger.is_some();
1073 let mut lir_mapping_metadata = if should_compute_lir_metadata {
1074 Some(Vec::with_capacity(nodes.len()))
1075 } else {
1076 None
1077 };
1078
1079 let mut topo_iter = topological_order.into_iter().peekable();
1080 while let Some(lir_id) = topo_iter.next() {
1081 let node = nodes.remove(&lir_id).unwrap();
1082
1083 let metadata = if should_compute_lir_metadata {
1087 let operator = node.expr.humanize(&DummyHumanizer);
1088
1089 let operator = if topo_iter.peek().is_none() {
1091 match &binding {
1092 BindingInfo::Body { in_let: true } => format!("Returning {operator}"),
1093 BindingInfo::Body { in_let: false } => operator,
1094 BindingInfo::Let { id, last: true } => {
1095 format!("With {id} = {operator}")
1096 }
1097 BindingInfo::Let { id, last: false } => {
1098 format!("{id} = {operator}")
1099 }
1100 BindingInfo::LetRec { id, last: true } => {
1101 format!("With Recursive {id} = {operator}")
1102 }
1103 BindingInfo::LetRec { id, last: false } => {
1104 format!("{id} = {operator}")
1105 }
1106 }
1107 } else {
1108 operator
1109 };
1110
1111 let operator_id_start = self.scope.worker().peek_identifier();
1112 Some((operator, operator_id_start))
1113 } else {
1114 None
1115 };
1116
1117 let mut bundle = self.render_plan_expr(node.expr, &collections);
1118
1119 if let Some((operator, operator_id_start)) = metadata {
1120 let operator_id_end = self.scope.worker().peek_identifier();
1121 let operator_span = (operator_id_start, operator_id_end);
1122
1123 if let Some(lir_mapping_metadata) = &mut lir_mapping_metadata {
1124 lir_mapping_metadata.push((
1125 lir_id,
1126 LirMetadata::new(operator, node.parent, node.nesting, operator_span),
1127 ))
1128 }
1129 }
1130
1131 self.log_operator_hydration(&mut bundle, lir_id);
1132
1133 collections.insert(lir_id, bundle);
1134 }
1135
1136 if let Some(lir_mapping_metadata) = lir_mapping_metadata {
1137 self.log_lir_mapping(object_id, lir_mapping_metadata);
1138 }
1139
1140 collections
1141 .remove(&root_id)
1142 .expect("LetFreePlan invariant (1)")
1143 }
1144
1145 fn render_plan_expr(
1152 &self,
1153 expr: render_plan::Expr,
1154 collections: &BTreeMap<LirId, CollectionBundle<'scope, T>>,
1155 ) -> CollectionBundle<'scope, T> {
1156 use render_plan::Expr::*;
1157
1158 let expect_input = |id| {
1159 collections
1160 .get(&id)
1161 .cloned()
1162 .unwrap_or_else(|| panic!("missing input collection: {id}"))
1163 };
1164
1165 match expr {
1166 Constant { rows } => {
1167 let (rows, errs) = match rows {
1169 Ok(rows) => (rows, Vec::new()),
1170 Err(e) => (Vec::new(), vec![e]),
1171 };
1172
1173 let as_of_frontier = self.as_of_frontier.clone();
1175 let until = self.until.clone();
1176 let ok_collection = rows
1177 .into_iter()
1178 .filter_map(move |(row, mut time, diff)| {
1179 time.advance_by(as_of_frontier.borrow());
1180 if !until.less_equal(&time) {
1181 Some((
1182 row,
1183 <T as Refines<mz_repr::Timestamp>>::to_inner(time),
1184 diff,
1185 ))
1186 } else {
1187 None
1188 }
1189 })
1190 .to_stream(self.scope)
1191 .as_collection();
1192
1193 let mut error_time: mz_repr::Timestamp = Timestamp::minimum();
1194 error_time.advance_by(self.as_of_frontier.borrow());
1195 let err_collection = errs
1196 .into_iter()
1197 .map(move |e| {
1198 (
1199 DataflowErrorSer::from(e),
1200 <T as Refines<mz_repr::Timestamp>>::to_inner(error_time),
1201 Diff::ONE,
1202 )
1203 })
1204 .to_stream(self.scope)
1205 .as_collection();
1206
1207 CollectionBundle::from_collections(ok_collection, err_collection)
1208 }
1209 Get { id, keys, plan } => {
1210 let mut collection = self
1213 .lookup_id(id)
1214 .unwrap_or_else(|| panic!("Get({:?}) not found at render time", id));
1215 match plan {
1216 mz_compute_types::plan::GetPlan::PassArrangements => {
1217 assert!(
1219 keys.arranged
1220 .iter()
1221 .all(|(key, _, _)| collection.arranged.contains_key(key))
1222 );
1223 assert!(keys.raw <= collection.collection.is_some());
1224 collection.arranged.retain(|key, _value| {
1226 keys.arranged.iter().any(|(key2, _, _)| key2 == key)
1227 });
1228 collection
1229 }
1230 mz_compute_types::plan::GetPlan::Arrangement(key, row, mfp) => {
1231 let (oks, errs) = collection.as_collection_core(
1232 mfp,
1233 Some((key, row)),
1234 self.until.clone(),
1235 &self.config_set,
1236 );
1237 CollectionBundle::from_collections(oks, errs)
1238 }
1239 mz_compute_types::plan::GetPlan::Collection(mfp) => {
1240 let (oks, errs) = collection.as_collection_core(
1241 mfp,
1242 None,
1243 self.until.clone(),
1244 &self.config_set,
1245 );
1246 CollectionBundle::from_collections(oks, errs)
1247 }
1248 }
1249 }
1250 Mfp {
1251 input,
1252 mfp,
1253 input_key_val,
1254 } => {
1255 let input = expect_input(input);
1256 if mfp.is_identity() {
1258 input
1259 } else {
1260 let (oks, errs) = input.as_collection_core(
1261 mfp,
1262 input_key_val,
1263 self.until.clone(),
1264 &self.config_set,
1265 );
1266 CollectionBundle::from_collections(oks, errs)
1267 }
1268 }
1269 FlatMap {
1270 input_key,
1271 input,
1272 exprs,
1273 func,
1274 mfp_after: mfp,
1275 } => {
1276 let input = expect_input(input);
1277 self.render_flat_map(input_key, input, exprs, func, mfp)
1278 }
1279 Join { inputs, plan } => {
1280 let inputs = inputs.into_iter().map(expect_input).collect();
1281 match plan {
1282 mz_compute_types::plan::join::JoinPlan::Linear(linear_plan) => {
1283 self.render_join(inputs, linear_plan)
1284 }
1285 mz_compute_types::plan::join::JoinPlan::Delta(delta_plan) => {
1286 self.render_delta_join(inputs, delta_plan)
1287 }
1288 }
1289 }
1290 Reduce {
1291 input_key,
1292 input,
1293 key_val_plan,
1294 plan,
1295 mfp_after,
1296 temporal_bucketing_strategy,
1297 } => {
1298 let input = expect_input(input);
1299 let mfp_option = (!mfp_after.is_identity()).then_some(mfp_after);
1300 self.render_reduce(
1301 input_key,
1302 input,
1303 key_val_plan,
1304 plan,
1305 mfp_option,
1306 temporal_bucketing_strategy,
1307 )
1308 }
1309 TopK {
1310 input,
1311 top_k_plan,
1312 temporal_bucketing_strategy,
1313 } => {
1314 let input = expect_input(input);
1315 self.render_topk(input, top_k_plan, temporal_bucketing_strategy)
1316 }
1317 Negate { input } => {
1318 let input = expect_input(input);
1319 let (oks, errs) = input.as_specific_collection(None, &self.config_set);
1320 CollectionBundle::from_collections(oks.negate(), errs)
1321 }
1322 Threshold {
1323 input,
1324 threshold_plan,
1325 } => {
1326 let input = expect_input(input);
1327 self.render_threshold(input, threshold_plan)
1328 }
1329 Union {
1330 inputs,
1331 consolidate_output,
1332 temporal_bucketing_strategies,
1333 } => {
1334 let mut oks = Vec::new();
1335 let mut errs = Vec::new();
1336 for (input, strategy) in inputs.into_iter().zip_eq(temporal_bucketing_strategies) {
1337 let (os, es) =
1338 expect_input(input).as_specific_collection(None, &self.config_set);
1339 let os = if matches!(strategy, ArrangementStrategy::TemporalBucketing)
1343 && ENABLE_COMPUTE_TEMPORAL_BUCKETING.get(&self.config_set)
1344 {
1345 let summary: mz_repr::Timestamp = TEMPORAL_BUCKETING_SUMMARY
1346 .get(&self.config_set)
1347 .try_into()
1348 .expect("must fit");
1349 T::maybe_apply_temporal_bucketing(
1350 os.inner,
1351 self.as_of_frontier.clone(),
1352 summary,
1353 )
1354 } else {
1355 os
1356 };
1357 oks.push(os);
1358 errs.push(es);
1359 }
1360 let mut oks = differential_dataflow::collection::concatenate(self.scope, oks);
1361 if consolidate_output {
1362 oks = CollectionExt::consolidate_named::<KeyBatcher<_, _, _>>(
1363 oks,
1364 "UnionConsolidation",
1365 )
1366 }
1367 let errs = differential_dataflow::collection::concatenate(self.scope, errs);
1368 CollectionBundle::from_collections(oks, errs)
1369 }
1370 ArrangeBy {
1371 input_key,
1372 input,
1373 input_mfp,
1374 forms: keys,
1375 strategy,
1376 } => {
1377 let input = expect_input(input);
1378 input.ensure_collections(
1379 keys,
1380 input_key,
1381 input_mfp,
1382 self.as_of_frontier.clone(),
1383 self.until.clone(),
1384 &self.config_set,
1385 strategy,
1386 )
1387 }
1388 }
1389 }
1390
1391 fn log_dataflow_global_id(&self, dataflow_index: usize, global_id: GlobalId) {
1392 if let Some(logger) = &self.compute_logger {
1393 logger.log(&ComputeEvent::DataflowGlobal(DataflowGlobal {
1394 dataflow_index,
1395 global_id,
1396 }));
1397 }
1398 }
1399
1400 fn log_lir_mapping(&self, global_id: GlobalId, mapping: Vec<(LirId, LirMetadata)>) {
1401 if let Some(logger) = &self.compute_logger {
1402 logger.log(&ComputeEvent::LirMapping(LirMapping { global_id, mapping }));
1403 }
1404 }
1405
1406 fn log_operator_hydration(&self, bundle: &mut CollectionBundle<'scope, T>, lir_id: LirId) {
1407 match bundle.arranged.values_mut().next() {
1427 Some(arrangement) => {
1428 use ArrangementFlavor::*;
1429
1430 match arrangement {
1431 Local(a, _) => {
1432 a.stream = self.log_operator_hydration_inner(a.stream.clone(), lir_id);
1433 }
1434 Trace(_, a, _) => {
1435 a.stream = self.log_operator_hydration_inner(a.stream.clone(), lir_id);
1436 }
1437 }
1438 }
1439 None => {
1440 let (oks, _) = bundle
1441 .collection
1442 .as_mut()
1443 .expect("CollectionBundle invariant");
1444 let stream = self.log_operator_hydration_inner(oks.inner.clone(), lir_id);
1445 *oks = stream.as_collection();
1446 }
1447 }
1448 }
1449
1450 fn log_operator_hydration_inner<D>(
1451 &self,
1452 stream: Stream<'scope, T, D>,
1453 lir_id: LirId,
1454 ) -> Stream<'scope, T, D>
1455 where
1456 D: timely::Container + Clone + 'static,
1457 {
1458 let Some(logger) = self.compute_logger.clone() else {
1459 return stream.clone(); };
1461
1462 let export_ids = self.export_ids.clone();
1463
1464 let mut hydration_frontier = Antichain::new();
1472 for time in self.as_of_frontier.iter() {
1473 if let Some(time) = time.try_step_forward() {
1474 hydration_frontier.insert(Refines::to_inner(time));
1475 }
1476 }
1477
1478 let name = format!("LogOperatorHydration ({lir_id})");
1479 stream.unary_frontier(Pipeline, &name, |_cap, _info| {
1480 let mut hydrated = false;
1481
1482 for &export_id in &export_ids {
1483 logger.log(&ComputeEvent::OperatorHydration(OperatorHydration {
1484 export_id,
1485 lir_id,
1486 hydrated,
1487 }));
1488 }
1489
1490 move |(input, frontier), output| {
1491 input.for_each(|cap, data| {
1493 output.session(&cap).give_container(data);
1494 });
1495
1496 if hydrated {
1497 return;
1498 }
1499
1500 if PartialOrder::less_equal(&hydration_frontier.borrow(), &frontier.frontier()) {
1501 hydrated = true;
1502
1503 for &export_id in &export_ids {
1504 logger.log(&ComputeEvent::OperatorHydration(OperatorHydration {
1505 export_id,
1506 lir_id,
1507 hydrated,
1508 }));
1509 }
1510 }
1511 }
1512 })
1513 }
1514}
1515
1516#[allow(dead_code)] pub trait RenderTimestamp: MzTimestamp + Default + Refines<mz_repr::Timestamp> {
1519 fn system_time(&mut self) -> &mut mz_repr::Timestamp;
1524 fn system_delay(delay: mz_repr::Timestamp) -> <Self as Timestamp>::Summary;
1526 fn event_time(&self) -> mz_repr::Timestamp;
1528 fn event_time_mut(&mut self) -> &mut mz_repr::Timestamp;
1530 fn event_delay(delay: mz_repr::Timestamp) -> <Self as Timestamp>::Summary;
1532 fn step_back(&self) -> Self;
1535}
1536
1537pub trait MaybeBucketByTime: Timestamp {
1544 fn maybe_apply_temporal_bucketing<'scope, D>(
1545 stream: StreamVec<'scope, Self, (D, Self, Diff)>,
1546 as_of: Antichain<mz_repr::Timestamp>,
1547 summary: mz_repr::Timestamp,
1548 ) -> VecCollection<'scope, Self, D, Diff>
1549 where
1550 D: differential_dataflow::ExchangeData
1551 + crate::typedefs::MzData
1552 + differential_dataflow::Hashable;
1553}
1554
1555impl RenderTimestamp for mz_repr::Timestamp {
1556 fn system_time(&mut self) -> &mut mz_repr::Timestamp {
1557 self
1558 }
1559 fn system_delay(delay: mz_repr::Timestamp) -> <Self as Timestamp>::Summary {
1560 delay
1561 }
1562 fn event_time(&self) -> mz_repr::Timestamp {
1563 *self
1564 }
1565 fn event_time_mut(&mut self) -> &mut mz_repr::Timestamp {
1566 self
1567 }
1568 fn event_delay(delay: mz_repr::Timestamp) -> <Self as Timestamp>::Summary {
1569 delay
1570 }
1571 fn step_back(&self) -> Self {
1572 self.saturating_sub(1)
1573 }
1574}
1575
1576impl MaybeBucketByTime for mz_repr::Timestamp {
1577 fn maybe_apply_temporal_bucketing<'scope, D>(
1578 stream: StreamVec<'scope, Self, (D, Self, Diff)>,
1579 as_of: Antichain<mz_repr::Timestamp>,
1580 summary: mz_repr::Timestamp,
1581 ) -> VecCollection<'scope, Self, D, Diff>
1582 where
1583 D: differential_dataflow::ExchangeData
1584 + crate::typedefs::MzData
1585 + differential_dataflow::Hashable,
1586 {
1587 stream
1588 .bucket::<CapacityContainerBuilder<_>>(as_of, summary)
1589 .as_collection()
1590 }
1591}
1592
1593impl RenderTimestamp for Product<mz_repr::Timestamp, PointStamp<u64>> {
1594 fn system_time(&mut self) -> &mut mz_repr::Timestamp {
1595 &mut self.outer
1596 }
1597 fn system_delay(delay: mz_repr::Timestamp) -> <Self as Timestamp>::Summary {
1598 Product::new(delay, Default::default())
1599 }
1600 fn event_time(&self) -> mz_repr::Timestamp {
1601 self.outer
1602 }
1603 fn event_time_mut(&mut self) -> &mut mz_repr::Timestamp {
1604 &mut self.outer
1605 }
1606 fn event_delay(delay: mz_repr::Timestamp) -> <Self as Timestamp>::Summary {
1607 Product::new(delay, Default::default())
1608 }
1609 fn step_back(&self) -> Self {
1610 let inner = self.inner.clone();
1614 let mut vec = inner.into_inner();
1615 for item in vec.iter_mut() {
1616 *item = item.saturating_sub(1);
1617 }
1618 Product::new(self.outer.saturating_sub(1), PointStamp::new(vec))
1619 }
1620}
1621
1622impl MaybeBucketByTime for Product<mz_repr::Timestamp, PointStamp<u64>> {
1623 fn maybe_apply_temporal_bucketing<'scope, D>(
1624 stream: StreamVec<'scope, Self, (D, Self, Diff)>,
1625 _as_of: Antichain<mz_repr::Timestamp>,
1626 _summary: mz_repr::Timestamp,
1627 ) -> VecCollection<'scope, Self, D, Diff>
1628 where
1629 D: differential_dataflow::ExchangeData
1630 + crate::typedefs::MzData
1631 + differential_dataflow::Hashable,
1632 {
1633 stream.as_collection()
1635 }
1636}
1637
1638#[derive(Clone)]
1648pub(crate) struct StartSignal {
1649 fut: futures::future::Shared<oneshot::Receiver<Infallible>>,
1654 token_ref: Weak<RefCell<Box<dyn Any>>>,
1656}
1657
1658impl StartSignal {
1659 pub fn new() -> (Self, Rc<dyn Any>) {
1662 let (tx, rx) = oneshot::channel::<Infallible>();
1663 let token: Rc<RefCell<Box<dyn Any>>> = Rc::new(RefCell::new(Box::new(tx)));
1664 let signal = Self {
1665 fut: rx.shared(),
1666 token_ref: Rc::downgrade(&token),
1667 };
1668 (signal, token)
1669 }
1670
1671 pub fn has_fired(&self) -> bool {
1672 self.token_ref.strong_count() == 0
1673 }
1674
1675 pub fn into_send_future(self) -> impl Future<Output = ()> + Send {
1680 use futures::FutureExt;
1681 self.fut.map(|_| ())
1682 }
1683
1684 pub fn drop_on_fire(&self, to_drop: Box<dyn Any>) {
1685 if let Some(token) = self.token_ref.upgrade() {
1686 let mut token = token.borrow_mut();
1687 let inner = std::mem::replace(&mut *token, Box::new(()));
1688 *token = Box::new((inner, to_drop));
1689 }
1690 }
1691}
1692
1693impl Future for StartSignal {
1694 type Output = ();
1695
1696 fn poll(mut self: Pin<&mut Self>, cx: &mut std::task::Context<'_>) -> Poll<Self::Output> {
1697 self.fut.poll_unpin(cx).map(|_| ())
1698 }
1699}
1700
1701pub(crate) trait WithStartSignal {
1703 fn with_start_signal(self, signal: StartSignal) -> Self;
1708}
1709
1710impl<'scope, Tr> WithStartSignal for Arranged<'scope, Tr>
1711where
1712 Tr: TraceReader<Time: RenderTimestamp> + Clone,
1713{
1714 fn with_start_signal(self, signal: StartSignal) -> Self {
1715 Arranged {
1716 stream: self.stream.with_start_signal(signal),
1717 trace: self.trace,
1718 }
1719 }
1720}
1721
1722impl<'scope, T: Timestamp, D> WithStartSignal for Stream<'scope, T, D>
1723where
1724 D: timely::Container + Clone + 'static,
1725{
1726 fn with_start_signal(self, signal: StartSignal) -> Self {
1727 let activations = self.scope().activations();
1728 self.unary(Pipeline, "StartSignal", |_cap, info| {
1729 let token = Box::new(ActivateOnDrop::new((), info.address, activations));
1730 signal.drop_on_fire(token);
1731
1732 let mut stash = Vec::new();
1733
1734 move |input, output| {
1735 if !signal.has_fired() {
1737 input.for_each(|cap, data| stash.push((cap, std::mem::take(data))));
1738 return;
1739 }
1740
1741 for (cap, mut data) in std::mem::take(&mut stash) {
1743 output.session(&cap).give_container(&mut data);
1744 }
1745
1746 input.for_each(|cap, data| {
1748 output.session(&cap).give_container(data);
1749 });
1750 }
1751 })
1752 }
1753}
1754
1755fn suppress_early_progress<'scope, T: Timestamp, D>(
1777 stream: Stream<'scope, T, D>,
1778 as_of: Antichain<T>,
1779) -> Stream<'scope, T, D>
1780where
1781 D: Data + timely::Container,
1782{
1783 stream.unary_frontier(Pipeline, "SuppressEarlyProgress", |default_cap, _info| {
1784 let mut early_cap = Some(default_cap);
1785
1786 move |(input, frontier), output| {
1787 input.for_each_time(|data_cap, data| {
1788 if as_of.less_than(data_cap.time()) {
1789 let mut session = output.session(&data_cap);
1790 for data in data {
1791 session.give_container(data);
1792 }
1793 } else {
1794 let cap = early_cap.as_ref().expect("early_cap can't be dropped yet");
1795 let mut session = output.session(&cap);
1796 for data in data {
1797 session.give_container(data);
1798 }
1799 }
1800 });
1801
1802 if !PartialOrder::less_equal(&frontier.frontier(), &as_of.borrow()) {
1803 early_cap.take();
1804 }
1805 }
1806 })
1807}
1808
1809trait LimitProgress<T: Timestamp> {
1811 fn limit_progress(
1839 self,
1840 handle: MzProbeHandle<T>,
1841 slack_ms: u64,
1842 limit: Option<usize>,
1843 upper: Antichain<T>,
1844 name: String,
1845 ) -> Self;
1846}
1847
1848impl<'scope, D, R> LimitProgress<mz_repr::Timestamp>
1851 for StreamVec<'scope, mz_repr::Timestamp, (D, mz_repr::Timestamp, R)>
1852where
1853 D: Clone + 'static,
1854 R: Clone + 'static,
1855{
1856 fn limit_progress(
1857 self,
1858 handle: MzProbeHandle<mz_repr::Timestamp>,
1859 slack_ms: u64,
1860 limit: Option<usize>,
1861 upper: Antichain<mz_repr::Timestamp>,
1862 name: String,
1863 ) -> Self {
1864 let scope = self.scope();
1865 let stream =
1866 self.unary_frontier(Pipeline, &format!("LimitProgress({name})"), |_cap, info| {
1867 let mut pending_times: BTreeSet<mz_repr::Timestamp> = BTreeSet::new();
1869 let mut retained_cap: Option<Capability<mz_repr::Timestamp>> = None;
1871
1872 let activator = scope.activator_for(info.address);
1873 handle.activate(activator.clone());
1874
1875 move |(input, frontier), output| {
1876 input.for_each(|cap, data| {
1877 for time in data
1878 .iter()
1879 .flat_map(|(_, time, _)| u64::from(time).checked_add(slack_ms))
1880 {
1881 let rounded_time = if slack_ms == 0 {
1885 time
1886 } else {
1887 (time / slack_ms).saturating_add(1).saturating_mul(slack_ms)
1888 };
1889 if !upper.less_than(&rounded_time.into()) {
1890 pending_times.insert(rounded_time.into());
1891 }
1892 }
1893 output.session(&cap).give_container(data);
1894 if retained_cap.as_ref().is_none_or(|c| {
1895 !c.time().less_than(cap.time()) && !upper.less_than(cap.time())
1896 }) {
1897 retained_cap = Some(cap.retain(0));
1898 }
1899 });
1900
1901 handle.with_frontier(|f| {
1902 while pending_times
1903 .first()
1904 .map_or(false, |retained_time| !f.less_than(&retained_time))
1905 {
1906 let _ = pending_times.pop_first();
1907 }
1908 });
1909
1910 while limit.map_or(false, |limit| pending_times.len() > limit) {
1911 let _ = pending_times.pop_first();
1912 }
1913
1914 match (retained_cap.as_mut(), pending_times.first()) {
1915 (Some(cap), Some(first)) => cap.downgrade(first),
1916 (_, None) => retained_cap = None,
1917 _ => {}
1918 }
1919
1920 if frontier.is_empty() {
1921 retained_cap = None;
1922 pending_times.clear();
1923 }
1924
1925 if !pending_times.is_empty() {
1926 tracing::debug!(
1927 name,
1928 info.global_id,
1929 pending_times = %PendingTimesDisplay(pending_times.iter().cloned()),
1930 frontier = ?frontier.frontier().get(0),
1931 probe = ?handle.with_frontier(|f| f.get(0).cloned()),
1932 ?upper,
1933 "pending times",
1934 );
1935 }
1936 }
1937 });
1938 stream
1939 }
1940}
1941
1942struct PendingTimesDisplay<T>(T);
1945
1946impl<T> std::fmt::Display for PendingTimesDisplay<T>
1947where
1948 T: IntoIterator<Item = mz_repr::Timestamp> + Clone,
1949{
1950 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1951 let mut iter = self.0.clone().into_iter();
1952 write!(f, "[")?;
1953 if let Some(first) = iter.next() {
1954 write!(f, "{}", first)?;
1955 let mut last = u64::from(first);
1956 for time in iter {
1957 write!(f, ", +{}", u64::from(time) - last)?;
1958 last = u64::from(time);
1959 }
1960 }
1961 write!(f, "]")?;
1962 Ok(())
1963 }
1964}
1965
1966#[derive(Clone, Copy, Debug)]
1969struct Pairer {
1970 split_arity: usize,
1971}
1972
1973impl Pairer {
1974 fn new(split_arity: usize) -> Self {
1976 Self { split_arity }
1977 }
1978
1979 fn merge<'a, I1, I2>(&self, first: I1, second: I2) -> Row
1981 where
1982 I1: IntoIterator<Item = Datum<'a>>,
1983 I2: IntoIterator<Item = Datum<'a>>,
1984 {
1985 SharedRow::pack(first.into_iter().chain(second))
1986 }
1987
1988 fn split<'a>(&self, datum_iter: impl IntoIterator<Item = Datum<'a>>) -> (Row, Row) {
1990 let mut datum_iter = datum_iter.into_iter();
1991 let mut row_builder = SharedRow::get();
1992 let first = row_builder.pack_using(datum_iter.by_ref().take(self.split_arity));
1993 let second = row_builder.pack_using(datum_iter);
1994 (first, second)
1995 }
1996}