andygrove commented on code in PR #6835:
URL: https://github.com/apache/datafusion-comet/pull/6835#discussion_r4288784010
##########
spark/src/main/scala/org/apache/spark/sql/comet/operators.scala:
##########
@@ -2885,6 +2887,56 @@ trait CometHashJoin {
val leftKeys = join.leftKeys.map(exprToProto(_, join.left.output))
val rightKeys = join.rightKeys.map(exprToProto(_, join.right.output))
+ // Spark reports the streamed side's order as the join's, and drops a sort
above the join
+ // that it satisfies. DataFusion's hash join keeps the unmatched probe
rows of a right outer
+ // join in place only when the probe input declares an order, and
otherwise moves them after
+ // the matched rows of each batch. Comet builds on the build side, so a
left outer join built
+ // on the right and a right outer join built on the left both run as that
right outer join,
+ // and an input sorted outside the native plan, a relation sorted before
it was cached for
+ // instance, declares no order to DataFusion. Send the order so the native
join declares it
+ // on the probe input; see SortedInputExec. Every other join type keeps
the probe order as
+ // it is, except the null-aware anti join: Comet builds that one on its
streamed side, so its
+ // output is not in the order Spark reports for it, and the planner
declares no order above
+ // it.
+ //
+ // The order Spark reports can be false to begin with. A Comet join that
may move its
+ // streamed rows reports it all the same: the null-aware anti join, a
right outer join like
+ // this one whose probe input declared no order, and the nested loop join,
which keeps the
+ // order of neither input. A cache reports the order of the plan that
wrote it, so the rows
+ // of such a join can reach this one through a cache, as a JVM input that
the native planner
+ // cannot look behind. Declaring the order there would let an operator
above rely on it, an
+ // aggregate in sorted mode for instance, which then splits groups. So
nothing is declared
+ // when such a join feeds the streamed side with no sort in between, or
when a leaf reports
+ // an order that cannot be checked; see mayReportFalseOrder.
+ //
+ // Declaring any order makes DataFusion keep the probe rows in place, so
the longest prefix
+ // of the order that Comet can express is enough. A key counts as
expressible when the native
+ // sort accepts its type, as `supportedSortType` decides for a sort: a
collated string is
+ // ordered by its collation in Spark but would be declared in raw byte
order. Nothing here can
+ // tell whether an operator above relies on the order, so a join whose
first key Comet cannot
+ // express stays in Spark.
+ val (streamed, _) = streamedSide(join.buildSide, join.left, join.right)
+ val streamedSortOrders =
+ if (runsAsRightOuterJoin(join.joinType, join.buildSide) &&
+ join.outputOrdering.nonEmpty && !mayReportFalseOrder(streamed)) {
+ val orders = join.outputOrdering.iterator
+ .takeWhile(order => supportedScalarSortElementType(order.dataType))
Review Comment:
You're right, and the collated string was the only case that check was for.
It is `!hasNonDefaultStringCollation(order.dataType)` now, so a collated string
at any depth still ends the prefix, `CometSortOrder` still declines a nested
key under a non-default null order, and arrays and structs go through
`exprToProto` again. The new test "outer hash joins keep a streamed order that
starts with an array or a struct" runs both outer shapes over a cache sorted by
an array or a struct column, and with the old check the join falls back and the
test fails. I also fixed the key-type test's comment. The test reads the sorted
rows from a cache rather than from a `sortWithinPartitions` right below the
join, because Spark's `EliminateSorts` drops a sort under a join. Fixed in
b50b2cdce6.
##########
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:
Added in c52fbb863b as "outer hash joins keep a streamed double order with
signed zeros and NaN". The cache is sorted by a double key holding -Infinity,
-1.5, -0.0, 0.0, 1.5, Infinity and NaN. For both outer shapes the test checks
that no sort is left and the partitions are in Spark's order, that the grouped
rows match Spark's with each key once, and that three top-k queries above the
join return Spark's rows in order. `ORDER BY k, v LIMIT 600` ends among the
zeros, `ORDER BY k DESC, v LIMIT 300` takes every NaN, and `ORDER BY k LIMIT
572` stops where the zeros start.
So float and double keys can stay in `streamedSortOrders`. The planner
builds the declared order through `create_sort_expr`, which wraps a float or
double key in `NormalizeNaNAndZero`, and floats nested in an array or a struct
in `NormalizeNestedFloats`, as it does for a native sort. The declared key ties
-0.0 with 0.0 and puts NaN above every other value, as Spark does. A top-k also
can't stop early on the declaration today. `CometTakeOrderedAndProjectExec`
runs its per-partition top-k as a native plan of its own over the join's
output, which arrives through a scan that declares nothing, and
`CometLocalTopKExec` only fuses with a native scan. When the join's reported
order already satisfies the top-k's keys, as with `ORDER BY k LIMIT n`, it
takes the first rows of each partition, which is right because the join keeps
them in order.
##########
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:
It did come back misordered. Spark reports a bucketed table's sort columns
as the scan's order only with
`spark.sql.legacy.bucketedTableScan.outputOrdering` on and one file per bucket.
With that set, a broadcast left outer join over a table written by one task
with `bucketBy(2, "k").sortBy("k")` and read by the native scan sent the order,
lost the sort above it, and returned unsorted partitions. Fixed in 204440adee.
The planner's walk now counts a scan of files with one file group as a source
whose order holds, because it reads that group's files one after the other. The
new test "outer hash joins keep the order of a sorted bucketed table read by
the native scan" covers both outer shapes, and a planner test covers the
one-file-group check.
The native range was the same shape, and 65f406fe0f fixed it after
comphead's comment by declaring the range's order at the source. The other
native leaves, `CometIcebergNativeScanExec` and `CometCsvNativeScanExec`,
report no order, so nothing from here goes into #6922. The contributor guide
now says that a leaf that reports an order has to be a source the planner knows
or declare its order itself, and the comment on `with_declared_order` says so
too. The operators.md sentence is softened in 939a05a934: the join keeps the
order where Comet can show that it holds, and the last paragraph points to
#6922 for the caches it still leaves undeclared.
--
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]
