mz_adapter/

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

//! Optimizer interface to the adapter and coordinator code.
//!
//! The goal of this crate is to abstract optimizer specifics behind a
//! high-level interface that is ready to be consumed by the coordinator code in
//! a future-proof way (that is, the API is taking the upcoming evolution of
//! these components into account).
//!
//! The contents of this crate should have minimal dependencies to the rest of
//! the coordinator code so we can pull them out as a separate crate in the
//! future without too much effort.
//!
//! The main type in this module is a very simple [`Optimize`] trait which
//! allows us to adhere to the following principles:
//!
//! - Implementors of this trait are structs that encapsulate all context
//!   required to optimize a statement of type `T` end-to-end (for example
//!   [`materialized_view::Optimizer`] for `T` = `MaterializedView`).
//! - Each struct implements [`Optimize`] once for each optimization stage. The
//!   `From` type represents the input of the stage and `Self::To` the
//!   associated stage output. This allows to have more than one entrypoints to
//!   a pipeline.
//! - The concrete types used for stage results are opaque structs that are
//!   specific to the pipeline of that statement type.
//!   - We use different structs even if two statement types might have
//!     structurally identical intermediate results. This ensures that client
//!     code cannot first execute some optimization stages for one type and then
//!     some stages for a different type.
//!   - The only way to construct such a struct is by running the [`Optimize`]
//!     stage that produces it. This ensures that client code cannot interfere
//!     with the pipeline.
//!   - In general, the internals of these structs can be accessed only behind a
//!     shared reference. This ensures that client code can look up information
//!     from intermediate stages but cannot modify it.
//!   - Timestamp selection is modeled as a conversion between structs that are
//!     adjacent in the pipeline using a method called `resolve`.
//!   - The struct representing the result of the final stage of the
//!     optimization pipeline can be destructed to access its internals with a
//!     method called `unapply`.
//! - The `Send` trait bounds on the `Self` and `From` types ensure that
//!   [`Optimize`] instances can be passed to different threads (this is
//!   required of off-thread optimization).
//!
//! For details, see the `20230714_optimizer_interface.md` design doc in this
//! repository.

pub mod copy_to;
pub mod dataflows;
pub mod index;
pub mod materialized_view;
pub mod peek;
pub mod subscribe;
pub mod view;

use std::fmt::Debug;

use mz_adapter_types::connection::ConnectionId;
use mz_adapter_types::dyncfgs::PERSIST_FAST_PATH_ORDER;
use mz_catalog::memory::objects::{CatalogCollectionEntry, CatalogEntry, Index};
use mz_compute_types::ComputeInstanceId;
use mz_compute_types::dataflows::DataflowDescription;
use mz_compute_types::dyncfgs::SUBSCRIBE_SNAPSHOT_OPTIMIZATION;
use mz_compute_types::plan::Plan;
use mz_controller_types::ClusterId;
use mz_expr::{EvalError, MirRelationExpr, OptimizedMirRelationExpr, UnmaterializableFunc};
use mz_ore::stack::RecursionLimitError;
use mz_repr::adt::timestamp::TimestampError;
use mz_repr::optimize::{OptimizerFeatureOverrides, OptimizerFeatures, OverrideFrom};
use mz_repr::{CatalogItemId, GlobalId};
use mz_sql::names::{FullItemName, QualifiedItemName};
use mz_sql::plan::{HirRelationExpr, PlanError};
use mz_sql::session::metadata::SessionMetadata;
use mz_sql::session::vars::SystemVars;
use mz_transform::{MaybeShouldPanic, StatisticsOracle, TransformCtx, TransformError};

use crate::TimestampContext;

// Alias types
// -----------

/// A type for a [`DataflowDescription`] backed by `Mir~` plans. Used internally
/// by the optimizer implementations.
type MirDataflowDescription = DataflowDescription<OptimizedMirRelationExpr>;
/// A type for a [`DataflowDescription`] backed by `Lir~` plans. Used internally
/// by the optimizer implementations.
type LirDataflowDescription = DataflowDescription<Plan>;

// Core API
// --------

/// A trait that represents an optimization stage.
///
/// The trait is implemented by structs that encapsulate the context needed to
/// run an end-to-end optimization pipeline for a specific statement type
/// (`Index`, `View`, `MaterializedView`, `Subscribe`, `Select`).
///
/// Each implementation represents a concrete optimization stage for a fixed
/// statement type that consumes an input of type `From` and produces output of
/// type `Self::To`.
pub trait Optimize<From> {
    type To;

    /// Execute the optimization stage, transforming the input plan of type
    /// `From` to an output plan of type `To`.
    fn optimize(&mut self, plan: From) -> Result<Self::To, OptimizerError>;

    /// Like [`Self::optimize`], but additionally ensures that panics occurring
    /// in the [`Self::optimize`] call are caught and demoted to an
    /// [`OptimizerError::Internal`] error.
    ///
    /// Additionally, if the result of the optimization is an error (not a panic) that indicates we
    /// should panic, then panic.
    #[mz_ore::instrument(target = "optimizer", level = "debug", name = "optimize")]
    fn catch_unwind_optimize(&mut self, plan: From) -> Result<Self::To, OptimizerError> {
        mz_transform::catch_unwind_optimize(|| self.optimize(plan))
    }
}

// One-shot peek optimizer dispatch
// --------------------------------

/// The optimizer driving a one-shot statement through the peek sequencing state
/// machine.
///
/// `SELECT` and `EXPLAIN` are optimized by the [`peek::Optimizer`], while `COPY
/// TO` is optimized by the [`copy_to::Optimizer`]. Both share the same
/// surrounding state machine (timestamp selection, read holds, off-thread
/// optimization, …), so this enum lets that shared machinery carry either
/// optimizer without caring which one it is. The variants are kept distinct
/// (rather than abstracted behind a trait object) because the downstream stages
/// need to recover the concrete optimizer and its statement-specific result.
#[derive(Debug)]
pub enum PeekOptimizer {
    /// Optimizer for `SELECT` and `EXPLAIN` statements.
    Select(peek::Optimizer),
    /// Optimizer for `COPY TO` statements.
    CopyTo(copy_to::Optimizer),
}

/// The global LIR plan produced by [`PeekOptimizer::optimize`], tagged with the
/// path that produced it.
#[derive(Debug)]
pub enum PeekGlobalLirPlan {
    /// The result of the `SELECT`/`EXPLAIN` pipeline.
    Select(peek::GlobalLirPlan),
    /// The result of the `COPY TO` pipeline.
    CopyTo(copy_to::GlobalLirPlan),
}

impl PeekOptimizer {
    /// The cluster that will run the optimized dataflow.
    pub fn cluster_id(&self) -> ComputeInstanceId {
        match self {
            PeekOptimizer::Select(optimizer) => optimizer.cluster_id(),
            PeekOptimizer::CopyTo(optimizer) => optimizer.cluster_id(),
        }
    }

    /// Runs the full one-shot optimization pipeline end-to-end:
    ///
    /// 1. HIR ⇒ MIR lowering and local MIR optimization,
    /// 2. timestamp resolution,
    /// 3. global MIR optimization, MIR ⇒ LIR lowering, and global LIR
    ///    optimization.
    ///
    /// The pipeline shape is identical for both variants; only the concrete
    /// (statement-specific) plan types differ, so the steps are shared via
    /// [`optimize_oneshot`].
    pub fn optimize(
        &mut self,
        raw_expr: HirRelationExpr,
        timestamp_ctx: TimestampContext,
        session: &dyn SessionMetadata,
        stats: Box<dyn StatisticsOracle>,
    ) -> Result<PeekGlobalLirPlan, OptimizerError> {
        match self {
            PeekOptimizer::Select(optimizer) => {
                let plan = optimize_oneshot(optimizer, raw_expr, |local_mir_plan| {
                    local_mir_plan.resolve(timestamp_ctx, session, stats)
                })?;
                Ok(PeekGlobalLirPlan::Select(plan))
            }
            PeekOptimizer::CopyTo(optimizer) => {
                let plan = optimize_oneshot(optimizer, raw_expr, |local_mir_plan| {
                    local_mir_plan.resolve(timestamp_ctx, session, stats)
                })?;
                Ok(PeekGlobalLirPlan::CopyTo(plan))
            }
        }
    }

    /// Consumes `self`, returning the inner [`peek::Optimizer`] if this is the
    /// `SELECT`/`EXPLAIN` path and `None` otherwise.
    pub fn into_select(self) -> Option<peek::Optimizer> {
        match self {
            PeekOptimizer::Select(optimizer) => Some(optimizer),
            PeekOptimizer::CopyTo(_) => None,
        }
    }

    /// Consumes `self`, returning the inner [`copy_to::Optimizer`] if this is
    /// the `COPY TO` path and `None` otherwise.
    pub fn into_copy_to(self) -> Option<copy_to::Optimizer> {
        match self {
            PeekOptimizer::CopyTo(optimizer) => Some(optimizer),
            PeekOptimizer::Select(_) => None,
        }
    }
}

/// Runs the shared one-shot optimization pipeline for a single optimizer.
///
/// This factors out the (otherwise duplicated) HIR ⇒ local MIR ⇒ resolve ⇒
/// global LIR sequence that is common to the `SELECT`/`EXPLAIN` and `COPY TO`
/// paths. The `resolve` closure attaches the timestamp/session/stats context to
/// the local plan; it is path-specific only in the concrete plan type it
/// operates on.
fn optimize_oneshot<O, LocalPlan, ResolvedPlan, GlobalPlan>(
    optimizer: &mut O,
    raw_expr: HirRelationExpr,
    resolve: impl FnOnce(LocalPlan) -> ResolvedPlan,
) -> Result<GlobalPlan, OptimizerError>
where
    O: Optimize<HirRelationExpr, To = LocalPlan> + Optimize<ResolvedPlan, To = GlobalPlan>,
{
    // HIR ⇒ MIR lowering and MIR optimization (local).
    let local_mir_plan = optimizer.catch_unwind_optimize(raw_expr)?;
    // Attach resolved context required to continue the pipeline.
    let resolved_mir_plan = resolve(local_mir_plan);
    // MIR optimization (global), MIR ⇒ LIR lowering, and LIR optimization (global).
    optimizer.catch_unwind_optimize(resolved_mir_plan)
}

// Optimizer configuration
// -----------------------

/// Feature flags for the optimizer.
///
/// To add a new feature flag, do the following steps:
///
/// 1. To make the flag available to all stages in our [`Optimize`] pipelines
///    and allow engineers to set a system-wide override:
///    1. Add the flag to the `optimizer_feature_flags!(...)` macro call.
///    2. Add the flag to the `feature_flags!(...)` macro call and extend the
///       `From<&SystemVars>` implementation for [`OptimizerFeatures`].
///
/// 2. To enable `EXPLAIN ... WITH(...)` overrides which will allow engineers to
///    inspect plan differences before deploying the optimizer changes:
///    1. Add the flag to the `ExplainPlanOptionName` definition.
///    2. Add the flag to the `generate_extracted_config!(ExplainPlanOption,
///       ...)` macro call.
///    3. Extend the `TryFrom<ExplainPlanOptionExtracted>` implementation for
///       [`mz_repr::explain::ExplainConfig`].
///
/// 3. To enable `CLUSTER ... FEATURES(...)` overrides which will allow
///    engineers to experiment with runtime differences before deploying the
///    optimizer changes:
///    1. Add the flag to the `ClusterFeatureName` definition.
///    2. Add the flag to the `generate_extracted_config!(ClusterFeature, ...)`
///       macro call.
///    3. Extend the `let optimizer_feature_overrides = ...` call in
///       `plan_create_cluster`.
#[derive(Clone, Debug)]
pub struct OptimizerConfig {
    /// The mode in which the optimizer runs.
    pub mode: OptimizeMode,
    /// If the [`GlobalId`] is set the optimizer works in "replan" mode.
    ///
    /// This means that it will not consider catalog items (more specifically
    /// indexes) with [`GlobalId`] greater or equal than the one provided here.
    pub replan: Option<GlobalId>,
    /// Show the slow path plan even if a fast path plan was created. Useful for debugging.
    /// Enforced if `timing` is set.
    pub no_fast_path: bool,
    // If set, allow some additional queries down the Persist fast path when we believe
    // the orderings are compatible.
    persist_fast_path_order: bool,
    // Enable calculating with_snapshot metadata for subscribes.
    subscribe_snapshot_optimization: bool,
    /// Optimizer feature flags.
    pub features: OptimizerFeatures,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub enum OptimizeMode {
    /// A mode where the optimized statement is executed.
    Execute,
    /// A mode where the optimized statement is explained.
    Explain,
}

impl From<&SystemVars> for OptimizerConfig {
    fn from(vars: &SystemVars) -> Self {
        Self {
            mode: OptimizeMode::Execute,
            replan: None,
            no_fast_path: false,
            persist_fast_path_order: PERSIST_FAST_PATH_ORDER.get(vars.dyncfgs()),
            subscribe_snapshot_optimization: SUBSCRIBE_SNAPSHOT_OPTIMIZATION.get(vars.dyncfgs()),
            features: OptimizerFeatures::from(vars),
        }
    }
}

/// Override [`OptimizerConfig::features`] from [`OptimizerFeatureOverrides`].
impl OverrideFrom<OptimizerFeatureOverrides> for OptimizerConfig {
    fn override_from(mut self, overrides: &OptimizerFeatureOverrides) -> Self {
        self.features = self.features.override_from(overrides);
        self
    }
}

/// [`OptimizerConfig`] overrides coming from an [`ExplainContext`].
impl OverrideFrom<ExplainContext> for OptimizerConfig {
    fn override_from(mut self, ctx: &ExplainContext) -> Self {
        let ExplainContext::Plan(ctx) = ctx else {
            return self; // Return immediately for all other contexts.
        };

        // Override general parameters.
        self.mode = OptimizeMode::Explain;
        self.replan = ctx.replan;
        self.no_fast_path = ctx.config.no_fast_path;

        // Override feature flags that can be enabled in the EXPLAIN config.
        self.features = self.features.override_from(&ctx.config.features);

        // Return the final result.
        self
    }
}

impl From<&OptimizerConfig> for mz_sql::plan::HirToMirConfig {
    fn from(config: &OptimizerConfig) -> Self {
        Self {
            enable_new_outer_join_lowering: config.features.enable_new_outer_join_lowering,
            enable_variadic_left_join_lowering: config.features.enable_variadic_left_join_lowering,
            enable_cast_elimination: config.features.enable_cast_elimination,
            enable_simplify_quantified_comparisons: config
                .features
                .enable_simplify_quantified_comparisons,
        }
    }
}

// OptimizerCatalog
// ===============

pub trait OptimizerCatalog: Debug + Send + Sync {
    fn get_entry(&self, id: &GlobalId) -> CatalogCollectionEntry;
    fn get_entry_by_item_id(&self, id: &CatalogItemId) -> &CatalogEntry;
    fn resolve_full_name(
        &self,
        name: &QualifiedItemName,
        conn_id: Option<&ConnectionId>,
    ) -> FullItemName;

    /// Returns all indexes on the given object and cluster known in the
    /// catalog.
    fn get_indexes_on(
        &self,
        id: GlobalId,
        cluster: ClusterId,
    ) -> Box<dyn Iterator<Item = (GlobalId, &Index)> + '_>;
}

// OptimizerError
// ===============

/// Error types that can be generated during optimization.
#[derive(Debug, thiserror::Error)]
pub enum OptimizerError {
    #[error("{0}")]
    PlanError(#[from] PlanError),
    #[error("{0}")]
    RecursionLimitError(#[from] RecursionLimitError),
    #[error("{0}")]
    TransformError(#[from] TransformError),
    #[error("{0}")]
    EvalError(#[from] EvalError),
    #[error("cannot materialize call to {0}")]
    UnmaterializableFunction(UnmaterializableFunc),
    #[error("cannot call {func} in {context} ")]
    UncallableFunction {
        func: UnmaterializableFunc,
        context: &'static str,
    },
    #[error("access to function {0} is restricted")]
    RestrictedFunction(UnmaterializableFunc),
    #[error("{0}")]
    UnsupportedTemporalExpression(String),
    /// This is a specific kind of internal error. It's distinct from `Internal`, because we want to
    /// catch it and swallow it in some cases.
    #[error("internal optimizer error: MfpPlan couldn't be converted into SafeMfpPlan")]
    InternalUnsafeMfpPlan(String),
    #[error("internal optimizer error: {0}")]
    Internal(String),
}

impl From<String> for OptimizerError {
    fn from(msg: String) -> Self {
        Self::Internal(msg)
    }
}

impl OptimizerError {
    pub fn detail(&self) -> Option<String> {
        match self {
            Self::UnmaterializableFunction(UnmaterializableFunc::CurrentTimestamp) => {
                Some("See: https://materialize.com/docs/sql/functions/now_and_mz_now/".into())
            }
            Self::RestrictedFunction(_) => Some(
                "Access to system catalog objects is restricted for this role. \
                Contact your administrator if you need access."
                    .into(),
            ),
            _ => None,
        }
    }

    pub fn hint(&self) -> Option<String> {
        match self {
            Self::UnmaterializableFunction(UnmaterializableFunc::CurrentTimestamp) => {
                Some("In temporal filters `mz_now()` may work instead.".into())
            }
            _ => None,
        }
    }
}

impl From<TimestampError> for OptimizerError {
    fn from(value: TimestampError) -> Self {
        OptimizerError::EvalError(EvalError::from(value))
    }
}

impl From<anyhow::Error> for OptimizerError {
    fn from(value: anyhow::Error) -> Self {
        OptimizerError::Internal(value.to_string())
    }
}

impl MaybeShouldPanic for OptimizerError {
    fn should_panic(&self) -> Option<String> {
        match self {
            OptimizerError::TransformError(TransformError::CallerShouldPanic(msg)) => {
                Some(msg.to_string())
            }
            _ => None,
        }
    }
}

// Tracing helpers
// ---------------

#[mz_ore::instrument(target = "optimizer", level = "debug", name = "local")]
fn optimize_mir_local(
    expr: MirRelationExpr,
    ctx: &mut TransformCtx,
) -> Result<OptimizedMirRelationExpr, OptimizerError> {
    #[allow(deprecated)]
    let optimizer = mz_transform::Optimizer::logical_optimizer(ctx);
    let expr = optimizer.optimize(expr, ctx)?;

    // Trace the result of this phase.
    mz_repr::explain::trace_plan(expr.as_inner());

    Ok::<_, OptimizerError>(expr)
}

/// This is just a wrapper around [mz_transform::Optimizer::constant_optimizer],
/// running it, and tracing the result plan.
#[mz_ore::instrument(target = "optimizer", level = "debug", name = "constant")]
fn optimize_mir_constant(
    expr: MirRelationExpr,
    ctx: &mut TransformCtx,
    limit: bool,
) -> Result<MirRelationExpr, OptimizerError> {
    let optimizer = mz_transform::Optimizer::constant_optimizer(ctx, limit);
    let expr = optimizer.optimize(expr, ctx)?;

    // Trace the result of this phase.
    mz_repr::explain::trace_plan(expr.as_inner());

    Ok::<_, OptimizerError>(expr.0)
}

macro_rules! trace_plan {
    (at: $span:literal, $plan:expr) => {
        tracing::debug_span!(target: "optimizer", $span).in_scope(|| {
            mz_repr::explain::trace_plan($plan);
        });
    }
}

use trace_plan;

use crate::coord::ExplainContext;