sql: do not evaluate strict UDFs if any input values are NULL

A UDF can have one of two behaviors when it is invoked with NULL inputs:

  1. If the UDF is `CALLED ON NULL INPUT` (the default) then the
     function is evaluated regardless of whether or not any of the input
     values are NULL.
  2. If the UDF `RETURNS NULL ON NULL INPUT` or is `STRICT` then the
     function is not evaluated if any of the input values are NULL.
     Instead, the function directly results in NULL.

This commit implements these two behaviors.

In the future, we can add a normalization rule that folds a strict UDF
if any of its inputs are constant NULL values.

Release note: None

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

@@ -737,5 +737,12 @@ func (b *Builder) buildUDF(ctx *buildScalarCtx, scalar opt.ScalarExpr) (tree.Typ
 		}
 		return plan, nil
 	}
-	return tree.NewTypedRoutineExpr(udf.Name, inputExprs, planFn, len(udf.Body), udf.Typ), nil
+	return tree.NewTypedRoutineExpr(
+		udf.Name,
+		inputExprs,
+		planFn,
+		len(udf.Body),
+		udf.Typ,
+		udf.CalledOnNullInput,
+	), nil
 }

pkg/sql/opt/exec/execbuilder/testdata/udf

@@ -59,3 +59,39 @@ Scan /Table/106/1/1/0
 Scan /Table/106/1/2/0
 Scan /Table/106/1/1/0
 Scan /Table/106/1/2/0
+
+statement ok
+CREATE FUNCTION fetch_a_of_1(i INT) RETURNS INT CALLED ON NULL INPUT LANGUAGE SQL AS $$
+  SELECT a FROM t WHERE k = 1
+$$
+
+statement ok
+CREATE FUNCTION fetch_a_of_1_strict(i INT) RETURNS INT RETURNS NULL ON NULL INPUT LANGUAGE SQL AS $$
+  SELECT a FROM t WHERE k = 1
+$$
+
+statement ok
+CREATE FUNCTION fetch_a_of_2_strict(i INT, j INT) RETURNS INT STRICT LANGUAGE SQL AS $$
+  SELECT a FROM t WHERE k = 2
+$$
+
+# When the function is CALLED ON NULL INPUT then it should be evaluated
+# regardless of whether or not any of its inputs are NULL. The trace proves that
+# the function is evaluated. It shows the scan performed by the statement in the
+# function body.
+query T kvtrace
+SELECT fetch_a_of_1(NULL::INT)
+----
+Scan /Table/106/1/1/0
+
+# When the function RETURNS NULL ON NULL INPUT or STRICT then it should not be
+# evaluated if any of its inputs are NULL. The empty traces prove that the
+# function is not evaluated. No scan is performed for the statement in the
+# function body.
+query T kvtrace
+SELECT fetch_a_of_1_strict(NULL::INT)
+----
+
+query T kvtrace
+SELECT fetch_a_of_2_strict(1, NULL::INT)
+----

pkg/sql/opt/ops/scalar.opt

@@ -1245,6 +1245,11 @@ define UDFPrivate {
     # Volatility is the user-provided volatility of the function given during
     # CREATE FUNCTION.
     Volatility Volatility
+
+    # CalledOnNullInput is true if the function should be called when any of its
+    # inputs are NULL. If false, the function will not be evaluated in the
+    # presence of NULL inputs, and will instead evaluate directly to NULL.
+    CalledOnNullInput bool
 }
 
 # KVOptions is a set of KVOptionItems that specify arbitrary keys and values

pkg/sql/opt/optbuilder/scalar.go

@@ -672,11 +672,12 @@ func (b *Builder) buildUDF(
 	out = b.factory.ConstructUDF(
 		input,
 		&memo.UDFPrivate{
-			Name:       def.Name,
-			Body:       rels,
-			Typ:        f.ResolvedType(),
-			Volatility: o.Volatility,
-			ArgCols:    argCols,
+			Name:              def.Name,
+			ArgCols:           argCols,
+			Body:              rels,
+			Typ:               f.ResolvedType(),
+			Volatility:        o.Volatility,
+			CalledOnNullInput: o.CalledOnNullInput,
 		},
 	)
 	return b.finishBuildScalar(f, out, inScope, outScope, outCol)

pkg/sql/opt/testutils/testcat/testdata/udf

@@ -55,3 +55,13 @@ SHOW CREATE FUNCTION e
 ----
 FUNCTION e(i: int) -> int [immutable, called-on-null-input=false]
  └── SELECT i
+
+exec-ddl
+CREATE FUNCTION f() RETURNS INT RETURNS NULL ON NULL INPUT LANGUAGE SQL AS 'SELECT 1'
+----
+
+exec-ddl
+SHOW CREATE FUNCTION f
+----
+FUNCTION f() -> int [volatile, called-on-null-input=false]
+ └── SELECT 1

pkg/sql/routine.go

@@ -24,6 +24,16 @@ import (
 func (p *planner) EvalRoutineExpr(
 	ctx context.Context, expr *tree.RoutineExpr, input tree.Datums,
 ) (result tree.Datum, err error) {
+	// If the routine should not be called on null input, then directly return
+	// NULL if any of the datums in the input are NULL.
+	if !expr.CalledOnNullInput {
+		for i := range input {
+			if input[i] == tree.DNull {
+				return tree.DNull, nil
+			}
+		}
+	}
+
 	retTypes := []*types.T{expr.ResolvedType()}
 
 	// The result of the routine is the result of the last statement. The result

pkg/sql/sem/tree/routine.go

@@ -50,19 +50,30 @@ type RoutineExpr struct {
 	// Typ is the type of the routine's result.
 	Typ *types.T
 
+	// CalledOnNullInput is true if the function should be called when any of
+	// its inputs are NULL. If false, the function will not be evaluated in the
+	// presence of null inputs, and will instead evaluate directly to NULL.
+	CalledOnNullInput bool
+
 	name string
 }
 
 // NewTypedRoutineExpr returns a new RoutineExpr that is well-typed.
 func NewTypedRoutineExpr(
-	name string, input TypedExprs, planFn RoutinePlanFn, numStmts int, typ *types.T,
+	name string,
+	input TypedExprs,
+	planFn RoutinePlanFn,
+	numStmts int,
+	typ *types.T,
+	calledOnNullInput bool,
 ) *RoutineExpr {
 	return &RoutineExpr{
-		Input:    input,
-		PlanFn:   planFn,
-		NumStmts: numStmts,
-		Typ:      typ,
-		name:     name,
+		Input:             input,
+		PlanFn:            planFn,
+		NumStmts:          numStmts,
+		Typ:               typ,
+		CalledOnNullInput: calledOnNullInput,
+		name:              name,
 	}
 }