Method render_flat_map

src/compute/src/render/flat_map.rs:29–127

Applies a `TableFunc` to every row, followed by an `mfp`.

(
        &self,
        input_key: Option<Vec<MirScalarExpr>>,
        input: CollectionBundle<'scope, T>,
        exprs: Vec<MirScalarExpr>,
        func: TableFunc,
        mfp: MapFilterProject,
 

impl<'scope, T: crate::render::RenderTimestamp> Context<'scope, T> {
    /// Applies a `TableFunc` to every row, followed by an `mfp`.
    pub fn render_flat_map(
        &self,
        input_key: Option<Vec<MirScalarExpr>>,
        input: CollectionBundle<'scope, T>,
        exprs: Vec<MirScalarExpr>,
        func: TableFunc,
        mfp: MapFilterProject,
    ) -> CollectionBundle<'scope, T> {
        let until = self.until.clone();
        let mfp_plan = mfp.into_plan().expect("MapFilterProject planning failed");
        let (ok_collection, err_collection) =
            input.as_specific_collection(input_key.as_deref(), &self.config_set);
        let stream = ok_collection.inner;
        let scope = input.scope();

        // Budget to limit the number of rows processed in a single invocation.
        //
        // The current implementation can only yield between input batches, but not from within
        // a batch. A `generate_series` can still cause unavailability if it generates many rows.
        let budget = COMPUTE_FLAT_MAP_FUEL.get(&self.config_set);

        let (oks, errs) = stream.unary_fallible(Pipeline, "FlatMapStage", move |_, info| {
            let activator = scope.activator_for(info.address);
            let mut queue = VecDeque::new();
            Box::new(move |input, ok_output, err_output| {
                let mut datums = DatumVec::new();
                let mut datums_mfp = DatumVec::new();

                // Buffer for extensions to `input_row`.
                let mut table_func_output = Vec::new();

                let mut budget = budget;

                input.for_each(|cap, data| {
                    queue.push_back((cap.retain(0), cap.retain(1), std::mem::take(data)))
                });

                while let Some((ok_cap, err_cap, data)) = queue.pop_front() {
                    let mut ok_session = ok_output.session_with_builder(&ok_cap);
                    let mut err_session = err_output.session_with_builder(&err_cap);

                    'input: for (input_row, time, diff) in data {
                        let temp_storage = RowArena::new();

                        // Unpack datums for expression evaluation.
                        let datums_local = datums.borrow_with(&input_row);
                        let args = exprs
                            .iter()
                            .map(|e| e.eval(&datums_local, &temp_storage))
                            .collect::<Result<Vec<_>, _>>();
                        let args = match args {
                            Ok(args) => args,
                            Err(e) => {
                                err_session.give((e.into(), time, diff));
                                continue 'input;
                            }
                        };
                        let mut extensions = match func.eval(&args, &temp_storage) {