Skip to main content

mz_compute/
render.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10//! Renders a plan into a timely/differential dataflow computation.
11//!
12//! ## Error handling
13//!
14//! Timely and differential have no idioms for computations that can error. The
15//! philosophy is, reasonably, to define the semantics of the computation such
16//! that errors are unnecessary: e.g., by using wrap-around semantics for
17//! integer overflow.
18//!
19//! Unfortunately, SQL semantics are not nearly so elegant, and require errors
20//! in myriad cases. The classic example is a division by zero, but invalid
21//! input for casts, overflowing integer operations, and dozens of other
22//! functions need the ability to produce errors ar runtime.
23//!
24//! At the moment, only *scalar* expression evaluation can fail, so only
25//! operators that evaluate scalar expressions can fail. At the time of writing,
26//! that includes map, filter, reduce, and join operators. Constants are a bit
27//! of a special case: they can be either a constant vector of rows *or* a
28//! constant, singular error.
29//!
30//! The approach taken is to build two parallel trees of computation: one for
31//! the rows that have been successfully evaluated (the "oks tree"), and one for
32//! the errors that have been generated (the "errs tree"). For example:
33//!
34//! ```text
35//!    oks1  errs1       oks2  errs2
36//!      |     |           |     |
37//!      |     |           |     |
38//!   project  |           |     |
39//!      |     |           |     |
40//!      |     |           |     |
41//!     map    |           |     |
42//!      |\    |           |     |
43//!      | \   |           |     |
44//!      |  \  |           |     |
45//!      |   \ |           |     |
46//!      |    \|           |     |
47//!   project  +           +     +
48//!      |     |          /     /
49//!      |     |         /     /
50//!    join ------------+     /
51//!      |     |             /
52//!      |     | +----------+
53//!      |     |/
54//!     oks   errs
55//! ```
56//!
57//! The project operation cannot fail, so errors from errs1 are propagated
58//! directly. Map operators are fallible and so can inject additional errors
59//! into the stream. Join operators combine the errors from each of their
60//! inputs.
61//!
62//! The semantics of the error stream are minimal. From the perspective of SQL,
63//! a dataflow is considered to be in an error state if there is at least one
64//! element in the final errs collection. The error value returned to the user
65//! is selected arbitrarily; SQL only makes provisions to return one error to
66//! the user at a time. There are plans to make the err collection accessible to
67//! end users, so they can see all errors at once.
68//!
69//! To make errors transient, simply ensure that the operator can retract any
70//! produced errors when corrected data arrives. To make errors permanent, write
71//! the operator such that it never retracts the errors it produced. Future work
72//! will likely want to introduce some sort of sort order for errors, so that
73//! permanent errors are returned to the user ahead of transient errors—probably
74//! by introducing a new error type a la:
75//!
76//! ```no_run
77//! # struct EvalError;
78//! # struct SourceError;
79//! enum DataflowError {
80//!     Transient(EvalError),
81//!     Permanent(SourceError),
82//! }
83//! ```
84//!
85//! If the error stream is empty, the oks stream must be correct. If the error
86//! stream is non-empty, then there are no semantics for the oks stream. This is
87//! sufficient to support SQL in its current form, but is likely to be
88//! unsatisfactory long term. We suspect that we can continue to imbue the oks
89//! stream with semantics if we are very careful in describing what data should
90//! and should not be produced upon encountering an error. Roughly speaking, the
91//! oks stream could represent the correct result of the computation where all
92//! rows that caused an error have been pruned from the stream. There are
93//! strange and confusing questions here around foreign keys, though: what if
94//! the optimizer proves that a particular key must exist in a collection, but
95//! the key gets pruned away because its row participated in a scalar expression
96//! evaluation that errored?
97//!
98//! In the meantime, it is probably wise for operators to keep the oks stream
99//! roughly "as correct as possible" even when errors are present in the errs
100//! stream. This reduces the amount of recomputation that must be performed
101//! if/when the errors are retracted.
102
103use 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 ::columnar::{Columnar as ColumnarData, Index as ColumnarIndex, Push as ColumnarPush};
114use differential_dataflow::dynamic::pointstamp::PointStamp;
115use differential_dataflow::lattice::Lattice;
116use differential_dataflow::operators::arrange::Arranged;
117use differential_dataflow::operators::arrange::ShutdownButton;
118use differential_dataflow::operators::iterate::Variable;
119use differential_dataflow::trace::cursor::{BatchCursor, BatchDiff, BatchKey, BatchVal};
120use differential_dataflow::trace::{BatchReader, Cursor, Navigable, TraceReader};
121use differential_dataflow::{AsCollection, Collection, Data, VecCollection};
122use futures::FutureExt;
123use futures::channel::oneshot;
124use itertools::Itertools;
125use mz_compute_types::dataflows::{DataflowDescription, IndexDesc};
126use mz_compute_types::dyncfgs::{
127    COMPUTE_APPLY_COLUMN_DEMANDS, COMPUTE_LOGICAL_BACKPRESSURE_INFLIGHT_SLACK,
128    COMPUTE_LOGICAL_BACKPRESSURE_MAX_RETAINED_CAPABILITIES, ENABLE_COMPUTE_LOGICAL_BACKPRESSURE,
129    ENABLE_COMPUTE_TEMPORAL_BUCKETING, ENABLE_ERROR_DISTINCT, SUBSCRIBE_SNAPSHOT_OPTIMIZATION,
130    TEMPORAL_BUCKETING_SUMMARY,
131};
132use mz_compute_types::plan::render_plan::{
133    self, BindStage, LetBind, LetFreePlan, RecBind, RenderPlan,
134};
135use mz_compute_types::plan::scalar::LirScalarExpr;
136use mz_compute_types::plan::{ArrangementStrategy, LirId};
137use mz_expr::{EvalError, Id, LocalId, permutation_for_arrangement};
138use mz_persist_client::operators::shard_source::{ErrorHandler, SnapshotMode};
139use mz_repr::explain::DummyHumanizer;
140use mz_repr::fixed_length::ExtendDatums;
141use mz_repr::{Datum, DatumVec, Diff, GlobalId, ReprRelationType, Row, RowArena, SharedRow};
142use mz_storage_operators::persist_source;
143use mz_storage_types::controller::CollectionMetadata;
144use mz_timely_util::columnar::Column;
145use mz_timely_util::columnation::ColumnationChunker;
146use mz_timely_util::operator::{CollectionExt, StreamExt};
147use mz_timely_util::probe::{Handle as MzProbeHandle, ProbeNotify};
148use mz_timely_util::scope_label::ScopeExt;
149use timely::PartialOrder;
150use timely::dataflow::channels::pact::Pipeline;
151use timely::dataflow::operators::core::to_stream::ToStreamBuilder;
152use timely::dataflow::operators::vec::ToStream;
153use timely::dataflow::operators::vec::{BranchWhen, Filter};
154use timely::dataflow::operators::{Capability, Operator, Probe, probe};
155use timely::dataflow::{Scope, Stream, StreamVec};
156use timely::order::{Product, TotalOrder};
157use timely::progress::timestamp::Refines;
158use timely::progress::{Antichain, Timestamp};
159use timely::scheduling::ActivateOnDrop;
160use timely::worker::Worker as TimelyWorker;
161
162use crate::arrangement::manager::TraceBundle;
163use crate::compute_state::ComputeState;
164use crate::extensions::arrange::{KeyCollection, MzArrange};
165use crate::extensions::reduce::MzReduce;
166use crate::extensions::temporal_bucket::TemporalBucketing;
167use crate::logging::compute::{
168    ComputeEvent, DataflowGlobal, LirMapping, LirMetadata, LogDataflowErrors, OperatorHydration,
169};
170use crate::render::columnar::{
171    columnar_consolidate, columnar_negate, columnar_to_vec, concat_many, vec_to_columnar,
172};
173use crate::render::context::{ArrangementFlavor, Context};
174use crate::render::errors::DataflowErrorSer;
175use crate::typedefs::{ErrBatcher, ErrBuilder, ErrSpine, KeyBatcher, MzTimestamp};
176use mz_row_spine::{DatumSeq, RowRowBatcher, RowRowBuilder};
177use mz_timely_util::columnar::consolidate::ConsolidatingColumnBuilder;
178
179pub(crate) mod columnar;
180pub mod context;
181pub(crate) mod errors;
182mod flat_map;
183mod join;
184mod reduce;
185pub mod sinks;
186mod threshold;
187mod top_k;
188
189pub use context::CollectionBundle;
190pub use join::LinearJoinSpec;
191
192/// Guard that presses a differential [`ShutdownButton`] when dropped.
193///
194/// Dropping this guard releases the imported trace's capabilities.
195struct PressOnDrop<T>(ShutdownButton<T>);
196
197impl<T> Drop for PressOnDrop<T> {
198    fn drop(&mut self) {
199        self.0.press();
200    }
201}
202
203/// Assemble the "compute"  side of a dataflow, i.e. all but the sources.
204///
205/// This method imports sources from provided assets, and then builds the remaining
206/// dataflow using "compute-local" assets like shared arrangements, and producing
207/// both arrangements and sinks.
208pub fn build_compute_dataflow(
209    timely_worker: &mut TimelyWorker,
210    compute_state: &mut ComputeState,
211    dataflow: DataflowDescription<RenderPlan, CollectionMetadata>,
212    start_signal: StartSignal,
213    until: Antichain<mz_repr::Timestamp>,
214    dataflow_expiration: Antichain<mz_repr::Timestamp>,
215) {
216    // Mutually recursive view definitions require special handling.
217    let recursive = dataflow
218        .objects_to_build
219        .iter()
220        .any(|object| object.plan.is_recursive());
221
222    // Determine indexes to export, and their dependencies.
223    let indexes = dataflow
224        .index_exports
225        .iter()
226        .map(|(idx_id, (idx, _typ))| (*idx_id, dataflow.depends_on(idx.on_id), idx.as_lir()))
227        .collect::<Vec<_>>();
228
229    // Determine sinks to export, and their dependencies.
230    let sinks = dataflow
231        .sink_exports
232        .iter()
233        .map(|(sink_id, sink)| (*sink_id, dataflow.depends_on(sink.from), sink.clone()))
234        .collect::<Vec<_>>();
235
236    let worker_logging = timely_worker.logger_for("timely").map(Into::into);
237    let apply_demands = COMPUTE_APPLY_COLUMN_DEMANDS.get(&compute_state.worker_config);
238    let subscribe_snapshot_optimization =
239        SUBSCRIBE_SNAPSHOT_OPTIMIZATION.get(&compute_state.worker_config);
240
241    let name = format!("Dataflow: {}", dataflow.debug_name);
242    let input_name = format!("InputRegion: {}", dataflow.debug_name);
243    let build_name = format!("BuildRegion: {}", dataflow.debug_name);
244
245    timely_worker.dataflow_core(&name, worker_logging, Box::new(()), |_, scope| {
246        let scope = scope.with_label();
247
248        // The scope.clone() occurs to allow import in the region.
249        // We build a region here to establish a pattern of a scope inside the dataflow,
250        // so that other similar uses (e.g. with iterative scopes) do not require weird
251        // alternate type signatures.
252        let mut imported_sources = Vec::new();
253        let mut tokens: BTreeMap<_, Rc<dyn Any>> = BTreeMap::new();
254        let output_probe = MzProbeHandle::default();
255
256        scope.clone().region_named(&input_name, |region| {
257            // Import declared sources into the rendering context.
258            for (source_id, import) in dataflow.source_imports.iter() {
259                region.region_named(&format!("Source({:?})", source_id), |inner| {
260                    let mut read_schema = None;
261                    let mut mfp = import.desc.arguments.operators.clone().map(|mut ops| {
262                        // If enabled, we read from Persist with a `RelationDesc` that
263                        // omits uneeded columns.
264                        if apply_demands {
265                            let demands = ops.demand();
266                            let new_desc = import
267                                .desc
268                                .storage_metadata
269                                .relation_desc
270                                .apply_demand(&demands);
271                            let new_arity = demands.len();
272                            let remap: BTreeMap<_, _> = demands
273                                .into_iter()
274                                .enumerate()
275                                .map(|(new, old)| (old, new))
276                                .collect();
277                            ops.permute_fn(|old_idx| remap[&old_idx], new_arity);
278                            read_schema = Some(new_desc);
279                        }
280
281                        mz_expr::MfpPlan::create_from(ops)
282                            .expect("Linear operators should always be valid")
283                    });
284
285                    let snapshot_mode = if import.with_snapshot || !subscribe_snapshot_optimization
286                    {
287                        SnapshotMode::Include
288                    } else {
289                        compute_state.metrics.inc_subscribe_snapshot_optimization();
290                        SnapshotMode::Exclude
291                    };
292                    let suppress_early_progress_as_of = dataflow.as_of.clone();
293
294                    // Note: For correctness, we require that sources only emit times advanced by
295                    // `dataflow.as_of`. `persist_source` is documented to provide this guarantee.
296                    let (mut ok_stream, err_stream, token) = persist_source::persist_source::<
297                        DataflowErrorSer,
298                        ConsolidatingColumnBuilder<Row, mz_repr::Timestamp, Diff>,
299                    >(
300                        inner,
301                        *source_id,
302                        Arc::clone(&compute_state.persist_clients),
303                        &compute_state.txns_ctx,
304                        import.desc.storage_metadata.clone(),
305                        read_schema,
306                        dataflow.as_of.clone(),
307                        snapshot_mode,
308                        until.clone(),
309                        mfp.as_mut(),
310                        compute_state.dataflow_max_inflight_bytes(),
311                        start_signal.clone().into_send_future(),
312                        ErrorHandler::Halt("compute_import"),
313                    );
314
315                    // If `mfp` is non-identity, we need to apply what remains.
316                    // For the moment, assert that it is either trivial or `None`.
317                    assert!(mfp.map(|x| x.is_identity()).unwrap_or(true));
318
319                    // To avoid a memory spike during arrangement hydration (database-issues#6368), need to
320                    // ensure that the first frontier we report into the dataflow is beyond the
321                    // `as_of`.
322                    if let Some(as_of) = suppress_early_progress_as_of {
323                        ok_stream = suppress_early_progress(ok_stream, as_of);
324                    }
325
326                    if ENABLE_COMPUTE_LOGICAL_BACKPRESSURE.get(&compute_state.worker_config) {
327                        // Apply logical backpressure to the source.
328                        let limit = COMPUTE_LOGICAL_BACKPRESSURE_MAX_RETAINED_CAPABILITIES
329                            .get(&compute_state.worker_config);
330                        let slack = COMPUTE_LOGICAL_BACKPRESSURE_INFLIGHT_SLACK
331                            .get(&compute_state.worker_config)
332                            .as_millis()
333                            .try_into()
334                            .expect("must fit");
335
336                        let stream = ok_stream.limit_progress(
337                            output_probe.clone(),
338                            slack,
339                            limit,
340                            import.upper.clone(),
341                            name.clone(),
342                        );
343                        ok_stream = stream;
344                    }
345
346                    // Attach a probe reporting the input frontier.
347                    let input_probe =
348                        compute_state.input_probe_for(*source_id, dataflow.export_ids());
349                    ok_stream = ok_stream.probe_with(&input_probe);
350
351                    let (oks, errs) = (
352                        ok_stream
353                            .as_collection()
354                            .leave_region(region)
355                            .leave_region(scope),
356                        err_stream
357                            .as_collection()
358                            .leave_region(region)
359                            .leave_region(scope),
360                    );
361
362                    imported_sources.push((mz_expr::Id::Global(*source_id), (oks, errs)));
363
364                    // Associate returned tokens with the source identifier.
365                    tokens.insert(*source_id, Rc::new(token));
366                });
367            }
368        });
369
370        // If there exists a recursive expression, we'll need to use a non-region scope,
371        // in order to support additional timestamp coordinates for iteration.
372        if recursive {
373            scope.clone().iterative::<PointStamp<u64>, _, _>(|region| {
374                let mut context = Context::for_dataflow_in(
375                    &dataflow,
376                    region.clone(),
377                    compute_state,
378                    until,
379                    dataflow_expiration,
380                );
381
382                for (id, (oks, errs)) in imported_sources.into_iter() {
383                    let bundle = crate::render::CollectionBundle::from_edge(
384                        oks.enter(region),
385                        errs.enter(region),
386                    );
387                    // Associate collection bundle with the source identifier.
388                    context.insert_id(id, bundle);
389                }
390
391                // Import declared indexes into the rendering context.
392                for (idx_id, idx) in &dataflow.index_imports {
393                    let input_probe = compute_state.input_probe_for(*idx_id, dataflow.export_ids());
394                    let snapshot_mode = if idx.with_snapshot || !subscribe_snapshot_optimization {
395                        SnapshotMode::Include
396                    } else {
397                        compute_state.metrics.inc_subscribe_snapshot_optimization();
398                        SnapshotMode::Exclude
399                    };
400                    context.import_index(
401                        scope,
402                        compute_state,
403                        &mut tokens,
404                        input_probe,
405                        *idx_id,
406                        &idx.desc.as_lir(),
407                        &idx.typ,
408                        snapshot_mode,
409                        start_signal.clone(),
410                    );
411                }
412
413                // Build declared objects.
414                for object in dataflow.objects_to_build {
415                    let bundle = context.scope.clone().region_named(
416                        &format!("BuildingObject({:?})", object.id),
417                        |region| {
418                            let depends = object.plan.depends();
419                            let in_let = object.plan.is_recursive();
420                            context
421                                .enter_region(region, Some(&depends))
422                                .render_recursive_plan(
423                                    object.id,
424                                    0,
425                                    object.plan,
426                                    // recursive plans _must_ have bodies in a let
427                                    BindingInfo::Body { in_let },
428                                )
429                                .leave_region(context.scope)
430                        },
431                    );
432                    let global_id = object.id;
433
434                    context.log_dataflow_global_id(
435                        *bundle
436                            .scope()
437                            .addr()
438                            .first()
439                            .expect("Dataflow root id must exist"),
440                        global_id,
441                    );
442                    context.insert_id(Id::Global(object.id), bundle);
443                }
444
445                // Export declared indexes.
446                for (idx_id, dependencies, idx) in indexes {
447                    context.export_index_iterative(
448                        scope,
449                        compute_state,
450                        &tokens,
451                        dependencies,
452                        idx_id,
453                        &idx,
454                        &output_probe,
455                    );
456                }
457
458                // Export declared sinks.
459                for (sink_id, dependencies, sink) in sinks {
460                    context.export_sink(
461                        compute_state,
462                        &tokens,
463                        dependencies,
464                        sink_id,
465                        &sink,
466                        start_signal.clone(),
467                        &output_probe,
468                        scope,
469                    );
470                }
471            });
472        } else {
473            scope.clone().region_named(&build_name, |region| {
474                let mut context = Context::for_dataflow_in(
475                    &dataflow,
476                    region.clone(),
477                    compute_state,
478                    until,
479                    dataflow_expiration,
480                );
481
482                for (id, (oks, errs)) in imported_sources.into_iter() {
483                    let bundle = crate::render::CollectionBundle::from_edge(
484                        oks.enter_region(region),
485                        errs.enter_region(region),
486                    );
487                    // Associate collection bundle with the source identifier.
488                    context.insert_id(id, bundle);
489                }
490
491                // Import declared indexes into the rendering context.
492                for (idx_id, idx) in &dataflow.index_imports {
493                    let input_probe = compute_state.input_probe_for(*idx_id, dataflow.export_ids());
494                    let snapshot_mode = if idx.with_snapshot || !subscribe_snapshot_optimization {
495                        SnapshotMode::Include
496                    } else {
497                        compute_state.metrics.inc_subscribe_snapshot_optimization();
498                        SnapshotMode::Exclude
499                    };
500                    context.import_index(
501                        scope,
502                        compute_state,
503                        &mut tokens,
504                        input_probe,
505                        *idx_id,
506                        &idx.desc.as_lir(),
507                        &idx.typ,
508                        snapshot_mode,
509                        start_signal.clone(),
510                    );
511                }
512
513                // Build declared objects.
514                for object in dataflow.objects_to_build {
515                    let bundle = context.scope.clone().region_named(
516                        &format!("BuildingObject({:?})", object.id),
517                        |region| {
518                            let depends = object.plan.depends();
519                            context
520                                .enter_region(region, Some(&depends))
521                                .render_plan(object.id, object.plan)
522                                .leave_region(context.scope)
523                        },
524                    );
525                    let global_id = object.id;
526                    context.log_dataflow_global_id(
527                        *bundle
528                            .scope()
529                            .addr()
530                            .first()
531                            .expect("Dataflow root id must exist"),
532                        global_id,
533                    );
534                    context.insert_id(Id::Global(object.id), bundle);
535                }
536
537                // Export declared indexes.
538                for (idx_id, dependencies, idx) in indexes {
539                    context.export_index(
540                        compute_state,
541                        &tokens,
542                        dependencies,
543                        idx_id,
544                        &idx,
545                        &output_probe,
546                    );
547                }
548
549                // Export declared sinks.
550                for (sink_id, dependencies, sink) in sinks {
551                    context.export_sink(
552                        compute_state,
553                        &tokens,
554                        dependencies,
555                        sink_id,
556                        &sink,
557                        start_signal.clone(),
558                        &output_probe,
559                        scope,
560                    );
561                }
562            });
563        }
564    });
565}
566
567// This implementation block allows child timestamps to vary from parent timestamps,
568// but requires the parent timestamp to be `repr::Timestamp`.
569impl<'g, T> Context<'g, T>
570where
571    T: Refines<mz_repr::Timestamp> + RenderTimestamp,
572{
573    /// Import the collection from the arrangement, discarding batches from the snapshot.
574    /// (This does not guarantee that no records from the snapshot are included; the assumption is
575    /// that we'll filter those out later if necessary.)
576    fn import_filtered_index_collection<
577        'outer,
578        Tr: TraceReader<Time = mz_repr::Timestamp, Batch: Navigable> + Clone,
579        V: Data,
580    >(
581        &self,
582        arranged: Arranged<'outer, Tr>,
583        start_signal: StartSignal,
584        mut logic: impl FnMut(BatchKey<'_, Tr>, BatchVal<'_, Tr>) -> V + 'static,
585    ) -> VecCollection<'g, T, V, BatchDiff<Tr>>
586    where
587        // This is implied by the fact that the outer timestamp = mz_repr::Timestamp, but it's essential
588        // for our batch-level filtering to be safe, so we document it here regardless.
589        mz_repr::Timestamp: TotalOrder,
590        BatchCursor<Tr>: Cursor<Time = mz_repr::Timestamp>,
591    {
592        let oks = arranged.stream.with_start_signal(start_signal).filter({
593            let as_of = self.as_of_frontier.clone();
594            move |b| !<Antichain<mz_repr::Timestamp> as PartialOrder>::less_equal(b.upper(), &as_of)
595        });
596        Arranged::<'outer, Tr>::flat_map_batches(oks, move |a, b| [logic(a, b)]).enter(self.scope)
597    }
598
599    pub(crate) fn import_index<'outer>(
600        &mut self,
601        outer: Scope<'outer, mz_repr::Timestamp>,
602        compute_state: &mut ComputeState,
603        tokens: &mut BTreeMap<GlobalId, Rc<dyn std::any::Any>>,
604        input_probe: probe::Handle<mz_repr::Timestamp>,
605        idx_id: GlobalId,
606        idx: &IndexDesc<LirScalarExpr>,
607        typ: &ReprRelationType,
608        snapshot_mode: SnapshotMode,
609        start_signal: StartSignal,
610    ) {
611        if let Some(traces) = compute_state.traces.get_mut(&idx_id) {
612            assert!(
613                PartialOrder::less_equal(&traces.compaction_frontier(), &self.as_of_frontier),
614                "Index {idx_id} has been allowed to compact beyond the dataflow as_of"
615            );
616
617            let token = traces.to_drop().clone();
618
619            let (mut oks, ok_button) = traces.oks_mut().import_frontier_core(
620                outer,
621                &format!("Index({}, {:?})", idx.on_id, idx.key),
622                self.as_of_frontier.clone(),
623                self.until.clone(),
624            );
625
626            oks.stream = oks.stream.probe_with(&input_probe);
627
628            let (err_arranged, err_button) = traces.errs_mut().import_frontier_core(
629                outer,
630                &format!("ErrIndex({}, {:?})", idx.on_id, idx.key),
631                self.as_of_frontier.clone(),
632                self.until.clone(),
633            );
634
635            let bundle = match snapshot_mode {
636                SnapshotMode::Include => {
637                    let ok_arranged = oks
638                        .enter(self.scope)
639                        .with_start_signal(start_signal.clone());
640                    let err_arranged = err_arranged
641                        .enter(self.scope)
642                        .with_start_signal(start_signal);
643                    CollectionBundle::from_expressions(
644                        idx.key.clone(),
645                        ArrangementFlavor::Trace(idx_id, ok_arranged, err_arranged),
646                    )
647                }
648                SnapshotMode::Exclude => {
649                    // When we import an index without a snapshot, we have two balancing considerations:
650                    // - It's easy to filter out irrelevant batches from the stream, but hard to filter them out from an arrangement.
651                    //   (The `TraceFrontier` wrapper allows us to set an "until" frontier, but not a lower.)
652                    // - We do not actually need to reference the arrangement in this dataflow, since all operators that use the arrangement
653                    //   (joins, reduces, etc.) also require the snapshot data.
654                    // So: when the snapshot is excluded, we import only the (filtered) collection itself and ignore the arrangement.
655                    let oks = {
656                        let mut datums = DatumVec::new();
657                        let (permutation, _thinning) =
658                            permutation_for_arrangement(&idx.key, typ.arity());
659                        self.import_filtered_index_collection(
660                            oks,
661                            start_signal.clone(),
662                            move |k: DatumSeq, v: DatumSeq| {
663                                let temp_storage = RowArena::new();
664                                let mut datums_borrow = datums.borrow();
665                                k.extend_datums(&temp_storage, &mut datums_borrow, None);
666                                v.extend_datums(&temp_storage, &mut datums_borrow, None);
667                                SharedRow::pack(permutation.iter().map(|i| datums_borrow[*i]))
668                            },
669                        )
670                    };
671                    let errs = self.import_filtered_index_collection(
672                        err_arranged,
673                        start_signal,
674                        |e, _| e.clone(),
675                    );
676                    // The filtered index collection is row-shaped and already
677                    // consolidated, so the encode here is non-consolidating.
678                    CollectionBundle::from_edge(vec_to_columnar(oks), errs)
679                }
680            };
681            self.update_id(Id::Global(idx.on_id), bundle);
682            tokens.insert(
683                idx_id,
684                Rc::new((PressOnDrop(ok_button), PressOnDrop(err_button), token)),
685            );
686        } else {
687            panic!(
688                "import of index {} failed while building dataflow {}",
689                idx_id, self.dataflow_id
690            );
691        }
692    }
693}
694
695// This implementation block requires the scopes have the same timestamp as the trace manager.
696// That makes some sense, because we are hoping to deposit an arrangement in the trace manager.
697impl<'g> Context<'g, mz_repr::Timestamp> {
698    pub(crate) fn export_index(
699        &self,
700        compute_state: &mut ComputeState,
701        tokens: &BTreeMap<GlobalId, Rc<dyn std::any::Any>>,
702        dependency_ids: BTreeSet<GlobalId>,
703        idx_id: GlobalId,
704        idx: &IndexDesc<LirScalarExpr>,
705        output_probe: &MzProbeHandle<mz_repr::Timestamp>,
706    ) {
707        // put together tokens that belong to the export
708        let mut needed_tokens = Vec::new();
709        for dep_id in dependency_ids {
710            if let Some(token) = tokens.get(&dep_id) {
711                needed_tokens.push(Rc::clone(token));
712            }
713        }
714        let bundle = self.lookup_id(Id::Global(idx_id)).unwrap_or_else(|| {
715            panic!(
716                "Arrangement alarmingly absent! id: {:?}",
717                Id::Global(idx_id)
718            )
719        });
720
721        let key = &idx.key;
722        match bundle.arrangement(key) {
723            Some(ArrangementFlavor::Local(mut oks, mut errs)) => {
724                // NOTE: Do not give an exported arrangement a second reader that holds a trace
725                // handle, such as a `reduce`. Such a reader pins the shared spine's physical
726                // frontier at its own lagging progress, and `ArrangementManager::maintenance` can
727                // then no longer advance it, so batches pile up in `Spine::pending`. A cursor is
728                // only checked for straddling over pending batches, so an importing dataflow's
729                // `cursor_through` eventually panics with `upper` straddles batch. Watching
730                // `errs.stream` in `output_probe` does not help, and neither does discarding the
731                // reader's output. Stream-level readers like `as_collection` are unaffected. This is
732                // why error multiplicity is not collapsed here, leaving multiplicity that crosses
733                // an index boundary unbounded. TODO(CPU-209): bound it without a trace reader.
734
735                // Ensure that the frontier does not advance past the expiration time, if set.
736                // Otherwise, we might write down incorrect data.
737                if let Some(&expiration) = self.dataflow_expiration.as_option() {
738                    oks.stream = oks.stream.expire_stream_at(
739                        &format!("{}_export_index_oks", self.debug_name),
740                        expiration,
741                    );
742                    errs.stream = errs.stream.expire_stream_at(
743                        &format!("{}_export_index_errs", self.debug_name),
744                        expiration,
745                    );
746                }
747
748                oks.stream = oks.stream.probe_notify_with(vec![output_probe.clone()]);
749
750                // Attach logging of dataflow errors.
751                if let Some(logger) = compute_state.compute_logger.clone() {
752                    errs.stream = errs.stream.log_dataflow_errors(logger, idx_id);
753                }
754
755                compute_state.traces.set(
756                    idx_id,
757                    TraceBundle::new(oks.trace, errs.trace).with_drop(needed_tokens),
758                );
759            }
760            Some(ArrangementFlavor::Trace(gid, _, _)) => {
761                // Duplicate of existing arrangement with id `gid`, so
762                // just create another handle to that arrangement.
763                let trace = compute_state.traces.get(&gid).unwrap().clone();
764                compute_state.traces.set(idx_id, trace);
765            }
766            None => {
767                println!("collection available: {:?}", bundle.collection.is_none());
768                println!(
769                    "keys available: {:?}",
770                    bundle.arranged.keys().collect::<Vec<_>>()
771                );
772                panic!(
773                    "Arrangement alarmingly absent! id: {:?}, keys: {:?}",
774                    Id::Global(idx_id),
775                    key
776                );
777            }
778        };
779    }
780}
781
782// This implementation block requires the scopes have the same timestamp as the trace manager.
783// That makes some sense, because we are hoping to deposit an arrangement in the trace manager.
784impl<'g, T> Context<'g, T>
785where
786    T: RenderTimestamp,
787{
788    pub(crate) fn export_index_iterative<'outer>(
789        &self,
790        outer: Scope<'outer, mz_repr::Timestamp>,
791        compute_state: &mut ComputeState,
792        tokens: &BTreeMap<GlobalId, Rc<dyn std::any::Any>>,
793        dependency_ids: BTreeSet<GlobalId>,
794        idx_id: GlobalId,
795        idx: &IndexDesc<LirScalarExpr>,
796        output_probe: &MzProbeHandle<mz_repr::Timestamp>,
797    ) {
798        // put together tokens that belong to the export
799        let mut needed_tokens = Vec::new();
800        for dep_id in dependency_ids {
801            if let Some(token) = tokens.get(&dep_id) {
802                needed_tokens.push(Rc::clone(token));
803            }
804        }
805        let bundle = self.lookup_id(Id::Global(idx_id)).unwrap_or_else(|| {
806            panic!(
807                "Arrangement alarmingly absent! id: {:?}",
808                Id::Global(idx_id)
809            )
810        });
811
812        let key = &idx.key;
813        match bundle.arrangement(key) {
814            Some(ArrangementFlavor::Local(oks, errs)) => {
815                // TODO: The following as_collection/leave/arrange sequence could be optimized.
816                //   * Combine as_collection and leave into a single function.
817                //   * Use columnar to extract columns from the batches to implement leave.
818                let mut oks = oks
819                    .as_collection(|k, v| (k.to_row(), v.to_row()))
820                    .leave(outer)
821                    .mz_arrange::<
822                        ColumnationChunker<_>,
823                        RowRowBatcher<_, _>,
824                        RowRowBuilder<_, _>,
825                        _,
826                    >(
827                        "Arrange export iterative",
828                    );
829
830                let mut errs = errs
831                    .as_collection(|k, v| (k.clone(), v.clone()))
832                    .leave(outer)
833                    .mz_arrange::<ColumnationChunker<_>, ErrBatcher<_, _>, ErrBuilder<_, _>, _>(
834                        "Arrange export iterative err",
835                    );
836
837                // Ensure that the frontier does not advance past the expiration time, if set.
838                // Otherwise, we might write down incorrect data.
839                if let Some(&expiration) = self.dataflow_expiration.as_option() {
840                    oks.stream = oks.stream.expire_stream_at(
841                        &format!("{}_export_index_iterative_oks", self.debug_name),
842                        expiration,
843                    );
844                    errs.stream = errs.stream.expire_stream_at(
845                        &format!("{}_export_index_iterative_err", self.debug_name),
846                        expiration,
847                    );
848                }
849
850                oks.stream = oks.stream.probe_notify_with(vec![output_probe.clone()]);
851
852                // Attach logging of dataflow errors.
853                if let Some(logger) = compute_state.compute_logger.clone() {
854                    errs.stream = errs.stream.log_dataflow_errors(logger, idx_id);
855                }
856
857                compute_state.traces.set(
858                    idx_id,
859                    TraceBundle::new(oks.trace, errs.trace).with_drop(needed_tokens),
860                );
861            }
862            Some(ArrangementFlavor::Trace(gid, _, _)) => {
863                // Duplicate of existing arrangement with id `gid`, so
864                // just create another handle to that arrangement.
865                let trace = compute_state.traces.get(&gid).unwrap().clone();
866                compute_state.traces.set(idx_id, trace);
867            }
868            None => {
869                println!("collection available: {:?}", bundle.collection.is_none());
870                println!(
871                    "keys available: {:?}",
872                    bundle.arranged.keys().collect::<Vec<_>>()
873                );
874                panic!(
875                    "Arrangement alarmingly absent! id: {:?}, keys: {:?}",
876                    Id::Global(idx_id),
877                    key,
878                );
879            }
880        };
881    }
882}
883
884/// Information about bindings, tracked in `render_recursive_plan` and
885/// `render_plan`, to be passed to `render_letfree_plan`.
886///
887/// `render_letfree_plan` uses these to produce nice output (e.g., `With ...
888/// Returning ...`) for local bindings in the `mz_lir_mapping` output.
889enum BindingInfo {
890    Body { in_let: bool },
891    Let { id: LocalId, last: bool },
892    LetRec { id: LocalId, last: bool },
893}
894
895impl<'scope> Context<'scope, Product<mz_repr::Timestamp, PointStamp<u64>>> {
896    /// Renders a plan to a differential dataflow, producing the collection of results.
897    ///
898    /// This method allows for `plan` to contain [`RecBind`]s, and is planned
899    /// in the context of `level` pre-existing iteration coordinates.
900    ///
901    /// This method recursively descends [`RecBind`] values, establishing nested scopes for each
902    /// and establishing the appropriate recursive dependencies among the bound variables.
903    /// Once all [`RecBind`]s have been rendered it calls in to `render_plan` which will error if
904    /// further [`RecBind`]s are found.
905    ///
906    /// The method requires that all variables conclude with a physical representation that
907    /// contains a collection (i.e. a non-arrangement), and it will panic otherwise.
908    fn render_recursive_plan(
909        &mut self,
910        object_id: GlobalId,
911        level: usize,
912        plan: RenderPlan,
913        binding: BindingInfo,
914    ) -> CollectionBundle<'scope, Product<mz_repr::Timestamp, PointStamp<u64>>> {
915        for BindStage { lets, recs } in plan.binds {
916            // Render the let bindings in order.
917            let mut let_iter = lets.into_iter().peekable();
918            while let Some(LetBind { id, value }) = let_iter.next() {
919                let bundle =
920                    self.scope
921                        .clone()
922                        .region_named(&format!("Binding({:?})", id), |region| {
923                            let depends = value.depends();
924                            let last = let_iter.peek().is_none();
925                            let binding = BindingInfo::Let { id, last };
926                            self.enter_region(region, Some(&depends))
927                                .render_letfree_plan(object_id, value, binding)
928                                .leave_region(self.scope)
929                        });
930                let bundle = self.distinct_binding_errs(bundle);
931                self.insert_id(Id::Local(id), bundle);
932            }
933
934            let rec_ids: Vec<_> = recs.iter().map(|r| r.id).collect();
935
936            // A binding's `Variable` serves the `Get`s rendered before the rec
937            // loop binds the real value, which are exactly the values of
938            // `recs[0..=i]`. A binding no such value reads has no use for a
939            // `Variable`-backed bundle, and installing one would build a
940            // re-encode that repacks the whole collection once per iteration
941            // with nothing to consume it.
942            let mut variable_read = BTreeSet::new();
943            let mut read_so_far = BTreeSet::new();
944            for rec in recs.iter() {
945                read_so_far.extend(rec.value.depends());
946                if read_so_far.contains(&Id::Local(rec.id)) {
947                    variable_read.insert(rec.id);
948                }
949            }
950
951            // Define variables for rec bindings.
952            // It is important that we only use the `Variable` until the object is bound.
953            // At that point, all subsequent uses should have access to the object itself.
954            let mut variables = BTreeMap::new();
955            for id in rec_ids.iter() {
956                use differential_dataflow::dynamic::feedback_summary;
957                let inner = feedback_summary::<u64>(level + 1, 1);
958                let (oks_v, oks_collection) =
959                    Variable::new(self.scope, Product::new(Default::default(), inner.clone()));
960                let (err_v, err_collection) =
961                    Variable::new(self.scope, Product::new(Default::default(), inner));
962
963                if variable_read.contains(id) {
964                    // The feedback `Variable` stays `Vec`, so each iteration crosses
965                    // the container boundary twice, encoded here for the readers and
966                    // decoded where the value is fed back. The encode is a stateless,
967                    // timestamp-agnostic pass-through, so it leaves the iterative
968                    // frontier and the fixpoint alone.
969                    self.insert_id(
970                        Id::Local(*id),
971                        CollectionBundle::from_edge(
972                            vec_to_columnar(oks_collection),
973                            err_collection,
974                        ),
975                    );
976                }
977                variables.insert(Id::Local(*id), (oks_v, err_v));
978            }
979            // The decoded value is kept so the extraction below reuses it rather than
980            // decoding the same stream twice.
981            let mut decoded_oks = BTreeMap::new();
982            let mut rec_iter = recs.into_iter().peekable();
983            while let Some(RecBind { id, value, limit }) = rec_iter.next() {
984                let last = rec_iter.peek().is_none();
985                let binding = BindingInfo::LetRec { id, last };
986                let bundle = self.render_recursive_plan(object_id, level + 1, value, binding);
987                // We need to ensure that the raw collection exists, but do not have enough information
988                // here to cause that to happen.
989                let (oks, mut err) = bundle.collection.clone().unwrap();
990                let oks = columnar_to_vec(oks);
991                decoded_oks.insert(id, oks.clone());
992                // Collapses what forward reads see. `err_v` below feeds reads rendered before this
993                // binding and is collapsed separately; without this, a `Get` in a later rec binding
994                // or in the body resolves to the bundle stored here and compounds level over level,
995                // which is exactly what the collapse prevents for non-recursive bindings.
996                let bundle = self.distinct_binding_errs(bundle);
997                self.insert_id(Id::Local(id), bundle);
998                let (oks_v, err_v) = variables.remove(&Id::Local(id)).unwrap();
999
1000                // Set oks variable to `oks` but consolidated to ensure iteration ceases at fixed point.
1001                let mut oks = CollectionExt::consolidate_named::<KeyBatcher<_, _, _>>(
1002                    oks,
1003                    "LetRecConsolidation",
1004                );
1005
1006                if let Some(limit) = limit {
1007                    // We swallow the results of the `max_iter`th iteration, because
1008                    // these results would go into the `max_iter + 1`th iteration.
1009                    let (in_limit, over_limit) =
1010                        oks.inner.branch_when(move |Product { inner: ps, .. }| {
1011                            // The iteration number, or if missing a zero (as trailing zeros are truncated).
1012                            let iteration_index = *ps.get(level).unwrap_or(&0);
1013                            // The pointstamp starts counting from 0, so we need to add 1.
1014                            iteration_index + 1 >= limit.max_iters.into()
1015                        });
1016                    oks = VecCollection::new(in_limit);
1017                    if !limit.return_at_limit {
1018                        err = err.concat(VecCollection::new(over_limit).map(move |_data| {
1019                            DataflowErrorSer::from(EvalError::LetRecLimitExceeded(
1020                                format!("{}", limit.max_iters.get()).into(),
1021                            ))
1022                        }));
1023                    }
1024                }
1025
1026                // Set err variable to the distinct elements of `err`.
1027                // Distinctness is important, as we otherwise might add the same error each iteration,
1028                // say if the limit of `oks` has an error. This would result in non-terminating rather
1029                // than a clean report of the error. The trade-off is that we lose information about
1030                // multiplicities of errors, but .. this seems to be the better call.
1031                let err: KeyCollection<_, _, _> = err.into();
1032                let errs = err
1033                    .mz_arrange::<
1034                        ColumnationChunker<_>,
1035                        ErrBatcher<_, _>,
1036                        ErrBuilder<_, _>,
1037                        ErrSpine<_, _>,
1038                    >("Arrange recursive err")
1039                    .mz_reduce_abelian::<_, ErrBuilder<_, _>, ErrSpine<_, _>, _>(
1040                        "Distinct recursive err",
1041                        move |_k, _s, t| t.push(((), Diff::ONE)),
1042                    )
1043                    .as_collection(|k, _| k.clone());
1044
1045                oks_v.set(oks);
1046                err_v.set(errs);
1047            }
1048            // Now extract each of the rec bindings into the outer scope.
1049            for id in rec_ids.into_iter() {
1050                let bundle = self.remove_id(Id::Local(id)).unwrap();
1051                let (_, err) = bundle.collection.unwrap();
1052                let oks = decoded_oks
1053                    .remove(&id)
1054                    .expect("rec binding decoded while rendering above");
1055                // `leave_dynamic` has already stripped the iteration coordinate, so
1056                // this encode runs in the parent scope.
1057                self.insert_id(
1058                    Id::Local(id),
1059                    CollectionBundle::from_edge(
1060                        vec_to_columnar(oks.leave_dynamic(level + 1)),
1061                        err.leave_dynamic(level + 1),
1062                    ),
1063                );
1064            }
1065        }
1066
1067        self.render_letfree_plan(object_id, plan.body, binding)
1068    }
1069}
1070
1071impl<'scope, T: RenderTimestamp + MaybeBucketByTime> Context<'scope, T> {
1072    /// Renders a non-recursive plan to a differential dataflow, producing the collection of
1073    /// results.
1074    ///
1075    /// The return type reflects the uncertainty about the data representation, perhaps
1076    /// as a stream of data, perhaps as an arrangement, perhaps as a stream of batches.
1077    ///
1078    /// # Panics
1079    ///
1080    /// Panics if the given plan contains any [`RecBind`]s. Recursive plans must be rendered using
1081    /// `render_recursive_plan` instead.
1082    fn render_plan(
1083        &mut self,
1084        object_id: GlobalId,
1085        plan: RenderPlan,
1086    ) -> CollectionBundle<'scope, T> {
1087        let mut in_let = false;
1088        for BindStage { lets, recs } in plan.binds {
1089            assert!(recs.is_empty());
1090
1091            let mut let_iter = lets.into_iter().peekable();
1092            while let Some(LetBind { id, value }) = let_iter.next() {
1093                // if we encounter a single let, the body is in a let
1094                in_let = true;
1095                let bundle =
1096                    self.scope
1097                        .clone()
1098                        .region_named(&format!("Binding({:?})", id), |region| {
1099                            let depends = value.depends();
1100                            let last = let_iter.peek().is_none();
1101                            let binding = BindingInfo::Let { id, last };
1102                            self.enter_region(region, Some(&depends))
1103                                .render_letfree_plan(object_id, value, binding)
1104                                .leave_region(self.scope)
1105                        });
1106                let bundle = self.distinct_binding_errs(bundle);
1107                self.insert_id(Id::Local(id), bundle);
1108            }
1109        }
1110
1111        self.scope.clone().region_named("Main Body", |region| {
1112            let depends = plan.body.depends();
1113            self.enter_region(region, Some(&depends))
1114                .render_letfree_plan(object_id, plan.body, BindingInfo::Body { in_let })
1115                .leave_region(self.scope)
1116        })
1117    }
1118
1119    /// Collapses a binding's error multiplicities.
1120    ///
1121    /// Applied to every binding, not only the multiply-read ones. Gating on the reference count
1122    /// would be a pure optimization, since collapsing a binding one `Get` reads is harmless, and
1123    /// there is almost nothing to gate: `NormalizeLets` inlines single-use bindings, so the ones
1124    /// reaching rendering are shared. See [`CollectionBundle::distinct_errs`] for why the collapse
1125    /// is needed at all, and why a binding's definition is the place for it rather than the
1126    /// multi-input operators where the duplicate copies happen to meet again.
1127    fn distinct_binding_errs(
1128        &self,
1129        bundle: CollectionBundle<'scope, T>,
1130    ) -> CollectionBundle<'scope, T> {
1131        if ENABLE_ERROR_DISTINCT.get(&self.config_set) {
1132            bundle.distinct_errs()
1133        } else {
1134            bundle
1135        }
1136    }
1137
1138    /// Renders a let-free plan to a differential dataflow, producing the collection of results.
1139    fn render_letfree_plan(
1140        &self,
1141        object_id: GlobalId,
1142        plan: LetFreePlan,
1143        binding: BindingInfo,
1144    ) -> CollectionBundle<'scope, T> {
1145        let (mut nodes, root_id, topological_order) = plan.destruct();
1146
1147        // Rendered collections by their `LirId`.
1148        let mut collections = BTreeMap::new();
1149
1150        // Mappings to send along.
1151        // To save overhead, we'll only compute mappings when we need to,
1152        // which means things get gated behind options. Unfortunately, that means we
1153        // have several `Option<...>` types that are _all_ `Some` or `None` together,
1154        // but there's no convenient way to express the invariant.
1155        let should_compute_lir_metadata = self.compute_logger.is_some();
1156        let mut lir_mapping_metadata = if should_compute_lir_metadata {
1157            Some(Vec::with_capacity(nodes.len()))
1158        } else {
1159            None
1160        };
1161
1162        let mut topo_iter = topological_order.into_iter().peekable();
1163        while let Some(lir_id) = topo_iter.next() {
1164            let node = nodes.remove(&lir_id).unwrap();
1165
1166            // TODO(mgree) need ExprHumanizer in DataflowDescription to get nice column names
1167            // ActiveComputeState can't have a catalog reference, so we'll need to capture the names
1168            // in some other structure and have that structure impl ExprHumanizer
1169            let metadata = if should_compute_lir_metadata {
1170                let operator = node.expr.humanize(&DummyHumanizer);
1171
1172                // mark the last operator in topo order with any binding decoration
1173                let operator = if topo_iter.peek().is_none() {
1174                    match &binding {
1175                        BindingInfo::Body { in_let: true } => format!("Returning {operator}"),
1176                        BindingInfo::Body { in_let: false } => operator,
1177                        BindingInfo::Let { id, last: true } => {
1178                            format!("With {id} = {operator}")
1179                        }
1180                        BindingInfo::Let { id, last: false } => {
1181                            format!("{id} = {operator}")
1182                        }
1183                        BindingInfo::LetRec { id, last: true } => {
1184                            format!("With Recursive {id} = {operator}")
1185                        }
1186                        BindingInfo::LetRec { id, last: false } => {
1187                            format!("{id} = {operator}")
1188                        }
1189                    }
1190                } else {
1191                    operator
1192                };
1193
1194                let operator_id_start = self.scope.worker().peek_identifier();
1195                Some((operator, operator_id_start))
1196            } else {
1197                None
1198            };
1199
1200            let mut bundle = self.render_plan_expr(node.expr, &collections);
1201
1202            if let Some((operator, operator_id_start)) = metadata {
1203                let operator_id_end = self.scope.worker().peek_identifier();
1204                let operator_span = (operator_id_start, operator_id_end);
1205
1206                if let Some(lir_mapping_metadata) = &mut lir_mapping_metadata {
1207                    lir_mapping_metadata.push((
1208                        lir_id,
1209                        LirMetadata::new(operator, node.parent, node.nesting, operator_span),
1210                    ))
1211                }
1212            }
1213
1214            self.log_operator_hydration(&mut bundle, lir_id);
1215
1216            collections.insert(lir_id, bundle);
1217        }
1218
1219        if let Some(lir_mapping_metadata) = lir_mapping_metadata {
1220            self.log_lir_mapping(object_id, lir_mapping_metadata);
1221        }
1222
1223        collections
1224            .remove(&root_id)
1225            .expect("LetFreePlan invariant (1)")
1226    }
1227
1228    /// Renders a [`render_plan::Expr`], producing the collection of results.
1229    ///
1230    /// # Panics
1231    ///
1232    /// Panics if any of the expr's inputs is not found in `collections`.
1233    /// Callers must ensure that input nodes have been rendered previously.
1234    fn render_plan_expr(
1235        &self,
1236        expr: render_plan::Expr,
1237        collections: &BTreeMap<LirId, CollectionBundle<'scope, T>>,
1238    ) -> CollectionBundle<'scope, T> {
1239        use render_plan::Expr::*;
1240
1241        let expect_input = |id| {
1242            collections
1243                .get(&id)
1244                .cloned()
1245                .unwrap_or_else(|| panic!("missing input collection: {id}"))
1246        };
1247
1248        match expr {
1249            Constant { rows } => {
1250                // Produce both rows and errs to avoid conditional dataflow construction.
1251                let (rows, errs) = match rows {
1252                    Ok(rows) => (rows, Vec::new()),
1253                    Err(e) => (Vec::new(), vec![e]),
1254                };
1255
1256                // We should advance times in constant collections to start from `as_of`.
1257                let as_of_frontier = self.as_of_frontier.clone();
1258                let until = self.until.clone();
1259                // Advancing times to `as_of` can collapse distinct times onto one, so
1260                // rows the planner left distinct can become duplicates. The
1261                // `ConsolidatingColumnBuilder` folds those within the batch.
1262                let ok_collection = rows
1263                    .into_iter()
1264                    .filter_map(move |(row, mut time, diff)| {
1265                        time.advance_by(as_of_frontier.borrow());
1266                        if !until.less_equal(&time) {
1267                            Some((
1268                                row.0,
1269                                <T as Refines<mz_repr::Timestamp>>::to_inner(time),
1270                                diff,
1271                            ))
1272                        } else {
1273                            None
1274                        }
1275                    })
1276                    .to_stream_with_builder::<_, ConsolidatingColumnBuilder<Row, T, Diff>>(
1277                        self.scope,
1278                    )
1279                    .as_collection();
1280
1281                let mut error_time: mz_repr::Timestamp = Timestamp::minimum();
1282                error_time.advance_by(self.as_of_frontier.borrow());
1283                let err_collection = errs
1284                    .into_iter()
1285                    .map(move |e| {
1286                        (
1287                            DataflowErrorSer::from(e),
1288                            <T as Refines<mz_repr::Timestamp>>::to_inner(error_time),
1289                            Diff::ONE,
1290                        )
1291                    })
1292                    .to_stream(self.scope)
1293                    .as_collection();
1294
1295                CollectionBundle::from_edge(ok_collection, err_collection)
1296            }
1297            Get { id, keys, plan } => {
1298                // Recover the collection from `self` and then apply `mfp` to it.
1299                // If `mfp` happens to be trivial, we can just return the collection.
1300                let mut collection = self
1301                    .lookup_id(id)
1302                    .unwrap_or_else(|| panic!("Get({:?}) not found at render time", id));
1303                match plan {
1304                    mz_compute_types::plan::GetPlan::PassArrangements => {
1305                        // Assert that each of `keys` are present in `collection`.
1306                        assert!(
1307                            keys.arranged
1308                                .iter()
1309                                .all(|(key, _, _)| collection.arranged.contains_key(key))
1310                        );
1311                        assert!(keys.raw <= collection.collection.is_some());
1312                        // Retain only those keys we want to import.
1313                        collection.arranged.retain(|key, _value| {
1314                            keys.arranged.iter().any(|(key2, _, _)| key2 == key)
1315                        });
1316                        collection
1317                    }
1318                    mz_compute_types::plan::GetPlan::Arrangement(key, row, mfp) => {
1319                        let (oks, errs) = collection.as_collection_core(
1320                            mfp,
1321                            Some((key, row)),
1322                            self.until.clone(),
1323                        );
1324                        CollectionBundle::from_edge(oks, errs)
1325                    }
1326                    mz_compute_types::plan::GetPlan::Collection(mfp) => {
1327                        let (oks, errs) =
1328                            collection.as_collection_core(mfp, None, self.until.clone());
1329                        CollectionBundle::from_edge(oks, errs)
1330                    }
1331                }
1332            }
1333            Mfp {
1334                input,
1335                mfp,
1336                input_key_val,
1337            } => {
1338                let input = expect_input(input);
1339                // If `mfp` is non-trivial, we should apply it and produce a collection.
1340                if mfp.is_identity() {
1341                    input
1342                } else {
1343                    let (oks, errs) =
1344                        input.as_collection_core(mfp, input_key_val, self.until.clone());
1345                    CollectionBundle::from_edge(oks, errs)
1346                }
1347            }
1348            FlatMap {
1349                input_key,
1350                input,
1351                exprs,
1352                func,
1353                mfp_after: mfp,
1354            } => {
1355                let input = expect_input(input);
1356                self.render_flat_map(input_key, input, exprs, func, mfp)
1357            }
1358            Join { inputs, plan } => {
1359                let inputs = inputs.into_iter().map(expect_input).collect();
1360                match plan {
1361                    mz_compute_types::plan::join::JoinPlan::Linear(linear_plan) => {
1362                        self.render_join(inputs, linear_plan)
1363                    }
1364                    mz_compute_types::plan::join::JoinPlan::Delta(delta_plan) => {
1365                        self.render_delta_join(inputs, delta_plan)
1366                    }
1367                }
1368            }
1369            Reduce {
1370                input_key,
1371                input,
1372                key_val_plan,
1373                plan,
1374                mfp_after,
1375                temporal_bucketing_strategy,
1376            } => {
1377                let input = expect_input(input);
1378                let mfp_option = (!mfp_after.is_identity()).then_some(mfp_after);
1379                self.render_reduce(
1380                    input_key,
1381                    input,
1382                    key_val_plan,
1383                    plan,
1384                    mfp_option,
1385                    temporal_bucketing_strategy,
1386                )
1387            }
1388            TopK {
1389                input,
1390                top_k_plan,
1391                temporal_bucketing_strategy,
1392            } => {
1393                let input = expect_input(input);
1394                self.render_topk(input, top_k_plan, temporal_bucketing_strategy)
1395            }
1396            Negate { input } => {
1397                let input = expect_input(input);
1398                let (oks, errs) = input
1399                    .collection
1400                    .clone()
1401                    .expect("Negate input must be an unarranged collection");
1402                CollectionBundle::from_edge(columnar_negate(oks), errs)
1403            }
1404            Threshold {
1405                input,
1406                threshold_plan,
1407            } => {
1408                let input = expect_input(input);
1409                self.render_threshold(input, threshold_plan)
1410            }
1411            Union {
1412                inputs,
1413                consolidate_output,
1414                temporal_bucketing_strategies,
1415            } => {
1416                let mut oks = Vec::new();
1417                let mut errs = Vec::new();
1418                for (input, strategy) in inputs.into_iter().zip_eq(temporal_bucketing_strategies) {
1419                    let (os, es) = expect_input(input)
1420                        .collection
1421                        .clone()
1422                        .expect("Union input must be an unarranged collection");
1423                    // Apply per-input temporal bucketing. No-op for `Direct`.
1424                    // Only consolidating Unions carry non-`Direct` strategies;
1425                    // see the `Union` arm of `lower_mir_expr_stack_safe`.
1426                    let os = if matches!(strategy, ArrangementStrategy::TemporalBucketing)
1427                        && ENABLE_COMPUTE_TEMPORAL_BUCKETING.get(&self.config_set)
1428                    {
1429                        let summary: mz_repr::Timestamp = TEMPORAL_BUCKETING_SUMMARY
1430                            .get(&self.config_set)
1431                            .try_into()
1432                            .expect("must fit");
1433                        T::maybe_apply_temporal_bucketing(
1434                            os.inner,
1435                            self.as_of_frontier.clone(),
1436                            summary,
1437                        )
1438                    } else {
1439                        os
1440                    };
1441                    oks.push(os);
1442                    errs.push(es);
1443                }
1444                let oks = concat_many(self.scope, oks);
1445                let oks = if consolidate_output {
1446                    columnar_consolidate(oks, "UnionConsolidation")
1447                } else {
1448                    oks
1449                };
1450                let errs = differential_dataflow::collection::concatenate(self.scope, errs);
1451                CollectionBundle::from_edge(oks, errs)
1452            }
1453            ArrangeBy {
1454                input_key,
1455                input,
1456                input_mfp,
1457                forms: keys,
1458                strategy,
1459            } => {
1460                let input = expect_input(input);
1461                input.ensure_collections(
1462                    keys,
1463                    input_key,
1464                    input_mfp,
1465                    self.as_of_frontier.clone(),
1466                    self.until.clone(),
1467                    &self.config_set,
1468                    strategy,
1469                )
1470            }
1471        }
1472    }
1473
1474    fn log_dataflow_global_id(&self, dataflow_index: usize, global_id: GlobalId) {
1475        if let Some(logger) = &self.compute_logger {
1476            logger.log(&ComputeEvent::DataflowGlobal(DataflowGlobal {
1477                dataflow_index,
1478                global_id,
1479            }));
1480        }
1481    }
1482
1483    fn log_lir_mapping(&self, global_id: GlobalId, mapping: Vec<(LirId, LirMetadata)>) {
1484        if let Some(logger) = &self.compute_logger {
1485            logger.log(&ComputeEvent::LirMapping(LirMapping { global_id, mapping }));
1486        }
1487    }
1488
1489    fn log_operator_hydration(&self, bundle: &mut CollectionBundle<'scope, T>, lir_id: LirId) {
1490        // A `CollectionBundle` can contain more than one collection, which makes it not obvious to
1491        // which we should attach the logging operator.
1492        //
1493        // We could attach to each collection and track the lower bound of output frontiers.
1494        // However, that would be of limited use because we expect all collections to hydrate at
1495        // roughly the same time: The `ArrangeBy` operator is not fueled, so as soon as it sees the
1496        // frontier of the unarranged collection advance, it will perform all work necessary to
1497        // also advance its own frontier. We don't expect significant delays between frontier
1498        // advancements of the unarranged and arranged collections, so attaching the logging
1499        // operator to any one of them should produce accurate results.
1500        //
1501        // If the `CollectionBundle` contains both unarranged and arranged representations it is
1502        // beneficial to attach the logging operator to one of the arranged representation to avoid
1503        // unnecessary cloning of data. The unarranged collection feeds into the arrangements, so
1504        // if we attached the logging operator to it, we would introduce a fork in its output
1505        // stream, which would necessitate that all output data is cloned. In contrast, we can hope
1506        // that the output streams of the arrangements don't yet feed into anything else, so
1507        // attaching a (pass-through) logging operator does not introduce a fork.
1508
1509        match bundle.arranged.values_mut().next() {
1510            Some(arrangement) => {
1511                use ArrangementFlavor::*;
1512
1513                match arrangement {
1514                    Local(a, _) => {
1515                        a.stream = self.log_operator_hydration_inner(a.stream.clone(), lir_id);
1516                    }
1517                    Trace(_, a, _) => {
1518                        a.stream = self.log_operator_hydration_inner(a.stream.clone(), lir_id);
1519                    }
1520                }
1521            }
1522            None => {
1523                let (oks, _) = bundle
1524                    .collection
1525                    .as_mut()
1526                    .expect("CollectionBundle invariant");
1527                let stream = self.log_operator_hydration_inner(oks.inner.clone(), lir_id);
1528                *oks = stream.as_collection();
1529            }
1530        }
1531    }
1532
1533    fn log_operator_hydration_inner<D>(
1534        &self,
1535        stream: Stream<'scope, T, D>,
1536        lir_id: LirId,
1537    ) -> Stream<'scope, T, D>
1538    where
1539        D: timely::Container + Clone + 'static,
1540    {
1541        let Some(logger) = self.compute_logger.clone() else {
1542            return stream.clone(); // hydration logging disabled
1543        };
1544
1545        let export_ids = self.export_ids.clone();
1546
1547        // Convert the dataflow as-of into a frontier we can compare with input frontiers.
1548        //
1549        // We (somewhat arbitrarily) define operators in iterative scopes to be hydrated when their
1550        // frontier advances to an outer time that's greater than the `as_of`. Comparing
1551        // `refine(as_of) < input_frontier` would find the moment when the first iteration was
1552        // complete, which is not what we want. We want `refine(as_of + 1) <= input_frontier`
1553        // instead.
1554        let mut hydration_frontier = Antichain::new();
1555        for time in self.as_of_frontier.iter() {
1556            if let Some(time) = time.try_step_forward() {
1557                hydration_frontier.insert(Refines::to_inner(time));
1558            }
1559        }
1560
1561        let name = format!("LogOperatorHydration ({lir_id})");
1562        stream.unary_frontier(Pipeline, &name, |_cap, _info| {
1563            let mut hydrated = false;
1564
1565            for &export_id in &export_ids {
1566                logger.log(&ComputeEvent::OperatorHydration(OperatorHydration {
1567                    export_id,
1568                    lir_id,
1569                    hydrated,
1570                }));
1571            }
1572
1573            move |(input, frontier), output| {
1574                // Pass through inputs.
1575                input.for_each(|cap, data| {
1576                    output.session(&cap).give_container(data);
1577                });
1578
1579                if hydrated {
1580                    return;
1581                }
1582
1583                if PartialOrder::less_equal(&hydration_frontier.borrow(), &frontier.frontier()) {
1584                    hydrated = true;
1585
1586                    for &export_id in &export_ids {
1587                        logger.log(&ComputeEvent::OperatorHydration(OperatorHydration {
1588                            export_id,
1589                            lir_id,
1590                            hydrated,
1591                        }));
1592                    }
1593                }
1594            }
1595        })
1596    }
1597}
1598
1599#[allow(dead_code)] // Some of the methods on this trait are unused, but useful to have.
1600/// A timestamp type that can be used for operations within MZ's dataflow layer.
1601pub trait RenderTimestamp: MzTimestamp + Default + Refines<mz_repr::Timestamp> {
1602    /// The system timestamp component of the timestamp.
1603    ///
1604    /// This is useful for manipulating the system time, as when delaying
1605    /// updates for subsequent cancellation, as with monotonic reduction.
1606    fn system_time(&mut self) -> &mut mz_repr::Timestamp;
1607    /// Effects a system delay in terms of the timestamp summary.
1608    fn system_delay(delay: mz_repr::Timestamp) -> <Self as Timestamp>::Summary;
1609    /// The event timestamp component of the timestamp.
1610    fn event_time(&self) -> mz_repr::Timestamp;
1611    /// The event timestamp component of the timestamp, as a mutable reference.
1612    fn event_time_mut(&mut self) -> &mut mz_repr::Timestamp;
1613    /// Effects an event delay in terms of the timestamp summary.
1614    fn event_delay(delay: mz_repr::Timestamp) -> <Self as Timestamp>::Summary;
1615    /// Steps the timestamp back so that logical compaction to the output will
1616    /// not conflate `self` with any historical times.
1617    fn step_back(&self) -> Self;
1618}
1619
1620/// Apply temporal bucketing to a stream when the timestamp type supports it.
1621///
1622/// Sibling to [`RenderTimestamp`]: bucketing is an arrangement-time concern, not a
1623/// general property of a render timestamp, so the dispatch lives in its own trait.
1624/// Total-ordered timestamps perform real bucketing; partially-ordered timestamps
1625/// (e.g. `Product<…>` in iterative scopes) implement this as a no-op.
1626pub trait MaybeBucketByTime: Timestamp + ColumnarData {
1627    /// Buckets a columnar dataflow edge, keeping it columnar.
1628    fn maybe_apply_temporal_bucketing<'scope, D>(
1629        stream: Stream<'scope, Self, Column<(D, Self, Diff)>>,
1630        as_of: Antichain<mz_repr::Timestamp>,
1631        summary: mz_repr::Timestamp,
1632    ) -> Collection<'scope, Self, Column<(D, Self, Diff)>>
1633    where
1634        D: differential_dataflow::ExchangeData
1635            + crate::typedefs::MzData
1636            + differential_dataflow::Hashable
1637            + ColumnarData,
1638        for<'a> ::columnar::Ref<'a, D>: Copy + Ord + std::hash::Hash,
1639        for<'a> ::columnar::Ref<'a, Self>: Copy + Ord,
1640        for<'a> ::columnar::Ref<'a, Diff>: Ord,
1641        for<'a> <(D, Self, Diff) as ColumnarData>::Container:
1642            ColumnarPush<&'a (D, Self, Diff)> + ColumnarPush<::columnar::Ref<'a, (D, Self, Diff)>>;
1643
1644    /// Buckets a `Vec` stream, keeping it `Vec`.
1645    ///
1646    /// For a consumer that re-encodes what it reads, where a `Vec` hands it moved
1647    /// allocations rather than copied bytes. The reduce key-value path is the one
1648    /// such caller, since its bucketed output feeds an arrangement.
1649    fn maybe_apply_temporal_bucketing_vec<'scope, D>(
1650        stream: StreamVec<'scope, Self, (D, Self, Diff)>,
1651        as_of: Antichain<mz_repr::Timestamp>,
1652        summary: mz_repr::Timestamp,
1653    ) -> VecCollection<'scope, Self, D, Diff>
1654    where
1655        D: differential_dataflow::ExchangeData
1656            + crate::typedefs::MzData
1657            + differential_dataflow::Hashable
1658            + ColumnarData,
1659        for<'a> ::columnar::Ref<'a, D>: Copy + Ord + std::hash::Hash,
1660        for<'a> ::columnar::Ref<'a, Self>: Copy + Ord,
1661        for<'a> ::columnar::Ref<'a, Diff>: Ord,
1662        for<'a> <(D, Self, Diff) as ColumnarData>::Container:
1663            ColumnarPush<&'a (D, Self, Diff)> + ColumnarPush<::columnar::Ref<'a, (D, Self, Diff)>>;
1664}
1665
1666impl RenderTimestamp for mz_repr::Timestamp {
1667    fn system_time(&mut self) -> &mut mz_repr::Timestamp {
1668        self
1669    }
1670    fn system_delay(delay: mz_repr::Timestamp) -> <Self as Timestamp>::Summary {
1671        delay
1672    }
1673    fn event_time(&self) -> mz_repr::Timestamp {
1674        *self
1675    }
1676    fn event_time_mut(&mut self) -> &mut mz_repr::Timestamp {
1677        self
1678    }
1679    fn event_delay(delay: mz_repr::Timestamp) -> <Self as Timestamp>::Summary {
1680        delay
1681    }
1682    fn step_back(&self) -> Self {
1683        self.saturating_sub(1)
1684    }
1685}
1686
1687impl MaybeBucketByTime for mz_repr::Timestamp {
1688    fn maybe_apply_temporal_bucketing<'scope, D>(
1689        stream: Stream<'scope, Self, Column<(D, Self, Diff)>>,
1690        as_of: Antichain<mz_repr::Timestamp>,
1691        summary: mz_repr::Timestamp,
1692    ) -> Collection<'scope, Self, Column<(D, Self, Diff)>>
1693    where
1694        D: differential_dataflow::ExchangeData
1695            + crate::typedefs::MzData
1696            + differential_dataflow::Hashable
1697            + ColumnarData,
1698        for<'a> ::columnar::Ref<'a, D>: Copy + Ord + std::hash::Hash,
1699        for<'a> ::columnar::Ref<'a, Self>: Copy + Ord,
1700        for<'a> ::columnar::Ref<'a, Diff>: Ord,
1701        for<'a> <(D, Self, Diff) as ColumnarData>::Container:
1702            ColumnarPush<&'a (D, Self, Diff)> + ColumnarPush<::columnar::Ref<'a, (D, Self, Diff)>>,
1703    {
1704        stream.bucket(as_of, summary).as_collection()
1705    }
1706
1707    fn maybe_apply_temporal_bucketing_vec<'scope, D>(
1708        stream: StreamVec<'scope, Self, (D, Self, Diff)>,
1709        as_of: Antichain<mz_repr::Timestamp>,
1710        summary: mz_repr::Timestamp,
1711    ) -> VecCollection<'scope, Self, D, Diff>
1712    where
1713        D: differential_dataflow::ExchangeData
1714            + crate::typedefs::MzData
1715            + differential_dataflow::Hashable
1716            + ColumnarData,
1717        for<'a> ::columnar::Ref<'a, D>: Copy + Ord + std::hash::Hash,
1718        for<'a> ::columnar::Ref<'a, Self>: Copy + Ord,
1719        for<'a> ::columnar::Ref<'a, Diff>: Ord,
1720        for<'a> <(D, Self, Diff) as ColumnarData>::Container:
1721            ColumnarPush<&'a (D, Self, Diff)> + ColumnarPush<::columnar::Ref<'a, (D, Self, Diff)>>,
1722    {
1723        stream.bucket(as_of, summary).as_collection()
1724    }
1725}
1726
1727impl RenderTimestamp for Product<mz_repr::Timestamp, PointStamp<u64>> {
1728    fn system_time(&mut self) -> &mut mz_repr::Timestamp {
1729        &mut self.outer
1730    }
1731    fn system_delay(delay: mz_repr::Timestamp) -> <Self as Timestamp>::Summary {
1732        Product::new(delay, Default::default())
1733    }
1734    fn event_time(&self) -> mz_repr::Timestamp {
1735        self.outer
1736    }
1737    fn event_time_mut(&mut self) -> &mut mz_repr::Timestamp {
1738        &mut self.outer
1739    }
1740    fn event_delay(delay: mz_repr::Timestamp) -> <Self as Timestamp>::Summary {
1741        Product::new(delay, Default::default())
1742    }
1743    fn step_back(&self) -> Self {
1744        // It is necessary to step back both coordinates of a product,
1745        // and when one is a `PointStamp` that also means all coordinates
1746        // of the pointstamp.
1747        let inner = self.inner.clone();
1748        let mut vec = inner.into_inner();
1749        for item in vec.iter_mut() {
1750            *item = item.saturating_sub(1);
1751        }
1752        Product::new(self.outer.saturating_sub(1), PointStamp::new(vec))
1753    }
1754}
1755
1756impl MaybeBucketByTime for Product<mz_repr::Timestamp, PointStamp<u64>> {
1757    fn maybe_apply_temporal_bucketing<'scope, D>(
1758        stream: Stream<'scope, Self, Column<(D, Self, Diff)>>,
1759        _as_of: Antichain<mz_repr::Timestamp>,
1760        _summary: mz_repr::Timestamp,
1761    ) -> Collection<'scope, Self, Column<(D, Self, Diff)>>
1762    where
1763        D: differential_dataflow::ExchangeData
1764            + crate::typedefs::MzData
1765            + differential_dataflow::Hashable
1766            + ColumnarData,
1767        for<'a> ::columnar::Ref<'a, D>: Copy + Ord + std::hash::Hash,
1768        for<'a> ::columnar::Ref<'a, Self>: Copy + Ord,
1769        for<'a> ::columnar::Ref<'a, Diff>: Ord,
1770        for<'a> <(D, Self, Diff) as ColumnarData>::Container:
1771            ColumnarPush<&'a (D, Self, Diff)> + ColumnarPush<::columnar::Ref<'a, (D, Self, Diff)>>,
1772    {
1773        // TODO: Implement bucketing on outer timestamp for iterative scopes.
1774        stream.as_collection()
1775    }
1776
1777    fn maybe_apply_temporal_bucketing_vec<'scope, D>(
1778        stream: StreamVec<'scope, Self, (D, Self, Diff)>,
1779        _as_of: Antichain<mz_repr::Timestamp>,
1780        _summary: mz_repr::Timestamp,
1781    ) -> VecCollection<'scope, Self, D, Diff>
1782    where
1783        D: differential_dataflow::ExchangeData
1784            + crate::typedefs::MzData
1785            + differential_dataflow::Hashable
1786            + ColumnarData,
1787        for<'a> ::columnar::Ref<'a, D>: Copy + Ord + std::hash::Hash,
1788        for<'a> ::columnar::Ref<'a, Self>: Copy + Ord,
1789        for<'a> ::columnar::Ref<'a, Diff>: Ord,
1790        for<'a> <(D, Self, Diff) as ColumnarData>::Container:
1791            ColumnarPush<&'a (D, Self, Diff)> + ColumnarPush<::columnar::Ref<'a, (D, Self, Diff)>>,
1792    {
1793        // TODO: Implement bucketing on outer timestamp for iterative scopes.
1794        stream.as_collection()
1795    }
1796}
1797
1798/// A signal that can be awaited by operators to suspend them prior to startup.
1799///
1800/// Creating a signal also yields a token, dropping of which causes the signal to fire.
1801///
1802/// `StartSignal` is designed to be usable by both async and sync Timely operators.
1803///
1804///  * Async operators can simply `await` it.
1805///  * Sync operators should register an [`ActivateOnDrop`] value via [`StartSignal::drop_on_fire`]
1806///    and then check `StartSignal::has_fired()` on each activation.
1807#[derive(Clone)]
1808pub(crate) struct StartSignal {
1809    /// A future that completes when the signal fires.
1810    ///
1811    /// The inner type is `Infallible` because no data is ever expected on this channel. Instead the
1812    /// signal is activated by dropping the corresponding `Sender`.
1813    fut: futures::future::Shared<oneshot::Receiver<Infallible>>,
1814    /// A weak reference to the token, to register drop-on-fire values.
1815    token_ref: Weak<RefCell<Box<dyn Any>>>,
1816}
1817
1818impl StartSignal {
1819    /// Create a new `StartSignal` and a corresponding token that activates the signal when
1820    /// dropped.
1821    pub fn new() -> (Self, Rc<dyn Any>) {
1822        let (tx, rx) = oneshot::channel::<Infallible>();
1823        let token: Rc<RefCell<Box<dyn Any>>> = Rc::new(RefCell::new(Box::new(tx)));
1824        let signal = Self {
1825            fut: rx.shared(),
1826            token_ref: Rc::downgrade(&token),
1827        };
1828        (signal, token)
1829    }
1830
1831    pub fn has_fired(&self) -> bool {
1832        self.token_ref.strong_count() == 0
1833    }
1834
1835    /// Returns a Send-safe future that completes when the signal fires.
1836    ///
1837    /// Unlike `StartSignal` itself, the returned future does not retain a reference to the token,
1838    /// so it cannot be used for `drop_on_fire` or `has_fired` checks.
1839    pub fn into_send_future(self) -> impl Future<Output = ()> + Send {
1840        use futures::FutureExt;
1841        self.fut.map(|_| ())
1842    }
1843
1844    pub fn drop_on_fire(&self, to_drop: Box<dyn Any>) {
1845        if let Some(token) = self.token_ref.upgrade() {
1846            let mut token = token.borrow_mut();
1847            let inner = std::mem::replace(&mut *token, Box::new(()));
1848            *token = Box::new((inner, to_drop));
1849        }
1850    }
1851}
1852
1853impl Future for StartSignal {
1854    type Output = ();
1855
1856    fn poll(mut self: Pin<&mut Self>, cx: &mut std::task::Context<'_>) -> Poll<Self::Output> {
1857        self.fut.poll_unpin(cx).map(|_| ())
1858    }
1859}
1860
1861/// Extension trait to attach a `StartSignal` to operator outputs.
1862pub(crate) trait WithStartSignal {
1863    /// Delays data and progress updates until the start signal has fired.
1864    ///
1865    /// Note that this operator needs to buffer all incoming data, so it has some memory footprint,
1866    /// depending on the amount and shape of its inputs.
1867    fn with_start_signal(self, signal: StartSignal) -> Self;
1868}
1869
1870impl<'scope, Tr> WithStartSignal for Arranged<'scope, Tr>
1871where
1872    Tr: TraceReader<Time: RenderTimestamp> + Clone,
1873{
1874    fn with_start_signal(self, signal: StartSignal) -> Self {
1875        Arranged {
1876            stream: self.stream.with_start_signal(signal),
1877            trace: self.trace,
1878        }
1879    }
1880}
1881
1882impl<'scope, T: Timestamp, D> WithStartSignal for Stream<'scope, T, D>
1883where
1884    D: timely::Container + Clone + 'static,
1885{
1886    fn with_start_signal(self, signal: StartSignal) -> Self {
1887        let activations = self.scope().activations();
1888        self.unary(Pipeline, "StartSignal", |_cap, info| {
1889            let token = Box::new(ActivateOnDrop::new((), info.address, activations));
1890            signal.drop_on_fire(token);
1891
1892            let mut stash = Vec::new();
1893
1894            move |input, output| {
1895                // Stash incoming updates as long as the start signal has not fired.
1896                if !signal.has_fired() {
1897                    input.for_each(|cap, data| stash.push((cap, std::mem::take(data))));
1898                    return;
1899                }
1900
1901                // Release any data we might still have stashed.
1902                for (cap, mut data) in std::mem::take(&mut stash) {
1903                    output.session(&cap).give_container(&mut data);
1904                }
1905
1906                // Pass through all remaining input data.
1907                input.for_each(|cap, data| {
1908                    output.session(&cap).give_container(data);
1909                });
1910            }
1911        })
1912    }
1913}
1914
1915/// Suppress progress messages for times before the given `as_of`.
1916///
1917/// This operator exists specifically to work around a memory spike we'd otherwise see when
1918/// hydrating arrangements (database-issues#6368). The memory spike happens because when the `arrange_core`
1919/// operator observes a frontier advancement without data it inserts an empty batch into the spine.
1920/// When it later inserts the snapshot batch into the spine, an empty batch is already there and
1921/// the spine initiates a merge of these batches, which requires allocating a new batch the size of
1922/// the snapshot batch.
1923///
1924/// The strategy to avoid the spike is to prevent the insertion of that initial empty batch by
1925/// ensuring that the first frontier advancement downstream `arrange_core` operators observe is
1926/// beyond the `as_of`, so the snapshot data has already been collected.
1927///
1928/// To ensure this, this operator needs to take two measures:
1929///  * Keep around a minimum capability until the input announces progress beyond the `as_of`.
1930///  * Reclock all updates emitted at times not beyond the `as_of` to the minimum time.
1931///
1932/// The second measure requires elaboration: If we wouldn't reclock snapshot updates, they might
1933/// still be upstream of `arrange_core` operators when those get to know about us dropping the
1934/// minimum capability. The in-flight snapshot updates would hold back the input frontiers of
1935/// `arrange_core` operators to the `as_of`, which would cause them to insert empty batches.
1936fn suppress_early_progress<'scope, T: Timestamp, D>(
1937    stream: Stream<'scope, T, D>,
1938    as_of: Antichain<T>,
1939) -> Stream<'scope, T, D>
1940where
1941    D: timely::Container + Clone,
1942{
1943    stream.unary_frontier(Pipeline, "SuppressEarlyProgress", |default_cap, _info| {
1944        let mut early_cap = Some(default_cap);
1945
1946        move |(input, frontier), output| {
1947            input.for_each_time(|data_cap, data| {
1948                if as_of.less_than(data_cap.time()) {
1949                    let mut session = output.session(&data_cap);
1950                    for data in data {
1951                        session.give_container(data);
1952                    }
1953                } else {
1954                    let cap = early_cap.as_ref().expect("early_cap can't be dropped yet");
1955                    let mut session = output.session(&cap);
1956                    for data in data {
1957                        session.give_container(data);
1958                    }
1959                }
1960            });
1961
1962            if !PartialOrder::less_equal(&frontier.frontier(), &as_of.borrow()) {
1963                early_cap.take();
1964            }
1965        }
1966    })
1967}
1968
1969/// Extension trait for [`Stream`] to selectively limit progress.
1970trait LimitProgress<T: Timestamp> {
1971    /// Limit the progress of the stream until its frontier reaches the given `upper` bound. Expects
1972    /// the implementation to observe times in data, and release capabilities based on the probe's
1973    /// frontier, after applying `slack` to round up timestamps.
1974    ///
1975    /// The implementation of this operator is subtle to avoid regressions in the rest of the
1976    /// system. Specifically joins hold back compaction on the other side of the join, so we need to
1977    /// make sure we release capabilities as soon as possible. This is why we only limit progress
1978    /// for times before the `upper`, which is the time until which the source can distinguish
1979    /// updates at the time of rendering. Once we make progress to the `upper`, we need to release
1980    /// our capability.
1981    ///
1982    /// This isn't perfect, and can result in regressions if on of the inputs lags behind. We could
1983    /// consider using the join of the uppers, i.e, use lower bound upper of all available inputs.
1984    ///
1985    /// Once the input frontier reaches `[]`, the implementation must release any capability to
1986    /// allow downstream operators to release resources.
1987    ///
1988    /// The implementation should limit the number of pending times to `limit` if it is `Some` to
1989    /// avoid unbounded memory usage.
1990    ///
1991    /// * `handle` is a probe installed on the dataflow's outputs as late as possible, but before
1992    ///   any timestamp rounding happens (c.f., `REFRESH EVERY` materialized views).
1993    /// * `slack_ms` is the number of milliseconds to round up timestamps to.
1994    /// * `name` is a human-readable name for the operator.
1995    /// * `limit` is the maximum number of pending times to keep around.
1996    /// * `upper` is the upper bound of the stream's frontier until which the implementation can
1997    ///   retain a capability.
1998    fn limit_progress(
1999        self,
2000        handle: MzProbeHandle<T>,
2001        slack_ms: u64,
2002        limit: Option<usize>,
2003        upper: Antichain<T>,
2004        name: String,
2005    ) -> Self;
2006}
2007
2008/// Reads the times of the records a container holds, for [`LimitProgress`].
2009trait RecordTimes {
2010    /// Call `f` once per record, with that record's time.
2011    fn for_each_time(&self, f: impl FnMut(mz_repr::Timestamp));
2012}
2013
2014impl<D, R> RecordTimes for Vec<(D, mz_repr::Timestamp, R)> {
2015    fn for_each_time(&self, mut f: impl FnMut(mz_repr::Timestamp)) {
2016        for (_, time, _) in self {
2017            f(*time);
2018        }
2019    }
2020}
2021
2022impl<D, R> RecordTimes for Column<(D, mz_repr::Timestamp, R)>
2023where
2024    D: ColumnarData,
2025    R: ColumnarData,
2026    (D, mz_repr::Timestamp, R): ColumnarData<
2027        Container = (
2028            D::Container,
2029            <mz_repr::Timestamp as ColumnarData>::Container,
2030            R::Container,
2031        ),
2032    >,
2033{
2034    fn for_each_time(&self, mut f: impl FnMut(mz_repr::Timestamp)) {
2035        for time in self.borrow().1.into_index_iter() {
2036            f(time);
2037        }
2038    }
2039}
2040
2041// TODO: We could make this generic over a `T` that can be converted to and from a u64 millisecond
2042// number.
2043impl<'scope, C> LimitProgress<mz_repr::Timestamp> for Stream<'scope, mz_repr::Timestamp, C>
2044where
2045    C: timely::Container + Clone + RecordTimes,
2046{
2047    fn limit_progress(
2048        self,
2049        handle: MzProbeHandle<mz_repr::Timestamp>,
2050        slack_ms: u64,
2051        limit: Option<usize>,
2052        upper: Antichain<mz_repr::Timestamp>,
2053        name: String,
2054    ) -> Self {
2055        let scope = self.scope();
2056        let stream =
2057            self.unary_frontier(Pipeline, &format!("LimitProgress({name})"), |_cap, info| {
2058                // Times that we've observed on our input.
2059                let mut pending_times: BTreeSet<mz_repr::Timestamp> = BTreeSet::new();
2060                // Capability for the lower bound of `pending_times`, if any.
2061                let mut retained_cap: Option<Capability<mz_repr::Timestamp>> = None;
2062
2063                let activator = scope.activator_for(info.address);
2064                handle.activate(activator.clone());
2065
2066                move |(input, frontier), output| {
2067                    input.for_each(|cap, data| {
2068                        data.for_each_time(|time| {
2069                            let Some(time) = u64::from(time).checked_add(slack_ms) else {
2070                                return;
2071                            };
2072                            // `slack_ms == 0` means no rounding; otherwise round up to the next
2073                            // multiple of `slack_ms`. Avoids a divide-by-zero panic when the
2074                            // operator is configured without slack.
2075                            let rounded_time = if slack_ms == 0 {
2076                                time
2077                            } else {
2078                                (time / slack_ms).saturating_add(1).saturating_mul(slack_ms)
2079                            };
2080                            if !upper.less_than(&rounded_time.into()) {
2081                                pending_times.insert(rounded_time.into());
2082                            }
2083                        });
2084                        output.session(&cap).give_container(data);
2085                        if retained_cap.as_ref().is_none_or(|c| {
2086                            !c.time().less_than(cap.time()) && !upper.less_than(cap.time())
2087                        }) {
2088                            retained_cap = Some(cap.retain(0));
2089                        }
2090                    });
2091
2092                    handle.with_frontier(|f| {
2093                        while pending_times
2094                            .first()
2095                            .map_or(false, |retained_time| !f.less_than(&retained_time))
2096                        {
2097                            let _ = pending_times.pop_first();
2098                        }
2099                    });
2100
2101                    while limit.map_or(false, |limit| pending_times.len() > limit) {
2102                        let _ = pending_times.pop_first();
2103                    }
2104
2105                    match (retained_cap.as_mut(), pending_times.first()) {
2106                        (Some(cap), Some(first)) => cap.downgrade(first),
2107                        (_, None) => retained_cap = None,
2108                        _ => {}
2109                    }
2110
2111                    if frontier.is_empty() {
2112                        retained_cap = None;
2113                        pending_times.clear();
2114                    }
2115
2116                    if !pending_times.is_empty() {
2117                        tracing::debug!(
2118                            name,
2119                            info.global_id,
2120                            pending_times = %PendingTimesDisplay(pending_times.iter().cloned()),
2121                            frontier = ?frontier.frontier().get(0),
2122                            probe = ?handle.with_frontier(|f| f.get(0).cloned()),
2123                            ?upper,
2124                            "pending times",
2125                        );
2126                    }
2127                }
2128            });
2129        stream
2130    }
2131}
2132
2133/// A formatter for an iterator of timestamps that displays the first element, and subsequently
2134/// the difference between timestamps.
2135struct PendingTimesDisplay<T>(T);
2136
2137impl<T> std::fmt::Display for PendingTimesDisplay<T>
2138where
2139    T: IntoIterator<Item = mz_repr::Timestamp> + Clone,
2140{
2141    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
2142        let mut iter = self.0.clone().into_iter();
2143        write!(f, "[")?;
2144        if let Some(first) = iter.next() {
2145            write!(f, "{}", first)?;
2146            let mut last = u64::from(first);
2147            for time in iter {
2148                write!(f, ", +{}", u64::from(time) - last)?;
2149                last = u64::from(time);
2150            }
2151        }
2152        write!(f, "]")?;
2153        Ok(())
2154    }
2155}
2156
2157/// Helper to merge pairs of datum iterators into a row or split a datum iterator
2158/// into two rows, given the arity of the first component.
2159#[derive(Clone, Copy, Debug)]
2160struct Pairer {
2161    split_arity: usize,
2162}
2163
2164impl Pairer {
2165    /// Creates a pairer with knowledge of the arity of first component in the pair.
2166    fn new(split_arity: usize) -> Self {
2167        Self { split_arity }
2168    }
2169
2170    /// Merges a pair of datum iterators creating a `Row` instance.
2171    fn merge<'a, I1, I2>(&self, first: I1, second: I2) -> Row
2172    where
2173        I1: IntoIterator<Item = Datum<'a>>,
2174        I2: IntoIterator<Item = Datum<'a>>,
2175    {
2176        SharedRow::pack(first.into_iter().chain(second))
2177    }
2178
2179    /// Splits a datum iterator into a pair of `Row` instances.
2180    fn split<'a>(&self, datum_iter: impl IntoIterator<Item = Datum<'a>>) -> (Row, Row) {
2181        let mut datum_iter = datum_iter.into_iter();
2182        let mut row_builder = SharedRow::get();
2183        let first = row_builder.pack_using(datum_iter.by_ref().take(self.split_arity));
2184        let second = row_builder.pack_using(datum_iter);
2185        (first, second)
2186    }
2187}