Skip to content

Commit b9995f5

Browse files
andygroveandygrove
committed
refactor: move Grace Hash Join tests into separate CometGraceHashJoinSuite
1 parent fb1a746 commit b9995f5

2 files changed

Lines changed: 297 additions & 254 deletions

File tree

‎spark/src/test/scala/org/apache/comet/exec/CometGraceHashJoinSuite.scala‎

Lines changed: 296 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,296 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
20+
package org.apache.comet.exec
21+
22+
import scala.util.Random
23+
24+
import org.scalactic.source.Position
25+
import org.scalatest.Tag
26+
27+
import org.apache.spark.sql.CometTestBase
28+
import org.apache.spark.sql.comet.CometHashJoinExec
29+
import org.apache.spark.sql.internal.SQLConf
30+
import org.apache.spark.sql.types.{DataTypes, Decimal, StructField, StructType}
31+
32+
import org.apache.comet.CometConf
33+
import org.apache.comet.testing.{DataGenOptions, ParquetGenerator}
34+
35+
class CometGraceHashJoinSuite extends CometTestBase {
36+
37+
import testImplicits._
38+
39+
override protected def test(testName: String, testTags: Tag*)(testFun: => Any)(implicit
40+
pos: Position): Unit = {
41+
super.test(testName, testTags: _*) {
42+
withSQLConf(CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "true") {
43+
testFun
44+
}
45+
}
46+
}
47+
48+
// Common SQL config for Grace Hash Join tests
49+
private val graceHashJoinConf: Seq[(String, String)] = Seq(
50+
CometConf.COMET_EXEC_GRACE_HASH_JOIN_NUM_PARTITIONS.key -> "4",
51+
"spark.sql.join.forceApplyShuffledHashJoin" -> "true",
52+
SQLConf.PREFER_SORTMERGEJOIN.key -> "false",
53+
SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
54+
SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1")
55+
56+
test("Grace HashJoin - all join types") {
57+
withSQLConf(graceHashJoinConf: _*) {
58+
withParquetTable((0 until 100).map(i => (i, i % 5)), "tbl_a") {
59+
withParquetTable((0 until 100).map(i => (i % 10, i + 2)), "tbl_b") {
60+
// Inner join
61+
checkSparkAnswer(
62+
sql(
63+
"SELECT /*+ SHUFFLE_HASH(tbl_a) */ * FROM tbl_a JOIN tbl_b ON tbl_a._2 = tbl_b._1"))
64+
65+
// Left join
66+
checkSparkAnswer(sql(
67+
"SELECT /*+ SHUFFLE_HASH(tbl_a) */ * FROM tbl_a LEFT JOIN tbl_b ON tbl_a._2 = tbl_b._1"))
68+
69+
// Right join
70+
checkSparkAnswer(sql(
71+
"SELECT /*+ SHUFFLE_HASH(tbl_a) */ * FROM tbl_a RIGHT JOIN tbl_b ON tbl_a._2 = tbl_b._1"))
72+
73+
// Full outer join
74+
checkSparkAnswer(sql(
75+
"SELECT /*+ SHUFFLE_HASH(tbl_a) */ * FROM tbl_a FULL JOIN tbl_b ON tbl_a._2 = tbl_b._1"))
76+
77+
// Left semi join
78+
checkSparkAnswer(sql(
79+
"SELECT /*+ SHUFFLE_HASH(tbl_b) */ * FROM tbl_a LEFT SEMI JOIN tbl_b ON tbl_a._2 = tbl_b._1"))
80+
81+
// Left anti join
82+
checkSparkAnswer(sql(
83+
"SELECT /*+ SHUFFLE_HASH(tbl_b) */ * FROM tbl_a LEFT ANTI JOIN tbl_b ON tbl_a._2 = tbl_b._1"))
84+
}
85+
}
86+
}
87+
}
88+
89+
test("Grace HashJoin - with filter condition") {
90+
withSQLConf(graceHashJoinConf: _*) {
91+
withParquetTable((0 until 100).map(i => (i, i % 5)), "tbl_a") {
92+
withParquetTable((0 until 100).map(i => (i % 10, i + 2)), "tbl_b") {
93+
checkSparkAnswer(
94+
sql("SELECT /*+ SHUFFLE_HASH(tbl_a) */ * FROM tbl_a JOIN tbl_b " +
95+
"ON tbl_a._2 = tbl_b._1 AND tbl_a._1 > tbl_b._2"))
96+
97+
checkSparkAnswer(
98+
sql("SELECT /*+ SHUFFLE_HASH(tbl_a) */ * FROM tbl_a LEFT JOIN tbl_b " +
99+
"ON tbl_a._2 = tbl_b._1 AND tbl_a._1 > tbl_b._2"))
100+
101+
checkSparkAnswer(
102+
sql("SELECT /*+ SHUFFLE_HASH(tbl_a) */ * FROM tbl_a FULL JOIN tbl_b " +
103+
"ON tbl_a._2 = tbl_b._1 AND tbl_a._1 > tbl_b._2"))
104+
}
105+
}
106+
}
107+
}
108+
109+
test("Grace HashJoin - various data types") {
110+
withSQLConf(graceHashJoinConf: _*) {
111+
// String keys
112+
withParquetTable((0 until 50).map(i => (s"key_${i % 10}", i)), "str_a") {
113+
withParquetTable((0 until 50).map(i => (s"key_${i % 5}", i * 2)), "str_b") {
114+
checkSparkAnswer(
115+
sql(
116+
"SELECT /*+ SHUFFLE_HASH(str_a) */ * FROM str_a JOIN str_b ON str_a._1 = str_b._1"))
117+
}
118+
}
119+
120+
// Decimal keys
121+
withParquetTable((0 until 50).map(i => (Decimal(i % 10), i)), "dec_a") {
122+
withParquetTable((0 until 50).map(i => (Decimal(i % 5), i * 2)), "dec_b") {
123+
checkSparkAnswer(
124+
sql(
125+
"SELECT /*+ SHUFFLE_HASH(dec_a) */ * FROM dec_a JOIN dec_b ON dec_a._1 = dec_b._1"))
126+
}
127+
}
128+
}
129+
}
130+
131+
test("Grace HashJoin - empty tables") {
132+
withSQLConf(graceHashJoinConf: _*) {
133+
withParquetTable(Seq.empty[(Int, Int)], "empty_a") {
134+
withParquetTable((0 until 10).map(i => (i, i)), "nonempty_b") {
135+
// Empty left side
136+
checkSparkAnswer(sql(
137+
"SELECT /*+ SHUFFLE_HASH(empty_a) */ * FROM empty_a JOIN nonempty_b ON empty_a._1 = nonempty_b._1"))
138+
139+
// Empty left with left join
140+
checkSparkAnswer(sql(
141+
"SELECT /*+ SHUFFLE_HASH(empty_a) */ * FROM empty_a LEFT JOIN nonempty_b ON empty_a._1 = nonempty_b._1"))
142+
143+
// Empty right side
144+
checkSparkAnswer(sql(
145+
"SELECT /*+ SHUFFLE_HASH(nonempty_b) */ * FROM nonempty_b JOIN empty_a ON nonempty_b._1 = empty_a._1"))
146+
147+
// Empty right with right join
148+
checkSparkAnswer(sql(
149+
"SELECT /*+ SHUFFLE_HASH(nonempty_b) */ * FROM nonempty_b RIGHT JOIN empty_a ON nonempty_b._1 = empty_a._1"))
150+
}
151+
}
152+
}
153+
}
154+
155+
test("Grace HashJoin - self join") {
156+
withSQLConf(graceHashJoinConf: _*) {
157+
withParquetTable((0 until 50).map(i => (i, i % 10)), "self_tbl") {
158+
checkSparkAnswer(
159+
sql("SELECT /*+ SHUFFLE_HASH(a) */ * FROM self_tbl a JOIN self_tbl b ON a._2 = b._2"))
160+
}
161+
}
162+
}
163+
164+
test("Grace HashJoin - build side selection") {
165+
withSQLConf(graceHashJoinConf: _*) {
166+
withParquetTable((0 until 100).map(i => (i, i % 5)), "tbl_a") {
167+
withParquetTable((0 until 100).map(i => (i % 10, i + 2)), "tbl_b") {
168+
// Build left (hint on left table)
169+
checkSparkAnswer(
170+
sql(
171+
"SELECT /*+ SHUFFLE_HASH(tbl_a) */ * FROM tbl_a JOIN tbl_b ON tbl_a._2 = tbl_b._1"))
172+
173+
// Build right (hint on right table)
174+
checkSparkAnswer(
175+
sql(
176+
"SELECT /*+ SHUFFLE_HASH(tbl_b) */ * FROM tbl_a JOIN tbl_b ON tbl_a._2 = tbl_b._1"))
177+
178+
// Left join build right
179+
checkSparkAnswer(sql(
180+
"SELECT /*+ SHUFFLE_HASH(tbl_b) */ * FROM tbl_a LEFT JOIN tbl_b ON tbl_a._2 = tbl_b._1"))
181+
182+
// Right join build left
183+
checkSparkAnswer(sql(
184+
"SELECT /*+ SHUFFLE_HASH(tbl_a) */ * FROM tbl_a RIGHT JOIN tbl_b ON tbl_a._2 = tbl_b._1"))
185+
}
186+
}
187+
}
188+
}
189+
190+
test("Grace HashJoin - plan shows CometHashJoinExec") {
191+
withSQLConf(graceHashJoinConf: _*) {
192+
withParquetTable((0 until 50).map(i => (i, i % 5)), "tbl_a") {
193+
withParquetTable((0 until 50).map(i => (i % 10, i + 2)), "tbl_b") {
194+
val df = sql(
195+
"SELECT /*+ SHUFFLE_HASH(tbl_a) */ * FROM tbl_a JOIN tbl_b ON tbl_a._2 = tbl_b._1")
196+
checkSparkAnswerAndOperator(df, Seq(classOf[CometHashJoinExec]))
197+
}
198+
}
199+
}
200+
}
201+
202+
test("Grace HashJoin - multiple key columns") {
203+
withSQLConf(graceHashJoinConf: _*) {
204+
withParquetTable((0 until 50).map(i => (i, i % 5, i % 3)), "multi_a") {
205+
withParquetTable((0 until 50).map(i => (i % 10, i % 5, i % 3)), "multi_b") {
206+
checkSparkAnswer(
207+
sql("SELECT /*+ SHUFFLE_HASH(multi_a) */ * FROM multi_a JOIN multi_b " +
208+
"ON multi_a._2 = multi_b._2 AND multi_a._3 = multi_b._3"))
209+
}
210+
}
211+
}
212+
}
213+
214+
// Schema with types that work well as join keys (no NaN/float issues)
215+
private val fuzzJoinSchema = StructType(
216+
Seq(
217+
StructField("c_int", DataTypes.IntegerType),
218+
StructField("c_long", DataTypes.LongType),
219+
StructField("c_str", DataTypes.StringType),
220+
StructField("c_date", DataTypes.DateType),
221+
StructField("c_dec", DataTypes.createDecimalType(10, 2)),
222+
StructField("c_short", DataTypes.ShortType),
223+
StructField("c_bool", DataTypes.BooleanType)))
224+
225+
private val joinTypes =
226+
Seq("JOIN", "LEFT JOIN", "RIGHT JOIN", "FULL JOIN", "LEFT SEMI JOIN", "LEFT ANTI JOIN")
227+
228+
test("Grace HashJoin fuzz - all join types with generated data") {
229+
val dataGenOptions =
230+
DataGenOptions(allowNull = true, generateNegativeZero = false, generateNaN = false)
231+
232+
withSQLConf(graceHashJoinConf: _*) {
233+
withTempPath { dir =>
234+
val path1 = s"${dir.getAbsolutePath}/fuzz_left"
235+
val path2 = s"${dir.getAbsolutePath}/fuzz_right"
236+
val random = new Random(42)
237+
238+
withSQLConf(CometConf.COMET_ENABLED.key -> "false") {
239+
ParquetGenerator
240+
.makeParquetFile(random, spark, path1, fuzzJoinSchema, 200, dataGenOptions)
241+
ParquetGenerator
242+
.makeParquetFile(random, spark, path2, fuzzJoinSchema, 200, dataGenOptions)
243+
}
244+
245+
spark.read.parquet(path1).createOrReplaceTempView("fuzz_l")
246+
spark.read.parquet(path2).createOrReplaceTempView("fuzz_r")
247+
248+
for (jt <- joinTypes) {
249+
// Join on int column
250+
checkSparkAnswer(sql(
251+
s"SELECT /*+ SHUFFLE_HASH(fuzz_l) */ * FROM fuzz_l $jt fuzz_r ON fuzz_l.c_int = fuzz_r.c_int"))
252+
253+
// Join on string column
254+
checkSparkAnswer(sql(
255+
s"SELECT /*+ SHUFFLE_HASH(fuzz_l) */ * FROM fuzz_l $jt fuzz_r ON fuzz_l.c_str = fuzz_r.c_str"))
256+
257+
// Join on decimal column
258+
checkSparkAnswer(sql(
259+
s"SELECT /*+ SHUFFLE_HASH(fuzz_l) */ * FROM fuzz_l $jt fuzz_r ON fuzz_l.c_dec = fuzz_r.c_dec"))
260+
}
261+
}
262+
}
263+
}
264+
265+
test("Grace HashJoin fuzz - with spilling") {
266+
val dataGenOptions =
267+
DataGenOptions(allowNull = true, generateNegativeZero = false, generateNaN = false)
268+
269+
// Use very small memory pool to force spilling
270+
withSQLConf(
271+
(graceHashJoinConf ++ Seq(
272+
CometConf.COMET_ONHEAP_MEMORY_OVERHEAD.key -> "10000000",
273+
CometConf.COMET_EXEC_GRACE_HASH_JOIN_NUM_PARTITIONS.key -> "8")): _*) {
274+
withTempPath { dir =>
275+
val path1 = s"${dir.getAbsolutePath}/spill_left"
276+
val path2 = s"${dir.getAbsolutePath}/spill_right"
277+
val random = new Random(99)
278+
279+
withSQLConf(CometConf.COMET_ENABLED.key -> "false") {
280+
ParquetGenerator
281+
.makeParquetFile(random, spark, path1, fuzzJoinSchema, 500, dataGenOptions)
282+
ParquetGenerator
283+
.makeParquetFile(random, spark, path2, fuzzJoinSchema, 500, dataGenOptions)
284+
}
285+
286+
spark.read.parquet(path1).createOrReplaceTempView("spill_l")
287+
spark.read.parquet(path2).createOrReplaceTempView("spill_r")
288+
289+
for (jt <- joinTypes) {
290+
checkSparkAnswer(sql(
291+
s"SELECT /*+ SHUFFLE_HASH(spill_l) */ * FROM spill_l $jt spill_r ON spill_l.c_int = spill_r.c_int"))
292+
}
293+
}
294+
}
295+
}
296+
}

0 commit comments

Comments
 (0)