This is an automated email from the ASF dual-hosted git repository.
comphead pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new af20d2d874 feat: support `LargeList` in `cardinality` (#8726)
af20d2d874 is described below
commit af20d2d874c359031dfb4f6828a9a55891ffae06
Author: Alex Huang <[email protected]>
AuthorDate: Sat Jan 6 00:55:26 2024 +0800
feat: support `LargeList` in `cardinality` (#8726)
* support LargeList in cardinality
---
datafusion/physical-expr/src/array_expressions.rs | 21 +++++++++++++++---
datafusion/sqllogictest/test_files/array.slt | 26 +++++++++++++++++++++++
2 files changed, 44 insertions(+), 3 deletions(-)
diff --git a/datafusion/physical-expr/src/array_expressions.rs
b/datafusion/physical-expr/src/array_expressions.rs
index 15330af640..78c490d5db 100644
--- a/datafusion/physical-expr/src/array_expressions.rs
+++ b/datafusion/physical-expr/src/array_expressions.rs
@@ -2114,16 +2114,31 @@ pub fn cardinality(args: &[ArrayRef]) ->
Result<ArrayRef> {
return exec_err!("cardinality expects one argument");
}
- let list_array = as_list_array(&args[0])?.clone();
+ match &args[0].data_type() {
+ DataType::List(_) => {
+ let list_array = as_list_array(&args[0])?;
+ generic_list_cardinality::<i32>(list_array)
+ }
+ DataType::LargeList(_) => {
+ let list_array = as_large_list_array(&args[0])?;
+ generic_list_cardinality::<i64>(list_array)
+ }
+ other => {
+ exec_err!("cardinality does not support type '{:?}'", other)
+ }
+ }
+}
- let result = list_array
+fn generic_list_cardinality<O: OffsetSizeTrait>(
+ array: &GenericListArray<O>,
+) -> Result<ArrayRef> {
+ let result = array
.iter()
.map(|arr| match compute_array_dims(arr)? {
Some(vector) => Ok(Some(vector.iter().map(|x|
x.unwrap()).product::<u64>())),
None => Ok(None),
})
.collect::<Result<UInt64Array>>()?;
-
Ok(Arc::new(result) as ArrayRef)
}
diff --git a/datafusion/sqllogictest/test_files/array.slt
b/datafusion/sqllogictest/test_files/array.slt
index 7cee615a57..774d67b4fd 100644
--- a/datafusion/sqllogictest/test_files/array.slt
+++ b/datafusion/sqllogictest/test_files/array.slt
@@ -3320,18 +3320,33 @@ select cardinality(make_array(1, 2, 3, 4, 5)),
cardinality([1, 3, 5]), cardinali
----
5 3 5
+query III
+select cardinality(arrow_cast(make_array(1, 2, 3, 4, 5), 'LargeList(Int64)')),
cardinality(arrow_cast([1, 3, 5], 'LargeList(Int64)')),
cardinality(arrow_cast(make_array('h', 'e', 'l', 'l', 'o'), 'LargeList(Utf8)'));
+----
+5 3 5
+
# cardinality scalar function #2
query II
select cardinality(make_array([1, 2], [3, 4], [5, 6])),
cardinality(array_repeat(array_repeat(array_repeat(3, 3), 2), 3));
----
6 18
+query I
+select cardinality(arrow_cast(make_array([1, 2], [3, 4], [5, 6]),
'LargeList(List(Int64))'));
+----
+6
+
# cardinality scalar function #3
query II
select cardinality(make_array()), cardinality(make_array(make_array()))
----
NULL 0
+query II
+select cardinality(arrow_cast(make_array(), 'LargeList(Null)')),
cardinality(arrow_cast(make_array(make_array()), 'LargeList(List(Null))'))
+----
+NULL 0
+
# cardinality with columns
query III
select cardinality(column1), cardinality(column2), cardinality(column3) from
arrays;
@@ -3344,6 +3359,17 @@ NULL 3 4
4 NULL 1
4 3 NULL
+query III
+select cardinality(column1), cardinality(column2), cardinality(column3) from
large_arrays;
+----
+4 3 5
+4 3 5
+4 3 5
+4 3 3
+NULL 3 4
+4 NULL 1
+4 3 NULL
+
## array_remove (aliases: `list_remove`)
# array_remove scalar function #1
