-
Notifications
You must be signed in to change notification settings - Fork 232
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Expose row count statistics in GpuShuffleExchangeExec #1855
Changes from 4 commits
92815db
ee28ce4
0688a90
0c5ee37
f9dc37b
aa040f1
e59946b
71a62b1
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -97,6 +97,24 @@ class AdaptiveQueryExecSuite | |
collectWithSubqueries(plan)(ShimLoader.getSparkShims.reusedExchangeExecPfn) | ||
} | ||
|
||
test("get row counts from executed shuffle query stages") { | ||
assumeSpark301orLater | ||
|
||
skewJoinTest { spark => | ||
val (_, innerAdaptivePlan) = runAdaptiveAndVerifyResult( | ||
spark, | ||
"SELECT * FROM skewData1 join skewData2 ON key1 = key2") | ||
val innerSmj = findTopLevelGpuShuffleHashJoin(innerAdaptivePlan) | ||
val shuffleExchanges = ShimLoader.getSparkShims | ||
.findOperators(innerAdaptivePlan, _.isInstanceOf[ShuffleQueryStageExec]) | ||
.map(_.asInstanceOf[ShuffleQueryStageExec]) | ||
assert(shuffleExchanges.nonEmpty) | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. nit: maybe add a message to the assert for a more descriptive exception There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. We haven't generally been using assertion messages. I have updated the assertion to make it more specific though. |
||
val shim = ShimLoader.getSparkShims | ||
val stats = shuffleExchanges.map(e => shim.getQueryStageRuntimeStatistics(e)) | ||
assert(stats.forall(_.rowCount.contains(1000))) | ||
} | ||
} | ||
|
||
test("skewed inner join optimization") { | ||
skewJoinTest { spark => | ||
val (_, innerAdaptivePlan) = runAdaptiveAndVerifyResult( | ||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
nit:
Statistics(0)
since None is the default for rowcountThere was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Fixed