ARROW-10640: [C++] A "where" kernel to combine two arrays based on a mask #10377
Conversation
lidavidm
left a comment
There was a problem hiding this comment.
Just some quick initial feedback. I think the overall approach looks fine. Testing this thoroughly will be important.
For a different PR I added a DatumFromJSON to make it easier to test functions that accept a lot of permutations of scalar/array arguments. I could pull that out separately if that would make testing easier here. (That way you don't need a lot of overloads and boilerplate to test all the possible input combinations.)
| return Status::OK(); | ||
| } | ||
|
|
||
| template <typename Type, bool swap = false, typename Enable = void> |
There was a problem hiding this comment.
My idea is to reuse the impl for the cases like, cond, left: Array, right: Scalar and cond, left: Sca;ar, right: Array. In the second scenario, I can swap left and right and invert the cond without changing the loop mechanism.
|
|
||
| // nulls will be promoted as follows | ||
| // cond.val && (cond.data && left.val || ~cond.data && right.val) | ||
| Status promote_nulls(KernelContext* ctx, const ArrayData& cond, const ArrayData& left, |
There was a problem hiding this comment.
nit: use CamelCase (PromoteNulls)
| template <typename Type> | ||
| struct ResolveExec { | ||
| static Status Exec(KernelContext* ctx, const ExecBatch& batch, Datum* out) { | ||
| if (batch.length == 0) return Status::OK(); |
There was a problem hiding this comment.
Do you really need to check for this? You can assume you have three arguments.
| }; | ||
|
|
||
| template <typename Type, bool swap> | ||
| struct IfElseFunctor<Type, swap, enable_if_t<is_boolean_type<Type>::value>> { |
There was a problem hiding this comment.
nit: type_traits.h already has an enable_if_boolean
There was a problem hiding this comment.
(same with enable_if_number etc)
| const T* left_data = left.GetValues<T>(1); | ||
| int64_t offset = cond.offset; | ||
|
|
||
| // todo this can be improved by intrinsics. ex: _mm*_mask_store_e* (vmovdqa*) |
There was a problem hiding this comment.
I would assume memcpy already does this for you.
There was a problem hiding this comment.
Do you mean, load with mask?
There was a problem hiding this comment.
Sorry, I think I misunderstood the optimization you're thinking about. How would SIMD help here?
There was a problem hiding this comment.
Say, you first copy right to ouput. Then, cond becomes a mask to store left onto output. For that there are specialized SIMD instructions.
https://software.intel.com/sites/landingpage/IntrinsicsGuide/#cats=Store&text=mask_store&expand=5564,5566
So, we can drop the BitBlock objects, and remove all the loops and memcpy inside the while loop. We'd have to handle the alignment though.
| } | ||
| }; | ||
|
|
||
| void AddPrimitiveKernels(const std::shared_ptr<ScalarFunction>& scalar_function, |
There was a problem hiding this comment.
A lot of these function names are a little generic, can we rename them to reflect that they're for IfElse? We might end up putting other similar kernels here.
That would be great @lidavidm . If you could send me the PR# that would be great! :-) |
See #10386 |
|
I am closing this PR because there are some major refactors and it would be better to review them fresh. |
Adding a preliminary impl for an
if_else(cond: Datum, left: Datum, right: Datum)function. It works as follows,nullvalues will be promoted to the output.Current commit only provides only Array-Array-Array kernels for the function.