tantivy/jitexpr/examples/basic.rs

use std::collections::HashMap;
use std::error::Error;
use std::sync::Arc;

use jitexpr::ast::{Function, InferredTypeSet, Literal, UntypedExpr, infer_types};
use jitexpr::compile::{CompiledFn, CompiledFnCtx, compile};
use jitexpr::types::{VarType, VariableValue};

fn main() -> Result<(), Box<dyn Error>> {
    // A simple expression that goes:
    // my_col + 1
    let untyped_expr = UntypedExpr::new_fn_call(
        Function::Add,
        vec![
            UntypedExpr::variable("my_col"),
            // A float literal must be finite, so the conversion is fallible.
            UntypedExpr::literal(Literal::try_from(1.0f64)?),
        ],
    )?;

    // Infer types does not return specific types, but instead a set of acceptable
    // types for each variables.
    let inferred_types: HashMap<&str, InferredTypeSet> = infer_types(&untyped_expr)?;
    assert_eq!(
        inferred_types.get("my_col").unwrap(),
        &InferredTypeSet::NUMERICAL
    );

    // This is then up to us to decide the actual type for each variable.
    // In tantivy, this means picking the first column with a type in inferred_types.
    //
    // If none match then we should use the VarType::None.
    let variable_types: HashMap<&str, VarType> =
        std::iter::once(("my_col", VarType::F64)).collect();

    let compiled_fn: Arc<CompiledFn> = compile(&untyped_expr, &variable_types)?;
    let mut compiled_fn_ctx = CompiledFnCtx::new(compiled_fn);

    // We use a nullable wrapper around the value union to pass typed variables.
    // For present values, it is up to us to populate the correct union member.
    // Not doing so is UB.
    let input: Box<[VariableValue]> = vec![VariableValue::from(1.2f64)].into_boxed_slice();
    let output = unsafe { compiled_fn_ctx.call(&input[..]) };
    assert_eq!(unsafe { output.as_f64() }, Some(1.2f64 + 1.0f64));

    Ok(())
}