Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Implement array_pop_back function #7348

Merged
merged 2 commits into from
Aug 23, 2023
Merged

Implement array_pop_back function #7348

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
8fb9dcc
feat: support `array_pop_back` function
tanruixiang Aug 21, 2023
4c00af7
Merge branch 'main' into support_array_pop_back
tanruixiang Aug 23, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions datafusion/expr/src/built_in_function.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -134,6 +134,8 @@ pub enum BuiltinScalarFunction {
ArrayHasAll,
/// array_has_any
ArrayHasAny,
/// array_pop_back
ArrayPopBack,
/// array_dims
ArrayDims,
/// array_element
Expand Down Expand Up @@ -367,6 +369,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,
Expand Down Expand Up @@ -548,6 +551,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))))
Expand Down Expand Up @@ -825,6 +829,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())
}
Expand Down Expand Up @@ -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",
Expand Down
9 changes: 9 additions & 0 deletions datafusion/expr/src/expr_fn.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -1081,6 +1089,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);
Expand Down
161 changes: 161 additions & 0 deletions datafusion/physical-expr/src/array_expressions.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -601,6 +601,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> {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@tanruixiang Do you want to implement array_pop_front function (See the ticket: #6996)? It has the similar syntax and it's a quite improvement your contribution in the project.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, I would love to implement array_pop_front , do I need to implement array_pop_front in this pr? Or in the next pr.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Next pr.

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];
Expand Down Expand Up @@ -1994,6 +2010,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]
Expand Down
4 changes: 3 additions & 1 deletion datafusion/physical-expr/src/functions.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -446,10 +446,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))
}
Expand Down
1 change: 1 addition & 0 deletions datafusion/proto/proto/datafusion.proto
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -597,6 +597,7 @@ enum ScalarFunction {
Flatten = 112;
Isnan = 113;
Iszero = 114;
ArrayPopBack = 115;
}

message ScalarFunctionNode {
Expand Down
3 changes: 3 additions & 0 deletions datafusion/proto/src/generated/pbjson.rs
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

3 changes: 3 additions & 0 deletions datafusion/proto/src/generated/prost.rs
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

9 changes: 8 additions & 1 deletion datafusion/proto/src/logical_plan/from_proto.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,6 @@ use datafusion_common::{
internal_err, Column, DFField, DFSchema, DFSchemaRef, DataFusionError,
OwnedTableReference, Result, ScalarValue,
};
use datafusion_expr::expr::{Alias, Placeholder};
use datafusion_expr::{
abs, acos, acosh, array, array_append, array_concat, array_dims, array_element,
array_has, array_has_all, array_has_any, array_length, array_ndims, array_position,
Expand All @@ -59,6 +58,10 @@ use datafusion_expr::{
JoinConstraint, JoinType, Like, Operator, TryCast, WindowFrame, WindowFrameBound,
WindowFrameUnits,
};
use datafusion_expr::{
array_pop_back,
expr::{Alias, Placeholder},
};
use std::sync::Arc;

#[derive(Debug)]
Expand Down Expand Up @@ -460,6 +463,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,
Expand Down Expand Up @@ -1268,6 +1272,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)?,
Expand Down
1 change: 1 addition & 0 deletions datafusion/proto/src/logical_plan/to_proto.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1459,6 +1459,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,
Expand Down
Loading