Merge pull request #30666 from ggevay/insert-opt

Create a special optimizer pipeline for constant INSERTs

Author: ggevay

misc/python/materialize/feature_benchmark/scenarios/benchmark_main.py

@@ -311,7 +311,10 @@ def benchmark(self) -> MeasurementSource:
 
 
 class InsertMultiRow(DML):
-    """Measure the time it takes for a single multi-row INSERT statement to return."""
+    """Measure the time it takes for a single multi-row INSERT statement to return.
+    When `sequence_insert` calls `constant_optimizer`, it should be able to reach a constant. Otherwise, we run the full
+    logical optimizer, which makes this test show a regression.
+    """
 
     SCALE = 4  # FATAL:  request larger than 2.0 MB
 

src/adapter/src/catalog/state.rs

@@ -1200,13 +1200,12 @@ impl CatalogState {
                     }
                     Some(_) | None => {
                         let optimizer_features = optimizer_config.features.clone();
-                        // Build an optimizer for this VIEW.
                         // TODO(aalexandrov): ideally this should be a materialized_view::Optimizer.
                         let mut optimizer = optimize::view::Optimizer::new(optimizer_config, None);
 
                         let raw_expr = materialized_view.expr;
                         let optimized_expr = match optimizer.optimize(raw_expr.clone()) {
-                            Ok(optimzed_expr) => optimzed_expr,
+                            Ok(optimized_expr) => optimized_expr,
                             Err(err) => return Err((err.into(), cached_expr)),
                         };
 

src/adapter/src/client.rs

@@ -741,7 +741,6 @@ impl SessionClient {
 
         // Collect optimizer parameters.
         let optimizer_config = optimize::OptimizerConfig::from(conn_catalog.system_vars());
-        // Build an optimizer for this VIEW.
         let mut optimizer = optimize::view::Optimizer::new(optimizer_config, None);
 
         let result: Result<_, AdapterError> =

src/adapter/src/coord/sequencer/inner.rs

@@ -2647,7 +2647,7 @@ impl Coordinator {
             // Collect optimizer parameters.
             let optimizer_config = optimize::OptimizerConfig::from(self.catalog().system_config());
 
-            // Build an optimizer for this VIEW.
+            // (`optimize::view::Optimizer` has a special case for constant queries.)
             let mut optimizer = optimize::view::Optimizer::new(optimizer_config, None);
 
             // HIR ⇒ MIR lowering and MIR ⇒ MIR optimization (local)

src/adapter/src/coord/sequencer/inner/explain_timestamp.rs

@@ -130,7 +130,6 @@ impl Coordinator {
         // Collect optimizer parameters.
         let optimizer_config = optimize::OptimizerConfig::from(self.catalog().system_config());
 
-        // Build an optimizer for this VIEW.
         let mut optimizer = optimize::view::Optimizer::new(optimizer_config, None);
 
         let span = Span::current();

src/adapter/src/optimize.rs

@@ -359,6 +359,22 @@ fn optimize_mir_local(
     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,
+) -> Result<MirRelationExpr, OptimizerError> {
+    let optimizer = mz_transform::Optimizer::constant_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.0)
+}
+
 macro_rules! trace_plan {
     (at: $span:literal, $plan:expr) => {
         tracing::debug_span!(target: "optimizer", $span).in_scope(|| {

src/adapter/src/optimize/view.rs

@@ -7,7 +7,13 @@
 // the Business Source License, use of this software will be governed
 // by the Apache License, Version 2.0.
 
-//! Optimizer implementation for `CREATE VIEW` statements.
+//! An Optimizer that
+//! 1. Optimistically calls `optimize_mir_constant`.
+//! 2. Then, if we haven't arrived at a constant, it calls `optimize_mir_local`, i.e., the
+//!    logical optimizer.
+//!
+//! This is used for `CREATE VIEW` statements and in various other situations where no physical
+//! optimization is needed, such as for `INSERT` statements.
 
 use std::time::Instant;
 
@@ -18,7 +24,10 @@ use mz_transform::dataflow::DataflowMetainfo;
 use mz_transform::typecheck::{empty_context, SharedContext as TypecheckContext};
 use mz_transform::TransformCtx;
 
-use crate::optimize::{optimize_mir_local, trace_plan, Optimize, OptimizerConfig, OptimizerError};
+use crate::optimize::{
+    optimize_mir_constant, optimize_mir_local, trace_plan, Optimize, OptimizerConfig,
+    OptimizerError,
+};
 
 pub struct Optimizer {
     /// A typechecking context to use throughout the optimizer pipeline.
@@ -52,13 +61,24 @@ impl Optimize<HirRelationExpr> for Optimizer {
         trace_plan!(at: "raw", &expr);
 
         // HIR ⇒ MIR lowering and decorrelation
-        let expr = expr.lower(&self.config, self.metrics.as_ref())?;
+        let mut expr = expr.lower(&self.config, self.metrics.as_ref())?;
 
-        // MIR ⇒ MIR optimization (local)
         let mut df_meta = DataflowMetainfo::default();
         let mut transform_ctx =
             TransformCtx::local(&self.config.features, &self.typecheck_ctx, &mut df_meta);
-        let expr = optimize_mir_local(expr, &mut transform_ctx)?;
+
+        // First, we run a very simple optimizer pipeline, which only folds constants. This takes
+        // care of constant INSERTs. (This optimizer is also used for INSERTs, not just VIEWs.)
+        expr = optimize_mir_constant(expr, &mut transform_ctx)?;
+
+        // MIR ⇒ MIR optimization (local)
+        let expr = if expr.as_const().is_some() {
+            // No need to optimize further, because we already have a constant.
+            OptimizedMirRelationExpr(expr)
+        } else {
+            // Call the real optimization.
+            optimize_mir_local(expr, &mut transform_ctx)?
+        };
 
         if let Some(metrics) = &self.metrics {
             metrics.observe_e2e_optimization_time("view", time.elapsed());

src/transform/src/lib.rs

@@ -552,9 +552,11 @@ pub fn fuse_and_collapse_fixpoint() -> Fixpoint {
     }
 }
 
-/// Does constant folding idempotently. This needs to call `FoldConstants` together with
-/// `NormalizeLets` in a fixpoint loop, because currently `FoldConstants` doesn't inline CTEs, so
-/// these two need to alternate until fixpoint.
+/// Does constant folding to a fixpoint: An expression all of whose leaves are constants, of size
+/// small enough to be inlined and folded should reach a single `MirRelationExpr::Constant`.
+///
+/// This needs to call `FoldConstants` together with `NormalizeLets` in a fixpoint loop, because
+/// currently `FoldConstants` doesn't inline CTEs, so these two need to alternate until fixpoint.
 ///
 /// Also note that `FoldConstants` can break the normalized form by removing all references to a
 /// Let.
@@ -819,6 +821,15 @@ impl Optimizer {
         }
     }
 
+    /// Builds a tiny optimizer, which just folds constants. For more details, see
+    /// [fold_constants_fixpoint].
+    pub fn constant_optimizer(_ctx: &mut TransformCtx) -> Self {
+        Self {
+            name: "fast_path_optimizer",
+            transforms: vec![Box::new(fold_constants_fixpoint())],
+        }
+    }
+
     /// Optimizes the supplied relation expression.
     ///
     /// These optimizations are performed with no information about available arrangements,