parthchandra commented on code in PR #6835:
URL: https://github.com/apache/datafusion-comet/pull/6835#discussion_r4265887756
##########
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:
[question] When the serde sends a non-empty `streamed_sort_orders` but the
probe neither already declares an order nor passes `keeps_jvm_input_order`,
this returns the probe unchanged with no declaration and no fallback. Spark has
already dropped the sort above the join by then, so if that ever happened the
output would be silently misordered rather than falling back. The serde
`mayReportFalseOrder` and this walk encode the same policy twice over two
representations. I couldn't construct a case where they disagree in this
direction, so this is about robustness - could a short `debug_assert` or a
comment record the invariant that the serde never sends an order this walk
can't place?
##########
native/core/src/execution/operators/sorted_input.rs:
##########
@@ -0,0 +1,124 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! An input that declares the order its rows arrive in.
+
+use std::fmt::Formatter;
+use std::sync::Arc;
+
+use datafusion::common::tree_node::TreeNodeRecursion;
+use datafusion::common::{internal_err, Result};
+use datafusion::execution::TaskContext;
+use datafusion::physical_expr::{EquivalenceProperties, LexOrdering,
PhysicalExpr};
+use datafusion::physical_plan::execution_plan::CardinalityEffect;
+use datafusion::physical_plan::{
+ apply_expression_roots, DisplayAs, DisplayFormatType, ExecutionPlan,
PlanProperties,
+ SendableRecordBatchStream,
+};
+
+/// Passes `input` through unchanged and declares to DataFusion the order its
rows arrive in.
+///
+/// Spark's plan knows the order of an input that enters the native plan from
the JVM, a relation
+/// sorted before it was cached for instance, but the `ScanExec` it arrives
through declares none,
+/// and DataFusion's hash join keeps the unmatched probe rows of a right outer
join in place only
+/// when the probe input declares an order. Otherwise it moves them after the
matched rows of each
+/// batch, which loses an order Spark has already dropped a sort for. The
planner puts this only
+/// on that join's probe input, with the order Spark reports for it, and only
when the rows reach
+/// that input from the JVM through operators that keep their order. The join
keeps the order in
+/// its own output, so the operators above it in the native plan see it too:
an aggregate grouped
+/// by the key runs in sorted, streaming mode, a sort aggregate takes the
order as already
+/// satisfied, and a top-k sort whose keys start with it stops reading early.
A wrong order would
+/// therefore give wrong groups or wrong rows, not only rows out of order.
+#[derive(Debug)]
+pub struct SortedInputExec {
+ input: Arc<dyn ExecutionPlan>,
+ ordering: LexOrdering,
+ cache: Arc<PlanProperties>,
+}
+
+impl SortedInputExec {
+ pub fn new(input: Arc<dyn ExecutionPlan>, ordering: LexOrdering) -> Self {
+ let input_properties = input.properties();
+ let cache = Arc::new(PlanProperties::new(
+ EquivalenceProperties::new_with_orderings(input.schema(),
[ordering.clone()]),
Review Comment:
[nit] `EquivalenceProperties::new_with_orderings(input.schema(),
[ordering])` builds equivalence properties from scratch and drops the input's
existing ones (constants from filters, equivalence classes from join keys). For
a node that only adds an ordering claim, starting from
`input.equivalence_properties().clone()` and adding the ordering would keep
that information available to operators above. Not a correctness problem, just
lost optimization context.
##########
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:
[question] The streamed order is serialized with `exprToProto` only, but the
native Sort serde first gates the same kind of keys through `supportedSortType`
(`QueryPlanSerde.scala:1339`), which rejects non-default string collation and
unsupported sort-key types. A sort key that is a collated string and is not the
join key gets past the join-key collation check at line 2487, serializes fine
(`supportedDataType` allows any string by default), and then gets declared with
DataFusion's raw-byte ordering. It looks safe today only because every operator
that would compare the key falls back for collated strings on its own (sorted
aggregate grouping around line 1982, native sort via `supportedSortType`), so
the join output is only ever placed by probe row position. But that safety is
accidental. Could `streamedSortOrders` reuse `supportedSortType` here so the
intent is explicit and this doesn't break if a future change lets a sorted
aggregate accept a collated group key?
##########
spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala:
##########
@@ -3288,4 +3290,262 @@ 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] =
+ df.queryExecution.toRdd
+ .mapPartitions { it =>
+ val keys = it.map(_.getInt(0)).toArray
+ Iterator(keys.sameElements(keys.sorted))
+ }
+ .collect()
+
+ test("a broadcast outer hash join over a sorted cache keeps its order with
default settings") {
+ // A cache written with no exchange in its plan is not adaptive, so with
AQE on, as it is by
+ // default, the scan reports the cached plan's order just as it does with
AQE off. Only the
+ // native cache scan is switched on here, and the broadcast join is the
one Spark plans for
+ // a small build side.
+ withSQLConf(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").sortWithinPartitions("_1").cache()
+ streamed.count()
+ val build = spark.table("small")
+ val df = streamed
+ .join(build, streamed("_1") === build("_1"), "left_outer")
+ .select(streamed("_1").as("k"), streamed("_2").as("v"),
build("_2").as("w"))
+ .sortWithinPartitions("k")
+ val expected = (0 until 10000).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(_))
+ // The plan AQE chose is final once the query has run.
+ val plan = df.queryExecution.executedPlan
+ assert(plan.isInstanceOf[AdaptiveSparkPlanExec], plan)
+ val joins = collect(plan) { case j: CometBroadcastHashJoinExec =>
j }
+ assert(joins.size == 1 && joins.head.joinType == LeftOuter, plan)
+ assert(joins.head.nativeOp.getHashJoin.getStreamedSortOrdersCount
> 0, plan)
+ assert(collect(plan) { case s @ (_: CometSortExec | _: SortExec)
=> s }.isEmpty, plan)
+ assert(collect(plan) { case s: CometInMemoryTableScanExec => s
}.size == 1, plan)
+ assert(sortedPartitions(df).forall(identity), plan)
+ } finally {
+ spark.catalog.clearCache()
+ }
+ }
+ }
+ }
+ }
+
+ test("an outer hash join over a native cache scan of a NOT IN keeps each
group once") {
+ // The native cache scan reports the order of the plan that wrote the
cache. A NOT IN is not in
+ // the order Spark reports for it, since Comet builds it on its streamed
side, so the outer
+ // join that reads the cache must not declare the order, or the aggregates
above it run in
+ // sorted mode over unsorted rows and split groups. Sparse keys and small
batches make each
+ // group span several. A cache with an exchange in its plan is written
under AQE, and the scan
+ // reads the final plan behind the adaptive one, 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()
+ assert(
+ !sortedPartitions(spark.table("anti")).forall(identity),
+ "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
+ def rows(): Seq[(Int, Long, Long, Long)] =
+ sql(query)
+ .collect()
+ .map(r => (r.getInt(0), r.getLong(1), r.getLong(2),
r.getLong(3)))
+ .toSeq
+ .sorted
+ val cometRows = rows()
+ val plan = sql(query).queryExecution.executedPlan
+ assert(
+ collect(plan) { case s: CometInMemoryTableScanExec => s
}.size == 1,
+ s"the outer join does not read the native cache scan (AQE
$adaptive):\n$plan")
+ // withSQLConf returns Unit on Spark 3.x, so Spark's rows are
assigned inside it.
+ var sparkRows = Seq.empty[(Int, Long, Long, Long)]
+ withSQLConf(CometConf.COMET_ENABLED.key -> "false") {
+ sparkRows = rows()
+ }
+ assert(
+ cometRows.map(_._1).distinct.size == cometRows.size,
+ "a key is grouped twice")
+ assert(cometRows == sparkRows)
+ } finally {
+ spark.catalog.clearCache()
+ spark.catalog.dropTempView("anti")
+ spark.catalog.dropTempView("cached")
+ }
+ }
+ }
+ }
+ }
+ }
+ }
+
+ test("an aggregate on the key above a hash join over a sorted cache groups
each key once") {
+ // The order the join declares on its probe input reaches past the join,
so the aggregate
+ // above it in the same native plan groups in sorted, streaming mode. That
is right while
+ // Spark's reported order holds, and a wrong order would split a key
across groups, which only
+ // a comparison of the rows as a multiset catches. Small batches make each
group span several.
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ 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((0 until 40 by 2).map(m => (m * 1000, m)), "small") {
+ try {
+ val streamed =
+ spark.table("big").repartition(2,
$"_1").sortWithinPartitions("_1").cache()
+ streamed.count()
+ val build = spark.table("small")
+
+ // Fifty rows per key, and a match in the build side for every
even key.
+ val expected = (0 until 40).map { m =>
+ Row(m * 1000, 50L, 50L * m + 49000L, if (m % 2 == 0) 50L else 0L)
+ }
+
+ for {
+ (hint, joinClass) <- Seq(
+ "shuffle_hash" -> classOf[CometHashJoinExec],
+ "broadcast" -> classOf[CometBroadcastHashJoinExec])
+ (joinType, buildSide) <- Seq("left_outer" -> BuildRight,
"right_outer" -> BuildLeft)
+ } {
+ 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)
+ }
+ val df = joined
+ .groupBy(streamed("_1"))
+ .agg(count(lit(1)), sum(streamed("_2")), count(build("_2")))
+ val plan = df.queryExecution.executedPlan
+ val joins: Seq[(SparkPlan, BuildSide)] = collect(plan) {
+ case j: CometHashJoinExec => (j, j.buildSide)
+ case j: CometBroadcastHashJoinExec => (j, j.buildSide)
+ }
+ assert(
+ joins.map { case (j, side) => (j.getClass, side) } ==
Seq((joinClass, buildSide)),
+ plan)
+ // A final aggregate after a shuffle would merge a key a partial
aggregate split, so
+ // both aggregates have to sit above the join with no exchange
in between.
+ def aboveJoin(p: SparkPlan): Boolean = p.find(_ eq
joins.head._1).isDefined
+ assert(
+ collect(plan) { case a: CometHashAggregateExec => a
}.count(aboveJoin) == 2,
Review Comment:
[question] This "aggregate on the key above a hash join" test checks the
join class and build side, two hash aggregates above the join, no exchange, and
the rows - but none of those actually require the streamed-order declaration.
The no-exchange property comes from the pre-cache `repartition(2, $"_1")`, not
from the order, and the rows are correct in plain hash-aggregate mode too, so
this test passes with or without the order being declared. The other tests
assert `nativeOp.getHashJoin.getStreamedSortOrdersCount > 0`. Could this one do
the same (and/or assert the partitions are sorted) so it proves the
sorted-order path is what's being exercised?
##########
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 =
Review Comment:
[test coverage] Every test here keys and orders on an `Int`.
`SortedInputExec` passes rows through and never compares values, so the order
is physically Spark's and only the aggregate's group-boundary detection
compares the key, which lowers the risk, but the serialization and the
sorted-aggregate grouping are still untested for non-integer keys. Could you
add a decimal key (with boundary values), a string key (including non-ASCII /
Asian characters), and a timestamp key under a non-UTC session timezone (plus a
pre-1970 value) to the parametrized tests? That would exercise `exprToProto` /
`create_sort_expr` of the `SortOrder` and the prefix logic across types and
confirm the sorted aggregate groups non-integer keys the same as Spark.
Complex-type keys (struct/array/map) are unreachable because `exprToProto`
rejects them by default, so those are out of scope - worth a one-line test
comment saying why.
--
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]
