Support Arrays for the Map scalar functions #11712
Conversation
dharanad
changed the title
dharanad
changed the title
dharanad
changed the title
|
Sorry, this change took more time than expected. The good thing is I deepened my understanding of Arrow. |
| @@ -202,3 +226,128 @@ fn get_element_type(data_type: &DataType) -> datafusion_common::Result<&DataType | |||
| ), | |||
| } | |||
| } | |||
|
|
|||
| /// Helper function to create MapArray from array of values to support arrays for Map scalar function | |||
|
I hope to review this later today |
| let mut key_array_vec = vec![]; | ||
| let mut value_array_vec = vec![]; | ||
| for (k, v) in keys.iter().zip(values.iter()) { | ||
| running_offset = running_offset.add(O::usize_as(k.len())); | ||
| offset_buffer.push(running_offset); | ||
| key_array_vec.push(k.as_ref()); | ||
| value_array_vec.push(v.as_ref()); | ||
| } | ||
|
|
||
| // concatenate all the arrays | ||
| let flattened_keys = arrow::compute::concat(key_array_vec.as_ref())?; | ||
| if flattened_keys.null_count() > 0 { | ||
| return exec_err!("keys cannot be null"); | ||
| } | ||
| let flattened_values = arrow::compute::concat(value_array_vec.as_ref())?; |
There was a problem hiding this comment.
Could we use arrow::array::MutableArrayData to avoid concatenation here 🤔 ?
There was a problem hiding this comment.
Under the hood, compute::concat uses MutableArrayData. Once this PR gets merged, i can look if i can improve this by using MutableArrayData and not using concat at all.
Co-authored-by: Alex Huang <[email protected]>
|
Thanks @dharanad and @jayzhan211 |
|
Thank You @alamb @goldmedal @jayzhan211 @Weijun-H . |
That describes almost all of my experiences working on this project. I am glad others get the same thrill |
Which issue does this PR close?
Closes #11436
Rationale for this change
What changes are included in this PR?
Add new function to handle support arrays for the Map scalar function.
Are these changes tested?
Existing testcases
Are there any user-facing changes?
No breaking changes