Skip to content

Commit

Permalink
Remove ArrayFunctionArguments struct
Browse files Browse the repository at this point in the history
  • Loading branch information
@jkosh44
jkosh44 committed Feb 12, 2025
1 parent 6c74609 commit 133b120
Show file tree
Hide file tree
Showing 6 changed files with 37 additions and 81 deletions.
78 changes: 20 additions & 58 deletions datafusion/expr-common/src/signature.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -230,7 +230,7 @@ pub enum ArrayFunctionSignature {
/// A function takes at least one List/LargeList/FixedSizeList argument.
Array {
/// A full list of the arguments accepted by this function.
arguments: ArrayFunctionArguments,
arguments: Vec<ArrayFunctionArgument>,
/// Additional information about how array arguments should be coerced.
array_coercion: Option<ListCoercion>,
},
Expand All @@ -246,7 +246,13 @@ impl Display for ArrayFunctionSignature {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
ArrayFunctionSignature::Array { arguments, .. } => {
write!(f, "{arguments}")
for (idx, argument) in arguments.iter().enumerate() {
write!(f, "{argument}")?;
if idx != arguments.len() - 1 {
write!(f, ", ")?;
}
}
Ok(())
}
ArrayFunctionSignature::RecursiveArray => {
write!(f, "recursive_array")
Expand All @@ -258,44 +264,6 @@ impl Display for ArrayFunctionSignature {
}
}

/// A wrapper around a vec of [`ArrayFunctionArgument`], to ensure that the vec has at least one
/// array element.
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash)]
pub struct ArrayFunctionArguments {
/// A full list of the arguments accepted by a function.
arguments: Vec<ArrayFunctionArgument>,
}

impl ArrayFunctionArguments {
/// Returns an error if there are no [`ArrayFunctionArgument::Array`] arguments.
pub fn new(arguments: Vec<ArrayFunctionArgument>) -> Result<Self, &'static str> {
if !arguments
.iter()
.any(|arg| *arg == ArrayFunctionArgument::Array)
{
Err("missing array argument")
} else {
Ok(Self { arguments })
}
}

pub fn inner(&self) -> &[ArrayFunctionArgument] {
self.arguments.as_slice()
}
}

impl Display for ArrayFunctionArguments {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
for (idx, argument) in self.arguments.iter().enumerate() {
write!(f, "{argument}")?;
if idx != self.arguments.len() - 1 {
write!(f, ", ")?;
}
}
Ok(())
}
}

#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash)]
pub enum ArrayFunctionArgument {
/// A non-list or list argument. The list dimensions should be one less than the Array's list
Expand Down Expand Up @@ -622,11 +590,10 @@ impl Signature {
Signature {
type_signature: TypeSignature::ArraySignature(
ArrayFunctionSignature::Array {
arguments: ArrayFunctionArguments::new(vec![
arguments: vec![
ArrayFunctionArgument::Array,
ArrayFunctionArgument::Element,
])
.expect("contains array"),
],
array_coercion,
},
),
Expand All @@ -641,20 +608,18 @@ impl Signature {
Signature {
type_signature: TypeSignature::OneOf(vec![
TypeSignature::ArraySignature(ArrayFunctionSignature::Array {
arguments: ArrayFunctionArguments::new(vec![
arguments: vec![
ArrayFunctionArgument::Array,
ArrayFunctionArgument::Element,
])
.expect("contains array"),
],
array_coercion: array_coercion.clone(),
}),
TypeSignature::ArraySignature(ArrayFunctionSignature::Array {
arguments: ArrayFunctionArguments::new(vec![
arguments: vec![
ArrayFunctionArgument::Array,
ArrayFunctionArgument::Element,
ArrayFunctionArgument::Index,
])
.expect("contains array"),
],
array_coercion,
}),
]),
Expand All @@ -669,11 +634,10 @@ impl Signature {
Signature {
type_signature: TypeSignature::ArraySignature(
ArrayFunctionSignature::Array {
arguments: ArrayFunctionArguments::new(vec![
arguments: vec![
ArrayFunctionArgument::Element,
ArrayFunctionArgument::Array,
])
.expect("contains array"),
],
array_coercion,
},
),
Expand All @@ -688,11 +652,10 @@ impl Signature {
Signature {
type_signature: TypeSignature::ArraySignature(
ArrayFunctionSignature::Array {
arguments: ArrayFunctionArguments::new(vec![
arguments: vec![
ArrayFunctionArgument::Array,
ArrayFunctionArgument::Index,
])
.expect("contains array"),
],
array_coercion,
},
),
Expand All @@ -704,10 +667,9 @@ impl Signature {
Signature {
type_signature: TypeSignature::ArraySignature(
ArrayFunctionSignature::Array {
arguments: ArrayFunctionArguments::new(vec![
arguments: vec![
ArrayFunctionArgument::Array,
])
.expect("contains array"),
],
array_coercion,
},
),
Expand Down
2 changes: 1 addition & 1 deletion datafusion/expr/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -71,7 +71,7 @@ pub use datafusion_expr_common::columnar_value::ColumnarValue;
pub use datafusion_expr_common::groups_accumulator::{EmitTo, GroupsAccumulator};
pub use datafusion_expr_common::operator::Operator;
pub use datafusion_expr_common::signature::{
ArrayFunctionArgument, ArrayFunctionArguments, ArrayFunctionSignature, Signature,
ArrayFunctionArgument, ArrayFunctionSignature, Signature,
TypeSignature, TypeSignatureClass, Volatility, TIMEZONE_WILDCARD,
};
pub use datafusion_expr_common::type_coercion::binary;
Expand Down
2 changes: 1 addition & 1 deletion datafusion/expr/src/type_coercion/functions.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -703,7 +703,7 @@ fn get_valid_types(
TypeSignature::Exact(valid_types) => vec![valid_types.clone()],
TypeSignature::ArraySignature(ref function_signature) => match function_signature {
ArrayFunctionSignature::Array { arguments, array_coercion, } => {
array_valid_types(function_name, current_types, arguments.inner(), array_coercion.as_ref())?
array_valid_types(function_name, current_types, arguments, array_coercion.as_ref())?
}
ArrayFunctionSignature::RecursiveArray => {
if current_types.len() != 1 {
Expand Down
7 changes: 3 additions & 4 deletions datafusion/functions-nested/src/cardinality.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@ use datafusion_common::cast::{as_large_list_array, as_list_array, as_map_array};
use datafusion_common::Result;
use datafusion_common::{exec_err, plan_err};
use datafusion_expr::{
ArrayFunctionArgument, ArrayFunctionArguments, ArrayFunctionSignature, ColumnarValue,
ArrayFunctionArgument, ArrayFunctionSignature, ColumnarValue,
Documentation, ScalarUDFImpl, Signature, TypeSignature, Volatility,
};
use datafusion_macros::user_doc;
Expand All @@ -50,10 +50,9 @@ impl Cardinality {
signature: Signature::one_of(
vec![
TypeSignature::ArraySignature(ArrayFunctionSignature::Array {
arguments: ArrayFunctionArguments::new(vec![
arguments: vec![
ArrayFunctionArgument::Array,
])
.expect("contains array"),
],
array_coercion: None,
}),
TypeSignature::ArraySignature(ArrayFunctionSignature::MapArray),
Expand Down
12 changes: 5 additions & 7 deletions datafusion/functions-nested/src/extract.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,7 @@ use datafusion_common::{
exec_err, internal_datafusion_err, plan_err, DataFusionError, Result,
};
use datafusion_expr::{
ArrayFunctionArgument, ArrayFunctionArguments, ArrayFunctionSignature, Expr,
ArrayFunctionArgument, ArrayFunctionSignature, Expr,
TypeSignature,
};
use datafusion_expr::{
Expand Down Expand Up @@ -335,22 +335,20 @@ impl ArraySlice {
signature: Signature::one_of(
vec![
TypeSignature::ArraySignature(ArrayFunctionSignature::Array {
arguments: ArrayFunctionArguments::new(vec![
arguments: vec![
ArrayFunctionArgument::Array,
ArrayFunctionArgument::Index,
ArrayFunctionArgument::Index,
])
.expect("contains array"),
],
array_coercion: None,
}),
TypeSignature::ArraySignature(ArrayFunctionSignature::Array {
arguments: ArrayFunctionArguments::new(vec![
arguments: vec![
ArrayFunctionArgument::Array,
ArrayFunctionArgument::Index,
ArrayFunctionArgument::Index,
ArrayFunctionArgument::Index,
])
.expect("contains array"),
],
array_coercion: None,
}),
],
Expand Down
17 changes: 7 additions & 10 deletions datafusion/functions-nested/src/replace.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@ use datafusion_common::cast::as_int64_array;
use datafusion_common::utils::ListCoercion;
use datafusion_common::{exec_err, Result};
use datafusion_expr::{
ArrayFunctionArgument, ArrayFunctionArguments, ArrayFunctionSignature, ColumnarValue,
ArrayFunctionArgument, ArrayFunctionSignature, ColumnarValue,
Documentation, ScalarUDFImpl, Signature, TypeSignature, Volatility,
};
use datafusion_macros::user_doc;
Expand Down Expand Up @@ -96,12 +96,11 @@ impl ArrayReplace {
signature: Signature {
type_signature: TypeSignature::ArraySignature(
ArrayFunctionSignature::Array {
arguments: ArrayFunctionArguments::new(vec![
arguments: vec![
ArrayFunctionArgument::Array,
ArrayFunctionArgument::Element,
ArrayFunctionArgument::Element,
])
.expect("contains array"),
],
array_coercion: Some(ListCoercion::FixedSizedListToList),
},
),
Expand Down Expand Up @@ -178,13 +177,12 @@ impl ArrayReplaceN {
signature: Signature {
type_signature: TypeSignature::ArraySignature(
ArrayFunctionSignature::Array {
arguments: ArrayFunctionArguments::new(vec![
arguments: vec![
ArrayFunctionArgument::Array,
ArrayFunctionArgument::Element,
ArrayFunctionArgument::Element,
ArrayFunctionArgument::Index,
])
.expect("contains array"),
],
array_coercion: Some(ListCoercion::FixedSizedListToList),
},
),
Expand Down Expand Up @@ -260,12 +258,11 @@ impl ArrayReplaceAll {
signature: Signature {
type_signature: TypeSignature::ArraySignature(
ArrayFunctionSignature::Array {
arguments: ArrayFunctionArguments::new(vec![
arguments: vec![
ArrayFunctionArgument::Array,
ArrayFunctionArgument::Element,
ArrayFunctionArgument::Element,
])
.expect("contains array"),
],
array_coercion: Some(ListCoercion::FixedSizedListToList),
},
),
Expand Down

0 comments on commit 133b120

Please sign in to comment.