Skip to content

[SPARK-23913][SQL] Add array_intersect function #21102

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
kiszk wants to merge 21 commits into from
Closed

[SPARK-23913][SQL] Add array_intersect function #21102

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
21 commits
Select commit Hold shift + click to select a range
bc5a02d
initial commit
kiszk Apr 18, 2018
c7f230d
simplification
kiszk Apr 18, 2018
e96ab2d
fix pyspark test failure
kiszk Apr 19, 2018
49bc319
update
kiszk Jul 10, 2018
4ec43b1
add missing file
kiszk Jul 10, 2018
2c30351
fix test failure
kiszk Jul 12, 2018
74103ce
update
kiszk Jul 12, 2018
387d0ca
fix compilation failure
kiszk Jul 12, 2018
e374aa9
fix test failures
kiszk Jul 13, 2018
e31f7e6
minor refactoring
kiszk Jul 14, 2018
9818ba9
update
kiszk Jul 16, 2018
352743f
address review comment
kiszk Jul 16, 2018
89a828b
avoid to use OpenHashSet.remove method
kiszk Jul 21, 2018
90fecd5
update dataType setting
kiszk Jul 22, 2018
663358e
update
kiszk Jul 27, 2018
3ef99b8
optimize nullchecks in generated code
kiszk Jul 29, 2018
3462b98
update
kiszk Jul 31, 2018
cf742d4
refactor methods that are used in Except and Intersect
kiszk Aug 2, 2018
6fba1ee
update
kiszk Aug 3, 2018
ce755e2
address review comments
kiszk Aug 6, 2018
33781b6
address review comment
kiszk Aug 6, 2018
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
address review comments
  • Loading branch information
@kiszk
kiszk committed Aug 6, 2018
commit ce755e2b049ca000d6da754654b792e181e6d904
144 changes: 74 additions & 70 deletions ...alyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/collectionOperations.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4043,7 +4043,7 @@ object ArrayUnion {
array2, without duplicates.
""",
examples = """
Examples:Fun
Examples:
> SELECT _FUNC_(array(1, 2, 3), array(1, 3, 5));
array(1, 3)
""",
Expand All @@ -4060,81 +4060,89 @@ case class ArrayIntersect(left: Expression, right: Expression) extends ArraySetL
@transient lazy val evalIntersect: (ArrayData, ArrayData) => ArrayData = {
if (elementTypeSupportEquals) {
(array1, array2) =>
val hs = new OpenHashSet[Any]
val hsResult = new OpenHashSet[Any]
var foundNullElement = false
var i = 0
while (i < array2.numElements()) {
if (array2.isNullAt(i)) {
foundNullElement = true
} else {
val elem = array2.get(i, elementType)
hs.add(elem)
}
i += 1
}
val arrayBuffer = new scala.collection.mutable.ArrayBuffer[Any]
i = 0
while (i < array1.numElements()) {
if (array1.isNullAt(i)) {
if (foundNullElement) {
arrayBuffer += null
foundNullElement = false
if (array1.numElements() != 0 && array2.numElements() != 0) {
val hs = new OpenHashSet[Any]
val hsResult = new OpenHashSet[Any]
var foundNullElement = false
var i = 0
while (i < array2.numElements()) {
if (array2.isNullAt(i)) {
foundNullElement = true
} else {
val elem = array2.get(i, elementType)
hs.add(elem)
}
} else {
val elem = array1.get(i, elementType)
if (hs.contains(elem) && !hsResult.contains(elem)) {
arrayBuffer += elem
hsResult.add(elem)
i += 1
}
val arrayBuffer = new scala.collection.mutable.ArrayBuffer[Any]
i = 0
while (i < array1.numElements()) {
if (array1.isNullAt(i)) {
if (foundNullElement) {
arrayBuffer += null
foundNullElement = false
}
} else {
val elem = array1.get(i, elementType)
if (hs.contains(elem) && !hsResult.contains(elem)) {
arrayBuffer += elem
hsResult.add(elem)
}
}
i += 1
}
i += 1
new GenericArrayData(arrayBuffer)
} else {
new GenericArrayData(Seq.empty)
Copy link
Member

@ueshin ueshin Aug 6, 2018

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: Array.empty or Array.emptyObjectArray?

}
new GenericArrayData(arrayBuffer)
} else {
(array1, array2) =>
val arrayBuffer = new scala.collection.mutable.ArrayBuffer[Any]
var alreadySeenNull = false
var i = 0
while (i < array1.numElements()) {
var found = false
val elem1 = array1.get(i, elementType)
if (array1.isNullAt(i)) {
if (!alreadySeenNull) {
if (array1.numElements() != 0 && array2.numElements() != 0) {
val arrayBuffer = new scala.collection.mutable.ArrayBuffer[Any]
var alreadySeenNull = false
var i = 0
while (i < array1.numElements()) {
var found = false
val elem1 = array1.get(i, elementType)
if (array1.isNullAt(i)) {
if (!alreadySeenNull) {
var j = 0
while (!found && j < array2.numElements()) {
found = array2.isNullAt(j)
j += 1
}
// array2 is scanned only once for null element
alreadySeenNull = true
}
} else {
var j = 0
while (!found && j < array2.numElements()) {
found = array2.isNullAt(j)
j += 1
}
// array2 is scanned only once for null element
alreadySeenNull = true
}
} else {
var j = 0
while (!found && j < array2.numElements()) {
if (!array2.isNullAt(j)) {
val elem2 = array2.get(j, elementType)
if (ordering.equiv(elem1, elem2)) {
// check whether elem1 is already stored in arrayBuffer
var foundArrayBuffer = false
var k = 0
while (!foundArrayBuffer && k < arrayBuffer.size) {
val va = arrayBuffer(k)
foundArrayBuffer = (va != null) && ordering.equiv(va, elem1)
k += 1
if (!array2.isNullAt(j)) {
val elem2 = array2.get(j, elementType)
if (ordering.equiv(elem1, elem2)) {
// check whether elem1 is already stored in arrayBuffer
var foundArrayBuffer = false
var k = 0
while (!foundArrayBuffer && k < arrayBuffer.size) {
val va = arrayBuffer(k)
foundArrayBuffer = (va != null) && ordering.equiv(va, elem1)
k += 1
}
found = !foundArrayBuffer
}
found = !foundArrayBuffer
}
j += 1
}
j += 1
}
if (found) {
arrayBuffer += elem1
}
i += 1
}
if (found) {
arrayBuffer += elem1
}
i += 1
new GenericArrayData(arrayBuffer)
} else {
new GenericArrayData(Seq.empty)
}
new GenericArrayData(arrayBuffer)
}
}

Expand Down Expand Up @@ -4162,9 +4170,8 @@ case class ArrayIntersect(left: Expression, right: Expression) extends ArraySetL
val classTag = s"scala.reflect.ClassTag$$.MODULE$$.$hsTypeName()"
val hashSet = ctx.freshName("hashSet")
val hashSetResult = ctx.freshName("hashSetResult")
val arrayBuilder = "scala.collection.mutable.ArrayBuilder"
val arrayBuilder = classOf[mutable.ArrayBuilder[_]].getName
val arrayBuilderClass = s"$arrayBuilder$$of$ptName"
val arrayBuilderClassTag = s"scala.reflect.ClassTag$$.MODULE$$.$ptName()"

def withArray2NullCheck(body: String): String =
if (right.dataType.asInstanceOf[ArrayType].containsNull) {
Expand Down Expand Up @@ -4250,8 +4257,7 @@ case class ArrayIntersect(left: Expression, right: Expression) extends ArraySetL
|for (int $i = 0; $i < $array2.numElements(); $i++) {
| $writeArray2ToHashSet
|}
|$arrayBuilderClass $builder =
| ($arrayBuilderClass)$arrayBuilder.make($arrayBuilderClassTag);
|$arrayBuilderClass $builder = new $arrayBuilderClass();
|int $size = 0;
|for (int $i = 0; $i < $array1.numElements(); $i++) {
| $processArray1
Expand Down Expand Up @@ -4396,9 +4402,8 @@ case class ArrayExcept(left: Expression, right: Expression) extends ArraySetLike
val openHashSet = classOf[OpenHashSet[_]].getName
val classTag = s"scala.reflect.ClassTag$$.MODULE$$.$hsTypeName()"
val hashSet = ctx.freshName("hashSet")
val arrayBuilder = "scala.collection.mutable.ArrayBuilder"
val arrayBuilder = classOf[mutable.ArrayBuilder[_]].getName
val arrayBuilderClass = s"$arrayBuilder$$of$ptName"
val arrayBuilderClassTag = s"scala.reflect.ClassTag$$.MODULE$$.$ptName()"

def withArray2NullCheck(body: String): String =
if (right.dataType.asInstanceOf[ArrayType].containsNull) {
Expand Down Expand Up @@ -4474,8 +4479,7 @@ case class ArrayExcept(left: Expression, right: Expression) extends ArraySetLike
|for (int $i = 0; $i < $array2.numElements(); $i++) {
| $writeArray2ToHashSet
|}
|$arrayBuilderClass $builder =
| ($arrayBuilderClass)$arrayBuilder.make($arrayBuilderClassTag);
|$arrayBuilderClass $builder = new $arrayBuilderClass();
|int $size = 0;
|for (int $i = 0; $i < $array1.numElements(); $i++) {
| $processArray1
Expand Down
12 changes: 6 additions & 6 deletions sql/core/src/test/scala/org/apache/spark/sql/DataFrameFunctionsSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1679,26 +1679,26 @@ class DataFrameFunctionsSuite extends QueryTest with SharedSQLContext {
val df6 = Seq((null, null)).toDF("a", "b")
intercept[AnalysisException] {
Copy link
Member

@maropu maropu Aug 6, 2018

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could you also check the error message?

df6.select(array_intersect($"a", $"b"))
}
}.getMessage.contains("data type mismatch")
intercept[AnalysisException] {
df6.selectExpr("array_intersect(a, b)")
}
}.getMessage.contains("data type mismatch")

val df7 = Seq((Array(1), Array("a"))).toDF("a", "b")
intercept[AnalysisException] {
df7.select(array_intersect($"a", $"b"))
}
}.getMessage.contains("data type mismatch")
intercept[AnalysisException] {
df7.selectExpr("array_intersect(a, b)")
}
}.getMessage.contains("data type mismatch")

val df8 = Seq((null, Array("a"))).toDF("a", "b")
intercept[AnalysisException] {
df8.select(array_intersect($"a", $"b"))
}
}.getMessage.contains("data type mismatch")
intercept[AnalysisException] {
df8.selectExpr("array_intersect(a, b)")
}
}.getMessage.contains("data type mismatch")
}

test("transform function - array for primitive type not containing null") {
Expand Down