Merge #62924

62924: opt: add FoldFunctionWithNullArg normalization rule r=mgartner a=mgartner

This commit adds the `FoldFunctionWithNullArg` normalization rule which
folds a function that does not allow Null arguments to Null when one of
the arguments is Null. See the documentation for the new rule for more
details.

Release note (performance improvement): The optimizer now folds
functions to `NULL` when the function does not allow `NULL` arguments
and one of the arguments is a `NULL` constant. As a result, more
efficient query plans will be produced for queries with these types of
function calls.


Co-authored-by: Marcus Gartner <marcus@cockroachlabs.com>

pkg/sql/opt/memo/testdata/stats/join

@@ -1497,7 +1497,6 @@ FROM
 full-join (cross)
  ├── columns: a:1(int) b:2(int) c:3(int) a:5(int) b:6(int) c:7(int)
  ├── multiplicity: left-rows(exactly-one), right-rows(one-or-more)
- ├── immutable
  ├── stats: [rows=100]
  ├── key: (1,2)
  ├── fd: (1,2)-->(3,5-7)
@@ -1513,20 +1512,7 @@ full-join (cross)
  │    ├── key: ()
  │    └── fd: ()-->(5-7)
  └── filters
-      └── is [type=bool, immutable, subquery]
-           ├── function: not_like_escape [type=bool]
-           │    ├── '' [type=string]
-           │    ├── CAST(NULL AS STRING) [type=string]
-           │    └── cast: STRING [type=string]
-           │         └── subquery [type=unknown]
-           │              └── values
-           │                   ├── columns: "?column?":9(unknown)
-           │                   ├── cardinality: [1 - 1]
-           │                   ├── stats: [rows=1]
-           │                   ├── key: ()
-           │                   ├── fd: ()-->(9)
-           │                   └── (NULL,) [type=tuple{unknown}]
-           └── false [type=bool]
+      └── false [type=bool, constraints=(contradiction; tight)]
 
 expr
 (SemiJoin

pkg/sql/opt/norm/fold_constants_funcs.go

@@ -541,6 +541,34 @@ func (c *CustomFuncs) FoldColumnAccess(input opt.ScalarExpr, idx memo.TupleOrdin
 	return nil
 }
 
+// CanFoldFunctionWithNullArg returns true if the given function can be folded
+// to Null when any of its arguments are Null. A function can be folded to Null
+// in this case if all of the following are true:
+//
+//   1. It does not allow Null arguments (NullableArgs=false).
+//   2. It is a normal function, not an aggregate, window, or generator.
+//
+// See FoldFunctionWithNullArg for more details.
+func (c *CustomFuncs) CanFoldFunctionWithNullArg(private *memo.FunctionPrivate) bool {
+	return !private.Properties.NullableArgs &&
+		private.Properties.Class == tree.NormalClass
+}
+
+// HasNullArg returns true if one of args is Null.
+func (c *CustomFuncs) HasNullArg(args memo.ScalarListExpr) bool {
+	for i := range args {
+		if args[i].Op() == opt.NullOp {
+			return true
+		}
+	}
+	return false
+}
+
+// FunctionReturnType returns the return type of the given function.
+func (c *CustomFuncs) FunctionReturnType(private *memo.FunctionPrivate) *types.T {
+	return private.Typ
+}
+
 // FoldFunction evaluates a function expression with constant inputs. It
 // returns a constant expression as long as the function is contained in the
 // FoldFunctionAllowlist, and the evaluation causes no error.

pkg/sql/opt/norm/rules/fold_constants.opt

@@ -160,6 +160,40 @@ $result
 =>
 $result
 
+# FoldFunctionWithNullArg folds a Function to Null when one of its arguments is
+# Null and all of the following are true:
+#
+#   1. The function does not allow Null arguments (NullableArgs=false).
+#   2. The function is a normal function not an aggregate, window, or generator.
+#
+# It is safe to fold functions to Null in this case because a function with
+# NullableArgs=false would never error with a Null argument, even if the other
+# args are invalid. For example, calling encode with NULL bytes and an invalid
+# encoding format does not error:
+#
+#     SELECT encode(NULL::BYTES, 'foo')
+#       => NULL
+#
+# Stable and volatile functions that rely on context or produce side-effects can
+# also be folded to Null in this case because a function with NullableArgs=false
+# is never evaluated if any of its arguments are Null. The function results
+# directly in Null without being invoked, so it is guaranteed not to rely on
+# context or produce side-effects. See FunctionProperties.NullableArgs for more
+# details.
+#
+# FoldFunctionWithNullArg is defined before FoldFunction so that we can avoid
+# the overhead of evaluating the function in FoldFunction if it has any Null
+# arguments.
+[FoldFunctionWithNullArg, Normalize]
+(Function
+    $args:*
+    $private:* &
+        (CanFoldFunctionWithNullArg $private) &
+        (HasNullArg $args)
+)
+=>
+(Null (FunctionReturnType $private))
+
 # FoldFunction is similar to FoldBinary, but it involves a function with
 # constant inputs. As with FoldBinary, FoldFunction applies as long as the
 # evaluation would not cause an error. Additionally, only certain functions
@@ -174,7 +208,7 @@ $result
 =>
 $result
 
-# FoldEqualsAnyNull, conversts a scalar ANY operation to NULL if the right-hand
+# FoldEqualsAnyNull, converts a scalar ANY operation to NULL if the right-hand
 # side tuple is NULL, e.g. x = ANY(NULL::int[]). See #42562.
 [FoldEqualsAnyNull, Normalize]
 (AnyScalar * (Null) *)

pkg/sql/opt/norm/testdata/rules/fold_constants

@@ -813,6 +813,97 @@ values
  ├── fd: ()-->(1)
  └── ('2017-05-10 00:00:00+00:00',)
 
+# --------------------------------------------------
+# FoldFunctionWithNullArgs
+# --------------------------------------------------
+
+norm expect=FoldFunctionWithNullArg
+SELECT lpad(NULL::STRING, i) FROM a
+----
+project
+ ├── columns: lpad:8
+ ├── fd: ()-->(8)
+ ├── scan a
+ └── projections
+      └── CAST(NULL AS STRING) [as=lpad:8]
+
+norm expect=FoldFunctionWithNullArg
+SELECT lpad(s, NULL::INT) FROM a
+----
+project
+ ├── columns: lpad:8
+ ├── fd: ()-->(8)
+ ├── scan a
+ └── projections
+      └── CAST(NULL AS STRING) [as=lpad:8]
+
+# Fold a stable function.
+norm expect=FoldFunctionWithNullArg
+SELECT date_trunc(s, NULL::TIMESTAMPTZ) FROM a
+----
+project
+ ├── columns: date_trunc:8
+ ├── fd: ()-->(8)
+ ├── scan a
+ └── projections
+      └── CAST(NULL AS TIMESTAMPTZ) [as=date_trunc:8]
+
+# Fold a volatile function.
+norm expect=FoldFunctionWithNullArg
+SELECT nextval(NULL::STRING) FROM a
+----
+project
+ ├── columns: nextval:8
+ ├── fd: ()-->(8)
+ ├── scan a
+ └── projections
+      └── CAST(NULL AS INT8) [as=nextval:8]
+
+# Do not fold an aggregate function.
+norm expect-not=FoldFunctionWithNullArg
+SELECT string_agg(s, NULL::STRING) FROM a
+----
+scalar-group-by
+ ├── columns: string_agg:9
+ ├── cardinality: [1 - 1]
+ ├── key: ()
+ ├── fd: ()-->(9)
+ ├── project
+ │    ├── columns: column8:8 s:4
+ │    ├── fd: ()-->(8)
+ │    ├── scan a
+ │    │    └── columns: s:4
+ │    └── projections
+ │         └── CAST(NULL AS STRING) [as=column8:8]
+ └── aggregations
+      └── string-agg [as=string_agg:9, outer=(4,8)]
+           ├── s:4
+           └── column8:8
+
+# Do not fold a function that allows null arguments.
+norm expect-not=FoldFunctionWithNullArg
+SELECT concat(s, NULL::STRING) FROM a
+----
+project
+ ├── columns: concat:8
+ ├── immutable
+ ├── scan a
+ │    └── columns: s:4
+ └── projections
+      └── concat(s:4, CAST(NULL AS STRING)) [as=concat:8, outer=(4), immutable]
+
+# Do not fold a function without a Null argument.
+norm expect-not=FoldFunctionWithNullArg
+SELECT lpad(s, 5) FROM a
+----
+project
+ ├── columns: lpad:8
+ ├── immutable
+ ├── scan a
+ │    └── columns: s:4
+ └── projections
+      └── lpad(s:4, 5) [as=lpad:8, outer=(4), immutable]
+
 # --------------------------------------------------
 # FoldFunction
 # --------------------------------------------------

pkg/sql/opt/xform/testdata/rules/select

@@ -4832,10 +4832,10 @@ CREATE TABLE t60527 (
 )
 ----
 
-# TODO(mgartner): Functions with NullableArgs=false and a NULL argument can be
-# normalized to NULL. This would simplify this query plan to an empty Values
-# expression.
-opt
+# FoldFunctionWithNullArg is disabled to prevent the expression from being
+# normalized to an empty Values expression and avoiding
+# GenerateInvertedIndexScans altogether.
+opt disable=FoldFunctionWithNullArg
 SELECT * FROM t60527 WHERE (ST_Covers(NULL::GEOGRAPHY, g) AND ST_DWithin(NULL::GEOGRAPHY, g, 1))
 ----
 select
@@ -4856,10 +4856,10 @@ CREATE TABLE t62686 (
 )
 ----
 
-# TODO(mgartner): Functions with NullableArgs=false and a NULL argument can be
-# normalized to NULL. This would simplify this query plan to an empty Values
-# expression.
-opt
+# FoldFunctionWithNullArg is disabled to prevent the expression from being
+# normalized to an empty Values expression and avoiding
+# GenerateInvertedIndexScans altogether.
+opt disable=FoldFunctionWithNullArg
 SELECT * FROM t62686 WHERE ST_DFullyWithin(c, ST_GeomFromText('POINT(1 1)'), NULL::FLOAT8)
 ----
 select

pkg/sql/sem/tree/function_definition.go

@@ -51,10 +51,17 @@ type FunctionProperties struct {
 	// considered undocumented.
 	Private bool
 
-	// NullableArgs is set to true when a function's definition can
-	// handle NULL arguments. When set, the function will be given the
-	// chance to see NULL arguments. When not, the function will
-	// evaluate directly to NULL in the presence of any NULL arguments.
+	// NullableArgs is set to true when a function's definition can handle NULL
+	// arguments. When set to true, the function will be given the chance to see NULL
+	// arguments.
+	//
+	// When set to false, the function will directly result in NULL in the
+	// presence of any NULL arguments without evaluating the function's
+	// implementation defined in Overload.Fn. Therefore, if the function is
+	// expected to produce side-effects with a NULL argument, NullableArgs must
+	// be true. Note that if this behavior changes so that NullableArgs=false
+	// functions can produce side-effects, the FoldFunctionWithNullArg optimizer
+	// rule must be changed to avoid folding those functions.
 	//
 	// NOTE: when set, a function should be prepared for any of its arguments to
 	// be NULL and should act accordingly.