Simple Functions Preview

[findepi]

This is currently a preview PR for #12635 to share status of the work.

Based on add_one function, this creates add_one_udf (by convention, configurable):

#[excalibur_function]
fn add_one(a: u64) -> u64 {
    a + 1
}

The generated ScalarUDFImpl is fully functional, with signature derived from the Rust-level signature of the decorated function:

#[test]
fn simple_function_signature() {
    let udf = add_one_udf();
    assert_eq!(udf.name(), "add_one");

    assert_eq!(
        udf.signature(),
        &Signature::coercible(
            vec![TypeSignatureClass::Native(Arc::new(NativeType::UInt64))],
            Volatility::Immutable
        )
    );
    let return_type = udf.return_type(&[DataType::UInt64]).unwrap();
    assert_eq!(return_type, DataType::UInt64);
}

#[test]
fn simple_function_invoke_array() {
    let udf = add_one_udf();

    let invoke_args = vec![ColumnarValue::Array(Arc::new(UInt64Array::from(vec![
        1000, 2000, 3000, 4000, 5000,
    ])))];
    let ColumnarValue::Array(result_array) = udf
        .invoke_with_args(ScalarFunctionArgs {
            args: invoke_args,
            number_rows: 5,
            return_type: &DataType::UInt64,
        })
        .unwrap()
    else {
        panic!("Expected array result");
    };

    assert_eq!(
        &*result_array,
        &UInt64Array::from(vec![1001, 2001, 3001, 4001, 5001])
    );
}

#[test]
fn simple_function_invoke_scalar() {
    let udf = add_one_udf();

    let invoke_args = vec![ColumnarValue::Scalar(ScalarValue::UInt64(Some(1000)))];
    let ColumnarValue::Array(result_array) = udf
        .invoke_with_args(ScalarFunctionArgs {
            args: invoke_args,
            number_rows: 1,
            return_type: &DataType::UInt64,
        })
        .unwrap()
    else {
        panic!("Expected array result");
    };

    assert_eq!(&*result_array, &UInt64Array::from(vec![1001]));
}

[findepi]

i now also have support for

various numeric

#[excalibur_function]
fn add(a: i32, b: u32) -> i64 {
    a as i64 + b as i64
}

nullable function arguments

#[excalibur_function]
fn first_non_null(a: Option<i32>, b: Option<i32>) -> Option<i32> {
    a.or(b)
}

nullable results

#[excalibur_function]
fn try_div(a: i32, b: i32) -> Option<i32> {
    if b == 0 {
        None
    } else {
        Some(a / b)
    }
}

functions that can throw

#[excalibur_function]
fn maybe_fail(fail: bool) -> Result<bool> {
    if fail {
        exec_err!("This test function just failed")
    } else {
        Ok(true)
    }
}

and string arguments

#[excalibur_function]
fn character_length(s: &str) -> u64 {
    s.chars().count() as u64
}

will update this PR soon (updated)

[jayzhan211]

#[excalibur_function]
fn add(a: i32, b: u32) -> i64 {
    a as i64 + b as i64
}

Is it possible to define a function using generics? Otherwise, will we need to duplicate the function for different types?

[findepi]

Is it possible to define a function using generics?

currently not.

Otherwise, will we need to duplicate the function for different types?

currently yes; which is easy with a macro.
of course it would be nicer to support such functions directly. I don't know yet how to do this.

[findepi]

i learned a bunch during last few days, being thrown at pages like below, which are good read

• https://github.com/dtolnay/case-studies
• https://blog.rust-lang.org/2022/10/28/gats-stabilization.html
• GATs: Decide whether to have defaults for where Self: 'a rust-lang/rust#87479
• https://docs.rs/lending-iterator/latest/lending_iterator
• https://docs.rs/higher-order-closure/latest/higher_order_closure & https://github.com/danielhenrymantilla/higher-order-closure.rs
• compiler trait solver gets stuck (infinite loop or tail recursion?), related to GAT and associated type equality constraint rust-lang/rust#134370; Cargo Build gets stuck compiling my code rust-lang/rust#135050

The result is a slightly smarter approach to function argument dispatch so that

• a type parameter doesn't dictate decoder type; e.g. string may be coming from utf8 or utf8view arrays
• a type parameter can be a reference (&str) without major hacks.

[findepi]

Added support for out types:

#[excalibur_function]
fn concat(a: &str, b: &str, out: &mut impl std::fmt::Write) -> Result<ValuePresence> {
    if a.is_empty() && b.is_empty() {
        // Be like Oracle
        return Ok(ValuePresence::Null);
    }
    out.write_str(a)?;
    out.write_str(b)?;
    Ok(ValuePresence::Value)
}

[findepi]

Marked this as non-draft. It would be great to have reviews here.

[findepi]

I bencharked

#[excalibur_function]
fn character_length(s: &str) -> i32 {
    s.chars().count() as i32
}

comparing it with the character_length implementation we ship today

datafusion/datafusion/functions/src/unicode/character_length.rs

Lines 48 to 175 in f31ca5b

	#[derive(Debug)]
	pub struct CharacterLengthFunc {
	signature: Signature,
	aliases: Vec<String>,
	}
	impl Default for CharacterLengthFunc {
	fn default() -> Self {
	Self::new()
	}
	}
	impl CharacterLengthFunc {
	pub fn new() -> Self {
	use DataType::*;
	Self {
	signature: Signature::uniform(
	1,
	vec![Utf8, LargeUtf8, Utf8View],
	Volatility::Immutable,
	),
	aliases: vec![String::from("length"), String::from("char_length")],
	}
	}
	}
	impl ScalarUDFImpl for CharacterLengthFunc {
	fn as_any(&self) -> &dyn Any {
	self
	}
	fn name(&self) -> &str {
	"character_length"
	}
	fn signature(&self) -> &Signature {
	&self.signature
	}
	fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
	utf8_to_int_type(&arg_types[0], "character_length")
	}
	fn invoke_batch(
	&self,
	args: &[ColumnarValue],
	_number_rows: usize,
	) -> Result<ColumnarValue> {
	make_scalar_function(character_length, vec![])(args)
	}
	fn aliases(&self) -> &[String] {
	&self.aliases
	}
	fn documentation(&self) -> Option<&Documentation> {
	self.doc()
	}
	}
	/// Returns number of characters in the string.
	/// character_length('josé') = 4
	/// The implementation counts UTF-8 code points to count the number of characters
	fn character_length(args: &[ArrayRef]) -> Result<ArrayRef> {
	match args[0].data_type() {
	DataType::Utf8 => {
	let string_array = args[0].as_string::<i32>();
	character_length_general::<Int32Type, _>(string_array)
	}
	DataType::LargeUtf8 => {
	let string_array = args[0].as_string::<i64>();
	character_length_general::<Int64Type, _>(string_array)
	}
	DataType::Utf8View => {
	let string_array = args[0].as_string_view();
	character_length_general::<Int32Type, _>(string_array)
	}
	_ => unreachable!("CharacterLengthFunc"),
	}
	}
	fn character_length_general<'a, T, V>(array: V) -> Result<ArrayRef>
	where
	T: ArrowPrimitiveType,
	T::Native: OffsetSizeTrait,
	V: StringArrayType<'a>,
	{
	let mut builder = PrimitiveBuilder::<T>::with_capacity(array.len());
	// String characters are variable length encoded in UTF-8, counting the
	// number of chars requires expensive decoding, however checking if the
	// string is ASCII only is relatively cheap.
	// If strings are ASCII only, count bytes instead.
	let is_array_ascii_only = array.is_ascii();
	if array.null_count() == 0 {
	if is_array_ascii_only {
	for i in 0..array.len() {
	let value = array.value(i);
	builder.append_value(T::Native::usize_as(value.len()));
	}
	} else {
	for i in 0..array.len() {
	let value = array.value(i);
	builder.append_value(T::Native::usize_as(value.chars().count()));
	}
	}
	} else if is_array_ascii_only {
	for i in 0..array.len() {
	if array.is_null(i) {
	builder.append_null();
	} else {
	let value = array.value(i);
	builder.append_value(T::Native::usize_as(value.len()));
	}
	}
	} else {
	for i in 0..array.len() {
	if array.is_null(i) {
	builder.append_null();
	} else {
	let value = array.value(i);
	builder.append_value(T::Native::usize_as(value.chars().count()));
	}
	}
	}
	Ok(Arc::new(builder.finish()) as ArrayRef)
	}

I don't have expectations that generic implementation will be faster than 50x¹ longer, hand-written and optimized code.
Rather, I wanted to check what can be a performance difference between simple 4-liner implementation and one that has dedicated code paths for ascii/unicode, nulls/non-nulls and Utf8/Utf8View combinations. I.e. anything but not "simple".

The standard implementation is a bit faster at short strings

The generic implementation is a bit faster at long ASCII strings

Footnotes

1. including unit tests code, which is natural necessity in a non-trivial implementation, but excluding docs ↩

[jayzhan211]

I don't have expectations that generic implementation will be faster than 50x1 longer, hand-written and optimized code.
Rather, I wanted to check what can be a performance difference between simple 4-liner implementation and one that has dedicated code paths for ascii/unicode, nulls/non-nulls and Utf8/Utf8View combinations. I.e. anything but not "simple".

Interesting, I will review it when I have time

[jayzhan211]

Rather, I wanted to check what can be a performance difference between simple 4-liner implementation and one that has dedicated code paths for ascii/unicode, nulls/non-nulls and Utf8/Utf8View combinations. I.e. anything but not "simple".

I don't find it outperform the existing function 🫤

Is this something possible to support but not yet?

#[excalibur_function]
fn concat(a: &str, b: &str) -> Result<String> {
    if a.is_empty() && b.is_empty() {
        // Be like Oracle
        return Ok(String::new());
    }

    Ok(format!("{}{}", a, b))
}

When I try concat with null, it doesn't output what I expect. How do we handle null in excalibur_function?

#[excalibur_function]
fn concat(a: &str, b: &str, out: &mut impl std::fmt::Write) -> Result<ValuePresence> {
    if a.is_empty() && b.is_empty() {
        // Be like Oracle
        return Ok(ValuePresence::Null);
    }
    out.write_str(a)?;
    out.write_str(b)?;
    Ok(ValuePresence::Value)
}

#[test]
fn test_concat() {
    let udf = concat_udf();

    let n_rows = 3;
    let invoke_args = vec![
        ColumnarValue::Array(Arc::new(StringArray::from(vec![Some("hello"), None, Some("hello")]))),
        ColumnarValue::Scalar(ScalarValue::Utf8(Some("world".to_string()))),
    ];
    let time_start = time::Instant::now();
    let ColumnarValue::Array(result_array) = udf
        .invoke_with_args(ScalarFunctionArgs {
            args: invoke_args,
            number_rows: n_rows,
            return_type: &DataType::Utf8,
        })
        .unwrap()
    else {
        panic!("Expected array result");
    };

    println!(
        "Time elapsed in character_length_array() is: {:?}",
        time_start.elapsed()
    );

    assert_eq!(
        &*result_array,
        &StringArray::from(vec!["helloworld", "world", "helloworld"])
    );
}

Similar to character_length, if I want to return 0 for null case, is it possible?
Either return null or 0 should be possible

#[excalibur_function]
fn character_length(s: &str) -> i32 {
    s.chars().count() as i32
}

[findepi]

Is this something possible to support but not yet?

#[excalibur_function]
fn concat(a: &str, b: &str) -> Result<String> {

support for String return type is easy to add.
It's not there due to inherent performance hit -- the rustc almost never can inline such String allocation

(However, the StringViewArray API is almost imperfect, since it disallows writing to the builder, requiring AsRef<str> and forrcing the allocation.)

When I try concat with null, it doesn't output what I expect. How do we handle null in excalibur_function?

Depends on "what i expect".

Generally the excalibur functions follow the "return null on null input" convention which is default for SQL.
DataFusion's builtin concat is "buggy" because it treats null input arg as empty.
Implementing a function that should be called also for null input requires declaring it, by wrapping relevant arguments in Option<...>.

if I want to return 0 for null case, is it possible?
Either return null or 0 should be possible

yes, just use eg Option<i32> for "nullable i32" function result