Skip to content

Commit 3ffeb52

Browse files
tanruixiangtanruixiang
authored
feat: support array_pop_back function (#7348)
1 parent bfe3e42 commit 3ffeb52

File tree

12 files changed

+293
-2
lines changed

12 files changed

+293
-2
lines changed

‎datafusion/expr/src/built_in_function.rs

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -134,6 +134,8 @@ pub enum BuiltinScalarFunction {
134134
ArrayHasAll,
135135
/// array_has_any
136136
ArrayHasAny,
137+
/// array_pop_back
138+
ArrayPopBack,
137139
/// array_dims
138140
ArrayDims,
139141
/// array_element
@@ -370,6 +372,7 @@ impl BuiltinScalarFunction {
370372
BuiltinScalarFunction::ArrayElement => Volatility::Immutable,
371373
BuiltinScalarFunction::ArrayLength => Volatility::Immutable,
372374
BuiltinScalarFunction::ArrayNdims => Volatility::Immutable,
375+
BuiltinScalarFunction::ArrayPopBack => Volatility::Immutable,
373376
BuiltinScalarFunction::ArrayPosition => Volatility::Immutable,
374377
BuiltinScalarFunction::ArrayPositions => Volatility::Immutable,
375378
BuiltinScalarFunction::ArrayPrepend => Volatility::Immutable,
@@ -552,6 +555,7 @@ impl BuiltinScalarFunction {
552555
},
553556
BuiltinScalarFunction::ArrayLength => Ok(UInt64),
554557
BuiltinScalarFunction::ArrayNdims => Ok(UInt64),
558+
BuiltinScalarFunction::ArrayPopBack => Ok(input_expr_types[0].clone()),
555559
BuiltinScalarFunction::ArrayPosition => Ok(UInt64),
556560
BuiltinScalarFunction::ArrayPositions => {
557561
Ok(List(Arc::new(Field::new("item", UInt64, true))))
@@ -823,6 +827,7 @@ impl BuiltinScalarFunction {
823827
// for now, the list is small, as we do not have many built-in functions.
824828
match self {
825829
BuiltinScalarFunction::ArrayAppend => Signature::any(2, self.volatility()),
830+
BuiltinScalarFunction::ArrayPopBack => Signature::any(1, self.volatility()),
826831
BuiltinScalarFunction::ArrayConcat => {
827832
Signature::variadic_any(self.volatility())
828833
}
@@ -1333,6 +1338,7 @@ fn aliases(func: &BuiltinScalarFunction) -> &'static [&'static str] {
13331338
}
13341339
BuiltinScalarFunction::ArrayLength => &["array_length", "list_length"],
13351340
BuiltinScalarFunction::ArrayNdims => &["array_ndims", "list_ndims"],
1341+
BuiltinScalarFunction::ArrayPopBack => &["array_pop_back", "list_pop_back"],
13361342
BuiltinScalarFunction::ArrayPosition => &[
13371343
"array_position",
13381344
"list_position",

‎datafusion/expr/src/expr_fn.rs

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -545,6 +545,14 @@ scalar_expr!(
545545
array element,
546546
"appends an element to the end of an array."
547547
);
548+
549+
scalar_expr!(
550+
ArrayPopBack,
551+
array_pop_back,
552+
array,
553+
"returns the array without the last element."
554+
);
555+
548556
nary_scalar_expr!(ArrayConcat, array_concat, "concatenates arrays.");
549557
scalar_expr!(
550558
ArrayHas,
@@ -1087,6 +1095,7 @@ mod test {
10871095
test_scalar_expr!(FromUnixtime, from_unixtime, unixtime);
10881096

10891097
test_scalar_expr!(ArrayAppend, array_append, array, element);
1098+
test_scalar_expr!(ArrayPopBack, array_pop_back, array);
10901099
test_unary_scalar_expr!(ArrayDims, array_dims);
10911100
test_scalar_expr!(ArrayLength, array_length, array, dimension);
10921101
test_unary_scalar_expr!(ArrayNdims, array_ndims);

‎datafusion/physical-expr/src/array_expressions.rs

Lines changed: 161 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -599,6 +599,22 @@ pub fn array_slice(args: &[ArrayRef]) -> Result<ArrayRef> {
599599
define_array_slice(list_array, key, extra_key, false)
600600
}
601601

602+
pub fn array_pop_back(args: &[ArrayRef]) -> Result<ArrayRef> {
603+
let list_array = as_list_array(&args[0])?;
604+
let key = vec![0; list_array.len()];
605+
let extra_key: Vec<_> = list_array
606+
.iter()
607+
.map(|x| x.map_or(0, |arr| arr.len() as i64 - 1))
608+
.collect();
609+
610+
define_array_slice(
611+
list_array,
612+
&Int64Array::from(key),
613+
&Int64Array::from(extra_key),
614+
false,
615+
)
616+
}
617+
602618
macro_rules! append {
603619
($ARRAY:expr, $ELEMENT:expr, $ARRAY_TYPE:ident) => {{
604620
let mut offsets: Vec<i32> = vec![0];
@@ -2005,6 +2021,151 @@ mod tests {
20052021
);
20062022
}
20072023

2024+
#[test]
2025+
fn test_array_pop_back() {
2026+
// array_pop_back([1, 2, 3, 4]) = [1, 2, 3]
2027+
let list_array = return_array().into_array(1);
2028+
let arr = array_pop_back(&[list_array])
2029+
.expect("failed to initialize function array_pop_back");
2030+
let result =
2031+
as_list_array(&arr).expect("failed to initialize function array_pop_back");
2032+
assert_eq!(
2033+
&[1, 2, 3],
2034+
result
2035+
.value(0)
2036+
.as_any()
2037+
.downcast_ref::<Int64Array>()
2038+
.unwrap()
2039+
.values()
2040+
);
2041+
2042+
// array_pop_back([1, 2, 3]) = [1, 2]
2043+
let list_array = Arc::new(result.clone());
2044+
let arr = array_pop_back(&[list_array])
2045+
.expect("failed to initialize function array_pop_back");
2046+
let result =
2047+
as_list_array(&arr).expect("failed to initialize function array_pop_back");
2048+
assert_eq!(
2049+
&[1, 2],
2050+
result
2051+
.value(0)
2052+
.as_any()
2053+
.downcast_ref::<Int64Array>()
2054+
.unwrap()
2055+
.values()
2056+
);
2057+
2058+
// array_pop_back([1, 2]) = [1]
2059+
let list_array = Arc::new(result.clone());
2060+
let arr = array_pop_back(&[list_array])
2061+
.expect("failed to initialize function array_pop_back");
2062+
let result =
2063+
as_list_array(&arr).expect("failed to initialize function array_pop_back");
2064+
assert_eq!(
2065+
&[1],
2066+
result
2067+
.value(0)
2068+
.as_any()
2069+
.downcast_ref::<Int64Array>()
2070+
.unwrap()
2071+
.values()
2072+
);
2073+
2074+
// array_pop_back([1]) = []
2075+
let list_array = Arc::new(result.clone());
2076+
let arr = array_pop_back(&[list_array])
2077+
.expect("failed to initialize function array_pop_back");
2078+
let result =
2079+
as_list_array(&arr).expect("failed to initialize function array_pop_back");
2080+
assert_eq!(
2081+
&[],
2082+
result
2083+
.value(0)
2084+
.as_any()
2085+
.downcast_ref::<Int64Array>()
2086+
.unwrap()
2087+
.values()
2088+
);
2089+
// array_pop_back([]) = []
2090+
let list_array = Arc::new(result.clone());
2091+
let arr = array_pop_back(&[list_array])
2092+
.expect("failed to initialize function array_pop_back");
2093+
let result =
2094+
as_list_array(&arr).expect("failed to initialize function array_pop_back");
2095+
assert_eq!(
2096+
&[],
2097+
result
2098+
.value(0)
2099+
.as_any()
2100+
.downcast_ref::<Int64Array>()
2101+
.unwrap()
2102+
.values()
2103+
);
2104+
2105+
// array_pop_back([1, NULL, 3, NULL]) = [1, NULL, 3]
2106+
let list_array = return_array_with_nulls().into_array(1);
2107+
let arr = array_pop_back(&[list_array])
2108+
.expect("failed to initialize function array_pop_back");
2109+
let result =
2110+
as_list_array(&arr).expect("failed to initialize function array_pop_back");
2111+
assert_eq!(3, result.values().len());
2112+
assert_eq!(
2113+
&[false, true, false],
2114+
&[
2115+
result.values().is_null(0),
2116+
result.values().is_null(1),
2117+
result.values().is_null(2)
2118+
]
2119+
);
2120+
}
2121+
#[test]
2122+
fn test_nested_array_pop_back() {
2123+
// array_pop_back([[1, 2, 3, 4], [5, 6, 7, 8]]) = [[1, 2, 3, 4]]
2124+
let list_array = return_nested_array().into_array(1);
2125+
let arr = array_pop_back(&[list_array])
2126+
.expect("failed to initialize function array_slice");
2127+
let result =
2128+
as_list_array(&arr).expect("failed to initialize function array_slice");
2129+
assert_eq!(
2130+
&[1, 2, 3, 4],
2131+
result
2132+
.value(0)
2133+
.as_any()
2134+
.downcast_ref::<ListArray>()
2135+
.unwrap()
2136+
.value(0)
2137+
.as_any()
2138+
.downcast_ref::<Int64Array>()
2139+
.unwrap()
2140+
.values()
2141+
);
2142+
2143+
// array_pop_back([[1, 2, 3, 4]]) = []
2144+
let list_array = Arc::new(result.clone());
2145+
let arr = array_pop_back(&[list_array])
2146+
.expect("failed to initialize function array_pop_back");
2147+
let result =
2148+
as_list_array(&arr).expect("failed to initialize function array_pop_back");
2149+
assert!(result
2150+
.value(0)
2151+
.as_any()
2152+
.downcast_ref::<ListArray>()
2153+
.unwrap()
2154+
.is_empty());
2155+
// array_pop_back([]) = []
2156+
let list_array = Arc::new(result.clone());
2157+
let arr = array_pop_back(&[list_array])
2158+
.expect("failed to initialize function array_pop_back");
2159+
let result =
2160+
as_list_array(&arr).expect("failed to initialize function array_pop_back");
2161+
assert!(result
2162+
.value(0)
2163+
.as_any()
2164+
.downcast_ref::<ListArray>()
2165+
.unwrap()
2166+
.is_empty());
2167+
}
2168+
20082169
#[test]
20092170
fn test_array_slice() {
20102171
// array_slice([1, 2, 3, 4], 1, 3) = [1, 2, 3]

‎datafusion/physical-expr/src/functions.rs

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -449,10 +449,12 @@ pub fn create_physical_fun(
449449
BuiltinScalarFunction::Flatten => {
450450
Arc::new(|args| make_scalar_function(array_expressions::flatten)(args))
451451
}
452-
453452
BuiltinScalarFunction::ArrayNdims => {
454453
Arc::new(|args| make_scalar_function(array_expressions::array_ndims)(args))
455454
}
455+
BuiltinScalarFunction::ArrayPopBack => {
456+
Arc::new(|args| make_scalar_function(array_expressions::array_pop_back)(args))
457+
}
456458
BuiltinScalarFunction::ArrayPosition => {
457459
Arc::new(|args| make_scalar_function(array_expressions::array_position)(args))
458460
}

‎datafusion/proto/proto/datafusion.proto

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -598,6 +598,7 @@ enum ScalarFunction {
598598
Isnan = 113;
599599
Iszero = 114;
600600
ArrayEmpty = 115;
601+
ArrayPopBack = 116;
601602
}
602603

603604
message ScalarFunctionNode {

‎datafusion/proto/src/generated/pbjson.rs

Lines changed: 3 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎datafusion/proto/src/generated/prost.rs

Lines changed: 3 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎datafusion/proto/src/logical_plan/from_proto.rs

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -59,7 +59,7 @@ use datafusion_expr::{
5959
WindowFrameUnits,
6060
};
6161
use datafusion_expr::{
62-
array_empty,
62+
array_empty, array_pop_back,
6363
expr::{Alias, Placeholder},
6464
};
6565
use std::sync::Arc;
@@ -464,6 +464,7 @@ impl From<&protobuf::ScalarFunction> for BuiltinScalarFunction {
464464
ScalarFunction::Flatten => Self::Flatten,
465465
ScalarFunction::ArrayLength => Self::ArrayLength,
466466
ScalarFunction::ArrayNdims => Self::ArrayNdims,
467+
ScalarFunction::ArrayPopBack => Self::ArrayPopBack,
467468
ScalarFunction::ArrayPosition => Self::ArrayPosition,
468469
ScalarFunction::ArrayPositions => Self::ArrayPositions,
469470
ScalarFunction::ArrayPrepend => Self::ArrayPrepend,
@@ -1272,6 +1273,9 @@ pub fn parse_expr(
12721273
parse_expr(&args[0], registry)?,
12731274
parse_expr(&args[1], registry)?,
12741275
)),
1276+
ScalarFunction::ArrayPopBack => {
1277+
Ok(array_pop_back(parse_expr(&args[0], registry)?))
1278+
}
12751279
ScalarFunction::ArrayPrepend => Ok(array_prepend(
12761280
parse_expr(&args[0], registry)?,
12771281
parse_expr(&args[1], registry)?,

‎datafusion/proto/src/logical_plan/to_proto.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1460,6 +1460,7 @@ impl TryFrom<&BuiltinScalarFunction> for protobuf::ScalarFunction {
14601460
BuiltinScalarFunction::Flatten => Self::Flatten,
14611461
BuiltinScalarFunction::ArrayLength => Self::ArrayLength,
14621462
BuiltinScalarFunction::ArrayNdims => Self::ArrayNdims,
1463+
BuiltinScalarFunction::ArrayPopBack => Self::ArrayPopBack,
14631464
BuiltinScalarFunction::ArrayPosition => Self::ArrayPosition,
14641465
BuiltinScalarFunction::ArrayPositions => Self::ArrayPositions,
14651466
BuiltinScalarFunction::ArrayPrepend => Self::ArrayPrepend,

0 commit comments

Comments
 (0)