Add set based operations for arrays: array_intersect, array_union, array_except, and arrays_overlap for running on GPU

integration_tests/src/main/python/array_test.py

@@ -394,3 +394,80 @@ def test_array_max_q1():
     def q1(spark):
         return spark.sql('SELECT ARRAY_MAX(TRANSFORM(ARRAY_REPEAT(STRUCT(1, 2), 0), s -> s.col2))')
     assert_gpu_and_cpu_are_equal_collect(q1)
+
+
+@pytest.mark.parametrize('data_gen', [byte_gen, short_gen, int_gen, long_gen,
+    FloatGen(special_cases=[]), DoubleGen(special_cases=[]), string_gen, boolean_gen, date_gen, timestamp_gen], ids=idfn)
+def test_array_intersect(data_gen):
+    gen = StructGen(
+        [('a', ArrayGen(data_gen, nullable=False)),
+        ('b', ArrayGen(data_gen, nullable=False))],
+        nullable=False)
+
+    assert_gpu_and_cpu_are_equal_collect(
+        lambda spark: gen_df(spark, gen).selectExpr(
+            'sort_array(array_intersect(a, b))',
+            'sort_array(array_intersect(b, a))',
+            'sort_array(array_intersect(a, array()))',
+            'sort_array(array_intersect(array(), b))',
+            'sort_array(array_intersect(a, a))',
+        )
+    )
+
+@pytest.mark.parametrize('data_gen', [byte_gen, short_gen, int_gen, long_gen,
+    FloatGen(special_cases=[]), DoubleGen(special_cases=[]), string_gen, boolean_gen, date_gen, timestamp_gen], ids=idfn)
+def test_array_union(data_gen):
+    gen = StructGen(
+        [('a', ArrayGen(data_gen, nullable=False)),
+        ('b', ArrayGen(data_gen, nullable=False))],
+        nullable=False)
+
+    # The 4th item in this integration test here is left commmented out here
+    # There is an issue with running that item with collect(), which affects
+    # the integration test here.
+    # See this issue (https://github.com/NVIDIA/spark-rapids/issues/5957) 
+    assert_gpu_and_cpu_are_equal_collect(
+        lambda spark: gen_df(spark, gen).selectExpr(
+            'sort_array(array_union(a, b))',
+            'sort_array(array_union(b, a))',
+            'sort_array(array_union(a, array()))',
+            # 'sort_array(array_union(array(), b))', # see https://github.com/NVIDIA/spark-rapids/issues/5957
+            'sort_array(array_union(a, a))',
+        )
+    )
+
+@pytest.mark.parametrize('data_gen', [byte_gen, short_gen, int_gen, long_gen,
+    FloatGen(special_cases=[]), DoubleGen(special_cases=[]), string_gen, boolean_gen, date_gen, timestamp_gen], ids=idfn)
+def test_array_except(data_gen):
+    gen = StructGen(
+        [('a', ArrayGen(data_gen, nullable=False)),
+        ('b', ArrayGen(data_gen, nullable=False))],
+        nullable=False)
+
+    assert_gpu_and_cpu_are_equal_collect(
+        lambda spark: gen_df(spark, gen).selectExpr(
+            'sort_array(array_except(a, b))',
+            'sort_array(array_except(b, a))',
+            'sort_array(array_except(a, array()))',
+            'sort_array(array_except(array(), b))',
+            'sort_array(array_except(a, a))',
+        )
+    )
+
+@pytest.mark.parametrize('data_gen', [byte_gen, short_gen, int_gen, long_gen,
+    FloatGen(special_cases=[]), DoubleGen(special_cases=[]), string_gen, boolean_gen, date_gen, timestamp_gen], ids=idfn)
+def test_arrays_overlap(data_gen):
+    gen = StructGen(
+        [('a', ArrayGen(data_gen, nullable=False)),
+        ('b', ArrayGen(data_gen, nullable=False))],
+        nullable=False)
+
+    assert_gpu_and_cpu_are_equal_collect(
+        lambda spark: gen_df(spark, gen).selectExpr(
+            'arrays_overlap(a, b)',
+            'arrays_overlap(b, a)',
+            'arrays_overlap(a, array())',
+            'arrays_overlap(array(), b)',
+            'arrays_overlap(a, a)',
+        )
+    )

sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuOverrides.scala

@@ -2907,7 +2907,75 @@ object GpuOverrides extends Logging {
         }
       }
     ),
-
+    expr[ArrayExcept](
+      "Returns an array of the elements in array1 but not in array2, without duplicates",
+      ExprChecks.binaryProject(
+        TypeSig.ARRAY.nested(TypeSig.commonCudfTypes + TypeSig.DECIMAL_128 + TypeSig.NULL),
+        TypeSig.ARRAY.nested(TypeSig.all),
+        ("array1",
+            TypeSig.ARRAY.nested(TypeSig.commonCudfTypes + TypeSig.DECIMAL_128 + TypeSig.NULL),
+            TypeSig.ARRAY.nested(TypeSig.all)),
+        ("array2",
+            TypeSig.ARRAY.nested(TypeSig.commonCudfTypes + TypeSig.DECIMAL_128 + TypeSig.NULL),
+            TypeSig.ARRAY.nested(TypeSig.all))),
+      (in, conf, p, r) => new BinaryExprMeta[ArrayExcept](in, conf, p, r) {
+        override def convertToGpu(lhs: Expression, rhs: Expression): GpuExpression = {
+          GpuArrayExcept(lhs, rhs)
+        }
+      }
+    ),
+    expr[ArrayIntersect](
+      "Returns an array of the elements in the intersection of array1 and array2, without" +
+        " duplicates",
+      ExprChecks.binaryProject(
+        TypeSig.ARRAY.nested(TypeSig.commonCudfTypes + TypeSig.DECIMAL_128 + TypeSig.NULL),
+        TypeSig.ARRAY.nested(TypeSig.all),
+        ("array1",
+            TypeSig.ARRAY.nested(TypeSig.commonCudfTypes + TypeSig.DECIMAL_128 + TypeSig.NULL),
+            TypeSig.ARRAY.nested(TypeSig.all)),
+        ("array2",
+            TypeSig.ARRAY.nested(TypeSig.commonCudfTypes + TypeSig.DECIMAL_128 + TypeSig.NULL),
+            TypeSig.ARRAY.nested(TypeSig.all))),
+      (in, conf, p, r) => new BinaryExprMeta[ArrayIntersect](in, conf, p, r) {
+        override def convertToGpu(lhs: Expression, rhs: Expression): GpuExpression = {
+          GpuArrayIntersect(lhs, rhs)
+        }
+      }
+    ),
+    expr[ArrayUnion](
+      "Returns an array of the elements in the union of array1 and array2, without duplicates.",
+      ExprChecks.binaryProject(
+        TypeSig.ARRAY.nested(TypeSig.commonCudfTypes + TypeSig.DECIMAL_128 + TypeSig.NULL),
+        TypeSig.ARRAY.nested(TypeSig.all),
+        ("array1",
+            TypeSig.ARRAY.nested(TypeSig.commonCudfTypes + TypeSig.DECIMAL_128 + TypeSig.NULL),
+            TypeSig.ARRAY.nested(TypeSig.all)),
+        ("array2",
+            TypeSig.ARRAY.nested(TypeSig.commonCudfTypes + TypeSig.DECIMAL_128 + TypeSig.NULL),
+            TypeSig.ARRAY.nested(TypeSig.all))),
+      (in, conf, p, r) => new BinaryExprMeta[ArrayUnion](in, conf, p, r) {
+        override def convertToGpu(lhs: Expression, rhs: Expression): GpuExpression = {
+          GpuArrayUnion(lhs, rhs)
+        }
+      }
+    ),
+    expr[ArraysOverlap](
+      "Returns true if a1 contains at least a non-null element present also in a2. If the arrays " +
+      "have no common element and they are both non-empty and either of them contains a null " + 
+      "element null is returned, false otherwise.",
+      ExprChecks.binaryProject(TypeSig.BOOLEAN, TypeSig.BOOLEAN,
+        ("array1",
+            TypeSig.ARRAY.nested(TypeSig.commonCudfTypes + TypeSig.DECIMAL_128 + TypeSig.NULL),
+            TypeSig.ARRAY.nested(TypeSig.all)),
+        ("array2",
+            TypeSig.ARRAY.nested(TypeSig.commonCudfTypes + TypeSig.DECIMAL_128 + TypeSig.NULL),
+            TypeSig.ARRAY.nested(TypeSig.all))),
+      (in, conf, p, r) => new BinaryExprMeta[ArraysOverlap](in, conf, p, r) {
+        override def convertToGpu(lhs: Expression, rhs: Expression): GpuExpression = {
+          GpuArraysOverlap(lhs, rhs)
+        }
+      }
+    ),
     expr[TransformKeys](
       "Transform keys in a map using a transform function",
       ExprChecks.projectOnly(TypeSig.MAP.nested(TypeSig.commonCudfTypes + TypeSig.DECIMAL_128 +

sql-plugin/src/main/scala/org/apache/spark/sql/rapids/collectionOperations.scala

@@ -679,6 +679,171 @@ case class GpuArraysZip(children: Seq[Expression]) extends GpuExpression with Sh
   }
 }
 
+trait GpuArraySetOperation extends GpuBinaryExpression {
+
+  override def doColumnar(lhs: GpuScalar, rhs: GpuColumnVector): ColumnVector = {
+    withResource(GpuColumnVector.from(lhs, rhs.getRowCount.toInt, lhs.dataType)) { left =>
+      doColumnar(left, rhs)
+    }
+  }
+
+  override def doColumnar(lhs: GpuColumnVector, rhs: GpuScalar): ColumnVector = {
+    withResource(GpuColumnVector.from(rhs, lhs.getRowCount.toInt, rhs.dataType)) { right =>
+      doColumnar(lhs, right)
+    }
+  }
+
+  override def doColumnar(numRows: Int, lhs: GpuScalar, rhs: GpuScalar): ColumnVector = {
+    withResource(GpuColumnVector.from(lhs, numRows, lhs.dataType)) { left =>
+      withResource(GpuColumnVector.from(rhs, numRows, rhs.dataType)) { right =>
+        doColumnar(left, right)
+      }
+    }
+  }
+}
+
+case class GpuArrayExcept(left: Expression, right: Expression)
+    extends GpuArraySetOperation with ExpectsInputTypes {
+
+  override def inputTypes: Seq[AbstractDataType] = Seq(ArrayType, ArrayType)
+
+  override def checkInputDataTypes(): TypeCheckResult =
+    (left.dataType, right.dataType) match {
+    case (ArrayType(ldt, _), ArrayType(rdt, _)) =>
+      if (ldt.sameType(rdt)) {
+        TypeCheckResult.TypeCheckSuccess
+      } else {
+        TypeCheckResult.TypeCheckFailure(
+          s"Array_intersect requires both array params to have the same subType: $ldt != $rdt")
+      }
+    case dt =>
+      TypeCheckResult.TypeCheckFailure(s"$prettyName only supports array input, but found $dt")
+  }
+
+  override def dataType: DataType = left.dataType
+
+  override def nullable: Boolean = false
+
+  override def doColumnar(lhs: GpuColumnVector, rhs: GpuColumnVector): ColumnVector = {
+    ColumnView.setDifference(lhs.getBase, rhs.getBase)
+  }
+}
+
+case class GpuArrayIntersect(left: Expression, right: Expression)
+    extends GpuArraySetOperation with ExpectsInputTypes {
+
+  override def inputTypes: Seq[AbstractDataType] = Seq(ArrayType, ArrayType)
+
+  override def checkInputDataTypes(): TypeCheckResult =
+    (left.dataType, right.dataType) match {
+    case (ArrayType(ldt, _), ArrayType(rdt, _)) =>
+      if (ldt.sameType(rdt)) {
+        TypeCheckResult.TypeCheckSuccess
+      } else {
+        TypeCheckResult.TypeCheckFailure(
+          s"Array_intersect requires both array params to have the same subType: $ldt != $rdt")
+      }
+    case dt =>
+      TypeCheckResult.TypeCheckFailure(s"$prettyName only supports array input, but found $dt")
+  }
+
+  override def dataType: DataType = left.dataType
+
+  override def nullable: Boolean = false
+
+  override def doColumnar(lhs: GpuColumnVector, rhs: GpuColumnVector): ColumnVector = {
+    ColumnView.setIntersect(lhs.getBase, rhs.getBase)
+  }
+}
+
+case class GpuArrayUnion(left: Expression, right: Expression)
+    extends GpuArraySetOperation with ExpectsInputTypes {
+
+  override def inputTypes: Seq[AbstractDataType] = Seq(ArrayType, ArrayType)
+
+  override def checkInputDataTypes(): TypeCheckResult =
+    (left.dataType, right.dataType) match {
+    case (ArrayType(ldt, _), ArrayType(rdt, _)) =>
+      if (ldt.sameType(rdt)) {
+        TypeCheckResult.TypeCheckSuccess
+      } else {
+        TypeCheckResult.TypeCheckFailure(
+          s"Array_union requires both array params to have the same subType: $ldt != $rdt")
+      }
+    case dt =>
+      TypeCheckResult.TypeCheckFailure(s"$prettyName only supports array input, but found $dt")
+  }
+
+  override def dataType: DataType = left.dataType
+
+  override def nullable: Boolean = false
+
+  override def doColumnar(lhs: GpuColumnVector, rhs: GpuColumnVector): ColumnVector = {
+    ColumnView.setUnion(lhs.getBase, rhs.getBase)
+  }
+}
+
+case class GpuArraysOverlap(left: Expression, right: Expression)
+    extends GpuArraySetOperation with ExpectsInputTypes {
+
+  override def inputTypes: Seq[AbstractDataType] = Seq(ArrayType, ArrayType)
+
+  override def checkInputDataTypes(): TypeCheckResult =
+    (left.dataType, right.dataType) match {
+    case (ArrayType(ldt, _), ArrayType(rdt, _)) =>
+      if (ldt.sameType(rdt)) {
+        TypeCheckResult.TypeCheckSuccess
+      } else {
+        TypeCheckResult.TypeCheckFailure(
+          s"Array_union requires both array params to have the same subType: $ldt != $rdt")
+      }
+    case dt =>
+      TypeCheckResult.TypeCheckFailure(s"$prettyName only supports array input, but found $dt")
+  }
+
+  override def dataType: DataType = BooleanType
+
+  override def nullable: Boolean = true
+
+  override def doColumnar(lhs: GpuColumnVector, rhs: GpuColumnVector): ColumnVector = {
+    // Spark's arrays_overlap() function will return null when both arrays are non-empty,
+    // they have no-non null overlapping elements and either array contains at least one null
+    // value, so we have to handle that here.
+    val bothNonEmpty = withResource(lhs.getBase.countElements) { leftCount => 
+      withResource(rhs.getBase.countElements) { rightCount => 
+        withResource(Scalar.fromInt(0)) { zero =>
+          withResource(leftCount.greaterThan(zero)) { leftNonEmpty =>
+            withResource(rightCount.greaterThan(zero)) { rightNonEmpty => 
+              leftNonEmpty.and(rightNonEmpty)
+            }
+          }
+        }
+      }
+    }
+    val eitherHasNulls = withResource(lhs.getBase.listContainsNulls) { leftHasNulls =>
+      withResource(rhs.getBase.listContainsNulls) { rightHasNulls => 
+        leftHasNulls.or(rightHasNulls)
+      }
+    }
+    withResource(bothNonEmpty) { _ =>
+      withResource(eitherHasNulls) { _ => 
+        withResource(ColumnView.listOverlap(lhs.getBase, rhs.getBase)) { overlaps => 
+          withResource(overlaps.not) { notOverlap =>
+            withResource(bothNonEmpty.and(eitherHasNulls)) { nonEmptyAndHasNull =>
+              withResource(nonEmptyAndHasNull.and(notOverlap)) { returnNull =>
+                returnNull.ifElse(
+                  GpuColumnVector.columnVectorFromNull(lhs.getRowCount.toInt, BooleanType),
+                  overlaps
+                )
+              }
+            }
+          }
+        }
+      }
+    }
+  }
+}
+
 class GpuSequenceMeta(
     expr: Sequence,
     conf: RapidsConf,