opt: drop provably-non-null COALESCE operands

When a leading COALESCE argument is known non-null from the input's
NotNullCols, later arguments aren't needed. Uses NotNullCols to trim
leading COALESCE args in Project and Select.

Example: `SELECT COALESCE(i, j) FROM ij` with `i NOT NULL` no longer
keeps the coalesce in the plan.

Implementation details:
- Add SimplifyCoalesceProject and SimplifyCoalesceSelect norm rules that
  call SimplifyCoalesceInScalar on each projection/filter using the
  input's not-null columns.
- Introduce scalarContainsSimplifiableCoalesce (in scalar_funcs.go) as
  the match guard; it walks the full scalar tree so that nested Coalesce
  expressions like COALESCE(a, COALESCE(b, c)) are not missed when only
  the inner Coalesce is simplifiable.
- CanSimplifyCoalesce mirrors simplifyCoalesce's logic with a direct
  scan, avoiding expression construction on the hot path. The final
  return is true (not false) because when the loop exhausts all-null
  constant args, simplifyCoalesce returns args[last] unwrapped, which
  is a genuine simplification.
- SimplifyCoalesceInScalar recurses into the simplified result via
  c.f.Replace so nested Coalesces are handled in a single pass.
- The private simplifyCoalesce helper iterates over all args (including
  the last) so ExprIsNeverNull can fire at any position.

Test plan: `bazel test //pkg/sql/opt/norm:norm_test --test_filter=TestNormRules`

Fixes #103596

Release note: None

Author: lohitkolluri

pkg/sql/opt/norm/project_funcs.go

@@ -1008,3 +1008,34 @@ func (c *CustomFuncs) HasVolatileProjection(projections memo.ProjectionsExpr) bo
 	}
 	return false
 }
+
+// CanSimplifyCoalesceInProjections returns true if any projection contains a
+// Coalesce expression that can be simplified using the input's not-null columns.
+func (c *CustomFuncs) CanSimplifyCoalesceInProjections(
+	projections memo.ProjectionsExpr, input memo.RelExpr,
+) bool {
+	notNullCols := c.NotNullCols(input)
+	for i := range projections {
+		if c.scalarContainsSimplifiableCoalesce(projections[i].Element, notNullCols) {
+			return true
+		}
+	}
+	return false
+}
+
+// SimplifyCoalesceInProjections simplifies Coalesce expressions in projections
+// using the input's not-null columns.
+func (c *CustomFuncs) SimplifyCoalesceInProjections(
+	projections memo.ProjectionsExpr, input memo.RelExpr,
+) memo.ProjectionsExpr {
+	notNullCols := c.NotNullCols(input)
+	newProjections := make(memo.ProjectionsExpr, len(projections))
+	for i := range projections {
+		p := &projections[i]
+		newProjections[i] = c.f.ConstructProjectionsItem(
+			c.SimplifyCoalesceInScalar(p.Element, notNullCols),
+			p.Col,
+		)
+	}
+	return newProjections
+}

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

@@ -331,6 +331,23 @@ $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:* &
+        (CanSimplifyCoalesceInProjections $projections $input)
+    $passthrough:*
+)
+=>
+(Project
+    $input
+    (SimplifyCoalesceInProjections $projections $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,17 @@
 =>
 (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.
+[SimplifyCoalesceSelect, Normalize, LowPriority]
+(Select
+    $input:*
+    $filters:* & (CanSimplifyCoalesceInFilters $filters $input)
+)
+=>
+(Select $input (SimplifyCoalesceInFilters $filters $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)
+}
+
+// scalarContainsSimplifiableCoalesce 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) scalarContainsSimplifiableCoalesce(
+	e opt.ScalarExpr, notNullCols opt.ColSet,
+) bool {
+	found := false
+	var replace ReplaceFunc
+	replace = func(e opt.Expr) opt.Expr {
+		if co, ok := e.(*memo.CoalesceExpr); ok {
+			if c.CanSimplifyCoalesce(co.Args, notNullCols) {
+				found = true
+				return co
+			}
+			// Even if the outer Coalesce is not simplifiable, recurse into its
+			// args to find simplifiable nested Coalesce expressions.
+			return c.f.Replace(co, replace)
+		}
+		return c.f.Replace(e, replace)
+	}
+	replace(e)
+	return found
+}
+
 // 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,33 @@ func (c *CustomFuncs) addConjuncts(
 func (c *CustomFuncs) ForDuplicateRemoval(private *memo.OrdinalityPrivate) (ok bool) {
 	return private.ForDuplicateRemoval
 }
+
+// CanSimplifyCoalesceInFilters returns true if any filter condition contains a
+// Coalesce expression that can be simplified using the input's not-null columns.
+func (c *CustomFuncs) CanSimplifyCoalesceInFilters(
+	filters memo.FiltersExpr, input memo.RelExpr,
+) bool {
+	notNullCols := c.NotNullCols(input)
+	for i := range filters {
+		if c.scalarContainsSimplifiableCoalesce(filters[i].Condition, notNullCols) {
+			return true
+		}
+	}
+	return false
+}
+
+// SimplifyCoalesceInFilters simplifies Coalesce expressions in filter conditions
+// using the input's not-null columns.
+func (c *CustomFuncs) SimplifyCoalesceInFilters(
+	filters memo.FiltersExpr, input memo.RelExpr,
+) memo.FiltersExpr {
+	notNullCols := c.NotNullCols(input)
+	newFilters := make(memo.FiltersExpr, len(filters))
+	for i := range filters {
+		newFilters[i] = c.f.ConstructFiltersItem(
+			c.SimplifyCoalesceInScalar(filters[i].Condition, notNullCols),
+		)
+	}
+	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
 (

pkg/sql/opt/norm/testdata/rules/scalar

@@ -2,6 +2,10 @@ exec-ddl
 CREATE TABLE a (k INT PRIMARY KEY, i INT, f FLOAT, s STRING, arr int[])
 ----
 
+exec-ddl
+CREATE TABLE ij (i INT NOT NULL, j INT)
+----
+
 exec-ddl
 CREATE TABLE xy (x INT PRIMARY KEY, y INT)
 ----
@@ -172,7 +176,9 @@ project
  └── projections
       └── COALESCE(s:4, s:4 || 'foo') [as=coalesce:8, outer=(4), immutable]
 
-# Trailing null can't be removed.
+# Trailing null can't be removed when leading arg is a non-constant (nullable
+# variable): the original SimplifyCoalesce rule only fires when the leading arg
+# is a constant, so no folding happens here.
 norm
 SELECT COALESCE(i, NULL, NULL) FROM a
 ----
@@ -183,6 +189,44 @@ project
  └── projections
       └── COALESCE(i:2, CAST(NULL AS INT8), CAST(NULL AS INT8)) [as=coalesce:8, outer=(2)]
 
+# All-null COALESCE collapses to a single null via SimplifyCoalesce (the scalar
+# rule fires when the leading arg is constant). CanSimplifyCoalesce correctly
+# returns true when the loop exhausts all-null args, matching simplifyCoalesce's
+# behaviour of returning args[last] (unwrapping the COALESCE).
+norm expect=SimplifyCoalesce
+SELECT COALESCE(NULL::INT, NULL::INT) FROM a
+----
+project
+ ├── columns: coalesce:8
+ ├── fd: ()-->(8)
+ ├── scan a
+ └── projections
+      └── CAST(NULL AS INT8) [as=coalesce:8]
+
+# --------------------------------------------------
+# SimplifyCoalesceProject
+# --------------------------------------------------
+
+norm expect=SimplifyCoalesceProject
+SELECT COALESCE(i, j) FROM ij
+----
+project
+ ├── columns: coalesce:6!null
+ ├── scan ij
+ │    └── columns: i:1!null
+ └── projections
+      └── i:1 [as=coalesce:6, outer=(1)]
+
+norm expect=SimplifyCoalesceProject
+SELECT COALESCE(i, NULL, NULL) FROM ij
+----
+project
+ ├── columns: coalesce:6!null
+ ├── scan ij
+ │    └── columns: i:1!null
+ └── projections
+      └── i:1 [as=coalesce:6, outer=(1)]
+
 norm expect=SimplifyCoalesce
 SELECT COALESCE((1, 2, 3), (2, 3, 4)) FROM a
 ----
@@ -193,6 +237,69 @@ project
  └── projections
       └── (1, 2, 3) [as=coalesce:8]
 
+# Negative test: should not fire when no arg is provably non-null.
+norm expect-not=SimplifyCoalesceProject
+SELECT COALESCE(i, f) FROM a
+----
+project
+ ├── columns: coalesce:8
+ ├── immutable
+ ├── scan a
+ │    └── columns: i:2 f:3
+ └── projections
+      └── COALESCE(i:2::FLOAT8, f:3) [as=coalesce:8, outer=(2,3), immutable]
+
+# Middle argument is provably non-null: leading NULLs are stripped, then the
+# NOT NULL arg is returned directly.
+norm expect=SimplifyCoalesceProject
+SELECT COALESCE(NULL::INT, i, j) FROM ij
+----
+project
+ ├── columns: coalesce:6!null
+ ├── scan ij
+ │    └── columns: i:1!null
+ └── projections
+      └── i:1 [as=coalesce:6, outer=(1)]
+
+# Nested COALESCE: outer is not directly simplifiable, but the inner
+# COALESCE(i, j) can be simplified to i (since i is NOT NULL). The fix
+# ensures the inner Coalesce is not missed.
+norm expect=SimplifyCoalesceProject
+SELECT COALESCE(j, COALESCE(i, j)) FROM ij
+----
+project
+ ├── columns: coalesce:6
+ ├── scan ij
+ │    └── columns: i:1!null j:2
+ └── projections
+      └── COALESCE(j:2, i:1) [as=coalesce:6, outer=(1,2)]
+
+# COALESCE inside an IF expression: the inner nested Coalesce is simplified
+# via SimplifyCoalesceProject even though the outer Coalesce is not itself
+# directly simplifiable.
+norm expect=SimplifyCoalesceProject
+SELECT IF(j > 0, COALESCE(j, COALESCE(i, j)), j) FROM ij
+----
+project
+ ├── columns: if:6
+ ├── scan ij
+ │    └── columns: i:1!null j:2
+ └── projections
+      └── CASE j:2 > 0 WHEN true THEN COALESCE(j:2, i:1) ELSE j:2 END [as=if:6, outer=(1,2)]
+
+# SimplifyCoalesceProject does not fire when the first arg is a non-constant
+# nullable variable: CanSimplifyCoalesce exits early with false because i is
+# not a constant, and no leading nulls were stripped (i > 0 = false).
+norm expect-not=SimplifyCoalesceProject
+SELECT COALESCE(i, NULL, NULL) FROM a
+----
+project
+ ├── columns: coalesce:8
+ ├── scan a
+ │    └── columns: i:2
+ └── projections
+      └── COALESCE(i:2, CAST(NULL AS INT8), CAST(NULL AS INT8)) [as=coalesce:8, outer=(2)]
+
 
 # --------------------------------------------------
 # EliminateCast