[SPARK-23922][SQL] Add arrays_overlap function

python/pyspark/sql/functions.py

@@ -1834,6 +1834,20 @@ def array_contains(col, value):
     return Column(sc._jvm.functions.array_contains(_to_java_column(col), value))
 
 
+@since(2.4)
+def arrays_overlap(a1, a2):
+    """
+    Collection function: returns true if the arrays contain any common non-null element; if not,
+    returns null if any of the arrays contains a null element and false otherwise.
+
+    >>> df = spark.createDataFrame([(["a", "b"], ["b", "c"]), (["a"], ["b", "c"])], ['x', 'y'])
+    >>> df.select(arrays_overlap(df.x, df.y).alias("overlap")).collect()
+    [Row(overlap=True), Row(overlap=False)]
+    """
+    sc = SparkContext._active_spark_context
+    return Column(sc._jvm.functions.arrays_overlap(_to_java_column(a1), _to_java_column(a2)))
+
+
 @since(2.4)
 def slice(x, start, length):
     """

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala

@@ -401,6 +401,7 @@ object FunctionRegistry {
     // collection functions
     expression[CreateArray]("array"),
     expression[ArrayContains]("array_contains"),
+    expression[ArraysOverlap]("arrays_overlap"),
     expression[ArrayJoin]("array_join"),
     expression[ArrayPosition]("array_position"),
     expression[ArraySort]("array_sort"),

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/collectionOperations.scala

@@ -18,15 +18,45 @@ package org.apache.spark.sql.catalyst.expressions
 
 import java.util.Comparator
 
+import scala.collection.mutable
+
 import org.apache.spark.sql.catalyst.InternalRow
-import org.apache.spark.sql.catalyst.analysis.TypeCheckResult
+import org.apache.spark.sql.catalyst.analysis.{TypeCheckResult, TypeCoercion}
 import org.apache.spark.sql.catalyst.expressions.ArraySortLike.NullOrder
 import org.apache.spark.sql.catalyst.expressions.codegen._
 import org.apache.spark.sql.catalyst.util.{ArrayData, GenericArrayData, MapData, TypeUtils}
 import org.apache.spark.sql.types._
 import org.apache.spark.unsafe.array.ByteArrayMethods
 import org.apache.spark.unsafe.types.{ByteArray, UTF8String}
 
+/**
+ * Base trait for [[BinaryExpression]]s with two arrays of the same element type and implicit
+ * casting.
+ */
+trait BinaryArrayExpressionWithImplicitCast extends BinaryExpression
+  with ImplicitCastInputTypes {
+
+  protected lazy val elementType: DataType = inputTypes.head.asInstanceOf[ArrayType].elementType
+
+  override def inputTypes: Seq[AbstractDataType] = {
+    TypeCoercion.findWiderTypeForTwo(left.dataType, right.dataType) match {
+      case Some(arrayType) => Seq(arrayType, arrayType)
+      case _ => Seq.empty
+    }
+  }
+
+  override def checkInputDataTypes(): TypeCheckResult = {
+    (left.dataType, right.dataType) match {
+      case (ArrayType(e1, _), ArrayType(e2, _)) if e1.sameType(e2) =>
+        TypeCheckResult.TypeCheckSuccess
+      case _ => TypeCheckResult.TypeCheckFailure(s"input to function $prettyName should have " +
+        s"been two ${ArrayType.simpleString}s with same element type, but it's " +
+        s"[${left.dataType.simpleString}, ${right.dataType.simpleString}]")
+    }
+  }
+}
+
+
 /**
  * Given an array or map, returns its size. Returns -1 if null.
  */
@@ -529,6 +559,157 @@ case class ArrayContains(left: Expression, right: Expression)
   override def prettyName: String = "array_contains"
 }
 
+/**
+ * Checks if the two arrays contain at least one common element.
+ */
+// scalastyle:off line.size.limit
+@ExpressionDescription(
+  usage = "_FUNC_(a1, a2) - Returns true if a1 contains at least an element present also in a2. If the arrays have no common element and either of them contains a null element null is returned, false otherwise.",
+  examples = """
+    Examples:
+      > SELECT _FUNC_(array(1, 2, 3), array(3, 4, 5));
+       true
+  """, since = "2.4.0")
+// scalastyle:off line.size.limit
+case class ArraysOverlap(left: Expression, right: Expression)
+  extends BinaryArrayExpressionWithImplicitCast {
+
+  override def dataType: DataType = BooleanType
+
+  override def nullable: Boolean = {
+    left.nullable || right.nullable || left.dataType.asInstanceOf[ArrayType].containsNull ||
+      right.dataType.asInstanceOf[ArrayType].containsNull
+  }
+
+  override def nullSafeEval(a1: Any, a2: Any): Any = {
+    var hasNull = false
+    val arr1 = a1.asInstanceOf[ArrayData]
+    val arr2 = a2.asInstanceOf[ArrayData]
+    val (bigger, smaller, biggerDt) = if (arr1.numElements() > arr2.numElements()) {
+      (arr1, arr2, left.dataType.asInstanceOf[ArrayType])
+    } else {
+      (arr2, arr1, right.dataType.asInstanceOf[ArrayType])
+    }
+    if (smaller.numElements() > 0) {
+      val smallestSet = new mutable.HashSet[Any]
+      smaller.foreach(elementType, (_, v) =>
+        if (v == null) {
+          hasNull = true
+        } else {
+          smallestSet += v
+        })
+      bigger.foreach(elementType, (_, v1) =>
+        if (v1 == null) {
+          hasNull = true
+        } else if (smallestSet.contains(v1)) {
+          return true
+        }
+      )
+    } else if (containsNull(bigger, biggerDt)) {
+      hasNull = true
+    }
+    if (hasNull) {
+      null
+    } else {
+      false
+    }
+  }
+
+  def containsNull(arr: ArrayData, dt: ArrayType): Boolean = {
+    if (dt.containsNull) {
+      var i = 0
+      var hasNull = false
+      while (i < arr.numElements && !hasNull) {
+        hasNull = arr.isNullAt(i)
+        i += 1
+      }
+      hasNull
+    } else {
+      false
+    }
+  }
+
+  override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
+    nullSafeCodeGen(ctx, ev, (a1, a2) => {
+      val i = ctx.freshName("i")
+      val smaller = ctx.freshName("smallerArray")
+      val bigger = ctx.freshName("biggerArray")
+      val getFromSmaller = CodeGenerator.getValue(smaller, elementType, i)
+      val getFromBigger = CodeGenerator.getValue(bigger, elementType, i)
+      val smallerEmptyCode = if (inputTypes.exists(_.asInstanceOf[ArrayType].containsNull)) {
+        s"""
+           |else {
+           |  for (int $i = 0; $i < $bigger.numElements(); $i ++) {
+           |    if ($bigger.isNullAt($i)) {
+           |      ${ev.isNull} = true;
+           |      break;
+           |    }
+           |  }
+           |}
+         """.stripMargin
+      } else {
+        ""
+      }
+      val javaElementClass = CodeGenerator.boxedType(elementType)
+      val javaSet = classOf[java.util.HashSet[_]].getName
+      val set2 = ctx.freshName("set")
+      val addToSetFromSmallerCode = nullSafeElementCodegen(
+        smaller, i, s"$set2.add($getFromSmaller);", s"${ev.isNull} = true;")
+      val elementIsInSetCode = nullSafeElementCodegen(
+        bigger,
+        i,
+        s"""
+           |if ($set2.contains($getFromBigger)) {
+           |  ${ev.isNull} = false;
+           |  ${ev.value} = true;
+           |  break;
+           |}
+           |""".stripMargin,
+        s"${ev.isNull} = true;")
+      s"""
+         |ArrayData $smaller;
+         |ArrayData $bigger;
+         |if ($a1.numElements() > $a2.numElements()) {
+         |  $bigger = $a1;
+         |  $smaller = $a2;
+         |} else {
+         |  $smaller = $a1;
+         |  $bigger = $a2;
+         |}
+         |if ($smaller.numElements() > 0) {
+         |  $javaSet<$javaElementClass> $set2 = new $javaSet<$javaElementClass>();
+         |  for (int $i = 0; $i < $smaller.numElements(); $i ++) {
+         |    $addToSetFromSmallerCode
+         |  }
+         |  for (int $i = 0; $i < $bigger.numElements(); $i ++) {
+         |    $elementIsInSetCode
+         |  }
+         |} $smallerEmptyCode
+         |""".stripMargin
+    })
+  }
+
+  def nullSafeElementCodegen(
+      arrayVar: String,
+      index: String,
+      code: String,
+      isNullCode: String): String = {
+    if (inputTypes.exists(_.asInstanceOf[ArrayType].containsNull)) {
+      s"""
+         |if ($arrayVar.isNullAt($index)) {
+         |  $isNullCode
+         |} else {
+         |  $code
+         |}
+         |""".stripMargin
+    } else {
+      code
+    }
+  }
+
+  override def prettyName: String = "arrays_overlap"
+}
+
 /**
  * Slices an array according to the requested start index and length
  */

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/CollectionExpressionsSuite.scala

@@ -136,6 +136,37 @@ class CollectionExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper
     checkEvaluation(ArrayContains(a3, Literal.create(null, StringType)), null)
   }
 
+  test("ArraysOverlap") {
+    val a0 = Literal.create(Seq(1, 2, 3), ArrayType(IntegerType))
+    val a1 = Literal.create(Seq(4, 5, 3), ArrayType(IntegerType))
+    val a2 = Literal.create(Seq(null, 5, 6), ArrayType(IntegerType))
+    val a3 = Literal.create(Seq(7, 8), ArrayType(IntegerType))
+    val a4 = Literal.create(Seq.empty[Int], ArrayType(IntegerType))
+
+    val a5 = Literal.create(Seq[String](null, ""), ArrayType(StringType))
+    val a6 = Literal.create(Seq[String]("", "abc"), ArrayType(StringType))
+    val a7 = Literal.create(Seq[String]("def", "ghi"), ArrayType(StringType))
+
+    checkEvaluation(ArraysOverlap(a0, a1), true)
+    checkEvaluation(ArraysOverlap(a0, a2), null)
+    checkEvaluation(ArraysOverlap(a1, a2), true)
+    checkEvaluation(ArraysOverlap(a1, a3), false)
+    checkEvaluation(ArraysOverlap(a0, a4), false)
+    checkEvaluation(ArraysOverlap(a2, a4), null)
+    checkEvaluation(ArraysOverlap(a4, a2), null)
+
+    checkEvaluation(ArraysOverlap(a5, a6), true)
+    checkEvaluation(ArraysOverlap(a5, a7), null)
+    checkEvaluation(ArraysOverlap(a6, a7), false)
+
+    // null handling
+    checkEvaluation(ArraysOverlap(Literal.create(null, ArrayType(IntegerType)), a0), null)
+    checkEvaluation(ArraysOverlap(a0, Literal.create(null, ArrayType(IntegerType))), null)
+    checkEvaluation(ArraysOverlap(
+      Literal.create(Seq(null), ArrayType(IntegerType)),
+      Literal.create(Seq(null), ArrayType(IntegerType))), null)
+  }
+
   test("Slice") {
     val a0 = Literal.create(Seq(1, 2, 3, 4, 5, 6), ArrayType(IntegerType))
     val a1 = Literal.create(Seq[String]("a", "b", "c", "d"), ArrayType(StringType))

sql/core/src/main/scala/org/apache/spark/sql/functions.scala

@@ -3039,6 +3039,16 @@ object functions {
     ArrayContains(column.expr, Literal(value))
   }
 
+  /**
+   * Returns `true` if `a1` and `a2` have at least one non-null element in common. If not and
+   * any of the arrays contains a `null`, it returns `null`. It returns `false` otherwise.
+   * @group collection_funcs
+   * @since 2.4.0
+   */
+  def arrays_overlap(a1: Column, a2: Column): Column = withExpr {
+    ArraysOverlap(a1.expr, a2.expr)
+  }
+
   /**
    * Returns an array containing all the elements in `x` from index `start` (or starting from the
    * end if `start` is negative) with the specified `length`.

sql/core/src/test/scala/org/apache/spark/sql/DataFrameFunctionsSuite.scala

@@ -442,6 +442,29 @@ class DataFrameFunctionsSuite extends QueryTest with SharedSQLContext {
     )
   }
 
+  test("arrays_overlap function") {
+    val df = Seq(
+      (Seq[Option[Int]](Some(1), Some(2)), Seq[Option[Int]](Some(-1), Some(10))),
+      (Seq.empty[Option[Int]], Seq[Option[Int]](Some(-1), None)),
+      (Seq[Option[Int]](Some(3), Some(2)), Seq[Option[Int]](Some(1), Some(2)))
+    ).toDF("a", "b")
+
+    val answer = Seq(Row(false), Row(null), Row(true))
+
+    checkAnswer(df.select(arrays_overlap(df("a"), df("b"))), answer)
+    checkAnswer(df.selectExpr("arrays_overlap(a, b)"), answer)
+
+    checkAnswer(sql("select arrays_overlap(array(1, 2, 3), array('a', 'b', 'c'))"), Row(false))
+
+    intercept[AnalysisException] {
+      sql("select arrays_overlap(array(array(1)), array('a'))")
+    }
+
+    intercept[AnalysisException] {
+      sql("select arrays_overlap(null, null)")
+    }
+  }
+
   test("slice function") {
     val df = Seq(
       Seq(1, 2, 3),