tustvold commented on code in PR #6384:
URL: https://github.com/apache/arrow-datafusion/pull/6384#discussion_r1213005254
##########
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> {
Review Comment:
It would perhaps be nicer to use a combination of
https://docs.rs/arrow-array/latest/arrow_array/array/struct.GenericListArray.html#method.try_new
and https://docs.rs/arrow-select/latest/arrow_select/concat/index.html
--
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]
