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 name references 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 is safe because computed columns only permit immutable functions,
and CockroachDB currently restricts UDFs in computed columns to
single-statement SQL-language functions — exactly the class of UDFs
that can be mechanically inlined.

Fixes: #157195
Informs: #147472

Release note (bug fix): Fixed a bug where IMPORT INTO would fail
when the destination table had computed columns defined using
user-defined functions, reporting "function <oid> not found".
Computed columns referencing immutable single-statement SQL UDFs
now work correctly with IMPORT INTO.

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

Author: mw5h

pkg/sql/catalog/schemaexpr/computed_column.go

@@ -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,132 @@ 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) (bool, tree.Expr, 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,
+			)
+		}
+		selClause, ok := sel.Select.(*tree.SelectClause)
+		if !ok || len(selClause.Exprs) != 1 {
+			return false, nil, errors.Newf(
+				"expected single-expression SELECT in UDF body",
+			)
+		}
+		bodyExpr := selClause.Exprs[0].Expr
+
+		// Build a mapping from parameter names to the actual argument
+		// expressions from the call site. The body references parameters
+		// by name (e.g. "x"), and we replace those with the
+		// already-typed argument expressions (e.g. IndexedVar).
+		paramMap := make(map[string]tree.Expr, len(fn.RoutineParams))
+		for i, p := range fn.RoutineParams {
+			if i < len(funcExpr.Exprs) {
+				paramMap[string(p.Name)] = funcExpr.Exprs[i]
+			}
+		}
+
+		// Substitute parameter references in the body.
+		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
+					}
+				}
+				return true, inner, nil
+			},
+		)
+		if 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)

pkg/sql/logictest/testdata/logic_test/udf_unsupported

@@ -6,25 +6,85 @@ CREATE FUNCTION test_tbl_f() RETURNS INT IMMUTABLE LANGUAGE SQL AS $$ SELECT 1 $
 statement ok
 CREATE TABLE test_tbl_t (a INT PRIMARY KEY, b INT);
 
-# Insert a row to verify that backfills that use UDFs are blocked without internal errors.
+# Insert rows to verify that backfills that use UDFs produce correct values.
 statement ok
-INSERT INTO test_tbl_t VALUES (1, 1);
+INSERT INTO test_tbl_t VALUES (1, 10), (2, 20), (3, 30);
 
-statement error pgcode 0A000 unimplemented: cannot evaluate function in this context
+# UDFs in computed columns are supported via expression inlining.
+statement ok
 ALTER TABLE test_tbl_t ADD COLUMN c int AS (test_tbl_f()) stored;
 
+# Verify backfilled values are correct.
+query III rowsort
+SELECT * FROM test_tbl_t
+----
+1  10  1
+2  20  1
+3  30  1
+
 statement error pgcode 0A000 unimplemented: cannot evaluate function in this context
-ALTER TABLE test_tbl_t ADD COLUMN c int DEFAULT (test_tbl_f());
+ALTER TABLE test_tbl_t ADD COLUMN d int DEFAULT (test_tbl_f());
 
 statement error pgcode 0A000 unimplemented: cannot evaluate function in this context
 CREATE INDEX t_idx_partial ON test_tbl_t(b) WHERE test_tbl_f() > 0;
 
-statement error pgcode 0A000 unimplemented: cannot evaluate function in this context
+# UDFs in expression indexes are supported via expression inlining.
+statement ok
 CREATE INDEX idx_b ON test_tbl_t (test_tbl_f());
 
 subtest end
 
 
+subtest backfill_udf_computed_column
+
+# Test that backfilling a computed column that references a UDF with
+# parameters produces correct values.
+statement ok
+CREATE FUNCTION bf_double(x INT) RETURNS INT IMMUTABLE LANGUAGE SQL AS $$ SELECT x * 2 $$;
+
+statement ok
+CREATE TABLE bf_t (a INT PRIMARY KEY);
+
+statement ok
+INSERT INTO bf_t VALUES (1), (2), (3), (4), (5);
+
+statement ok
+ALTER TABLE bf_t ADD COLUMN b INT AS (bf_double(a)) STORED;
+
+query II rowsort
+SELECT * FROM bf_t
+----
+1  2
+2  4
+3  6
+4  8
+5  10
+
+# Test with a strict UDF (RETURNS NULL ON NULL INPUT) on a nullable
+# column. Rows with NULL input should produce NULL output.
+statement ok
+CREATE FUNCTION bf_strict(x INT) RETURNS INT IMMUTABLE RETURNS NULL ON NULL INPUT LANGUAGE SQL AS $$ SELECT x + 100 $$;
+
+statement ok
+CREATE TABLE bf_strict_t (a INT PRIMARY KEY, b INT);
+
+statement ok
+INSERT INTO bf_strict_t VALUES (1, 10), (2, NULL), (3, 30), (4, NULL);
+
+statement ok
+ALTER TABLE bf_strict_t ADD COLUMN c INT AS (bf_strict(b)) STORED;
+
+query III rowsort
+SELECT * FROM bf_strict_t
+----
+1  10    110
+2  NULL  NULL
+3  30    130
+4  NULL  NULL
+
+subtest end
+
+
 subtest cross_db
 
 statement ok