You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@arrow.apache.org by "Weijun-H (via GitHub)" <gi...@apache.org> on 2024/01/04 09:52:22 UTC
[PR] feat: support `array_resize` [arrow-datafusion]
Weijun-H opened a new pull request, #8744:
URL: https://github.com/apache/arrow-datafusion/pull/8744
## Which issue does this PR close?
<!--
We generally require a GitHub issue to be filed for all bug fixes and enhancements and this helps us generate change logs for our releases. You can link an issue to this PR using the GitHub syntax. For example `Closes #123` indicates that this PR will close issue #123.
-->
Closes #7194
## Rationale for this change
<!--
Why are you proposing this change? If this is already explained clearly in the issue then this section is not needed.
Explaining clearly why changes are proposed helps reviewers understand your changes and offer better suggestions for fixes.
-->
## What changes are included in this PR?
<!--
There is no need to duplicate the description in the issue here but it is sometimes worth providing a summary of the individual changes in this PR.
-->
## Are these changes tested?
<!--
We typically require tests for all PRs in order to:
1. Prevent the code from being accidentally broken by subsequent changes
2. Serve as another way to document the expected behavior of the code
If tests are not included in your PR, please explain why (for example, are they covered by existing tests)?
-->
## Are there any user-facing changes?
<!--
If there are user-facing changes then we may require documentation to be updated before approving the PR.
-->
<!--
If there are any breaking changes to public APIs, please add the `api change` label.
-->
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: support `array_resize` [arrow-datafusion]
Posted by "jayzhan211 (via GitHub)" <gi...@apache.org>.
jayzhan211 commented on code in PR #8744:
URL: https://github.com/apache/arrow-datafusion/pull/8744#discussion_r1441689180
##########
package.json:
##########
@@ -0,0 +1 @@
+{}
Review Comment:
🤔 Why are these 3 files needed
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: support `array_resize` [arrow-datafusion]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on code in PR #8744:
URL: https://github.com/apache/arrow-datafusion/pull/8744#discussion_r1444074085
##########
datafusion/physical-expr/src/array_expressions.rs:
##########
@@ -2560,6 +2562,101 @@ pub fn array_distinct(args: &[ArrayRef]) -> Result<ArrayRef> {
}
}
+pub fn array_resize(arg: &[ArrayRef]) -> Result<ArrayRef> {
+ if arg.len() < 2 || arg.len() > 3 {
+ return exec_err!("array_resize needs two or three arguments");
+ }
+
+ let new_len = as_int64_array(&arg[1])?;
+ let new_element = if arg.len() == 3 {
+ Some(arg[2].clone())
+ } else {
+ None
+ };
+
+ match &arg[0].data_type() {
+ DataType::List(field) => {
+ let array = as_list_array(&arg[0])?;
+ general_list_resize::<i32>(array, new_len, field, new_element)
+ }
+ DataType::LargeList(field) => {
+ let array = as_large_list_array(&arg[0])?;
+ general_list_resize::<i64>(array, new_len, field, new_element)
+ }
+ _ => exec_err!("array_resize only support list array"),
Review Comment:
```suggestion
_ => exec_err!("array_resize only support LIST datatype"),
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: support `array_resize` [arrow-datafusion]
Posted by "jayzhan211 (via GitHub)" <gi...@apache.org>.
jayzhan211 commented on code in PR #8744:
URL: https://github.com/apache/arrow-datafusion/pull/8744#discussion_r1447426362
##########
datafusion/physical-expr/src/array_expressions.rs:
##########
@@ -2611,6 +2618,98 @@ pub fn array_distinct(args: &[ArrayRef]) -> Result<ArrayRef> {
}
}
+/// array_resize SQL function
+pub fn array_resize(arg: &[ArrayRef]) -> Result<ArrayRef> {
+ if arg.len() < 2 || arg.len() > 3 {
+ return exec_err!("array_resize needs two or three arguments");
+ }
+
+ let new_len = as_int64_array(&arg[1])?;
+ let new_element = if arg.len() == 3 {
+ Some(arg[2].clone())
+ } else {
+ None
+ };
+
+ match &arg[0].data_type() {
+ DataType::List(field) => {
+ let array = as_list_array(&arg[0])?;
+ general_list_resize::<i32>(array, new_len, field, new_element)
+ }
+ DataType::LargeList(field) => {
+ let array = as_large_list_array(&arg[0])?;
+ general_list_resize::<i64>(array, new_len, field, new_element)
+ }
+ array_type => exec_err!("array_resize does not support type '{array_type:?}'."),
+ }
+}
+
+/// array_resize keep the original array and append the default element to the end
+fn general_list_resize<O: OffsetSizeTrait>(
+ array: &GenericListArray<O>,
+ count_array: &Int64Array,
+ field: &FieldRef,
+ default_element: Option<ArrayRef>,
+) -> Result<ArrayRef>
+where
+ O: TryInto<i64>,
+{
+ let data_type = array.value_type();
+
+ let values = array.values();
+ let original_data = values.to_data();
+
+ // create default element array
+ let default_element = if let Some(default_element) = default_element {
+ default_element
+ } else {
+ let null_scalar = ScalarValue::try_from(&data_type)?;
+ null_scalar.to_array_of_size(original_data.len())?
+ };
+ let default_value_data = default_element.to_data();
+
+ // create a mutable array to store the original data
+ let capacity = Capacities::Array(original_data.len() + default_value_data.len());
+ let mut offsets = vec![O::usize_as(0)];
+ let mut mutable = MutableArrayData::with_capacities(
+ vec![&original_data, &default_value_data],
+ false,
+ capacity,
+ );
+
+ for (row_index, offset_window) in array.offsets().windows(2).enumerate() {
+ let count = count_array.value(row_index).to_usize().ok_or_else(|| {
+ exec_datafusion_err!("array_resize: failed to convert size to usize")
Review Comment:
does type conversion more like internal_err?
`wasn't expected/anticipated by the implementation
and that is most likely a bug (the error message even encourages users
to open a bug report)`
unlike array.len() != number, which is `error` from user.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: support `array_resize` [arrow-datafusion]
Posted by "Weijun-H (via GitHub)" <gi...@apache.org>.
Weijun-H commented on PR #8744:
URL: https://github.com/apache/arrow-datafusion/pull/8744#issuecomment-1878117609
> array_concat(array_slice(array, 0, min(array_length, count)), array_repeat(default_array, max(count-array_length), 0))
I think this approach is not a better idea, as supporting the column for `max(count-array_length)` in an argument is hard. However, the MutableArray could be another direction to consider, just like you did for other functions. I will try it as a follow-up PR.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: support `array_resize` [arrow-datafusion]
Posted by "jayzhan211 (via GitHub)" <gi...@apache.org>.
jayzhan211 commented on code in PR #8744:
URL: https://github.com/apache/arrow-datafusion/pull/8744#discussion_r1442517967
##########
datafusion/common/src/utils.rs:
##########
@@ -492,6 +496,58 @@ pub fn list_ndims(data_type: &DataType) -> u64 {
}
}
+/// Create an new empty array based on the given data type.
+pub fn empty_list(data_type: &DataType) -> Result<ArrayRef> {
+ match data_type {
+ DataType::Boolean => Ok(Arc::new(BooleanArray::from(vec![None]))),
+ DataType::UInt8 => Ok(Arc::new(UInt8Array::from(vec![None]))),
+ DataType::UInt16 => Ok(Arc::new(UInt16Array::from(vec![None]))),
+ DataType::UInt32 => Ok(Arc::new(UInt32Array::from(vec![None]))),
+ DataType::UInt64 => Ok(Arc::new(UInt64Array::from(vec![None]))),
+ DataType::Int8 => Ok(Arc::new(Int8Array::from(vec![None]))),
+ DataType::Int16 => Ok(Arc::new(Int16Array::from(vec![None]))),
+ DataType::Int32 => Ok(Arc::new(Int32Array::from(vec![None]))),
+ DataType::Int64 => Ok(Arc::new(Int64Array::from(vec![None]))),
+ DataType::Float32 => Ok(Arc::new(Float32Array::from(vec![None]))),
+ DataType::Float64 => Ok(Arc::new(Float64Array::from(vec![None]))),
+ DataType::List(field) => {
+ let value = empty_list(field.data_type())?;
+ Ok(Arc::new(ListArray::try_new(
+ field.clone(),
+ OffsetBuffer::new(vec![0, 0].into()),
+ value,
+ None,
+ )?))
+ }
+ DataType::LargeList(field) => {
+ let value = empty_list(field.data_type())?;
+
+ Ok(Arc::new(LargeListArray::try_new(
+ field.clone(),
+ OffsetBuffer::new(vec![0, 0].into()),
+ value,
+ None,
+ )?))
+ }
+ DataType::Null => Ok(Arc::new(NullArray::new(1))),
+ _ => _internal_err!("Unsupported data type for empty list: {:?}", data_type),
+ }
+}
+
+// DataType::Boolean,
Review Comment:
I think you forgot to cleanup this
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: support `array_resize` [arrow-datafusion]
Posted by "jayzhan211 (via GitHub)" <gi...@apache.org>.
jayzhan211 commented on code in PR #8744:
URL: https://github.com/apache/arrow-datafusion/pull/8744#discussion_r1446966314
##########
datafusion/physical-expr/src/array_expressions.rs:
##########
@@ -2611,6 +2618,95 @@ pub fn array_distinct(args: &[ArrayRef]) -> Result<ArrayRef> {
}
}
+/// array_resize SQL function
+pub fn array_resize(arg: &[ArrayRef]) -> Result<ArrayRef> {
+ if arg.len() < 2 || arg.len() > 3 {
+ return exec_err!("array_resize needs two or three arguments");
+ }
+
+ let new_len = as_int64_array(&arg[1])?;
+ let new_element = if arg.len() == 3 {
+ Some(arg[2].clone())
+ } else {
+ None
+ };
+
+ match &arg[0].data_type() {
+ DataType::List(field) => {
+ let array = as_list_array(&arg[0])?;
+ general_list_resize::<i32>(array, new_len, field, new_element)
+ }
+ DataType::LargeList(field) => {
+ let array = as_large_list_array(&arg[0])?;
+ general_list_resize::<i64>(array, new_len, field, new_element)
+ }
+ array_type => exec_err!("array_resize does not support type '{array_type:?}'."),
+ }
+}
+
+/// array_resize keep the original array and append the default element to the end
+fn general_list_resize<O: OffsetSizeTrait>(
+ array: &GenericListArray<O>,
+ count_array: &Int64Array,
+ field: &FieldRef,
+ default_element: Option<ArrayRef>,
+) -> Result<ArrayRef>
+where
+ O: TryInto<i64>,
+{
+ let mut default_element_array = vec![];
+
+ let data_type = array.value_type();
+ let default_element = if let Some(default_element) = default_element {
+ default_element
+ } else {
+ let null_scalar = ScalarValue::try_from(&data_type)?;
+ null_scalar.to_array_of_size(1)?
Review Comment:
Should this be `to_array_of_size(row_len)`?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: support `array_resize` [arrow-datafusion]
Posted by "jayzhan211 (via GitHub)" <gi...@apache.org>.
jayzhan211 commented on PR #8744:
URL: https://github.com/apache/arrow-datafusion/pull/8744#issuecomment-1884252351
> > array_concat(array_slice(array, 0, min(array_length, count)), array_repeat(default_array, max(count-array_length), 0))
>
> I think this approach is not a better idea, as supporting the column for `max(count-array_length)` in an argument is hard. However, the MutableArray could be another direction to consider, just like you did for other functions. I will try it as a follow-up PR. @jayzhan211
I think manually create array for `max(count - array_length, 0)` is just one iteration for the array. The same as `min(array_length, count)`. After creating those two arrays, we can apply `array_concat(array_slice, array_repeat)` easily.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: support `array_resize` [arrow-datafusion]
Posted by "jayzhan211 (via GitHub)" <gi...@apache.org>.
jayzhan211 commented on code in PR #8744:
URL: https://github.com/apache/arrow-datafusion/pull/8744#discussion_r1441655668
##########
datafusion/physical-expr/src/array_expressions.rs:
##########
@@ -2560,6 +2561,103 @@ pub fn array_distinct(args: &[ArrayRef]) -> Result<ArrayRef> {
}
}
+pub fn array_resize(arg: &[ArrayRef]) -> Result<ArrayRef> {
+ if arg.len() < 2 || arg.len() > 3 {
+ return exec_err!("array_resize needs two or three arguments");
+ }
+
+ let new_len = as_int64_array(&arg[1])?;
+ let new_element = if arg.len() == 3 {
+ Some(arg[2].clone())
+ } else {
+ None
+ };
+
+ match &arg[0].data_type() {
+ DataType::List(field) => {
+ let array = as_list_array(&arg[0])?;
+ general_list_resize::<i32>(array, new_len, field, new_element)
+ }
+ DataType::LargeList(field) => {
+ let array = as_large_list_array(&arg[0])?;
+ general_list_resize::<i64>(array, new_len, field, new_element)
+ }
+ _ => internal_err!("array_resize only support list array"),
+ }
+}
+
+fn general_list_resize<O: OffsetSizeTrait>(
+ array: &GenericListArray<O>,
+ count_array: &Int64Array,
+ field: &FieldRef,
+ new_element: Option<ArrayRef>,
Review Comment:
`default_element` might be more concise?
##########
datafusion/physical-expr/src/array_expressions.rs:
##########
@@ -2560,6 +2561,103 @@ pub fn array_distinct(args: &[ArrayRef]) -> Result<ArrayRef> {
}
}
+pub fn array_resize(arg: &[ArrayRef]) -> Result<ArrayRef> {
+ if arg.len() < 2 || arg.len() > 3 {
+ return exec_err!("array_resize needs two or three arguments");
+ }
+
+ let new_len = as_int64_array(&arg[1])?;
+ let new_element = if arg.len() == 3 {
+ Some(arg[2].clone())
+ } else {
+ None
+ };
+
+ match &arg[0].data_type() {
+ DataType::List(field) => {
+ let array = as_list_array(&arg[0])?;
+ general_list_resize::<i32>(array, new_len, field, new_element)
+ }
+ DataType::LargeList(field) => {
+ let array = as_large_list_array(&arg[0])?;
+ general_list_resize::<i64>(array, new_len, field, new_element)
+ }
+ _ => internal_err!("array_resize only support list array"),
+ }
+}
+
+fn general_list_resize<O: OffsetSizeTrait>(
+ array: &GenericListArray<O>,
+ count_array: &Int64Array,
+ field: &FieldRef,
+ new_element: Option<ArrayRef>,
+) -> Result<ArrayRef> {
+ let mut offsets = vec![O::usize_as(0)];
+ let mut new_arrays = vec![];
+
+ let dt = array.value_type();
+ let converter = RowConverter::new(vec![SortField::new(dt.clone())])?;
+ let new_element = if let Some(new_element) = new_element {
+ new_element
+ } else {
+ empty_list(&dt)?
+ };
+ let rows = converter.convert_columns(&[new_element.clone()])?;
+
+ for (index, arr) in array.iter().enumerate() {
+ let count = count_array.value(index).to_usize().ok_or_else(|| {
+ DataFusionError::Execution(
Review Comment:
nit: `exec_err`
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: support `array_resize` [arrow-datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb merged PR #8744:
URL: https://github.com/apache/arrow-datafusion/pull/8744
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: support `array_resize` [arrow-datafusion]
Posted by "jayzhan211 (via GitHub)" <gi...@apache.org>.
jayzhan211 commented on code in PR #8744:
URL: https://github.com/apache/arrow-datafusion/pull/8744#discussion_r1442519085
##########
datafusion/physical-expr/src/array_expressions.rs:
##########
@@ -2560,6 +2562,101 @@ pub fn array_distinct(args: &[ArrayRef]) -> Result<ArrayRef> {
}
}
+pub fn array_resize(arg: &[ArrayRef]) -> Result<ArrayRef> {
+ if arg.len() < 2 || arg.len() > 3 {
+ return exec_err!("array_resize needs two or three arguments");
+ }
+
+ let new_len = as_int64_array(&arg[1])?;
+ let new_element = if arg.len() == 3 {
+ Some(arg[2].clone())
+ } else {
+ None
+ };
+
+ match &arg[0].data_type() {
+ DataType::List(field) => {
+ let array = as_list_array(&arg[0])?;
+ general_list_resize::<i32>(array, new_len, field, new_element)
+ }
+ DataType::LargeList(field) => {
+ let array = as_large_list_array(&arg[0])?;
+ general_list_resize::<i64>(array, new_len, field, new_element)
+ }
+ _ => exec_err!("array_resize only support list array"),
+ }
+}
+
+fn general_list_resize<O: OffsetSizeTrait>(
+ array: &GenericListArray<O>,
+ count_array: &Int64Array,
+ field: &FieldRef,
+ default_element: Option<ArrayRef>,
+) -> Result<ArrayRef> {
+ let mut offsets = vec![O::usize_as(0)];
+ let mut new_arrays = vec![];
+
+ let dt = array.value_type();
+ let converter = RowConverter::new(vec![SortField::new(dt.clone())])?;
+ let default_element = if let Some(default_element) = default_element {
+ default_element
+ } else {
+ empty_list(&dt)?
+ };
+ let rows = converter.convert_columns(&[default_element.clone()])?;
+
+ for (index, arr) in array.iter().enumerate() {
+ let count = count_array.value(index).to_usize().ok_or_else(|| {
+ exec_datafusion_err!("array_resize: failed to convert size to usize")
+ })?;
+ let row = if rows.size() > 0 {
Review Comment:
maybe also rename row to `default_value`?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: support `array_resize` [arrow-datafusion]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on PR #8744:
URL: https://github.com/apache/arrow-datafusion/pull/8744#issuecomment-1880178023
thanks @Weijun-H I'll review it tomorrow more closely if no one else beats me.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: support `array_resize` [arrow-datafusion]
Posted by "jayzhan211 (via GitHub)" <gi...@apache.org>.
jayzhan211 commented on code in PR #8744:
URL: https://github.com/apache/arrow-datafusion/pull/8744#discussion_r1446962462
##########
datafusion/physical-expr/src/array_expressions.rs:
##########
@@ -2611,6 +2618,95 @@ pub fn array_distinct(args: &[ArrayRef]) -> Result<ArrayRef> {
}
}
+/// array_resize SQL function
+pub fn array_resize(arg: &[ArrayRef]) -> Result<ArrayRef> {
+ if arg.len() < 2 || arg.len() > 3 {
+ return exec_err!("array_resize needs two or three arguments");
+ }
+
+ let new_len = as_int64_array(&arg[1])?;
+ let new_element = if arg.len() == 3 {
+ Some(arg[2].clone())
+ } else {
+ None
+ };
+
+ match &arg[0].data_type() {
+ DataType::List(field) => {
+ let array = as_list_array(&arg[0])?;
+ general_list_resize::<i32>(array, new_len, field, new_element)
+ }
+ DataType::LargeList(field) => {
+ let array = as_large_list_array(&arg[0])?;
+ general_list_resize::<i64>(array, new_len, field, new_element)
+ }
+ array_type => exec_err!("array_resize does not support type '{array_type:?}'."),
+ }
+}
+
+/// array_resize keep the original array and append the default element to the end
+fn general_list_resize<O: OffsetSizeTrait>(
+ array: &GenericListArray<O>,
+ count_array: &Int64Array,
+ field: &FieldRef,
+ default_element: Option<ArrayRef>,
+) -> Result<ArrayRef>
+where
+ O: TryInto<i64>,
+{
+ let mut default_element_array = vec![];
+
+ let data_type = array.value_type();
+ let default_element = if let Some(default_element) = default_element {
+ default_element
+ } else {
+ let null_scalar = ScalarValue::try_from(&data_type)?;
+ null_scalar.to_array_of_size(1)?
+ };
+
+ // create a mutable array to store the original data
+ let values = array.values();
+ let original_data = values.to_data();
+ let capacity = Capacities::Array(original_data.len());
+ let mut offsets = vec![O::usize_as(0)];
+ let mut mutable =
+ MutableArrayData::with_capacities(vec![&original_data], false, capacity);
+
+ for (row_index, offset_window) in array.offsets().windows(2).enumerate() {
+ let count = count_array.value(row_index).to_usize().ok_or_else(|| {
+ exec_datafusion_err!("array_resize: failed to convert size to usize")
+ })?;
+ let count = O::usize_as(count);
+ let start = offset_window[0];
+ let end = if start + count > offset_window[1] {
+ let value = (start + count - offset_window[1]).try_into().map_err(|_| {
+ exec_datafusion_err!("array_resize: failed to convert size to i64")
+ })?;
+ default_element_array.push(Some(value));
+ offset_window[1]
+ } else {
+ default_element_array.push(None);
+ start + count
+ };
+ mutable.extend(0, (start).to_usize().unwrap(), (end).to_usize().unwrap());
+ offsets.push(offsets[row_index] + end - start);
+ }
+
+ let default_element_array = Arc::new(Int64Array::from(default_element_array));
+ let default_element_array =
+ general_repeat::<O>(&default_element, &default_element_array)?;
+
+ let data = mutable.freeze();
+ let original_part = Arc::new(GenericListArray::<O>::try_new(
+ field.clone(),
+ OffsetBuffer::<O>::new(offsets.into()),
+ arrow_array::make_array(data),
+ None,
+ )?);
+
+ array_concat(&[original_part, default_element_array])
Review Comment:
We can avoid array_concat if we construct with MutableArray,
```rust
let mut mutable =
MutableArrayData::with_capacities(vec![&original_data, default_value_data], false, capacity);
```
Then, extend(1, ...) for default value
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: support `array_resize` [arrow-datafusion]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on code in PR #8744:
URL: https://github.com/apache/arrow-datafusion/pull/8744#discussion_r1445010056
##########
datafusion/physical-expr/src/array_expressions.rs:
##########
@@ -2560,6 +2562,101 @@ pub fn array_distinct(args: &[ArrayRef]) -> Result<ArrayRef> {
}
}
+pub fn array_resize(arg: &[ArrayRef]) -> Result<ArrayRef> {
+ if arg.len() < 2 || arg.len() > 3 {
+ return exec_err!("array_resize needs two or three arguments");
+ }
+
+ let new_len = as_int64_array(&arg[1])?;
+ let new_element = if arg.len() == 3 {
+ Some(arg[2].clone())
+ } else {
+ None
+ };
+
+ match &arg[0].data_type() {
+ DataType::List(field) => {
+ let array = as_list_array(&arg[0])?;
+ general_list_resize::<i32>(array, new_len, field, new_element)
+ }
+ DataType::LargeList(field) => {
+ let array = as_large_list_array(&arg[0])?;
+ general_list_resize::<i64>(array, new_len, field, new_element)
+ }
+ array_type => exec_err!("array_resize does not support type '{array_type:?}'."),
+ }
+}
+
+fn general_list_resize<O: OffsetSizeTrait>(
+ array: &GenericListArray<O>,
+ count_array: &Int64Array,
+ field: &FieldRef,
+ default_element: Option<ArrayRef>,
+) -> Result<ArrayRef> {
+ let mut offsets = vec![O::usize_as(0)];
+ let mut new_arrays = vec![];
+
+ let dt = array.value_type();
+ let converter = RowConverter::new(vec![SortField::new(dt.clone())])?;
+ let default_element = if let Some(default_element) = default_element {
+ default_element
+ } else {
+ empty_list(&dt)?
+ };
+ let rows = converter.convert_columns(&[default_element.clone()])?;
+
+ for (index, arr) in array.iter().enumerate() {
+ let count = count_array.value(index).to_usize().ok_or_else(|| {
+ exec_datafusion_err!("array_resize: failed to convert size to usize")
+ })?;
+ let default_value = if rows.size() > 0 {
+ rows.row(index)
+ } else {
+ rows.row(0)
+ };
+ match arr {
+ Some(arr) => {
+ let values = converter.convert_columns(&[arr])?;
+ let rows = values.iter().collect::<Vec<_>>();
+
+ let remain_count = min(rows.len(), count);
+ let mut rows = rows[..remain_count].to_vec();
+ let new_rows = vec![&default_value; count - remain_count];
+ rows.extend(new_rows);
+
+ let last_offset = offsets.last().copied().unwrap();
Review Comment:
Please handle the potential error on `.unwrap`
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: support `array_resize` [arrow-datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #8744:
URL: https://github.com/apache/arrow-datafusion/pull/8744#discussion_r1446623400
##########
datafusion/physical-expr/src/array_expressions.rs:
##########
@@ -2560,6 +2562,101 @@ pub fn array_distinct(args: &[ArrayRef]) -> Result<ArrayRef> {
}
}
+pub fn array_resize(arg: &[ArrayRef]) -> Result<ArrayRef> {
+ if arg.len() < 2 || arg.len() > 3 {
+ return exec_err!("array_resize needs two or three arguments");
+ }
+
+ let new_len = as_int64_array(&arg[1])?;
+ let new_element = if arg.len() == 3 {
+ Some(arg[2].clone())
+ } else {
+ None
+ };
+
+ match &arg[0].data_type() {
+ DataType::List(field) => {
+ let array = as_list_array(&arg[0])?;
+ general_list_resize::<i32>(array, new_len, field, new_element)
+ }
+ DataType::LargeList(field) => {
+ let array = as_large_list_array(&arg[0])?;
+ general_list_resize::<i64>(array, new_len, field, new_element)
+ }
+ array_type => exec_err!("array_resize does not support type '{array_type:?}'."),
+ }
+}
+
+fn general_list_resize<O: OffsetSizeTrait>(
+ array: &GenericListArray<O>,
+ count_array: &Int64Array,
+ field: &FieldRef,
+ default_element: Option<ArrayRef>,
+) -> Result<ArrayRef> {
+ let mut offsets = vec![O::usize_as(0)];
+ let mut new_arrays = vec![];
+
+ let dt = array.value_type();
+ let converter = RowConverter::new(vec![SortField::new(dt.clone())])?;
+ let default_element = if let Some(default_element) = default_element {
+ default_element
+ } else {
+ empty_list(&dt)?
+ };
+ let rows = converter.convert_columns(&[default_element.clone()])?;
+
+ for (index, arr) in array.iter().enumerate() {
Review Comment:
Instead of `RowConverter` which will copy the data twice, you can probably use `MutableArrayData` to create the output directly
##########
datafusion/physical-expr/src/array_expressions.rs:
##########
@@ -2560,6 +2562,101 @@ pub fn array_distinct(args: &[ArrayRef]) -> Result<ArrayRef> {
}
}
+pub fn array_resize(arg: &[ArrayRef]) -> Result<ArrayRef> {
+ if arg.len() < 2 || arg.len() > 3 {
+ return exec_err!("array_resize needs two or three arguments");
+ }
+
+ let new_len = as_int64_array(&arg[1])?;
+ let new_element = if arg.len() == 3 {
+ Some(arg[2].clone())
+ } else {
+ None
+ };
+
+ match &arg[0].data_type() {
+ DataType::List(field) => {
+ let array = as_list_array(&arg[0])?;
+ general_list_resize::<i32>(array, new_len, field, new_element)
+ }
+ DataType::LargeList(field) => {
+ let array = as_large_list_array(&arg[0])?;
+ general_list_resize::<i64>(array, new_len, field, new_element)
+ }
+ array_type => exec_err!("array_resize does not support type '{array_type:?}'."),
+ }
+}
+
+fn general_list_resize<O: OffsetSizeTrait>(
+ array: &GenericListArray<O>,
+ count_array: &Int64Array,
+ field: &FieldRef,
+ default_element: Option<ArrayRef>,
+) -> Result<ArrayRef> {
+ let mut offsets = vec![O::usize_as(0)];
+ let mut new_arrays = vec![];
+
+ let dt = array.value_type();
+ let converter = RowConverter::new(vec![SortField::new(dt.clone())])?;
+ let default_element = if let Some(default_element) = default_element {
Review Comment:
I think if you used ScalarValue here (rather than Option<ArrayRef>) you can probably avoid empty_list entirely
##########
datafusion/common/src/utils.rs:
##########
@@ -492,6 +496,44 @@ pub fn list_ndims(data_type: &DataType) -> u64 {
}
}
+/// Create an new empty array based on the given data type.
Review Comment:
Technically this list isn't empty, is it ? It contains a single `Null` value
Also, it seems like this handles all data types, not just a list. I wonder if the same thing could be done by making a null scalar like
```rust
let null_scalar = ScalarValue::try_from(&data_type)?;
null_scalar.to_array_of_size(1)
```
🤔
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: support `array_resize` [arrow-datafusion]
Posted by "jayzhan211 (via GitHub)" <gi...@apache.org>.
jayzhan211 commented on code in PR #8744:
URL: https://github.com/apache/arrow-datafusion/pull/8744#discussion_r1446964517
##########
datafusion/physical-expr/src/array_expressions.rs:
##########
@@ -2611,6 +2618,95 @@ pub fn array_distinct(args: &[ArrayRef]) -> Result<ArrayRef> {
}
}
+/// array_resize SQL function
+pub fn array_resize(arg: &[ArrayRef]) -> Result<ArrayRef> {
+ if arg.len() < 2 || arg.len() > 3 {
+ return exec_err!("array_resize needs two or three arguments");
+ }
+
+ let new_len = as_int64_array(&arg[1])?;
+ let new_element = if arg.len() == 3 {
+ Some(arg[2].clone())
+ } else {
+ None
+ };
+
+ match &arg[0].data_type() {
+ DataType::List(field) => {
+ let array = as_list_array(&arg[0])?;
+ general_list_resize::<i32>(array, new_len, field, new_element)
+ }
+ DataType::LargeList(field) => {
+ let array = as_large_list_array(&arg[0])?;
+ general_list_resize::<i64>(array, new_len, field, new_element)
+ }
+ array_type => exec_err!("array_resize does not support type '{array_type:?}'."),
+ }
+}
+
+/// array_resize keep the original array and append the default element to the end
+fn general_list_resize<O: OffsetSizeTrait>(
+ array: &GenericListArray<O>,
+ count_array: &Int64Array,
+ field: &FieldRef,
+ default_element: Option<ArrayRef>,
+) -> Result<ArrayRef>
+where
+ O: TryInto<i64>,
+{
+ let mut default_element_array = vec![];
+
+ let data_type = array.value_type();
+ let default_element = if let Some(default_element) = default_element {
+ default_element
+ } else {
+ let null_scalar = ScalarValue::try_from(&data_type)?;
+ null_scalar.to_array_of_size(1)?
+ };
+
+ // create a mutable array to store the original data
+ let values = array.values();
+ let original_data = values.to_data();
+ let capacity = Capacities::Array(original_data.len());
+ let mut offsets = vec![O::usize_as(0)];
+ let mut mutable =
+ MutableArrayData::with_capacities(vec![&original_data], false, capacity);
+
+ for (row_index, offset_window) in array.offsets().windows(2).enumerate() {
+ let count = count_array.value(row_index).to_usize().ok_or_else(|| {
+ exec_datafusion_err!("array_resize: failed to convert size to usize")
+ })?;
+ let count = O::usize_as(count);
+ let start = offset_window[0];
+ let end = if start + count > offset_window[1] {
+ let value = (start + count - offset_window[1]).try_into().map_err(|_| {
+ exec_datafusion_err!("array_resize: failed to convert size to i64")
+ })?;
+ default_element_array.push(Some(value));
+ offset_window[1]
+ } else {
+ default_element_array.push(None);
+ start + count
+ };
+ mutable.extend(0, (start).to_usize().unwrap(), (end).to_usize().unwrap());
+ offsets.push(offsets[row_index] + end - start);
+ }
+
+ let default_element_array = Arc::new(Int64Array::from(default_element_array));
+ let default_element_array =
+ general_repeat::<O>(&default_element, &default_element_array)?;
Review Comment:
repeat can be replaced with repeating extend() calls
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org