andygrove commented on code in PR #6835:
URL: https://github.com/apache/datafusion-comet/pull/6835#discussion_r4285108267
##########
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:
Took the first route. `CometSampleExec` no longer reports its child's order,
so it matches Spark's `SampleExec`, which overrides only `outputPartitioning`
in 3.4.3, 3.5.9, 4.0.4, 4.1.3 and 4.2.0. A join over a sample now sends no
order, the cache it writes reports none, and the sort above that cache stays. I
left the native `SampleExec` alone, since reporting `maintains_input_order`
there would only make a claim that Spark's own operator doesn't.
The test is rewritten to pin that: the cached join sent no order, the cache
reports none, and `sortWithinPartitions` over it returns sorted partitions.
With the sample gone I haven't found another shape where the serde sends an
order and the planner declares nothing, so no test reaches that check now, and
it stays as a guard. The `adding_a_new_operator.md` section says what an
operator that reports an order has to do. Fixed in a2ad967e98 and ddce25a6fe.
##########
spark/src/main/scala/org/apache/spark/sql/comet/operators.scala:
##########
@@ -2864,6 +2866,51 @@ 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. 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 = if (join.buildSide == BuildRight) join.left else join.right
+ val streamedSortOrders = (join.joinType, join.buildSide) match {
+ case (LeftOuter, BuildRight) | (RightOuter, BuildLeft)
+ if join.outputOrdering.nonEmpty && !mayReportFalseOrder(streamed) =>
+ val orders = join.outputOrdering.iterator
Review Comment:
Done in a2ad967e98, with one difference. `supportedSortType` tags the join
with a fallback reason as a side effect, and here a key it rejects only
shortens the prefix, so `doConvert` asks `supportedScalarSortElementType` about
each key, the check `supportedSortType` uses for a single column. It rejects a
collated string and any type the native sort does not take, and a join whose
first key is rejected falls back through the existing path.
##########
native/core/src/execution/planner.rs:
##########
@@ -2937,6 +2951,63 @@ impl PhysicalPlanner {
Ok(Arc::new(prepared))
}
+ /// `plan`, declaring `sort_orders` as the order its rows arrive in,
unless it already
+ /// declares one. Spark's plan knows the order of an input that enters the
native plan from
+ /// the JVM, but the `ScanExec` it arrives through declares none. Only
that order is declared:
+ /// when an operator in between may have moved rows, `plan` is left as it
is.
+ fn with_declared_order(
+ &self,
+ plan: Arc<dyn ExecutionPlan>,
+ sort_orders: &[Expr],
+ ) -> Result<Arc<dyn ExecutionPlan>, ExecutionError> {
+ if plan.output_ordering().is_some() ||
!Self::keeps_jvm_input_order(&plan) {
Review Comment:
I put the invariant in a comment on `with_declared_order` and in the new
"Reporting an Output Ordering" section of `adding_a_new_operator.md`, and left
out the `debug_assert`. The shape viirya found, the native sample, is fixed,
but I can't show that no other operator reports an order without saying it
keeps it, and I'd rather not turn a shape nobody has hit yet into a panic.
Would you rather have the `debug_assert` anyway, so that such an operator fails
a test early? Fixed in a2ad967e98.
--
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]
