Initial Implementation of array_intersect

Veeupup · Veeupup · commit 522f698f5426 · 2023-11-08T00:50:22.000+08:00
Signed-off-by: veeupup &lt;code@tanweime.com&gt;
diff --git a/datafusion/core/tests/sqllogictests/test_files/array.slt b/datafusion/core/tests/sqllogictests/test_files/array.slt
@@ -1695,6 +1695,15 @@ select array_has_all(make_array(1,2,3), make_array(1,3)),
 ----
 true false true false false false true true false false true false true
 
+query ????
+SELECT  array_intersect(make_array(1,2,3), make_array(2,3,4)),
+        array_intersect(make_array(1,3,5), make_array(2,4,6)),
+        array_intersect(make_array('aa','bb','cc'), make_array('cc','aa','dd')), 
+        array_intersect(make_array(true, false), make_array(true))
+;
+----
+[2, 3] [] [cc, aa] [true]
+
 query BBBB
 select list_has_all(make_array(1,2,3), make_array(4,5,6)),
        list_has_all(make_array(1,2,3), make_array(1,2)),
diff --git a/datafusion/expr/src/built_in_function.rs b/datafusion/expr/src/built_in_function.rs
@@ -153,6 +153,8 @@ pub enum BuiltinScalarFunction {
     ArrayReplaceAll,
     /// array_to_string
     ArrayToString,
+    /// array_intersect
+    ArrayIntersect,
     /// cardinality
     Cardinality,
     /// construct an array from columns
@@ -359,6 +361,7 @@ impl BuiltinScalarFunction {
             BuiltinScalarFunction::ArrayReplaceN => Volatility::Immutable,
             BuiltinScalarFunction::ArrayReplaceAll => Volatility::Immutable,
             BuiltinScalarFunction::ArrayToString => Volatility::Immutable,
+            BuiltinScalarFunction::ArrayIntersect => Volatility::Immutable,
             BuiltinScalarFunction::Cardinality => Volatility::Immutable,
             BuiltinScalarFunction::MakeArray => Volatility::Immutable,
             BuiltinScalarFunction::TrimArray => Volatility::Immutable,
@@ -543,6 +546,34 @@ impl BuiltinScalarFunction {
             BuiltinScalarFunction::ArrayReplaceN => Ok(input_expr_types[0].clone()),
             BuiltinScalarFunction::ArrayReplaceAll => Ok(input_expr_types[0].clone()),
             BuiltinScalarFunction::ArrayToString => Ok(Utf8),
+            BuiltinScalarFunction::ArrayIntersect => {
+                if input_expr_types.len() < 2 || input_expr_types.len() > 2 {
+                    Err(DataFusionError::Internal(format!(
+                        "The {self} function must have two arrays as parameters"
+                    )))
+                } else {
+                    match (&input_expr_types[0], &input_expr_types[1]) {
+                        (List(l_field), List(r_field)) => {
+                            if !l_field.data_type().equals_datatype(r_field.data_type()) {
+                                Err(DataFusionError::Internal(format!(
+                                    "The {self} function array data type not equal, [0]: {:?}, [1]: {:?}",
+                                    l_field.data_type(), r_field.data_type()
+                                )))
+                            } else {
+                                Ok(List(Arc::new(Field::new(
+                                    "item",
+                                    l_field.data_type().clone(),
+                                    true,
+                                ))))
+                            }
+                        }
+                        _ => Err(DataFusionError::Internal(format!(
+                            "The {} parameters should be array, [0]: {:?}, [1]: {:?}",
+                            self, input_expr_types[0], input_expr_types[1]
+                        ))),
+                    }
+                }
+            }
             BuiltinScalarFunction::Cardinality => Ok(UInt64),
             BuiltinScalarFunction::MakeArray => match input_expr_types.len() {
                 0 => Ok(List(Arc::new(Field::new("item", Null, true)))),
@@ -834,6 +865,7 @@ impl BuiltinScalarFunction {
             BuiltinScalarFunction::ArrayToString => {
                 Signature::variadic_any(self.volatility())
             }
+            BuiltinScalarFunction::ArrayIntersect => Signature::any(2, self.volatility()),
             BuiltinScalarFunction::Cardinality => Signature::any(1, self.volatility()),
             BuiltinScalarFunction::MakeArray => {
                 Signature::variadic_any(self.volatility())
@@ -1324,6 +1356,7 @@ fn aliases(func: &BuiltinScalarFunction) -> &'static [&'static str] {
         BuiltinScalarFunction::Cardinality => &["cardinality"],
         BuiltinScalarFunction::MakeArray => &["make_array", "make_list"],
         BuiltinScalarFunction::TrimArray => &["trim_array"],
+        BuiltinScalarFunction::ArrayIntersect => &["array_intersect", "list_interact"],
     }
 }
 
diff --git a/datafusion/expr/src/expr_fn.rs b/datafusion/expr/src/expr_fn.rs
@@ -654,6 +654,12 @@ scalar_expr!(
     array n,
     "removes the last n elements from the array."
 );
+scalar_expr!(
+    ArrayIntersect,
+    array_intersect,
+    first_array second_array,
+    "Returns an array of the elements in the intersection of array1 and array2."
+);
 
 // string functions
 scalar_expr!(Ascii, ascii, chr, "ASCII code value of the character");
diff --git a/datafusion/physical-expr/src/array_expressions.rs b/datafusion/physical-expr/src/array_expressions.rs
@@ -27,6 +27,7 @@ use datafusion_common::cast::{as_generic_string_array, as_int64_array, as_list_a
 use datafusion_common::ScalarValue;
 use datafusion_common::{DataFusionError, Result};
 use datafusion_expr::ColumnarValue;
+use hashbrown::{HashMap, HashSet};
 use itertools::Itertools;
 use std::sync::Arc;
 
@@ -1820,6 +1821,90 @@ pub fn array_has_all(args: &[ArrayRef]) -> Result<ArrayRef> {
     Ok(Arc::new(boolean_builder.finish()))
 }
 
+macro_rules! array_intersect_normal {
+    ($FIRST_ARRAY:expr, $SECOND_ARRAY:expr, $DATA_TYPE:expr, $ARRAY_TYPE:ident, $BUILDER:ident) => {{
+        let mut offsets: Vec<i32> = vec![0];
+        let mut values =
+            downcast_arg!(new_empty_array(&$DATA_TYPE), $ARRAY_TYPE).clone();
+
+        for (first_arr, second_arr) in $FIRST_ARRAY.iter().zip($SECOND_ARRAY.iter()) {
+            let last_offset: i32 = offsets.last().copied().ok_or_else(|| {
+                DataFusionError::Internal(format!("offsets should not be empty"))
+            })?;
+            match (first_arr, second_arr) {
+                (Some(first_arr), Some(second_arr)) => {
+                    let first_arr = downcast_arg!(first_arr, $ARRAY_TYPE);
+                    // TODO(veeupup): maybe use stack-implemented map to avoid heap memory allocation
+                    let first_set = first_arr.iter().dedup().flatten().collect::<HashSet<_>>();
+                    let second_arr = downcast_arg!(second_arr, $ARRAY_TYPE);
+
+                    let mut builder = $BUILDER::new();
+                    for elem in second_arr.iter().dedup().flatten() {
+                        if first_set.contains(&elem) {
+                            builder.append_value(elem);
+                        }
+                    }
+
+                    let arr = builder.finish();
+                    values = downcast_arg!(
+                        compute::concat(&[
+                            &values,
+                            &arr
+                        ])?
+                        .clone(),
+                        $ARRAY_TYPE
+                    )
+                    .clone();
+                    offsets.push(last_offset + arr.len() as i32);
+                },
+                _ => {
+                    offsets.push(last_offset);
+                }
+            }
+        }
+        let field = Arc::new(Field::new("item", $DATA_TYPE, true));
+
+        Ok(Arc::new(ListArray::try_new(
+            field,
+            OffsetBuffer::new(offsets.into()),
+            Arc::new(values),
+            None,
+        )?))
+
+    }};
+}
+
+/// array_intersect SQL function
+pub fn array_intersect(args: &[ArrayRef]) -> Result<ArrayRef> {
+    assert_eq!(args.len(), 2);
+
+    let first_array = as_list_array(&args[0])?;
+    let second_array = as_list_array(&args[1])?;
+
+    match (first_array.value_type(), second_array.value_type()) {
+        // (DataType::List(_), DataType::List(_)) => concat_internal(args)?,
+        (DataType::Utf8, DataType::Utf8) => array_intersect_normal!(first_array, second_array, DataType::Utf8, StringArray, StringBuilder),
+        (DataType::LargeUtf8, DataType::LargeUtf8) => array_intersect_normal!(first_array, second_array, DataType::LargeUtf8, LargeStringArray, LargeStringBuilder),
+        (DataType::Boolean, DataType::Boolean) => array_intersect_normal!(first_array, second_array, DataType::Boolean, BooleanArray, BooleanBuilder),
+        // (DataType::Float32, DataType::Float32) => array_intersect_normal!(arr, element, Float32Array),
+        // (DataType::Float64, DataType::Float64) => array_intersect_normal!(arr, element, Float64Array),
+        (DataType::Int8, DataType::Int8) => array_intersect_normal!(first_array, second_array, DataType::Int8, Int8Array, Int8Builder),
+        (DataType::Int16, DataType::Int16) => array_intersect_normal!(first_array, second_array, DataType::Int16, Int16Array, Int16Builder),
+        (DataType::Int32, DataType::Int32) => array_intersect_normal!(first_array, second_array, DataType::Int32, Int32Array, Int32Builder),
+        (DataType::Int64, DataType::Int64) => array_intersect_normal!(first_array, second_array, DataType::Int64, Int64Array, Int64Builder),
+        (DataType::UInt8, DataType::UInt8) => array_intersect_normal!(first_array, second_array, DataType::UInt8, UInt8Array, UInt8Builder),
+        (DataType::UInt16, DataType::UInt16) => array_intersect_normal!(first_array, second_array, DataType::UInt16, UInt16Array, UInt16Builder),
+        (DataType::UInt32, DataType::UInt32) => array_intersect_normal!(first_array, second_array, DataType::UInt32, UInt32Array, UInt32Builder),
+        (DataType::UInt64, DataType::UInt64) => array_intersect_normal!(first_array, second_array, DataType::UInt64, UInt64Array, UInt64Builder),
+        // (DataType::Null, _) => return Ok(array(&[ColumnarValue::Array(args[1].clone())])?.into_array(1)),
+        (first_value_dt, second_value_dt) => {
+            Err(DataFusionError::NotImplemented(format!(
+                "array_intersect is not implemented for '{first_value_dt:?}' and '{second_value_dt:?}'",
+            )))
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
diff --git a/datafusion/physical-expr/src/functions.rs b/datafusion/physical-expr/src/functions.rs
@@ -465,6 +465,9 @@ pub fn create_physical_fun(
         BuiltinScalarFunction::ArrayToString => Arc::new(|args| {
             make_scalar_function(array_expressions::array_to_string)(args)
         }),
+        BuiltinScalarFunction::ArrayIntersect => Arc::new(|args| {
+            make_scalar_function(array_expressions::array_intersect)(args)
+        }),
         BuiltinScalarFunction::Cardinality => {
             Arc::new(|args| make_scalar_function(array_expressions::cardinality)(args))
         }
diff --git a/datafusion/proto/proto/datafusion.proto b/datafusion/proto/proto/datafusion.proto
@@ -577,6 +577,7 @@ enum ScalarFunction {
   ArrayReplaceN = 108;
   ArrayRemoveAll = 109;
   ArrayReplaceAll = 110;
+  ArrayIntersect = 111;
 }
 
 message ScalarFunctionNode {
diff --git a/datafusion/proto/src/generated/prost.rs b/datafusion/proto/src/generated/prost.rs
diff --git a/datafusion/proto/src/logical_plan/from_proto.rs b/datafusion/proto/src/logical_plan/from_proto.rs
@@ -469,6 +469,7 @@ impl From<&protobuf::ScalarFunction> for BuiltinScalarFunction {
             ScalarFunction::ArrayReplaceN => Self::ArrayReplaceN,
             ScalarFunction::ArrayReplaceAll => Self::ArrayReplaceAll,
             ScalarFunction::ArrayToString => Self::ArrayToString,
+            ScalarFunction::ArrayIntersect => Self::ArrayIntersect,
             ScalarFunction::Cardinality => Self::Cardinality,
             ScalarFunction::Array => Self::MakeArray,
             ScalarFunction::TrimArray => Self::TrimArray,
diff --git a/datafusion/proto/src/logical_plan/to_proto.rs b/datafusion/proto/src/logical_plan/to_proto.rs
@@ -1417,6 +1417,7 @@ impl TryFrom<&BuiltinScalarFunction> for protobuf::ScalarFunction {
             BuiltinScalarFunction::ArrayReplaceN => Self::ArrayReplaceN,
             BuiltinScalarFunction::ArrayReplaceAll => Self::ArrayReplaceAll,
             BuiltinScalarFunction::ArrayToString => Self::ArrayToString,
+            BuiltinScalarFunction::ArrayIntersect => Self::ArrayIntersect,
             BuiltinScalarFunction::Cardinality => Self::Cardinality,
             BuiltinScalarFunction::MakeArray => Self::Array,
             BuiltinScalarFunction::TrimArray => Self::TrimArray,
diff --git a/docs/source/user-guide/expressions.md b/docs/source/user-guide/expressions.md
@@ -200,6 +200,7 @@ Unlike to some databases the math functions in Datafusion works the same way as
 | array_replace_n(array, from, to, max) | Replaces the first `max` occurrences of the specified element with another specified element. `array_replace_n([1, 2, 2, 3, 2, 1, 4], 2, 5, 2) -> [1, 5, 5, 3, 2, 1, 4]` |
 | array_replace_all(array, from, to)    | Replaces all occurrences of the specified element with another specified element. `array_replace_all([1, 2, 2, 3, 2, 1, 4], 2, 5) -> [1, 5, 5, 3, 5, 1, 4]`              |
 | array_to_string(array, delimeter)     | Converts each element to its text representation. `array_to_string([1, 2, 3, 4], ',') -> 1,2,3,4`                                                                        |
+| array_intersect(array1, array2)        | Returns an array of the elements in the intersection of array1 and array2. `array_intersect([1, 2, 3, 4], [5, 6, 3, 4]) -> [3, 4]`                                        |
 | cardinality(array)                    | Returns the total number of elements in the array. `cardinality([[1, 2, 3], [4, 5, 6]]) -> 6`                                                                            |
 | make_array(value1, [value2 [, ...]])  | Returns an Arrow array using the specified input expressions. `make_array(1, 2, 3) -> [1, 2, 3]`                                                                         |
 | trim_array(array, n)                  | Removes the last n elements from the array.                                                                                                                              |

Original file line number	Diff line number	Diff line change
`@@ -577,6 +577,7 @@ enum ScalarFunction {`
`577`	`577`	`ArrayReplaceN = 108;`
`578`	`578`	`ArrayRemoveAll = 109;`
`579`	`579`	`ArrayReplaceAll = 110;`
	`580`	`+ ArrayIntersect = 111;`
`580`	`581`	`}`
`581`	`582`
`582`	`583`	`message ScalarFunctionNode {`