inaos / iron-array

2 stars 0 forks source link

Implement the where function in expressions #536

Open FrancescAlted opened 2 years ago

FrancescAlted commented 2 years ago

This is required by Haizea. Here it is how numexpr implements it:

https://numexpr.readthedocs.io/projects/NumExpr3/en/latest/user_guide.html#supported-functions

Our implementation should be similar. Perhaps a good approach would be to implement it as an scalar UDF, e.g.:

@udf.scalar(lib="lib")
def wheref64(cond: bool, val1: udf.float64, val2: udf.float64) -> float:
    return val1 if cond else val2

@udf.scalar(lib="lib")
def wheref32(cond: bool, val1: udf.float32, val2: udf.float32) -> float:
    return val1 if cond else val2

and our expression would just be:

"a + lib.where(bool_arr, vals1, vals2)"
FrancescAlted commented 2 years ago

This requires mixing bool and float arrays. @jdavid do you think this is feasible?

FrancescAlted commented 2 years ago

Note that the existing support for the extended slice support with bools in https://github.com/inaos/iron-array-python/pull/164 (https://github.com/inaos/iron-array-python/pull/164/files#diff-1cd4d40b944c51195b323743ae454471a5688dcfc370096be3907d3a28e8ae17R46) should be enough for the time being, so no hurry in this one.