opt: simplify COALESCE using NOT NULL column derivation

pkg/sql/opt/norm/project_funcs.go

@@ -1008,3 +1008,23 @@ func (c *CustomFuncs) HasVolatileProjection(projections memo.ProjectionsExpr) bo
 	}
 	return false
 }
+
+// SimplifyCoalesceInProjections simplifies Coalesce expressions in projections
+// using the given not-null columns. Projections whose element is unchanged are
+// reused as-is to avoid unnecessary memo invalidation.
+func (c *CustomFuncs) SimplifyCoalesceInProjections(
+	projections memo.ProjectionsExpr, notNullCols opt.ColSet,
+) memo.ProjectionsExpr {
+	newProjections := make(memo.ProjectionsExpr, len(projections))
+	for i := range projections {
+		p := &projections[i]
+		simplified := c.SimplifyCoalesceInScalar(p.Element, notNullCols)
+		if simplified == p.Element {
+			// No change; reuse the original ProjectionsItem.
+			newProjections[i] = *p
+		} else {
+			newProjections[i] = c.f.ConstructProjectionsItem(simplified, p.Col)
+		}
+	}
+	return newProjections
+}

pkg/sql/opt/norm/rules/project.opt

@@ -331,6 +331,27 @@ $input
     $passthrough
 )
 
+# SimplifyCoalesceProject simplifies Coalesce projections by eliminating leading
+# operands that are guaranteed to be non-null according to the input's
+# relational properties.
+[SimplifyCoalesceProject, Normalize, LowPriority]
+(Project
+    $input:*
+    $projections:[
+            ...
+            (ProjectionsItem $scalar:*) &
+                (CanSimplifyCoalesceInScalar $scalar (NotNullCols $input))
+            ...
+        ]
+    $passthrough:*
+)
+=>
+(Project
+    $input
+    (SimplifyCoalesceInProjections $projections (NotNullCols $input))
+    $passthrough
+)
+
 # FoldIsNullProject folds "x IS NULL" projections to false if "x" is not null in
 # the Project's input. It matches if there is at least one projection that can
 # be folded, and it replaces all projections that can be folded.

pkg/sql/opt/norm/rules/select.opt

@@ -31,6 +31,23 @@
 =>
 (Select $input (SimplifyFilters $filters))
 
+# SimplifyCoalesceSelect simplifies Coalesce expressions in Select filters by
+# eliminating leading operands that are guaranteed to be non-null according to
+# the input's relational properties and null-rejecting columns from sibling
+# filter conditions.
+[SimplifyCoalesceSelect, Normalize, LowPriority]
+(Select
+    $input:*
+    $filters:[
+            ...
+            (FiltersItem $cond:*) &
+                (CanSimplifyCoalesceInFilters $filters $cond (NotNullCols $input))
+            ...
+        ]
+)
+=>
+(Select $input (SimplifyCoalesceInFilters $filters (NotNullCols $input)))
+
 # ConsolidateSelectFilters consolidates filters that constrain a single
 # variable. For example, filters x >= 5 and x <= 10 would be combined into a
 # single Range operation.

pkg/sql/opt/norm/scalar_funcs.go

@@ -74,25 +74,109 @@ func (c *CustomFuncs) ConstructSortedUniqueList(
 // SimplifyCoalesce discards any leading null operands, and then if the next
 // operand is a constant, replaces with that constant.
 func (c *CustomFuncs) SimplifyCoalesce(args memo.ScalarListExpr) opt.ScalarExpr {
-	for i := 0; i < len(args)-1; i++ {
-		item := args[i]
+	return c.simplifyCoalesce(args, opt.ColSet{})
+}
+
+func (c *CustomFuncs) simplifyCoalesce(
+	args memo.ScalarListExpr, notNullCols opt.ColSet,
+) opt.ScalarExpr {
+	// Iterate over all args (including the last) so that ExprIsNeverNull can
+	// fire for any position: once a provably-non-null arg is found, COALESCE
+	// will always return that arg, so we can replace the whole expression.
+	for i := 0; i < len(args); i++ {
+		if memo.ExprIsNeverNull(args[i], notNullCols) {
+			return args[i]
+		}
 
-		// If item is not a constant value, then its value may turn out to be
-		// null, so no more folding. Return operands from then on.
-		if !c.IsConstValueOrGroupOfConstValues(item) {
+		if !c.IsConstValueOrGroupOfConstValues(args[i]) {
+			if i >= len(args)-1 {
+				return args[i]
+			}
 			return c.f.ConstructCoalesce(args[i:])
 		}
 
-		if item.Op() != opt.NullOp {
-			return item
+		if args[i].Op() != opt.NullOp {
+			return args[i]
 		}
 	}
 
-	// All operands up to the last were null (or the last is the only operand),
-	// so return the last operand without the wrapping COALESCE function.
 	return args[len(args)-1]
 }
 
+// CanSimplifyCoalesce returns true if simplifyCoalesce would change the given
+// Coalesce expression. It mirrors the logic of simplifyCoalesce directly to
+// avoid the cost of constructing a new expression just to compare arg counts.
+func (c *CustomFuncs) CanSimplifyCoalesce(args memo.ScalarListExpr, notNullCols opt.ColSet) bool {
+	for i, arg := range args {
+		if memo.ExprIsNeverNull(arg, notNullCols) {
+			return true
+		}
+		if !c.IsConstValueOrGroupOfConstValues(arg) {
+			// Non-constant arg at position i blocks further simplification. The
+			// expression changes only if leading nulls (i > 0) were already
+			// stripped.
+			return i > 0
+		}
+		if arg.Op() != opt.NullOp {
+			// Non-null constant: simplifyCoalesce returns it directly.
+			return true
+		}
+		// NullOp constant: will be stripped; continue to next arg.
+	}
+	// The loop exhausted all args, meaning every arg was a null constant.
+	// simplifyCoalesce returns args[last] directly (unwrapping the COALESCE),
+	// which is a simplification whenever there is more than one arg. The
+	// single-arg case is already handled by EliminateCoalesce before this
+	// function is called, so we can safely return true here.
+	return true
+}
+
+// SimplifyCoalesceInScalar recursively simplifies Coalesce expressions in the
+// given scalar expression tree using the provided not-null columns. It handles
+// nested Coalesce expressions by recursing into the simplified result.
+func (c *CustomFuncs) SimplifyCoalesceInScalar(
+	e opt.ScalarExpr, notNullCols opt.ColSet,
+) opt.ScalarExpr {
+	var replace ReplaceFunc
+	replace = func(e opt.Expr) opt.Expr {
+		if co, ok := e.(*memo.CoalesceExpr); ok {
+			simplified := c.simplifyCoalesce(co.Args, notNullCols)
+			// Recurse into the simplified result so that nested Coalesce
+			// expressions (e.g. COALESCE(a, COALESCE(b, c))) are also handled
+			// when the outer simplification does not fire.
+			return c.f.Replace(simplified, replace)
+		}
+		return c.f.Replace(e, replace)
+	}
+	return replace(e).(opt.ScalarExpr)
+}
+
+// CanSimplifyCoalesceInScalar returns true if the scalar expression tree
+// contains any Coalesce expression that can be simplified using notNullCols. It
+// recurses into Coalesce args so that nested Coalesce expressions are not
+// missed when the outer Coalesce is not itself simplifiable.
+func (c *CustomFuncs) CanSimplifyCoalesceInScalar(e opt.ScalarExpr, notNullCols opt.ColSet) bool {
+	if co, ok := e.(*memo.CoalesceExpr); ok {
+		if c.CanSimplifyCoalesce(co.Args, notNullCols) {
+			return true
+		}
+		for _, arg := range co.Args {
+			if c.CanSimplifyCoalesceInScalar(arg, notNullCols) {
+				return true
+			}
+		}
+		return false
+	}
+	for i, n := 0, e.ChildCount(); i < n; i++ {
+		if sc, ok := e.Child(i).(opt.ScalarExpr); ok {
+			if c.CanSimplifyCoalesceInScalar(sc, notNullCols) {
+				return true
+			}
+		}
+	}
+	return false
+}
+
 // IsConstValueEqual returns whether const1 and const2 are equal.
 func (c *CustomFuncs) IsConstValueEqual(const1, const2 opt.ScalarExpr) bool {
 	op1 := const1.Op()

pkg/sql/opt/norm/select_funcs.go

@@ -419,3 +419,71 @@ func (c *CustomFuncs) addConjuncts(
 func (c *CustomFuncs) ForDuplicateRemoval(private *memo.OrdinalityPrivate) (ok bool) {
 	return private.ForDuplicateRemoval
 }
+
+// computeFilterNotNullCols returns per-filter NOT NULL column sets and their
+// total union for the given filters. It extracts null-rejecting columns from
+// each filter item's constraints. Used by both CanSimplifyCoalesceInFilters and
+// SimplifyCoalesceInFilters to avoid duplicating constraint extraction.
+func computeFilterNotNullCols(
+	c *CustomFuncs, filters memo.FiltersExpr,
+) (perFilter []opt.ColSet, total opt.ColSet) {
+	perFilter = make([]opt.ColSet, len(filters))
+	for i := range filters {
+		constraints := filters[i].ScalarProps().Constraints
+		if constraints != nil {
+			constraints.ExtractNotNullCols(c.f.ctx, c.f.evalCtx, &perFilter[i])
+		}
+	}
+	for _, fc := range perFilter {
+		total = total.Union(fc)
+	}
+	return
+}
+
+// CanSimplifyCoalesceInFilters returns true if any filter condition contains a
+// Coalesce expression that can be simplified using the input's NOT NULL columns
+// plus null-rejecting columns from other filter conditions in the same set.
+// The _ opt.ScalarExpr parameter is unused; it is required by the optgen rule
+// which binds $cond but this function works with the full $filters set.
+func (c *CustomFuncs) CanSimplifyCoalesceInFilters(
+	filters memo.FiltersExpr, _ opt.ScalarExpr, notNullCols opt.ColSet,
+) bool {
+	perFilter, total := computeFilterNotNullCols(c, filters)
+	for i := range filters {
+		diff := total.Difference(perFilter[i])
+		enriched := notNullCols
+		if !diff.Empty() {
+			enriched = notNullCols.Union(diff)
+		}
+		if c.CanSimplifyCoalesceInScalar(filters[i].Condition, enriched) {
+			return true
+		}
+	}
+	return false
+}
+
+// SimplifyCoalesceInFilters simplifies Coalesce expressions in filter
+// conditions using the given not-null columns, including null-rejecting columns
+// derived from other filter conditions. Filters whose condition is unchanged
+// are reused as-is to avoid unnecessary memo invalidation.
+func (c *CustomFuncs) SimplifyCoalesceInFilters(
+	filters memo.FiltersExpr, notNullCols opt.ColSet,
+) memo.FiltersExpr {
+	perFilter, total := computeFilterNotNullCols(c, filters)
+	newFilters := make(memo.FiltersExpr, len(filters))
+	for i := range filters {
+		f := &filters[i]
+		diff := total.Difference(perFilter[i])
+		enriched := notNullCols
+		if !diff.Empty() {
+			enriched = notNullCols.Union(diff)
+		}
+		simplified := c.SimplifyCoalesceInScalar(f.Condition, enriched)
+		if simplified == f.Condition {
+			newFilters[i] = *f
+		} else {
+			newFilters[i] = c.f.ConstructFiltersItem(simplified)
+		}
+	}
+	return newFilters
+}

pkg/sql/opt/norm/testdata/rules/decorrelate

@@ -3006,7 +3006,7 @@ group-by (hash)
            └── y:2
 
 # Right input of SemiJoin is Project.
-norm expect=TryDecorrelateSemiJoin
+norm expect=TryDecorrelateSemiJoin disable=SimplifyCoalesceProject
 SELECT k FROM a
 WHERE EXISTS
 (