Wip wop

Author: edmondop

datafusion/expr/src/expr_fn.rs

@@ -706,7 +706,6 @@ scalar_expr!(
 );
 nary_scalar_expr!(ArrayUnion, array_union, "returns an array of the elements in the union of array1 and array2 without duplicates.");
 
-
 scalar_expr!(
     Cardinality,
     cardinality,

datafusion/physical-expr/src/array_expressions.rs

@@ -18,20 +18,19 @@
 //! Array expressions
 
 use arrow::array::*;
-use arrow_array::types::Int64Type;
 use arrow::buffer::OffsetBuffer;
 use arrow::compute;
 use arrow::datatypes::{DataType, Field, UInt64Type};
-use arrow::row::{RowConverter, SortField, Row};
+use arrow::row::{RowConverter, SortField};
 use arrow_buffer::NullBuffer;
 use core::any::type_name;
 use datafusion_common::cast::{as_generic_string_array, as_int64_array, as_list_array};
 use datafusion_common::{exec_err, internal_err, not_impl_err, plan_err, ScalarValue};
 use datafusion_common::{DataFusionError, Result};
 use datafusion_expr::ColumnarValue;
 use itertools::Itertools;
-use std::sync::Arc;
 use std::collections::HashSet;
+use std::sync::Arc;
 
 macro_rules! downcast_arg {
     ($ARG:expr, $ARRAY_TYPE:ident) => {{
@@ -1488,7 +1487,7 @@ macro_rules! to_string {
 //         let arr = list_arr.value(i);
 //         let i64arr = as_primitive_array::<Int64Type>(&arr);
 //         for v in i64arr.iter() {
-            
+
 //             // v is Option<i64>
 //         }
 //     }
@@ -1504,15 +1503,15 @@ macro_rules! to_string {
 //     converted.iter().filter(|row| dedup.insert(*row)).for_each(|row| distinct_rows.push(row));
 //     let dedup =  row_converter.convert_rows(&distinct_rows)?;
 //     let res =  make_array(dedup.as_slice())?;
-    
+
 //     Ok(res)
 // }
 
-fn union_generic_lists<OffsetSize: OffsetSizeTrait>(l: &GenericListArray<OffsetSize>, r: &GenericListArray<OffsetSize>) -> Result<GenericListArray<OffsetSize>, DataFusionError>{
-    let converter = RowConverter::new(vec![
-        SortField::new(
-            l.data_type().clone()
-    )])?;
+fn union_generic_lists<OffsetSize: OffsetSizeTrait>(
+    l: &GenericListArray<OffsetSize>,
+    r: &GenericListArray<OffsetSize>,
+) -> Result<GenericListArray<OffsetSize>, DataFusionError> {
+    let converter = RowConverter::new(vec![SortField::new(l.data_type().clone())])?;
     let mut dedup = HashSet::new();
     let nulls = NullBuffer::union(l.nulls(), r.nulls());
     let field = Arc::new(Field::new("union", l.data_type().to_owned(), false));
@@ -1529,13 +1528,13 @@ fn union_generic_lists<OffsetSize: OffsetSizeTrait>(l: &GenericListArray<OffsetS
     for (l_w, r_w) in l.offsets().windows(2).zip(r.offsets().windows(2)) {
         let l_slice = l_w[0].as_usize()..l_w[1].as_usize();
         let r_slice = r_w[0].as_usize()..r_w[1].as_usize();
-        for i in l_slice{
+        for i in l_slice {
             dedup.insert(l_values.row(i));
         }
-        for i in r_slice{
+        for i in r_slice {
             dedup.insert(r_values.row(i));
         }
-    
+
         rows.extend(dedup.iter());
         offsets.push(OffsetSize::usize_as(rows.len()));
         dedup.clear();
@@ -1544,45 +1543,48 @@ fn union_generic_lists<OffsetSize: OffsetSizeTrait>(l: &GenericListArray<OffsetS
     let values = converter.convert_rows(rows)?;
     let offsets = OffsetBuffer::new(offsets.into());
     let result = values[0].clone();
-    Ok(GenericListArray::<OffsetSize>::new(field, offsets, result, nulls))
+    Ok(GenericListArray::<OffsetSize>::new(
+        field, offsets, result, nulls,
+    ))
 }
 
-
 /// Array_union SQL function
 pub fn array_union(args: &[ArrayRef]) -> Result<ArrayRef> {
     if args.len() != 2 {
-        return exec_err!("array_union needs two arguments")
+        return exec_err!("array_union needs two arguments");
     }
     let array1 = &args[0];
-    let array2= &args[1];
+    let array2 = &args[1];
     check_datatypes("array_union", &[&array1, &array2])?;
-    match (array1.data_type(), array2.data_type()){
-        (DataType::Null, _) => {
-            Ok(array2.clone())
-        },
-        (_, DataType::Null) => {
-            Ok(array1.clone())
+    match (array1.data_type(), array2.data_type()) {
+        (DataType::Null, _) => Ok(array2.clone()),
+        (_, DataType::Null) => Ok(array1.clone()),
+        (DataType::List(field), _) => {
+            if field.data_type().equals_datatype(&DataType::Int32) {
+                let list1 = array1.as_list::<i32>();
+                let list2 = array2.as_list::<i32>();
+                let result = union_generic_lists::<i32>(list1, list2)?;
+                Ok(result.values().clone())
+            } else if field.data_type().equals_datatype(&DataType::Int64) {
+                eprintln!("{:?}", array1);
+                let list1 = array1.as_list::<i64>();
+                let list2 = array2.as_list::<i64>();
+                let result = union_generic_lists::<i64>(list1, list2)?;
+                Ok(result.values().clone())
+            } else {
+                return internal_err!(
+                    "array_union only support list with offsets of type int32 and int64"
+                );
+            }
         }
-        (DataType::Int32,    DataType::Int32) =>  {
-            let list1 = array1.as_list();
-            let list2 = array2.as_list();
-            let result = union_generic_lists::<i32>(list1, list2)?;
-            Ok(result.values().clone())
-        },
-        (DataType::Int64,    DataType::Int64) =>  {
-            let list1 = array1.as_list();
-            let list2 = array2.as_list();
-            let result =  union_generic_lists::<i64>(list1, list2)?;
-            Ok(result.values().clone())
-        },
         _ => {
-            return internal_err!("array_union only support list with offsets of type int32 and int64");
-        },
-
+            return internal_err!(
+                "array_union only support list with offsets of type int32 and int64"
+            );
+        }
     }
 }
 
-
 /// Array_to_string SQL function
 pub fn array_to_string(args: &[ArrayRef]) -> Result<ArrayRef> {
     let arr = &args[0];
@@ -1694,7 +1696,6 @@ pub fn array_to_string(args: &[ArrayRef]) -> Result<ArrayRef> {
     Ok(Arc::new(StringArray::from(res)))
 }
 
-
 /// Cardinality SQL function
 pub fn cardinality(args: &[ArrayRef]) -> Result<ArrayRef> {
     let list_array = as_list_array(&args[0])?.clone();