You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@arrow.apache.org by ja...@apache.org on 2024/03/24 00:57:19 UTC
(arrow-datafusion) branch main updated: Issue-9750 - Extract array_length, array_reverse and array_sort functions from functions-array subcrate' s kernels and udf containers (#9751)
This is an automated email from the ASF dual-hosted git repository.
jayzhan 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 6eda0e2e35 Issue-9750 - Extract array_length, array_reverse and array_sort functions from functions-array subcrate' s kernels and udf containers (#9751)
6eda0e2e35 is described below
commit 6eda0e2e35c99b64e107f860598085b8e06b2277
Author: Eren Avsarogullari <er...@apache.org>
AuthorDate: Sat Mar 23 17:57:13 2024 -0700
Issue-9750 - Extract array_length, array_reverse and array_sort functions from functions-array subcrate' s kernels and udf containers (#9751)
---
datafusion/functions-array/src/kernels.rs | 232 +-----------------------------
datafusion/functions-array/src/length.rs | 158 ++++++++++++++++++++
datafusion/functions-array/src/lib.rs | 15 +-
datafusion/functions-array/src/range.rs | 19 ++-
datafusion/functions-array/src/reverse.rs | 147 +++++++++++++++++++
datafusion/functions-array/src/sort.rs | 176 +++++++++++++++++++++++
datafusion/functions-array/src/udf.rs | 167 ---------------------
7 files changed, 507 insertions(+), 407 deletions(-)
diff --git a/datafusion/functions-array/src/kernels.rs b/datafusion/functions-array/src/kernels.rs
index 9b743fa913..1fb3abd529 100644
--- a/datafusion/functions-array/src/kernels.rs
+++ b/datafusion/functions-array/src/kernels.rs
@@ -18,19 +18,18 @@
//! implementation kernels for array functions
use arrow::array::{
- Array, ArrayRef, BooleanArray, Capacities, GenericListArray, Int64Array,
- LargeListArray, ListArray, MutableArrayData, OffsetSizeTrait, UInt64Array,
+ Array, ArrayRef, BooleanArray, Capacities, GenericListArray, Int64Array, ListArray,
+ MutableArrayData, OffsetSizeTrait, UInt64Array,
};
use arrow::compute;
use arrow::datatypes::{DataType, Field, UInt64Type};
use arrow_array::new_null_array;
-use arrow_buffer::{ArrowNativeType, BooleanBufferBuilder, NullBuffer, OffsetBuffer};
+use arrow_buffer::{ArrowNativeType, OffsetBuffer};
use arrow_schema::FieldRef;
-use arrow_schema::SortOptions;
use datafusion_common::cast::{
as_generic_list_array, as_int64_array, as_large_list_array, as_list_array,
- as_null_array, as_string_array,
+ as_null_array,
};
use datafusion_common::{
exec_err, internal_datafusion_err, DataFusionError, Result, ScalarValue,
@@ -190,64 +189,6 @@ fn general_array_empty<O: OffsetSizeTrait>(array: &ArrayRef) -> Result<ArrayRef>
Ok(Arc::new(builder))
}
-/// Returns the length of a concrete array dimension
-fn compute_array_length(
- arr: Option<ArrayRef>,
- dimension: Option<i64>,
-) -> Result<Option<u64>> {
- let mut current_dimension: i64 = 1;
- let mut value = match arr {
- Some(arr) => arr,
- None => return Ok(None),
- };
- let dimension = match dimension {
- Some(value) => {
- if value < 1 {
- return Ok(None);
- }
-
- value
- }
- None => return Ok(None),
- };
-
- loop {
- if current_dimension == dimension {
- return Ok(Some(value.len() as u64));
- }
-
- match value.data_type() {
- DataType::List(..) => {
- value = downcast_arg!(value, ListArray).value(0);
- current_dimension += 1;
- }
- DataType::LargeList(..) => {
- value = downcast_arg!(value, LargeListArray).value(0);
- current_dimension += 1;
- }
- _ => return Ok(None),
- }
- }
-}
-
-/// Dispatch array length computation based on the offset type.
-fn general_array_length<O: OffsetSizeTrait>(array: &[ArrayRef]) -> Result<ArrayRef> {
- let list_array = as_generic_list_array::<O>(&array[0])?;
- let dimension = if array.len() == 2 {
- as_int64_array(&array[1])?.clone()
- } else {
- Int64Array::from_value(1, list_array.len())
- };
-
- let result = list_array
- .iter()
- .zip(dimension.iter())
- .map(|(arr, dim)| compute_array_length(arr, dim))
- .collect::<Result<UInt64Array>>()?;
-
- Ok(Arc::new(result) as ArrayRef)
-}
-
/// Array_repeat SQL function
pub fn array_repeat(args: &[ArrayRef]) -> Result<ArrayRef> {
if args.len() != 2 {
@@ -394,19 +335,6 @@ fn general_list_repeat<O: OffsetSizeTrait>(
)?))
}
-/// Array_length SQL function
-pub fn array_length(args: &[ArrayRef]) -> Result<ArrayRef> {
- if args.len() != 1 && args.len() != 2 {
- return exec_err!("array_length expects one or two arguments");
- }
-
- match &args[0].data_type() {
- DataType::List(_) => general_array_length::<i32>(args),
- DataType::LargeList(_) => general_array_length::<i64>(args),
- array_type => exec_err!("array_length does not support type '{array_type:?}'"),
- }
-}
-
/// array_resize SQL function
pub fn array_resize(arg: &[ArrayRef]) -> Result<ArrayRef> {
if arg.len() < 2 || arg.len() > 3 {
@@ -501,89 +429,6 @@ where
)?))
}
-/// Array_sort SQL function
-pub fn array_sort(args: &[ArrayRef]) -> Result<ArrayRef> {
- if args.is_empty() || args.len() > 3 {
- return exec_err!("array_sort expects one to three arguments");
- }
-
- let sort_option = match args.len() {
- 1 => None,
- 2 => {
- let sort = as_string_array(&args[1])?.value(0);
- Some(SortOptions {
- descending: order_desc(sort)?,
- nulls_first: true,
- })
- }
- 3 => {
- let sort = as_string_array(&args[1])?.value(0);
- let nulls_first = as_string_array(&args[2])?.value(0);
- Some(SortOptions {
- descending: order_desc(sort)?,
- nulls_first: order_nulls_first(nulls_first)?,
- })
- }
- _ => return exec_err!("array_sort expects 1 to 3 arguments"),
- };
-
- let list_array = as_list_array(&args[0])?;
- let row_count = list_array.len();
-
- let mut array_lengths = vec![];
- let mut arrays = vec![];
- let mut valid = BooleanBufferBuilder::new(row_count);
- for i in 0..row_count {
- if list_array.is_null(i) {
- array_lengths.push(0);
- valid.append(false);
- } else {
- let arr_ref = list_array.value(i);
- let arr_ref = arr_ref.as_ref();
-
- let sorted_array = compute::sort(arr_ref, sort_option)?;
- array_lengths.push(sorted_array.len());
- arrays.push(sorted_array);
- valid.append(true);
- }
- }
-
- // Assume all arrays have the same data type
- let data_type = list_array.value_type();
- let buffer = valid.finish();
-
- let elements = arrays
- .iter()
- .map(|a| a.as_ref())
- .collect::<Vec<&dyn Array>>();
-
- let list_arr = ListArray::new(
- Arc::new(Field::new("item", data_type, true)),
- OffsetBuffer::from_lengths(array_lengths),
- Arc::new(compute::concat(elements.as_slice())?),
- Some(NullBuffer::new(buffer)),
- );
- Ok(Arc::new(list_arr))
-}
-
-fn order_desc(modifier: &str) -> Result<bool> {
- match modifier.to_uppercase().as_str() {
- "DESC" => Ok(true),
- "ASC" => Ok(false),
- _ => exec_err!("the second parameter of array_sort expects DESC or ASC"),
- }
-}
-
-fn order_nulls_first(modifier: &str) -> Result<bool> {
- match modifier.to_uppercase().as_str() {
- "NULLS FIRST" => Ok(true),
- "NULLS LAST" => Ok(false),
- _ => exec_err!(
- "the third parameter of array_sort expects NULLS FIRST or NULLS LAST"
- ),
- }
-}
-
// Create new offsets that are euqiavlent to `flatten` the array.
fn get_offsets_for_flatten<O: OffsetSizeTrait>(
offsets: OffsetBuffer<O>,
@@ -652,72 +497,3 @@ pub fn flatten(args: &[ArrayRef]) -> Result<ArrayRef> {
}
}
}
-
-/// array_reverse SQL function
-pub fn array_reverse(arg: &[ArrayRef]) -> Result<ArrayRef> {
- if arg.len() != 1 {
- return exec_err!("array_reverse needs one argument");
- }
-
- match &arg[0].data_type() {
- DataType::List(field) => {
- let array = as_list_array(&arg[0])?;
- general_array_reverse::<i32>(array, field)
- }
- DataType::LargeList(field) => {
- let array = as_large_list_array(&arg[0])?;
- general_array_reverse::<i64>(array, field)
- }
- DataType::Null => Ok(arg[0].clone()),
- array_type => exec_err!("array_reverse does not support type '{array_type:?}'."),
- }
-}
-
-fn general_array_reverse<O: OffsetSizeTrait>(
- array: &GenericListArray<O>,
- field: &FieldRef,
-) -> Result<ArrayRef>
-where
- O: TryFrom<i64>,
-{
- 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 nulls = vec![];
- let mut mutable =
- MutableArrayData::with_capacities(vec![&original_data], false, capacity);
-
- for (row_index, offset_window) in array.offsets().windows(2).enumerate() {
- // skip the null value
- if array.is_null(row_index) {
- nulls.push(false);
- offsets.push(offsets[row_index] + O::one());
- mutable.extend(0, 0, 1);
- continue;
- } else {
- nulls.push(true);
- }
-
- let start = offset_window[0];
- let end = offset_window[1];
-
- let mut index = end - O::one();
- let mut cnt = 0;
-
- while index >= start {
- mutable.extend(0, index.to_usize().unwrap(), index.to_usize().unwrap() + 1);
- index = index - O::one();
- cnt += 1;
- }
- offsets.push(offsets[row_index] + O::usize_as(cnt));
- }
-
- let data = mutable.freeze();
- Ok(Arc::new(GenericListArray::<O>::try_new(
- field.clone(),
- OffsetBuffer::<O>::new(offsets.into()),
- arrow_array::make_array(data),
- Some(nulls.into()),
- )?))
-}
diff --git a/datafusion/functions-array/src/length.rs b/datafusion/functions-array/src/length.rs
new file mode 100644
index 0000000000..e8e3611317
--- /dev/null
+++ b/datafusion/functions-array/src/length.rs
@@ -0,0 +1,158 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! [`ScalarUDFImpl`] definitions for array_length function.
+
+use crate::utils::{downcast_arg, make_scalar_function};
+use arrow_array::{
+ Array, ArrayRef, Int64Array, LargeListArray, ListArray, OffsetSizeTrait, UInt64Array,
+};
+use arrow_schema::DataType;
+use arrow_schema::DataType::{FixedSizeList, LargeList, List, UInt64};
+use core::any::type_name;
+use datafusion_common::cast::{as_generic_list_array, as_int64_array};
+use datafusion_common::DataFusionError;
+use datafusion_common::{exec_err, plan_err};
+use datafusion_expr::expr::ScalarFunction;
+use datafusion_expr::{ColumnarValue, Expr, ScalarUDFImpl, Signature, Volatility};
+use std::any::Any;
+use std::sync::Arc;
+
+make_udf_function!(
+ ArrayLength,
+ array_length,
+ array,
+ "returns the length of the array dimension.",
+ array_length_udf
+);
+
+#[derive(Debug)]
+pub(super) struct ArrayLength {
+ signature: Signature,
+ aliases: Vec<String>,
+}
+impl ArrayLength {
+ pub fn new() -> Self {
+ Self {
+ signature: Signature::variadic_any(Volatility::Immutable),
+ aliases: vec![String::from("array_length"), String::from("list_length")],
+ }
+ }
+}
+
+impl ScalarUDFImpl for ArrayLength {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+ fn name(&self) -> &str {
+ "array_length"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, arg_types: &[DataType]) -> datafusion_common::Result<DataType> {
+ Ok(match arg_types[0] {
+ List(_) | LargeList(_) | FixedSizeList(_, _) => UInt64,
+ _ => {
+ return plan_err!("The array_length function can only accept List/LargeList/FixedSizeList.");
+ }
+ })
+ }
+
+ fn invoke(&self, args: &[ColumnarValue]) -> datafusion_common::Result<ColumnarValue> {
+ make_scalar_function(array_length_inner)(args)
+ }
+
+ fn aliases(&self) -> &[String] {
+ &self.aliases
+ }
+}
+
+/// Array_length SQL function
+pub fn array_length_inner(args: &[ArrayRef]) -> datafusion_common::Result<ArrayRef> {
+ if args.len() != 1 && args.len() != 2 {
+ return exec_err!("array_length expects one or two arguments");
+ }
+
+ match &args[0].data_type() {
+ List(_) => general_array_length::<i32>(args),
+ LargeList(_) => general_array_length::<i64>(args),
+ array_type => exec_err!("array_length does not support type '{array_type:?}'"),
+ }
+}
+
+/// Dispatch array length computation based on the offset type.
+fn general_array_length<O: OffsetSizeTrait>(
+ array: &[ArrayRef],
+) -> datafusion_common::Result<ArrayRef> {
+ let list_array = as_generic_list_array::<O>(&array[0])?;
+ let dimension = if array.len() == 2 {
+ as_int64_array(&array[1])?.clone()
+ } else {
+ Int64Array::from_value(1, list_array.len())
+ };
+
+ let result = list_array
+ .iter()
+ .zip(dimension.iter())
+ .map(|(arr, dim)| compute_array_length(arr, dim))
+ .collect::<datafusion_common::Result<UInt64Array>>()?;
+
+ Ok(Arc::new(result) as ArrayRef)
+}
+
+/// Returns the length of a concrete array dimension
+fn compute_array_length(
+ arr: Option<ArrayRef>,
+ dimension: Option<i64>,
+) -> datafusion_common::Result<Option<u64>> {
+ let mut current_dimension: i64 = 1;
+ let mut value = match arr {
+ Some(arr) => arr,
+ None => return Ok(None),
+ };
+ let dimension = match dimension {
+ Some(value) => {
+ if value < 1 {
+ return Ok(None);
+ }
+
+ value
+ }
+ None => return Ok(None),
+ };
+
+ loop {
+ if current_dimension == dimension {
+ return Ok(Some(value.len() as u64));
+ }
+
+ match value.data_type() {
+ List(..) => {
+ value = downcast_arg!(value, ListArray).value(0);
+ current_dimension += 1;
+ }
+ LargeList(..) => {
+ value = downcast_arg!(value, LargeListArray).value(0);
+ current_dimension += 1;
+ }
+ _ => return Ok(None),
+ }
+ }
+}
diff --git a/datafusion/functions-array/src/lib.rs b/datafusion/functions-array/src/lib.rs
index 6ed77e5d17..f4ca5408aa 100644
--- a/datafusion/functions-array/src/lib.rs
+++ b/datafusion/functions-array/src/lib.rs
@@ -34,12 +34,15 @@ mod core;
mod except;
mod extract;
mod kernels;
+mod length;
mod position;
mod range;
mod remove;
mod replace;
+mod reverse;
mod rewrite;
mod set_ops;
+mod sort;
mod string;
mod udf;
mod utils;
@@ -64,6 +67,7 @@ pub mod expr_fn {
pub use super::extract::array_pop_back;
pub use super::extract::array_pop_front;
pub use super::extract::array_slice;
+ pub use super::length::array_length;
pub use super::position::array_position;
pub use super::position::array_positions;
pub use super::range::gen_series;
@@ -74,19 +78,18 @@ pub mod expr_fn {
pub use super::replace::array_replace;
pub use super::replace::array_replace_all;
pub use super::replace::array_replace_n;
+ pub use super::reverse::array_reverse;
pub use super::set_ops::array_distinct;
pub use super::set_ops::array_intersect;
pub use super::set_ops::array_union;
+ pub use super::sort::array_sort;
pub use super::string::array_to_string;
pub use super::string::string_to_array;
pub use super::udf::array_dims;
pub use super::udf::array_empty;
- pub use super::udf::array_length;
pub use super::udf::array_ndims;
pub use super::udf::array_repeat;
pub use super::udf::array_resize;
- pub use super::udf::array_reverse;
- pub use super::udf::array_sort;
pub use super::udf::cardinality;
pub use super::udf::flatten;
}
@@ -114,12 +117,12 @@ pub fn register_all(registry: &mut dyn FunctionRegistry) -> Result<()> {
array_has::array_has_all_udf(),
array_has::array_has_any_udf(),
udf::array_empty_udf(),
- udf::array_length_udf(),
+ length::array_length_udf(),
udf::flatten_udf(),
- udf::array_sort_udf(),
+ sort::array_sort_udf(),
udf::array_repeat_udf(),
udf::array_resize_udf(),
- udf::array_reverse_udf(),
+ reverse::array_reverse_udf(),
set_ops::array_distinct_udf(),
set_ops::array_intersect_udf(),
set_ops::array_union_udf(),
diff --git a/datafusion/functions-array/src/range.rs b/datafusion/functions-array/src/range.rs
index 7dfce71332..176a5617d5 100644
--- a/datafusion/functions-array/src/range.rs
+++ b/datafusion/functions-array/src/range.rs
@@ -22,6 +22,7 @@ use arrow::datatypes::{DataType, Field};
use arrow_buffer::{BooleanBufferBuilder, NullBuffer, OffsetBuffer};
use std::any::Any;
+use crate::utils::make_scalar_function;
use arrow_array::types::{Date32Type, IntervalMonthDayNanoType};
use arrow_array::Date32Array;
use arrow_schema::IntervalUnit::MonthDayNano;
@@ -85,10 +86,13 @@ impl ScalarUDFImpl for Range {
}
fn invoke(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
- let args = ColumnarValue::values_to_arrays(args)?;
match args[0].data_type() {
- DataType::Int64 => gen_range_inner(&args, false).map(ColumnarValue::Array),
- DataType::Date32 => gen_range_date(&args, false).map(ColumnarValue::Array),
+ DataType::Int64 => {
+ make_scalar_function(|args| gen_range_inner(args, false))(args)
+ }
+ DataType::Date32 => {
+ make_scalar_function(|args| gen_range_date(args, false))(args)
+ }
_ => {
exec_err!("unsupported type for range")
}
@@ -151,10 +155,13 @@ impl ScalarUDFImpl for GenSeries {
}
fn invoke(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
- let args = ColumnarValue::values_to_arrays(args)?;
match args[0].data_type() {
- DataType::Int64 => gen_range_inner(&args, true).map(ColumnarValue::Array),
- DataType::Date32 => gen_range_date(&args, true).map(ColumnarValue::Array),
+ DataType::Int64 => {
+ make_scalar_function(|args| gen_range_inner(args, true))(args)
+ }
+ DataType::Date32 => {
+ make_scalar_function(|args| gen_range_date(args, true))(args)
+ }
_ => {
exec_err!("unsupported type for range")
}
diff --git a/datafusion/functions-array/src/reverse.rs b/datafusion/functions-array/src/reverse.rs
new file mode 100644
index 0000000000..7eb9e53dee
--- /dev/null
+++ b/datafusion/functions-array/src/reverse.rs
@@ -0,0 +1,147 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! [`ScalarUDFImpl`] definitions for array_reverse function.
+
+use crate::utils::make_scalar_function;
+use arrow::array::{Capacities, MutableArrayData};
+use arrow_array::{Array, ArrayRef, GenericListArray, OffsetSizeTrait};
+use arrow_buffer::OffsetBuffer;
+use arrow_schema::{DataType, FieldRef};
+use datafusion_common::cast::{as_large_list_array, as_list_array};
+use datafusion_common::exec_err;
+use datafusion_expr::expr::ScalarFunction;
+use datafusion_expr::{ColumnarValue, Expr, ScalarUDFImpl, Signature, Volatility};
+use std::any::Any;
+use std::sync::Arc;
+
+make_udf_function!(
+ ArrayReverse,
+ array_reverse,
+ array,
+ "reverses the order of elements in the array.",
+ array_reverse_udf
+);
+
+#[derive(Debug)]
+pub(super) struct ArrayReverse {
+ signature: Signature,
+ aliases: Vec<String>,
+}
+
+impl ArrayReverse {
+ pub fn new() -> Self {
+ Self {
+ signature: Signature::any(1, Volatility::Immutable),
+ aliases: vec!["array_reverse".to_string(), "list_reverse".to_string()],
+ }
+ }
+}
+
+impl ScalarUDFImpl for ArrayReverse {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+ fn name(&self) -> &str {
+ "array_reverse"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, arg_types: &[DataType]) -> datafusion_common::Result<DataType> {
+ Ok(arg_types[0].clone())
+ }
+
+ fn invoke(&self, args: &[ColumnarValue]) -> datafusion_common::Result<ColumnarValue> {
+ make_scalar_function(array_reverse_inner)(args)
+ }
+
+ fn aliases(&self) -> &[String] {
+ &self.aliases
+ }
+}
+
+/// array_reverse SQL function
+pub fn array_reverse_inner(arg: &[ArrayRef]) -> datafusion_common::Result<ArrayRef> {
+ if arg.len() != 1 {
+ return exec_err!("array_reverse needs one argument");
+ }
+
+ match &arg[0].data_type() {
+ DataType::List(field) => {
+ let array = as_list_array(&arg[0])?;
+ general_array_reverse::<i32>(array, field)
+ }
+ DataType::LargeList(field) => {
+ let array = as_large_list_array(&arg[0])?;
+ general_array_reverse::<i64>(array, field)
+ }
+ DataType::Null => Ok(arg[0].clone()),
+ array_type => exec_err!("array_reverse does not support type '{array_type:?}'."),
+ }
+}
+
+fn general_array_reverse<O: OffsetSizeTrait>(
+ array: &GenericListArray<O>,
+ field: &FieldRef,
+) -> datafusion_common::Result<ArrayRef>
+where
+ O: TryFrom<i64>,
+{
+ 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 nulls = vec![];
+ let mut mutable =
+ MutableArrayData::with_capacities(vec![&original_data], false, capacity);
+
+ for (row_index, offset_window) in array.offsets().windows(2).enumerate() {
+ // skip the null value
+ if array.is_null(row_index) {
+ nulls.push(false);
+ offsets.push(offsets[row_index] + O::one());
+ mutable.extend(0, 0, 1);
+ continue;
+ } else {
+ nulls.push(true);
+ }
+
+ let start = offset_window[0];
+ let end = offset_window[1];
+
+ let mut index = end - O::one();
+ let mut cnt = 0;
+
+ while index >= start {
+ mutable.extend(0, index.to_usize().unwrap(), index.to_usize().unwrap() + 1);
+ index = index - O::one();
+ cnt += 1;
+ }
+ offsets.push(offsets[row_index] + O::usize_as(cnt));
+ }
+
+ let data = mutable.freeze();
+ Ok(Arc::new(GenericListArray::<O>::try_new(
+ field.clone(),
+ OffsetBuffer::<O>::new(offsets.into()),
+ arrow_array::make_array(data),
+ Some(nulls.into()),
+ )?))
+}
diff --git a/datafusion/functions-array/src/sort.rs b/datafusion/functions-array/src/sort.rs
new file mode 100644
index 0000000000..2f3fa33e68
--- /dev/null
+++ b/datafusion/functions-array/src/sort.rs
@@ -0,0 +1,176 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! [`ScalarUDFImpl`] definitions for array_sort function.
+
+use crate::utils::make_scalar_function;
+use arrow::compute;
+use arrow_array::{Array, ArrayRef, ListArray};
+use arrow_buffer::{BooleanBufferBuilder, NullBuffer, OffsetBuffer};
+use arrow_schema::DataType::{FixedSizeList, LargeList, List};
+use arrow_schema::{DataType, Field, SortOptions};
+use datafusion_common::cast::{as_list_array, as_string_array};
+use datafusion_common::exec_err;
+use datafusion_expr::expr::ScalarFunction;
+use datafusion_expr::{ColumnarValue, Expr, ScalarUDFImpl, Signature, Volatility};
+use std::any::Any;
+use std::sync::Arc;
+
+make_udf_function!(
+ ArraySort,
+ array_sort,
+ array desc null_first,
+ "returns sorted array.",
+ array_sort_udf
+);
+
+#[derive(Debug)]
+pub(super) struct ArraySort {
+ signature: Signature,
+ aliases: Vec<String>,
+}
+
+impl ArraySort {
+ pub fn new() -> Self {
+ Self {
+ signature: Signature::variadic_any(Volatility::Immutable),
+ aliases: vec!["array_sort".to_string(), "list_sort".to_string()],
+ }
+ }
+}
+
+impl ScalarUDFImpl for ArraySort {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+ fn name(&self) -> &str {
+ "array_sort"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, arg_types: &[DataType]) -> datafusion_common::Result<DataType> {
+ match &arg_types[0] {
+ List(field) | FixedSizeList(field, _) => Ok(List(Arc::new(Field::new(
+ "item",
+ field.data_type().clone(),
+ true,
+ )))),
+ LargeList(field) => Ok(LargeList(Arc::new(Field::new(
+ "item",
+ field.data_type().clone(),
+ true,
+ )))),
+ _ => exec_err!(
+ "Not reachable, data_type should be List, LargeList or FixedSizeList"
+ ),
+ }
+ }
+
+ fn invoke(&self, args: &[ColumnarValue]) -> datafusion_common::Result<ColumnarValue> {
+ make_scalar_function(array_sort_inner)(args)
+ }
+
+ fn aliases(&self) -> &[String] {
+ &self.aliases
+ }
+}
+
+/// Array_sort SQL function
+pub fn array_sort_inner(args: &[ArrayRef]) -> datafusion_common::Result<ArrayRef> {
+ if args.is_empty() || args.len() > 3 {
+ return exec_err!("array_sort expects one to three arguments");
+ }
+
+ let sort_option = match args.len() {
+ 1 => None,
+ 2 => {
+ let sort = as_string_array(&args[1])?.value(0);
+ Some(SortOptions {
+ descending: order_desc(sort)?,
+ nulls_first: true,
+ })
+ }
+ 3 => {
+ let sort = as_string_array(&args[1])?.value(0);
+ let nulls_first = as_string_array(&args[2])?.value(0);
+ Some(SortOptions {
+ descending: order_desc(sort)?,
+ nulls_first: order_nulls_first(nulls_first)?,
+ })
+ }
+ _ => return exec_err!("array_sort expects 1 to 3 arguments"),
+ };
+
+ let list_array = as_list_array(&args[0])?;
+ let row_count = list_array.len();
+
+ let mut array_lengths = vec![];
+ let mut arrays = vec![];
+ let mut valid = BooleanBufferBuilder::new(row_count);
+ for i in 0..row_count {
+ if list_array.is_null(i) {
+ array_lengths.push(0);
+ valid.append(false);
+ } else {
+ let arr_ref = list_array.value(i);
+ let arr_ref = arr_ref.as_ref();
+
+ let sorted_array = compute::sort(arr_ref, sort_option)?;
+ array_lengths.push(sorted_array.len());
+ arrays.push(sorted_array);
+ valid.append(true);
+ }
+ }
+
+ // Assume all arrays have the same data type
+ let data_type = list_array.value_type();
+ let buffer = valid.finish();
+
+ let elements = arrays
+ .iter()
+ .map(|a| a.as_ref())
+ .collect::<Vec<&dyn Array>>();
+
+ let list_arr = ListArray::new(
+ Arc::new(Field::new("item", data_type, true)),
+ OffsetBuffer::from_lengths(array_lengths),
+ Arc::new(compute::concat(elements.as_slice())?),
+ Some(NullBuffer::new(buffer)),
+ );
+ Ok(Arc::new(list_arr))
+}
+
+fn order_desc(modifier: &str) -> datafusion_common::Result<bool> {
+ match modifier.to_uppercase().as_str() {
+ "DESC" => Ok(true),
+ "ASC" => Ok(false),
+ _ => exec_err!("the second parameter of array_sort expects DESC or ASC"),
+ }
+}
+
+fn order_nulls_first(modifier: &str) -> datafusion_common::Result<bool> {
+ match modifier.to_uppercase().as_str() {
+ "NULLS FIRST" => Ok(true),
+ "NULLS LAST" => Ok(false),
+ _ => exec_err!(
+ "the third parameter of array_sort expects NULLS FIRST or NULLS LAST"
+ ),
+ }
+}
diff --git a/datafusion/functions-array/src/udf.rs b/datafusion/functions-array/src/udf.rs
index 1567031057..943950447b 100644
--- a/datafusion/functions-array/src/udf.rs
+++ b/datafusion/functions-array/src/udf.rs
@@ -86,70 +86,6 @@ impl ScalarUDFImpl for ArrayDims {
}
}
-make_udf_function!(
- ArraySort,
- array_sort,
- array desc null_first,
- "returns sorted array.",
- array_sort_udf
-);
-
-#[derive(Debug)]
-pub(super) struct ArraySort {
- signature: Signature,
- aliases: Vec<String>,
-}
-
-impl ArraySort {
- pub fn new() -> Self {
- Self {
- signature: Signature::variadic_any(Volatility::Immutable),
- aliases: vec!["array_sort".to_string(), "list_sort".to_string()],
- }
- }
-}
-
-impl ScalarUDFImpl for ArraySort {
- fn as_any(&self) -> &dyn Any {
- self
- }
- fn name(&self) -> &str {
- "array_sort"
- }
-
- fn signature(&self) -> &Signature {
- &self.signature
- }
-
- fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
- use DataType::*;
- match &arg_types[0] {
- List(field) | FixedSizeList(field, _) => Ok(List(Arc::new(Field::new(
- "item",
- field.data_type().clone(),
- true,
- )))),
- LargeList(field) => Ok(LargeList(Arc::new(Field::new(
- "item",
- field.data_type().clone(),
- true,
- )))),
- _ => exec_err!(
- "Not reachable, data_type should be List, LargeList or FixedSizeList"
- ),
- }
- }
-
- fn invoke(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
- let args = ColumnarValue::values_to_arrays(args)?;
- crate::kernels::array_sort(&args).map(ColumnarValue::Array)
- }
-
- fn aliases(&self) -> &[String] {
- &self.aliases
- }
-}
-
make_udf_function!(
ArrayResize,
array_resize,
@@ -420,60 +356,6 @@ impl ScalarUDFImpl for ArrayRepeat {
}
}
-make_udf_function!(
- ArrayLength,
- array_length,
- array,
- "returns the length of the array dimension.",
- array_length_udf
-);
-
-#[derive(Debug)]
-pub(super) struct ArrayLength {
- signature: Signature,
- aliases: Vec<String>,
-}
-impl ArrayLength {
- pub fn new() -> Self {
- Self {
- signature: Signature::variadic_any(Volatility::Immutable),
- aliases: vec![String::from("array_length"), String::from("list_length")],
- }
- }
-}
-
-impl ScalarUDFImpl for ArrayLength {
- fn as_any(&self) -> &dyn Any {
- self
- }
- fn name(&self) -> &str {
- "array_length"
- }
-
- fn signature(&self) -> &Signature {
- &self.signature
- }
-
- fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
- use DataType::*;
- Ok(match arg_types[0] {
- List(_) | LargeList(_) | FixedSizeList(_, _) => UInt64,
- _ => {
- return plan_err!("The array_length function can only accept List/LargeList/FixedSizeList.");
- }
- })
- }
-
- fn invoke(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
- let args = ColumnarValue::values_to_arrays(args)?;
- crate::kernels::array_length(&args).map(ColumnarValue::Array)
- }
-
- fn aliases(&self) -> &[String] {
- &self.aliases
- }
-}
-
make_udf_function!(
Flatten,
flatten,
@@ -541,52 +423,3 @@ impl ScalarUDFImpl for Flatten {
&self.aliases
}
}
-
-make_udf_function!(
- ArrayReverse,
- array_reverse,
- array,
- "reverses the order of elements in the array.",
- array_reverse_udf
-);
-
-#[derive(Debug)]
-pub(super) struct ArrayReverse {
- signature: Signature,
- aliases: Vec<String>,
-}
-
-impl crate::udf::ArrayReverse {
- pub fn new() -> Self {
- Self {
- signature: Signature::any(1, Volatility::Immutable),
- aliases: vec!["array_reverse".to_string(), "list_reverse".to_string()],
- }
- }
-}
-
-impl ScalarUDFImpl for crate::udf::ArrayReverse {
- fn as_any(&self) -> &dyn Any {
- self
- }
- fn name(&self) -> &str {
- "array_reserse"
- }
-
- fn signature(&self) -> &Signature {
- &self.signature
- }
-
- fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
- Ok(arg_types[0].clone())
- }
-
- fn invoke(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
- let args = ColumnarValue::values_to_arrays(args)?;
- crate::kernels::array_reverse(&args).map(ColumnarValue::Array)
- }
-
- fn aliases(&self) -> &[String] {
- &self.aliases
- }
-}