mz_compute_types/

dataflows.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.

//! Types for describing dataflows.

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

use mz_expr::{
    CollectionPlan, MapFilterProject, MfpPlan, MirRelationExpr, MirScalarExpr,
    OptimizedMirRelationExpr,
};
use mz_ore::collections::CollectionExt;
use mz_ore::soft_assert_or_log;
use mz_repr::refresh_schedule::RefreshSchedule;
use mz_repr::{GlobalId, ReprRelationType, SqlRelationType, Timestamp};
use mz_storage_types::time_dependence::TimeDependence;
use serde::{Deserialize, Serialize};
use timely::progress::Antichain;

use crate::plan::LirRelationExpr;
use crate::plan::render_plan::RenderPlan;
use crate::plan::scalar::{LirScalarExpr, lses_from_mses};
use crate::sinks::{ComputeSinkConnection, ComputeSinkDesc};
use crate::sources::{SourceInstanceArguments, SourceInstanceDesc};

/// A description of a dataflow to construct and results to surface.
#[derive(Clone, Debug, Serialize, Deserialize, Eq, PartialEq)]
pub struct DataflowDescription<P, S: 'static = (), O = MapFilterProject> {
    /// Sources instantiations made available to the dataflow pair with monotonicity information.
    pub source_imports: BTreeMap<GlobalId, SourceImport<S, O>>,
    /// Indexes made available to the dataflow.
    /// (id of index, import)
    pub index_imports: BTreeMap<GlobalId, IndexImport>,
    /// Views and indexes to be built and stored in the local context.
    /// Objects must be built in the specific order, as there may be
    /// dependencies of later objects on prior identifiers.
    pub objects_to_build: Vec<BuildDesc<P>>,
    /// Indexes to be made available to be shared with other dataflows
    /// (id of new index, description of index, relationtype of base source/view/table)
    pub index_exports: BTreeMap<GlobalId, (IndexDesc<MirScalarExpr>, ReprRelationType)>,
    /// sinks to be created
    /// (id of new sink, description of sink)
    pub sink_exports: BTreeMap<GlobalId, ComputeSinkDesc<S>>,
    /// An optional frontier to which inputs should be advanced.
    ///
    /// If this is set, it should override the default setting determined by
    /// the upper bound of `since` frontiers contributing to the dataflow.
    /// It is an error for this to be set to a frontier not beyond that default.
    pub as_of: Option<Antichain<Timestamp>>,
    /// Frontier beyond which the dataflow should not execute.
    /// Specifically, updates at times greater or equal to this frontier are suppressed.
    /// This is often set to `as_of + 1` to enable "batch" computations.
    /// Note that frontier advancements might still happen to times that are after the `until`,
    /// only data is suppressed. (This is consistent with how frontier advancements can also
    /// happen before the `as_of`.)
    pub until: Antichain<Timestamp>,
    /// The initial as_of when the collection is first created. Filled only for materialized views.
    /// Note that this doesn't change upon restarts.
    pub initial_storage_as_of: Option<Antichain<Timestamp>>,
    /// The schedule of REFRESH materialized views.
    pub refresh_schedule: Option<RefreshSchedule>,
    /// Human-readable name
    pub debug_name: String,
    /// Description of how the dataflow's progress relates to wall-clock time. None for unknown.
    pub time_dependence: Option<TimeDependence>,
}

impl<P, S, O> DataflowDescription<P, S, O> {
    /// Tests if the dataflow refers to a single timestamp, namely
    /// that `as_of` has a single coordinate and that the `until`
    /// value corresponds to the `as_of` value plus one, or `as_of`
    /// is the maximum timestamp and is thus single.
    pub fn is_single_time(&self) -> bool {
        // TODO: this would be much easier to check if `until` was a strict lower bound,
        // and we would be testing that `until == as_of`.

        let until = &self.until;

        // IF `as_of` is not set at all this can't be a single time dataflow.
        let Some(as_of) = self.as_of.as_ref() else {
            return false;
        };
        // Ensure that as_of <= until.
        soft_assert_or_log!(
            timely::PartialOrder::less_equal(as_of, until),
            "expected `as_of ≤ until`, got `{as_of:?} ≰ {until:?}`",
        );
        // IF `as_of` is not a single timestamp this can't be a single time dataflow.
        let Some(as_of) = as_of.as_option() else {
            return false;
        };
        // With `as_of = MAX` the `as_of <= until` invariant leaves only two valid
        // frontiers, and both are handled correctly by the comparison below:
        //   * `until = {}` (unbounded) yields `None == None`, so single time. `MAX`
        //     is the last timestamp, so an unbounded dataflow starting there sees
        //     exactly one time.
        //   * `until = {MAX}` is the empty range produced by e.g.
        //     `SUBSCRIBE ... AS OF MAX UP TO MAX`. `try_step_forward` on `MAX`
        //     overflows to `None`, so `None == Some(MAX)` is false: not single time.
        as_of.try_step_forward().as_ref() == until.as_option()
    }
}

impl LirDataflowDescription {
    /// Check invariants expected to be true about `DataflowDescription`s.
    pub fn check_invariants(&self) -> Result<(), String> {
        let mut plans: Vec<_> = self.objects_to_build.iter().map(|o| &o.plan).collect();
        let mut lir_ids = BTreeSet::new();

        while let Some(plan) = plans.pop() {
            let lir_id = plan.lir_id;
            if !lir_ids.insert(lir_id) {
                return Err(format!(
                    "duplicate `LirId` in `DataflowDescription`: {lir_id}"
                ));
            }
            plans.extend(plan.node.children());
        }

        Ok(())
    }
}

impl DataflowDescription<OptimizedMirRelationExpr, ()> {
    /// Imports a previously exported index.
    ///
    /// This method makes available an index previously exported as `id`, identified
    /// to the query by `description` (which names the view the index arranges, and
    /// the keys by which it is arranged).
    pub fn import_index(
        &mut self,
        id: GlobalId,
        desc: IndexDesc<MirScalarExpr>,
        typ: ReprRelationType,
        monotonic: bool,
    ) {
        self.index_imports.insert(
            id,
            IndexImport {
                desc,
                typ,
                monotonic,
                with_snapshot: true,
            },
        );
    }

    /// Imports a source and makes it available as `id`.
    pub fn import_source(&mut self, id: GlobalId, typ: SqlRelationType, monotonic: bool) {
        // Import the source with no linear operators applied to it.
        // They may be populated by whole-dataflow optimization.
        // Similarly, we require the snapshot by default, though optimization may choose to skip it.
        self.source_imports.insert(
            id,
            SourceImport {
                desc: SourceInstanceDesc {
                    storage_metadata: (),
                    arguments: SourceInstanceArguments { operators: None },
                    typ,
                },
                monotonic,
                with_snapshot: true,
                upper: Antichain::new(),
            },
        );
    }

    /// Binds to `id` the relation expression `plan`.
    pub fn insert_plan(&mut self, id: GlobalId, plan: OptimizedMirRelationExpr) {
        self.objects_to_build.push(BuildDesc { id, plan });
    }

    /// Exports as `id` an index described by `description`.
    ///
    /// Future uses of `import_index` in other dataflow descriptions may use `id`,
    /// as long as this dataflow has not been terminated in the meantime.
    pub fn export_index(
        &mut self,
        id: GlobalId,
        description: IndexDesc<MirScalarExpr>,
        on_type: ReprRelationType,
    ) {
        // We first create a "view" named `id` that ensures that the
        // data are correctly arranged and available for export.
        self.insert_plan(
            id,
            OptimizedMirRelationExpr::declare_optimized(MirRelationExpr::ArrangeBy {
                input: Box::new(MirRelationExpr::global_get(
                    description.on_id,
                    on_type.clone(),
                )),
                keys: vec![description.key.clone()],
            }),
        );
        self.index_exports.insert(id, (description, on_type));
    }

    /// Exports as `id` a sink described by `description`.
    pub fn export_sink(&mut self, id: GlobalId, description: ComputeSinkDesc<()>) {
        self.sink_exports.insert(id, description);
    }

    /// Returns true iff `id` is already imported.
    pub fn is_imported(&self, id: &GlobalId) -> bool {
        self.objects_to_build.iter().any(|bd| &bd.id == id)
            || self.index_imports.keys().any(|i| i == id)
            || self.source_imports.keys().any(|i| i == id)
    }

    /// The number of columns associated with an identifier in the dataflow.
    pub fn arity_of(&self, id: &GlobalId) -> usize {
        for (source_id, source_import) in self.source_imports.iter() {
            let source = &source_import.desc;
            if source_id == id {
                return source.typ.arity();
            }
        }
        for IndexImport { desc, typ, .. } in self.index_imports.values() {
            if &desc.on_id == id {
                return typ.arity();
            }
        }
        for desc in self.objects_to_build.iter() {
            if &desc.id == id {
                return desc.plan.arity();
            }
        }
        panic!("GlobalId {} not found in MirDataflowDescription", id);
    }

    /// Calls r and s on any sub-members of those types in self. Halts at the first error return.
    pub fn visit_children<R, S, E>(&mut self, r: R, s: S) -> Result<(), E>
    where
        R: Fn(&mut OptimizedMirRelationExpr) -> Result<(), E>,
        S: Fn(&mut MirScalarExpr) -> Result<(), E>,
    {
        for BuildDesc { plan, .. } in &mut self.objects_to_build {
            r(plan)?;
        }
        for source_import in self.source_imports.values_mut() {
            let Some(mfp) = source_import.desc.arguments.operators.as_mut() else {
                continue;
            };
            for expr in mfp.expressions.iter_mut() {
                s(expr)?;
            }
            for (_, expr) in mfp.predicates.iter_mut() {
                s(expr)?;
            }
        }
        Ok(())
    }
}

impl<P, S, O> DataflowDescription<P, S, O> {
    /// Applies `f` to the operators of every source import, failing on the
    /// first error.
    pub fn try_map_source_operators<O2, E>(
        self,
        mut f: impl FnMut(O) -> Result<O2, E>,
    ) -> Result<DataflowDescription<P, S, O2>, E> {
        let DataflowDescription {
            source_imports,
            index_imports,
            objects_to_build,
            index_exports,
            sink_exports,
            as_of,
            until,
            initial_storage_as_of,
            refresh_schedule,
            debug_name,
            time_dependence,
        } = self;
        let source_imports = source_imports
            .into_iter()
            .map(|(id, import)| {
                let SourceImport {
                    desc:
                        SourceInstanceDesc {
                            arguments: SourceInstanceArguments { operators },
                            storage_metadata,
                            typ,
                        },
                    monotonic,
                    with_snapshot,
                    upper,
                } = import;
                let desc = SourceInstanceDesc {
                    arguments: SourceInstanceArguments {
                        operators: operators.map(&mut f).transpose()?,
                    },
                    storage_metadata,
                    typ,
                };
                let import = SourceImport {
                    desc,
                    monotonic,
                    with_snapshot,
                    upper,
                };
                Ok((id, import))
            })
            .collect::<Result<_, E>>()?;
        Ok(DataflowDescription {
            source_imports,
            index_imports,
            objects_to_build,
            index_exports,
            sink_exports,
            as_of,
            until,
            initial_storage_as_of,
            refresh_schedule,
            debug_name,
            time_dependence,
        })
    }

    /// Creates a new dataflow description with a human-readable name.
    pub fn new(name: String) -> Self {
        Self {
            source_imports: Default::default(),
            index_imports: Default::default(),
            objects_to_build: Vec::new(),
            index_exports: Default::default(),
            sink_exports: Default::default(),
            as_of: Default::default(),
            until: Antichain::new(),
            initial_storage_as_of: None,
            refresh_schedule: None,
            debug_name: name,
            time_dependence: None,
        }
    }

    /// Sets the `as_of` frontier to the supplied argument.
    ///
    /// This method allows the dataflow to indicate a frontier up through
    /// which all times should be advanced. This can be done for at least
    /// two reasons: 1. correctness and 2. performance.
    ///
    /// Correctness may require an `as_of` to ensure that historical detail
    /// is consolidated at representative times that do not present specific
    /// detail that is not specifically correct. For example, updates may be
    /// compacted to times that are no longer the source times, but instead
    /// some byproduct of when compaction was executed; we should not present
    /// those specific times as meaningfully different from other equivalent
    /// times.
    ///
    /// Performance may benefit from an aggressive `as_of` as it reduces the
    /// number of distinct moments at which collections vary. Differential
    /// dataflow will refresh its outputs at each time its inputs change and
    /// to moderate that we can minimize the volume of distinct input times
    /// as much as possible.
    ///
    /// Generally, one should consider setting `as_of` at least to the `since`
    /// frontiers of contributing data sources and as aggressively as the
    /// computation permits.
    pub fn set_as_of(&mut self, as_of: Antichain<Timestamp>) {
        self.as_of = Some(as_of);
    }

    /// Records the initial `as_of` of the storage collection associated with a materialized view.
    pub fn set_initial_as_of(&mut self, initial_as_of: Antichain<Timestamp>) {
        self.initial_storage_as_of = Some(initial_as_of);
    }

    /// Identifiers of imported objects (indexes and sources).
    pub fn import_ids(&self) -> impl Iterator<Item = GlobalId> + Clone + '_ {
        self.imported_index_ids().chain(self.imported_source_ids())
    }

    /// Identifiers of imported indexes.
    pub fn imported_index_ids(&self) -> impl Iterator<Item = GlobalId> + Clone + '_ {
        self.index_imports.keys().copied()
    }

    /// Identifiers of imported sources.
    pub fn imported_source_ids(&self) -> impl Iterator<Item = GlobalId> + Clone + '_ {
        self.source_imports.keys().copied()
    }

    /// Whether `id` names an import of this dataflow, index or source.
    pub fn is_import(&self, id: &GlobalId) -> bool {
        self.index_imports.contains_key(id) || self.source_imports.contains_key(id)
    }

    /// Identifiers of exported objects (indexes and sinks).
    pub fn export_ids(&self) -> impl Iterator<Item = GlobalId> + Clone + '_ {
        self.exported_index_ids().chain(self.exported_sink_ids())
    }

    /// Identifiers of exported indexes.
    pub fn exported_index_ids(&self) -> impl Iterator<Item = GlobalId> + Clone + '_ {
        self.index_exports.keys().copied()
    }

    /// Identifiers of exported sinks.
    pub fn exported_sink_ids(&self) -> impl Iterator<Item = GlobalId> + Clone + '_ {
        self.sink_exports.keys().copied()
    }

    /// Identifiers of exported persist sinks.
    pub fn persist_sink_ids(&self) -> impl Iterator<Item = GlobalId> + '_ {
        self.sink_exports
            .iter()
            .filter_map(|(id, desc)| match desc.connection {
                ComputeSinkConnection::MaterializedView(_) => Some(*id),
                _ => None,
            })
    }

    /// Identifiers of exported subscribe sinks.
    pub fn subscribe_ids(&self) -> impl Iterator<Item = GlobalId> + '_ {
        self.sink_exports
            .iter()
            .filter_map(|(id, desc)| match desc.connection {
                ComputeSinkConnection::Subscribe(_) => Some(*id),
                _ => None,
            })
    }

    /// Identifiers of exported copy to sinks.
    pub fn copy_to_ids(&self) -> impl Iterator<Item = GlobalId> + '_ {
        self.sink_exports
            .iter()
            .filter_map(|(id, desc)| match desc.connection {
                ComputeSinkConnection::CopyToS3Oneshot(_) => Some(*id),
                _ => None,
            })
    }

    /// Produce a `Display`able value containing the import IDs of this dataflow.
    pub fn display_import_ids(&self) -> impl fmt::Display + '_ {
        use mz_ore::str::{bracketed, separated};
        bracketed("[", "]", separated(", ", self.import_ids()))
    }

    /// Produce a `Display`able value containing the export IDs of this dataflow.
    pub fn display_export_ids(&self) -> impl fmt::Display + '_ {
        use mz_ore::str::{bracketed, separated};
        bracketed("[", "]", separated(", ", self.export_ids()))
    }

    /// Whether this dataflow installs transient collections.
    pub fn is_transient(&self) -> bool {
        self.export_ids().all(|id| id.is_transient())
    }

    /// Returns the description of the object to build with the specified
    /// identifier.
    ///
    /// # Panics
    ///
    /// Panics if `id` is not present in `objects_to_build` exactly once.
    pub fn build_desc(&self, id: GlobalId) -> &BuildDesc<P> {
        let mut builds = self.objects_to_build.iter().filter(|build| build.id == id);
        let build = builds
            .next()
            .unwrap_or_else(|| panic!("object to build id {id} unexpectedly missing"));
        assert!(builds.next().is_none());
        build
    }

    /// Returns the id of the dataflow's sink export.
    ///
    /// # Panics
    ///
    /// Panics if the dataflow has no sink exports or has more than one.
    pub fn sink_id(&self) -> GlobalId {
        let sink_exports = &self.sink_exports;
        let sink_id = sink_exports.keys().into_element();
        *sink_id
    }
}

impl<P, S, O> DataflowDescription<P, S, O>
where
    P: CollectionPlan,
{
    /// Computes the set of identifiers upon which the specified collection
    /// identifier depends.
    ///
    /// `collection_id` must specify a valid object in `objects_to_build`.
    ///
    /// This method includes identifiers for e.g. intermediate views, and should be filtered
    /// if one only wants sources and indexes.
    ///
    /// This method is safe for mutually recursive view definitions.
    pub fn depends_on(&self, collection_id: GlobalId) -> BTreeSet<GlobalId> {
        let mut out = BTreeSet::new();
        self.depends_on_into(collection_id, &mut out);
        out
    }

    /// Like `depends_on`, but appends to an existing `BTreeSet`.
    pub fn depends_on_into(&self, collection_id: GlobalId, out: &mut BTreeSet<GlobalId>) {
        out.insert(collection_id);
        if self.source_imports.contains_key(&collection_id) {
            // The collection is provided by an imported source. Report the
            // dependency on the source.
            out.insert(collection_id);
            return;
        }

        // NOTE(benesch): we're not smart enough here to know *which* index
        // for the collection will be used, if one exists, so we have to report
        // the dependency on all of them.
        let mut found_index = false;
        for (index_id, IndexImport { desc, .. }) in &self.index_imports {
            if desc.on_id == collection_id {
                // The collection is provided by an imported index. Report the
                // dependency on the index.
                out.insert(*index_id);
                found_index = true;
            }
        }
        if found_index {
            return;
        }

        // The collection is not provided by a source or imported index.
        // It must be a collection whose plan we have handy. Recurse.
        let build = self.build_desc(collection_id);
        for id in build.plan.depends_on() {
            if !out.contains(&id) {
                self.depends_on_into(id, out)
            }
        }
    }

    /// Computes the set of imports upon which the specified collection depends.
    ///
    /// This method behaves like `depends_on` but filters out internal dependencies that are not
    /// included in the dataflow imports.
    pub fn depends_on_imports(&self, collection_id: GlobalId) -> BTreeSet<GlobalId> {
        let deps = self.depends_on(collection_id);
        deps.into_iter().filter(|id| self.is_import(id)).collect()
    }

    /// Computes the set of imports the dataflow's exports read, meaning the imports reachable from
    /// an index export's `on_id` or a sink export's `from`.
    ///
    /// A description that has been through the optimizer answers [`Self::import_ids`] here, because
    /// the optimizer prunes the import list to what the exports read. The two come apart while a
    /// description is still being assembled, and this is what the prune and the assertion guarding
    /// it are both defined in terms of. The answer covers the dataflow as a whole, so an import only
    /// one export reads is still reported, and a dataflow with no exports reports none.
    ///
    /// NOTE: On the index side this over-approximates. [`Self::depends_on`] cannot tell which index
    /// on a collection a plan will use, so reaching a collection reports every index imported on it.
    /// Pruning index imports needs the exact usage information the MIR pipeline collects, not this.
    ///
    /// Panics for an export naming a collection that is neither an import nor built exactly once
    /// here, which is [`Self::depends_on`]'s precondition on its argument. Rendering resolves the
    /// same ids, so a description that trips this does not survive being built either.
    pub fn used_import_ids(&self) -> BTreeSet<GlobalId> {
        let mut deps = BTreeSet::new();
        for (index_desc, _typ) in self.index_exports.values() {
            self.depends_on_into(index_desc.on_id, &mut deps);
        }
        for sink_desc in self.sink_exports.values() {
            self.depends_on_into(sink_desc.from, &mut deps);
        }
        deps.retain(|id| self.is_import(id));
        deps
    }
}

impl<S> RenderDataflowDescription<S>
where
    S: Clone + PartialEq,
{
    /// Determine if a dataflow description is compatible with this dataflow description.
    ///
    /// Compatible dataflows have structurally equal exports, imports, and objects to build. The
    /// `as_of` of the receiver has to be less equal the `other` `as_of`.
    ///
    /// Note that this method performs normalization as part of the structural equality checking,
    /// which involves cloning both `self` and `other`. It is therefore relatively expensive and
    /// should only be used on cold code paths.
    ///
    // TODO: The semantics of this function are only useful for command reconciliation at the moment.
    pub fn compatible_with(&self, other: &Self) -> bool {
        let old = self.as_comparable();
        let new = other.as_comparable();

        let equality = old.index_exports == new.index_exports
            && old.sink_exports == new.sink_exports
            && old.objects_to_build == new.objects_to_build
            && old.index_imports == new.index_imports
            && old.source_imports == new.source_imports
            && old.time_dependence == new.time_dependence;

        let partial = if let (Some(old_as_of), Some(new_as_of)) = (&old.as_of, &new.as_of) {
            timely::PartialOrder::less_equal(old_as_of, new_as_of)
        } else {
            false
        };

        equality && partial
    }

    /// Returns a `DataflowDescription` that has the same structure as `self` and can be
    /// structurally compared to other `DataflowDescription`s.
    ///
    /// The function normalizes several properties. It replaces transient `GlobalId`s
    /// that are only used internally (i.e. not imported nor exported) with consecutive IDs
    /// starting from `t1`. It replaces the source import's `upper` by a dummy value.
    fn as_comparable(&self) -> Self {
        let external_ids: BTreeSet<_> = self.import_ids().chain(self.export_ids()).collect();

        let mut id_counter = 0;
        let mut replacements = BTreeMap::new();

        let mut maybe_replace = |id: GlobalId| {
            if id.is_transient() && !external_ids.contains(&id) {
                *replacements.entry(id).or_insert_with(|| {
                    id_counter += 1;
                    GlobalId::Transient(id_counter)
                })
            } else {
                id
            }
        };

        let mut source_imports = self.source_imports.clone();
        for import in source_imports.values_mut() {
            import.upper = Antichain::new();
        }

        let mut objects_to_build = self.objects_to_build.clone();
        for object in &mut objects_to_build {
            object.id = maybe_replace(object.id);
            object.plan.replace_ids(&mut maybe_replace);
        }

        let mut index_exports = self.index_exports.clone();
        for (desc, _typ) in index_exports.values_mut() {
            desc.on_id = maybe_replace(desc.on_id);
        }

        let mut sink_exports = self.sink_exports.clone();
        for desc in sink_exports.values_mut() {
            desc.from = maybe_replace(desc.from);
        }

        DataflowDescription {
            source_imports,
            index_imports: self.index_imports.clone(),
            objects_to_build,
            index_exports,
            sink_exports,
            as_of: self.as_of.clone(),
            until: self.until.clone(),
            initial_storage_as_of: self.initial_storage_as_of.clone(),
            refresh_schedule: self.refresh_schedule.clone(),
            debug_name: self.debug_name.clone(),
            time_dependence: self.time_dependence.clone(),
        }
    }
}

/// A dataflow of optimized MIR, as the optimizer builds and rewrites it.
pub type MirDataflowDescription = DataflowDescription<OptimizedMirRelationExpr>;

/// A dataflow lowered to LIR.
///
/// Its source operators are already planned. Rendering evaluates them as
/// given, so lowering is the only place they are planned.
pub type LirDataflowDescription<S = ()> =
    DataflowDescription<LirRelationExpr, S, MfpPlan<LirScalarExpr>>;

/// A dataflow ready to render.
pub type RenderDataflowDescription<S> = DataflowDescription<RenderPlan, S, MfpPlan<LirScalarExpr>>;

/// An index storing processed updates so they can be queried
/// or reused in other computations
#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize, Hash)]
pub struct IndexDesc<E> {
    /// Identity of the collection the index is on.
    pub on_id: GlobalId,
    /// Expressions to be arranged, in order of decreasing primacy.
    pub key: Vec<E>,
}

impl IndexDesc<MirScalarExpr> {
    /// Translate an index description from MIR to LIR.
    pub fn as_lir(&self) -> IndexDesc<LirScalarExpr> {
        let on_id = self.on_id.clone();
        let key = lses_from_mses(&self.key);

        IndexDesc { on_id, key }
    }
}

/// Information about an imported index, and how it will be used by the dataflow.
#[derive(Clone, Debug, Serialize, Deserialize, Eq, PartialEq)]
pub struct IndexImport {
    /// Description of index.
    pub desc: IndexDesc<MirScalarExpr>,
    /// Schema and keys of the object the index is on.
    pub typ: ReprRelationType,
    /// Whether the index will supply monotonic data.
    pub monotonic: bool,
    /// Whether this import must include the snapshot data.
    pub with_snapshot: bool,
}

/// Information about an imported source, and how it will be used by the dataflow.
#[derive(Clone, Debug, Serialize, Deserialize, Eq, PartialEq)]
pub struct SourceImport<S: 'static = (), O = MapFilterProject> {
    /// Description of the source instance to import.
    pub desc: SourceInstanceDesc<S, O>,
    /// Whether the source will supply monotonic data.
    pub monotonic: bool,
    /// Whether this import must include the snapshot data.
    pub with_snapshot: bool,
    /// The initial known upper frontier for the source.
    pub upper: Antichain<Timestamp>,
}

/// An association of a global identifier to an expression.
#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub struct BuildDesc<P> {
    /// TODO(database-issues#7533): Add documentation.
    pub id: GlobalId,
    /// TODO(database-issues#7533): Add documentation.
    pub plan: P,
}

#[cfg(test)]
mod tests {
    use mz_expr::{AccessStrategy, Id, MirRelationExpr, UnmaterializableFunc, func};
    use mz_repr::optimize::OptimizerFeatures;
    use mz_repr::{RelationDesc, ReprRelationType, ReprScalarType, SqlRelationType};

    use crate::plan::scalar::mfp_plan_mir_to_lir;
    use crate::sinks::{ComputeSinkConnection, ComputeSinkDesc, SubscribeSinkConnection};
    use crate::sources::{SourceInstanceArguments, SourceInstanceDesc};

    use super::*;

    const READ: GlobalId = GlobalId::User(1);
    const UNREAD: GlobalId = GlobalId::User(2);
    const VIEW: GlobalId = GlobalId::Transient(1);
    const SINK: GlobalId = GlobalId::Transient(2);
    const INDEX: GlobalId = GlobalId::Transient(3);

    fn typ() -> ReprRelationType {
        ReprRelationType::new(vec![ReprScalarType::Int64.nullable(false)])
    }

    /// A dataflow importing `READ` and `UNREAD`, building `VIEW` from `plan`, and exporting a
    /// subscribe sink over it.
    fn dataflow(plan: MirRelationExpr) -> MirDataflowDescription {
        let source_import = || SourceImport {
            desc: SourceInstanceDesc {
                arguments: SourceInstanceArguments { operators: None },
                storage_metadata: (),
                typ: SqlRelationType::from_repr(&typ()),
            },
            monotonic: false,
            with_snapshot: true,
            upper: Antichain::from_elem(Timestamp::MIN),
        };

        let mut df = MirDataflowDescription::new("test".to_string());
        df.source_imports.insert(READ, source_import());
        df.source_imports.insert(UNREAD, source_import());
        df.objects_to_build.push(BuildDesc {
            id: VIEW,
            plan: OptimizedMirRelationExpr::declare_optimized(plan),
        });
        df.sink_exports.insert(
            SINK,
            ComputeSinkDesc {
                from: VIEW,
                from_desc: RelationDesc::new(SqlRelationType::from_repr(&typ()), ["c"]),
                connection: ComputeSinkConnection::Subscribe(SubscribeSinkConnection {
                    output: Vec::new(),
                }),
                with_snapshot: true,
                up_to: Antichain::new(),
                non_null_assertions: Vec::new(),
                refresh_schedule: None,
                from_arrangement: None,
            },
        );
        df
    }

    #[mz_ore::test]
    fn used_import_ids_reports_only_read_imports() {
        let df = dataflow(MirRelationExpr::Get {
            id: Id::Global(READ),
            typ: typ(),
            access_strategy: AccessStrategy::Persist,
        });

        assert_eq!(df.used_import_ids(), BTreeSet::from([READ]));
    }

    /// An export the optimizer folded to a constant reads nothing, even though the imports it was
    /// folded from are still there. This is the shape that must not be reported as reading them.
    #[mz_ore::test]
    fn used_import_ids_is_empty_for_a_constant_export() {
        let df = dataflow(MirRelationExpr::Constant {
            rows: Ok(Vec::new()),
            typ: typ(),
        });

        assert_eq!(df.used_import_ids(), BTreeSet::new());
    }

    /// Index exports are walked from the collection they are on, just like sink exports are from
    /// the collection they are from.
    #[mz_ore::test]
    fn used_import_ids_covers_index_exports() {
        let mut df = dataflow(MirRelationExpr::Constant {
            rows: Ok(Vec::new()),
            typ: typ(),
        });
        let other_view = GlobalId::Transient(4);
        df.objects_to_build.push(BuildDesc {
            id: other_view,
            plan: OptimizedMirRelationExpr::declare_optimized(MirRelationExpr::Get {
                id: Id::Global(UNREAD),
                typ: typ(),
                access_strategy: AccessStrategy::Persist,
            }),
        });
        df.index_exports.insert(
            INDEX,
            (
                IndexDesc {
                    on_id: other_view,
                    key: Vec::new(),
                },
                typ(),
            ),
        );

        assert_eq!(df.used_import_ids(), BTreeSet::from([UNREAD]));
    }

    /// An imported index is reached through the collection it is on, and reported by the id of the
    /// index rather than that of the collection.
    #[mz_ore::test]
    fn used_import_ids_reports_imported_indexes_by_index_id() {
        let indexed_view = GlobalId::User(3);
        let imported_index = GlobalId::User(4);

        let mut df = dataflow(MirRelationExpr::Get {
            id: Id::Global(indexed_view),
            typ: typ(),
            access_strategy: AccessStrategy::Index(Vec::new()),
        });
        df.index_imports.insert(
            imported_index,
            IndexImport {
                desc: IndexDesc {
                    on_id: indexed_view,
                    key: Vec::new(),
                },
                typ: typ(),
                monotonic: false,
                with_snapshot: true,
            },
        );

        assert_eq!(df.used_import_ids(), BTreeSet::from([imported_index]));
    }

    /// Builds a minimal dataflow description carrying only the `as_of` and
    /// `until` frontiers that `is_single_time` inspects.
    fn desc_with_bounds(
        as_of: Antichain<Timestamp>,
        until: Antichain<Timestamp>,
    ) -> DataflowDescription<OptimizedMirRelationExpr, ()> {
        let mut desc = DataflowDescription::new("test".to_string());
        desc.set_as_of(as_of);
        desc.until = until;
        desc
    }

    #[mz_ore::test]
    fn is_single_time_classification() {
        // Ordinary single-time dataflow: `until = as_of + 1`.
        assert!(
            desc_with_bounds(
                Antichain::from_elem(5.into()),
                Antichain::from_elem(6.into())
            )
            .is_single_time()
        );

        // `as_of = MAX` with an empty `until` is single time: `MAX` is the last
        // timestamp, so an unbounded dataflow starting there sees exactly one time.
        assert!(
            desc_with_bounds(Antichain::from_elem(Timestamp::MAX), Antichain::new())
                .is_single_time()
        );

        // `as_of = MAX` with `until = {MAX}` is an empty range, produced by
        // `SUBSCRIBE ... AS OF MAX UP TO MAX`. It must not be classified as single
        // time, and it must not trip an invariant assertion (see CPU-169).
        assert!(
            !desc_with_bounds(
                Antichain::from_elem(Timestamp::MAX),
                Antichain::from_elem(Timestamp::MAX)
            )
            .is_single_time()
        );
    }

    #[mz_ore::test]
    fn lowering_plans_source_operators() {
        // `#1 = #0::mz_timestamp`, filtered by `mz_now() < #1` and
        // `mz_now() = #0::mz_timestamp`.
        let mz_now = || MirScalarExpr::CallUnmaterializable(UnmaterializableFunc::MzNow);
        let cast = || MirScalarExpr::column(0).call_unary(func::CastInt64ToMzTimestamp);
        let mfp = MapFilterProject::new(1).map([cast()]).filter([
            mz_now().call_binary(MirScalarExpr::column(1), func::Lt),
            mz_now().call_binary(cast(), func::Eq),
        ]);
        let planned = MfpPlan::create_from(mfp.clone()).expect("plannable");

        let mut df = dataflow(MirRelationExpr::Get {
            id: Id::Global(READ),
            typ: typ(),
            access_strategy: AccessStrategy::Persist,
        });
        let unread = df.source_imports.get_mut(&UNREAD).expect("imported");
        unread.desc.arguments.operators = Some(mfp);

        let lowered = LirRelationExpr::finalize_dataflow(df, &OptimizerFeatures::default(), None)
            .expect("lowers");
        let operators = &lowered.source_imports[&UNREAD].desc.arguments.operators;
        assert_eq!(operators.as_ref(), Some(&mfp_plan_mir_to_lir(planned)));
    }

    /// MIR predicate pushdown copies `mz_now() <= #0` into the source's operators and keeps it in
    /// the `Filter` above the read. Hoisting the read's copy into the source must not leave the
    /// source evaluating the same bound twice.
    #[mz_ore::test]
    fn lowering_merges_pushed_down_temporal_predicate() {
        let mz_now = || MirScalarExpr::CallUnmaterializable(UnmaterializableFunc::MzNow);
        let cast = || MirScalarExpr::column(0).call_unary(func::CastInt64ToMzTimestamp);
        let le = || mz_now().call_binary(cast(), func::Lte);

        let mut df = dataflow(
            MirRelationExpr::Get {
                id: Id::Global(READ),
                typ: typ(),
                access_strategy: AccessStrategy::Persist,
            }
            .filter([le()]),
        );
        let read = df.source_imports.get_mut(&READ).expect("imported");
        read.desc.arguments.operators = Some(MapFilterProject::new(1).filter([le()]));

        let lowered = LirRelationExpr::finalize_dataflow(df, &OptimizerFeatures::default(), None)
            .expect("lowers");
        let operators = lowered.source_imports[&READ]
            .desc
            .arguments
            .operators
            .as_ref()
            .expect("planned operators");
        let (_safe, lower, upper) = operators.as_parts();
        assert!(lower.is_empty(), "{operators:?}");
        assert_eq!(upper.len(), 1, "{operators:?}");
    }
}