Skip to content

Commit d15d051

Browse files
andygroveandygrove
authored
perf: Improve query planning to more reliably fall back to columnar shuffle when native shuffle is not supported (apache#1209)
1 parent be48839 commit d15d051

File tree

356 files changed

+30950
-28502
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

356 files changed

+30950
-28502
lines changed

‎common/src/main/scala/org/apache/comet/CometConf.scala‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -254,6 +254,13 @@ object CometConf extends ShimCometConf {
254254
.checkValues(Set("native", "jvm", "auto"))
255255
.createWithDefault("auto")
256256

257+
val COMET_SHUFFLE_FALLBACK_TO_COLUMNAR: ConfigEntry[Boolean] =
258+
conf(s"$COMET_EXEC_CONFIG_PREFIX.shuffle.fallbackToColumnar")
259+
.doc("Whether to try falling back to columnar shuffle when native shuffle is not supported")
260+
.internal()
261+
.booleanConf
262+
.createWithDefault(false)
263+
257264
val COMET_EXEC_BROADCAST_FORCE_ENABLED: ConfigEntry[Boolean] =
258265
conf(s"$COMET_EXEC_CONFIG_PREFIX.broadcast.enabled")
259266
.doc(

‎spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala‎

Lines changed: 79 additions & 54 deletions
Original file line numberDiff line numberDiff line change
@@ -788,68 +788,93 @@ class CometSparkSessionExtensions
788788
}
789789

790790
// Native shuffle for Comet operators
791-
case s: ShuffleExchangeExec
792-
if isCometShuffleEnabled(conf) &&
793-
isCometNativeShuffleMode(conf) &&
794-
QueryPlanSerde.supportPartitioning(s.child.output, s.outputPartitioning)._1 =>
795-
logInfo("Comet extension enabled for Native Shuffle")
791+
case s: ShuffleExchangeExec =>
792+
val nativePrecondition = isCometShuffleEnabled(conf) &&
793+
isCometNativeShuffleMode(conf) &&
794+
QueryPlanSerde.supportPartitioning(s.child.output, s.outputPartitioning)._1
796795

797-
val newOp = transform1(s)
798-
newOp match {
799-
case Some(nativeOp) =>
800-
// Switch to use Decimal128 regardless of precision, since Arrow native execution
801-
// doesn't support Decimal32 and Decimal64 yet.
802-
conf.setConfString(CometConf.COMET_USE_DECIMAL_128.key, "true")
803-
val cometOp = CometShuffleExchangeExec(s, shuffleType = CometNativeShuffle)
804-
CometSinkPlaceHolder(nativeOp, s, cometOp)
805-
case None =>
806-
s
807-
}
796+
val nativeShuffle: Option[SparkPlan] =
797+
if (nativePrecondition) {
798+
val newOp = transform1(s)
799+
newOp match {
800+
case Some(nativeOp) =>
801+
// Switch to use Decimal128 regardless of precision, since Arrow native execution
802+
// doesn't support Decimal32 and Decimal64 yet.
803+
conf.setConfString(CometConf.COMET_USE_DECIMAL_128.key, "true")
804+
val cometOp = CometShuffleExchangeExec(s, shuffleType = CometNativeShuffle)
805+
Some(CometSinkPlaceHolder(nativeOp, s, cometOp))
806+
case None =>
807+
None
808+
}
809+
} else {
810+
None
811+
}
808812

809-
// Columnar shuffle for regular Spark operators (not Comet) and Comet operators
810-
// (if configured).
811-
// If the child of ShuffleExchangeExec is also a ShuffleExchangeExec, we should not
812-
// convert it to CometColumnarShuffle,
813-
case s: ShuffleExchangeExec
814-
if isCometShuffleEnabled(conf) && isCometJVMShuffleMode(conf) &&
813+
// this is a temporary workaround because some Spark SQL tests fail
814+
// when we enable COMET_SHUFFLE_FALLBACK_TO_COLUMNAR due to valid bugs
815+
// that we had not previously seen
816+
val tryColumnarNext =
817+
!nativePrecondition || (nativePrecondition && nativeShuffle.isEmpty &&
818+
COMET_SHUFFLE_FALLBACK_TO_COLUMNAR.get(conf))
819+
820+
val nativeOrColumnarShuffle = if (nativeShuffle.isDefined) {
821+
nativeShuffle
822+
} else if (tryColumnarNext) {
823+
// Columnar shuffle for regular Spark operators (not Comet) and Comet operators
824+
// (if configured).
825+
// If the child of ShuffleExchangeExec is also a ShuffleExchangeExec, we should not
826+
// convert it to CometColumnarShuffle,
827+
if (isCometShuffleEnabled(conf) && isCometJVMShuffleMode(conf) &&
815828
QueryPlanSerde.supportPartitioningTypes(s.child.output, s.outputPartitioning)._1 &&
816-
!isShuffleOperator(s.child) =>
817-
logInfo("Comet extension enabled for JVM Columnar Shuffle")
829+
!isShuffleOperator(s.child)) {
818830

819-
val newOp = QueryPlanSerde.operator2Proto(s)
820-
newOp match {
821-
case Some(nativeOp) =>
822-
s.child match {
823-
case n if n.isInstanceOf[CometNativeExec] || !n.supportsColumnar =>
824-
val cometOp = CometShuffleExchangeExec(s, shuffleType = CometColumnarShuffle)
825-
CometSinkPlaceHolder(nativeOp, s, cometOp)
826-
case _ =>
827-
s
831+
val newOp = QueryPlanSerde.operator2Proto(s)
832+
newOp match {
833+
case Some(nativeOp) =>
834+
s.child match {
835+
case n if n.isInstanceOf[CometNativeExec] || !n.supportsColumnar =>
836+
val cometOp =
837+
CometShuffleExchangeExec(s, shuffleType = CometColumnarShuffle)
838+
Some(CometSinkPlaceHolder(nativeOp, s, cometOp))
839+
case _ =>
840+
None
841+
}
842+
case None =>
843+
None
828844
}
829-
case None =>
830-
s
845+
} else {
846+
None
847+
}
848+
} else {
849+
None
831850
}
832851

833-
case s: ShuffleExchangeExec =>
834-
val isShuffleEnabled = isCometShuffleEnabled(conf)
835-
val outputPartitioning = s.outputPartitioning
836-
val reason = getCometShuffleNotEnabledReason(conf).getOrElse("no reason available")
837-
val msg1 = createMessage(!isShuffleEnabled, s"Comet shuffle is not enabled: $reason")
838-
val columnarShuffleEnabled = isCometJVMShuffleMode(conf)
839-
val msg2 = createMessage(
840-
isShuffleEnabled && !columnarShuffleEnabled && !QueryPlanSerde
841-
.supportPartitioning(s.child.output, outputPartitioning)
842-
._1,
843-
"Native shuffle: " +
844-
s"${QueryPlanSerde.supportPartitioning(s.child.output, outputPartitioning)._2}")
845-
val msg3 = createMessage(
846-
isShuffleEnabled && columnarShuffleEnabled && !QueryPlanSerde
852+
if (nativeOrColumnarShuffle.isDefined) {
853+
nativeOrColumnarShuffle.get
854+
} else {
855+
val isShuffleEnabled = isCometShuffleEnabled(conf)
856+
val outputPartitioning = s.outputPartitioning
857+
val reason = getCometShuffleNotEnabledReason(conf).getOrElse("no reason available")
858+
val msg1 = createMessage(!isShuffleEnabled, s"Comet shuffle is not enabled: $reason")
859+
val columnarShuffleEnabled = isCometJVMShuffleMode(conf)
860+
val msg2 = createMessage(
861+
isShuffleEnabled && !columnarShuffleEnabled && !QueryPlanSerde
862+
.supportPartitioning(s.child.output, outputPartitioning)
863+
._1,
864+
"Native shuffle: " +
865+
s"${QueryPlanSerde.supportPartitioning(s.child.output, outputPartitioning)._2}")
866+
val typeInfo = QueryPlanSerde
847867
.supportPartitioningTypes(s.child.output, outputPartitioning)
848-
._1,
849-
"JVM shuffle: " +
850-
s"${QueryPlanSerde.supportPartitioningTypes(s.child.output, outputPartitioning)._2}")
851-
withInfo(s, Seq(msg1, msg2, msg3).flatten.mkString(","))
852-
s
868+
._2
869+
val msg3 = createMessage(
870+
isShuffleEnabled && columnarShuffleEnabled && !QueryPlanSerde
871+
.supportPartitioningTypes(s.child.output, outputPartitioning)
872+
._1,
873+
"JVM shuffle: " +
874+
s"$typeInfo")
875+
withInfo(s, Seq(msg1, msg2, msg3).flatten.mkString(","))
876+
s
877+
}
853878

854879
case op =>
855880
op match {

‎spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4-spark3_5/q10/explain.txt‎

Lines changed: 63 additions & 59 deletions
Original file line numberDiff line numberDiff line change
@@ -1,49 +1,50 @@
11
== Physical Plan ==
2-
TakeOrderedAndProject (45)
3-
+- * HashAggregate (44)
4-
+- Exchange (43)
5-
+- * HashAggregate (42)
6-
+- * Project (41)
7-
+- * BroadcastHashJoin Inner BuildRight (40)
8-
:- * Project (35)
9-
: +- * BroadcastHashJoin Inner BuildRight (34)
10-
: :- * Project (28)
11-
: : +- * Filter (27)
12-
: : +- * BroadcastHashJoin ExistenceJoin(exists#1) BuildRight (26)
13-
: : :- * BroadcastHashJoin ExistenceJoin(exists#2) BuildRight (19)
14-
: : : :- * CometColumnarToRow (12)
15-
: : : : +- CometBroadcastHashJoin (11)
16-
: : : : :- CometFilter (2)
17-
: : : : : +- CometScan parquet spark_catalog.default.customer (1)
18-
: : : : +- CometBroadcastExchange (10)
19-
: : : : +- CometProject (9)
20-
: : : : +- CometBroadcastHashJoin (8)
21-
: : : : :- CometScan parquet spark_catalog.default.store_sales (3)
22-
: : : : +- CometBroadcastExchange (7)
23-
: : : : +- CometProject (6)
24-
: : : : +- CometFilter (5)
25-
: : : : +- CometScan parquet spark_catalog.default.date_dim (4)
26-
: : : +- BroadcastExchange (18)
27-
: : : +- * CometColumnarToRow (17)
28-
: : : +- CometProject (16)
29-
: : : +- CometBroadcastHashJoin (15)
30-
: : : :- CometScan parquet spark_catalog.default.web_sales (13)
31-
: : : +- ReusedExchange (14)
32-
: : +- BroadcastExchange (25)
33-
: : +- * CometColumnarToRow (24)
34-
: : +- CometProject (23)
35-
: : +- CometBroadcastHashJoin (22)
36-
: : :- CometScan parquet spark_catalog.default.catalog_sales (20)
37-
: : +- ReusedExchange (21)
38-
: +- BroadcastExchange (33)
39-
: +- * CometColumnarToRow (32)
40-
: +- CometProject (31)
41-
: +- CometFilter (30)
42-
: +- CometScan parquet spark_catalog.default.customer_address (29)
43-
+- BroadcastExchange (39)
44-
+- * CometColumnarToRow (38)
45-
+- CometFilter (37)
46-
+- CometScan parquet spark_catalog.default.customer_demographics (36)
2+
TakeOrderedAndProject (46)
3+
+- * HashAggregate (45)
4+
+- * CometColumnarToRow (44)
5+
+- CometColumnarExchange (43)
6+
+- * HashAggregate (42)
7+
+- * Project (41)
8+
+- * BroadcastHashJoin Inner BuildRight (40)
9+
:- * Project (35)
10+
: +- * BroadcastHashJoin Inner BuildRight (34)
11+
: :- * Project (28)
12+
: : +- * Filter (27)
13+
: : +- * BroadcastHashJoin ExistenceJoin(exists#1) BuildRight (26)
14+
: : :- * BroadcastHashJoin ExistenceJoin(exists#2) BuildRight (19)
15+
: : : :- * CometColumnarToRow (12)
16+
: : : : +- CometBroadcastHashJoin (11)
17+
: : : : :- CometFilter (2)
18+
: : : : : +- CometScan parquet spark_catalog.default.customer (1)
19+
: : : : +- CometBroadcastExchange (10)
20+
: : : : +- CometProject (9)
21+
: : : : +- CometBroadcastHashJoin (8)
22+
: : : : :- CometScan parquet spark_catalog.default.store_sales (3)
23+
: : : : +- CometBroadcastExchange (7)
24+
: : : : +- CometProject (6)
25+
: : : : +- CometFilter (5)
26+
: : : : +- CometScan parquet spark_catalog.default.date_dim (4)
27+
: : : +- BroadcastExchange (18)
28+
: : : +- * CometColumnarToRow (17)
29+
: : : +- CometProject (16)
30+
: : : +- CometBroadcastHashJoin (15)
31+
: : : :- CometScan parquet spark_catalog.default.web_sales (13)
32+
: : : +- ReusedExchange (14)
33+
: : +- BroadcastExchange (25)
34+
: : +- * CometColumnarToRow (24)
35+
: : +- CometProject (23)
36+
: : +- CometBroadcastHashJoin (22)
37+
: : :- CometScan parquet spark_catalog.default.catalog_sales (20)
38+
: : +- ReusedExchange (21)
39+
: +- BroadcastExchange (33)
40+
: +- * CometColumnarToRow (32)
41+
: +- CometProject (31)
42+
: +- CometFilter (30)
43+
: +- CometScan parquet spark_catalog.default.customer_address (29)
44+
+- BroadcastExchange (39)
45+
+- * CometColumnarToRow (38)
46+
+- CometFilter (37)
47+
+- CometScan parquet spark_catalog.default.customer_demographics (36)
4748

4849

4950
(1) CometScan parquet spark_catalog.default.customer
@@ -243,50 +244,53 @@ Functions [1]: [partial_count(1)]
243244
Aggregate Attributes [1]: [count#31]
244245
Results [9]: [cd_gender#23, cd_marital_status#24, cd_education_status#25, cd_purchase_estimate#26, cd_credit_rating#27, cd_dep_count#28, cd_dep_employed_count#29, cd_dep_college_count#30, count#32]
245246

246-
(43) Exchange
247+
(43) CometColumnarExchange
247248
Input [9]: [cd_gender#23, cd_marital_status#24, cd_education_status#25, cd_purchase_estimate#26, cd_credit_rating#27, cd_dep_count#28, cd_dep_employed_count#29, cd_dep_college_count#30, count#32]
248-
Arguments: hashpartitioning(cd_gender#23, cd_marital_status#24, cd_education_status#25, cd_purchase_estimate#26, cd_credit_rating#27, cd_dep_count#28, cd_dep_employed_count#29, cd_dep_college_count#30, 5), ENSURE_REQUIREMENTS, [plan_id=5]
249+
Arguments: hashpartitioning(cd_gender#23, cd_marital_status#24, cd_education_status#25, cd_purchase_estimate#26, cd_credit_rating#27, cd_dep_count#28, cd_dep_employed_count#29, cd_dep_college_count#30, 5), ENSURE_REQUIREMENTS, CometColumnarShuffle, [plan_id=5]
249250

250-
(44) HashAggregate [codegen id : 6]
251+
(44) CometColumnarToRow [codegen id : 6]
252+
Input [9]: [cd_gender#23, cd_marital_status#24, cd_education_status#25, cd_purchase_estimate#26, cd_credit_rating#27, cd_dep_count#28, cd_dep_employed_count#29, cd_dep_college_count#30, count#32]
253+
254+
(45) HashAggregate [codegen id : 6]
251255
Input [9]: [cd_gender#23, cd_marital_status#24, cd_education_status#25, cd_purchase_estimate#26, cd_credit_rating#27, cd_dep_count#28, cd_dep_employed_count#29, cd_dep_college_count#30, count#32]
252256
Keys [8]: [cd_gender#23, cd_marital_status#24, cd_education_status#25, cd_purchase_estimate#26, cd_credit_rating#27, cd_dep_count#28, cd_dep_employed_count#29, cd_dep_college_count#30]
253257
Functions [1]: [count(1)]
254258
Aggregate Attributes [1]: [count(1)#33]
255259
Results [14]: [cd_gender#23, cd_marital_status#24, cd_education_status#25, count(1)#33 AS cnt1#34, cd_purchase_estimate#26, count(1)#33 AS cnt2#35, cd_credit_rating#27, count(1)#33 AS cnt3#36, cd_dep_count#28, count(1)#33 AS cnt4#37, cd_dep_employed_count#29, count(1)#33 AS cnt5#38, cd_dep_college_count#30, count(1)#33 AS cnt6#39]
256260

257-
(45) TakeOrderedAndProject
261+
(46) TakeOrderedAndProject
258262
Input [14]: [cd_gender#23, cd_marital_status#24, cd_education_status#25, cnt1#34, cd_purchase_estimate#26, cnt2#35, cd_credit_rating#27, cnt3#36, cd_dep_count#28, cnt4#37, cd_dep_employed_count#29, cnt5#38, cd_dep_college_count#30, cnt6#39]
259263
Arguments: 100, [cd_gender#23 ASC NULLS FIRST, cd_marital_status#24 ASC NULLS FIRST, cd_education_status#25 ASC NULLS FIRST, cd_purchase_estimate#26 ASC NULLS FIRST, cd_credit_rating#27 ASC NULLS FIRST, cd_dep_count#28 ASC NULLS FIRST, cd_dep_employed_count#29 ASC NULLS FIRST, cd_dep_college_count#30 ASC NULLS FIRST], [cd_gender#23, cd_marital_status#24, cd_education_status#25, cnt1#34, cd_purchase_estimate#26, cnt2#35, cd_credit_rating#27, cnt3#36, cd_dep_count#28, cnt4#37, cd_dep_employed_count#29, cnt5#38, cd_dep_college_count#30, cnt6#39]
260264

261265
===== Subqueries =====
262266

263267
Subquery:1 Hosting operator id = 3 Hosting Expression = ss_sold_date_sk#7 IN dynamicpruning#8
264-
BroadcastExchange (50)
265-
+- * CometColumnarToRow (49)
266-
+- CometProject (48)
267-
+- CometFilter (47)
268-
+- CometScan parquet spark_catalog.default.date_dim (46)
268+
BroadcastExchange (51)
269+
+- * CometColumnarToRow (50)
270+
+- CometProject (49)
271+
+- CometFilter (48)
272+
+- CometScan parquet spark_catalog.default.date_dim (47)
269273

270274

271-
(46) CometScan parquet spark_catalog.default.date_dim
275+
(47) CometScan parquet spark_catalog.default.date_dim
272276
Output [3]: [d_date_sk#9, d_year#10, d_moy#11]
273277
Batched: true
274278
Location [not included in comparison]/{warehouse_dir}/date_dim]
275279
PushedFilters: [IsNotNull(d_year), IsNotNull(d_moy), EqualTo(d_year,2002), GreaterThanOrEqual(d_moy,1), LessThanOrEqual(d_moy,4), IsNotNull(d_date_sk)]
276280
ReadSchema: struct<d_date_sk:int,d_year:int,d_moy:int>
277281

278-
(47) CometFilter
282+
(48) CometFilter
279283
Input [3]: [d_date_sk#9, d_year#10, d_moy#11]
280284
Condition : (((((isnotnull(d_year#10) AND isnotnull(d_moy#11)) AND (d_year#10 = 2002)) AND (d_moy#11 >= 1)) AND (d_moy#11 <= 4)) AND isnotnull(d_date_sk#9))
281285

282-
(48) CometProject
286+
(49) CometProject
283287
Input [3]: [d_date_sk#9, d_year#10, d_moy#11]
284288
Arguments: [d_date_sk#9], [d_date_sk#9]
285289

286-
(49) CometColumnarToRow [codegen id : 1]
290+
(50) CometColumnarToRow [codegen id : 1]
287291
Input [1]: [d_date_sk#9]
288292

289-
(50) BroadcastExchange
293+
(51) BroadcastExchange
290294
Input [1]: [d_date_sk#9]
291295
Arguments: HashedRelationBroadcastMode(List(cast(input[0, int, true] as bigint)),false), [plan_id=6]
292296

0 commit comments

Comments
 (0)