Skip to main content

mz_compute/render/join/
mz_join_core.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//! A fork of DD's `JoinCore::join_core`.
11//!
12//! Currently, compute rendering knows two implementations for linear joins:
13//!
14//!  * Differential's `JoinCore::join_core`
15//!  * A Materialize fork thereof, called `mz_join_core`
16//!
17//! `mz_join_core` exists to solve a responsiveness problem with the DD implementation.
18//! DD's join is only able to yield between keys. When computing a large cross-join or a highly
19//! skewed join, this can result in loss of interactivity when the join operator refuses to yield
20//! control for multiple seconds or longer, which in turn causes degraded user experience.
21//! `mz_join_core` resolves the loss-of-interactivity issue by also yielding within keys.
22//!
23//! For the moment, we keep both implementations around, selectable through feature flags.
24//! Eventually, we hope that `mz_join_core` proves itself sufficiently to become the only join
25//! implementation.
26
27use std::cell::Cell;
28use std::cell::RefCell;
29use std::cmp::Ordering;
30use std::collections::VecDeque;
31use std::marker::PhantomData;
32use std::pin::Pin;
33use std::rc::Rc;
34use std::time::Instant;
35
36use differential_dataflow::Data;
37use differential_dataflow::consolidation::{consolidate_from, consolidate_updates};
38use differential_dataflow::lattice::Lattice;
39use differential_dataflow::operators::arrange::arrangement::Arranged;
40use differential_dataflow::trace::cursor::{
41    BatchCursor, BatchKey, BatchVal, CursorList, cursor_list,
42};
43use differential_dataflow::trace::{Cursor, Navigable, TraceReader};
44use mz_ore::future::yield_now;
45use mz_repr::Diff;
46use timely::container::PushInto;
47use timely::dataflow::Stream;
48use timely::dataflow::channels::pact::Pipeline;
49use timely::dataflow::operators::generic::OutputBuilderSession;
50use timely::dataflow::operators::{CapabilitySet, Operator};
51use timely::{ContainerBuilder, PartialOrder};
52use tracing::trace;
53
54/// Joins two arranged collections with the same key type.
55///
56/// Each matching pair of records `(key, val1)` and `(key, val2)` are subjected to the `result` function,
57/// which produces something implementing `IntoIterator`, where the output collection will have an entry for
58/// every value returned by the iterator.
59pub(super) fn mz_join_core<'scope, T, Tr1, Tr2, L, I, YFn, CB>(
60    arranged1: Arranged<'scope, Tr1>,
61    arranged2: Arranged<'scope, Tr2>,
62    result: L,
63    yield_fn: YFn,
64) -> Stream<'scope, T, CB::Container>
65where
66    T: timely::progress::Timestamp + Lattice,
67    Tr1: TraceReader<Batch: Navigable, Time = T> + Clone + 'static,
68    Tr2: TraceReader<Batch: Navigable, Time = T> + Clone + 'static,
69    BatchCursor<Tr1>: Cursor<Time = T, Diff = Diff>,
70    for<'a> BatchCursor<Tr2>: Cursor<Key<'a> = BatchKey<'a, Tr1>, Time = T, Diff = Diff>,
71    L: FnMut(BatchKey<'_, Tr1>, BatchVal<'_, Tr1>, BatchVal<'_, Tr2>) -> I + 'static,
72    I: IntoIterator<Item: Data> + 'static,
73    YFn: Fn(Instant, usize) -> bool + 'static,
74    CB: ContainerBuilder + PushInto<(I::Item, T, Diff)> + 'static,
75{
76    let scope = arranged1.stream.scope();
77    let mut trace1 = arranged1.trace.clone();
78    let mut trace2 = arranged2.trace.clone();
79
80    arranged1.stream.binary_frontier::<_, CB, _, _, _, _>(
81        arranged2.stream,
82        Pipeline,
83        Pipeline,
84        "Join",
85        move |capability, info| {
86            let operator_id = info.global_id;
87
88            // Acquire an activator to reschedule the operator when it has unfinished work.
89            let activator = scope.activator_for(info.address);
90
91            // Our initial invariants are that for each trace, physical compaction is less or equal the trace's upper bound.
92            // These invariants ensure that we can reference observed batch frontiers from `_start_upper` onward, as long as
93            // we maintain our physical compaction capabilities appropriately. These assertions are tested as we load up the
94            // initial work for the two traces, and before the operator is constructed.
95
96            // Acknowledged frontier for each input.
97            // These two are used exclusively to track batch boundaries on which we may want/need to call `batches_through`.
98            // They will drive our physical compaction of each trace, and we want to maintain at all times that each is beyond
99            // the physical compaction frontier of their corresponding trace.
100            // Should we ever *drop* a trace, these are 1. much harder to maintain correctly, but 2. no longer used.
101            use timely::progress::frontier::Antichain;
102            let mut acknowledged1 = Antichain::from_elem(<T>::minimum());
103            let mut acknowledged2 = Antichain::from_elem(<T>::minimum());
104
105            // deferred work of batches from each input.
106            let result_fn = Rc::new(RefCell::new(result));
107            // A trace hands out a `CursorList` over its batches' cursors, while an individual
108            // batch hands out a single `BatchCursor`. Input 1's per-batch work joins one batch1
109            // cursor against trace2's merged cursor, and vice versa for input 2.
110            let mut todo1 = Work::<BatchCursor<Tr1>, CursorList<BatchCursor<Tr2>>, _, _>::new(
111                Rc::clone(&result_fn),
112            );
113            let mut todo2 =
114                Work::<CursorList<BatchCursor<Tr1>>, BatchCursor<Tr2>, _, _>::new(result_fn);
115
116            // We'll unload the initial batches here, to put ourselves in a less non-deterministic state to start.
117            trace1.map_spans(|batch1| {
118                trace!(
119                    operator_id,
120                    input = 1,
121                    lower = ?batch1.lower().elements(),
122                    upper = ?batch1.upper().elements(),
123                    has_updates = batch1.has_updates(),
124                    "pre-loading batch",
125                );
126
127                acknowledged1.clone_from(batch1.upper());
128                // No `todo1` work here, because we haven't accepted anything into `batches2` yet.
129                // It is effectively "empty", because we choose to drain `trace1` before `trace2`.
130                // Once we start streaming batches in, we will need to respond to new batches from
131                // `input1` with logic that would have otherwise been here. Check out the next loop
132                // for the structure.
133            });
134            // At this point, `ack1` should exactly equal `trace1.read_upper()`, as they are both determined by
135            // iterating through batches and capturing the upper bound. This is a great moment to assert that
136            // `trace1`'s physical compaction frontier is before the frontier of completed times in `trace1`.
137            // TODO: in the case that this does not hold, instead "upgrade" the physical compaction frontier.
138            assert!(PartialOrder::less_equal(
139                &trace1.get_physical_compaction(),
140                &acknowledged1.borrow()
141            ));
142
143            trace!(
144                operator_id,
145                input = 1,
146                acknowledged1 = ?acknowledged1.elements(),
147                "pre-loading finished",
148            );
149
150            // We capture batch2's updates first and establish work second to avoid taking a
151            // `RefCell` lock on both traces at the same time, as they could be the same trace and
152            // this would panic.
153            let mut batch2_updates = Vec::new();
154            trace2.map_spans(|batch2| {
155                trace!(
156                    operator_id,
157                    input = 2,
158                    lower = ?batch2.lower().elements(),
159                    upper = ?batch2.upper().elements(),
160                    has_updates = batch2.has_updates(),
161                    "pre-loading batch",
162                );
163
164                acknowledged2.clone_from(batch2.upper());
165                // Empty batches carry no updates, and have nothing to join.
166                if let Some(updates2) = &batch2.inner {
167                    batch2_updates.push(updates2.clone());
168                }
169            });
170            // At this point, `ack2` should exactly equal `trace2.read_upper()`, as they are both determined by
171            // iterating through batches and capturing the upper bound. This is a great moment to assert that
172            // `trace2`'s physical compaction frontier is before the frontier of completed times in `trace2`.
173            // TODO: in the case that this does not hold, instead "upgrade" the physical compaction frontier.
174            assert!(PartialOrder::less_equal(
175                &trace2.get_physical_compaction(),
176                &acknowledged2.borrow()
177            ));
178
179            // Batches wholly at or before these frontiers were joined by the start-up loading
180            // above, and arriving input batches are ignored up to them. Beyond them, every
181            // non-empty arriving batch must be joined, even once `acknowledged` has advanced past
182            // it. `advance_upper` consults the shared trace, whose merges may have consolidated an
183            // in-flight batch's updates away once logical compaction equates an add/remove pair's
184            // times. The trace's emptiness there is valid only for readers at or beyond the
185            // compaction frontier, while our consumers may read finer times, so the raw batch still
186            // owes them its updates. Testing against these fixed boundaries rather than the mutable
187            // `acknowledged` frontiers keeps such batches from being dropped. See
188            // TimelyDataflow/differential-dataflow#801.
189            let preload_upper1 = acknowledged1.clone();
190            let preload_upper2 = acknowledged2.clone();
191
192            // Load up deferred work using the trace2 updates captured just above.
193            for updates2 in batch2_updates.into_iter() {
194                trace!(
195                    operator_id,
196                    input = 2,
197                    acknowledged1 = ?acknowledged1.elements(),
198                    "deferring work for batch",
199                );
200
201                // It is safe to ask for `ack1` because we have confirmed it to be in advance of `distinguish_since`.
202                let (trace1_cursor, trace1_storage) =
203                    cursor_list(trace1.batches_through(acknowledged1.borrow()).unwrap());
204                // We could downgrade the capability here, but doing so is a bit complicated mathematically.
205                // TODO: downgrade the capability by searching out the one time in `batch2.lower()` and not
206                // in `batch2.upper()`. Only necessary for non-empty batches, as empty batches may not have
207                // that property.
208                todo2.push(
209                    trace1_cursor,
210                    trace1_storage,
211                    updates2.cursor(),
212                    updates2,
213                    CapabilitySet::from_elem(capability.clone()),
214                    capability.time().clone(),
215                );
216            }
217
218            trace!(
219                operator_id,
220                input = 2,
221                acknowledged2 = ?acknowledged2.elements(),
222                "pre-loading finished",
223            );
224
225            // Droppable handles to shared trace data structures.
226            let mut trace1_option = Some(trace1);
227            let mut trace2_option = Some(trace2);
228
229            move |(input1, frontier1), (input2, frontier2), output| {
230                // 1. Consuming input.
231                //
232                // The join computation repeatedly accepts batches of updates from each of its inputs.
233                //
234                // For each accepted batch, it prepares a work-item to join the batch against previously "accepted"
235                // updates from its other input. It is important to track which updates have been accepted, because
236                // we use a shared trace and there may be updates present that are in advance of this accepted bound.
237                //
238                // Batches are accepted: 1. in bulk at start-up (above), 2. as we observe them in the input stream,
239                // and 3. if the trace can confirm a region of empty space directly following our accepted bound.
240                // This last case is a consequence of our inability to transmit empty batches, as they may be formed
241                // in the absence of timely dataflow capabilities.
242
243                // Drain input 1, prepare work.
244                input1.for_each(|capability, data| {
245                    let trace2 = trace2_option
246                        .as_mut()
247                        .expect("we only drop a trace in response to the other input emptying");
248                    let capability = capability.retain_stamp(0);
249                    let meet = capability_meet(&capability);
250                    for batch1 in data.drain(..) {
251                        // The preload boundary and `acknowledged1` are both uppers of whole
252                        // stream batches or of trace merges of them, so an arriving batch lies
253                        // wholly on one side of each. A batch spanning the preload boundary would
254                        // be partly joined twice, one spanning `acknowledged1` mis-accounted.
255                        mz_ore::soft_assert_or_log!(
256                            PartialOrder::less_equal(batch1.upper(), &preload_upper1)
257                                || PartialOrder::less_equal(&preload_upper1, batch1.lower()),
258                            "batch spans the preload boundary",
259                        );
260                        mz_ore::soft_assert_or_log!(
261                            PartialOrder::less_equal(&acknowledged1, batch1.lower())
262                                || PartialOrder::less_equal(batch1.upper(), &acknowledged1),
263                            "batch spans the acknowledged frontier",
264                        );
265                        // Ignore any pre-loaded data, which was joined at start-up. This tests the
266                        // fixed preload boundary, not `acknowledged1`, which `advance_upper` can
267                        // move past an in-flight batch whose updates trace merges consolidated away.
268                        // Such a batch must still be joined.
269                        if !PartialOrder::less_equal(batch1.upper(), &preload_upper1) {
270                            trace!(
271                                operator_id,
272                                input = 1,
273                                lower = ?batch1.lower().elements(),
274                                upper = ?batch1.upper().elements(),
275                                has_updates = batch1.has_updates(),
276                                "loading batch",
277                            );
278
279                            if let Some(updates1) = &batch1.inner {
280                                trace!(
281                                    operator_id,
282                                    input = 1,
283                                    acknowledged2 = ?acknowledged2.elements(),
284                                    "deferring work for batch",
285                                );
286
287                                // It is safe to ask for `ack2` as we validated that it was at least `get_physical_compaction()`
288                                // at start-up, and have held back physical compaction ever since.
289                                let (trace2_cursor, trace2_storage) = cursor_list(
290                                    trace2.batches_through(acknowledged2.borrow()).unwrap(),
291                                );
292                                todo1.push(
293                                    updates1.cursor(),
294                                    updates1.clone(),
295                                    trace2_cursor,
296                                    trace2_storage,
297                                    capability.clone(),
298                                    meet.clone(),
299                                );
300                            }
301
302                            // To update `acknowledged1` we might presume that `batch1.lower` should equal it, but we
303                            // may have skipped over empty batches. Still, the batches are in-order, and we should be
304                            // able to just assume the most recent `batch1.upper`, unless `advance_upper` already moved
305                            // `acknowledged1` past this batch, in which case we keep the further frontier.
306                            if PartialOrder::less_equal(&acknowledged1, batch1.lower()) {
307                                mz_ore::soft_assert_or_log!(
308                                    PartialOrder::less_equal(&acknowledged1, batch1.upper()),
309                                    "acknowledged1 {:?} regressed past batch1 upper {:?}",
310                                    acknowledged1.elements(),
311                                    batch1.upper().elements(),
312                                );
313                                acknowledged1.clone_from(batch1.upper());
314                            }
315
316                            trace!(
317                                operator_id,
318                                input = 1,
319                                acknowledged1 = ?acknowledged1.elements(),
320                                "batch acknowledged",
321                            );
322                        }
323                    }
324                });
325
326                // Drain input 2, prepare work.
327                input2.for_each(|capability, data| {
328                    let trace1 = trace1_option
329                        .as_mut()
330                        .expect("we only drop a trace in response to the other input emptying");
331                    let capability = capability.retain_stamp(0);
332                    let meet = capability_meet(&capability);
333                    for batch2 in data.drain(..) {
334                        // The preload boundary and `acknowledged2` are both uppers of whole
335                        // stream batches or of trace merges of them, so an arriving batch lies
336                        // wholly on one side of each. A batch spanning the preload boundary would
337                        // be partly joined twice, one spanning `acknowledged2` mis-accounted.
338                        mz_ore::soft_assert_or_log!(
339                            PartialOrder::less_equal(batch2.upper(), &preload_upper2)
340                                || PartialOrder::less_equal(&preload_upper2, batch2.lower()),
341                            "batch spans the preload boundary",
342                        );
343                        mz_ore::soft_assert_or_log!(
344                            PartialOrder::less_equal(&acknowledged2, batch2.lower())
345                                || PartialOrder::less_equal(batch2.upper(), &acknowledged2),
346                            "batch spans the acknowledged frontier",
347                        );
348                        // Ignore any pre-loaded data, which was joined at start-up. This tests the
349                        // fixed preload boundary, not `acknowledged2`, which `advance_upper` can
350                        // move past an in-flight batch whose updates trace merges consolidated away.
351                        // Such a batch must still be joined.
352                        if !PartialOrder::less_equal(batch2.upper(), &preload_upper2) {
353                            trace!(
354                                operator_id,
355                                input = 2,
356                                lower = ?batch2.lower().elements(),
357                                upper = ?batch2.upper().elements(),
358                                has_updates = batch2.has_updates(),
359                                "loading batch",
360                            );
361
362                            if let Some(updates2) = &batch2.inner {
363                                trace!(
364                                    operator_id,
365                                    input = 2,
366                                    acknowledged1 = ?acknowledged1.elements(),
367                                    "deferring work for batch",
368                                );
369
370                                // It is safe to ask for `ack1` as we validated that it was at least `get_physical_compaction()`
371                                // at start-up, and have held back physical compaction ever since.
372                                let (trace1_cursor, trace1_storage) = cursor_list(
373                                    trace1.batches_through(acknowledged1.borrow()).unwrap(),
374                                );
375                                todo2.push(
376                                    trace1_cursor,
377                                    trace1_storage,
378                                    updates2.cursor(),
379                                    updates2.clone(),
380                                    capability.clone(),
381                                    meet.clone(),
382                                );
383                            }
384
385                            // To update `acknowledged2` we might presume that `batch2.lower` should equal it, but we
386                            // may have skipped over empty batches. Still, the batches are in-order, and we should be
387                            // able to just assume the most recent `batch2.upper`, unless `advance_upper` already moved
388                            // `acknowledged2` past this batch, in which case we keep the further frontier.
389                            if PartialOrder::less_equal(&acknowledged2, batch2.lower()) {
390                                mz_ore::soft_assert_or_log!(
391                                    PartialOrder::less_equal(&acknowledged2, batch2.upper()),
392                                    "acknowledged2 {:?} regressed past batch2 upper {:?}",
393                                    acknowledged2.elements(),
394                                    batch2.upper().elements(),
395                                );
396                                acknowledged2.clone_from(batch2.upper());
397                            }
398
399                            trace!(
400                                operator_id,
401                                input = 2,
402                                acknowledged2 = ?acknowledged2.elements(),
403                                "batch acknowledged",
404                            );
405                        }
406                    }
407                });
408
409                // Advance acknowledged frontiers through any empty regions that we may not receive as batches.
410                if let Some(trace1) = trace1_option.as_mut() {
411                    trace!(
412                        operator_id,
413                        input = 1,
414                        acknowledged1 = ?acknowledged1.elements(),
415                        "advancing trace upper",
416                    );
417                    trace1.advance_upper(&mut acknowledged1);
418                }
419                if let Some(trace2) = trace2_option.as_mut() {
420                    trace!(
421                        operator_id,
422                        input = 2,
423                        acknowledged2 = ?acknowledged2.elements(),
424                        "advancing trace upper",
425                    );
426                    trace2.advance_upper(&mut acknowledged2);
427                }
428
429                // 2. Join computation.
430                //
431                // For each of the inputs, we do some amount of work (measured in terms of number
432                // of output records produced). This is meant to yield control to allow downstream
433                // operators to consume and reduce the output, but it it also means to provide some
434                // degree of responsiveness. There is a potential risk here that if we fall behind
435                // then the increasing queues hold back physical compaction of the underlying traces
436                // which results in unintentionally quadratic processing time (each batch of either
437                // input must scan all batches from the other input).
438
439                // Perform some amount of outstanding work for input 1.
440                trace!(
441                    operator_id,
442                    input = 1,
443                    work_left = todo1.remaining(),
444                    "starting work"
445                );
446                todo1.process(output, &yield_fn);
447                trace!(
448                    operator_id,
449                    input = 1,
450                    work_left = todo1.remaining(),
451                    "ceasing work",
452                );
453
454                // Perform some amount of outstanding work for input 2.
455                trace!(
456                    operator_id,
457                    input = 2,
458                    work_left = todo2.remaining(),
459                    "starting work"
460                );
461                todo2.process(output, &yield_fn);
462                trace!(
463                    operator_id,
464                    input = 2,
465                    work_left = todo2.remaining(),
466                    "ceasing work",
467                );
468
469                // Re-activate operator if work remains.
470                if !todo1.is_empty() || !todo2.is_empty() {
471                    activator.activate();
472                }
473
474                // 3. Trace maintenance.
475                //
476                // Importantly, we use `input.frontier()` here rather than `acknowledged` to track
477                // the progress of an input, because should we ever drop one of the traces we will
478                // lose the ability to extract information from anything other than the input.
479                // For example, if we dropped `trace2` we would not be able to use `advance_upper`
480                // to keep `acknowledged2` up to date wrt empty batches, and would hold back logical
481                // compaction of `trace1`.
482
483                // Maintain `trace1`. Drop if `input2` is empty, or advance based on future needs.
484                if let Some(trace1) = trace1_option.as_mut() {
485                    if frontier2.is_empty() {
486                        trace!(operator_id, input = 1, "dropping trace handle");
487                        trace1_option = None;
488                    } else {
489                        trace!(
490                            operator_id,
491                            input = 1,
492                            logical = ?*frontier2.frontier(),
493                            physical = ?acknowledged1.elements(),
494                            "advancing trace compaction",
495                        );
496
497                        // Allow `trace1` to compact logically up to the frontier we may yet receive,
498                        // in the opposing input (`input2`). All `input2` times will be beyond this
499                        // frontier, and joined times only need to be accurate when advanced to it.
500                        trace1.set_logical_compaction(frontier2.frontier());
501                        // Allow `trace1` to compact physically up to the upper bound of batches we
502                        // have received in its input (`input1`). We will not require a cursor that
503                        // is not beyond this bound.
504                        trace1.set_physical_compaction(acknowledged1.borrow());
505                    }
506                }
507
508                // Maintain `trace2`. Drop if `input1` is empty, or advance based on future needs.
509                if let Some(trace2) = trace2_option.as_mut() {
510                    if frontier1.is_empty() {
511                        trace!(operator_id, input = 2, "dropping trace handle");
512                        trace2_option = None;
513                    } else {
514                        trace!(
515                            operator_id,
516                            input = 2,
517                            logical = ?*frontier1.frontier(),
518                            physical = ?acknowledged2.elements(),
519                            "advancing trace compaction",
520                        );
521
522                        // Allow `trace2` to compact logically up to the frontier we may yet receive,
523                        // in the opposing input (`input1`). All `input1` times will be beyond this
524                        // frontier, and joined times only need to be accurate when advanced to it.
525                        trace2.set_logical_compaction(frontier1.frontier());
526                        // Allow `trace2` to compact physically up to the upper bound of batches we
527                        // have received in its input (`input2`). We will not require a cursor that
528                        // is not beyond this bound.
529                        trace2.set_physical_compaction(acknowledged2.borrow());
530                    }
531                }
532            }
533        },
534    )
535}
536
537/// The lattice meet of the times of `capabilities`.
538///
539/// A message in a partially ordered scope can carry several capabilities. Their meet lower bounds
540/// every time in the message's batches, so it is a valid time to advance the opposing trace's
541/// updates to. Panics if `capabilities` is empty.
542fn capability_meet<T: Lattice + timely::progress::Timestamp>(capabilities: &CapabilitySet<T>) -> T {
543    capabilities
544        .iter()
545        .map(|c| c.time().clone())
546        .reduce(|a, b| a.meet(&b))
547        .expect("message stamped with no capabilities")
548}
549
550/// Work collected by the join operator.
551///
552/// The join operator enqueues new work here first, and then processes it at a controlled rate,
553/// potentially yielding control to the Timely runtime in between. This allows it to avoid OOMs,
554/// caused by buffering massive amounts of data at the output, and loss of interactivity.
555///
556/// Collected work can be reduced by calling the `process` method.
557struct Work<C1, C2, D, L>
558where
559    C1: Cursor,
560    C2: Cursor,
561{
562    /// Pending work.
563    todo: VecDeque<(Pin<Box<dyn Future<Output = ()>>>, CapabilitySet<C1::Time>)>,
564    /// A function that transforms raw join matches into join results.
565    result_fn: Rc<RefCell<L>>,
566    /// A buffer holding the join results.
567    ///
568    /// Written by the work futures, drained by `Work::process`.
569    output: Rc<RefCell<Vec<(D, C1::Time, Diff)>>>,
570    /// The number of join results produced by work futures.
571    ///
572    /// Used with `yield_fn` to inform when `Work::process` should yield.
573    produced: Rc<Cell<usize>>,
574
575    _cursors: PhantomData<(C1, C2)>,
576}
577
578impl<C1, C2, D, L, I> Work<C1, C2, D, L>
579where
580    C1: Cursor<Diff = Diff> + 'static,
581    C2: for<'a> Cursor<Key<'a> = C1::Key<'a>, Time = C1::Time, Diff = Diff> + 'static,
582    D: Data,
583    L: FnMut(C1::Key<'_>, C1::Val<'_>, C2::Val<'_>) -> I + 'static,
584    I: IntoIterator<Item = D> + 'static,
585{
586    fn new(result_fn: Rc<RefCell<L>>) -> Self {
587        Self {
588            todo: Default::default(),
589            result_fn,
590            output: Default::default(),
591            produced: Default::default(),
592            _cursors: PhantomData,
593        }
594    }
595
596    /// Return the amount of remaining work chunks.
597    fn remaining(&self) -> usize {
598        self.todo.len()
599    }
600
601    /// Return whether there is any work pending.
602    fn is_empty(&self) -> bool {
603        self.remaining() == 0
604    }
605
606    /// Append some pending work.
607    ///
608    /// Every time in the freshly arrived batch must be greater or equal to some element of
609    /// `capability`, and to `meet`. The join results are advanced by `meet`.
610    fn push(
611        &mut self,
612        cursor1: C1,
613        storage1: C1::Storage,
614        cursor2: C2,
615        storage2: C2::Storage,
616        capability: CapabilitySet<C1::Time>,
617        meet: C1::Time,
618    ) {
619        let fut = self.start_work(cursor1, storage1, cursor2, storage2, meet);
620
621        self.todo.push_back((Box::pin(fut), capability));
622    }
623
624    /// Process pending work until none is remaining or `yield_fn` requests a yield.
625    fn process<CB, YFn>(
626        &mut self,
627        output: &mut OutputBuilderSession<'_, C1::Time, CB>,
628        yield_fn: YFn,
629    ) where
630        CB: ContainerBuilder + PushInto<(D, C1::Time, Diff)>,
631        YFn: Fn(Instant, usize) -> bool,
632    {
633        let start_time = Instant::now();
634        self.produced.set(0);
635
636        let waker = futures::task::noop_waker();
637        let mut ctx = std::task::Context::from_waker(&waker);
638
639        while let Some((mut fut, cap)) = self.todo.pop_front() {
640            // Drive the work future until it's done or it's time to yield.
641            let mut done = false;
642            let mut should_yield = false;
643            while !done && !should_yield {
644                done = fut.as_mut().poll(&mut ctx).is_ready();
645                should_yield = yield_fn(start_time, self.produced.get());
646            }
647
648            // Drain the produced join results.
649            let mut output_buf = self.output.borrow_mut();
650
651            // Consolidating here is important when the join closure produces data that
652            // consolidates well, for example when projecting columns.
653            let old_len = output_buf.len();
654            consolidate_updates(&mut output_buf);
655            let recovered = old_len - output_buf.len();
656            self.produced.update(|x| x - recovered);
657
658            output
659                .session_with_builder(&cap)
660                .give_iterator(output_buf.drain(..));
661
662            if done {
663                // We have finished processing a chunk of work. Use this opportunity to truncate
664                // the output buffer, so we don't keep excess memory allocated forever.
665                *output_buf = Default::default();
666            } else if !done {
667                // Still work to do in this chunk.
668                self.todo.push_front((fut, cap));
669            }
670
671            if should_yield {
672                break;
673            }
674        }
675    }
676
677    /// Start the work of joining the updates produced by the given cursors.
678    ///
679    /// This method returns a `Future` that can be polled to make progress on the join work.
680    /// Returning a future allows us to implement the logic using async/await syntax where we can
681    /// conveniently pause the work at any point by calling `yield_now().await`. We are allowed to
682    /// hold references across yield points, which is something we wouldn't get with a hand-rolled
683    /// state machine implementation.
684    fn start_work(
685        &self,
686        mut cursor1: C1,
687        storage1: C1::Storage,
688        mut cursor2: C2,
689        storage2: C2::Storage,
690        meet: C1::Time,
691    ) -> impl Future<Output = ()> + use<C1, C2, D, L, I> {
692        let result_fn = Rc::clone(&self.result_fn);
693        let output = Rc::clone(&self.output);
694        let produced = Rc::clone(&self.produced);
695
696        async move {
697            let mut joiner = Joiner::new(result_fn, output, produced, meet);
698
699            while let Some(key1) = cursor1.get_key(&storage1)
700                && let Some(key2) = cursor2.get_key(&storage2)
701            {
702                match key1.cmp(&key2) {
703                    Ordering::Less => cursor1.seek_key(&storage1, key2),
704                    Ordering::Greater => cursor2.seek_key(&storage2, key1),
705                    Ordering::Equal => {
706                        joiner
707                            .join_key(key1, &mut cursor1, &storage1, &mut cursor2, &storage2)
708                            .await;
709
710                        cursor1.step_key(&storage1);
711                        cursor2.step_key(&storage2);
712                    }
713                }
714            }
715        }
716    }
717}
718
719/// Type that knows how to perform the core join logic.
720///
721/// The joiner implements two join strategies:
722///
723///  * The "simple" strategy produces a match for each combination of (val1, time1, val2, time2)
724///    found in the inputs. If there are multiple times in the input, it may produce matches for
725///    times in which one of the values wasn't present. These matches cancel each other out, so the
726///    result ends up correct.
727///  * The "linear scan over times" strategy sorts the input data by time and then steps through
728///    the input histories, producing matches for a pair of values only if both values where
729///    present at the same time.
730///
731/// The linear scan strategy avoids redundant work and is much more efficient than the simple
732/// strategy when many distinct times are present in the inputs. However, sorting the input data
733/// incurs some overhead, so we still prefer the simple variant when the input data is small.
734struct Joiner<'a, C1, C2, D, L>
735where
736    C1: Cursor,
737    C2: Cursor,
738{
739    /// A function that transforms raw join matches into join results.
740    result_fn: Rc<RefCell<L>>,
741    /// A buffer holding the join results.
742    output: Rc<RefCell<Vec<(D, C1::Time, Diff)>>>,
743    /// The number of join results produced.
744    produced: Rc<Cell<usize>>,
745    /// A time to which all join results should be advanced.
746    meet: C1::Time,
747
748    /// Buffer for edit histories from the first input.
749    history1: ValueHistory<'a, C1>,
750    /// Buffer for edit histories from the second input.
751    history2: ValueHistory<'a, C2>,
752}
753
754impl<'a, C1, C2, D, L, I> Joiner<'a, C1, C2, D, L>
755where
756    C1: Cursor<Diff = Diff>,
757    C2: Cursor<Key<'a> = C1::Key<'a>, Time = C1::Time, Diff = Diff>,
758    D: Data,
759    L: FnMut(C1::Key<'_>, C1::Val<'_>, C2::Val<'_>) -> I + 'static,
760    I: IntoIterator<Item = D> + 'static,
761{
762    fn new(
763        result_fn: Rc<RefCell<L>>,
764        output: Rc<RefCell<Vec<(D, C1::Time, Diff)>>>,
765        produced: Rc<Cell<usize>>,
766        meet: C1::Time,
767    ) -> Self {
768        Self {
769            result_fn,
770            output,
771            produced,
772            meet,
773            history1: ValueHistory::new(),
774            history2: ValueHistory::new(),
775        }
776    }
777
778    /// Produce matches for the values of a single key.
779    async fn join_key(
780        &mut self,
781        key: C1::Key<'_>,
782        cursor1: &mut C1,
783        storage1: &'a C1::Storage,
784        cursor2: &mut C2,
785        storage2: &'a C2::Storage,
786    ) {
787        self.history1.edits.load(cursor1, storage1, &self.meet);
788        self.history2.edits.load(cursor2, storage2, &self.meet);
789
790        // If the input data is small, use the simple strategy.
791        //
792        // TODO: This conditional is taken directly from DD. We should check if it might make sense
793        //       to do something different, like using the simple strategy always when the number
794        //       of distinct times is small.
795        if self.history1.edits.len() < 10 || self.history2.edits.len() < 10 {
796            self.join_key_simple(key);
797            yield_now().await;
798        } else {
799            self.join_key_linear_time_scan(key).await;
800        }
801    }
802
803    /// Produce matches for the values of a single key, using the simple strategy.
804    ///
805    /// This strategy is only meant to be used for small inputs, so we don't bother including yield
806    /// points or optimizations.
807    fn join_key_simple(&self, key: C1::Key<'_>) {
808        let mut result_fn = self.result_fn.borrow_mut();
809        let mut output = self.output.borrow_mut();
810
811        for (v1, t1, r1) in self.history1.edits.iter() {
812            for (v2, t2, r2) in self.history2.edits.iter() {
813                for data in result_fn(key, v1, v2) {
814                    output.push((data, t1.join(t2), r1 * r2));
815                    self.produced.update(|x| x + 1);
816                }
817            }
818        }
819    }
820
821    /// Produce matches for the values of a single key, using a linear scan through times.
822    async fn join_key_linear_time_scan(&mut self, key: C1::Key<'_>) {
823        let history1 = &mut self.history1;
824        let history2 = &mut self.history2;
825
826        history1.replay();
827        history2.replay();
828
829        // TODO: It seems like there is probably a good deal of redundant `advance_buffer_by`
830        //       in here. If a time is ever repeated, for example, the call will be identical
831        //       and accomplish nothing. If only a single record has been added, it may not
832        //       be worth the time to collapse (advance, re-sort) the data when a linear scan
833        //       is sufficient.
834
835        // Join the next entry in `history1`.
836        let work_history1 = |history1: &mut ValueHistory<C1>, history2: &mut ValueHistory<C2>| {
837            let mut result_fn = self.result_fn.borrow_mut();
838            let mut output = self.output.borrow_mut();
839
840            let (t1, meet, v1, r1) = history1.get().unwrap();
841            history2.advance_past_by(meet);
842            for &(v2, ref t2, r2) in &history2.past {
843                for data in result_fn(key, v1, v2) {
844                    output.push((data, t1.join(t2), r1 * r2));
845                    self.produced.update(|x| x + 1);
846                }
847            }
848            history1.step();
849        };
850
851        // Join the next entry in `history2`.
852        let work_history2 = |history1: &mut ValueHistory<C1>, history2: &mut ValueHistory<C2>| {
853            let mut result_fn = self.result_fn.borrow_mut();
854            let mut output = self.output.borrow_mut();
855
856            let (t2, meet, v2, r2) = history2.get().unwrap();
857            history1.advance_past_by(meet);
858            for &(v1, ref t1, r1) in &history1.past {
859                for data in result_fn(key, v1, v2) {
860                    output.push((data, t1.join(t2), r1 * r2));
861                    self.produced.update(|x| x + 1);
862                }
863            }
864            history2.step();
865        };
866
867        while let Some(time1) = history1.get_time()
868            && let Some(time2) = history2.get_time()
869        {
870            if time1 < time2 {
871                work_history1(history1, history2)
872            } else {
873                work_history2(history1, history2)
874            };
875            yield_now().await;
876        }
877
878        while !history1.is_empty() {
879            work_history1(history1, history2);
880            yield_now().await;
881        }
882        while !history2.is_empty() {
883            work_history2(history1, history2);
884            yield_now().await;
885        }
886    }
887}
888
889/// An accumulation of (value, time, diff) updates.
890///
891/// Deduplicated values are stored in `values`. Each entry includes the end index of the
892/// corresponding range in `edits`. The edits stored for a value are consolidated.
893struct EditList<'a, C: Cursor> {
894    values: Vec<(C::Val<'a>, usize)>,
895    edits: Vec<(C::Time, Diff)>,
896}
897
898impl<'a, C> EditList<'a, C>
899where
900    C: Cursor<Diff = Diff>,
901{
902    fn len(&self) -> usize {
903        self.edits.len()
904    }
905
906    /// Load the updates in the given cursor.
907    ///
908    /// Steps over values, but not over keys.
909    fn load(&mut self, cursor: &mut C, storage: &'a C::Storage, meet: &C::Time) {
910        self.values.clear();
911        self.edits.clear();
912
913        let mut edit_idx = 0;
914        while let Some(value) = cursor.get_val(storage) {
915            cursor.map_times(storage, |time, diff| {
916                let mut time = C::owned_time(time);
917                time.join_assign(meet);
918                self.edits.push((time, C::owned_diff(diff)));
919            });
920
921            consolidate_from(&mut self.edits, edit_idx);
922
923            if self.edits.len() > edit_idx {
924                edit_idx = self.edits.len();
925                self.values.push((value, edit_idx));
926            }
927
928            cursor.step_val(storage);
929        }
930    }
931
932    /// Iterate over the contained updates.
933    fn iter(&self) -> impl Iterator<Item = (C::Val<'a>, &C::Time, Diff)> {
934        self.values
935            .iter()
936            .enumerate()
937            .flat_map(|(idx, (value, end))| {
938                let start = if idx == 0 { 0 } else { self.values[idx - 1].1 };
939                let edits = &self.edits[start..*end];
940                edits.iter().map(|(time, diff)| (*value, time, *diff))
941            })
942    }
943}
944
945/// A history for replaying updates in time order.
946struct ValueHistory<'a, C: Cursor> {
947    /// Unsorted updates to replay.
948    edits: EditList<'a, C>,
949    /// Time-sorted updates that have not been stepped over yet.
950    ///
951    /// Entries are (time, meet, value_idx, diff).
952    future: Vec<(C::Time, C::Time, usize, Diff)>,
953    /// Rolled-up updates that have been stepped over.
954    past: Vec<(C::Val<'a>, C::Time, Diff)>,
955}
956
957impl<'a, C> ValueHistory<'a, C>
958where
959    C: Cursor,
960{
961    /// Create a new empty `ValueHistory`.
962    fn new() -> Self {
963        Self {
964            edits: EditList {
965                values: Default::default(),
966                edits: Default::default(),
967            },
968            future: Default::default(),
969            past: Default::default(),
970        }
971    }
972
973    /// Return whether there are updates left to step over.
974    fn is_empty(&self) -> bool {
975        self.future.is_empty()
976    }
977
978    /// Return the next update.
979    fn get(&self) -> Option<(&C::Time, &C::Time, C::Val<'a>, Diff)> {
980        self.future.last().map(|(t, m, v, r)| {
981            let (value, _) = self.edits.values[*v];
982            (t, m, value, *r)
983        })
984    }
985
986    /// Return the time of the next update.
987    fn get_time(&self) -> Option<&C::Time> {
988        self.future.last().map(|(t, _, _, _)| t)
989    }
990
991    /// Populate `future` with the updates stored in `edits`.
992    fn replay(&mut self) {
993        self.future.clear();
994        self.past.clear();
995
996        let values = &self.edits.values;
997        let edits = &self.edits.edits;
998        for (idx, (_, end)) in values.iter().enumerate() {
999            let start = if idx == 0 { 0 } else { values[idx - 1].1 };
1000            for edit_idx in start..*end {
1001                let (time, diff) = &edits[edit_idx];
1002                self.future.push((time.clone(), time.clone(), idx, *diff));
1003            }
1004        }
1005
1006        self.future.sort_by(|x, y| y.cmp(x));
1007
1008        for idx in 1..self.future.len() {
1009            self.future[idx].1 = self.future[idx].1.meet(&self.future[idx - 1].1);
1010        }
1011    }
1012
1013    /// Advance the history by moving the next entry from `future` into `past`.
1014    fn step(&mut self) {
1015        let (time, _, value_idx, diff) = self.future.pop().unwrap();
1016        let (value, _) = self.edits.values[value_idx];
1017        self.past.push((value, time, diff));
1018    }
1019
1020    /// Advance all times in `past` by `meet`.
1021    fn advance_past_by(&mut self, meet: &C::Time) {
1022        for (_, time, _) in &mut self.past {
1023            time.join_assign(meet);
1024        }
1025        consolidate_updates(&mut self.past);
1026    }
1027}