Support array_distinct function. #8268
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -31,8 +31,8 @@ use arrow_buffer::NullBuffer; | |
|
|
||
| use arrow_schema::{FieldRef, SortOptions}; | ||
| use datafusion_common::cast::{ | ||
| as_generic_list_array, as_generic_string_array, as_int64_array, as_large_list_array, | ||
| as_list_array, as_null_array, as_string_array, | ||
| }; | ||
| use datafusion_common::utils::{array_into_list_array, list_ndims}; | ||
| use datafusion_common::{ | ||
|
|
@@ -2111,6 +2111,66 @@ pub fn array_intersect(args: &[ArrayRef]) -> Result<ArrayRef> { | |
| } | ||
| } | ||
|
|
||
| pub fn general_array_distinct<OffsetSize: OffsetSizeTrait>( | ||
| array: &GenericListArray<OffsetSize>, | ||
| field: &FieldRef, | ||
| ) -> Result<ArrayRef> { | ||
| let dt = array.value_type(); | ||
| let mut offsets = Vec::with_capacity(array.len()); | ||
| offsets.push(OffsetSize::usize_as(0)); | ||
| let mut new_arrays = Vec::with_capacity(array.len()); | ||
| let converter = RowConverter::new(vec![SortField::new(dt.clone())])?; | ||
| // distinct for each list in ListArray | ||
| for arr in array.iter().flatten() { | ||
| let values = converter.convert_columns(&[arr])?; | ||
|
There was a problem hiding this comment. why not distinct array in columnar way, There was a problem hiding this comment.
It is great to distinct array without row converter, but I don't think we can do that without downcast to exact arr then do the distinction. Is there any recommended way? There was a problem hiding this comment. I also don't have another idea other than downcast arr, I was just wondering if it is worth to downcast to exact arr. There was a problem hiding this comment.
Downcasting to the exact array type can result in faster code in many cases, as the rust compiler can make specialized implemenations for each type. However, there are a lot of DataTypes, including nested ones like Dict, List, Struct, etc so making specialized implementations often requires a lot of work The row converter handles all the types internally. What we have typically done in the past with DataFusion is to use non type specific code like |
||
| // sort elements in list and remove duplicates | ||
| let rows = values.iter().sorted().dedup().collect::<Vec<_>>(); | ||
| let last_offset: OffsetSize = offsets.last().copied().unwrap(); | ||
| offsets.push(last_offset + OffsetSize::usize_as(rows.len())); | ||
| let arrays = converter.convert_rows(rows)?; | ||
| let array = match arrays.get(0) { | ||
| Some(array) => array.clone(), | ||
| None => { | ||
| return internal_err!("array_distinct: failed to get array from rows") | ||
| } | ||
| }; | ||
| new_arrays.push(array); | ||
| } | ||
| let offsets = OffsetBuffer::new(offsets.into()); | ||
| let new_arrays_ref = new_arrays.iter().map(|v| v.as_ref()).collect::<Vec<_>>(); | ||
| let values = compute::concat(&new_arrays_ref)?; | ||
| Ok(Arc::new(GenericListArray::<OffsetSize>::try_new( | ||
| field.clone(), | ||
| offsets, | ||
| values, | ||
| None, | ||
| )?)) | ||
| } | ||
|
|
||
| /// array_distinct SQL function | ||
| /// example: from list [1, 3, 2, 3, 1, 2, 4] to [1, 2, 3, 4] | ||
| pub fn array_distinct(args: &[ArrayRef]) -> Result<ArrayRef> { | ||
| assert_eq!(args.len(), 1); | ||
|
|
||
| // handle null | ||
| if args[0].data_type() == &DataType::Null { | ||
| return Ok(args[0].clone()); | ||
| } | ||
|
|
||
| // handle for list & largelist | ||
| match args[0].data_type() { | ||
| DataType::List(field) => { | ||
| let array = as_list_array(&args[0])?; | ||
| general_array_distinct(array, field) | ||
|
Comment on lines
+2163
to
+2164
There was a problem hiding this comment. nit: put There was a problem hiding this comment. Sorry, I don't get your point. Iargelist differs with list, so I think it maybe better to handle it before generic function? There was a problem hiding this comment. I mean change the function signature: general_array_distinct<OffsetSize: OffsetSizeTrait>(
array: &ArrayRef,
field: &FieldRef,
)Then cast array in the fucntion There was a problem hiding this comment. I think you are referring to something like There was a problem hiding this comment. I think we can merge this PR as is and then add support for |
||
| } | ||
| DataType::LargeList(field) => { | ||
| let array = as_large_list_array(&args[0])?; | ||
| general_array_distinct(array, field) | ||
| } | ||
| _ => internal_err!("array_distinct only support list array"), | ||
| } | ||
| } | ||
|
|
||
| #[cfg(test)] | ||
| mod tests { | ||
| use super::*; | ||
|
|
||
Uh oh!
There was an error while loading. Please reload this page.