Planner Hooks

DataFusion cannot see into a plugin function. So a filter on the function prunes no files and no row groups. A planner hook tells Beacon what a call means. Then Beacon adds plain predicates that the file formats can use. All hooks are optional.

Builder method	Hook	Use
implies(f)	FunctionHooks::implied	A plain predicate that the call implies.
bounds(f)	FunctionHooks::bounds	The output interval for the argument intervals.
propagate(f)	FunctionHooks::propagate	The argument intervals for an output interval.

Beacon applies the hooks one time per query, after type coercion and before filter pushdown.

implies

example_in_box(lon, lat, min_lon, min_lat, max_lon, max_lat) is true when the point is in the box. Its implies hook gives a range for each column:

function("example_in_box")
    .args(vec![DataType::Float64; 6])
    .returns(DataType::Boolean)
    .implies(|args| {
        let lon = args[0].clone().between(args[2].clone(), args[4].clone());
        let lat = args[1].clone().between(args[3].clone(), args[5].clone());
        Ok(Some(Implied::Necessary(lon.and(lat))))
    })

So WHERE example_in_box(lon, lat, 2, 50, 8, 56) also gets lon BETWEEN 2 AND 8 AND lat BETWEEN 50 AND 56. Beacon uses that predicate to prune files with their statistics.

The hook gets the argument expressions. It returns one of these values:

Value	Meaning
Ok(Some(Implied::Necessary(expr)))	Beacon keeps the call and adds expr next to it.
Ok(Some(Implied::Exact(expr)))	Beacon replaces the call with expr.
Ok(None)	No rule for these arguments.

The rule must be correct:

• Necessary: each row where the call is true must also make expr true. Else rows disappear from the result.
• Exact: expr must give the same result as the call for each row.

bounds and propagate

example_add_one(x) gives x + 1. Its hooks work with intervals:

fn one() -> Result<Interval> {
    let one = ScalarValue::Float64(Some(1.0));
    Interval::try_new(one.clone(), one)
}

function("example_add_one")
    // ...
    .bounds(|args| Ok(Some(args[0].add(one()?)?)))
    .propagate(|output, args| {
        // An empty intersection means no row matches. No answer is always correct.
        let Some(x) = args[0].intersect(output.sub(one()?)?)? else {
            return Ok(None);
        };
        Ok(Some(vec![x]))
    })

• bounds gets the argument intervals and gives the output interval. x in [0, 10] gives [1, 11].
• propagate gets the output interval and the argument intervals. It gives the argument intervals that can make that output. The filter example_add_one(x) > 10 gives x > 9.

Beacon uses propagate for a filter f(..) or f(..) <op> <literal>. It adds col >= low and col <= high for each column argument. Beacon uses propagate only when implies gives no rule.

bounds and propagate also answer DataFusion interval analysis, through evaluate_bounds and propagate_constraints.

How Beacon combines the hooks

• Nested calls. In example_in_box(example_add_one(lon), lat, 3, 50, 9, 56), Beacon gets the range [3, 9] for the first argument. Then the propagate hook of example_add_one gives lon in [2, 8]. This also works when the two functions come from different plugins.
• Arithmetic and casts. Beacon follows + and - with a constant, * and / with a constant for float types, negation and casts. So example_add_one(x + 1) > 10 gives x > 8.
• Other conditions. A condition such as x >= 0 or x BETWEEN 2 AND 8 in the same filter sets the input interval of x. A propagate hook gets this interval. So a hook can relate two columns.
• Constants. Beacon folds a constant argument, such as CAST(2 AS DOUBLE), to a literal before implies sees it.
• Exact rules first. Beacon applies the Exact rules of calls in the arguments before implies sees them.
• Functions of other plugins. An expression crosses the plugin boundary as datafusion-proto bytes. Beacon decodes it with the functions of the session. So the arguments and the result of implies can call any function that Beacon knows. A function from another library is an FFI wrapper. Find it by its name, not with a downcast.

Rules

• bounds must hold every value that the function can give.
• propagate must hold every input value that can give the output.
• A wider interval is always correct. A narrower interval that is wrong removes rows from the result.
• A hook that gives Ok(None), an error or a panic gives no answer. The query does not fail.
• Interval arithmetic rounds float bounds outward. So [0, 10] + 1 can give [0.9999999999999999, 11.000000000000002]. Compare intervals with a tolerance in tests.

Hooks for a ScalarUDFImpl

Implement FunctionHooks for a type. Then wrap the function with HookedUdf::new:

use beacon_plugin_sdk::scalar_functions::{FunctionHooks, HookedUdf};

let udf = ScalarUDF::from(HookedUdf::new(Arc::new(ScalarUDF::from(MyUdf::new())), Arc::new(MyHooks)));

FunctionHooks::implied(&self, args, ctx) also gets the session TaskContext. It holds all functions that Beacon knows.

Test a hook

hooks_of(&udf) gives the hooks of a function. Call them directly:

let hooks = hooks_of(&udf).unwrap();
let ctx = SessionContext::new().task_ctx();
let implied = hooks.implied(&args, &ctx).unwrap();
let bounds = hooks.bounds(&[&interval(0.0, 10.0)]).unwrap();
let inputs = hooks.propagate(&interval(5.0, 9.0), &[&interval(0.0, 100.0)]).unwrap();

See the tests in native/src/lib.rs and Testing.