Merge pull request #17508 from aalexandrov/explain_remove_old

Closes https://github.com/MaterializeInc/materialize/issues/17360

Author: aalexandrov

src/adapter/src/catalog.rs

@@ -46,7 +46,7 @@ use mz_ore::metrics::MetricsRegistry;
 use mz_ore::now::{to_datetime, EpochMillis, NowFn};
 use mz_persist_client::cfg::PersistParameters;
 use mz_pgrepr::oid::FIRST_USER_OID;
-use mz_repr::{explain_new::ExprHumanizer, Diff, GlobalId, RelationDesc, ScalarType};
+use mz_repr::{explain::ExprHumanizer, Diff, GlobalId, RelationDesc, ScalarType};
 use mz_secrets::InMemorySecretsController;
 use mz_sql::ast::display::AstDisplay;
 use mz_sql::ast::Expr;

src/adapter/src/coord.rs

@@ -101,7 +101,7 @@ use mz_ore::tracing::OpenTelemetryContext;
 use mz_ore::{stack, task};
 use mz_persist_client::usage::StorageUsageClient;
 use mz_persist_client::ShardId;
-use mz_repr::explain_new::ExplainFormat;
+use mz_repr::explain::ExplainFormat;
 use mz_repr::{Datum, Diff, GlobalId, Row, Timestamp};
 use mz_secrets::SecretsController;
 use mz_sql::ast::{CreateSourceStatement, CreateSubsourceStatement, Raw, Statement};

src/adapter/src/coord/peek.rs

@@ -32,9 +32,8 @@ use mz_ore::cast::CastFrom;
 use mz_ore::str::Indent;
 use mz_ore::str::StrExt;
 use mz_ore::tracing::OpenTelemetryContext;
-use mz_repr::explain_new::{
-    fmt_text_constant_rows, separated_text, CompactScalarSeq, DisplayText, ExprHumanizer, Indices,
-};
+use mz_repr::explain::text::{fmt_text_constant_rows, separated_text, DisplayText};
+use mz_repr::explain::{CompactScalarSeq, ExprHumanizer, Indices};
 use mz_repr::{Diff, GlobalId, RelationType, Row};
 
 use crate::client::ConnectionId;
@@ -564,10 +563,9 @@ fn consolidate_constant_updates(rows: Vec<(Row, Diff)>) -> Vec<(Row, Diff)> {
 mod tests {
     use mz_expr::{func::IsNull, MapFilterProject, UnaryFunc};
     use mz_ore::str::Indent;
-    use mz_repr::{
-        explain_new::{text_string_at, DummyHumanizer, RenderingContext},
-        ColumnType, Datum, ScalarType,
-    };
+    use mz_repr::explain::text::text_string_at;
+    use mz_repr::explain::{DummyHumanizer, RenderingContext};
+    use mz_repr::{ColumnType, Datum, ScalarType};
 
     use super::*;
 

src/adapter/src/coord/sequencer.rs

@@ -36,7 +36,7 @@ use mz_expr::{
     OptimizedMirRelationExpr, RowSetFinishing,
 };
 use mz_ore::task;
-use mz_repr::explain_new::{ExplainFormat, Explainee};
+use mz_repr::explain::{ExplainFormat, Explainee};
 use mz_repr::{Datum, Diff, GlobalId, RelationDesc, Row, RowArena, Timestamp};
 use mz_sql::ast::{ExplainStage, IndexOptionName, ObjectType};
 use mz_sql::catalog::CatalogItem as SqlCatalogItem;
@@ -79,7 +79,7 @@ use crate::coord::{
     SinkConnectionReady, DEFAULT_LOGICAL_COMPACTION_WINDOW_TS,
 };
 use crate::error::AdapterError;
-use crate::explain_new::optimizer_trace::OptimizerTrace;
+use crate::explain::optimizer_trace::OptimizerTrace;
 use crate::metrics;
 use crate::notice::AdapterNotice;
 use crate::session::vars::{
@@ -2707,7 +2707,7 @@ impl Coordinator {
         plan: ExplainPlan,
     ) -> Result<ExecuteResponse, AdapterError> {
         use mz_compute_client::plan::Plan;
-        use mz_repr::explain_new::trace_plan;
+        use mz_repr::explain::trace_plan;
         use ExplainStage::*;
 
         let cluster = self.catalog.active_cluster(session)?.id;
@@ -2755,7 +2755,7 @@ impl Coordinator {
                             &optimized_plan,
                             &mut dataflow,
                         )?;
-                        mz_repr::explain_new::trace_plan(&dataflow);
+                        mz_repr::explain::trace_plan(&dataflow);
                         Ok(dataflow)
                     },
                 )?;

src/adapter/src/coord/timestamp_selection.rs

@@ -19,7 +19,7 @@ use tracing::{event, Level};
 
 use mz_compute_client::controller::ComputeInstanceId;
 use mz_expr::MirScalarExpr;
-use mz_repr::explain_new::ExprHumanizer;
+use mz_repr::explain::ExprHumanizer;
 use mz_repr::{RowArena, ScalarType, Timestamp, TimestampManipulation};
 use mz_sql::plan::QueryWhen;
 use mz_storage_client::types::sources::Timeline;

src/adapter/src/error.rs

@@ -19,7 +19,7 @@ use mz_compute_client::controller::error as compute_error;
 use mz_expr::{EvalError, UnmaterializableFunc};
 use mz_ore::stack::RecursionLimitError;
 use mz_ore::str::StrExt;
-use mz_repr::explain_new::ExplainError;
+use mz_repr::explain::ExplainError;
 use mz_repr::NotNullViolation;
 use mz_sql::plan::PlanError;
 use mz_storage_client::controller::StorageError;

src/adapter/src/explain_new/fast_path/mod.rs renamed to src/adapter/src/explain/fast_path/mod.rs

@@ -11,10 +11,8 @@
 
 use std::collections::BTreeMap;
 
-use mz_expr::explain_new::ExplainMultiPlan;
-use mz_repr::explain_new::{
-    AnnotatedPlan, Explain, ExplainConfig, ExplainError, UnsupportedFormat,
-};
+use mz_expr::explain::ExplainMultiPlan;
+use mz_repr::explain::{AnnotatedPlan, Explain, ExplainConfig, ExplainError, UnsupportedFormat};
 
 use crate::coord::peek::FastPathPlan;
 

src/adapter/src/explain_new/hir/mod.rs renamed to src/adapter/src/explain/hir/mod.rs

@@ -9,7 +9,7 @@
 
 //! `EXPLAIN` support for HIR structures.
 
-use mz_repr::explain_new::{Explain, ExplainConfig, ExplainError};
+use mz_repr::explain::{Explain, ExplainConfig, ExplainError};
 use mz_sql::plan::HirRelationExpr;
 
 use super::Explainable;

src/adapter/src/explain_new/lir/mod.rs renamed to src/adapter/src/explain/lir/mod.rs

@@ -11,7 +11,7 @@
 
 use mz_compute_client::plan::Plan;
 use mz_compute_client::types::dataflows::DataflowDescription;
-use mz_repr::explain_new::{Explain, ExplainConfig, ExplainError};
+use mz_repr::explain::{Explain, ExplainConfig, ExplainError};
 
 use super::Explainable;
 

src/adapter/src/explain_new/mir/mod.rs renamed to src/adapter/src/explain/mir/mod.rs

@@ -8,22 +8,22 @@
 // by the Apache License, Version 2.0.
 
 //! `EXPLAIN` support for MIR structures.
-
-use std::collections::BTreeMap;
+//!
+//! The specialized [`Explain`] implementation for an [`MirRelationExpr`]
+//! wrapped in an [`Explainable`] newtype struct allows us to interpret more
+//! [`ExplainConfig`] options. This is the case because attribute derivation and
+//! let normalization are defined in [`mz_transform`] and conssequently are not
+//! available for the default [`Explain`] implementation for [`MirRelationExpr`]
+//! in [`mz_expr`].
 
 use mz_compute_client::types::dataflows::DataflowDescription;
-use mz_expr::{
-    explain_new::{enforce_linear_chains, ExplainContext, ExplainMultiPlan, ExplainSinglePlan},
-    visit::Visit,
-    MirRelationExpr, OptimizedMirRelationExpr,
-};
-use mz_ore::{stack::RecursionLimitError, str::bracketed, str::separated};
-use mz_repr::explain_new::{
-    AnnotatedPlan, Attributes, Explain, ExplainConfig, ExplainError, UnsupportedFormat,
-};
-use mz_transform::attribute::{
-    Arity, DerivedAttributes, NonNegative, RelationType, SubtreeSize, UniqueKeys,
+use mz_expr::explain::{
+    enforce_linear_chains, ExplainContext, ExplainMultiPlan, ExplainSinglePlan,
 };
+use mz_expr::{MirRelationExpr, OptimizedMirRelationExpr};
+use mz_repr::explain::{Explain, ExplainConfig, ExplainError, UnsupportedFormat};
+use mz_transform::attribute::annotate_plan;
+use mz_transform::normalize_lets::normalize_lets;
 
 use super::Explainable;
 
@@ -68,12 +68,13 @@ impl<'a> Explainable<'a, MirRelationExpr> {
         // normalize the representation of nested Let bindings
         // and enforce sequential Let binding IDs
         if !config.raw_plans {
-            mz_transform::normalize_lets::normalize_lets(self.0)
-                .map_err(|e| ExplainError::UnknownError(e.to_string()))?;
+            normalize_lets(self.0).map_err(|e| ExplainError::UnknownError(e.to_string()))?;
         }
 
-        let plan = try_from(context, self.0)?;
-        Ok(ExplainSinglePlan { context, plan })
+        Ok(ExplainSinglePlan {
+            context,
+            plan: annotate_plan(self.0, context)?,
+        })
     }
 }
 
@@ -115,25 +116,26 @@ impl<'a> Explainable<'a, DataflowDescription<OptimizedMirRelationExpr>> {
             .iter_mut()
             .rev()
             .map(|build_desc| {
+                let plan = build_desc.plan.as_inner_mut();
+
                 // normalize the representation as linear chains
                 // (this implies !config.raw_plans by construction)
                 if config.linear_chains {
-                    enforce_linear_chains(build_desc.plan.as_inner_mut())?;
+                    enforce_linear_chains(plan)?;
                 };
                 // unless raw plans are explicitly requested
                 // normalize the representation of nested Let bindings
                 // and enforce sequential Let binding IDs
                 if !config.raw_plans {
-                    mz_transform::normalize_lets::normalize_lets(build_desc.plan.as_inner_mut())
-                        .map_err(|e| ExplainError::UnknownError(e.to_string()))?;
+                    normalize_lets(plan).map_err(|e| ExplainError::UnknownError(e.to_string()))?;
                 }
 
                 let id = context
                     .humanizer
                     .humanize_id(build_desc.id)
                     .unwrap_or_else(|| build_desc.id.to_string());
-                let plan = try_from(context, build_desc.plan.as_inner())?;
-                Ok((id, plan))
+
+                Ok((id, annotate_plan(plan, context)?))
             })
             .collect::<Result<Vec<_>, ExplainError>>()?;
 
@@ -159,79 +161,3 @@ impl<'a> Explainable<'a, DataflowDescription<OptimizedMirRelationExpr>> {
         })
     }
 }
-
-fn try_from<'a>(
-    context: &'a ExplainContext,
-    plan: &'a MirRelationExpr,
-) -> Result<AnnotatedPlan<'a, MirRelationExpr>, RecursionLimitError> {
-    let mut annotations = BTreeMap::<&MirRelationExpr, Attributes>::default();
-    let config = context.config;
-
-    if config.requires_attributes() {
-        // get the annotation keys
-        let subtree_refs = plan.post_order_vec();
-        // get the annotation values
-        let mut explain_attrs = DerivedAttributes::from(config);
-        plan.visit(&mut explain_attrs)?;
-
-        if config.subtree_size {
-            for (expr, attr) in std::iter::zip(
-                subtree_refs.iter(),
-                explain_attrs.remove_results::<SubtreeSize>().into_iter(),
-            ) {
-                let attrs = annotations.entry(expr).or_default();
-                attrs.subtree_size = Some(attr);
-            }
-        }
-        if config.non_negative {
-            for (expr, attr) in std::iter::zip(
-                subtree_refs.iter(),
-                explain_attrs.remove_results::<NonNegative>().into_iter(),
-            ) {
-                let attrs = annotations.entry(expr).or_default();
-                attrs.non_negative = Some(attr);
-            }
-        }
-
-        if config.arity {
-            for (expr, attr) in std::iter::zip(
-                subtree_refs.iter(),
-                explain_attrs.remove_results::<Arity>().into_iter(),
-            ) {
-                let attrs = annotations.entry(expr).or_default();
-                attrs.arity = Some(attr);
-            }
-        }
-
-        if config.types {
-            for (expr, types) in std::iter::zip(
-                subtree_refs.iter(),
-                explain_attrs.remove_results::<RelationType>().into_iter(),
-            ) {
-                let humanized_columns = types
-                    .into_iter()
-                    .map(|c| context.humanizer.humanize_column_type(&c))
-                    .collect::<Vec<_>>();
-                let attr = bracketed("(", ")", separated(", ", humanized_columns)).to_string();
-                let attrs = annotations.entry(expr).or_default();
-                attrs.types = Some(attr);
-            }
-        }
-
-        if config.keys {
-            for (expr, keys) in std::iter::zip(
-                subtree_refs.iter(),
-                explain_attrs.remove_results::<UniqueKeys>().into_iter(),
-            ) {
-                let formatted_keys = keys
-                    .into_iter()
-                    .map(|key_set| bracketed("[", "]", separated(", ", key_set)).to_string());
-                let attr = bracketed("(", ")", separated(", ", formatted_keys)).to_string();
-                let attrs = annotations.entry(expr).or_default();
-                attrs.keys = Some(attr);
-            }
-        }
-    }
-
-    Ok(AnnotatedPlan { plan, annotations })
-}