Skip to content
Closed
Changes from 1 commit
Commits
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
Simplified assertions for SortMerge and HashJoin
  • Loading branch information
vladanvasi-db committed Nov 8, 2024
commit 9851790db3e24a9ef44e99dfebb1b1955ffa6a2c
45 changes: 13 additions & 32 deletions sql/core/src/test/scala/org/apache/spark/sql/CollationSuite.scala
Original file line number Diff line number Diff line change
Expand Up @@ -51,29 +51,10 @@ class CollationSuite extends DatasourceV2SQLBase with AdaptiveSparkPlanHelper {

private def checkRightTypeOfJoinUsed(queryPlan: SparkPlan): Unit = {
assert(
// If sort merge join is forced, we should not see HashJoin in the plan.
isSortMergeForced ||
// If sort merge join is not forced, we should see HashJoin in the plan
// and not SortMergeJoin.
collectFirst(queryPlan) {
case _: HashJoin => ()
}.nonEmpty &&
collectFirst(queryPlan) {
case _: SortMergeJoinExec => ()
}.isEmpty
)

assert(
// If sort merge join is not forced, we should not see SortMergeJoin in the plan.
!isSortMergeForced ||
// If sort merge join is forced, we should see SortMergeJoin in the plan
// and not HashJoin.
collectFirst(queryPlan) {
case _: HashJoin => ()
}.isEmpty &&
collectFirst(queryPlan) {
case _: SortMergeJoinExec => ()
}.nonEmpty
collectFirst(queryPlan) {
case _: SortMergeJoinExec => assert(isSortMergeForced)
case _: HashJoin => assert(!isSortMergeForced)
}.nonEmpty
)
}

Expand Down Expand Up @@ -1413,7 +1394,7 @@ class CollationSuite extends DatasourceV2SQLBase with AdaptiveSparkPlanHelper {
for (codeGen <- Seq("NO_CODEGEN", "CODEGEN_ONLY")) {
val collationSetup = if (collation.isEmpty) "" else " COLLATE " + collation
val supportsBinaryEquality = collation.isEmpty || collation == "UNICODE" ||
CollationFactory.fetchCollation(collation).isUtf8BinaryType
CollationFactory.fetchCollation(collation).supportsBinaryEquality

test(s"Group by on map containing$collationSetup strings ($codeGen)") {
val tableName = "t"
Expand Down Expand Up @@ -1618,15 +1599,15 @@ class CollationSuite extends DatasourceV2SQLBase with AdaptiveSparkPlanHelper {

if (isSortMergeForced) {
// Only if collation doesn't support binary equality, collation key should be injected.
if (!CollationFactory.fetchCollation(t.collation).isUtf8BinaryType) {
if (!CollationFactory.fetchCollation(t.collation).supportsBinaryEquality) {
assert(queryPlan.toString().contains("collationkey"))
} else {
assert(!queryPlan.toString().contains("collationkey"))
}
}
else {
// Only if collation doesn't support binary equality, collation key should be injected.
if (!CollationFactory.fetchCollation(t.collation).isUtf8BinaryType) {
if (!CollationFactory.fetchCollation(t.collation).supportsBinaryEquality) {
assert(collectFirst(queryPlan) {
case b: HashJoin => b.leftKeys.head
}.head.isInstanceOf[CollationKey])
Expand Down Expand Up @@ -1681,15 +1662,15 @@ class CollationSuite extends DatasourceV2SQLBase with AdaptiveSparkPlanHelper {

if (isSortMergeForced) {
// Only if collation doesn't support binary equality, collation key should be injected.
if (!CollationFactory.fetchCollation(t.collation).isUtf8BinaryType) {
if (!CollationFactory.fetchCollation(t.collation).supportsBinaryEquality) {
assert(queryPlan.toString().contains("collationkey"))
} else {
assert(!queryPlan.toString().contains("collationkey"))
}
}
else {
// Only if collation doesn't support binary equality, collation key should be injected.
if (!CollationFactory.fetchCollation(t.collation).isUtf8BinaryType) {
if (!CollationFactory.fetchCollation(t.collation).supportsBinaryEquality) {
assert(collectFirst(queryPlan) {
case b: BroadcastHashJoinExec => b.leftKeys.head
}.head.asInstanceOf[ArrayTransform].function.asInstanceOf[LambdaFunction].
Expand Down Expand Up @@ -1749,15 +1730,15 @@ class CollationSuite extends DatasourceV2SQLBase with AdaptiveSparkPlanHelper {

if (isSortMergeForced) {
// Only if collation doesn't support binary equality, collation key should be injected.
if (!CollationFactory.fetchCollation(t.collation).isUtf8BinaryType) {
if (!CollationFactory.fetchCollation(t.collation).supportsBinaryEquality) {
assert(queryPlan.toString().contains("collationkey"))
} else {
assert(!queryPlan.toString().contains("collationkey"))
}
}
else {
// Only if collation doesn't support binary equality, collation key should be injected.
if (!CollationFactory.fetchCollation(t.collation).isUtf8BinaryType) {
if (!CollationFactory.fetchCollation(t.collation).supportsBinaryEquality) {
assert(collectFirst(queryPlan) {
case b: BroadcastHashJoinExec => b.leftKeys.head
}.head.asInstanceOf[ArrayTransform].function.
Expand Down Expand Up @@ -1813,7 +1794,7 @@ class CollationSuite extends DatasourceV2SQLBase with AdaptiveSparkPlanHelper {
checkRightTypeOfJoinUsed(queryPlan)

// Only if collation doesn't support binary equality, collation key should be injected.
if (!CollationFactory.fetchCollation(t.collation).isUtf8BinaryType) {
if (!CollationFactory.fetchCollation(t.collation).supportsBinaryEquality) {
assert(queryPlan.toString().contains("collationkey"))
} else {
assert(!queryPlan.toString().contains("collationkey"))
Expand Down Expand Up @@ -1870,7 +1851,7 @@ class CollationSuite extends DatasourceV2SQLBase with AdaptiveSparkPlanHelper {
checkRightTypeOfJoinUsed(queryPlan)

// Only if collation doesn't support binary equality, collation key should be injected.
if (!CollationFactory.fetchCollation(t.collation).isUtf8BinaryType) {
if (!CollationFactory.fetchCollation(t.collation).supportsBinaryEquality) {
assert(queryPlan.toString().contains("collationkey"))
} else {
assert(!queryPlan.toString().contains("collationkey"))
Expand Down