Create a special optimizer pipeline for constant INSERTs

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.
+    `constant_insert_optimizer` 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/coord/sequencer/inner.rs

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

src/adapter/src/optimize.rs

@@ -359,6 +359,22 @@ fn optimize_mir_local(
     Ok::<_, OptimizerError>(expr)
 }
 
+/// This calls a simple optimizer that is able to fully optimize constant INSERTs. For more details,
+/// see [mz_transform::Optimizer::constant_insert_optimizer].
+#[mz_ore::instrument(target = "optimizer", level = "debug", name = "constant_insert")]
+fn optimize_mir_constant_insert(
+    expr: MirRelationExpr,
+    ctx: &mut TransformCtx,
+) -> Result<MirRelationExpr, OptimizerError> {
+    let optimizer = mz_transform::Optimizer::constant_insert_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

@@ -19,7 +19,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_insert, optimize_mir_local, trace_plan, Optimize, OptimizerConfig,
+    OptimizerError,
+};
 
 pub struct Optimizer {
     /// A typechecking context to use throughout the optimizer pipeline.
@@ -53,14 +56,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())?;
 
         let mut df_meta = DataflowMetainfo::default();
         let mut transform_ctx =
             TransformCtx::local(&self.config.features, &self.typecheck_ctx, &mut df_meta);
 
+        // This optimizer is also used for INSERTs (in addition to VIEWs). Therefore, we first call
+        // a very simple optimizer, which should fully take care of constant INSERTs.
+        expr = optimize_mir_constant_insert(expr, &mut transform_ctx)?;
+
         // MIR ⇒ MIR optimization (local)
-        let expr = optimize_mir_local(expr, &mut transform_ctx)?;
+        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

@@ -819,6 +819,40 @@ impl Optimizer {
         }
     }
 
+    /// Builds a tiny optimizer, which is only suitable for optimizing the MIR plans that come from
+    /// constant INSERTs, e.g.
+    /// ```code
+    /// Return
+    ///   FlatMap wrap2(1, 2)
+    ///     Get l0
+    /// With
+    ///  cte l0 =
+    ///     Constant
+    ///       - ()
+    /// ```
+    /// which is planned for stuff like `INSERT INTO t VALUES (1,2);`.
+    ///
+    /// The important thing here is dealing with the `wrap` table function, which is always planned
+    /// for `VALUES`, and results in a Let binding in MIR. This is why we need to call
+    /// `NormalizeLets`, so that this Let binding is inlined. Then, `FoldConstants` is able to
+    /// evaluate the table function.
+    ///
+    /// The `InsertMultiRow` test is meant to make noise if this doesn't actually reach a constant
+    /// for a constant INSERT.
+    pub fn constant_insert_optimizer(_ctx: &mut TransformCtx) -> Self {
+        let transforms: Vec<Box<dyn Transform>> = vec![
+            Box::new(NormalizeLets::new(false)),
+            Box::new(canonicalization::ReduceScalars),
+            Box::new(FoldConstants {
+                limit: Some(FOLD_CONSTANTS_LIMIT),
+            }),
+        ];
+        Self {
+            name: "constant_insert_optimizer",
+            transforms,
+        }
+    }
+
     /// Optimizes the supplied relation expression.
     ///
     /// These optimizations are performed with no information about available arrangements,