feat: support array_pop_back function (#7348)

apache · Aug 23, 2023 · 3ffeb52 · 3ffeb52
1 parent bfe3e42
commit 3ffeb52
Show file tree

Hide file tree

Showing 12 changed files with 293 additions and 2 deletions.
diff --git a/datafusion/expr/src/built_in_function.rs b/datafusion/expr/src/built_in_function.rs
@@ -134,6 +134,8 @@ pub enum BuiltinScalarFunction {
     ArrayHasAll,
     /// array_has_any
     ArrayHasAny,
+    /// array_pop_back
+    ArrayPopBack,
     /// array_dims
     ArrayDims,
     /// array_element
@@ -370,6 +372,7 @@ impl BuiltinScalarFunction {
             BuiltinScalarFunction::ArrayElement => Volatility::Immutable,
             BuiltinScalarFunction::ArrayLength => Volatility::Immutable,
             BuiltinScalarFunction::ArrayNdims => Volatility::Immutable,
+            BuiltinScalarFunction::ArrayPopBack => Volatility::Immutable,
             BuiltinScalarFunction::ArrayPosition => Volatility::Immutable,
             BuiltinScalarFunction::ArrayPositions => Volatility::Immutable,
             BuiltinScalarFunction::ArrayPrepend => Volatility::Immutable,
@@ -552,6 +555,7 @@ impl BuiltinScalarFunction {
             },
             BuiltinScalarFunction::ArrayLength => Ok(UInt64),
             BuiltinScalarFunction::ArrayNdims => Ok(UInt64),
+            BuiltinScalarFunction::ArrayPopBack => Ok(input_expr_types[0].clone()),
             BuiltinScalarFunction::ArrayPosition => Ok(UInt64),
             BuiltinScalarFunction::ArrayPositions => {
                 Ok(List(Arc::new(Field::new("item", UInt64, true))))
@@ -823,6 +827,7 @@ impl BuiltinScalarFunction {
         // for now, the list is small, as we do not have many built-in functions.
         match self {
             BuiltinScalarFunction::ArrayAppend => Signature::any(2, self.volatility()),
+            BuiltinScalarFunction::ArrayPopBack => Signature::any(1, self.volatility()),
             BuiltinScalarFunction::ArrayConcat => {
                 Signature::variadic_any(self.volatility())
             }
@@ -1333,6 +1338,7 @@ fn aliases(func: &BuiltinScalarFunction) -> &'static [&'static str] {
         }
         BuiltinScalarFunction::ArrayLength => &["array_length", "list_length"],
         BuiltinScalarFunction::ArrayNdims => &["array_ndims", "list_ndims"],
+        BuiltinScalarFunction::ArrayPopBack => &["array_pop_back", "list_pop_back"],
         BuiltinScalarFunction::ArrayPosition => &[
             "array_position",
             "list_position",

diff --git a/datafusion/expr/src/expr_fn.rs b/datafusion/expr/src/expr_fn.rs
@@ -545,6 +545,14 @@ scalar_expr!(
     array element,
     "appends an element to the end of an array."
 );
+
+scalar_expr!(
+    ArrayPopBack,
+    array_pop_back,
+    array,
+    "returns the array without the last element."
+);
+
 nary_scalar_expr!(ArrayConcat, array_concat, "concatenates arrays.");
 scalar_expr!(
     ArrayHas,
@@ -1087,6 +1095,7 @@ mod test {
         test_scalar_expr!(FromUnixtime, from_unixtime, unixtime);
 
         test_scalar_expr!(ArrayAppend, array_append, array, element);
+        test_scalar_expr!(ArrayPopBack, array_pop_back, array);
         test_unary_scalar_expr!(ArrayDims, array_dims);
         test_scalar_expr!(ArrayLength, array_length, array, dimension);
         test_unary_scalar_expr!(ArrayNdims, array_ndims);

diff --git a/datafusion/physical-expr/src/array_expressions.rs b/datafusion/physical-expr/src/array_expressions.rs
@@ -599,6 +599,22 @@ pub fn array_slice(args: &[ArrayRef]) -> Result<ArrayRef> {
     define_array_slice(list_array, key, extra_key, false)
 }
 
+pub fn array_pop_back(args: &[ArrayRef]) -> Result<ArrayRef> {
+    let list_array = as_list_array(&args[0])?;
+    let key = vec![0; list_array.len()];
+    let extra_key: Vec<_> = list_array
+        .iter()
+        .map(|x| x.map_or(0, |arr| arr.len() as i64 - 1))
+        .collect();
+
+    define_array_slice(
+        list_array,
+        &Int64Array::from(key),
+        &Int64Array::from(extra_key),
+        false,
+    )
+}
+
 macro_rules! append {
     ($ARRAY:expr, $ELEMENT:expr, $ARRAY_TYPE:ident) => {{
         let mut offsets: Vec<i32> = vec![0];
@@ -2005,6 +2021,151 @@ mod tests {
         );
     }
 
+    #[test]
+    fn test_array_pop_back() {
+        // array_pop_back([1, 2, 3, 4]) = [1, 2, 3]
+        let list_array = return_array().into_array(1);
+        let arr = array_pop_back(&[list_array])
+            .expect("failed to initialize function array_pop_back");
+        let result =
+            as_list_array(&arr).expect("failed to initialize function array_pop_back");
+        assert_eq!(
+            &[1, 2, 3],
+            result
+                .value(0)
+                .as_any()
+                .downcast_ref::<Int64Array>()
+                .unwrap()
+                .values()
+        );
+
+        // array_pop_back([1, 2, 3]) = [1, 2]
+        let list_array = Arc::new(result.clone());
+        let arr = array_pop_back(&[list_array])
+            .expect("failed to initialize function array_pop_back");
+        let result =
+            as_list_array(&arr).expect("failed to initialize function array_pop_back");
+        assert_eq!(
+            &[1, 2],
+            result
+                .value(0)
+                .as_any()
+                .downcast_ref::<Int64Array>()
+                .unwrap()
+                .values()
+        );
+
+        // array_pop_back([1, 2]) = [1]
+        let list_array = Arc::new(result.clone());
+        let arr = array_pop_back(&[list_array])
+            .expect("failed to initialize function array_pop_back");
+        let result =
+            as_list_array(&arr).expect("failed to initialize function array_pop_back");
+        assert_eq!(
+            &[1],
+            result
+                .value(0)
+                .as_any()
+                .downcast_ref::<Int64Array>()
+                .unwrap()
+                .values()
+        );
+
+        // array_pop_back([1]) = []
+        let list_array = Arc::new(result.clone());
+        let arr = array_pop_back(&[list_array])
+            .expect("failed to initialize function array_pop_back");
+        let result =
+            as_list_array(&arr).expect("failed to initialize function array_pop_back");
+        assert_eq!(
+            &[],
+            result
+                .value(0)
+                .as_any()
+                .downcast_ref::<Int64Array>()
+                .unwrap()
+                .values()
+        );
+        // array_pop_back([]) = []
+        let list_array = Arc::new(result.clone());
+        let arr = array_pop_back(&[list_array])
+            .expect("failed to initialize function array_pop_back");
+        let result =
+            as_list_array(&arr).expect("failed to initialize function array_pop_back");
+        assert_eq!(
+            &[],
+            result
+                .value(0)
+                .as_any()
+                .downcast_ref::<Int64Array>()
+                .unwrap()
+                .values()
+        );
+
+        // array_pop_back([1, NULL, 3, NULL]) = [1, NULL, 3]
+        let list_array = return_array_with_nulls().into_array(1);
+        let arr = array_pop_back(&[list_array])
+            .expect("failed to initialize function array_pop_back");
+        let result =
+            as_list_array(&arr).expect("failed to initialize function array_pop_back");
+        assert_eq!(3, result.values().len());
+        assert_eq!(
+            &[false, true, false],
+            &[
+                result.values().is_null(0),
+                result.values().is_null(1),
+                result.values().is_null(2)
+            ]
+        );
+    }
+    #[test]
+    fn test_nested_array_pop_back() {
+        // array_pop_back([[1, 2, 3, 4], [5, 6, 7, 8]]) = [[1, 2, 3, 4]]
+        let list_array = return_nested_array().into_array(1);
+        let arr = array_pop_back(&[list_array])
+            .expect("failed to initialize function array_slice");
+        let result =
+            as_list_array(&arr).expect("failed to initialize function array_slice");
+        assert_eq!(
+            &[1, 2, 3, 4],
+            result
+                .value(0)
+                .as_any()
+                .downcast_ref::<ListArray>()
+                .unwrap()
+                .value(0)
+                .as_any()
+                .downcast_ref::<Int64Array>()
+                .unwrap()
+                .values()
+        );
+
+        // array_pop_back([[1, 2, 3, 4]]) = []
+        let list_array = Arc::new(result.clone());
+        let arr = array_pop_back(&[list_array])
+            .expect("failed to initialize function array_pop_back");
+        let result =
+            as_list_array(&arr).expect("failed to initialize function array_pop_back");
+        assert!(result
+            .value(0)
+            .as_any()
+            .downcast_ref::<ListArray>()
+            .unwrap()
+            .is_empty());
+        // array_pop_back([]) = []
+        let list_array = Arc::new(result.clone());
+        let arr = array_pop_back(&[list_array])
+            .expect("failed to initialize function array_pop_back");
+        let result =
+            as_list_array(&arr).expect("failed to initialize function array_pop_back");
+        assert!(result
+            .value(0)
+            .as_any()
+            .downcast_ref::<ListArray>()
+            .unwrap()
+            .is_empty());
+    }
+
     #[test]
     fn test_array_slice() {
         // array_slice([1, 2, 3, 4], 1, 3) = [1, 2, 3]

diff --git a/datafusion/physical-expr/src/functions.rs b/datafusion/physical-expr/src/functions.rs
@@ -449,10 +449,12 @@ pub fn create_physical_fun(
         BuiltinScalarFunction::Flatten => {
             Arc::new(|args| make_scalar_function(array_expressions::flatten)(args))
         }
-
         BuiltinScalarFunction::ArrayNdims => {
             Arc::new(|args| make_scalar_function(array_expressions::array_ndims)(args))
         }
+        BuiltinScalarFunction::ArrayPopBack => {
+            Arc::new(|args| make_scalar_function(array_expressions::array_pop_back)(args))
+        }
         BuiltinScalarFunction::ArrayPosition => {
             Arc::new(|args| make_scalar_function(array_expressions::array_position)(args))
         }

diff --git a/datafusion/proto/proto/datafusion.proto b/datafusion/proto/proto/datafusion.proto
@@ -598,6 +598,7 @@ enum ScalarFunction {
   Isnan = 113;
   Iszero = 114;
   ArrayEmpty = 115;
+  ArrayPopBack = 116;
 }
 
 message ScalarFunctionNode {

diff --git a/datafusion/proto/src/generated/pbjson.rs b/datafusion/proto/src/generated/pbjson.rs
diff --git a/datafusion/proto/src/generated/prost.rs b/datafusion/proto/src/generated/prost.rs
diff --git a/datafusion/proto/src/logical_plan/from_proto.rs b/datafusion/proto/src/logical_plan/from_proto.rs
@@ -59,7 +59,7 @@ use datafusion_expr::{
     WindowFrameUnits,
 };
 use datafusion_expr::{
-    array_empty,
+    array_empty, array_pop_back,
     expr::{Alias, Placeholder},
 };
 use std::sync::Arc;
@@ -464,6 +464,7 @@ impl From<&protobuf::ScalarFunction> for BuiltinScalarFunction {
             ScalarFunction::Flatten => Self::Flatten,
             ScalarFunction::ArrayLength => Self::ArrayLength,
             ScalarFunction::ArrayNdims => Self::ArrayNdims,
+            ScalarFunction::ArrayPopBack => Self::ArrayPopBack,
             ScalarFunction::ArrayPosition => Self::ArrayPosition,
             ScalarFunction::ArrayPositions => Self::ArrayPositions,
             ScalarFunction::ArrayPrepend => Self::ArrayPrepend,
@@ -1272,6 +1273,9 @@ pub fn parse_expr(
                     parse_expr(&args[0], registry)?,
                     parse_expr(&args[1], registry)?,
                 )),
+                ScalarFunction::ArrayPopBack => {
+                    Ok(array_pop_back(parse_expr(&args[0], registry)?))
+                }
                 ScalarFunction::ArrayPrepend => Ok(array_prepend(
                     parse_expr(&args[0], registry)?,
                     parse_expr(&args[1], registry)?,

diff --git a/datafusion/proto/src/logical_plan/to_proto.rs b/datafusion/proto/src/logical_plan/to_proto.rs
@@ -1460,6 +1460,7 @@ impl TryFrom<&BuiltinScalarFunction> for protobuf::ScalarFunction {
             BuiltinScalarFunction::Flatten => Self::Flatten,
             BuiltinScalarFunction::ArrayLength => Self::ArrayLength,
             BuiltinScalarFunction::ArrayNdims => Self::ArrayNdims,
+            BuiltinScalarFunction::ArrayPopBack => Self::ArrayPopBack,
             BuiltinScalarFunction::ArrayPosition => Self::ArrayPosition,
             BuiltinScalarFunction::ArrayPositions => Self::ArrayPositions,
             BuiltinScalarFunction::ArrayPrepend => Self::ArrayPrepend,