mz_compute/compute_state/
peek_result_iterator.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//! Code for extracting a peek result out of compute state/an arrangement.
7
8use std::iter::FusedIterator;
9use std::num::NonZeroI64;
10use std::ops::Range;
11
12use differential_dataflow::trace::cursor::{BatchCursor, BatchKey, CursorList};
13use differential_dataflow::trace::implementations::BatchContainer;
14use differential_dataflow::trace::{Cursor, Navigable, TraceReader};
15
16/// The merged cursor a [`TraceReader::cursor`] hands out over all of a trace's batches: a
17/// [`CursorList`] over the per-batch cursors.
18type TraceCursor<Tr> = CursorList<BatchCursor<Tr>>;
19/// Backing storage for a [`TraceCursor`]: the batches the cursor borrows from.
20type TraceStorage<Tr> = Vec<<Tr as TraceReader>::Batch>;
21use mz_ore::result::ResultExt;
22use mz_repr::fixed_length::ExtendDatums;
23use mz_repr::{DatumVec, Diff, GlobalId, Row, RowArena};
24use timely::order::PartialOrder;
25
26pub struct PeekResultIterator<Tr>
27where
28    Tr: TraceReader<Batch: Navigable>,
29{
30    // For debug/trace logging.
31    target_id: GlobalId,
32    cursor: TraceCursor<Tr>,
33    storage: TraceStorage<Tr>,
34    map_filter_project: mz_expr::SafeMfpPlan,
35    peek_timestamp: mz_repr::Timestamp,
36    row_builder: Row,
37    datum_vec: DatumVec,
38    literals: Option<Literals<Tr>>,
39}
40
41/// Helper to handle literals in peeks
42struct Literals<Tr: TraceReader<Batch: Navigable>> {
43    /// The literals in a container, sorted by `Ord`.
44    literals: <BatchCursor<Tr> as Cursor>::KeyContainer,
45    /// The range of the literals that are still available.
46    range: Range<usize>,
47    /// The current index in the literals.
48    current_index: Option<usize>,
49}
50
51impl<Tr> Literals<Tr>
52where
53    Tr: TraceReader<Batch: Navigable>,
54    BatchCursor<Tr>: Cursor<KeyContainer: BatchContainer<Owned: Ord>>,
55{
56    /// Construct a new `Literals` from a mutable slice of literals. Sorts contents.
57    fn new(
58        literals: &mut [<<BatchCursor<Tr> as Cursor>::KeyContainer as BatchContainer>::Owned],
59        cursor: &mut TraceCursor<Tr>,
60        storage: &TraceStorage<Tr>,
61    ) -> Self {
62        // We have to sort the literal constraints because cursor.seek_key can
63        // seek only forward.
64        literals.sort();
65        let mut container =
66            <BatchCursor<Tr> as Cursor>::KeyContainer::with_capacity(literals.len());
67        for constraint in literals {
68            container.push_own(constraint)
69        }
70        let range = 0..container.len();
71        let mut this = Self {
72            literals: container,
73            range,
74            current_index: None,
75        };
76        this.seek_next_literal_key(cursor, storage);
77        this
78    }
79
80    /// Returns the current literal, if any.
81    fn peek(&self) -> Option<BatchKey<'_, Tr>> {
82        self.current_index
83            .and_then(|index| self.literals.get(index))
84    }
85
86    /// Returns `true` if there are no more literals to process.
87    fn is_exhausted(&self) -> bool {
88        self.current_index.is_none()
89    }
90
91    /// Seeks the cursor to the next key of a matching literal, if any.
92    fn seek_next_literal_key(&mut self, cursor: &mut TraceCursor<Tr>, storage: &TraceStorage<Tr>) {
93        while let Some(index) = self.range.next() {
94            let literal = self.literals.get(index).expect("index out of bounds");
95            cursor.seek_key(storage, literal);
96            if cursor.get_key(storage).map_or(true, |key| key == literal) {
97                self.current_index = Some(index);
98                return;
99            }
100            // The cursor landed on a record that has a different key,
101            // meaning that there is no record whose key would match the
102            // current literal.
103        }
104        self.current_index = None;
105    }
106}
107
108/// An [Iterator] that extracts a peek result from a [TraceReader].
109///
110/// The iterator will apply a given `MapFilterProject` and obey literal
111/// constraints, if any.
112impl<Tr> PeekResultIterator<Tr>
113where
114    Tr: TraceReader<Batch: Navigable>,
115    for<'a> BatchCursor<Tr>: Cursor<
116            Key<'a>: ExtendDatums + Eq,
117            KeyContainer: BatchContainer<Owned = Row>,
118            Val<'a>: ExtendDatums,
119            TimeGat<'a>: PartialOrder<mz_repr::Timestamp>,
120            DiffGat<'a> = &'a Diff,
121        >,
122{
123    pub fn new(
124        target_id: GlobalId,
125        map_filter_project: mz_expr::SafeMfpPlan,
126        peek_timestamp: mz_repr::Timestamp,
127        literal_constraints: Option<&mut [Row]>,
128        trace_reader: &mut Tr,
129    ) -> Self {
130        let (mut cursor, storage) = trace_reader.cursor();
131        let literals = literal_constraints
132            .map(|constraints| Literals::new(constraints, &mut cursor, &storage));
133
134        Self {
135            target_id,
136            cursor,
137            storage,
138            map_filter_project,
139            peek_timestamp,
140            row_builder: Row::default(),
141            datum_vec: DatumVec::new(),
142            literals,
143        }
144    }
145
146    /// Returns `true` if the iterator has no more literals to process, or if there are no literals at all.
147    fn literals_exhausted(&self) -> bool {
148        self.literals.as_ref().map_or(false, Literals::is_exhausted)
149    }
150}
151
152impl<Tr> FusedIterator for PeekResultIterator<Tr>
153where
154    Tr: TraceReader<Batch: Navigable>,
155    for<'a> BatchCursor<Tr>: Cursor<
156            Key<'a>: ExtendDatums + Eq,
157            KeyContainer: BatchContainer<Owned = Row>,
158            Val<'a>: ExtendDatums,
159            TimeGat<'a>: PartialOrder<mz_repr::Timestamp>,
160            DiffGat<'a> = &'a Diff,
161        >,
162{
163}
164
165impl<Tr> Iterator for PeekResultIterator<Tr>
166where
167    Tr: TraceReader<Batch: Navigable>,
168    for<'a> BatchCursor<Tr>: Cursor<
169            Key<'a>: ExtendDatums + Eq,
170            KeyContainer: BatchContainer<Owned = Row>,
171            Val<'a>: ExtendDatums,
172            TimeGat<'a>: PartialOrder<mz_repr::Timestamp>,
173            DiffGat<'a> = &'a Diff,
174        >,
175{
176    type Item = Result<(Row, NonZeroI64), String>;
177
178    fn next(&mut self) -> Option<Self::Item> {
179        let result = loop {
180            if self.literals_exhausted() {
181                return None;
182            }
183
184            if !self.cursor.key_valid(&self.storage) {
185                return None;
186            }
187
188            if !self.cursor.val_valid(&self.storage) {
189                let exhausted = self.step_key();
190                if exhausted {
191                    return None;
192                }
193            }
194
195            match self.extract_current_row() {
196                Ok(Some(row)) => break Ok(row),
197                Ok(None) => {
198                    // Have to keep stepping and try with the next val.
199                    self.cursor.step_val(&self.storage);
200                }
201                Err(err) => break Err(err),
202            }
203        };
204
205        self.cursor.step_val(&self.storage);
206
207        Some(result)
208    }
209}
210
211impl<Tr> PeekResultIterator<Tr>
212where
213    Tr: TraceReader<Batch: Navigable>,
214    for<'a> BatchCursor<Tr>: Cursor<
215            Key<'a>: ExtendDatums + Eq,
216            KeyContainer: BatchContainer<Owned = Row>,
217            Val<'a>: ExtendDatums,
218            TimeGat<'a>: PartialOrder<mz_repr::Timestamp>,
219            DiffGat<'a> = &'a Diff,
220        >,
221{
222    /// Extracts and returns the row currently pointed at by our cursor. Returns
223    /// `Ok(None)` if our MapFilterProject evaluates to `None`. Also returns any
224    /// errors that arise from evaluating the MapFilterProject.
225    fn extract_current_row(&mut self) -> Result<Option<(Row, NonZeroI64)>, String> {
226        // TODO: This arena could be maintained and reused for longer,
227        // but it wasn't clear at what interval we should flush
228        // it to ensure we don't accidentally spike our memory use.
229        // This choice is conservative, and not the end of the world
230        // from a performance perspective.
231        let arena = RowArena::new();
232
233        let key_item = self.cursor.key(&self.storage);
234        let row_item = self.cursor.val(&self.storage);
235
236        // An optional literal that we might have added to the borrow. Needs to be declared
237        // before the borrow to ensure correct drop order.
238        let maybe_literal;
239        let mut borrow = self.datum_vec.borrow();
240        key_item.extend_datums(&arena, &mut borrow, None);
241        row_item.extend_datums(&arena, &mut borrow, None);
242
243        if let Some(literals) = &mut self.literals
244            && let Some(literal) = literals.peek()
245        {
246            // The peek was created from an IndexedFilter join. We have to add those columns
247            // here that the join would add in a dataflow.
248            maybe_literal = literal;
249            maybe_literal.extend_datums(&arena, &mut borrow, None);
250        }
251        if let Some(result) = self
252            .map_filter_project
253            .evaluate_into(&mut borrow, &arena, &mut self.row_builder)
254            .map(|row| row.cloned())
255            .map_err_to_string_with_causes()?
256        {
257            let mut copies = Diff::ZERO;
258            self.cursor.map_times(&self.storage, |time, diff| {
259                if time.less_equal(&self.peek_timestamp) {
260                    copies += diff;
261                }
262            });
263            let copies: i64 = if copies.is_negative() {
264                let row = &*borrow;
265                tracing::error!(
266                    target = %self.target_id, diff = %copies, ?row,
267                    "index peek encountered negative multiplicities in ok trace",
268                );
269                return Err(format!(
270                    "Invalid data in source, \
271                             saw retractions ({}) for row that does not exist: {:?}",
272                    -copies, row,
273                ));
274            } else {
275                copies.into_inner()
276            };
277            // if copies > 0 ... otherwise skip
278            if let Some(copies) = NonZeroI64::new(copies) {
279                Ok(Some((result, copies)))
280            } else {
281                Ok(None)
282            }
283        } else {
284            Ok(None)
285        }
286    }
287
288    /// Steps the key forward, respecting literal constraints.
289    ///
290    /// Returns `true` if we are exhausted.
291    fn step_key(&mut self) -> bool {
292        assert!(
293            !self.cursor.val_valid(&self.storage),
294            "must only step key when the vals for a key are exhausted"
295        );
296
297        if let Some(literals) = &mut self.literals {
298            literals.seek_next_literal_key(&mut self.cursor, &self.storage);
299
300            if literals.is_exhausted() {
301                return true;
302            }
303        } else {
304            self.cursor.step_key(&self.storage);
305        }
306
307        if !self.cursor.key_valid(&self.storage) {
308            // We're exhausted!
309            return true;
310        }
311
312        assert!(
313            self.cursor.val_valid(&self.storage),
314            "there must always be at least one val per key"
315        );
316
317        false
318    }
319}
mz_compute/compute_state/peek_result_iterator.rs

mz_compute/compute_state/
peek_result_iterator.rs
