Expose scalar UDF arguments and a custom bind-callback

scalar_udf.go

@@ -19,6 +19,7 @@
 	"context"
 	"database/sql"
 	"database/sql/driver"
+	"fmt"
 	"runtime"
 	"runtime/cgo"
 	"unsafe"
@@ -51,16 +52,33 @@
 
 // bindInfo holds bind data accessible during execution.
 type bindInfo struct {
-	connId uint64
+	connId         uint64
+	customBindData ScalarUDFArg
 }
 
+// ScalarUDFArg contains scalar UDF argument metadata and the optional argument.
+type ScalarUDFArg struct {
+	// Foldable is true, if the argument was folded into a value, else false.
+	Foldable bool
+	// Value is the folded argument value, or nil, if the argument is not foldable.
+	Value driver.Value
+}
+
+// CtxBindDataKeyType defines the key type of the context's optional bind data value.
+type CtxBindDataKeyType string
+
+// CtxBindDataKey defines the key of the context's optional bind data value.
+const CtxBindDataKey CtxBindDataKeyType = "ctx_bind_data_key"
+
 type (
 	// RowExecutorFn is the type for any row-based execution function.
 	// It takes the row values and returns the row execution result, or error.
 	RowExecutorFn func(values []driver.Value) (any, error)
 	// RowContextExecutorFn accepts a row-based execution function using a context.
 	// It takes a context and the row values, and returns the row execution result, or error.
 	RowContextExecutorFn func(ctx context.Context, values []driver.Value) (any, error)
+	// ScalarBinderFn takes a context and the scalar function's arguments.
+	ScalarBinderFn func(ctx context.Context, args []ScalarUDFArg) (ScalarUDFArg, error)
 )
 
 // ScalarFuncExecutor contains the functions to execute a user-defined scalar function.
@@ -70,6 +88,8 @@
 	RowExecutor RowExecutorFn
 	// RowContextExecutor accepts a row-based execution function of type RowContextExecutorFn.
 	RowContextExecutor RowContextExecutorFn
+	// Binder accepts a bind function of type ScalarBinderFn.
+	ScalarBinder ScalarBinderFn
 }
 
 // ScalarFunc is the user-defined scalar function interface.
@@ -194,6 +214,12 @@
 	bindDataPtr := mapping.ScalarFunctionGetBindData(functionInfo)
 	info := getPinned[*bindInfo](bindDataPtr)
 
+	// TODO: decide if we want another RowExecutor taking ScalarUDFArg as its input,
+	// TODO: or if we want to set it (like now) in the context.
+	ctx := function.ctxStore.load(info.connId)
+	ctx = context.WithValue(ctx, CtxBindDataKey, info.customBindData)
+	function.ctxStore.store(info.connId, ctx)
+
 	f := function.RowExecutor(info)
 	values := make([]driver.Value, len(inputChunk.columns))
 
@@ -263,17 +289,25 @@
 //export scalar_udf_bind_callback
 func scalar_udf_bind_callback(bindInfoPtr unsafe.Pointer) {
 	info := mapping.BindInfo{Ptr: bindInfoPtr}
+	extraInfo := mapping.BindGetExtraInfo(info)
+	function := getPinned[*scalarFuncContext](extraInfo)
 
-	// FIXME: Once available through the duckdb-go-bindings (and the C API),
-	// FIXME: we want to get extraInfo here, to access user-defined `func Binder(BindInfo) (any, any)` callbacks.
-	// FIXME: With these callbacks, we can set additional user-defined bind data in the context.
+	var clientCtx mapping.ClientContext
+	mapping.ScalarFunctionGetClientContext(info, &clientCtx)
+	defer mapping.DestroyClientContext(&clientCtx)
 
-	var ctx mapping.ClientContext
-	mapping.ScalarFunctionGetClientContext(info, &ctx)
-	defer mapping.DestroyClientContext(&ctx)
+	connId := mapping.ClientContextGetConnectionId(clientCtx)
+	data := bindInfo{connId: uint64(connId)}
 
-	id := mapping.ClientContextGetConnectionId(ctx)
-	data := bindInfo{connId: uint64(id)}
+	// Get any custom bind data by invoking the custom bind function.
+	if function.f.Executor().ScalarBinder != nil {
+		arg, err := function.bind(info, clientCtx, connId)
+		if err != nil {
+			mapping.ScalarFunctionBindSetError(info, err.Error())
+			return
+		}
+		data.customBindData = arg
+	}
 
 	// Set the copy callback of the bind info.
 	copyPtr := unsafe.Pointer(C.scalar_udf_bind_copy_callback_t(C.scalar_udf_bind_copy_callback))
@@ -292,6 +326,37 @@
 	mapping.ScalarFunctionSetBindData(info, unsafe.Pointer(&h), deleteCallbackPtr)
 }
 
+func (s *scalarFuncContext) bind(info mapping.BindInfo, clientCtx mapping.ClientContext, connId uint64) (ScalarUDFArg, error) {
+	ctx := s.ctxStore.load(connId)
+
+	argCount := mapping.ScalarFunctionBindGetArgumentCount(info)
+	var args []ScalarUDFArg
+	for i := 0; i < argCount; i++ {
+		func() {
+			expr := mapping.ScalarFunctionBindGetArgument(info, i)
+			defer mapping.DestroyExpression(&expr)
+			arg := ScalarUDFArg{
+				Foldable: mapping.ExpressionIsFoldable(expr),
+			}
+			if arg.Foldable {
+				var v mapping.Value
+				errorData := mapping.ExpressionFold(clientCtx, expr, &v)
+				// TODO: handle error data.
+				fmt.Println(errorData)
+				// TODO: generalize value.go getValue() to work here.
+				// TODO: we should be able to get the type info via the logical type,
+				// TODO: either directly from mapping.GetValueType (+recursion), or via the logical type (+recursion) of mapping.ExpressionReturnType,
+				// TODO: which should match the folded type.
+				var assumeThisIsTheExtractedValue any
+				arg.Value = assumeThisIsTheExtractedValue
+			}
+			args = append(args, arg)
+		}()
+	}
+
+	return s.f.Executor().ScalarBinder(ctx, args)
+}
+
 func registerInputParams(config ScalarFuncConfig, f mapping.ScalarFunction) error {
 	// Set variadic input parameters.
 	if config.VariadicTypeInfo != nil {

scalar_udf_test.go

@@ -6,13 +6,18 @@
 	"database/sql/driver"
 	"errors"
 	"fmt"
+	"strconv"
 	"testing"
 
 	"github.com/stretchr/testify/require"
 )
 
 var currentInfo TypeInfo
 
+type testCtxKeyType string
+
+const testCtxKey testCtxKeyType = "test_ctx_key"
+
 type (
 	simpleSUDF        struct{}
 	constantSUDF      struct{}
@@ -22,6 +27,7 @@
 	anyTypeSUDF       struct{}
 	unionTestSUDF     struct{}
 	getConnIdUDF      struct{}
+	easterEggUDF      struct{}
 	errExecutorSUDF   struct{}
 	errInputNilSUDF   struct{}
 	errResultNilSUDF  struct{}
@@ -70,23 +76,50 @@
 	return nil, errors.New("test invalid execution")
 }
 
-type testCtxKeyType string
-
-const testCtxKey testCtxKeyType = "my_conn_id"
-
 func getConnId(ctx context.Context, values []driver.Value) (any, error) {
 	if ctx == nil {
-		return nil, errors.New("context is nil")
+		return nil, errors.New("context is nil for getConnId")
 	}
 
 	id, ok := ctx.Value(testCtxKey).(uint64)
 	if !ok {
-		return nil, errors.New("context does not contain the connection id")
+		return nil, errors.New("context does not contain the connection id for getConnId")
 	}
 
 	return id, nil
 }
 
+func getEasterEgg(ctx context.Context, values []driver.Value) (any, error) {
+	if ctx == nil {
+		return nil, errors.New("context is nil for getEasterEgg")
+	}
+
+	bindData, ok := ctx.Value(CtxBindDataKey).(ScalarUDFArg)
+	if !ok {
+		return nil, errors.New("context does not contain the bind data for getEasterEgg")
+	}
+
+	count := uint64(0)
+	if bindData.foldable {
+		count += bindData.value.(uint64)
+	}
+	if values[0] != nil {
+		count += values[0].(uint64)
+	}
+
+	if count == 42 {
+		return "‧₊˚ ⋅ 𓐐𓎩 ‧₊˚ ⋅", nil
+	}
+	return strconv.Itoa(int(count)), nil
+}
+
+func bindEasterEgg(ctx context.Context, args []ScalarUDFArg) (ScalarUDFArg, error) {
+	if !args[1].foldable {
+		return ScalarUDFArg{}, errors.New("second argument must be foldable for bindEasterEgg")
+	}
+	return args[1], nil
+}
+
 func (*simpleSUDF) Config() ScalarFuncConfig {
 	return ScalarFuncConfig{[]TypeInfo{currentInfo, currentInfo}, currentInfo, nil, false, false}
 }
@@ -148,6 +181,19 @@
 	return ScalarFuncExecutor{RowContextExecutor: getConnId}
 }
 
+func (*easterEggUDF) Config() ScalarFuncConfig {
+	info, err := NewTypeInfo(TYPE_UBIGINT)
+	if err != nil {
+		panic(err)
+	}
+
+	return ScalarFuncConfig{[]TypeInfo{info, info}, currentInfo, nil, true, true}
+}
+
+func (*easterEggUDF) Executor() ScalarFuncExecutor {
+	return ScalarFuncExecutor{RowContextExecutor: getEasterEgg, ScalarBinder: bindEasterEgg}
+}
+
 func (*errExecutorSUDF) Config() ScalarFuncConfig {
 	scalarUDF := simpleSUDF{}
 	return scalarUDF.Config()
@@ -488,6 +534,39 @@
 	require.True(t, res)
 }
 
+func TestBindScalarUDF(t *testing.T) {
+	db := openDbWrapper(t, ``)
+	defer closeDbWrapper(t, db)
+
+	conn := openConnWrapper(t, db, context.Background())
+	defer closeConnWrapper(t, conn)
+
+	var err error
+	currentInfo, err = NewTypeInfo(TYPE_VARCHAR)
+	require.NoError(t, err)
+
+	var udf *easterEggUDF
+	err = RegisterScalarUDF(conn, "find_easter_egg", udf)
+	require.NoError(t, err)
+
+	var egg string
+	row := conn.QueryRowContext(context.Background(), `SELECT find_easter_egg(7, 8) AS egg`)
+	require.NoError(t, row.Scan(&egg))
+	require.Equal(t, egg, "15")
+
+	row = conn.QueryRowContext(context.Background(), `SELECT find_easter_egg(10, (44 - 12)::UBIGINT) AS egg`)
+	require.NoError(t, row.Scan(&egg))
+	require.Equal(t, egg, "‧₊˚ ⋅ 𓐐𓎩 ‧₊˚ ⋅")
+
+	row = conn.QueryRowContext(context.Background(), `SELECT find_easter_egg(NULL, NULL) AS egg`)
+	require.NoError(t, row.Scan(&egg))
+	require.Equal(t, egg, "0")
+
+	row = conn.QueryRowContext(context.Background(), `SELECT find_easter_egg(7, random()) AS egg`)
+	err = row.Scan(&egg)
+	require.ErrorContains(t, err, "second argument must be foldable for bindEasterEgg")
+}
+
 func TestErrScalarUDF(t *testing.T) {
 	db := openDbWrapper(t, ``)
 	defer closeDbWrapper(t, db)