Module tidb_query_codegen::rpn_function

Implementation of the rpn_fn attribute macro.

The macro should be applied to a function:

#[rpn_fn]
fn foo(x: Option<&u32>) -> Result<Option<u8>> {
    Ok(None)
}

The expanded function implements an operation in the coprocessor.

Arguments to macro

If neither varg or raw_varg are supplied, then the generated arguments follow from the supplied function’s arguments. Each argument must have a type Option<&T> for some T.

varg

The RPN operator takes a variable number of arguments. The arguments are passed as a &[Option<&T>]. E.g.,

#[rpn_fn(varg)]
pub fn foo(args: &[Option<&Int>]) -> Result<Option<Real>> {
    // Your RPN function logic
}

raw_varg

The RPN operator takes a variable number of arguments. The arguments are passed as a &[ScalarValueRef]. E.g.,

#[rpn_fn(raw_varg)]
pub fn foo(args: &[ScalarValueRef<'_>]) -> Result<Option<Real>> {
    // Your RPN function logic
}

Use raw_varg where the function takes a variable number of arguments and the types are not the same, for example, RPN function case_when.

max_args

The maximum number of arguments. The macro will generate code to check this as part of validation. Only valid if varg or raw_varg is also used. E.g., #[rpn_fn(varg, max_args = 2)]

min_args

The minimum number of arguments. The macro will generate code to check this as part of validation. Only valid if varg or raw_varg is also used. E.g., #[rpn_fn(varg, min_args = 2)]

extra_validator

A function name for custom validation code to be run when an operation is validated. The validator function should have the signature &tipb::Expr -> Result<()>. E.g., #[rpn_fn(raw_varg, extra_validator = json_object_validator)]

metadata_type

The type of the metadata structure defined in tipb. If metadata_mapper is not specified, the protobuf metadata structure will be used as the metadata directly.

metadata_mapper

A function name to construct a new metadata or transform a protobuf metadata structure into a desired form. The function signatures varies according to the existence of metadata_mapper and metadata_type as follows.

• metadata_mapper exists, metadata_type missing: fn(&mut tipb::Expr) -> T

Constructs a new metadata in type T.

• metadata_mapper exists, metadata_type exists: fn(MetaDataType, &mut tipb::Expr) -> T

Transforms a protobuf metadata type MetaDataType specified by metadata_type into a new type T.

capture

An array of argument names which are passed from the caller to the expanded function. The argument names must be in scope in the generated eval or run methods. Currently, that includes the following arguments (the supplied function must accept these arguments with the corresponding types, in addition to any other arguments):

• ctx: &mut expr::EvalContext
• output_rows: usize
• args: &[rpn_expr::RpnStackNode<'_>]
• extra: &mut rpn_expr::RpnFnCallExtra<'_>
• metadata: &T (where T is the type returned by your metadata constructor)

// This generates `with_context_fn_meta() -> RpnFnMeta`
#[rpn_fn(capture = [ctx])]
fn with_context(ctx: &mut EvalContext, param: Option<&Decimal>) -> Result<Option<Int>> {
    // Your RPN function logic
}

Generated code

Vararg functions

This includes varg and raw_varg.

The supplied function is preserved and a constructor function is generated with a _fn_meta suffix, e.g., #[rpn_fn] fn foo ... will preserve foo and generate foo_fn_meta. The constructor function returns an rpn_expr::RpnFnMeta value.

The constructor function will include code for validating the runtime arguments and running the function, pointers to these functions are stored in the result.

Non-vararg functions

Generate the following (examples assume a supplied function called foo_bar:

• A trait to represent the function (FooBar_Fn) with a single function eval.
  • An impl of that trait for all argument types which panics
  • An impl of that trait for the supported argument type which calls the supplied function.
• An evaluator struct (FooBar_Evaluator) which implements rpn_expr::function::Evaluator, which includes an eval method which dispatches to FooBar_Fn::eval.
• A constructor function similar to the vararg case.

The supplied function is preserved.

The supported argument type is represented as a type-level list, for example, a a function which takes two unsigned ints has an argument representation something like Arg<UInt, Arg<UInt, Null>>. See documentation in components/tidb_query_expr/src/types/function.rs for more details.

The _Fn trait can be customised by implementing it manually. For example, you are going to implement an RPN function called regex_match taking two arguments, the regex and the string to match. You want to build the regex only once if the first argument is a scalar. The code may look like:

fn regex_match_impl(regex: &Regex, text: Option<&Bytes>) -> Result<Option<i32>> {
    // match text
}

#[rpn_fn]
fn regex_match(regex: Option<&Bytes>, text: Option<&Bytes>) -> Result<Option<i32>> {
    let regex = build_regex(regex.cloned());
    regex_match_impl(&regex, text)
}

// Pay attention that the first argument is specialized to `ScalarArg`
impl<'a, Arg1> RegexMatch_Fn for Arg<ScalarArg<'a, Bytes>, Arg<Arg1, Null>>
where Arg1: RpnFnArg<Type = &'a Option<Bytes>> {
    fn eval(
        self,
        ctx: &mut EvalContext,
        output_rows: usize,
        args: &[RpnStackNode<'_>],
        extra: &mut RpnFnCallExtra<'_>,
        metadata: &(dyn Any + Send),
    ) -> Result<VectorValue> {
        let (regex, arg) = self.extract(0);
        let regex = build_regex(regex);
        let mut result = NotChunkedVec::with_capacity(output_rows);
        for row_index in 0..output_rows {
            let (text, _) = arg.extract(row_index);
            result.push(regex_match_impl(&regex, text)?);
        }
        Ok(Evaluable::cast_chunk_into_vector_value(result))
    }
}

If the RPN function accepts variable number of arguments and all arguments have the same eval type, like RPN function coalesce, you can use #[rpn_fn(varg)] like:

#[rpn_fn(varg)]
pub fn foo(args: &[Option<&Int>]) -> Result<Option<Real>> {
    // Your RPN function logic
}

Modules

kw

Structs

NormalRpnFn	Generates an RPN fn which is neither varg or raw_varg.
RawVargsRpnFn	Generates a raw_varg RPN fn.
RpnFnAttr	Parses an attribute like #[rpn_fn(varg, capture = [ctx, output_rows]).
RpnFnEvaluableType	Parses an evaluable type like Option<T>.
RpnFnRefEvaluableTypeWithOption	Parses an evaluable type like Option<&T>, Option<JsonRef>, Option<EnumRef>, Option<SetRef> or Option<BytesRef>.
RpnFnSignatureParam	Parses a function signature parameter like val: &Option<T> or val: &T. If input has &Option, set has_option to true; otherwise, set has_option to false. Caller can use has_option to check if input is valid.
RpnFnSignatureReturnGuardType	Parses a function signature return type like Result<SomeGuard>.
RpnFnSignatureReturnType	Parses a function signature return type like Result<Option<T>>.
ValidatorFnGenerator	Helper utility to generate RPN function validator function.
VargsRpnFn	Generates a varg RPN fn.
VargsRpnFnSignatureParam	Parses a function signature parameter like val: &[&Option<T>] or val: &[&T]. If input has &Option, set has_option to true; otherwise, set has_option to false. Caller can use has_option to check if input is valid.

Enums

RpnFnRefEvaluableType

Parses an evaluable type like &T, JsonRef or BytesRef. &T corresponds to Ref. JsonRef corresponds to Type.

Functions

generate_downcast_metadata
generate_init_metadata_fn
generate_metadata_type_checker
get_vargs_buf	Get corresponding VARGS buffer Json or JsonRef will be stored in VARG_PARAM_BUF_JSON_REF Bytes or BytesRef will be stored in VARG_PARAM_BUF_BYTES_REF
get_vectoried_type	Transform copr framework type into vectorized function type For example, Json in copr framework will be transformed into JsonRef before passing to vectorized functions.
is_bytes	Checks if parameter type is Bytes
is_enum	Checks if parameter type is Enum
is_json	Checks if parameter type is Json
is_set	Checks if parameter type is Set
transform	Entry point for the rpn_fn attribute.