kazuyukitanimura commented on code in PR #6835:
URL: https://github.com/apache/datafusion-comet/pull/6835#discussion_r4265230980
##########
spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala:
##########
@@ -2254,4 +2255,636 @@ class CometJoinSuite extends CometTestBase {
}
}
}
+
+ // Whether every partition of `df` is in `order`, which is bound to its
output.
+ private def partitionsInOrder(df: DataFrame, order: SortOrder*): Boolean = {
+ val ordering = new InterpretedOrdering(order)
+ df.queryExecution.toRdd.map(_.copy()).glom().collect().forall { part =>
+ part.zip(part.drop(1)).forall { case (a, b) => !ordering.gt(a, b) }
+ }
+ }
+
+ // The order of the integer key in the first output column.
+ private def keyOrder(direction: SortDirection): SortOrder =
+ SortOrder(BoundReference(0, IntegerType, nullable = true), direction)
+
+ test("outer hash joins keep a streamed side sorted outside the native plan
in order") {
+ // Spark drops the sort above the join because the join reports the
streamed side's order,
+ // and the streamed side arrives sorted through a JVM input, which
declares no order to
+ // DataFusion unless the planner declares it on the join's probe input; see
+ // CometHashJoin.doConvert. CometInMemoryCacheSuite covers the native
cache scan.
+ 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 -> "false",
+ CometConf.COMET_CONVERT_FROM_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")
+
+ // A residual condition, which keeps only the build rows past 3,
sends more probe rows
+ // down the join's unmatched path.
+ def expected(withCondition: Boolean): Seq[Row] = (0 until
10000).map { i =>
+ val k = i % 100
+ val matched = k % 10 == 0 && (!withCondition || k / 10 > 3)
+ Row(k, i, if (matched) k / 10 else null)
+ }
+
+ // Both shapes run as DataFusion's right join: Comet swaps a join
built on the right
+ // and leaves one built on the left as it is.
+ def outerJoin(hint: String, joinType: String, withCondition:
Boolean): DataFrame = {
+ val keys = streamed("_1") === build("_1")
+ val on = if (withCondition) keys && build("_2") > 3 else keys
+ val joined = if (joinType == "left_outer") {
+ streamed.join(build.hint(hint), on, joinType)
+ } else {
+ build.hint(hint).join(streamed, on, joinType)
+ }
+ joined
+ .select(streamed("_1").as("k"), streamed("_2").as("v"),
build("_2").as("w"))
+ .sortWithinPartitions("k")
+ }
+
+ for {
+ (hint, joinClass) <- Seq(
+ "shuffle_hash" -> classOf[CometHashJoinExec],
+ "broadcast" -> classOf[CometBroadcastHashJoinExec])
+ (joinType, buildSide) <- Seq("left_outer" -> BuildRight,
"right_outer" -> BuildLeft)
+ withCondition <- Seq(false, true)
+ } {
+ val df = outerJoin(hint, joinType, withCondition)
+ val plan = df.queryExecution.executedPlan
+ val joins: Seq[(Class[_], BuildSide)] = collect(plan) {
+ case j: CometHashJoinExec => (j.getClass, j.buildSide)
+ case j: CometBroadcastHashJoinExec => (j.getClass, j.buildSide)
+ }
+ assert(joins == Seq((joinClass, buildSide)), plan)
+ assert(collect(plan) { case c: CometSparkToColumnarExec => c
}.size == 1, plan)
+ assert(
+ collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s
}.isEmpty,
+ plan)
+ QueryTest.sameRows(expected(withCondition),
df.collect().toSeq).foreach(fail(_))
+ assert(partitionsInOrder(df, keyOrder(Ascending)), plan)
+ }
+ } finally {
+ spark.catalog.clearCache()
+ }
+ }
+ }
+ }
+ }
+
+ test("an outer hash join declares the part of the streamed order Comet can
express") {
+ // Declaring any order keeps the probe rows in place, so a prefix of the
order is enough.
+ // Nothing tells whether an operator above relies on the order, so a join
whose first key Comet
+ // cannot express stays in Spark rather than lose it. With Add disabled,
Comet cannot express
+ // `_1 + 1`.
+ 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 -> "false",
+ CometConf.COMET_CONVERT_FROM_IN_MEMORY_CACHE_ENABLED.key -> "true",
+ CometConf.getExprEnabledConfigKey("Add") -> "false") {
+ withParquetTable((0 until 1000).map(i => (i % 100, i)), "big") {
+ withParquetTable((0 until 10).map(i => (i * 10, i)), "small") {
+ def sortedCache(order: Column*): DataFrame = {
+ val cached =
+ spark.table("big").repartition(2,
$"_1").sortWithinPartitions(order: _*).cache()
+ cached.count()
+ cached
+ }
+ val build = spark.table("small")
+ try {
+ val unexpressed = sortedCache($"_1" + 1)
+ val (_, fallbackPlan) = checkSparkAnswerAndFallbackReason(
+ unexpressed.join(
+ build.hint("shuffle_hash"),
+ unexpressed("_1") === build("_1"),
+ "left_outer"),
+ "the streamed side's sort order, which this outer join has to
keep, cannot be " +
+ "expressed natively")
+ assert(
+ collect(fallbackPlan) { case j: ShuffledHashJoinExec => j }.size
== 1,
+ fallbackPlan)
+ assert(collect(fallbackPlan) { case j: CometHashJoinExec => j
}.isEmpty, fallbackPlan)
+
+ val prefixed = sortedCache($"_1", $"_1" + 1)
+ val df = prefixed
+ .join(build.hint("shuffle_hash"), prefixed("_1") ===
build("_1"), "left_outer")
+ .select(prefixed("_1").as("k"), prefixed("_2").as("v"),
build("_2").as("w"))
+ .sortWithinPartitions("k")
+ val plan = df.queryExecution.executedPlan
+ assert(collect(plan) { case j: CometHashJoinExec => j }.size == 1,
plan)
+ assert(collect(plan) { case s @ (_: CometSortExec | _: SortExec)
=> s }.isEmpty, plan)
+ val expected = (0 until 1000).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(_))
+ assert(partitionsInOrder(df, keyOrder(Ascending)), plan)
+ } finally {
+ spark.catalog.clearCache()
+ }
+ }
+ }
+ }
+ }
+
+ test("outer hash joins keep a descending, nulls last streamed order with
null keys") {
+ // Null keys match nothing, so all of them take the unmatched path, at the
end of each
+ // partition.
+ 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 -> "false",
+ CometConf.COMET_CONVERT_FROM_IN_MEMORY_CACHE_ENABLED.key -> "true") {
+ val data = (0 until 10000).map(i => (if (i % 7 == 0) None else Some(i %
100), i))
+ withParquetTable(data, "big") {
+ withParquetTable((0 until 10).map(i => (i * 10, i)), "small") {
+ try {
+ val streamed = spark
+ .table("big")
+ .repartition(2, $"_1")
+ .sortWithinPartitions($"_1".desc_nulls_last)
+ .cache()
+ streamed.count()
+ val build = spark.table("small").hint("shuffle_hash")
+ val expected = data.map { case (k, i) =>
+ Row(k.getOrElse(null), i, k.filter(_ % 10 == 0).map(_ /
10).getOrElse(null))
+ }
+ Seq("left_outer", "right_outer").foreach { joinType =>
+ val joined = if (joinType == "left_outer") {
+ streamed.join(build, streamed("_1") === build("_1"), joinType)
+ } else {
+ build.join(streamed, build("_1") === streamed("_1"), joinType)
+ }
+ val df = joined
+ .select(streamed("_1").as("k"), streamed("_2").as("v"),
build("_2").as("w"))
+ .sortWithinPartitions($"k".desc_nulls_last)
+ val plan = df.queryExecution.executedPlan
+ assert(collect(plan) { case j: CometHashJoinExec => j }.size ==
1, plan)
+ assert(
+ collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s
}.isEmpty,
+ plan)
+ QueryTest.sameRows(expected, df.collect().toSeq).foreach(fail(_))
+ assert(partitionsInOrder(df, keyOrder(Descending)), plan)
+ }
+ } finally {
+ spark.catalog.clearCache()
+ }
+ }
+ }
+ }
+ }
+
+ test("NOT IN below an outer hash join over a sorted cache keeps each group
once") {
+ // Spark reports the streamed side's order for a null-aware anti join, but
Comet builds that
+ // join on its streamed side, so the order does not hold above it.
Declaring it on the outer
+ // join's probe input would put the aggregates in sorted mode over
unsorted rows and split
+ // groups. Sparse keys keep DataFusion off its perfect hash build path,
which keeps the build
+ // batches in order and would hide that, and small batches spread each
partition over many.
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ 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()
+ val query =
+ """SELECT /*+ BROADCAST(b) */ a.k, count(*), sum(a.v),
count(b._2)
+ |FROM (SELECT * FROM cached WHERE k NOT IN (SELECT _1 FROM
excl)) a
+ |LEFT JOIN small b ON a.k = b._1
+ |GROUP BY a.k""".stripMargin
+ checkEachKeyGroupedOnce(query)
+ } finally {
+ spark.catalog.clearCache()
+ spark.catalog.dropTempView("cached")
+ }
+ }
+ }
+ }
+ }
+ }
+
+ test("an outer hash join over a cache of a NOT IN keeps each group once") {
+ // The same NOT IN as the test above, but cached, so its rows reach the
outer join from the
+ // JVM. The cache reports the order Spark reports for the null-aware anti
join, which its rows
+ // are not in, and the outer join's probe input is then a plain JVM input
as far as the native
+ // plan can see. The order must not be declared there either. A cache with
an exchange in its
+ // plan is written under AQE from Spark 3.5, and Spark's scan then reports
the order of the
+ // final plan behind the adaptive one, which reports none itself, 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()
+ val cachedPlan = cachedPlanOf(spark.table("anti"))
+ assert(
+ find(cachedPlan) {
+ case join: CometBroadcastHashJoinExec =>
join.isNullAwareAntiJoin
+ case _ => false
+ }.isDefined,
+ s"the cache was not written by a native NOT IN (AQE
$adaptive):\n$cachedPlan")
+ assert(
+ !partitionsInOrder(spark.table("anti"), keyOrder(Ascending)),
+ "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
+ checkEachKeyGroupedOnce(query)
+ } finally {
+ spark.catalog.clearCache()
+ spark.catalog.dropTempView("anti")
+ spark.catalog.dropTempView("cached")
+ }
+ }
+ }
+ }
+ }
+ }
+ }
+
+ test("an outer hash join over a checkpoint of a NOT IN keeps each group
once") {
+ // A checkpoint is read by an RDD scan that reports the order of the plan
it was taken from,
+ // as a cache does, but that plan is gone, so nothing can say whether the
order holds. Here it
+ // is the order Spark reports for the null-aware anti join, which the rows
are not in.
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ 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",
+ CometConf.COMET_CONVERT_FROM_RDD_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)")
+ .localCheckpoint()
+ .createOrReplaceTempView("anti")
+ val plan = spark.table("anti").queryExecution.executedPlan
+ val scans = collect(plan) { case scan: RDDScanExec => scan }
+ assert(
+ scans.size == 1 && scans.head.outputOrdering.nonEmpty,
+ s"the checkpoint does not report the NOT IN's order:\n$plan")
+ assert(
+ !partitionsInOrder(spark.table("anti"), keyOrder(Ascending)),
+ "the checkpointed 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
+ checkEachKeyGroupedOnce(query)
+ } finally {
+ spark.catalog.clearCache()
+ spark.catalog.dropTempView("anti")
+ spark.catalog.dropTempView("cached")
+ }
+ }
+ }
+ }
+ }
+ }
+
+ test("an outer hash join over a cache of an outer hash join with no declared
order") {
Review Comment:
The test "an outer hash join over a cache of an outer hash join with no
declared order" shows that a cache written through a native sample and an outer
join reports k ASC while its rows are out of order. Spark's own SampleExec
reports no output ordering in 3.4, 3.5 and 4.1, so with Comet off that cache
reports none and a consumer keeps its sort. With the native cache on by default
from 3.5, a sortWithinPartitions over this cache would now come back unsorted.
A sort-merge join on the same partitioning would get no sort and could lose
matches, as in #6898. I haven't run that second case. Could
CometSampleExec.outputOrdering report what Spark's does? Another route is to
override maintains_input_order in the native SampleExec, whose comment already
says it preserves ordering. That lets keeps_jvm_input_order declare through it
and makes the cache's claim true, but a consumer then still hits the first
shape in #6922. Either way, this test would need another way to write an outer
join w
ith no declared order.
--
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]
