Skip to content

Commit

Permalink
remove null checking
Browse files Browse the repository at this point in the history
  • Loading branch information
@Weijun-H
Weijun-H committed Feb 25, 2024
1 parent cb37528 commit a7f121e
Show file tree
Hide file tree
Showing 4 changed files with 90 additions and 130 deletions.
44 changes: 15 additions & 29 deletions datafusion/expr/src/built_in_function.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -924,63 +924,51 @@ impl BuiltinScalarFunction {
Signature::variadic_any(self.volatility())
}
BuiltinScalarFunction::ArrayAppend => {
Signature::array_and_element(false, self.volatility())
Signature::array_and_element(self.volatility())
}
BuiltinScalarFunction::MakeArray => {
// 0 or more arguments of arbitrary type
Signature::one_of(vec![VariadicEqual, Any(0)], self.volatility())
}
BuiltinScalarFunction::ArrayPopFront => {
Signature::array(true, self.volatility())
}
BuiltinScalarFunction::ArrayPopBack => {
Signature::array(true, self.volatility())
}
BuiltinScalarFunction::ArrayPopFront => Signature::array(self.volatility()),
BuiltinScalarFunction::ArrayPopBack => Signature::array(self.volatility()),
BuiltinScalarFunction::ArrayConcat => {
Signature::variadic_any(self.volatility())
}
BuiltinScalarFunction::ArrayDims => {
Signature::array(false, self.volatility())
}
BuiltinScalarFunction::ArrayEmpty => {
Signature::array(true, self.volatility())
}
BuiltinScalarFunction::ArrayDims => Signature::array(self.volatility()),
BuiltinScalarFunction::ArrayEmpty => Signature::array(self.volatility()),
BuiltinScalarFunction::ArrayElement => {
Signature::array_and_index(true, self.volatility())
Signature::array_and_index(self.volatility())
}
BuiltinScalarFunction::ArrayExcept => Signature::any(2, self.volatility()),
BuiltinScalarFunction::Flatten => Signature::array(true, self.volatility()),
BuiltinScalarFunction::Flatten => Signature::array(self.volatility()),
BuiltinScalarFunction::ArrayHasAll | BuiltinScalarFunction::ArrayHasAny => {
Signature::any(2, self.volatility())
}
BuiltinScalarFunction::ArrayHas => {
Signature::array_and_element(true, self.volatility())
Signature::array_and_element(self.volatility())
}
BuiltinScalarFunction::ArrayLength => {
Signature::variadic_any(self.volatility())
}
BuiltinScalarFunction::ArrayNdims => {
Signature::array(false, self.volatility())
}
BuiltinScalarFunction::ArrayDistinct => {
Signature::array(true, self.volatility())
}
BuiltinScalarFunction::ArrayNdims => Signature::array(self.volatility()),
BuiltinScalarFunction::ArrayDistinct => Signature::array(self.volatility()),
BuiltinScalarFunction::ArrayPosition => {
Signature::array_and_element_and_optional_index(self.volatility())
}
BuiltinScalarFunction::ArrayPositions => {
Signature::array_and_element(true, self.volatility())
Signature::array_and_element(self.volatility())
}
BuiltinScalarFunction::ArrayPrepend => {
Signature::element_and_array(false, self.volatility())
Signature::element_and_array(self.volatility())
}
BuiltinScalarFunction::ArrayRepeat => Signature::any(2, self.volatility()),
BuiltinScalarFunction::ArrayRemove => {
Signature::array_and_element(true, self.volatility())
Signature::array_and_element(self.volatility())
}
BuiltinScalarFunction::ArrayRemoveN => Signature::any(3, self.volatility()),
BuiltinScalarFunction::ArrayRemoveAll => {
Signature::array_and_element(true, self.volatility())
Signature::array_and_element(self.volatility())
}
BuiltinScalarFunction::ArrayReplace => Signature::any(3, self.volatility()),
BuiltinScalarFunction::ArrayReplaceN => Signature::any(4, self.volatility()),
Expand All @@ -994,9 +982,7 @@ impl BuiltinScalarFunction {

BuiltinScalarFunction::ArrayIntersect => Signature::any(2, self.volatility()),
BuiltinScalarFunction::ArrayUnion => Signature::any(2, self.volatility()),
BuiltinScalarFunction::Cardinality => {
Signature::array(false, self.volatility())
}
BuiltinScalarFunction::Cardinality => Signature::array(self.volatility()),
BuiltinScalarFunction::ArrayResize => {
Signature::variadic_any(self.volatility())
}
Expand Down
86 changes: 31 additions & 55 deletions datafusion/expr/src/signature.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -126,34 +126,29 @@ pub enum TypeSignature {
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub enum ArrayFunctionSignature {
/// Specialized Signature for ArrayAppend and similar functions
/// If `allow_null` is true, the function also accepts a single argument of type Null.
/// The first argument should be List/LargeList/FixedSizedList, and the second argument should be non-list or list.
/// The second argument's list dimension should be one dimension less than the first argument's list dimension.
/// List dimension of the List/LargeList is equivalent to the number of List.
/// List dimension of the non-list is 0.
ArrayAndElement(bool),
ArrayAndElement,
/// Specialized Signature for ArrayPrepend and similar functions
/// If `allow_null` is true, the function also accepts a single argument of type Null.
/// The first argument should be non-list or list, and the second argument should be List/LargeList.
/// The first argument's list dimension should be one dimension less than the second argument's list dimension.
ElementAndArray(bool),
ElementAndArray,
/// Specialized Signature for Array functions of the form (List/LargeList, Index)
/// If `allow_null` is true, the function also accepts a single argument of type Null.
/// The first argument should be List/LargeList/FixedSizedList, and the second argument should be Int64.
ArrayAndIndex(bool),
ArrayAndIndex,
/// Specialized Signature for Array functions of the form (List/LargeList, Element, Optional Index)
ArrayAndElementAndOptionalIndex,
/// Specialized Signature for ArrayEmpty and similar functions
/// The function takes a single argument that must be a List/LargeList/FixedSizeList
/// or something that can be coerced to one of those types.
/// If `allow_null` is true, the function also accepts a single argument of type Null.
Array(bool),
Array,
}

impl ArrayFunctionSignature {
/// Arguments to ArrayFunctionSignature
/// `current_types` - The data types of the arguments
/// `allow_null_coercion` - Whether null type coercion is allowed
/// Returns the valid types for the function signature
pub fn get_type_signature(
&self,
Expand All @@ -169,7 +164,7 @@ impl ArrayFunctionSignature {

let first_two_types = &current_types[0..2];
let mut valid_types =
array_append_or_prepend_valid_types(first_two_types, true, true)?;
array_append_or_prepend_valid_types(first_two_types, true)?;

// Early return if there are only 2 arguments
if current_types.len() == 2 {
Expand All @@ -193,7 +188,6 @@ impl ArrayFunctionSignature {
fn array_append_or_prepend_valid_types(
current_types: &[DataType],
is_append: bool,
allow_null_coercion: bool,
) -> Result<Vec<Vec<DataType>>> {
if current_types.len() != 2 {
return Ok(vec![vec![]]);
Expand All @@ -207,9 +201,6 @@ impl ArrayFunctionSignature {

// We follow Postgres on `array_append(Null, T)`, which is not valid.
if array_type.eq(&DataType::Null) {
if allow_null_coercion {
return Ok(vec![vec![array_type.clone(), elem_type.clone()]]);
}
return Ok(vec![vec![]]);
}

Expand Down Expand Up @@ -245,28 +236,20 @@ impl ArrayFunctionSignature {
_ => Ok(vec![vec![]]),
}
}
fn array_and_index(
current_types: &[DataType],
allow_null_coercion: bool,
) -> Result<Vec<Vec<DataType>>> {
fn array_and_index(current_types: &[DataType]) -> Result<Vec<Vec<DataType>>> {
if current_types.len() != 2 {
return Ok(vec![vec![]]);
}
let array_type = &current_types[0];
let array_type = array(&[array_type.clone()], allow_null_coercion)?;
let array_type = array(&[array_type.clone()])?;
if array_type[0].is_empty() {
Ok(vec![vec![]])
} else {
Ok(vec![vec![array_type[0][0].clone(), DataType::Int64]])
}
}
fn array(
current_types: &[DataType],
allow_null_coercion: bool,
) -> Result<Vec<Vec<DataType>>> {
if current_types.len() != 1
|| (current_types[0].is_null() && !allow_null_coercion)
{
fn array(current_types: &[DataType]) -> Result<Vec<Vec<DataType>>> {
if current_types.len() != 1 {
return Ok(vec![vec![]]);
}

Expand All @@ -279,47 +262,42 @@ impl ArrayFunctionSignature {
let array_type = coerced_fixed_size_list_to_list(array_type);
Ok(vec![vec![array_type]])
}
DataType::Null => Ok(vec![vec![array_type.clone()]]),
_ => Ok(vec![vec![]]),
}
}
match self {
ArrayFunctionSignature::ArrayAndElement(allow_null) => {
array_append_or_prepend_valid_types(current_types, true, *allow_null)
}
ArrayFunctionSignature::ElementAndArray(allow_null) => {
array_append_or_prepend_valid_types(current_types, false, *allow_null)
ArrayFunctionSignature::ArrayAndElement => {
array_append_or_prepend_valid_types(current_types, true)
}
ArrayFunctionSignature::ArrayAndIndex(allow_null) => {
array_and_index(current_types, *allow_null)
ArrayFunctionSignature::ElementAndArray => {
array_append_or_prepend_valid_types(current_types, false)
}
ArrayFunctionSignature::ArrayAndIndex => array_and_index(current_types),
ArrayFunctionSignature::ArrayAndElementAndOptionalIndex => {
array_element_and_optional_index(current_types)
}
ArrayFunctionSignature::Array(allow_null) => {
array(current_types, *allow_null)
}
ArrayFunctionSignature::Array => array(current_types),
}
}
}

impl std::fmt::Display for ArrayFunctionSignature {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
ArrayFunctionSignature::ArrayAndElement(allow_null) => {
write!(f, "ArrayAndElement({})", *allow_null)
ArrayFunctionSignature::ArrayAndElement => {
write!(f, "ArrayAndElement")
}
ArrayFunctionSignature::ArrayAndElementAndOptionalIndex => {
write!(f, "array, element, [index]")
}
ArrayFunctionSignature::ElementAndArray(allow_null) => {
write!(f, "ElementAndArray({})", *allow_null)
ArrayFunctionSignature::ElementAndArray => {
write!(f, "ElementAndArray")
}
ArrayFunctionSignature::ArrayAndIndex(allow_null) => {
write!(f, "ArrayAndIndex({})", *allow_null)
ArrayFunctionSignature::ArrayAndIndex => {
write!(f, "ArrayAndIndex")
}
ArrayFunctionSignature::Array(allow_null) => {
write!(f, "Array({})", *allow_null)
ArrayFunctionSignature::Array => {
write!(f, "Array")
}
}
}
Expand Down Expand Up @@ -458,10 +436,10 @@ impl Signature {
}
}
/// Specialized Signature for ArrayAppend and similar functions
pub fn array_and_element(allow_null: bool, volatility: Volatility) -> Self {
pub fn array_and_element(volatility: Volatility) -> Self {
Signature {
type_signature: TypeSignature::ArraySignature(
ArrayFunctionSignature::ArrayAndElement(allow_null),
ArrayFunctionSignature::ArrayAndElement,
),
volatility,
}
Expand All @@ -476,29 +454,27 @@ impl Signature {
}
}
/// Specialized Signature for ArrayPrepend and similar functions
pub fn element_and_array(allow_null: bool, volatility: Volatility) -> Self {
pub fn element_and_array(volatility: Volatility) -> Self {
Signature {
type_signature: TypeSignature::ArraySignature(
ArrayFunctionSignature::ElementAndArray(allow_null),
ArrayFunctionSignature::ElementAndArray,
),
volatility,
}
}
/// Specialized Signature for ArrayElement and similar functions
pub fn array_and_index(allow_null: bool, volatility: Volatility) -> Self {
pub fn array_and_index(volatility: Volatility) -> Self {
Signature {
type_signature: TypeSignature::ArraySignature(
ArrayFunctionSignature::ArrayAndIndex(allow_null),
ArrayFunctionSignature::ArrayAndIndex,
),
volatility,
}
}
/// Specialized Signature for ArrayEmpty and similar functions
pub fn array(allow_null: bool, volatility: Volatility) -> Self {
pub fn array(volatility: Volatility) -> Self {
Signature {
type_signature: TypeSignature::ArraySignature(ArrayFunctionSignature::Array(
allow_null,
)),
type_signature: TypeSignature::ArraySignature(ArrayFunctionSignature::Array),
volatility,
}
}
Expand Down
11 changes: 0 additions & 11 deletions datafusion/physical-expr/src/array_expressions.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -791,7 +791,6 @@ pub fn array_pop_front(args: &[ArrayRef]) -> Result<ArrayRef> {
let array = as_large_list_array(&args[0])?;
general_pop_front_list::<i64>(array)
}
DataType::Null => Ok(args[0].clone()),
_ => exec_err!(
"array_pop_front does not support type: {:?}",
array_data_type
Expand All @@ -815,7 +814,6 @@ pub fn array_pop_back(args: &[ArrayRef]) -> Result<ArrayRef> {
let array = as_large_list_array(&args[0])?;
general_pop_back_list::<i64>(array)
}
DataType::Null => Ok(args[0].clone()),
_ => exec_err!(
"array_pop_back does not support type: {:?}",
array_data_type
Expand Down Expand Up @@ -1480,10 +1478,6 @@ pub fn array_positions(args: &[ArrayRef]) -> Result<ArrayRef> {
check_datatypes("array_positions", &[arr.values(), element])?;
general_positions::<i64>(arr, element)
}
DataType::Null => Ok(new_null_array(
&DataType::List(Arc::new(Field::new("item", DataType::UInt64, true))),
1,
)),
array_type => {
exec_err!("array_positions does not support type '{array_type:?}'.")
}
Expand Down Expand Up @@ -1616,10 +1610,6 @@ fn array_remove_internal(
element_array: &ArrayRef,
arr_n: Vec<i64>,
) -> Result<ArrayRef> {
if array.data_type().is_null() {
return Ok(array.clone());
}

match array.data_type() {
DataType::List(_) => {
let list_array = array.as_list::<i32>();
Expand Down Expand Up @@ -2294,7 +2284,6 @@ pub fn array_has(args: &[ArrayRef]) -> Result<ArrayRef> {
DataType::LargeList(_) => {
general_array_has_dispatch::<i64>(&args[0], &args[1], ComparisonType::Single)
}
DataType::Null => Ok(new_null_array(&DataType::Boolean, 1)),
_ => exec_err!("array_has does not support type '{array_type:?}'."),
}
}
Expand Down
Loading

0 comments on commit a7f121e

Please sign in to comment.