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