fix: Use RDD partition index

[viirya]

Which issue does this PR close?

Closes #1113.

Rationale for this change

What changes are included in this PR?

How are these changes tested?

[codecov-commenter]

Codecov Report

Attention: Patch coverage is 91.66667% with 3 lines in your changes missing coverage. Please review.

Project coverage is 34.35%. Comparing base (b74bfe4) to head (df545c5).
Report is 11 commits behind head on main.

Files with missing lines	Patch %	Lines
...ain/scala/org/apache/comet/CometExecIterator.scala	75.00%	0 Missing and 1 partial ⚠️
...la/org/apache/spark/sql/comet/CometExecUtils.scala	66.66%	0 Missing and 1 partial ⚠️
...ark/sql/comet/CometTakeOrderedAndProjectExec.scala	75.00%	0 Missing and 1 partial ⚠️

Additional details and impacted files

@@             Coverage Diff              @@
##               main    #1112      +/-   ##
============================================
+ Coverage     34.33%   34.35%   +0.01%     
  Complexity      898      898              
============================================
  Files           115      115              
  Lines         42986    43416     +430     
  Branches       9369     9499     +130     
============================================
+ Hits          14761    14917     +156     
- Misses        25361    25585     +224     
- Partials       2864     2914      +50     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

---

🚨 Try these New Features:

• Flaky Tests Detection - Detect and resolve failed and flaky tests

[viirya]

Thanks @andygrove