Skip to content

Commit b50b2cd

Browse files
andygroveandygrove
committed
fix: send a streamed order that starts with an array or a struct
The check on the streamed order's keys took only the scalar types that the native sort accepts for a single column, so a join whose streamed order starts with an array or a struct went back to Spark. exprToProto takes those keys, and declines a nested key under a non-default null order itself. The check was there for the collated string, so it now rejects a key that holds one, at the top level or nested, as supportedSortType does for a sort.
1 parent c52fbb8 commit b50b2cd

2 files changed

Lines changed: 73 additions & 10 deletions

File tree

‎spark/src/main/scala/org/apache/spark/sql/comet/operators.scala‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -64,7 +64,7 @@ import org.apache.comet.rules.CometExecRule
6464
import org.apache.comet.serde.{CometCollectBuffer, CometOperatorSerde, Compatible, OperatorOuterClass, QueryContextInterner, SupportLevel, Unsupported}
6565
import org.apache.comet.serde.OperatorOuterClass.{AggregateMode => CometAggregateMode, Operator}
6666
import org.apache.comet.serde.QueryPlanSerde
67-
import org.apache.comet.serde.QueryPlanSerde.{aggExprToProto, exprToProto, isStringCollationType, supportedScalarSortElementType, supportedSortType}
67+
import org.apache.comet.serde.QueryPlanSerde.{aggExprToProto, exprToProto, hasNonDefaultStringCollation, isStringCollationType, supportedSortType}
6868
import org.apache.comet.serde.operator.CometSink
6969
import org.apache.comet.shims.MergeRowsMetricsShim
7070

@@ -2910,17 +2910,17 @@ trait CometHashJoin {
29102910
// an order that cannot be checked; see mayReportFalseOrder.
29112911
//
29122912
// Declaring any order makes DataFusion keep the probe rows in place, so the longest prefix
2913-
// of the order that Comet can express is enough. A key counts as expressible when the native
2914-
// sort accepts its type, as `supportedSortType` decides for a sort: a collated string is
2915-
// ordered by its collation in Spark but would be declared in raw byte order. Nothing here can
2916-
// tell whether an operator above relies on the order, so a join whose first key Comet cannot
2917-
// express stays in Spark.
2913+
// of the order that Comet can express is enough. A key that exprToProto takes counts as
2914+
// expressible, unless it holds a collated string, at the top level or nested, which Spark
2915+
// orders by its collation but the native side would declare in raw byte order. That is the
2916+
// check `supportedSortType` makes for a sort. Nothing here can tell whether an operator above
2917+
// relies on the order, so a join whose first key Comet cannot express stays in Spark.
29182918
val (streamed, _) = streamedSide(join.buildSide, join.left, join.right)
29192919
val streamedSortOrders =
29202920
if (runsAsRightOuterJoin(join.joinType, join.buildSide) &&
29212921
join.outputOrdering.nonEmpty && !mayReportFalseOrder(streamed)) {
29222922
val orders = join.outputOrdering.iterator
2923-
.takeWhile(order => supportedScalarSortElementType(order.dataType))
2923+
.takeWhile(order => !hasNonDefaultStringCollation(order.dataType))
29242924
.map(exprToProto(_, streamed.output))
29252925
.takeWhile(_.isDefined)
29262926
.flatten

‎spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala‎

Lines changed: 66 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,7 @@ import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec
4242
import org.apache.spark.sql.execution.datasources.SchemaColumnConvertNotSupportedException
4343
import org.apache.spark.sql.execution.exchange.{ReusedExchangeExec, ShuffleExchangeLike}
4444
import org.apache.spark.sql.execution.joins.{ShuffledHashJoinExec, SortMergeJoinExec}
45-
import org.apache.spark.sql.functions.{count, lit, sum}
45+
import org.apache.spark.sql.functions.{array, count, lit, struct, sum}
4646
import org.apache.spark.sql.internal.SQLConf
4747
import org.apache.spark.sql.types.{ArrayType, DoubleType, IntegerType, LongType, MetadataBuilder, StructField, StructType}
4848

@@ -2526,8 +2526,8 @@ class CometJoinSuite extends CometTestBase {
25262526
// SortedInputExec passes the rows through, so the order is physically Spark's, but the serde
25272527
// and the native sort expression handle the key, and the sorted aggregate above the join finds
25282528
// its group boundaries by comparing it. The timestamps run in a session time zone that is not
2529-
// UTC and include values before 1970. Struct, array and map keys are out of scope, since
2530-
// exprToProto rejects them by default and such a join stays in Spark.
2529+
// UTC and include values before 1970. Double, array and struct keys have tests of their own,
2530+
// and Spark cannot sort by a map.
25312531
val keyTypes = Seq(
25322532
"decimal(20,2)" -> Seq(
25332533
"-999999999999999999.99",
@@ -2682,6 +2682,69 @@ class CometJoinSuite extends CometTestBase {
26822682
}
26832683
}
26842684

2685+
test("outer hash joins keep a streamed order that starts with an array or a struct") {
2686+
// exprToProto takes an array or a struct sort key, so a join whose streamed order starts with
2687+
// one sends it, stays native, and the planner declares it on the JVM input. Spark drops a sort
2688+
// right below a join, so the sorted rows come from a cache.
2689+
withSQLConf(
2690+
SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
2691+
SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
2692+
CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "false",
2693+
CometConf.COMET_CONVERT_FROM_IN_MEMORY_CACHE_ENABLED.key -> "true") {
2694+
withParquetTable((0 until 2000).map(i => (i % 40, i)), "big") {
2695+
withParquetTable((0 until 40 by 2).map(k => (k, k / 2)), "small") {
2696+
val build = spark.table("small")
2697+
// An array of strings, the one array type that Comet reads from Spark's cache.
2698+
Seq(
2699+
"array" -> array(($"_1" % 7).cast("string"), $"_1".cast("string")),
2700+
"struct" -> struct($"_1" % 7, $"_1")).foreach { case (keyName, key) =>
2701+
try {
2702+
val streamed = spark
2703+
.table("big")
2704+
.select(key.as("s"), $"_1", $"_2")
2705+
.repartition(2, $"_1")
2706+
.sortWithinPartitions("s")
2707+
.cache()
2708+
streamed.count()
2709+
val order = SortOrder(
2710+
BoundReference(0, streamed.schema("s").dataType, nullable = true),
2711+
Ascending)
2712+
Seq("left_outer", "right_outer").foreach { joinType =>
2713+
def joined: DataFrame = {
2714+
val hinted = build.hint("broadcast")
2715+
if (joinType == "left_outer") {
2716+
streamed.join(hinted, streamed("_1") === build("_1"), joinType)
2717+
} else {
2718+
hinted.join(streamed, build("_1") === streamed("_1"), joinType)
2719+
}
2720+
}
2721+
def df: DataFrame = joined
2722+
.select(streamed("s"), streamed("_2").as("v"), build("_2").as("w"))
2723+
.sortWithinPartitions("s")
2724+
val plan = df.queryExecution.executedPlan
2725+
val context = s"$keyName $joinType\n$plan"
2726+
val joins = collect(plan) { case j: CometBroadcastHashJoinExec => j.nativeOp }
2727+
assert(joins.size == 1, context)
2728+
assert(joins.head.getHashJoin.getStreamedSortOrdersCount > 0, context)
2729+
assert(
2730+
collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s }.isEmpty,
2731+
context)
2732+
sameRowsAsSpark(df)
2733+
assert(partitionsInOrder(df, order), context)
2734+
sameRowsAsSpark(
2735+
joined
2736+
.groupBy(streamed("s"))
2737+
.agg(count(lit(1)), sum(streamed("_2")), count(build("_2"))))
2738+
}
2739+
} finally {
2740+
spark.catalog.clearCache()
2741+
}
2742+
}
2743+
}
2744+
}
2745+
}
2746+
}
2747+
26852748
// The setup several cache tests share: `big(_1, _2)` with 2000 rows over 40 sparse keys and
26862749
// `small` with a row for every even key, and `big` as `cached(k, v)`, partitioned on the key,
26872750
// sorted within each partition and cached, with 8-row batches and the native cache on. `views`

0 commit comments

Comments
 (0)