Scalar Functions

A scalar function gives one value for each row. A plugin has native functions, SQL functions or both.

Native functions

Write native functions in native/src/lib.rs. A function functions() returns all of them:

fn functions() -> Vec<ScalarUDF> {
    vec![add_one(), ScalarUDF::from(Hypot::new()), in_box(), round_to()]
}

beacon_plugin_sdk::export_plugin!(scalar = functions, formats = tsv::formats, tables = memory::tables);

Beacon calls functions() one time, when it loads the plugin. The SDK has two ways to write a function: the closure builder and ScalarUDFImpl.

The closure builder

beacon_plugin_sdk::scalar_functions::function builds a function from a closure. This is example_add_one without its hooks:

use beacon_plugin_sdk::scalar_functions::function;

fn add_one() -> ScalarUDF {
    function("example_add_one")
        .args([DataType::Float64])
        .returns(DataType::Float64)
        .alias("example_inc")
        .invoke_arrays(|arrays| {
            let Some(values) = arrays[0].as_any().downcast_ref::<Float64Array>() else {
                return exec_err!("example_add_one expects a Float64 argument");
            };
            let result: Float64Array = values.iter().map(|v| v.map(|v| v + 1.0)).collect();
            Ok(Arc::new(result))
        })
}

Method	Use
args(types)	Sets the exact argument types. DataFusion casts the arguments to these types.
returns(t)	Sets the return type. You must call it before invoke or invoke_arrays.
alias(name)	Adds one more name for the function.
stable()	The result can change between queries, but not in one query.
volatile()	The result can change for each call, such as a random value.
implies(f), bounds(f), propagate(f)	Planner hooks. See Planner hooks.
invoke_arrays(f)	Sets the body. The body gets arrays.
invoke(f)	Sets the body. The body gets each argument as an array or a scalar.

A function is Immutable by default: the same arguments always give the same result. invoke and invoke_arrays panic when returns is missing. Then the plugin does not load.

Read the arguments

Body	Arguments	Result
invoke_arrays(f)	&[ArrayRef]. Each constant argument becomes an array with one value per row.	ArrayRef
invoke(f)	&[ColumnarValue]. A constant argument, such as 2 in f(x, 2), is a ColumnarValue::Scalar.	ColumnarValue

Use invoke_arrays for a plain function. The body has one case only.

Use invoke for these cases:

• The function requires a constant argument.
• The function prepares a constant one time per batch, such as a parsed pattern or a factor.
• The function gives a scalar result for scalar arguments.

example_round(x, digits) uses invoke. It refuses a column for digits and computes 10^digits one time per batch:

.invoke(|args| {
    let ColumnarValue::Scalar(ScalarValue::Int64(digits)) = &args[1] else {
        return exec_err!("example_round expects a constant digits argument");
    };
    // ...
    match &args[0] {
        ColumnarValue::Scalar(ScalarValue::Float64(v)) => { /* one value */ }
        ColumnarValue::Array(array) => { /* one value per row */ }
        other => exec_err!("example_round expects a Float64 argument, got {other:?}"),
    }
})

The body of invoke must accept a scalar and an array for each argument.

ScalarUDFImpl

For full control, implement the DataFusion trait ScalarUDFImpl. Then wrap the type with ScalarUDF::from. example_hypot uses this way:

impl ScalarUDFImpl for Hypot {
    fn name(&self) -> &str { "example_hypot" }
    fn signature(&self) -> &Signature { &self.signature }
    fn return_type(&self, _: &[DataType]) -> Result<DataType> { Ok(DataType::Float64) }

    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
        let x = args.args[0].to_array(args.number_rows)?;
        let y = args.args[1].to_array(args.number_rows)?;
        // ...
    }
}

invoke_with_args gets each constant argument as a scalar. to_array(number_rows) changes a scalar to an array. Use it when the body needs arrays only.

Use ScalarUDFImpl for these cases:

• A signature that args cannot give, such as variadic arguments or more than one type.
• A return type that depends on the argument types.
• Other DataFusion methods, such as simplify or output_ordering.

To add planner hooks to a ScalarUDFImpl, see Planner hooks.

Errors and panics

• Return an error with exec_err!. The query fails with your message.
• The SDK catches a panic in a function body. The query fails with an error that names the function.
• A panic in name, aliases or signature is not safe. Do not panic in these methods or in Drop.
• A fault in native code, such as a segmentation fault, stops the Beacon server.

SQL functions

Write SQL functions in the files of [sql], such as sql/functions.sql:

CREATE FUNCTION example_add_two(x DOUBLE) RETURNS DOUBLE
RETURN example_add_one(example_add_one(x));

• A SQL file can hold only CREATE FUNCTION statements.
• A SQL function can call the native functions of the same plugin and all other Beacon functions.
• Beacon loads the files in the order of sql.files.

Names

• Start each name and each alias with the plugin name, such as example_.
• A name cannot be a built-in function, a user function or a function of another plugin. Then the install fails.
• A user cannot drop a plugin function with DROP FUNCTION. Use UNINSTALL <plugin>.

Next: Planner hooks.