Method harvest

src/transform/src/column_knowledge.rs:76–487

Harvest per-column knowledge. `knowledge_stack` is a pre-allocated vector but is expected not to contain any elements.

(
        &self,
        expr: &mut MirRelationExpr,
        knowledge: &mut BTreeMap<mz_expr::Id, Vec<DatumKnowledge>>,
        knowledge_stack: &mut Vec<DatumKnowledge>,
    )

    ///
    /// `knowledge_stack` is a pre-allocated vector but is expected not to contain any elements.
    fn harvest(
        &self,
        expr: &mut MirRelationExpr,
        knowledge: &mut BTreeMap<mz_expr::Id, Vec<DatumKnowledge>>,
        knowledge_stack: &mut Vec<DatumKnowledge>,
    ) -> Result<Vec<DatumKnowledge>, TransformError> {
        self.checked_recur(|_| {
            let result = match expr {
                MirRelationExpr::ArrangeBy { input, .. } => {
                    self.harvest(input, knowledge, knowledge_stack)
                }
                MirRelationExpr::Get { id, typ, .. } => {
                    Ok(knowledge.get(id).cloned().unwrap_or_else(|| {
                        typ.column_types.iter().map(DatumKnowledge::from).collect()
                    }))
                }
                MirRelationExpr::Constant { rows, typ } => {
                    // TODO: handle multi-row cases with some constant columns.
                    if let Ok([(row, _diff)]) = rows.as_deref() {
                        let knowledge = std::iter::zip(row.iter(), typ.column_types.iter())
                            .map(DatumKnowledge::from)
                            .collect();
                        Ok(knowledge)
                    } else {
                        Ok(typ.column_types.iter().map(DatumKnowledge::from).collect())
                    }
                }
                MirRelationExpr::Let { id, value, body } => {
                    let value_knowledge = self.harvest(value, knowledge, knowledge_stack)?;
                    let prior_knowledge =
                        knowledge.insert(mz_expr::Id::Local(id.clone()), value_knowledge);
                    let body_knowledge = self.harvest(body, knowledge, knowledge_stack)?;
                    knowledge.remove(&mz_expr::Id::Local(id.clone()));
                    if let Some(prior_knowledge) = prior_knowledge {
                        knowledge.insert(mz_expr::Id::Local(id.clone()), prior_knowledge);
                    }
                    Ok(body_knowledge)
                }
                MirRelationExpr::LetRec {
                    ids,
                    values,
                    limits,
                    body,
                } => {
                    // Set knowledge[i][j] = DatumKnowledge::bottom() for each
                    // column j and CTE i. This corresponds to the normal
                    // evaluation semantics where each recursive CTE is
                    // initialized to the empty collection.
                    for (id, value) in zip_eq(ids.iter(), values.iter()) {
                        let id = mz_expr::Id::Local(id.clone());
                        let knowledge_new = vec![DatumKnowledge::bottom(); value.arity()];
                        let knowledge_old = knowledge.insert(id, knowledge_new);
                        assert_none!(knowledge_old);
                    }

                    // Sum up the arity of all ids in the enclosing LetRec node.
                    let let_rec_arity = ids.iter().fold(0, |acc, id| {
                        let id = mz_expr::Id::Local(id.clone());