kazuyukitanimura commented on code in PR #6835:
URL: https://github.com/apache/datafusion-comet/pull/6835#discussion_r4286125701
##########
spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala:
##########
@@ -2284,4 +2286,678 @@ class CometJoinSuite extends CometTestBase {
}
}
}
+
+ 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()
+ }
+ }
+ }
+ }
+ }
+
+ // Compares the rows of a fresh `df` as a multiset with the ones Spark
returns for it. The rows
+ // of a key tie under the sort, so their order within it differs between the
two.
+ private def sameRowsAsSpark(df: => DataFrame): Unit = {
+ // withSQLConf returns Unit on Spark 3.x, so Spark's rows are assigned
inside it.
+ var sparkRows = Seq.empty[Row]
+ withSQLConf(CometConf.COMET_ENABLED.key -> "false") {
+ sparkRows = df.collect().toSeq
+ }
+ QueryTest.sameRows(sparkRows, df.collect().toSeq).foreach(fail(_))
+ }
+
+ test("outer hash joins keep the order of a native range") {
+ // Spark's RangeExec reports id ascending, or descending for a negative
step, and the sort above
+ // the join is dropped on that basis. The range is a native leaf, not a
JVM input, so the
+ // planner has to declare its order itself or the join moves its unmatched
rows. A broadcast
+ // join, since a shuffled one would put an exchange between the range and
the join.
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ CometConf.COMET_EXEC_RANGE_ENABLED.key -> "true") {
+ withParquetTable((0 until 40).map(k => (k.toLong, k)), "small") {
+ val build = spark.table("small")
+ for {
+ (start, end, step) <- Seq((0L, 2000L, 1L), (1999L, -1L, -1L))
+ joinType <- Seq("left_outer", "right_outer")
+ } {
+ val streamed = spark.range(start, end, step, 1)
+ val joined = if (joinType == "left_outer") {
+ streamed.join(build.hint("broadcast"), streamed("id") ===
build("_1"), joinType)
+ } else {
+ build.hint("broadcast").join(streamed, streamed("id") ===
build("_1"), joinType)
+ }
+ val direction = if (step > 0) Ascending else Descending
+ val df = joined
+ .select(streamed("id").as("k"), build("_2").as("w"))
+ .sortWithinPartitions(if (step > 0) $"k" else $"k".desc)
+ val plan = df.queryExecution.executedPlan
+ val context = s"$joinType step $step\n$plan"
+ assert(collect(plan) { case r: CometRangeExec => r }.size == 1,
context)
+ val joins = collect(plan) { case j: CometBroadcastHashJoinExec =>
j.nativeOp }
+ assert(joins.size == 1, context)
+ assert(joins.head.getHashJoin.getStreamedSortOrdersCount > 0,
context)
+ assert(
+ collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s
}.isEmpty,
+ context)
+ val expected =
+ (start until end by step).map(id => Row(id, if (id < 40) id.toInt
else null))
+ QueryTest.sameRows(expected, df.collect().toSeq).foreach(fail(_))
+ assert(
+ partitionsInOrder(
+ df,
+ SortOrder(BoundReference(0, LongType, nullable = true),
direction)),
+ context)
+ }
+ }
+ }
+ }
+
+ test("outer hash joins keep the streamed order for decimal, string and
timestamp keys") {
Review Comment:
Could we add a DoubleType key with -0.0, 0.0 and NaN? Spark treats -0.0 and
0.0 as equal when sorting, but DataFusion's top-k uses a total order where -0.0
< 0.0.
Would a native top-k above the join, such as ORDER BY k, v LIMIT n, stop
early on a declared [k ASC] and differ from Spark? If it does, should float and
double keys
be left out of streamedSortOrders, or should the join fall back for them?
##########
spark/src/main/scala/org/apache/spark/sql/comet/operators.scala:
##########
@@ -2876,12 +2921,44 @@ trait CometHashJoin {
else OperatorOuterClass.BuildSide.BuildRight)
.setNullAwareAntiJoin(isNullAwareAntiJoin)
.setDynamicFilterEnabled(CometConf.COMET_EXEC_JOIN_DYNAMIC_FILTER_ENABLED.get(join.conf))
+ .addAllStreamedSortOrders(streamedSortOrders.asJava)
condition.foreach(joinBuilder.setCondition)
Some(builder.setHashJoin(joinBuilder).build())
} else {
None
}
}
+
+ // Whether a Comet join that may move its streamed rows, while reporting
Spark's order for
+ // them, feeds the order that `plan` reports with no sort in between.
Whether an outer join
+ // declared the order on its probe input is up to the native planner, so
every one built that
+ // way counts. A cache reports the order of the plan that wrote it, so the
walk goes on into
+ // that plan. While CometExecRule converts a plan, a native plan's JVM
inputs sit behind a
+ // CometScanWrapper, a leaf, so the walk looks through that too.
+ private def mayReportFalseOrder(plan: SparkPlan): Boolean = plan match {
Review Comment:
Thanks for the careful walk here. The comment on with_declared_order says
the serde and the planner agree whenever an operator that reports an order
keeps it
natively. CometNativeScanExec reports originalPlan.outputOrdering and is a
CometLeafExec, so this returns false and the serde sends the order. On the
native side the
probe is a Parquet source, so keeps_jvm_input_order declares nothing and
there is no fallback. Does a sorted bucketed table with one file per bucket, as
the
streamed side of a left outer join followed by sortWithinPartitions, come
back misordered? comphead's native range comment is the same shape. This
matches main, but
#6922 doesn't list either one. Could you fix them or add both to #6922,
and soften the "Comet's native hash joins keep that order" sentence in
operators.md?
--
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]
