Uh oh!
There was an error while loading. Please reload this page.
[SPARK-38679][CORE] Expose the number partitions in a stage to TaskContext - #35995
[SPARK-38679][CORE] Expose the number partitions in a stage to TaskContext#35995vkorukanti wants to merge 5 commits into
Conversation
c86b83e to
ecf7112Comparejiangxb1987
commented
Mar 29, 2022
It looks like barrier execution can also use this api to simplify the implementation of |
7354419 to
b48a670Comparevkorukanti
commented
Mar 29, 2022
Thank you @jiangxb1987 for reviewing. Could you please take another look at the change? Had to exclude the new method from binary compatibility test. |
zsxwing
commented
Mar 29, 2022
@cloud-fan@Ngone51 Although this one adds a new API, it's a pretty straightforward change. It looks pretty safe to me to backport into 3.3. What do you think? Also cc @MaxGekk |
MaxGekk
commented
Mar 29, 2022
Since this is either not a bug fix nor in the allow list https://lists.apache.org/thread/zrd7lcm5f5f3md7wffjy7x6w2pdmxxp7, we cannot just silently merge to branch-3.3. @zsxwing@vkorukanti Could you write an email to the thread in the dev list, and explain why we need this in 3.3 and cannot postpone to 3.4. |
zsxwing
commented
Mar 29, 2022
@MaxGekk Thanks for the feedback. We will not merge this to 3.3. |
There was a problem hiding this comment.
We can remove lazy val numTasks in this file and use numPartitions() directly.
There was a problem hiding this comment.
Updated in the latest commit.
Add a new api to expose total partition count in a task. so that the task knows what fraction of the computation is doing.
With this extra information, users can generate 32bit unique int ids as below rather than using `monotonically_increasing_id` which generates 64bit long ids.
```scala
val rdd = ...
rdd.mapPartitions { rowsIter =>
val partitionId = TaskContext.get().partitionId()
val numPartitions = TaskContext.get().numPartitions()
var i = 0
rowsIter.map { row =>
val rowId = partitionId + i * numPartitions
i += 1
(rowId, row)
}
}
```
Test: Added new unit tests to verify the number of partitions retrieved from TaskContext is expected.b48a670 to
c975242Comparecloud-fan
commented
Mar 31, 2022
thanks, merging to master! |
What changes were proposed in this pull request?
Add a new api to expose total partition count in the stage belonging to the task in TaskContext,
Why are the changes needed?
Add a new api to expose total partition count in the stage belonging to the task in TaskContext, so that the task knows what fraction of the computation is doing.
With this extra information, users can generate 32bit unique int ids as below rather than using
monotonically_increasing_idwhich generates 64bit long ids.rdd.mapPartitions { rowsIter =>valpartitionId=TaskContext.get().partitionId() valnumPartitions=TaskContext.get().numPartitions() vari=0 rowsIter.map { row =>valrowId= partitionId + i * numPartitions i +=1 (rowId, row) } }Does this PR introduce any user-facing change?
Yes. We add a new API
TaskContext.numPartitions.How was this patch tested?
Added new unit tests to verify the number of partitions retrieved from TaskContext is expected.