opt: return RelExpr from (*memo).RootExpr

`(*memo).RootExpr` now returns a `RelExpr` instead of an `Expr`, helping
eliminate some type assertions.

Release note: None

Author: mgartner

pkg/sql/opt/exec/execbuilder/scalar.go

@@ -675,7 +675,7 @@ func (b *Builder) buildExistsSubquery(
 		stmtProps := []*physical.Required{{Presentation: physical.Presentation{aliasedCol}}}
 
 		// Create an wrapRootExprFn that wraps input in a Limit and a Project.
-		wrapRootExpr := func(f *norm.Factory, e memo.RelExpr) opt.Expr {
+		wrapRootExpr := func(f *norm.Factory, e memo.RelExpr) memo.RelExpr {
 			return f.ConstructProject(
 				f.ConstructLimit(
 					e,
@@ -1115,7 +1115,7 @@ func (b *Builder) initRoutineExceptionHandler(
 	blockState.ExceptionHandler = exceptionHandler
 }
 
-type wrapRootExprFn func(f *norm.Factory, e memo.RelExpr) opt.Expr
+type wrapRootExprFn func(f *norm.Factory, e memo.RelExpr) memo.RelExpr
 
 // buildRoutinePlanGenerator returns a tree.RoutinePlanFn that can plan the
 // statements in a routine that has one or more arguments.
@@ -1249,7 +1249,7 @@ func (b *Builder) buildRoutinePlanGenerator(
 			f.CopyAndReplace(originalMemo, stmt, props, replaceFn)
 
 			if wrapRootExpr != nil {
-				wrapped := wrapRootExpr(f, f.Memo().RootExpr().(memo.RelExpr)).(memo.RelExpr)
+				wrapped := wrapRootExpr(f, f.Memo().RootExpr())
 				f.Memo().SetRoot(wrapped, props)
 			}
 

pkg/sql/opt/memo/memo.go

@@ -132,7 +132,7 @@ type Memo struct {
 	// rootExpr is the root expression of the memo expression forest. It is set
 	// via a call to SetRoot. After optimization, it is set to be the root of the
 	// lowest cost tree in the forest.
-	rootExpr opt.Expr
+	rootExpr RelExpr
 
 	// rootProps are the physical properties required of the root memo expression.
 	// It is set via a call to SetRoot.
@@ -366,7 +366,7 @@ func (m *Memo) Metadata() *opt.Metadata {
 
 // RootExpr returns the root memo expression previously set via a call to
 // SetRoot.
-func (m *Memo) RootExpr() opt.Expr {
+func (m *Memo) RootExpr() RelExpr {
 	return m.rootExpr
 }
 
@@ -395,12 +395,7 @@ func (m *Memo) SetRoot(e RelExpr, phys *physical.Required) {
 // HasPlaceholders returns true if the memo contains at least one placeholder
 // operator.
 func (m *Memo) HasPlaceholders() bool {
-	rel, ok := m.rootExpr.(RelExpr)
-	if !ok {
-		panic(errors.AssertionFailedf("placeholders only supported when memo root is relational"))
-	}
-
-	return rel.Relational().HasPlaceholder
+	return m.rootExpr.Relational().HasPlaceholder
 }
 
 // IsStale returns true if the memo has been invalidated by changes to any of
@@ -555,8 +550,7 @@ func (m *Memo) ResetCost(e RelExpr, cost Cost) {
 func (m *Memo) IsOptimized() bool {
 	// The memo is optimized once the root expression has its physical properties
 	// assigned.
-	rel, ok := m.rootExpr.(RelExpr)
-	return ok && rel.RequiredPhysical() != nil
+	return m.rootExpr != nil && m.rootExpr.RequiredPhysical() != nil
 }
 
 // OptimizationCost returns a rough estimate of the cost of optimization of the

pkg/sql/opt/memo/memo_test.go

@@ -691,15 +691,15 @@ func TestStatsAvailable(t *testing.T) {
 
 	// Stats should not be available for any expression.
 	opttestutils.BuildQuery(t, &o, catalog, &evalCtx, "SELECT * FROM t WHERE a=1")
-	testNotAvailable(o.Memo().RootExpr().(memo.RelExpr))
+	testNotAvailable(o.Memo().RootExpr())
 
 	opttestutils.BuildQuery(t, &o, catalog, &evalCtx, "SELECT sum(a), b FROM t GROUP BY b")
-	testNotAvailable(o.Memo().RootExpr().(memo.RelExpr))
+	testNotAvailable(o.Memo().RootExpr())
 
 	opttestutils.BuildQuery(t, &o, catalog, &evalCtx,
 		"SELECT * FROM t AS t1, t AS t2 WHERE t1.a = t2.a AND t1.b = 5",
 	)
-	testNotAvailable(o.Memo().RootExpr().(memo.RelExpr))
+	testNotAvailable(o.Memo().RootExpr())
 
 	if _, err := catalog.ExecuteDDL(
 		`ALTER TABLE t INJECT STATISTICS '[
@@ -729,15 +729,15 @@ func TestStatsAvailable(t *testing.T) {
 
 	// Stats should be available for all expressions.
 	opttestutils.BuildQuery(t, &o, catalog, &evalCtx, "SELECT * FROM t WHERE a=1")
-	testAvailable(o.Memo().RootExpr().(memo.RelExpr))
+	testAvailable(o.Memo().RootExpr())
 
 	opttestutils.BuildQuery(t, &o, catalog, &evalCtx, "SELECT sum(a), b FROM t GROUP BY b")
-	testAvailable(o.Memo().RootExpr().(memo.RelExpr))
+	testAvailable(o.Memo().RootExpr())
 
 	opttestutils.BuildQuery(t, &o, catalog, &evalCtx,
 		"SELECT * FROM t AS t1, t AS t2 WHERE t1.a = t2.a AND t1.b = 5",
 	)
-	testAvailable(o.Memo().RootExpr().(memo.RelExpr))
+	testAvailable(o.Memo().RootExpr())
 }
 
 // traverseExpr is a helper function to recursively traverse a relational

pkg/sql/opt/norm/factory.go

@@ -420,7 +420,7 @@ func (f *Factory) AssignPlaceholders(from *memo.Memo) (err error) {
 		}
 		return f.CopyAndReplaceDefault(e, replaceFn)
 	}
-	f.CopyAndReplace(from, from.RootExpr().(memo.RelExpr), from.RootProps(), replaceFn)
+	f.CopyAndReplace(from, from.RootExpr(), from.RootProps(), replaceFn)
 
 	return nil
 }

pkg/sql/opt/norm/factory_test.go

@@ -101,7 +101,7 @@ func TestCopyAndReplace(t *testing.T) {
 		}
 		return o.Factory().CopyAndReplaceDefault(e, replaceFn)
 	}
-	o.Factory().CopyAndReplace(m, m.RootExpr().(memo.RelExpr), m.RootProps(), replaceFn)
+	o.Factory().CopyAndReplace(m, m.RootExpr(), m.RootProps(), replaceFn)
 
 	if e, err := o.Optimize(); err != nil {
 		t.Fatal(err)
@@ -146,7 +146,7 @@ func TestCopyAndReplaceWithScan(t *testing.T) {
 			replaceFn = func(e opt.Expr) opt.Expr {
 				return o.Factory().CopyAndReplaceDefault(e, replaceFn)
 			}
-			o.Factory().CopyAndReplace(m, m.RootExpr().(memo.RelExpr), m.RootProps(), replaceFn)
+			o.Factory().CopyAndReplace(m, m.RootExpr(), m.RootProps(), replaceFn)
 
 			if _, err := o.Optimize(); err != nil {
 				t.Fatal(err)

pkg/sql/opt/xform/optimizer.go

@@ -261,7 +261,7 @@ func (o *Optimizer) Optimize() (_ opt.Expr, err error) {
 	o.optimizeRootWithProps()
 
 	// Now optimize the entire expression tree.
-	root := o.mem.RootExpr().(memo.RelExpr)
+	root := o.mem.RootExpr()
 	rootProps := o.mem.RootProps()
 	o.optimizeGroup(root, rootProps)
 
@@ -922,10 +922,7 @@ func (o *Optimizer) ensureOptState(grp memo.RelExpr, required *physical.Required
 // properties required of it. This may trigger the creation of a new root and
 // new properties.
 func (o *Optimizer) optimizeRootWithProps() {
-	root, ok := o.mem.RootExpr().(memo.RelExpr)
-	if !ok {
-		panic(errors.AssertionFailedf("Optimize can only be called on relational root expressions"))
-	}
+	root := o.mem.RootExpr()
 	rootProps := o.mem.RootProps()
 
 	// [SimplifyRootOrdering]

pkg/sql/opt/xform/optimizer_test.go

@@ -113,7 +113,7 @@ func TestDetachMemoRace(t *testing.T) {
 			// Rewrite the filter to use a different column, which will trigger creation
 			// of new table statistics. If the statistics object is aliased, this will
 			// be racy.
-			f.CopyAndReplace(mem, mem.RootExpr().(memo.RelExpr), mem.RootProps(), replaceFn)
+			f.CopyAndReplace(mem, mem.RootExpr(), mem.RootProps(), replaceFn)
 			wg.Done()
 		}()
 	}

pkg/sql/opt/xform/placeholder_fast_path.go

@@ -45,7 +45,7 @@ func (o *Optimizer) TryPlaceholderFastPath() (ok bool, err error) {
 		}
 	}()
 
-	root := o.mem.RootExpr().(memo.RelExpr)
+	root := o.mem.RootExpr()
 
 	rootRelProps := root.Relational()
 	// We are dealing with a memo that still contains placeholders. The statistics

pkg/sql/plan_opt.go

@@ -590,7 +590,7 @@ func (opc *optPlanningCtx) reuseMemo(cachedMemo *memo.Memo) (*memo.Memo, error)
 	opc.flags.Set(planFlagOptimized)
 	mem := f.Memo()
 	if prep := opc.p.stmt.Prepared; opc.allowMemoReuse && prep != nil {
-		costWithOptimizationCost := mem.RootExpr().(memo.RelExpr).Cost()
+		costWithOptimizationCost := mem.RootExpr().Cost()
 		costWithOptimizationCost.Add(mem.OptimizationCost())
 		prep.Costs.AddCustom(costWithOptimizationCost)
 	}
@@ -772,7 +772,7 @@ func (opc *optPlanningCtx) fetchPreparedMemo(ctx context.Context) (_ *memo.Memo,
 			prep.IdealGenericPlan = true
 		case memoTypeGeneric:
 			prep.GenericMemo = newMemo
-			prep.Costs.SetGeneric(newMemo.RootExpr().(memo.RelExpr).Cost())
+			prep.Costs.SetGeneric(newMemo.RootExpr().Cost())
 			// Now that the cost of the generic plan is known, we need to
 			// re-evaluate the decision to use a generic or custom plan.
 			if !opc.chooseGenericPlan() {
@@ -961,11 +961,11 @@ func (opc *optPlanningCtx) runExecBuilder(
 	if opc.gf.Initialized() {
 		planTop.instrumentation.planGist = opc.gf.PlanGist()
 	}
-	planTop.instrumentation.costEstimate = mem.RootExpr().(memo.RelExpr).Cost().C
-	available := mem.RootExpr().(memo.RelExpr).Relational().Statistics().Available
+	planTop.instrumentation.costEstimate = mem.RootExpr().Cost().C
+	available := mem.RootExpr().Relational().Statistics().Available
 	planTop.instrumentation.statsAvailable = available
 	if available {
-		planTop.instrumentation.outputRows = mem.RootExpr().(memo.RelExpr).Relational().Statistics().RowCount
+		planTop.instrumentation.outputRows = mem.RootExpr().Relational().Statistics().RowCount
 	}
 
 	if stmt.ExpectedTypes != nil {
@@ -1045,7 +1045,7 @@ func (opc *optPlanningCtx) makeQueryIndexRecommendation(
 	f.FoldingControl().AllowStableFolds()
 	f.CopyAndReplace(
 		savedMemo,
-		savedMemo.RootExpr().(memo.RelExpr),
+		savedMemo.RootExpr(),
 		savedMemo.RootProps(),
 		f.CopyWithoutAssigningPlaceholders,
 	)
@@ -1066,7 +1066,7 @@ func (opc *optPlanningCtx) makeQueryIndexRecommendation(
 	opc.optimizer.Init(ctx, f.EvalContext(), opc.catalog)
 	f.CopyAndReplace(
 		savedMemo,
-		savedMemo.RootExpr().(memo.RelExpr),
+		savedMemo.RootExpr(),
 		savedMemo.RootProps(),
 		f.CopyWithoutAssigningPlaceholders,
 	)
@@ -1091,7 +1091,7 @@ func (opc *optPlanningCtx) makeQueryIndexRecommendation(
 	savedMemo.Metadata().UpdateTableMeta(origCtx, f.EvalContext(), optTables)
 	f.CopyAndReplace(
 		savedMemo,
-		savedMemo.RootExpr().(memo.RelExpr),
+		savedMemo.RootExpr(),
 		savedMemo.RootProps(),
 		f.CopyWithoutAssigningPlaceholders,
 	)

pkg/sql/reference_provider.go

@@ -132,7 +132,7 @@ func (f *referenceProviderFactory) NewReferenceProvider(
 	case *memo.CreateTriggerExpr:
 		planDeps, typeDeps, funcDeps, err = toPlanDependencies(t.Deps, t.TypeDeps, t.FuncDeps)
 	default:
-		return nil, errors.AssertionFailedf("unexpected root expression: %s", t.(memo.RelExpr).Op())
+		return nil, errors.AssertionFailedf("unexpected root expression: %s", t.Op())
 	}
 	if err != nil {
 		return nil, err