tustvold commented on code in PR #6384:
URL: https://github.com/apache/arrow-datafusion/pull/6384#discussion_r1213002787
##########
datafusion/physical-expr/src/array_expressions.rs:
##########
@@ -115,3 +149,1560 @@ pub fn array(values: &[ColumnarValue]) ->
Result<ColumnarValue> {
.collect();
Ok(ColumnarValue::Array(array_array(arrays.as_slice())?))
}
+
+macro_rules! downcast_arg {
+ ($ARG:expr, $ARRAY_TYPE:ident) => {{
+ $ARG.as_any().downcast_ref::<$ARRAY_TYPE>().ok_or_else(|| {
+ DataFusionError::Internal(format!(
+ "could not cast to {}",
+ type_name::<$ARRAY_TYPE>()
+ ))
+ })?
+ }};
+}
+
+macro_rules! append {
+ ($ARRAY:expr, $ELEMENT:expr, $ARRAY_TYPE:ident) => {{
+ let child_array =
+ downcast_arg!(downcast_arg!($ARRAY, ListArray).values(),
$ARRAY_TYPE);
+ let element = downcast_arg!($ELEMENT, $ARRAY_TYPE);
+ let concat = compute::concat(&[child_array, element])?;
+ let mut scalars = vec![];
+ for i in 0..concat.len() {
+ scalars.push(ColumnarValue::Scalar(ScalarValue::try_from_array(
+ &concat, i,
+ )?));
+ }
+ scalars
+ }};
+}
+
+/// Array_append SQL function
+pub fn array_append(args: &[ColumnarValue]) -> Result<ColumnarValue> {
+ if args.len() != 2 {
+ return Err(DataFusionError::Internal(format!(
+ "Array_append function requires two arguments, got {}",
+ args.len()
+ )));
+ }
+
+ let arr = match &args[0] {
+ ColumnarValue::Scalar(scalar) => scalar.to_array().clone(),
+ ColumnarValue::Array(arr) => arr.clone(),
+ };
+
+ let element = match &args[1] {
+ ColumnarValue::Scalar(scalar) => scalar.to_array().clone(),
+ _ => {
+ return Err(DataFusionError::Internal(
+ "Array_append function requires scalar element".to_string(),
+ ))
+ }
+ };
+
+ let data_type = arr.data_type();
+ let arrays = match data_type {
+ DataType::List(field) => {
+ match (field.data_type(), element.data_type()) {
+ (DataType::Utf8, DataType::Utf8) => append!(arr, element,
StringArray),
+ (DataType::LargeUtf8, DataType::LargeUtf8) => append!(arr,
element, LargeStringArray),
+ (DataType::Boolean, DataType::Boolean) => append!(arr,
element, BooleanArray),
+ (DataType::Float32, DataType::Float32) => append!(arr,
element, Float32Array),
+ (DataType::Float64, DataType::Float64) => append!(arr,
element, Float64Array),
+ (DataType::Int8, DataType::Int8) => append!(arr, element,
Int8Array),
+ (DataType::Int16, DataType::Int16) => append!(arr, element,
Int16Array),
+ (DataType::Int32, DataType::Int32) => append!(arr, element,
Int32Array),
+ (DataType::Int64, DataType::Int64) => append!(arr, element,
Int64Array),
+ (DataType::UInt8, DataType::UInt8) => append!(arr, element,
UInt8Array),
+ (DataType::UInt16, DataType::UInt16) => append!(arr, element,
UInt16Array),
+ (DataType::UInt32, DataType::UInt32) => append!(arr, element,
UInt32Array),
+ (DataType::UInt64, DataType::UInt64) => append!(arr, element,
UInt64Array),
+ (array_data_type, element_data_type) => {
+ return Err(DataFusionError::NotImplemented(format!(
+ "Array_append is not implemented for types
'{array_data_type:?}' and '{element_data_type:?}'."
+ )))
+ }
+ }
+ }
+ data_type => {
+ return Err(DataFusionError::Internal(format!(
+ "Array is not type '{data_type:?}'."
+ )))
+ }
+ };
+
+ array(arrays.as_slice())
+}
+
+macro_rules! prepend {
+ ($ARRAY:expr, $ELEMENT:expr, $ARRAY_TYPE:ident) => {{
+ let child_array =
+ downcast_arg!(downcast_arg!($ARRAY, ListArray).values(),
$ARRAY_TYPE);
+ let element = downcast_arg!($ELEMENT, $ARRAY_TYPE);
+ let concat = compute::concat(&[element, child_array])?;
+ let mut scalars = vec![];
+ for i in 0..concat.len() {
+ scalars.push(ColumnarValue::Scalar(ScalarValue::try_from_array(
+ &concat, i,
+ )?));
+ }
+ scalars
+ }};
+}
+
+/// Array_prepend SQL function
+pub fn array_prepend(args: &[ColumnarValue]) -> Result<ColumnarValue> {
+ if args.len() != 2 {
+ return Err(DataFusionError::Internal(format!(
+ "Array_prepend function requires two arguments, got {}",
+ args.len()
+ )));
+ }
+
+ let element = match &args[0] {
+ ColumnarValue::Scalar(scalar) => scalar.to_array().clone(),
+ _ => {
+ return Err(DataFusionError::Internal(
+ "Array_prepend function requires scalar element".to_string(),
+ ))
+ }
+ };
+
+ let arr = match &args[1] {
+ ColumnarValue::Scalar(scalar) => scalar.to_array().clone(),
+ ColumnarValue::Array(arr) => arr.clone(),
+ };
+
+ let data_type = arr.data_type();
+ let arrays = match data_type {
+ DataType::List(field) => {
+ match (field.data_type(), element.data_type()) {
+ (DataType::Utf8, DataType::Utf8) => prepend!(arr, element,
StringArray),
+ (DataType::LargeUtf8, DataType::LargeUtf8) => prepend!(arr,
element, LargeStringArray),
+ (DataType::Boolean, DataType::Boolean) => prepend!(arr,
element, BooleanArray),
+ (DataType::Float32, DataType::Float32) => prepend!(arr,
element, Float32Array),
+ (DataType::Float64, DataType::Float64) => prepend!(arr,
element, Float64Array),
+ (DataType::Int8, DataType::Int8) => prepend!(arr, element,
Int8Array),
+ (DataType::Int16, DataType::Int16) => prepend!(arr, element,
Int16Array),
+ (DataType::Int32, DataType::Int32) => prepend!(arr, element,
Int32Array),
+ (DataType::Int64, DataType::Int64) => prepend!(arr, element,
Int64Array),
+ (DataType::UInt8, DataType::UInt8) => prepend!(arr, element,
UInt8Array),
+ (DataType::UInt16, DataType::UInt16) => prepend!(arr, element,
UInt16Array),
+ (DataType::UInt32, DataType::UInt32) => prepend!(arr, element,
UInt32Array),
+ (DataType::UInt64, DataType::UInt64) => prepend!(arr, element,
UInt64Array),
+ (array_data_type, element_data_type) => {
+ return Err(DataFusionError::NotImplemented(format!(
+ "Array_prepend is not implemented for types
'{array_data_type:?}' and '{element_data_type:?}'."
+ )))
+ }
+ }
+ }
+ data_type => {
+ return Err(DataFusionError::Internal(format!(
+ "Array is not type '{data_type:?}'."
+ )))
+ }
+ };
+
+ array(arrays.as_slice())
+}
+
+/// Array_concat/Array_cat SQL function
+pub fn array_concat(args: &[ColumnarValue]) -> Result<ColumnarValue> {
+ let arrays: Vec<ArrayRef> = args
+ .iter()
+ .map(|x| match x {
+ ColumnarValue::Array(array) => array.clone(),
+ ColumnarValue::Scalar(scalar) => scalar.to_array().clone(),
+ })
+ .collect();
+ let data_type = arrays[0].data_type();
+ match data_type {
+ DataType::List(..) => {
+ let list_arrays =
Review Comment:
I think you could just call `to_data()` I'm not sure this needs to downcast
to ListArray
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
