schemaexpr: inline UDF calls in computed column expressions

Contexts that bypass the optimizer — IMPORT and schema change backfill —
cannot evaluate user-defined functions with SQL bodies. The eval.Expr
path returns an unimplemented error for functions where fn.Body != ""
because it lacks the optimizer and execbuilder infrastructure needed to
produce RoutineExpr plan generators.

Fix this by inlining UDF calls within MakeComputedExprs, the shared
entry point for both IMPORT and backfill computed column evaluation. For
each FuncExpr that references a UDF (fn.Body != ""), inlineUDFCalls:

1. Parses the SQL body and extracts the result expression.
2. Substitutes parameter references (named, ColumnItem, and ordinal
   $1/$2 placeholders) with the actual call-site argument expressions.
3. Wraps the result in a CASE expression for strict (RETURNS NULL ON
   NULL INPUT) functions to preserve null-handling semantics.
4. Type-checks the inlined expression against the function's return
   type.

This provides parity with non-UDF computed column expressions for the
constructs allowed in immutable UDF bodies: arithmetic, string
operations, type casts, CASE/COALESCE, immutable builtin calls, and
array constructors. Immutable UDFs cannot reference relations, so their
bodies are restricted to the same scalar expression language available
in non-UDF computed columns.

Limitations during backfill/IMPORT (the optimizer handles these fine in
CREATE TABLE and INSERT):
- Multi-statement SQL UDFs: cannot be inlined to a single expression
- PL/pgSQL UDFs: cannot be inlined (different language)
- UDFs with OUT/INOUT parameters: not supported by the inliner
- Nested UDF calls: the DistSQLFunctionResolver resolves UDFs by OID
  only; inlined body expressions reference functions by name, which
  only resolves builtins
- UDF bodies containing subqueries, CTEs, or FROM clauses: cannot be
  evaluated outside the optimizer

All of these produce explicit unimplemented errors during backfill.

Fixes: #157195
Informs: #147472

Release note (bug fix): Fixed a bug where IMPORT INTO and schema change
backfills would fail when the destination table had computed columns
defined using user-defined functions, reporting "function <oid> not
found". Computed columns referencing immutable single-expression SQL
UDFs now work correctly with IMPORT INTO and ALTER TABLE ADD COLUMN.
Multi-statement UDFs, PL/pgSQL UDFs, and nested UDF calls in computed
columns are not yet supported in these contexts.

Co-Authored-By: roachdev-claude <roachdev-claude-bot@cockroachlabs.com>

Author: mw5h

pkg/sql/catalog/schemaexpr/computed_column.go

@@ -323,8 +323,8 @@ func MakeComputedExprs(
 			continue
 		}
 
-		// Collect all column IDs that are referenced in the partial index
-		// predicate expression.
+		// Collect all column IDs that are referenced in the computed
+		// column expression.
 		colIDs, err := ExtractColumnIDs(tableDesc, exprs[compExprIdx])
 		if err != nil {
 			return nil, refColIDs, err
@@ -341,6 +341,16 @@ func MakeComputedExprs(
 			return nil, catalog.TableColSet{}, err
 		}
 
+		// Inline any UDF calls so that the expression can be evaluated
+		// by eval.Expr. Contexts that bypass the optimizer (IMPORT,
+		// backfill) cannot evaluate functions with SQL bodies directly.
+		// This must happen before the assignment cast so the cast
+		// operates on the resolved expression, not a FuncExpr wrapper.
+		typedExpr, err = inlineUDFCalls(ctx, typedExpr, semaCtx)
+		if err != nil {
+			return nil, catalog.TableColSet{}, err
+		}
+
 		// If the expression has a type that is not identical to the
 		// column's type, wrap the computed column expression in an assignment cast.
 		typedExpr, err = wrapWithAssignmentCast(ctx, typedExpr, col, semaCtx)
@@ -357,3 +367,184 @@ func MakeComputedExprs(
 	}
 	return computedExprs, refColIDs, nil
 }
+
+// inlineUDFCalls walks a TypedExpr tree and replaces FuncExpr nodes that
+// reference user-defined functions (those with SQL bodies) with the inlined
+// body expression. This is necessary because eval.Expr cannot evaluate
+// functions with SQL bodies — that capability is only available through
+// the optimizer. Contexts that evaluate computed column expressions without
+// the optimizer (IMPORT, schema change backfill) rely on this inlining.
+//
+// Only single-statement SQL-language UDFs can be inlined. Multi-statement
+// or PL/pgSQL functions will produce an error.
+func inlineUDFCalls(
+	ctx context.Context, expr tree.TypedExpr, semaCtx *tree.SemaContext,
+) (tree.TypedExpr, error) {
+	newExpr, err := tree.SimpleVisit(expr, func(e tree.Expr) (recurse bool, newExpr tree.Expr, err error) {
+		funcExpr, ok := e.(*tree.FuncExpr)
+		if !ok {
+			return true, e, nil
+		}
+		fn := funcExpr.ResolvedOverload()
+		if fn == nil || fn.Body == "" {
+			return true, e, nil
+		}
+		if fn.Language != tree.RoutineLangSQL {
+			return false, nil, unimplemented.Newf(
+				"computed_column_plpgsql_udf",
+				"PL/pgSQL user-defined functions in computed columns "+
+					"cannot be evaluated in this context",
+			)
+		}
+
+		// Parse the function body.
+		stmts, err := parserutils.Parse(fn.Body)
+		if err != nil {
+			return false, nil, errors.Wrap(err, "parsing UDF body for inlining")
+		}
+		if len(stmts) != 1 {
+			return false, nil, unimplemented.Newf(
+				"computed_column_multi_stmt_udf",
+				"multi-statement user-defined functions in computed columns "+
+					"cannot be evaluated in this context",
+			)
+		}
+
+		// Extract the result expression from the SELECT statement.
+		sel, ok := stmts[0].AST.(*tree.Select)
+		if !ok {
+			return false, nil, errors.Newf(
+				"expected SELECT in UDF body, got %T", stmts[0].AST,
+			)
+		}
+		if sel.With != nil {
+			return false, nil, unimplemented.Newf(
+				"computed_column_cte_udf",
+				"user-defined functions with CTEs in computed columns "+
+					"cannot be evaluated in this context",
+			)
+		}
+		selClause, ok := sel.Select.(*tree.SelectClause)
+		if !ok || len(selClause.Exprs) != 1 {
+			return false, nil, errors.Newf(
+				"expected single-expression SELECT in UDF body",
+			)
+		}
+		if len(selClause.From.Tables) > 0 {
+			return false, nil, unimplemented.Newf(
+				"computed_column_from_udf",
+				"user-defined functions with FROM clauses in computed columns "+
+					"cannot be evaluated in this context",
+			)
+		}
+		bodyExpr := selClause.Exprs[0].Expr
+
+		// Build a mapping from named parameters to the actual argument
+		// expressions from the call site. Unnamed parameters (those
+		// with empty names) are only referenceable via ordinal $1/$2
+		// placeholders, so they are excluded from this map.
+		paramMap := make(map[string]tree.Expr, len(fn.RoutineParams))
+		for i, p := range fn.RoutineParams {
+			if tree.IsOutParamClass(p.Class) {
+				return false, nil, unimplemented.Newf(
+					"computed_column_out_param_udf",
+					"user-defined functions with OUT or INOUT parameters "+
+						"in computed columns cannot be evaluated in this context",
+				)
+			}
+			if p.Name != "" && i < len(funcExpr.Exprs) {
+				paramMap[string(p.Name)] = funcExpr.Exprs[i]
+			}
+		}
+
+		// Substitute parameter references in the body. Named
+		// parameters appear as UnresolvedName or ColumnItem nodes;
+		// unnamed parameters use ordinal placeholders ($1, $2, etc.).
+		inlined, err := tree.SimpleVisit(
+			bodyExpr,
+			func(inner tree.Expr) (bool, tree.Expr, error) {
+				if name, ok := inner.(*tree.UnresolvedName); ok &&
+					name.NumParts == 1 {
+					if arg, exists := paramMap[name.Parts[0]]; exists {
+						return false, arg, nil
+					}
+				}
+				if ci, ok := inner.(*tree.ColumnItem); ok {
+					if arg, exists := paramMap[string(ci.ColumnName)]; exists {
+						return false, arg, nil
+					}
+				}
+				if ph, ok := inner.(*tree.Placeholder); ok {
+					idx := int(ph.Idx)
+					if idx < len(funcExpr.Exprs) {
+						return false, funcExpr.Exprs[idx], nil
+					}
+				}
+				return true, inner, nil
+			},
+		)
+		if err != nil {
+			return false, nil, err
+		}
+
+		// Verify the inlined body contains no constructs that cannot
+		// be evaluated outside the optimizer.
+		if _, err := tree.SimpleVisit(inlined, func(inner tree.Expr) (bool, tree.Expr, error) {
+			if _, ok := inner.(*tree.Subquery); ok {
+				return false, inner, unimplemented.Newf(
+					"computed_column_subquery_udf",
+					"user-defined functions with subqueries in computed columns "+
+						"cannot be evaluated in this context",
+				)
+			}
+			// Check for unreplaced parameter references. These
+			// indicate a parameter substitution bug or an
+			// unsupported reference pattern.
+			if ph, ok := inner.(*tree.Placeholder); ok {
+				return false, inner, errors.Newf(
+					"unresolved parameter reference $%d in inlined UDF body", ph.Idx+1,
+				)
+			}
+			return true, inner, nil
+		}); err != nil {
+			return false, nil, err
+		}
+
+		// For strict functions (RETURNS NULL ON NULL INPUT), wrap the
+		// inlined body to return NULL when any argument is NULL,
+		// preserving the function's null-handling semantics.
+		if !fn.CalledOnNullInput && len(funcExpr.Exprs) > 0 {
+			var nullCheck tree.TypedExpr
+			for i, arg := range funcExpr.Exprs {
+				isNull := &tree.IsNullExpr{Expr: arg}
+				if i == 0 {
+					nullCheck = isNull
+				} else {
+					nullCheck = tree.NewTypedOrExpr(nullCheck, isNull)
+				}
+			}
+			inlined = &tree.CaseExpr{
+				Whens: []*tree.When{{
+					Cond: nullCheck,
+					Val:  tree.DNull,
+				}},
+				Else: inlined,
+			}
+		}
+
+		// Type-check the inlined expression against the function's
+		// return type.
+		typedInlined, err := tree.TypeCheck(
+			ctx, inlined, semaCtx, funcExpr.ResolvedType(),
+		)
+		if err != nil {
+			return false, nil, errors.Wrap(err, "type-checking inlined UDF body")
+		}
+
+		return false, typedInlined, nil
+	})
+	if err != nil {
+		return nil, err
+	}
+	return newExpr.(tree.TypedExpr), nil
+}

pkg/sql/importer/import_stmt_test.go

@@ -1252,6 +1252,41 @@ END
 	}
 }
 
+// TestImportIntoComputedColumnWithUDF verifies that IMPORT INTO works when
+// the destination table has computed columns defined using user-defined
+// functions. Regression test for #157195.
+func TestImportIntoComputedColumnWithUDF(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+	ctx := context.Background()
+	baseDir, cleanup := testutils.TempDir(t)
+	defer cleanup()
+	tc := serverutils.StartCluster(
+		t, 1, base.TestClusterArgs{ServerArgs: base.TestServerArgs{ExternalIODir: baseDir}})
+	defer tc.Stopper().Stop(ctx)
+	conn := tc.ServerConn(0)
+	sqlDB := sqlutils.MakeSQLRunner(conn)
+
+	// Helper that writes CSV data and runs IMPORT INTO.
+	importCSV := func(t *testing.T, table, intoCols, csvData string) {
+		t.Helper()
+		f, err := os.CreateTemp(baseDir, "data")
+		require.NoError(t, err)
+		_, err = f.Write([]byte(csvData))
+		require.NoError(t, err)
+		require.NoError(t, f.Close())
+		sqlDB.Exec(t, fmt.Sprintf(
+			`IMPORT INTO %s (%s) CSV DATA ($1)`, table, intoCols,
+		), fmt.Sprintf("nodelocal://1/%s", filepath.Base(f.Name())))
+	}
+
+	sqlDB.Exec(t, `CREATE FUNCTION double_val(x INT) RETURNS INT IMMUTABLE LANGUAGE SQL AS $$ SELECT x * 2 $$`)
+	sqlDB.Exec(t, `CREATE TABLE t (a INT, b INT AS (double_val(a)) STORED)`)
+	importCSV(t, "t", "a", "1\n2\n3\n")
+	sqlDB.CheckQueryResults(t, `SELECT a, b FROM t ORDER BY a`,
+		[][]string{{"1", "2"}, {"2", "4"}, {"3", "6"}})
+}
+
 func TestImportRowLimit(t *testing.T) {
 	defer leaktest.AfterTest(t)()
 	defer log.Scope(t).Close(t)