mz_compute_types/

plan.rs

// Copyright Materialize, Inc. and contributors. All rights reserved.
//
// Use of this software is governed by the Business Source License
// included in the LICENSE file.
//
// As of the Change Date specified in that file, in accordance with
// the Business Source License, use of this software will be governed
// by the Apache License, Version 2.0.

//! An explicit representation of a rendering plan for provided dataflows.

#![warn(missing_debug_implementations)]

use std::collections::{BTreeMap, BTreeSet};

use columnar::Columnar;
use mz_expr::{
    CollectionPlan, EvalError, Id, LetRecLimit, LocalId, MapFilterProject, MfpPlan,
    OptimizedMirRelationExpr, SafeMfpPlan, TableFunc,
};
use mz_ore::metric;
use mz_ore::metrics::MetricsRegistry;
use mz_ore::metrics::raw::IntCounterVec;
use mz_ore::soft_assert_eq_no_log;
use mz_ore::str::Indent;
use mz_repr::explain::text::text_string_at;
use mz_repr::explain::{DummyHumanizer, ExplainConfig, ExprHumanizer, PlanRenderingContext};
use mz_repr::optimize::OptimizerFeatures;
use mz_repr::{Diff, GlobalId, StableRow, Timestamp};
use serde::{Deserialize, Serialize};

use crate::dataflows::{DataflowDescription, LirDataflowDescription};
use crate::plan::join::JoinPlan;
use crate::plan::reduce::{KeyValPlan, ReducePlan};
use crate::plan::scalar::LirScalarExpr;
use crate::plan::threshold::ThresholdPlan;
use crate::plan::top_k::TopKPlan;
use crate::plan::transform::{Transform, TransformConfig};

mod lowering;
pub mod pinned;

pub mod interpret;
pub mod join;
pub mod reduce;
pub mod render_plan;
pub mod scalar;
pub mod threshold;
pub mod top_k;
pub mod transform;

/// Metrics collected during MIR to LIR lowering.
#[derive(Debug, Clone)]
pub struct LoweringMetrics {
    /// Counts non-`None` results of `MapFilterProject::literal_constraints` during lowering,
    /// labeled by the call site (`"get"` or `"mfp"`).
    literal_constraints: IntCounterVec,
}

impl LoweringMetrics {
    /// Registers the lowering metrics into `registry`.
    pub fn register_into(registry: &MetricsRegistry) -> Self {
        Self {
            literal_constraints: registry.register(metric!(
                name: "mz_optimizer_lowering_literal_constraints_total",
                help: "How often the MFP-based literal-constraint detector succeeded, by call site.",
                var_labels: ["case"],
            )),
        }
    }

    /// Records that a `literal_constraints` call at `case` produced a usable constraint.
    pub fn inc_literal_constraints(&self, case: &str) {
        self.literal_constraints.with_label_values(&[case]).inc();
    }
}

/// The forms in which an operator's output is available.
///
/// These forms may include "raw", meaning as a streamed collection, but also any
/// number of "arranged" representations.
///
/// Each arranged representation is described by a `(to_key, permutation, thinning)`
/// triple, built by `permutation_for_arrangement`. `to_key` (length `K`) are the key
/// expressions over a row. `permutation` (length `A`, the raw/unthinned arity) maps
/// each row column to its position in the `(key, value)` concatenation. `thinning`
/// (length `M`) lists the row columns that form the value, in value order, so a value
/// datum at concatenation position `c >= K` came from row column `thinning[c - K]`.
///
/// This triple is unrelated to `KeyValRowMapping`'s `(to_key, to_val, to_row)` fields
/// despite the visual resemblance: `permutation` here plays the role of
/// `KeyValRowMapping::to_row`, and `thinning` plays the role of
/// `KeyValRowMapping::to_val`. Do not assume the same field order.
#[derive(
    Clone,
    Debug,
    Default,
    Deserialize,
    Eq,
    Ord,
    PartialEq,
    PartialOrd,
    Serialize
)]
pub struct AvailableCollections {
    /// Whether the collection exists in unarranged form.
    pub raw: bool,
    /// The list of available arrangements, each a `(to_key, permutation, thinning)`
    /// triple. See the struct-level documentation for field semantics.
    pub arranged: Vec<(Vec<LirScalarExpr>, Vec<usize>, Vec<usize>)>,
}

impl AvailableCollections {
    /// Represent a collection that has no arrangements.
    pub fn new_raw() -> Self {
        Self {
            raw: true,
            arranged: Vec::new(),
        }
    }

    /// Represent a collection that is arranged in the specified ways.
    pub fn new_arranged(arranged: Vec<(Vec<LirScalarExpr>, Vec<usize>, Vec<usize>)>) -> Self {
        assert!(
            !arranged.is_empty(),
            "Invariant violated: at least one collection must exist"
        );
        Self {
            raw: false,
            arranged,
        }
    }

    /// Get some arrangement, if one exists.
    pub fn arbitrary_arrangement(&self) -> Option<&(Vec<LirScalarExpr>, Vec<usize>, Vec<usize>)> {
        assert!(
            self.raw || !self.arranged.is_empty(),
            "Invariant violated: at least one collection must exist"
        );
        self.arranged.get(0)
    }
}

/// How to render the arrangements requested by an `ArrangeBy`.
///
/// Decided during LIR lowering and consumed by the renderer. The variant says what the
/// renderer will do, not what it knows about the input.
#[derive(
    Clone,
    Copy,
    Debug,
    Deserialize,
    Eq,
    Ord,
    PartialEq,
    PartialOrd,
    Serialize
)]
pub enum ArrangementStrategy {
    /// Form arrangements directly from the input collection.
    Direct,
    /// Insert temporal bucketing in front of the arrangement, to delay future-stamped
    /// updates (e.g., from `mz_now()` MFPs) until their bucket boundary releases them.
    TemporalBucketing,
}

impl std::fmt::Display for ArrangementStrategy {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            ArrangementStrategy::Direct => write!(f, "Direct"),
            ArrangementStrategy::TemporalBucketing => write!(f, "TemporalBucketing"),
        }
    }
}

/// An identifier for an LIR node.
#[derive(
    Clone,
    Copy,
    Debug,
    Deserialize,
    Eq,
    Ord,
    PartialEq,
    PartialOrd,
    Serialize,
    Columnar
)]
pub struct LirId(u64);

impl LirId {
    fn as_u64(&self) -> u64 {
        self.0
    }
}

impl From<LirId> for u64 {
    fn from(value: LirId) -> Self {
        value.as_u64()
    }
}

impl std::fmt::Display for LirId {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self.0)
    }
}

/// Version of the stable LIR serialization format.
///
/// Once a version has shipped, meaning a released Materialize durably stores
/// pinned plans in it, bump this when the serialized representation of
/// `PinnedDataflow` or anything it transitively contains changes, or when
/// a scalar function's declared properties change. Whether the current
/// version has shipped is recorded in [`LIR_VERSION_POLICY`]. Two snapshot
/// tests key off it: the schema snapshot in `tests/lir_schema.rs` against
/// `tests/snapshots/lir_v{LIR_VERSION}.json`, and the function property
/// registry in `tests/func_registry.rs` against
/// `tests/snapshots/func_registry.json` with a per-version digest in
/// `tests/snapshots/func_registry_digests.json`.
pub const LIR_VERSION: u64 = 1;

/// What a change to the stable LIR format must do at the current
/// [`LIR_VERSION`]. The snapshot tests print it when a snapshot goes stale.
///
/// This is the one place that records whether the current version has
/// shipped. Once a released Materialize durably stores pinned plans, replace
/// the text with the shipped policy: bump `LIR_VERSION` rather than rewriting
/// the shipped version's snapshots, so the old schema stays available to
/// migration tooling.
pub const LIR_VERSION_POLICY: &str = "\
The current LIR version has NOT shipped: no released Materialize stores pinned
LIR plans yet, so nothing depends on the checked-in snapshots. Regenerate them
in place and include the diff in your PR.

Once pinned plans are durably stored, a change like this must instead bump
LIR_VERSION in src/compute-types/src/plan.rs so it lands as a new version.
See doc/developer/design/20260826_pinned_lir.md.";

pub use constant_rows_serde::ConstantRows;

/// Serializes `LirRelationNode::Constant`'s rows through the named
/// [`ConstantRows`] mirror enum instead of std `Result`.
///
/// The stable LIR schema registry maps each container name to a single
/// format, and `Result` would clash with the differently instantiated
/// `Result` in `LirScalarExpr::Literal`. The mirror has the same variant
/// order as `Result`, so the encoded bytes are unchanged.
mod constant_rows_serde {
    use mz_expr::{EvalError, StableEvalError, StableEvalErrorRef};
    use mz_repr::{Diff, StableRow, Timestamp};
    use serde::{Deserialize, Deserializer, Serialize, Serializer};

    /// The serialized form of `LirRelationNode::Constant`'s rows.
    #[derive(Debug, Serialize, Deserialize)]
    pub enum ConstantRows {
        /// See `Result::Ok`.
        Ok(Vec<(StableRow, Timestamp, Diff)>),
        /// See `Result::Err`.
        Err(StableEvalError),
    }

    /// Borrowing mirror of [`ConstantRows`], to serialize without cloning.
    #[derive(Serialize)]
    #[serde(rename = "ConstantRows")]
    enum ConstantRowsRef<'a> {
        Ok(&'a Vec<(StableRow, Timestamp, Diff)>),
        Err(StableEvalErrorRef<'a>),
    }

    pub fn serialize<S: Serializer>(
        rows: &Result<Vec<(StableRow, Timestamp, Diff)>, EvalError>,
        serializer: S,
    ) -> Result<S::Ok, S::Error> {
        let mirror = match rows {
            Ok(rows) => ConstantRowsRef::Ok(rows),
            Err(err) => ConstantRowsRef::Err(StableEvalErrorRef(err)),
        };
        mirror.serialize(serializer)
    }

    pub fn deserialize<'de, D: Deserializer<'de>>(
        deserializer: D,
    ) -> Result<Result<Vec<(StableRow, Timestamp, Diff)>, EvalError>, D::Error> {
        Ok(match ConstantRows::deserialize(deserializer)? {
            ConstantRows::Ok(rows) => Ok(rows),
            ConstantRows::Err(err) => Err(err.0),
        })
    }
}

/// A rendering plan with as much conditional logic as possible removed.
#[derive(Clone, Debug, Deserialize, Eq, Ord, PartialEq, PartialOrd, Serialize)]
pub struct LirRelationExpr {
    /// A dataflow-local identifier.
    pub lir_id: LirId,
    /// The underlying operator.
    pub node: LirRelationNode,
}

/// The actual AST node of the `LirRelationExpr`.
#[derive(Clone, Debug, Deserialize, Eq, Ord, PartialEq, PartialOrd, Serialize)]
pub enum LirRelationNode {
    /// A collection containing a pre-determined collection.
    Constant {
        /// Explicit update triples for the collection.
        #[serde(with = "constant_rows_serde")]
        rows: Result<Vec<(StableRow, Timestamp, Diff)>, EvalError>,
    },
    /// A reference to a bound collection.
    ///
    /// This is commonly either an external reference to an existing source or
    /// maintained arrangement, or an internal reference to a `Let` identifier.
    Get {
        /// A global or local identifier naming the collection.
        id: Id,
        /// Arrangements that will be available.
        ///
        /// The collection will also be loaded if available, which it will
        /// not be for imported data, but which it may be for locally defined
        /// data.
        // TODO: Be more explicit about whether a collection is available,
        // although one can always produce it from an arrangement, and it
        // seems generally advantageous to do that instead (to avoid cloning
        // rows, by using `mfp` first on borrowed data).
        keys: AvailableCollections,
        /// The actions to take when introducing the collection.
        plan: GetPlan,
    },
    /// Binds `value` to `id`, and then results in `body` with that binding.
    ///
    /// This stage has the effect of sharing `value` across multiple possible
    /// uses in `body`, and is the only mechanism we have for sharing collection
    /// information across parts of a dataflow.
    ///
    /// The binding is not available outside of `body`.
    Let {
        /// The local identifier to be used, available to `body` as `Id::Local(id)`.
        id: LocalId,
        /// The collection that should be bound to `id`.
        value: Box<LirRelationExpr>,
        /// The collection that results, which is allowed to contain `Get` stages
        /// that reference `Id::Local(id)`.
        body: Box<LirRelationExpr>,
    },
    /// Binds `values` to `ids`, evaluates them potentially recursively, and returns `body`.
    ///
    /// All bindings are available to all bindings, and to `body`.
    /// The contents of each binding are initially empty, and then updated through a sequence
    /// of iterations in which each binding is updated in sequence, from the most recent values
    /// of all bindings.
    LetRec {
        /// The local identifiers to be used, available to `body` as `Id::Local(id)`.
        ids: Vec<LocalId>,
        /// The collection that should be bound to `id`.
        values: Vec<LirRelationExpr>,
        /// Maximum number of iterations. See further info on the MIR `LetRec`.
        limits: Vec<Option<LetRecLimit>>,
        /// The collection that results, which is allowed to contain `Get` stages
        /// that reference `Id::Local(id)`.
        body: Box<LirRelationExpr>,
    },
    /// Map, Filter, and Project operators.
    ///
    /// This stage contains work that we would ideally like to fuse to other plan
    /// stages, but for practical reasons cannot. For example: threshold, topk,
    /// and sometimes reduce stages are not able to absorb this operator.
    Mfp {
        /// The input collection.
        input: Box<LirRelationExpr>,
        /// Linear operator to apply to each record.
        mfp: MfpPlan<LirScalarExpr>,
        /// Whether the input is from an arrangement, and if so,
        /// whether we can seek to a specific value therein
        input_key_val: Option<(Vec<LirScalarExpr>, Option<StableRow>)>,
    },
    /// A variable number of output records for each input record.
    ///
    /// This stage is a bit of a catch-all for logic that does not easily fit in
    /// map stages. This includes table valued functions, but also functions of
    /// multiple arguments, and functions that modify the sign of updates.
    ///
    /// This stage allows a `MapFilterProject` operator to be fused to its output,
    /// and this can be very important as otherwise the output of `func` is just
    /// appended to the input record, for as many outputs as it has. This has the
    /// unpleasant default behavior of repeating potentially large records that
    /// are being unpacked, producing quadratic output in those cases. Instead,
    /// in these cases use a `mfp` member that projects away these large fields.
    FlatMap {
        /// The particular arrangement of the input we expect to use,
        /// if any
        input_key: Option<Vec<LirScalarExpr>>,
        /// The input collection.
        input: Box<LirRelationExpr>,
        /// Expressions that for each row prepare the arguments to `func`.
        exprs: Vec<LirScalarExpr>,
        /// The variable-record emitting function.
        func: TableFunc,
        /// Linear operator to apply to each record produced by `func`.
        mfp_after: MfpPlan<LirScalarExpr>,
    },
    /// A multiway relational equijoin, with fused map, filter, and projection.
    ///
    /// This stage performs a multiway join among `inputs`, using the equality
    /// constraints expressed in `plan`. The plan also describes the implementation
    /// strategy we will use, and any pushed down per-record work.
    Join {
        /// An ordered list of inputs that will be joined.
        inputs: Vec<LirRelationExpr>,
        /// Detailed information about the implementation of the join.
        ///
        /// This includes information about the implementation strategy, but also
        /// any map, filter, project work that we might follow the join with, but
        /// potentially pushed down into the implementation of the join.
        plan: JoinPlan,
    },
    /// Aggregation by key.
    Reduce {
        /// The particular arrangement of the input we expect to use,
        /// if any
        input_key: Option<Vec<LirScalarExpr>>,
        /// The input collection.
        input: Box<LirRelationExpr>,
        /// A plan for changing input records into key, value pairs.
        key_val_plan: KeyValPlan,
        /// A plan for performing the reduce.
        ///
        /// The implementation of reduction has several different strategies based
        /// on the properties of the reduction, and the input itself. Please check
        /// out the documentation for this type for more detail.
        plan: ReducePlan,
        /// An MFP that must be applied to results. The projection part of this
        /// MFP must preserve the key for the reduction; otherwise, the results
        /// become undefined. Additionally, the MFP is guaranteed to be free from
        /// temporal predicates so that it can be readily evaluated.
        mfp_after: SafeMfpPlan<LirScalarExpr>,
        /// Strategy for forming the internal input arrangement built by `Reduce`
        /// (materialized via `key_val_plan`).
        ///
        /// Set by the lowering from the input's `has_future_updates` flag. The
        /// renderer applies it to the keyed `(key, val)` stream feeding the
        /// reduce. See `render_reduce` for the rationale on why this is
        /// plumbed through `Reduce` rather than handled at the arrangement site.
        ///
        /// Note: unrelated to the hash buckets used by hierarchical reductions
        /// (e.g. `ReducePlan::Hierarchical`'s `buckets`), which are an internal
        /// sharding scheme for `min`/`max`-style aggregations. Here "bucketing"
        /// refers exclusively to temporal (time-domain) bucketing of
        /// future-stamped updates.
        temporal_bucketing_strategy: ArrangementStrategy,
    },
    /// Key-based "Top K" operator, retaining the first K records in each group.
    TopK {
        /// The input collection.
        input: Box<LirRelationExpr>,
        /// A plan for performing the Top-K.
        ///
        /// The implementation of reduction has several different strategies based
        /// on the properties of the reduction, and the input itself. Please check
        /// out the documentation for this type for more detail.
        top_k_plan: TopKPlan,
        /// Strategy for bucketing the input collection ahead of the Top-K operator.
        ///
        /// Set by the lowering from the input's `has_future_updates` flag. The
        /// renderer applies it to the per-row input stream at the top of
        /// `render_topk`, covering all three `TopKPlan` arms uniformly. See
        /// `LirRelationNode::Reduce::temporal_bucketing_strategy` for the underlying
        /// convention.
        temporal_bucketing_strategy: ArrangementStrategy,
    },
    /// Inverts the sign of each update.
    Negate {
        /// The input collection.
        input: Box<LirRelationExpr>,
    },
    /// Filters records that accumulate negatively.
    ///
    /// Although the operator suppresses updates, it is a stateful operator taking
    /// resources proportional to the number of records with non-zero accumulation.
    Threshold {
        /// The input collection.
        input: Box<LirRelationExpr>,
        /// A plan for performing the threshold.
        ///
        /// The implementation of reduction has several different strategies based
        /// on the properties of the reduction, and the input itself. Please check
        /// out the documentation for this type for more detail.
        threshold_plan: ThresholdPlan,
    },
    /// Adds the contents of the input collections.
    ///
    /// Importantly, this is *multiset* union, so the multiplicities of records will
    /// add. This is in contrast to *set* union, where the multiplicities would be
    /// capped at one. A set union can be formed with `Union` followed by `Reduce`
    /// implementing the "distinct" operator.
    Union {
        /// The input collections
        inputs: Vec<LirRelationExpr>,
        /// Whether to consolidate the output, e.g., cancel negated records.
        consolidate_output: bool,
        /// Per-input bucketing strategies. Lockstep with `inputs`: index `i` is the
        /// strategy applied to `inputs[i]` before concatenation.
        ///
        /// Set by the lowering from each input's `has_future_updates` flag. Only
        /// consolidating Unions (`consolidate_output: true`) carry non-`Direct`
        /// entries, because bucketing only pays off ahead of a consolidating
        /// downstream operator. See `LirRelationNode::Reduce::temporal_bucketing_strategy`
        /// for the underlying convention.
        temporal_bucketing_strategies: Vec<ArrangementStrategy>,
    },
    /// The `input` plan, but with additional arrangements.
    ///
    /// This operator does not change the logical contents of `input`, but ensures
    /// that certain arrangements are available in the results. This operator can
    /// be important for e.g. the `Join` stage which benefits from multiple arrangements
    /// or to cap a `LirRelationExpr` so that indexes can be exported.
    ArrangeBy {
        /// The key that must be used to access the input.
        input_key: Option<Vec<LirScalarExpr>>,
        /// The input collection.
        input: Box<LirRelationExpr>,
        /// The MFP that must be applied to the input.
        input_mfp: MfpPlan<LirScalarExpr>,
        /// A list of arrangement keys, and possibly a raw collection,
        /// that will be added to those of the input. Does not include
        /// any other existing arrangements.
        forms: AvailableCollections,
        /// How the renderer should form the arrangements requested by `forms`.
        strategy: ArrangementStrategy,
    },
}

impl LirRelationNode {
    /// Iterates through references to child expressions.
    pub fn children(&self) -> impl Iterator<Item = &LirRelationExpr> {
        let mut first = None;
        let mut second = None;
        let mut rest = None;
        let mut last = None;

        use LirRelationNode::*;
        match self {
            Constant { .. } | Get { .. } => (),
            Let { value, body, .. } => {
                first = Some(&**value);
                second = Some(&**body);
            }
            LetRec { values, body, .. } => {
                rest = Some(values);
                last = Some(&**body);
            }
            Mfp { input, .. }
            | FlatMap { input, .. }
            | Reduce { input, .. }
            | TopK { input, .. }
            | Negate { input, .. }
            | Threshold { input, .. }
            | ArrangeBy { input, .. } => {
                first = Some(&**input);
            }
            Join { inputs, .. } | Union { inputs, .. } => {
                rest = Some(inputs);
            }
        }

        first
            .into_iter()
            .chain(second)
            .chain(rest.into_iter().flatten())
            .chain(last)
    }

    /// Iterates through mutable references to child expressions.
    pub fn children_mut(&mut self) -> impl Iterator<Item = &mut LirRelationExpr> {
        let mut first = None;
        let mut second = None;
        let mut rest = None;
        let mut last = None;

        use LirRelationNode::*;
        match self {
            Constant { .. } | Get { .. } => (),
            Let { value, body, .. } => {
                first = Some(&mut **value);
                second = Some(&mut **body);
            }
            LetRec { values, body, .. } => {
                rest = Some(values);
                last = Some(&mut **body);
            }
            Mfp { input, .. }
            | FlatMap { input, .. }
            | Reduce { input, .. }
            | TopK { input, .. }
            | Negate { input, .. }
            | Threshold { input, .. }
            | ArrangeBy { input, .. } => {
                first = Some(&mut **input);
            }
            Join { inputs, .. } | Union { inputs, .. } => {
                rest = Some(inputs);
            }
        }

        first
            .into_iter()
            .chain(second)
            .chain(rest.into_iter().flatten())
            .chain(last)
    }
}

impl LirRelationNode {
    /// Attach an `lir_id` to a `LirRelationNode` to make a complete `LirRelationExpr`.
    pub fn as_plan(self, lir_id: LirId) -> LirRelationExpr {
        LirRelationExpr { lir_id, node: self }
    }
}

impl LirRelationExpr {
    /// Pretty-print this [LirRelationExpr] to a string.
    pub fn pretty(&self) -> String {
        let config = ExplainConfig::default();
        self.debug_explain(&config, None)
    }

    /// Pretty-print this [LirRelationExpr] to a string using a custom
    /// [ExplainConfig] and an optionally provided [ExprHumanizer].
    /// This is intended for debugging and tests, not users.
    pub fn debug_explain(
        &self,
        config: &ExplainConfig,
        humanizer: Option<&dyn ExprHumanizer>,
    ) -> String {
        text_string_at(self, || PlanRenderingContext {
            indent: Indent::default(),
            humanizer: humanizer.unwrap_or(&DummyHumanizer),
            annotations: BTreeMap::default(),
            config,
            ambiguous_ids: BTreeSet::default(),
        })
    }
}

/// How a `Get` stage will be rendered.
#[derive(Clone, Debug, Serialize, Deserialize, Eq, PartialEq, Ord, PartialOrd)]
pub enum GetPlan {
    /// Simply pass input arrangements on to the next stage.
    PassArrangements,
    /// Using the supplied key, optionally seek the row, and apply the MFP.
    Arrangement(
        Vec<LirScalarExpr>,
        Option<StableRow>,
        MfpPlan<LirScalarExpr>,
    ),
    /// Scan the input collection (unarranged) and apply the MFP.
    Collection(MfpPlan<LirScalarExpr>),
}

impl LirRelationExpr {
    /// Convert the dataflow description into one that uses render plans.
    #[mz_ore::instrument(
        target = "optimizer",
        level = "debug",
        fields(path.segment = "finalize_dataflow")
    )]
    pub fn finalize_dataflow(
        desc: DataflowDescription<OptimizedMirRelationExpr>,
        features: &OptimizerFeatures,
        metrics: Option<&LoweringMetrics>,
    ) -> Result<LirDataflowDescription, String> {
        fail::fail_point!("finalize_dataflow");

        // First, we lower the dataflow description from MIR to LIR. Lowering
        // also moves common parts of the MFPs pushed onto each source's reads into the source
        // itself (see `Context::refine_source_mfps`).
        let mut dataflow = Self::lower_dataflow(desc, features, metrics)?;

        // Note: `consolidate_output` for `Union` and per-input
        // `temporal_bucketing_strategies` are decided at lowering time (see the
        // `Union` arm of `lower_mir_expr_stack_safe`). The pre-existing
        // `refine_union_negate_consolidation` pass — which used to flip
        // `consolidate_output` to `true` for Unions with a `Negate` child — has
        // been folded into the lowering, since lowering is the only point where
        // the bucketing decision (which depends on `has_future_updates`) is
        // available.

        if dataflow.is_single_time() {
            // The relaxation of the `must_consolidate` flag performs an LIR-based
            // analysis and transform under checked recursion. By a similar argument
            // made in `from_mir`, we do not expect the recursion limit to be hit.
            // However, if that happens, we propagate an error to the caller.
            // To apply the transform, we first obtain monotonic source and index
            // global IDs and add them to a `TransformConfig` instance.
            let monotonic_ids = dataflow
                .source_imports
                .iter()
                .filter_map(|(id, source_import)| source_import.monotonic.then_some(*id))
                .chain(
                    dataflow
                        .index_imports
                        .iter()
                        .filter_map(|(_id, index_import)| {
                            if index_import.monotonic {
                                Some(index_import.desc.on_id)
                            } else {
                                None
                            }
                        }),
                )
                .collect::<BTreeSet<_>>();

            let config = TransformConfig { monotonic_ids };
            Self::refine_single_time_consolidation(&mut dataflow, &config)?;

            // For non-recursive delta joins in single-time dataflows, only the delta path for the
            // first relation produces updates: the other paths discard updates at the as-of, which
            // is the only time present. We keep just that path and, where the first input's
            // arrangement existed solely to seed it, drop the arrangement and consume the input as a
            // raw collection. This requires rewriting the path's initial closure to address the raw
            // row layout rather than the arranged `(key, value)` layout.
            for build_desc in dataflow.objects_to_build.iter_mut() {
                // Worklist of plan nodes. `LetRec` bodies are explored but `LetRec` values are not,
                // which excludes recursive (WMR) joins from this transform.
                let mut todo = vec![&mut build_desc.plan];
                while let Some(expr) = todo.pop() {
                    match &mut expr.node {
                        // TODO: also handle binary differential joins, which can likewise shed a
                        // bespoke arrangement on their first input.
                        LirRelationNode::Join {
                            inputs,
                            plan: JoinPlan::Delta(plan),
                        } => {
                            // Only the first relation's path survives at a single time.
                            plan.path_plans.truncate(1);

                            let source_relation = plan.path_plans[0].source_relation;
                            // Replace the source input's bespoke arrangement with a raw collection,
                            // but only when the surviving path's source is fed by an `ArrangeBy`
                            // that exists solely to build that arrangement. A source backed directly
                            // by an arranged import has no `ArrangeBy` node here, so this guard skips
                            // it and the path keeps reading it arranged.
                            if let Some(source_key) = plan.path_plans[0].source_key.clone() {
                                if let LirRelationNode::ArrangeBy { forms, .. } =
                                    &mut inputs[source_relation].node
                                {
                                    // Drop arrangement forms other than the source key, which the
                                    // remaining path no longer needs.
                                    forms.arranged.retain(|(key, _, _)| key == &source_key);
                                    if let Some((to_key, permutation, thinning)) =
                                        forms.arranged.pop()
                                    {
                                        // Make the input a raw collection and unset the source key.
                                        // Clearing every arrangement form is safe: `truncate(1)`
                                        // already dropped the sibling paths that were the only other
                                        // consumers, and this `ArrangeBy` is private to this join
                                        // input. What remains is the input's raw collection.
                                        forms.raw = true;
                                        forms.arranged.clear();
                                        plan.path_plans[0].source_key = None;

                                        // The initial closure addresses the arranged `(key, value)`
                                        // layout: columns `[0, K)` are key datums and columns
                                        // `[K, K + M)` are the thinned value datums. We rewrite it to
                                        // address the raw row instead.
                                        //
                                        // `to_key` (length `K`) are the key expressions over a row.
                                        // `permutation` (length `A`, the raw arity) maps each row
                                        // column to its position in the `(key, value)` concatenation.
                                        // `thinning` (length `M`) lists the row columns that form the
                                        // value.
                                        let key_len = to_key.len();
                                        let row_arity = permutation.len();
                                        let closure = &mut plan.path_plans[0].initial_closure;

                                        // Step 1: rewrite `ready_equivalences`, which reference the
                                        // arranged layout. A key column becomes its defining
                                        // expression. A value column becomes the row column it was
                                        // projected from.
                                        for class in closure.ready_equivalences.iter_mut() {
                                            for expr in class.iter_mut() {
                                                let mut todo = vec![expr];
                                                while let Some(expr) = todo.pop() {
                                                    if let LirScalarExpr::Column(c, _) = expr {
                                                        if let Some(key_expr) = to_key.get(*c) {
                                                            *expr = key_expr.clone();
                                                        } else {
                                                            *c = thinning[*c - key_len];
                                                        }
                                                    } else {
                                                        todo.extend(expr.children_mut());
                                                    }
                                                }
                                            }
                                        }

                                        // Step 2: rewrite the `before` MFP. Starting from a raw row,
                                        // materialize the key datums and project to the arranged
                                        // `(key, value)` layout the original MFP expects, then apply
                                        // it.
                                        let (m, f, p) = closure.before.as_map_filter_project();
                                        let mfp = MapFilterProject::new(row_arity)
                                            .map(to_key)
                                            .project(
                                                (row_arity..row_arity + key_len).chain(thinning),
                                            )
                                            .map(m)
                                            .filter(f)
                                            .project(p);
                                        closure.before =
                                            mfp.into_plan().unwrap().into_nontemporal().unwrap();
                                    }
                                }
                            }

                            todo.extend(inputs.iter_mut());
                        }
                        LirRelationNode::LetRec { body, .. } => {
                            todo.push(body);
                        }
                        x => {
                            todo.extend(x.children_mut());
                        }
                    }
                }
            }
        }

        soft_assert_eq_no_log!(dataflow.check_invariants(), Ok(()));

        mz_repr::explain::trace_plan(&dataflow);

        Ok(dataflow)
    }

    /// Lowers the dataflow description from MIR to LIR. To this end, the
    /// method collects all available arrangements and based on this information
    /// creates plans for every object to be built for the dataflow.
    #[mz_ore::instrument(
        target = "optimizer",
        level = "debug",
        fields(path.segment ="mir_to_lir")
    )]
    fn lower_dataflow(
        desc: DataflowDescription<OptimizedMirRelationExpr>,
        features: &OptimizerFeatures,
        metrics: Option<&LoweringMetrics>,
    ) -> Result<LirDataflowDescription, String> {
        let context = lowering::Context::new(desc.debug_name.clone(), features, metrics);
        let dataflow = context.lower(desc)?;

        mz_repr::explain::trace_plan(&dataflow);

        Ok(dataflow)
    }

    /// Refines the plans of objects to be built as part of a single-time `dataflow` to relax
    /// the setting of the `must_consolidate` attribute of monotonic operators, if necessary,
    /// whenever the input is deemed to be physically monotonic.
    #[mz_ore::instrument(
        target = "optimizer",
        level = "debug",
        fields(path.segment = "refine_single_time_consolidation")
    )]
    fn refine_single_time_consolidation(
        dataflow: &mut LirDataflowDescription,
        config: &TransformConfig,
    ) -> Result<(), String> {
        // We should only reach here if we have a one-shot SELECT query, i.e.,
        // a single-time dataflow.
        assert!(dataflow.is_single_time());

        let transform = transform::RelaxMustConsolidate;
        for build_desc in dataflow.objects_to_build.iter_mut() {
            transform
                .transform(config, &mut build_desc.plan)
                .map_err(|_| "Maximum recursion limit error in consolidation relaxation.")?;
        }
        mz_repr::explain::trace_plan(dataflow);
        Ok(())
    }
}

impl CollectionPlan for LirRelationNode {
    fn depends_on_into(&self, out: &mut BTreeSet<GlobalId>) {
        match self {
            LirRelationNode::Constant { rows: _ } => (),
            LirRelationNode::Get {
                id,
                keys: _,
                plan: _,
            } => match id {
                Id::Global(id) => {
                    out.insert(*id);
                }
                Id::Local(_) => (),
            },
            LirRelationNode::Let { id: _, value, body } => {
                value.depends_on_into(out);
                body.depends_on_into(out);
            }
            LirRelationNode::LetRec {
                ids: _,
                values,
                limits: _,
                body,
            } => {
                for value in values.iter() {
                    value.depends_on_into(out);
                }
                body.depends_on_into(out);
            }
            LirRelationNode::Join { inputs, plan: _ }
            | LirRelationNode::Union {
                inputs,
                consolidate_output: _,
                temporal_bucketing_strategies: _,
            } => {
                for input in inputs {
                    input.depends_on_into(out);
                }
            }
            LirRelationNode::Mfp {
                input,
                mfp: _,
                input_key_val: _,
            }
            | LirRelationNode::FlatMap {
                input_key: _,
                input,
                exprs: _,
                func: _,
                mfp_after: _,
            }
            | LirRelationNode::ArrangeBy {
                input_key: _,
                input,
                input_mfp: _,
                forms: _,
                strategy: _,
            }
            | LirRelationNode::Reduce {
                input_key: _,
                input,
                key_val_plan: _,
                plan: _,
                mfp_after: _,
                temporal_bucketing_strategy: _,
            }
            | LirRelationNode::TopK {
                input,
                top_k_plan: _,
                temporal_bucketing_strategy: _,
            }
            | LirRelationNode::Negate { input }
            | LirRelationNode::Threshold {
                input,
                threshold_plan: _,
            } => {
                input.depends_on_into(out);
            }
        }
    }
}

impl CollectionPlan for LirRelationExpr {
    fn depends_on_into(&self, out: &mut BTreeSet<GlobalId>) {
        self.node.depends_on_into(out);
    }
}

/// Returns bucket sizes, descending, suitable for hierarchical decomposition of an operator, based
/// on the expected number of rows that will have the same group key.
fn bucketing_of_expected_group_size(expected_group_size: Option<u64>) -> Vec<u64> {
    // NOTE(vmarcos): The fan-in of 16 defined below is used in the tuning advice built-in view
    // mz_introspection.mz_expected_group_size_advice.
    let mut buckets = vec![];
    let mut current = 16;

    // Plan for 4B records in the expected case if the user didn't specify a group size.
    let limit = expected_group_size.unwrap_or(4_000_000_000);

    // Distribute buckets in powers of 16, so that we can strike a balance between how many inputs
    // each layer gets from the preceding layer, while also limiting the number of layers.
    while current < limit {
        buckets.push(current);
        current = current.saturating_mul(16);
    }

    buckets.reverse();
    buckets
}