From ca260d99f17ef667b7f06d2da4a67255d27c94a9 Mon Sep 17 00:00:00 2001 From: Alex Huang Date: Thu, 4 Jan 2024 06:14:23 +0800 Subject: [PATCH] support LargeList in array_repeat (#8725) --- .../physical-expr/src/array_expressions.rs | 16 +++++--- datafusion/sqllogictest/test_files/array.slt | 37 +++++++++++++++++++ 2 files changed, 47 insertions(+), 6 deletions(-) diff --git a/datafusion/physical-expr/src/array_expressions.rs b/datafusion/physical-expr/src/array_expressions.rs index aad021610fcb..15330af640ae 100644 --- a/datafusion/physical-expr/src/array_expressions.rs +++ b/datafusion/physical-expr/src/array_expressions.rs @@ -1233,7 +1233,11 @@ pub fn array_repeat(args: &[ArrayRef]) -> Result { match element.data_type() { DataType::List(_) => { let list_array = as_list_array(element)?; - general_list_repeat(list_array, count_array) + general_list_repeat::(list_array, count_array) + } + DataType::LargeList(_) => { + let list_array = as_large_list_array(element)?; + general_list_repeat::(list_array, count_array) } _ => general_repeat(element, count_array), } @@ -1302,8 +1306,8 @@ fn general_repeat(array: &ArrayRef, count_array: &Int64Array) -> Result [[[1, 2, 3], [1, 2, 3]], [], [[6]]] /// ) /// ``` -fn general_list_repeat( - list_array: &ListArray, +fn general_list_repeat( + list_array: &GenericListArray, count_array: &Int64Array, ) -> Result { let data_type = list_array.data_type(); @@ -1335,9 +1339,9 @@ fn general_list_repeat( let data = mutable.freeze(); let repeated_array = arrow_array::make_array(data); - let list_arr = ListArray::try_new( + let list_arr = GenericListArray::::try_new( Arc::new(Field::new("item", value_type.clone(), true)), - OffsetBuffer::from_lengths(vec![original_data.len(); count]), + OffsetBuffer::::from_lengths(vec![original_data.len(); count]), repeated_array, None, )?; @@ -1354,7 +1358,7 @@ fn general_list_repeat( Ok(Arc::new(ListArray::try_new( Arc::new(Field::new("item", data_type.to_owned(), true)), - OffsetBuffer::from_lengths(lengths), + OffsetBuffer::::from_lengths(lengths), values, None, )?)) diff --git a/datafusion/sqllogictest/test_files/array.slt b/datafusion/sqllogictest/test_files/array.slt index a3b2c8cdf1e9..7cee615a5729 100644 --- a/datafusion/sqllogictest/test_files/array.slt +++ b/datafusion/sqllogictest/test_files/array.slt @@ -1957,6 +1957,15 @@ select ---- [[1], [1], [1], [1], [1]] [[1.1, 2.2, 3.3], [1.1, 2.2, 3.3], [1.1, 2.2, 3.3]] [[, ], [, ], [, ]] [[[1, 2], [3, 4]], [[1, 2], [3, 4]]] +query ???? +select + array_repeat(arrow_cast([1], 'LargeList(Int64)'), 5), + array_repeat(arrow_cast([1.1, 2.2, 3.3], 'LargeList(Float64)'), 3), + array_repeat(arrow_cast([null, null], 'LargeList(Null)'), 3), + array_repeat(arrow_cast([[1, 2], [3, 4]], 'LargeList(List(Int64))'), 2); +---- +[[1], [1], [1], [1], [1]] [[1.1, 2.2, 3.3], [1.1, 2.2, 3.3], [1.1, 2.2, 3.3]] [[, ], [, ], [, ]] [[[1, 2], [3, 4]], [[1, 2], [3, 4]]] + # array_repeat with columns #1 statement ok @@ -1967,6 +1976,16 @@ AS VALUES (3, 2, 2.2, 'rust', make_array(7)), (0, 3, 3.3, 'datafusion', make_array(8, 9)); +statement ok +CREATE TABLE large_array_repeat_table +AS SELECT + column1, + column2, + column3, + column4, + arrow_cast(column5, 'LargeList(Int64)') as column5 +FROM array_repeat_table; + query ?????? select array_repeat(column2, column1), @@ -1982,9 +2001,27 @@ from array_repeat_table; [2, 2, 2] [2.2, 2.2, 2.2] [rust, rust, rust] [[7], [7], [7]] [2, 2, 2] [[1], [1], [1]] [] [] [] [] [3, 3, 3] [] +query ?????? +select + array_repeat(column2, column1), + array_repeat(column3, column1), + array_repeat(column4, column1), + array_repeat(column5, column1), + array_repeat(column2, 3), + array_repeat(make_array(1), column1) +from large_array_repeat_table; +---- +[1] [1.1] [a] [[4, 5, 6]] [1, 1, 1] [[1]] +[, ] [, ] [, ] [, ] [, , ] [[1], [1]] +[2, 2, 2] [2.2, 2.2, 2.2] [rust, rust, rust] [[7], [7], [7]] [2, 2, 2] [[1], [1], [1]] +[] [] [] [] [3, 3, 3] [] + statement ok drop table array_repeat_table; +statement ok +drop table large_array_repeat_table; + ## array_concat (aliases: `array_cat`, `list_concat`, `list_cat`) # array_concat error