-
Notifications
You must be signed in to change notification settings - Fork 1.3k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Support Arrays for the Map scalar functions #11712
Conversation
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 |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
😍
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks @dharanad for working on this. I have some suggestions about error handling and adding more tests. Others look good to me.
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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Could we use arrow::array::MutableArrayData
to avoid concatenation here 🤔 ?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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]>
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
👍
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