-
Notifications
You must be signed in to change notification settings - Fork 1.1k
Description
Is your feature request related to a problem or challenge? Please describe what you are trying to do.
We are trying to implement more generic "scalar" kernels -- see discussion on #984
If one wants to call eq_scalar on a PrimtiveArray,
let array: Int8Array = vec![1,2,3]
.into_iter()
.map(Some)
.collect();
let comparison = arrow::compute::eq_scalar(&array, 2).unwrap();
let expected: BooleanArray = vec![false, true, false]
.into_iter()
.map(Some)
.collect();
assert_eq!(comparison, expected);In this case, since the data type (Int8Type ) is an associated type of the PrimitiveArray type, Rust can infer the correct type for argument 2 into the 2i8 needed for calling eq_scalar.
However as soon as you have an ArrayRef (aka dynamic array), it becomes challenging. To write a generic "compare this array to a scalar value" we would liek to write something like this:
/// Perform `left == right` operation on an array and a numeric scalar
/// value. Supports PrimtiveArrays, and DictionaryArrays that have primitive values
pub fn eq_dyn_scalar<T>(
left: &dyn Array,
right: T::Native,
) -> Result<BooleanArray>
where
T: ArrowNumericType,
{
...
}However, to actually call this function, the type of right needs to match T, but there is a potentially infinite number of concrete T types that could have a matching associated type Native, as explained by @shepmaster in #984 (comment)
Practically that means that all calls to eq_dyn_scalar would require type annotations which matched the values of the array, which largely defeats the purpose of eq_dyn_scalar
let array = builder.finish();
// still need the UInt8Type annotations
let a_eq = eq_dyn_scalar::<UInt8Type>(&array, 123u8).unwrap();
assert_eq!(
a_eq,
BooleanArray::from(vec![Some(true), None, Some(false)])
);Describe the solution you'd like
So the proposal from @shepmaster at #984 (comment) is to add a trait that maps in the reverse direction, so we could have something like:
struct UInt8Type;
struct Int64Type;
struct DictionaryArray<K>(K);
// ----
trait ArrowNumericType {
type Native;
}
trait IntoArrowNumericType {
type Arrow: ArrowNumericType<Native = Self>;
}
// ----
impl ArrowNumericType for UInt8Type {
type Native = u8;
}
impl IntoArrowNumericType for u8 {
type Arrow = UInt8Type;
}
impl ArrowNumericType for Int64Type {
type Native = i64;
}
impl IntoArrowNumericType for i64 {
type Arrow = Int64Type;
}
// ----
fn eq_dict_scalar<K, T>(left: &DictionaryArray<K>, right: T)
where
K: ArrowNumericType,
T: IntoArrowNumericType,
{
todo!()
}
// ----
fn usage<K>(d: &DictionaryArray<K>)
where
K: ArrowNumericType,
{
eq_dict_scalar(d, 8u8);
eq_dict_scalar(d, -64i64);
}Describe alternatives you've considered
A clear and concise description of any alternative solutions or features you've considered.
Additional context
Add any other context or screenshots about the feature request here.