From 9aaf1cd19ce6d74dd9046d6987b2f54f11268269 Mon Sep 17 00:00:00 2001 From: Eduard Karacharov Date: Sun, 27 Aug 2023 15:12:22 +0300 Subject: [PATCH] list_sort function --- datafusion/expr/src/built_in_function.rs | 6 ++++++ .../physical-expr/src/array_expressions.rs | 16 ++++++++++++++++ datafusion/physical-expr/src/functions.rs | 1 + datafusion/proto/src/generated/prost.rs | 2 ++ datafusion/proto/src/logical_plan/from_proto.rs | 1 + datafusion/proto/src/logical_plan/to_proto.rs | 1 + 6 files changed, 27 insertions(+) diff --git a/datafusion/expr/src/built_in_function.rs b/datafusion/expr/src/built_in_function.rs index de046cde8991..d7bffbff1935 100644 --- a/datafusion/expr/src/built_in_function.rs +++ b/datafusion/expr/src/built_in_function.rs @@ -168,6 +168,8 @@ pub enum BuiltinScalarFunction { ArrayReplaceAll, /// array_slice ArraySlice, + /// array_sort + ArraySort, /// array_to_string ArrayToString, /// cardinality @@ -385,6 +387,7 @@ impl BuiltinScalarFunction { BuiltinScalarFunction::ArrayReplaceAll => Volatility::Immutable, BuiltinScalarFunction::Flatten => Volatility::Immutable, BuiltinScalarFunction::ArraySlice => Volatility::Immutable, + BuiltinScalarFunction::ArraySort => Volatility::Immutable, BuiltinScalarFunction::ArrayToString => Volatility::Immutable, BuiltinScalarFunction::Cardinality => Volatility::Immutable, BuiltinScalarFunction::MakeArray => Volatility::Immutable, @@ -573,6 +576,7 @@ impl BuiltinScalarFunction { BuiltinScalarFunction::ArrayReplaceN => Ok(input_expr_types[0].clone()), BuiltinScalarFunction::ArrayReplaceAll => Ok(input_expr_types[0].clone()), BuiltinScalarFunction::ArraySlice => Ok(input_expr_types[0].clone()), + BuiltinScalarFunction::ArraySort => Ok(input_expr_types[0].clone()), BuiltinScalarFunction::ArrayToString => Ok(Utf8), BuiltinScalarFunction::Cardinality => Ok(UInt64), BuiltinScalarFunction::MakeArray => match input_expr_types.len() { @@ -857,6 +861,7 @@ impl BuiltinScalarFunction { Signature::any(3, self.volatility()) } BuiltinScalarFunction::ArraySlice => Signature::any(3, self.volatility()), + BuiltinScalarFunction::ArraySort => Signature::any(1, self.volatility()), BuiltinScalarFunction::ArrayToString => { Signature::variadic_any(self.volatility()) } @@ -1362,6 +1367,7 @@ fn aliases(func: &BuiltinScalarFunction) -> &'static [&'static str] { &["array_replace_all", "list_replace_all"] } BuiltinScalarFunction::ArraySlice => &["array_slice", "list_slice"], + BuiltinScalarFunction::ArraySort => &["array_sort", "list_sort"], BuiltinScalarFunction::ArrayToString => &[ "array_to_string", "list_to_string", diff --git a/datafusion/physical-expr/src/array_expressions.rs b/datafusion/physical-expr/src/array_expressions.rs index 06432b615c5d..41126657acff 100644 --- a/datafusion/physical-expr/src/array_expressions.rs +++ b/datafusion/physical-expr/src/array_expressions.rs @@ -22,6 +22,7 @@ use arrow::buffer::{Buffer, OffsetBuffer}; use arrow::compute; use arrow::datatypes::{DataType, Field, UInt64Type}; use arrow_buffer::NullBuffer; +use arrow_schema::SortOptions; use core::any::type_name; use datafusion_common::cast::{as_generic_string_array, as_int64_array, as_list_array}; use datafusion_common::{exec_err, internal_err, not_impl_err, plan_err, ScalarValue}; @@ -1864,6 +1865,21 @@ pub fn array_has_all(args: &[ArrayRef]) -> Result { Ok(Arc::new(boolean_builder.finish())) } +/// array_sort SQL function +pub fn array_sort(args: &[ArrayRef]) -> Result { + let list_array = as_list_array(&args[0])?; + let result = list_array.iter() + .map(|list| { + match list { + Some(values) => Ok(Some(compute::sort(&values, None)?)), + None => Ok(None), + } + }) + .collect::>>()?; + let v = ListArray::from_iter_primitive(result); + Ok(args[0].clone()) +} + #[cfg(test)] mod tests { use super::*; diff --git a/datafusion/physical-expr/src/functions.rs b/datafusion/physical-expr/src/functions.rs index e481687ae51b..54dc15e204e3 100644 --- a/datafusion/physical-expr/src/functions.rs +++ b/datafusion/physical-expr/src/functions.rs @@ -488,6 +488,7 @@ pub fn create_physical_fun( BuiltinScalarFunction::ArraySlice => { Arc::new(|args| make_scalar_function(array_expressions::array_slice)(args)) } + BuiltinScalarFunction::ArraySort => Arc::new(|args| make_scalar_function(array_expressions::array_sort)(args)), BuiltinScalarFunction::ArrayToString => Arc::new(|args| { make_scalar_function(array_expressions::array_to_string)(args) }), diff --git a/datafusion/proto/src/generated/prost.rs b/datafusion/proto/src/generated/prost.rs index 22a4beadb826..1898d1283b11 100644 --- a/datafusion/proto/src/generated/prost.rs +++ b/datafusion/proto/src/generated/prost.rs @@ -2379,6 +2379,7 @@ pub enum ScalarFunction { Iszero = 114, ArrayEmpty = 115, ArrayPopBack = 116, + ArraySort = 117, } impl ScalarFunction { /// String value of the enum field names used in the ProtoBuf definition. @@ -2504,6 +2505,7 @@ impl ScalarFunction { ScalarFunction::Iszero => "Iszero", ScalarFunction::ArrayEmpty => "ArrayEmpty", ScalarFunction::ArrayPopBack => "ArrayPopBack", + ScalarFunction::ArraySort => "ArraySort", } } /// Creates an enum from field names used in the ProtoBuf definition. diff --git a/datafusion/proto/src/logical_plan/from_proto.rs b/datafusion/proto/src/logical_plan/from_proto.rs index 7fbf2ff52c0d..9307a1fb16b0 100644 --- a/datafusion/proto/src/logical_plan/from_proto.rs +++ b/datafusion/proto/src/logical_plan/from_proto.rs @@ -476,6 +476,7 @@ impl From<&protobuf::ScalarFunction> for BuiltinScalarFunction { ScalarFunction::ArrayReplaceN => Self::ArrayReplaceN, ScalarFunction::ArrayReplaceAll => Self::ArrayReplaceAll, ScalarFunction::ArraySlice => Self::ArraySlice, + ScalarFunction::ArraySort => Self::ArraySort, ScalarFunction::ArrayToString => Self::ArrayToString, ScalarFunction::Cardinality => Self::Cardinality, ScalarFunction::Array => Self::MakeArray, diff --git a/datafusion/proto/src/logical_plan/to_proto.rs b/datafusion/proto/src/logical_plan/to_proto.rs index b90e3b0e4b0e..762686337006 100644 --- a/datafusion/proto/src/logical_plan/to_proto.rs +++ b/datafusion/proto/src/logical_plan/to_proto.rs @@ -1472,6 +1472,7 @@ impl TryFrom<&BuiltinScalarFunction> for protobuf::ScalarFunction { BuiltinScalarFunction::ArrayReplaceN => Self::ArrayReplaceN, BuiltinScalarFunction::ArrayReplaceAll => Self::ArrayReplaceAll, BuiltinScalarFunction::ArraySlice => Self::ArraySlice, + BuiltinScalarFunction::ArraySort => Self::ArraySort, BuiltinScalarFunction::ArrayToString => Self::ArrayToString, BuiltinScalarFunction::Cardinality => Self::Cardinality, BuiltinScalarFunction::MakeArray => Self::Array,