andygrove commented on code in PR #6835:
URL: https://github.com/apache/datafusion-comet/pull/6835#discussion_r4285210544
##########
spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala:
##########
@@ -3288,4 +3290,262 @@ class CometInMemoryCacheSuite extends CometTestBase {
}.nonEmpty)
}
}
+
+ test("hash joins over a sorted cache keep the streamed side's order") {
+ // CometInMemoryTableScan reports the cached plan's sort order, so Spark
drops the sort it
+ // would otherwise plan above a join that keeps that order. DataFusion's
hash join keeps the
+ // unmatched probe rows of a left outer join built on the right, or of a
right outer join built
+ // on the left, in place only when the probe input declares an order,
which the native scan
+ // over the cache does not, so the native join declares the cache's order
on its probe input.
+ // An inner join never moves probe rows.
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.SHUFFLE_PARTITIONS.key -> "2",
+ CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true") {
+ withParquetTable((0 until 10000).map(i => (i % 100, i)), "big") {
+ withParquetTable((0 until 10).map(i => (i * 10, i)), "small") {
+ try {
+ val streamed =
+ spark.table("big").repartition(2,
$"_1").sortWithinPartitions("_1").cache()
+ streamed.count()
+ val build = spark.table("small")
+
+ def query(hint: String, joinType: String): DataFrame = {
+ val joined = if (joinType == "right_outer") {
+ build.hint(hint).join(streamed, build("_1") ===
streamed("_1"), joinType)
+ } else {
+ streamed.join(build.hint(hint), streamed("_1") ===
build("_1"), joinType)
+ }
+ joined
+ .select(streamed("_1").as("k"), streamed("_2").as("v"),
build("_2").as("w"))
+ .sortWithinPartitions("k")
+ }
+
+ def joins(plan: SparkPlan): Seq[(Class[_], BuildSide)] =
collect(plan) {
+ case j: CometHashJoinExec => (j.getClass, j.buildSide)
+ case j: CometBroadcastHashJoinExec => (j.getClass, j.buildSide)
+ }
+
+ def sorts(plan: SparkPlan): Seq[SparkPlan] =
+ collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s }
+
+ val outerRows = (0 until 10000).map { i =>
+ val k = i % 100
+ Row(k, i, if (k % 10 == 0) k / 10 else null)
+ }
+ val innerRows = outerRows.filter(_.get(2) != null)
+
+ Seq(
+ "shuffle_hash" -> classOf[CometHashJoinExec],
+ "broadcast" -> classOf[CometBroadcastHashJoinExec]).foreach {
+ case (hint, joinClass) =>
+ Seq("left_outer" -> BuildRight, "right_outer" ->
BuildLeft).foreach {
+ case (joinType, buildSide) =>
+ val outer = query(hint, joinType)
+ val outerPlan = outer.queryExecution.executedPlan
+ assert(joins(outerPlan) == Seq((joinClass, buildSide)),
outerPlan)
+ assert(collect(outerPlan) { case s:
CometInMemoryTableScanExec =>
+ s
+ }.size == 1)
+ assert(sorts(outerPlan).isEmpty, outerPlan)
+ QueryTest.sameRows(outerRows,
outer.collect().toSeq).foreach(fail(_))
+ assert(sortedPartitions(outer).forall(identity), outerPlan)
+ }
+
+ val inner = query(hint, "inner")
+ val innerPlan = inner.queryExecution.executedPlan
+ assert(joins(innerPlan) == Seq((joinClass, BuildRight)),
innerPlan)
+ assert(sorts(innerPlan).isEmpty, innerPlan)
+ QueryTest.sameRows(innerRows,
inner.collect().toSeq).foreach(fail(_))
+ assert(sortedPartitions(inner).forall(identity), innerPlan)
+ }
+ } finally {
+ spark.catalog.clearCache()
+ }
+ }
+ }
+ }
+ }
+
+ // Whether each partition of `df` has the integer key in its first column in
ascending order.
+ private def sortedPartitions(df: DataFrame): Array[Boolean] =
+ df.queryExecution.toRdd
+ .mapPartitions { it =>
+ val keys = it.map(_.getInt(0)).toArray
+ Iterator(keys.sameElements(keys.sorted))
+ }
+ .collect()
+
+ test("a broadcast outer hash join over a sorted cache keeps its order with
default settings") {
+ // A cache written with no exchange in its plan is not adaptive, so with
AQE on, as it is by
+ // default, the scan reports the cached plan's order just as it does with
AQE off. Only the
+ // native cache scan is switched on here, and the broadcast join is the
one Spark plans for
+ // a small build side.
+ withSQLConf(CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true") {
+ withParquetTable((0 until 10000).map(i => (i % 100, i)), "big") {
+ withParquetTable((0 until 10).map(i => (i * 10, i)), "small") {
+ try {
+ val streamed =
spark.table("big").sortWithinPartitions("_1").cache()
+ streamed.count()
+ val build = spark.table("small")
+ val df = streamed
+ .join(build, streamed("_1") === build("_1"), "left_outer")
+ .select(streamed("_1").as("k"), streamed("_2").as("v"),
build("_2").as("w"))
+ .sortWithinPartitions("k")
+ val expected = (0 until 10000).map { i =>
+ val k = i % 100
+ Row(k, i, if (k % 10 == 0) k / 10 else null)
+ }
+ QueryTest.sameRows(expected, df.collect().toSeq).foreach(fail(_))
+ // The plan AQE chose is final once the query has run.
+ val plan = df.queryExecution.executedPlan
+ assert(plan.isInstanceOf[AdaptiveSparkPlanExec], plan)
+ val joins = collect(plan) { case j: CometBroadcastHashJoinExec =>
j }
+ assert(joins.size == 1 && joins.head.joinType == LeftOuter, plan)
+ assert(joins.head.nativeOp.getHashJoin.getStreamedSortOrdersCount
> 0, plan)
+ assert(collect(plan) { case s @ (_: CometSortExec | _: SortExec)
=> s }.isEmpty, plan)
+ assert(collect(plan) { case s: CometInMemoryTableScanExec => s
}.size == 1, plan)
+ assert(sortedPartitions(df).forall(identity), plan)
+ } finally {
+ spark.catalog.clearCache()
+ }
+ }
+ }
+ }
+ }
+
+ test("an outer hash join over a native cache scan of a NOT IN keeps each
group once") {
+ // The native cache scan reports the order of the plan that wrote the
cache. A NOT IN is not in
+ // the order Spark reports for it, since Comet builds it on its streamed
side, so the outer
+ // join that reads the cache must not declare the order, or the aggregates
above it run in
+ // sorted mode over unsorted rows and split groups. Sparse keys and small
batches make each
+ // group span several. A cache with an exchange in its plan is written
under AQE, and the scan
+ // reads the final plan behind the adaptive one, so this runs with AQE off
and on.
+ Seq(false, true).foreach { adaptive =>
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString,
+ SQLConf.SHUFFLE_PARTITIONS.key -> "2",
+ SQLConf.COLUMN_BATCH_SIZE.key -> "8",
+ CometConf.COMET_BATCH_SIZE.key -> "8",
+ CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true") {
+ withParquetTable((0 until 2000).map(i => ((i % 40) * 1000, i)), "big")
{
+ withParquetTable(Seq(1000, 10000, 20000, 100000).map(x =>
Tuple1(Option(x))), "excl") {
+ withParquetTable((0 until 40 by 2).map(k => (k * 1000, k)),
"small") {
+ try {
+ spark
+ .table("big")
+ .toDF("k", "v")
+ .repartition(2, $"k")
+ .sortWithinPartitions("k")
+ .cache()
+ .createOrReplaceTempView("cached")
+ spark.table("cached").count()
+ sql("SELECT * FROM cached WHERE k NOT IN (SELECT _1 FROM
excl)")
+ .cache()
+ .createOrReplaceTempView("anti")
+ spark.table("anti").count()
+ assert(
+ !sortedPartitions(spark.table("anti")).forall(identity),
+ "the cached NOT IN rows are in order, so nothing here can
split a group")
+ val query =
+ """SELECT /*+ BROADCAST(b) */ a.k, count(*), sum(a.v),
count(b._2)
+ |FROM anti a LEFT JOIN small b ON a.k = b._1
+ |GROUP BY a.k""".stripMargin
+ def rows(): Seq[(Int, Long, Long, Long)] =
+ sql(query)
+ .collect()
+ .map(r => (r.getInt(0), r.getLong(1), r.getLong(2),
r.getLong(3)))
+ .toSeq
+ .sorted
+ val cometRows = rows()
+ val plan = sql(query).queryExecution.executedPlan
+ assert(
+ collect(plan) { case s: CometInMemoryTableScanExec => s
}.size == 1,
+ s"the outer join does not read the native cache scan (AQE
$adaptive):\n$plan")
+ // withSQLConf returns Unit on Spark 3.x, so Spark's rows are
assigned inside it.
+ var sparkRows = Seq.empty[(Int, Long, Long, Long)]
+ withSQLConf(CometConf.COMET_ENABLED.key -> "false") {
+ sparkRows = rows()
+ }
+ assert(
+ cometRows.map(_._1).distinct.size == cometRows.size,
+ "a key is grouped twice")
+ assert(cometRows == sparkRows)
+ } finally {
+ spark.catalog.clearCache()
+ spark.catalog.dropTempView("anti")
+ spark.catalog.dropTempView("cached")
+ }
+ }
+ }
+ }
+ }
+ }
+ }
+
+ test("an aggregate on the key above a hash join over a sorted cache groups
each key once") {
+ // The order the join declares on its probe input reaches past the join,
so the aggregate
+ // above it in the same native plan groups in sorted, streaming mode. That
is right while
+ // Spark's reported order holds, and a wrong order would split a key
across groups, which only
+ // a comparison of the rows as a multiset catches. Small batches make each
group span several.
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.SHUFFLE_PARTITIONS.key -> "2",
+ SQLConf.COLUMN_BATCH_SIZE.key -> "8",
+ CometConf.COMET_BATCH_SIZE.key -> "8",
+ CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true") {
+ withParquetTable((0 until 2000).map(i => ((i % 40) * 1000, i)), "big") {
+ withParquetTable((0 until 40 by 2).map(m => (m * 1000, m)), "small") {
+ try {
+ val streamed =
+ spark.table("big").repartition(2,
$"_1").sortWithinPartitions("_1").cache()
+ streamed.count()
+ val build = spark.table("small")
+
+ // Fifty rows per key, and a match in the build side for every
even key.
+ val expected = (0 until 40).map { m =>
+ Row(m * 1000, 50L, 50L * m + 49000L, if (m % 2 == 0) 50L else 0L)
+ }
+
+ for {
+ (hint, joinClass) <- Seq(
+ "shuffle_hash" -> classOf[CometHashJoinExec],
+ "broadcast" -> classOf[CometBroadcastHashJoinExec])
+ (joinType, buildSide) <- Seq("left_outer" -> BuildRight,
"right_outer" -> BuildLeft)
+ } {
+ val joined = if (joinType == "right_outer") {
+ build.hint(hint).join(streamed, build("_1") ===
streamed("_1"), joinType)
+ } else {
+ streamed.join(build.hint(hint), streamed("_1") ===
build("_1"), joinType)
+ }
+ val df = joined
+ .groupBy(streamed("_1"))
+ .agg(count(lit(1)), sum(streamed("_2")), count(build("_2")))
+ val plan = df.queryExecution.executedPlan
+ val joins: Seq[(SparkPlan, BuildSide)] = collect(plan) {
+ case j: CometHashJoinExec => (j, j.buildSide)
+ case j: CometBroadcastHashJoinExec => (j, j.buildSide)
+ }
+ assert(
+ joins.map { case (j, side) => (j.getClass, side) } ==
Seq((joinClass, buildSide)),
+ plan)
+ // A final aggregate after a shuffle would merge a key a partial
aggregate split, so
+ // both aggregates have to sit above the join with no exchange
in between.
+ def aboveJoin(p: SparkPlan): Boolean = p.find(_ eq
joins.head._1).isDefined
+ assert(
+ collect(plan) { case a: CometHashAggregateExec => a
}.count(aboveJoin) == 2,
Review Comment:
You were right, and it was worse than the missing assertion. With the
planner's declaration removed the test still passed, because with 50 rows per
key a join that moves its unmatched rows still hands the aggregate the keys in
first-seen order. Fixed in 8e75130543: the test now also runs 1000 keys with
two rows each, where an 8-row batch holds alternately matched and unmatched
keys, so an undeclared join changes the order the aggregate first sees them in.
It asserts that the serde sent the order (`getStreamedSortOrdersCount > 0`) and
that every partition comes back in key order. With the declaration removed it
fails on the sorted-partitions assertion.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
