andygrove commented on code in PR #6835:
URL: https://github.com/apache/datafusion-comet/pull/6835#discussion_r4285735290
##########
native/core/src/execution/planner.rs:
##########
@@ -3057,6 +3071,72 @@ 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.
+ ///
+ /// The serde sends an order only when the streamed side's Comet operators
report it, and it
+ /// checks that nothing under them moves rows (`mayReportFalseOrder`),
while this walk asks
+ /// the native operators the same thing through `maintains_input_order`.
The two agree as
+ /// long as a Comet operator that reports an order keeps it natively and
says so. One that
+ /// does not lands here with `plan` unchanged, and the join then moves its
unmatched rows
+ /// after Spark has dropped the sort above it. Nothing here can fall back
at that point, so
+ /// the invariant is kept at the operator: see "Adding a New Operator" in
the contributor
+ /// guide.
Review Comment:
Good catch, and it was one more shape. I wrote the test before the fix: a
broadcast left or right outer join over `spark.range(0, 2000, 1, 1)`, and again
over a negative step, followed by `sortWithinPartitions("id")`. It sent the
order and left no sort in the plan. With the fix removed, the partitions come
back out of order.
`LazyMemoryExec::add_ordering` doesn't work for this in DataFusion 55.2,
though. It adds the ordering to the equivalence properties but doesn't refresh
the cached `output_ordering` on `PlanProperties`, so `output_ordering()` stays
`None` and the join's `right_side_ordered` check doesn't see it. So the planner
wraps the range in a `SortedInputExec` with `range_sort_expr(step)`, which
builds its properties from scratch. `SortedInputExec` now reports its input's
metrics, because wrapping the root hid the range's output rows metric and
`CometRangeExecSuite` caught it. A native range now declares its order for
every operator above it, not only a join. Fixed in 65f406fe0f and a4eb244de0.
##########
spark/src/main/scala/org/apache/spark/sql/comet/operators.scala:
##########
@@ -2897,12 +2948,98 @@ 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 the order that `plan` reports may be false, because something
that feeds it with no
+ // sort in between cannot be relied on to keep it. Comet's hash joins report
the order of their
+ // streamed side, which Spark's own joins keep, but three Comet joins do not
always: the
+ // null-aware anti join, which Comet builds on its streamed side; an outer
join that runs as
+ // DataFusion's right outer join, which keeps its probe rows in place only
if the planner
+ // declared an order on its probe input; and the nested loop join, which
DataFusion runs build
+ // row by build row. So does a leaf that reports an order from state this
walk cannot check, an
+ // RDD checkpointed from a plan that is gone for one.
+ //
+ // The walk follows the order, so a hash join, which passes on only its
streamed side's, is
+ // walked through that side alone. A cache reports the order of the plan
that wrote it, and the
+ // walk goes on into that plan, picked as Spark's scan picks it: an adaptive
plan reports no
+ // order, but the scan reads the final plan behind it once there is one.
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 {
+ case adaptive: AdaptiveSparkPlanExec if adaptive.isFinalPlan =>
+ mayReportFalseOrder(adaptive.executedPlan)
+ case _ if plan.outputOrdering.isEmpty => false
+ case _: SortExec | _: CometSortExec => false
+ case _: CometBroadcastNestedLoopJoinExec => true
+ case join: CometBroadcastHashJoinExec if join.isNullAwareAntiJoin => true
+ case join: CometBroadcastHashJoinExec =>
+ hashJoinMayReportFalseOrder(
+ join.nativeOp,
+ join.joinType,
+ join.buildSide,
+ join.left,
+ join.right)
+ case join: CometHashJoinExec =>
+ hashJoinMayReportFalseOrder(
+ join.nativeOp,
+ join.joinType,
+ join.buildSide,
+ join.left,
+ join.right)
+ case join: HashJoin =>
+ mayReportFalseOrder(if (join.buildSide == BuildRight) join.left else
join.right)
+ case wrapper: CometScanWrapper => mayReportFalseOrder(wrapper.originalPlan)
+ case scan: CometInMemoryTableScanExec =>
+ mayReportFalseOrder(scan.originalPlan.relation.cachedPlan)
+ case scan: InMemoryTableScanExec =>
mayReportFalseOrder(scan.relation.cachedPlan)
+ case stage: QueryStageExec => mayReportFalseOrder(stage.plan)
+ // A scan reports an order that Spark already relies on.
+ case _: DataSourceScanExec | _: DataSourceV2ScanExecBase | _: RangeExec |
_: CometLeafExec =>
+ false
+ case _ if plan.children.isEmpty => true
+ case _ => plan.children.exists(mayReportFalseOrder)
+ }
+
+ // Whether a Comet hash join that reports an order may not keep it. Only the
streamed side's
+ // order reaches the join's output, so a false order on the build side is of
no concern, and an
+ // outer join of the two shapes has to have declared the order on its probe
input as well.
+ private def hashJoinMayReportFalseOrder(
+ nativeOp: Operator,
+ joinType: JoinType,
+ buildSide: BuildSide,
+ left: SparkPlan,
+ right: SparkPlan): Boolean = {
+ val (streamed, streamedIndex) = if (buildSide == BuildRight) (left, 0)
else (right, 1)
+ mayReportFalseOrder(streamed) ||
+ (runsAsRightOuterJoin(joinType, buildSide) &&
!declaresStreamedOrder(nativeOp, streamedIndex))
+ }
+
+ // Whether the native planner declares the streamed order on the probe input
of the join that
+ // `nativeOp` describes. It does so when the serde sent the order and the
probe input reads from
+ // the JVM, which is the case when the join's plan has a scan there. A
native operator over that
+ // scan is declared as well, as far as the native walk goes, but that is for
the planner to say,
+ // so it counts as not declared.
+ private def declaresStreamedOrder(nativeOp: Operator, streamedIndex: Int):
Boolean =
+ nativeOp.getHashJoin.getStreamedSortOrdersCount > 0 &&
+ nativeOp.getChildrenCount > streamedIndex && {
+ val probe = nativeOp.getChildren(streamedIndex)
+ probe.hasScan || probe.hasShuffleScan
+ }
+
+ // A left outer join built on the right and a right outer join built on the
left both run as
+ // DataFusion's right outer join, which keeps its unmatched probe rows in
place only when the
+ // probe input declares an order.
+ private def runsAsRightOuterJoin(joinType: JoinType, buildSide: BuildSide):
Boolean =
+ (joinType, buildSide) match {
+ case (LeftOuter, BuildRight) | (RightOuter, BuildLeft) => true
+ case _ => false
Review Comment:
Done in 9c26cd8990. The `case` is now an `if` on `runsAsRightOuterJoin`, so
the pair is written once, and a `streamedSide` helper returns the streamed plan
and its child index for the three places that picked it.
##########
spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala:
##########
@@ -3288,4 +3290,288 @@ 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] =
Review Comment:
Done in a4eb244de0. `partitionsInOrder`, `keyOrder`,
`checkEachKeyGroupedOnce` and `cachedPlanOf` are in `CometTestBase`.
`checkEachKeyGroupedOnce` takes an optional plan check, which is how the
native-cache NOT IN test now uses it. The cache suite's `sortedPartitions` is
gone.
For the fixture, a `withSortedCache` helper sets up the settings, `big`,
`small` and the sorted `cached` view once, and drops the views afterwards. Four
of the seven tests use it: NOT IN below the outer join, the semi join, the
sample and the declared outer join. The other three differ in ways it doesn't
cover: the cached NOT IN runs with AQE off and on and builds a second cache,
the checkpoint test has no cache, and the inner join test builds two caches.
That took about 60 lines out of `CometJoinSuite`.
--
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]
