Skip to main content

mz_compute/render/join/
delta_join.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//! Delta join execution dataflow construction.
11//!
12//! Consult [DeltaJoinPlan] documentation for details.
13
14#![allow(clippy::op_ref)]
15
16use std::collections::BTreeSet;
17
18use differential_dataflow::consolidation::ConsolidatingContainerBuilder;
19use differential_dataflow::operators::arrange::Arranged;
20use differential_dataflow::trace::cursor::BatchCursor;
21use differential_dataflow::trace::implementations::BatchContainer;
22use differential_dataflow::trace::{Cursor, Navigable, TraceReader};
23use differential_dataflow::{AsCollection, VecCollection};
24use mz_compute_types::plan::join::JoinClosure;
25use mz_compute_types::plan::join::delta_join::{DeltaJoinPlan, DeltaPathPlan, DeltaStagePlan};
26use mz_compute_types::plan::scalar::LirScalarExpr;
27use mz_expr::Eval;
28use mz_repr::fixed_length::ExtendDatums;
29use mz_repr::{DatumVec, Diff, Row, RowArena, SharedRow};
30use mz_timely_util::operator::{CollectionExt, StreamExt};
31use timely::container::CapacityContainerBuilder;
32use timely::dataflow::channels::pact::Pipeline;
33use timely::dataflow::operators::OkErr;
34use timely::dataflow::operators::vec::Map;
35use timely::progress::Antichain;
36
37use crate::render::RenderTimestamp;
38use crate::render::columnar::{ColCollection, flat_map_datums, vec_to_columnar};
39use crate::render::context::{ArrangementFlavor, CollectionBundle, Context};
40use crate::render::errors::DataflowErrorSer;
41use crate::typedefs::{RowRowAgent, RowRowEnter};
42
43impl<'scope, T: RenderTimestamp> Context<'scope, T> {
44    /// Renders `MirRelationExpr:Join` using dogs^3 delta query dataflows.
45    ///
46    /// The join is followed by the application of `map_filter_project`, whose
47    /// implementation will be pushed in to the join pipeline if at all possible.
48    pub fn render_delta_join(
49        &self,
50        inputs: Vec<CollectionBundle<'scope, T>>,
51        join_plan: DeltaJoinPlan,
52    ) -> CollectionBundle<'scope, T> {
53        // We create a new region to contain the dataflow paths for the delta join.
54        let (oks, errs) = self.scope.clone().region_named("Join(Delta)", |inner| {
55            // Our plan is to iterate through each input relation, and attempt
56            // to find a plan that maximally uses existing keys (better: uses
57            // existing arrangements, to which we have access).
58            let mut join_results = Vec::new();
59
60            // An input bundle may carry collections no delta path consumes: arrangements keyed
61            // differently than any lookup, or a raw collection. Entering a collection into a region
62            // does work proportional to its data (and the raw collection is not reference counted),
63            // so we prune each bundle to what some path reads before bringing it into `inner`.
64            let mut arrangements = vec![BTreeSet::new(); inputs.len()];
65            let mut raw = vec![false; inputs.len()];
66            for path_plan in &join_plan.path_plans {
67                record_path_arrangements(
68                    &mut arrangements,
69                    &mut raw,
70                    path_plan.source_relation,
71                    path_plan.source_key.as_ref(),
72                    &path_plan.stage_plans,
73                );
74            }
75            let inputs = inputs
76                .iter()
77                .enumerate()
78                .map(|(index, cb)| {
79                    prune_bundle(cb, raw[index], &arrangements[index]).enter_region(inner)
80                })
81                .collect::<Vec<_>>();
82
83            // Collects error streams for the ambient scope, seeded with the errors the inputs
84            // arrive with. Every path reads every input, but an input's pre-existing errors belong
85            // in the output once, not once per path. See [`bundle_errs`].
86            let mut inner_errs: Vec<_> = inputs.iter().flat_map(bundle_errs).collect();
87
88            for path_plan in join_plan.path_plans {
89                // Deconstruct the stages of the path plan.
90                let DeltaPathPlan {
91                    source_relation,
92                    initial_closure,
93                    stage_plans,
94                    final_closure,
95                    source_key,
96                } = path_plan;
97
98                // This collection determines changes that result from updates inbound
99                // from `inputs[relation]` and reflects all strictly prior updates and
100                // concurrent updates from relations prior to `relation`.
101                let name = format!("delta path {}", source_relation);
102                let path_results = inner.clone().region_named(&name, |region| {
103                    // The plan is to move through each relation, starting from `relation` and in the order
104                    // indicated in `orders[relation]`. At each moment, we will have the columns from the
105                    // subset of relations encountered so far, and we will have applied as much as we can
106                    // of the filters in `equivalences` and the logic in `map_filter_project`, based on the
107                    // available columns.
108                    //
109                    // As we go, we will track the physical locations of each intended output column, as well
110                    // as the locations of intermediate results from partial application of `map_filter_project`.
111                    //
112                    // Just before we apply the `lookup` function to perform a join, we will first use our
113                    // available information to determine the filtering and logic that we can apply, and
114                    // introduce that in to the `lookup` logic to cause it to happen in that operator.
115
116                    // Prune each bundle to just what this path reads before entering the path's
117                    // region, for the reason noted at the join's region entry above.
118                    let mut path_arrangements = vec![BTreeSet::new(); inputs.len()];
119                    let mut path_raw = vec![false; inputs.len()];
120                    record_path_arrangements(
121                        &mut path_arrangements,
122                        &mut path_raw,
123                        source_relation,
124                        source_key.as_ref(),
125                        &stage_plans,
126                    );
127                    let bundles = inputs
128                        .iter()
129                        .enumerate()
130                        .map(|(index, cb)| {
131                            prune_bundle(cb, path_raw[index], &path_arrangements[index])
132                                .enter_region(region)
133                        })
134                        .collect::<Vec<_>>();
135
136                    // Collects error streams for the region scope. Concats before leaving.
137                    let mut region_errs = Vec::with_capacity(inputs.len());
138
139                    // Form the initial stream of updates that will hydrate the delta path.
140                    let (update_stream, err_stream) = build_update_stream(
141                        &bundles[source_relation],
142                        self.as_of_frontier.clone(),
143                        source_key,
144                        source_relation,
145                        initial_closure,
146                    );
147                    region_errs.push(err_stream);
148
149                    // Promote `time` to a datum element.
150                    //
151                    // The `half_join` operator manipulates as "data" a pair `(data, time)`,
152                    // while tracking the initial time `init_time` separately and without
153                    // modification. The initial value for both times is the initial time.
154                    let mut update_stream = update_stream
155                        .inner
156                        .map(|(v, t, d)| ((v, t.clone()), t, d))
157                        .as_collection();
158
159                    // Repeatedly update `update_stream` to reflect joins with more and more
160                    // other relations, in the specified order.
161                    for stage_plan in stage_plans {
162                        let DeltaStagePlan {
163                            lookup_relation,
164                            stream_key,
165                            stream_thinning,
166                            lookup_key,
167                            closure,
168                        } = stage_plan;
169
170                        // We require different logic based on the relative order of the two inputs.
171                        // If the `source` relation precedes the `lookup` relation, we present all
172                        // updates with less or equal `time`, and otherwise we present only updates
173                        // with strictly less `time`.
174                        //
175                        // We require demuxing over the two flavors of arrangement and over the
176                        // relative order of the inputs. Both are handled inside `build_halfjoin`.
177                        let (oks, errs) = build_halfjoin(
178                            update_stream,
179                            stream_key,
180                            stream_thinning,
181                            &bundles[lookup_relation],
182                            lookup_key,
183                            source_relation < lookup_relation,
184                            closure,
185                        );
186                        update_stream = oks;
187                        region_errs.push(errs);
188                    }
189
190                    // Delay updates as appropriate.
191                    //
192                    // The `half_join` operator maintains a time that we now discard (the `_`),
193                    // and replace with the `time` that is maintained with the data. The former
194                    // exists to pin a consistent total order on updates throughout the process,
195                    // while allowing `time` to vary upwards as a result of actions on time.
196                    let mut update_stream = update_stream
197                        .inner
198                        .map(|((row, time), _, diff)| (row, time, diff))
199                        .as_collection();
200
201                    // We have completed the join building, but may have work remaining.
202                    // For example, we may have expressions not pushed down (e.g. literals)
203                    // and projections that could not be applied (e.g. column repetition).
204                    if let Some(final_closure) = final_closure {
205                        let name = "DeltaJoinFinalization";
206                        type CB<C> = ConsolidatingContainerBuilder<C>;
207                        let (updates, errors) = update_stream
208                            .flat_map_fallible::<CB<_>, CB<_>, _, _, _, _>(name, {
209                                // Reuseable allocation for unpacking.
210                                let mut datums = DatumVec::new();
211                                move |row| {
212                                    let mut row_builder = SharedRow::get();
213                                    let temp_storage = RowArena::new();
214                                    let mut datums_local = datums.borrow_with(&row);
215                                    // TODO(mcsherry): re-use `row` allocation.
216                                    final_closure
217                                        .apply(&mut datums_local, &temp_storage, &mut row_builder)
218                                        .map(|row| row.cloned())
219                                        .map_err(DataflowErrorSer::from)
220                                        .transpose()
221                                }
222                            });
223
224                        update_stream = updates;
225                        region_errs.push(errors);
226                    }
227
228                    inner_errs.push(
229                        differential_dataflow::collection::concatenate(region, region_errs)
230                            .leave_region(inner),
231                    );
232                    update_stream.leave_region(inner)
233                });
234
235                join_results.push(path_results);
236            }
237
238            // Concatenate the results of each delta query as the accumulated results.
239            (
240                differential_dataflow::collection::concatenate(inner, join_results)
241                    .leave_region(self.scope),
242                differential_dataflow::collection::concatenate(inner, inner_errs)
243                    .leave_region(self.scope),
244            )
245        });
246        // The delta join is `Vec`-internal throughout, so the concatenated node output is
247        // encoded once here. Non-consolidating, because the per-path finalization already
248        // consolidated whatever it consolidates.
249        CollectionBundle::from_edge(vec_to_columnar(oks), errs)
250    }
251}
252
253/// Records what a single delta path reads from each input: each stage's `lookup_key` and, for the
254/// seed relation, either its `source_key` arrangement or, when `source_key` is `None`, its raw
255/// collection (flagged in `raw`).
256fn record_path_arrangements(
257    arrangements: &mut [BTreeSet<Vec<LirScalarExpr>>],
258    raw: &mut [bool],
259    source_relation: usize,
260    source_key: Option<&Vec<LirScalarExpr>>,
261    stage_plans: &[DeltaStagePlan],
262) {
263    match source_key {
264        Some(source_key) => {
265            arrangements[source_relation].insert(source_key.clone());
266        }
267        None => raw[source_relation] = true,
268    }
269    for stage_plan in stage_plans {
270        arrangements[stage_plan.lookup_relation].insert(stage_plan.lookup_key.clone());
271    }
272}
273
274/// Retains only the collections of `bundle` a delta path reads: the arrangements keyed by `keys`,
275/// and the raw collection if `raw`. Every other collection is dropped.
276fn prune_bundle<'scope, T: RenderTimestamp>(
277    bundle: &CollectionBundle<'scope, T>,
278    raw: bool,
279    keys: &BTreeSet<Vec<LirScalarExpr>>,
280) -> CollectionBundle<'scope, T> {
281    CollectionBundle {
282        collection: if raw { bundle.collection.clone() } else { None },
283        arranged: bundle
284            .arranged
285            .iter()
286            .filter(|(key, _)| keys.contains(*key))
287            .map(|(key, flavor)| (key.clone(), flavor.clone()))
288            .collect(),
289    }
290}
291
292/// The error collections of every collection and arrangement `bundle` holds.
293///
294/// A delta join reads every input from every one of its paths, but an input's pre-existing errors
295/// belong in the join's error output once. Propagating them per path would multiply their
296/// multiplicities by the number of paths, and because a join's output is another join's input, those
297/// factors compound multiplicatively through a nested plan until the `Diff` overflows. Error
298/// semantics depend only on presence, so the extra copies buy nothing.
299///
300/// Expects the pruned bundle (see [`prune_bundle`]), so that it yields errors only for the
301/// collections some path actually reads.
302///
303/// NOTE: This bounds an input's errors to one copy per retained form, not to one copy outright.
304/// Each form carries its own error collection, and an arrangement's is a distinct stream built from
305/// the raw one plus that key's key-formation errors, so an input the join reads under two lookup
306/// keys still contributes its errors twice. Reachable whenever a delta path set needs an
307/// error-carrying input arranged by more than one key.
308fn bundle_errs<'scope, T: RenderTimestamp>(
309    bundle: &CollectionBundle<'scope, T>,
310) -> Vec<VecCollection<'scope, T, DataflowErrorSer, Diff>> {
311    let mut collected = Vec::with_capacity(bundle.arranged.len() + 1);
312    if let Some((_oks, errs)) = &bundle.collection {
313        collected.push(errs.clone());
314    }
315    for flavor in bundle.arranged.values() {
316        let errs = match flavor {
317            ArrangementFlavor::Local(_oks, errs) => errs.clone().as_collection(|k, _v| k.clone()),
318            ArrangementFlavor::Trace(_id, _oks, errs) => {
319                errs.clone().as_collection(|k, _v| k.clone())
320            }
321        };
322        collected.push(errs);
323    }
324    collected
325}
326
327/// Constructs a `half_join` against the arrangement held by a collection bundle.
328///
329/// This wrapper demuxes over the two flavors of arrangement (dataflow-local or imported trace)
330/// that the bundle might hold for `lookup_key`, dispatching to the generic [`build_halfjoin_trace`]
331/// for each. `source_precedes_lookup` selects the tie-breaking comparison: `le` if the source
332/// relation precedes the lookup relation in the total order on relations, otherwise `lt`.
333///
334/// The returned error collection holds only the errors this stage produces. The errors `bundle`
335/// already carries are the caller's to propagate, once, rather than once per delta path that looks
336/// the input up. See [`bundle_errs`].
337fn build_halfjoin<'scope, T>(
338    updates: VecCollection<'scope, T, (Row, T), Diff>,
339    prev_key: Vec<LirScalarExpr>,
340    prev_thinning: Vec<usize>,
341    bundle: &CollectionBundle<'scope, T>,
342    lookup_key: Vec<LirScalarExpr>,
343    source_precedes_lookup: bool,
344    closure: JoinClosure,
345) -> (
346    VecCollection<'scope, T, (Row, T), Diff>,
347    VecCollection<'scope, T, DataflowErrorSer, Diff>,
348)
349where
350    T: RenderTimestamp,
351{
352    match bundle.arrangement(&lookup_key) {
353        Some(ArrangementFlavor::Local(oks, _errs)) => {
354            let (oks, errs2) = build_halfjoin_trace::<_, RowRowAgent<_, _>>(
355                updates,
356                oks,
357                prev_key,
358                prev_thinning,
359                !source_precedes_lookup,
360                closure,
361            );
362            (oks, errs2)
363        }
364        Some(ArrangementFlavor::Trace(_, oks, _errs)) => {
365            let (oks, errs2) = build_halfjoin_trace::<_, RowRowEnter<_, _, _>>(
366                updates,
367                oks,
368                prev_key,
369                prev_thinning,
370                !source_precedes_lookup,
371                closure,
372            );
373            (oks, errs2)
374        }
375        None => panic!("Arrangement promised by the planner is absent!"),
376    }
377}
378
379/// Constructs a `half_join` from supplied arguments.
380///
381/// This method exists to factor common logic from the code paths that are generic over the type of
382/// trace. `strict` breaks ties consistently: it is false when the source relation precedes the
383/// lookup relation in the total order on relations, and true otherwise. A delta query pairs a
384/// strict operator with a non-strict one, so each pair of matching updates interacts exactly once.
385///
386/// The input and output streams are of pairs `(data, time)` where the `time` component can be greater than
387/// the time of the update. This operator may manipulate `time` as part of this pair, but will not manipulate
388/// the time of the update. This is crucial for correctness, as the total order on times of updates is used
389/// to ensure that any two updates are matched at most once.
390fn build_halfjoin_trace<'scope, T, Tr>(
391    updates: VecCollection<'scope, T, (Row, T), Diff>,
392    trace: Arranged<'scope, Tr>,
393    prev_key: Vec<LirScalarExpr>,
394    prev_thinning: Vec<usize>,
395    strict: bool,
396    closure: JoinClosure,
397) -> (
398    VecCollection<'scope, T, (Row, T), Diff>,
399    VecCollection<'scope, T, DataflowErrorSer, Diff>,
400)
401where
402    T: RenderTimestamp,
403    Tr: TraceReader<Batch: Navigable, Time = T> + Clone + 'static,
404    for<'a> BatchCursor<Tr>: Cursor<
405            Val<'a>: ExtendDatums,
406            KeyContainer: BatchContainer<Owned = Row>,
407            Time = T,
408            Diff = Diff,
409        >,
410{
411    let name = "DeltaJoinKeyPreparation";
412    type CB<C> = CapacityContainerBuilder<C>;
413    let (updates, errs) = updates.map_fallible::<CB<_>, CB<_>, _, _, _>(name, {
414        // Reuseable allocation for unpacking.
415        let mut datums = DatumVec::new();
416        move |(row, time)| {
417            let temp_storage = RowArena::new();
418            let datums_local = datums.borrow_with(&row);
419            let mut row_builder = SharedRow::get();
420            row_builder.packer().try_extend(
421                prev_key
422                    .iter()
423                    .map(|e| e.eval(&datums_local, &temp_storage)),
424            )?;
425            let key = row_builder.clone();
426            row_builder
427                .packer()
428                .extend(prev_thinning.iter().map(|&c| datums_local[c]));
429            let row_value = row_builder.clone();
430
431            Ok((key, row_value, time))
432        }
433    });
434    let datums = DatumVec::new();
435
436    build_halfjoin_op(updates, trace, strict, closure, datums, errs)
437}
438
439/// Renders the half join of `updates` against `trace`, splitting errors from the closure off.
440fn build_halfjoin_op<'scope, T, Tr>(
441    updates: VecCollection<'scope, T, (Row, Row, T), Diff>,
442    trace: Arranged<'scope, Tr>,
443    strict: bool,
444    closure: JoinClosure,
445    mut datums: DatumVec,
446    errs: VecCollection<'scope, T, DataflowErrorSer, Diff>,
447) -> (
448    VecCollection<'scope, T, (Row, T), Diff>,
449    VecCollection<'scope, T, DataflowErrorSer, Diff>,
450)
451where
452    T: RenderTimestamp,
453    Tr: TraceReader<Batch: Navigable, Time = T> + Clone + 'static,
454    for<'a> BatchCursor<Tr>: Cursor<
455            Val<'a>: ExtendDatums,
456            KeyContainer: BatchContainer<Owned = Row>,
457            Time = T,
458            Diff = Diff,
459        >,
460{
461    type CB<C> = CapacityContainerBuilder<C>;
462
463    if closure.could_error() {
464        let (oks, errs2) =
465            differential_dogs3::operators::half_join::cursors::half_join_internal_unsafe(
466                updates,
467                trace,
468                |time, antichain| {
469                    antichain.insert(time.step_back());
470                },
471                strict,
472                // TODO(mcsherry): investigate/establish trade-offs here; time based had problems,
473                // in that we seem to yield too much and do too little work when we do.
474                |_timer, count| count > 1_000_000,
475                // TODO(mcsherry): consider `RefOrMut` in `half_join` interface to allow re-use.
476                move |session: &mut CB<Vec<_>>,
477                      key,
478                      stream_row,
479                      lookup_row,
480                      initial,
481                      diff1,
482                      output| {
483                    let mut row_builder = SharedRow::get();
484                    let temp_storage = RowArena::new();
485
486                    let mut datums_local = datums.borrow();
487                    datums_local.extend(key.iter());
488                    datums_local.extend(stream_row.iter());
489                    lookup_row.extend_datums(&temp_storage, &mut datums_local, None);
490
491                    let row = closure.apply(&mut datums_local, &temp_storage, &mut row_builder);
492
493                    for (time, diff2) in output.drain(..) {
494                        let row = row.as_ref().map(|row| row.cloned()).map_err(Clone::clone);
495                        let diff = diff1.clone() * diff2.clone();
496                        let data = ((row, time.clone()), initial.clone(), diff);
497                        use timely::container::PushInto;
498                        session.push_into(data);
499                    }
500                },
501            )
502            .ok_err(|(data_time, init_time, diff)| {
503                // TODO(mcsherry): consider `ok_err()` for `Collection`.
504                match data_time {
505                    (Ok(data), time) => Ok((data.map(|data| (data, time)), init_time, diff)),
506                    (Err(err), _time) => Err((DataflowErrorSer::from(err), init_time, diff)),
507                }
508            });
509
510        (
511            oks.as_collection().flat_map(|x| x),
512            errs.concat(errs2.as_collection()),
513        )
514    } else {
515        let oks = differential_dogs3::operators::half_join::cursors::half_join_internal_unsafe(
516            updates,
517            trace,
518            |time, antichain| {
519                antichain.insert(time.step_back());
520            },
521            strict,
522            // TODO(mcsherry): investigate/establish trade-offs here; time based had problems,
523            // in that we seem to yield too much and do too little work when we do.
524            |_timer, count| count > 1_000_000,
525            // TODO(mcsherry): consider `RefOrMut` in `half_join` interface to allow re-use.
526            move |session: &mut CB<Vec<_>>, key, stream_row, lookup_row, initial, diff1, output| {
527                if output.is_empty() {
528                    return;
529                }
530
531                let mut row_builder = SharedRow::get();
532                let temp_storage = RowArena::new();
533
534                let mut datums_local = datums.borrow();
535                datums_local.extend(key.iter());
536                datums_local.extend(stream_row.iter());
537                lookup_row.extend_datums(&temp_storage, &mut datums_local, None);
538
539                if let Some(row) = closure
540                    .apply(&mut datums_local, &temp_storage, &mut row_builder)
541                    .expect("Closure claimed to never error")
542                {
543                    for (time, diff2) in output.drain(..) {
544                        let diff = diff1.clone() * diff2.clone();
545                        use timely::container::PushInto;
546                        session.push_into(((row.clone(), time.clone()), initial.clone(), diff));
547                    }
548                }
549            },
550        );
551
552        (oks.as_collection(), errs)
553    }
554}
555
556/// Builds the initial update stream of a delta path from a collection bundle.
557///
558/// With a `source_key`, demuxes over the two flavors of arrangement the bundle might hold for that
559/// key, dispatching to the generic [`build_update_stream_trace`]. Without a `source_key`, the source
560/// relation is consumed as a raw (unarranged) collection via [`build_update_stream_stream`].
561///
562/// The returned error collection holds only the errors the initial closure produces. The errors
563/// `bundle` already carries are the caller's to propagate, once, rather than once per delta path.
564/// See [`bundle_errs`].
565fn build_update_stream<'scope, T>(
566    bundle: &CollectionBundle<'scope, T>,
567    as_of: Antichain<mz_repr::Timestamp>,
568    source_key: Option<Vec<LirScalarExpr>>,
569    source_relation: usize,
570    initial_closure: JoinClosure,
571) -> (
572    VecCollection<'scope, T, Row, Diff>,
573    VecCollection<'scope, T, DataflowErrorSer, Diff>,
574)
575where
576    T: RenderTimestamp,
577{
578    let Some(source_key) = source_key else {
579        // No source key means a single-time dataflow (e.g. a `SELECT`) whose plan was truncated to
580        // this one path, letting us hydrate the source from its raw collection instead of an
581        // arrangement.
582        let (oks, _errs) = bundle
583            .collection
584            .clone()
585            .expect("The unarranged collection doesn't exist.");
586        return build_update_stream_stream(oks, as_of, source_relation, initial_closure);
587    };
588    match bundle.arrangement(&source_key) {
589        Some(ArrangementFlavor::Local(oks, _errs)) => {
590            build_update_stream_trace::<_, RowRowAgent<_, _>>(
591                oks,
592                as_of,
593                source_relation,
594                initial_closure,
595            )
596        }
597        Some(ArrangementFlavor::Trace(_, oks, _errs)) => {
598            build_update_stream_trace::<_, RowRowEnter<_, _, _>>(
599                oks,
600                as_of,
601                source_relation,
602                initial_closure,
603            )
604        }
605        None => panic!("Arrangement promised by the planner is absent!"),
606    }
607}
608
609/// Builds the beginning of the update stream of a delta path.
610///
611/// At start-up time only the delta path for the first relation sees updates, since any updates fed to the
612/// other delta paths would be discarded anyway due to the tie-breaking logic that avoids double-counting
613/// updates happening at the same time on different relations.
614fn build_update_stream_trace<'scope, T, Tr>(
615    trace: Arranged<'scope, Tr>,
616    as_of: Antichain<mz_repr::Timestamp>,
617    source_relation: usize,
618    initial_closure: JoinClosure,
619) -> (
620    VecCollection<'scope, T, Row, Diff>,
621    VecCollection<'scope, T, DataflowErrorSer, Diff>,
622)
623where
624    T: RenderTimestamp,
625    for<'a, 'b> &'a T: PartialEq<<BatchCursor<Tr> as Cursor>::TimeGat<'b>>,
626    Tr: TraceReader<Batch: Navigable, Time = T> + Clone + 'static,
627    for<'a> BatchCursor<Tr>:
628        Cursor<Key<'a>: ExtendDatums, Val<'a>: ExtendDatums, Time = T, Diff = Diff>,
629{
630    let mut inner_as_of = Antichain::new();
631    for event_time in as_of.elements().iter() {
632        inner_as_of.insert(<T>::to_inner(event_time.clone()));
633    }
634
635    let (ok_stream, err_stream) =
636        trace
637            .stream
638            .unary_fallible(Pipeline, "UpdateStream", move |_, _| {
639                let mut datums = DatumVec::new();
640                Box::new(move |input, ok_output, err_output| {
641                    // Buffer to accumulate contributing (time, diff) pairs for each (key, val).
642                    let mut times_diffs = Vec::default();
643                    input.for_each(|time, data| {
644                        let mut row_builder = SharedRow::get();
645                        let mut ok_session = ok_output.session(&time);
646                        let mut err_session = err_output.session(&time);
647
648                        for span in data.iter() {
649                            // A span with no updates carries no batch to walk.
650                            let Some(batch) = &span.inner else {
651                                continue;
652                            };
653                            let mut cursor = batch.cursor();
654                            while let Some(key) = cursor.get_key(batch) {
655                                while let Some(val) = cursor.get_val(batch) {
656                                    // Collect contributing (time, diff) pairs before invoking the closure.
657                                    cursor.map_times(batch, |time, diff| {
658                                        if source_relation == 0
659                                            || inner_as_of.elements().iter().all(|e| e != time)
660                                        {
661                                            // TODO: Consolidate as we push, defensively.
662                                            times_diffs.push((
663                                                <BatchCursor<Tr> as Cursor>::owned_time(time),
664                                                <BatchCursor<Tr> as Cursor>::owned_diff(diff),
665                                            ));
666                                        }
667                                    });
668                                    differential_dataflow::consolidation::consolidate(
669                                        &mut times_diffs,
670                                    );
671                                    // The can not-uncommonly be empty, if the inbound updates cancel.
672                                    if !times_diffs.is_empty() {
673                                        let temp_storage = RowArena::new();
674
675                                        let mut datums_local = datums.borrow();
676                                        key.extend_datums(&temp_storage, &mut datums_local, None);
677                                        val.extend_datums(&temp_storage, &mut datums_local, None);
678
679                                        if !initial_closure.is_identity() {
680                                            match initial_closure
681                                                .apply(
682                                                    &mut datums_local,
683                                                    &temp_storage,
684                                                    &mut row_builder,
685                                                )
686                                                .map(|row| row.cloned())
687                                                .transpose()
688                                            {
689                                                Some(Ok(row)) => {
690                                                    for (time, diff) in times_diffs.drain(..) {
691                                                        ok_session.give((row.clone(), time, diff))
692                                                    }
693                                                }
694                                                Some(Err(err)) => {
695                                                    for (time, diff) in times_diffs.drain(..) {
696                                                        err_session.give((err.clone(), time, diff))
697                                                    }
698                                                }
699                                                None => {}
700                                            }
701                                        } else {
702                                            let row = {
703                                                row_builder.packer().extend(&*datums_local);
704                                                row_builder.clone()
705                                            };
706                                            for (time, diff) in times_diffs.drain(..) {
707                                                ok_session.give((row.clone(), time, diff));
708                                            }
709                                        }
710                                    }
711                                    times_diffs.clear();
712
713                                    cursor.step_val(batch);
714                                }
715                                cursor.step_key(batch);
716                            }
717                        }
718                    });
719                })
720            });
721
722    (
723        ok_stream.as_collection(),
724        err_stream.as_collection().map(DataflowErrorSer::from),
725    )
726}
727
728/// Builds the beginning of the update stream of a delta path from a raw collection.
729///
730/// This is the unarranged counterpart of [`build_update_stream_trace`]. Only the delta path for the
731/// first relation can be seeded from a raw collection, since the as-of filtering that the other
732/// paths rely on is only available from an arrangement's times. We assert that here.
733fn build_update_stream_stream<'scope, T>(
734    edge: ColCollection<'scope, T>,
735    _as_of: Antichain<mz_repr::Timestamp>,
736    source_relation: usize,
737    initial_closure: JoinClosure,
738) -> (
739    VecCollection<'scope, T, Row, Diff>,
740    VecCollection<'scope, T, DataflowErrorSer, Diff>,
741)
742where
743    T: RenderTimestamp,
744{
745    // The other paths discard updates at the as-of, so only the first relation's path can be seeded
746    // from a raw collection that carries no per-update times to filter on.
747    assert_eq!(source_relation, 0);
748
749    type CB<C> = ConsolidatingContainerBuilder<C>;
750    // The closure reads datums and builds a fresh row, so the input row is only
751    // ever borrowed. Reading it from the column directly keeps this path from
752    // materializing an owned `Row` per record.
753    let (oks, errs) = flat_map_datums::<_, CB<Vec<(Row, T, Diff)>>, _>(
754        edge,
755        "UpdateStreamCollection",
756        usize::MAX,
757        {
758            let mut datum_vec = DatumVec::new();
759            move |row_datums, time, diff, ok_session, err_session| {
760                let mut row_builder = SharedRow::get();
761                let temp_storage = RowArena::new();
762                // `JoinClosure::apply` unifies the lifetimes of `&self`, the datums,
763                // and the arena. Copying the datums into a local vec lets that
764                // lifetime shrink to this call. The copy moves datum references, not
765                // row data.
766                let mut datums = datum_vec.borrow();
767                datums.extend(row_datums.iter());
768                // `cloned` detaches the result from `temp_storage` and the shared row
769                // builder, both of which drop at the end of this call.
770                match initial_closure
771                    .apply(&mut datums, &temp_storage, &mut row_builder)
772                    .map(|row| row.cloned())
773                    .transpose()
774                {
775                    Some(Ok(row)) => {
776                        ok_session.give((row, time, diff));
777                        1
778                    }
779                    None => 0,
780                    Some(Err(e)) => {
781                        err_session.give((DataflowErrorSer::from(e), time, diff));
782                        1
783                    }
784                }
785            }
786        },
787    );
788    (oks.as_collection(), errs.as_collection())
789}