mz_compute/render/
flat_map.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
10use std::collections::VecDeque;
11
12use columnar::{Columnar, Index};
13use differential_dataflow::consolidation::ConsolidatingContainerBuilder;
14use mz_compute_types::dyncfgs::COMPUTE_FLAT_MAP_FUEL;
15use mz_compute_types::plan::scalar::LirScalarExpr;
16use mz_expr::TableFunc;
17use mz_expr::{Eval, MfpPlan};
18use mz_repr::{DatumVec, RowArena, SharedRow};
19use mz_repr::{Diff, Row, RowRef, Timestamp};
20use mz_timely_util::columnar::Column;
21use mz_timely_util::columnar::consolidate::ConsolidatingColumnBuilder;
22use mz_timely_util::operator::StreamExt;
23use timely::dataflow::channels::pact::Pipeline;
24use timely::dataflow::operators::Capability;
25use timely::dataflow::operators::generic::Session;
26use timely::dataflow::{Scope, Stream};
27use timely::progress::Antichain;
28
29use crate::render::RenderTimestamp;
30use crate::render::context::{CollectionBundle, Context};
31use crate::render::errors::DataflowErrorSer;
32
33impl<'scope, T: crate::render::RenderTimestamp> Context<'scope, T> {
34    /// Applies a `TableFunc` to every row, followed by an `mfp`.
35    pub fn render_flat_map(
36        &self,
37        input_key: Option<Vec<LirScalarExpr>>,
38        input: CollectionBundle<'scope, T>,
39        exprs: Vec<LirScalarExpr>,
40        func: TableFunc,
41        mfp_plan: MfpPlan<LirScalarExpr>,
42    ) -> CollectionBundle<'scope, T> {
43        let until = self.until.clone();
44        let scope = input.scope();
45
46        // Budget to limit the number of rows processed in a single invocation.
47        //
48        // The current implementation can only yield between input batches, but not from within
49        // a batch. A `generate_series` can still cause unavailability if it generates many rows.
50        let budget = COMPUTE_FLAT_MAP_FUEL.get(&self.config_set);
51
52        // The unarranged path reads the edge directly, so a columnar input is never
53        // decoded here. The keyed path materializes an existing arrangement, which
54        // `as_specific_collection` presents as a columnar edge.
55        let (edge, err_collection) = match input_key.as_deref() {
56            None => input
57                .collection
58                .clone()
59                .expect("The unarranged collection doesn't exist."),
60            Some(key) => input.as_specific_collection(Some(key)),
61        };
62
63        let (oks, errs) = flat_map_stage(edge.inner, scope, exprs, func, mfp_plan, until, budget);
64
65        use differential_dataflow::AsCollection;
66        let ok_collection = oks.as_collection();
67        let new_err_collection = errs.as_collection();
68        let err_collection = err_collection.concat(new_err_collection);
69        CollectionBundle::from_edge(ok_collection, err_collection)
70    }
71}
72
73/// Output ok-session container builder for [`flat_map_stage`].
74///
75/// Consolidating like the err builder, but emits `Column<(Row, T, Diff)>`. The mfp builds
76/// each output row fresh, so the owned give into staging is a move.
77type FlatMapOk<T> = ConsolidatingColumnBuilder<Row, T, Diff>;
78/// Output err-session container builder for [`flat_map_stage`].
79type FlatMapErr<T> = ConsolidatingContainerBuilder<Vec<(DataflowErrorSer, T, Diff)>>;
80
81/// The fueled FlatMap operator.
82///
83/// Each activation expands queued batches until the `budget` runs out, then re-activates
84/// and defers the rest of the queue, which bounds what one `generate_series` does before
85/// the worker yields.
86fn flat_map_stage<'scope, T>(
87    stream: Stream<'scope, T, Column<(Row, T, Diff)>>,
88    scope: Scope<'scope, T>,
89    exprs: Vec<LirScalarExpr>,
90    func: TableFunc,
91    mfp_plan: MfpPlan<LirScalarExpr>,
92    until: Antichain<Timestamp>,
93    budget: usize,
94) -> (
95    Stream<'scope, T, Column<(Row, T, Diff)>>,
96    Stream<'scope, T, Vec<(DataflowErrorSer, T, Diff)>>,
97)
98where
99    T: RenderTimestamp,
100{
101    stream.unary_fallible::<FlatMapOk<T>, FlatMapErr<T>, _, _>(
102        Pipeline,
103        "FlatMapStage",
104        move |_, info| {
105            let activator = scope.activator_for(info.address);
106            let mut queue = VecDeque::new();
107            Box::new(move |input, ok_output, err_output| {
108                let mut datums = DatumVec::new();
109                let mut datums_mfp = DatumVec::new();
110
111                // Buffer for extensions to `input_row`.
112                let mut table_func_output = Vec::new();
113
114                // Reused so a record's time and diff do not allocate per record.
115                let mut time = T::minimum();
116                let mut diff = Diff::ZERO;
117
118                let mut budget = budget;
119
120                input.for_each(|cap, data| {
121                    queue.push_back((cap.retain(0), cap.retain(1), std::mem::take(data)))
122                });
123
124                while let Some((ok_cap, err_cap, data)) = queue.pop_front() {
125                    let mut ok_session = ok_output.session_with_builder(&ok_cap);
126                    let mut err_session = err_output.session_with_builder(&err_cap);
127
128                    // Rows stay borrowed. The time and diff have to be owned to pass
129                    // them by reference, so they are copied into buffers rather than
130                    // built fresh: an iterative `T` owns a `PointStamp`'s allocation.
131                    for (input_row, t, d) in data.borrow().into_index_iter() {
132                        time.copy_from(t);
133                        diff.copy_from(d);
134                        process_flat_map_row(
135                            input_row,
136                            &time,
137                            &diff,
138                            &exprs,
139                            &func,
140                            &mfp_plan,
141                            &until,
142                            &mut datums,
143                            &mut datums_mfp,
144                            &mut table_func_output,
145                            &mut ok_session,
146                            &mut err_session,
147                            &mut budget,
148                        );
149                    }
150                    if budget == 0 {
151                        activator.activate();
152                        break;
153                    }
154                }
155            })
156        },
157    )
158}
159
160/// Expands one input record's table function and drains it through the mfp.
161///
162/// The expansion is chunked so [`drain_through_mfp`] amortizes the input-row decode.
163/// Argument or function evaluation errors emit to the err session and return early.
164fn process_flat_map_row<T>(
165    input_row: &RowRef,
166    time: &T,
167    diff: &Diff,
168    exprs: &[LirScalarExpr],
169    func: &TableFunc,
170    mfp_plan: &MfpPlan<LirScalarExpr>,
171    until: &Antichain<Timestamp>,
172    datums: &mut DatumVec,
173    datums_mfp: &mut DatumVec,
174    table_func_output: &mut Vec<(Row, Diff)>,
175    ok_session: &mut Session<'_, '_, T, FlatMapOk<T>, Capability<T>>,
176    err_session: &mut Session<'_, '_, T, FlatMapErr<T>, Capability<T>>,
177    budget: &mut usize,
178) where
179    T: RenderTimestamp,
180{
181    let temp_storage = RowArena::new();
182
183    // Unpack datums for expression evaluation.
184    let datums_local = datums.borrow_with(input_row);
185    let args = exprs
186        .iter()
187        .map(|e| e.eval(&datums_local, &temp_storage))
188        .collect::<Result<Vec<_>, _>>();
189    let args = match args {
190        Ok(args) => args,
191        Err(e) => {
192            err_session.give((e.into(), time.clone(), *diff));
193            return;
194        }
195    };
196    let mut extensions = match func.eval(&args, &temp_storage) {
197        Ok(exts) => exts.fuse(),
198        Err(e) => {
199            err_session.give((e.into(), time.clone(), *diff));
200            return;
201        }
202    };
203
204    // Draw additional columns out of the table func evaluation.
205    while let Some((extension, output_diff)) = extensions.next() {
206        table_func_output.push((extension, output_diff));
207        table_func_output.extend((&mut extensions).take(1023));
208        // We could consolidate `table_func_output`, but it seems unlikely to be productive.
209        drain_through_mfp(
210            input_row,
211            time,
212            diff,
213            datums_mfp,
214            table_func_output,
215            mfp_plan,
216            until,
217            ok_session,
218            err_session,
219            budget,
220        );
221        table_func_output.clear();
222    }
223}
224
225/// Drains a list of extensions to `input_row` through a supplied `MfpPlan` and into output buffers.
226///
227/// The method decodes `input_row`, and should be amortized across non-trivial `extensions`.
228fn drain_through_mfp<T>(
229    input_row: &RowRef,
230    input_time: &T,
231    input_diff: &Diff,
232    datum_vec: &mut DatumVec,
233    extensions: &[(Row, Diff)],
234    mfp_plan: &MfpPlan<LirScalarExpr>,
235    until: &Antichain<Timestamp>,
236    ok_output: &mut Session<'_, '_, T, FlatMapOk<T>, Capability<T>>,
237    err_output: &mut Session<'_, '_, T, FlatMapErr<T>, Capability<T>>,
238    budget: &mut usize,
239) where
240    T: RenderTimestamp,
241{
242    let temp_storage = RowArena::new();
243    let mut row_builder = SharedRow::get();
244
245    // This is not cheap, and is meant to be amortized across many `extensions`.
246    let mut datums_local = datum_vec.borrow_with(input_row);
247    let datums_len = datums_local.len();
248
249    let event_time = input_time.event_time().clone();
250
251    for (cols, diff) in extensions.iter() {
252        // Arrange `datums_local` to reflect the intended output pre-mfp.
253        datums_local.truncate(datums_len);
254        datums_local.extend(cols.iter());
255
256        let results = mfp_plan.evaluate(
257            &mut datums_local,
258            &temp_storage,
259            event_time,
260            *diff * *input_diff,
261            |time| !until.less_equal(time),
262            &mut row_builder,
263        );
264
265        for result in results {
266            *budget = budget.saturating_sub(1);
267            match result {
268                Ok((row, event_time, diff)) => {
269                    // Copy the whole time, and re-populate event time.
270                    let mut time = input_time.clone();
271                    *time.event_time_mut() = event_time;
272                    ok_output.give((row, time, diff));
273                }
274                Err((err, event_time, diff)) => {
275                    // Copy the whole time, and re-populate event time.
276                    let mut time = input_time.clone();
277                    *time.event_time_mut() = event_time;
278                    err_output.give((err, time, diff));
279                }
280            };
281        }
282    }
283}
284
285#[cfg(test)]
286mod tests {
287    use std::cell::RefCell;
288    use std::rc::Rc;
289
290    use differential_dataflow::input::Input;
291    use mz_expr::MapFilterProject;
292    use mz_repr::{Datum, ReprScalarType};
293    use timely::dataflow::operators::InspectCore;
294    use timely::dataflow::operators::capture::{Capture, Extract};
295
296    use super::*;
297    use crate::render::columnar::vec_to_columnar;
298
299    // `generate_series(1, stop, 1)`, reading `stop` from column 1, with an identity mfp.
300    fn flat_map_args() -> (Vec<LirScalarExpr>, TableFunc, MfpPlan<LirScalarExpr>) {
301        let exprs = vec![
302            LirScalarExpr::column(0),
303            LirScalarExpr::column(1),
304            LirScalarExpr::literal_ok(Datum::Int64(1), ReprScalarType::Int64),
305        ];
306        let func = TableFunc::GenerateSeriesInt64;
307        let mfp = MapFilterProject::<LirScalarExpr>::new(3)
308            .into_plan()
309            .expect("identity mfp");
310        (exprs, func, mfp)
311    }
312
313    fn input_row(stop: i64) -> Row {
314        Row::pack_slice(&[Datum::Int64(1), Datum::Int64(stop)])
315    }
316
317    /// Runs one input row at each of `batches` distinct timestamps through `flat_map_stage`,
318    /// stepping the worker manually. Returns the steps taken and the output-record count.
319    fn run_fueled(batches: u64, stop: i64, budget: usize) -> (usize, usize) {
320        let expected = usize::try_from(batches).unwrap() * usize::try_from(stop).unwrap();
321        timely::execute_directly(move |worker| {
322            let collected = Rc::new(RefCell::new(0usize));
323            let sink = Rc::clone(&collected);
324            let mut input = worker.dataflow::<Timestamp, _, _>(|scope| {
325                let (input, collection) = scope.new_collection();
326                let (exprs, func, mfp) = flat_map_args();
327                // Feed the operator the columnar edge it is given in production,
328                // so the fuel assertions below cover the shipped path.
329                let stream = vec_to_columnar(collection).inner;
330                let scope = stream.scope();
331                let (oks, _errs) =
332                    flat_map_stage(stream, scope, exprs, func, mfp, Antichain::new(), budget);
333                // Counted per container: a per-record `inspect` needs
334                // `&Container: IntoIterator`, which on macOS recurses through `objc2`'s
335                // blanket impls until the trait solver overflows.
336                oks.inspect_container(move |event| {
337                    if let Ok((_time, data)) = event {
338                        *sink.borrow_mut() += data.borrow().into_index_iter().count();
339                    }
340                });
341                input
342            });
343
344            // One batch per timestamp, so the per-batch output does not consolidate away.
345            for i in 0..batches {
346                input.advance_to(Timestamp::from(i));
347                input.update(input_row(stop), Diff::ONE);
348                input.flush();
349            }
350            input.advance_to(Timestamp::from(batches));
351            input.flush();
352
353            let mut steps = 0;
354            while *collected.borrow() < expected {
355                worker.step();
356                steps += 1;
357                assert!(steps < 10_000, "flat map did not converge");
358            }
359            (steps, *collected.borrow())
360        })
361    }
362
363    #[mz_ore::test]
364    fn flat_map_fuel_bounds_per_activation() {
365        let (fueled_steps, fueled_count) = run_fueled(4, 3, 1);
366        let (unfueled_steps, unfueled_count) = run_fueled(4, 3, usize::MAX);
367        assert_eq!(fueled_count, 12);
368        assert_eq!(unfueled_count, 12);
369        assert_eq!(
370            unfueled_steps, 1,
371            "unbounded budget drains in one activation"
372        );
373        assert!(
374            fueled_steps > unfueled_steps,
375            "fuel budget must spread work across activations: fueled={fueled_steps} unfueled={unfueled_steps}"
376        );
377    }
378
379    #[mz_ore::test]
380    fn flat_map_reads_columnar_input() {
381        // Several timestamps and a retraction, so time handling and a negative diff are
382        // both decoded.
383        let captured = timely::execute_directly(move |worker| {
384            worker.dataflow::<Timestamp, _, _>(|scope| {
385                let (mut input, collection) = scope.new_collection();
386                let (exprs, func, mfp) = flat_map_args();
387                let stream = vec_to_columnar(collection).inner;
388                let scope = stream.scope();
389                let (oks, _errs) = flat_map_stage(
390                    stream,
391                    scope,
392                    exprs,
393                    func,
394                    mfp,
395                    Antichain::new(),
396                    usize::MAX,
397                );
398                let captured = oks.capture();
399                // t=0: generate_series(1, 2); t=1: generate_series(1, 3);
400                // t=2: retract the t=0 row.
401                input.advance_to(Timestamp::from(0_u64));
402                input.update(input_row(2), Diff::ONE);
403                input.advance_to(Timestamp::from(1_u64));
404                input.update(input_row(3), Diff::ONE);
405                input.advance_to(Timestamp::from(2_u64));
406                input.update(input_row(2), -Diff::ONE);
407                input.advance_to(Timestamp::from(3_u64));
408                input.flush();
409                captured
410            })
411        });
412
413        let updates = extract_sorted_columns(captured);
414        assert!(!updates.is_empty());
415        assert!(
416            updates.iter().any(|(_, _, d)| *d < Diff::ZERO),
417            "the retraction must survive as a negative diff"
418        );
419    }
420
421    /// Decodes a capture of the columnar output into sorted `(row, time, diff)` updates.
422    fn extract_sorted_columns(
423        captured: std::sync::mpsc::Receiver<
424            timely::dataflow::operators::capture::Event<Timestamp, Column<(Row, Timestamp, Diff)>>,
425        >,
426    ) -> Vec<(Row, Timestamp, Diff)> {
427        let mut updates: Vec<(Row, Timestamp, Diff)> = captured
428            .extract()
429            .into_iter()
430            .flat_map(|(_, col)| {
431                col.borrow()
432                    .into_index_iter()
433                    .map(|(v, t, d)| {
434                        (
435                            Columnar::into_owned(v),
436                            Columnar::into_owned(t),
437                            Columnar::into_owned(d),
438                        )
439                    })
440                    .collect::<Vec<_>>()
441            })
442            .collect();
443        updates.sort();
444        updates
445    }
446
447    #[mz_ore::test]
448    fn flat_map_output_consolidates_within_batch() {
449        // Two input rows whose expansions overlap once the mfp projects away the
450        // differing `stop` column, so the output builder has duplicates to fold.
451        let captured = timely::execute_directly(move |worker| {
452            worker.dataflow::<Timestamp, _, _>(|scope| {
453                let (mut input, collection) = scope.new_collection();
454                let exprs = vec![
455                    LirScalarExpr::column(0),
456                    LirScalarExpr::column(1),
457                    LirScalarExpr::literal_ok(Datum::Int64(1), ReprScalarType::Int64),
458                ];
459                let func = TableFunc::GenerateSeriesInt64;
460                // Project to the generated value alone, collapsing the distinct prefixes.
461                let mfp = MapFilterProject::<LirScalarExpr>::new(3)
462                    .project(vec![2])
463                    .into_plan()
464                    .expect("project mfp");
465                let stream = vec_to_columnar(collection).inner;
466                let scope = stream.scope();
467                let (oks, _errs) = flat_map_stage(
468                    stream,
469                    scope,
470                    exprs,
471                    func,
472                    mfp,
473                    Antichain::new(),
474                    usize::MAX,
475                );
476                let captured = oks.capture();
477                // Both at t=0: {1, 2} and {1, 2, 3}, so 1 and 2 fold to a diff of two.
478                input.advance_to(Timestamp::from(0_u64));
479                input.update(input_row(2), Diff::ONE);
480                input.update(input_row(3), Diff::ONE);
481                input.advance_to(Timestamp::from(1_u64));
482                input.flush();
483                captured
484            })
485        });
486
487        let updates = extract_sorted_columns(captured);
488        let expected = vec![
489            (
490                Row::pack_slice(&[Datum::Int64(1)]),
491                Timestamp::from(0_u64),
492                Diff::from(2),
493            ),
494            (
495                Row::pack_slice(&[Datum::Int64(2)]),
496                Timestamp::from(0_u64),
497                Diff::from(2),
498            ),
499            (
500                Row::pack_slice(&[Datum::Int64(3)]),
501                Timestamp::from(0_u64),
502                Diff::ONE,
503            ),
504        ];
505        assert_eq!(updates, expected);
506    }
507}
mz_compute/render/flat_map.rs

mz_compute/render/
flat_map.rs
