feat: Improve shuffle metrics (second attempt)

[andygrove]

Which issue does this PR close?

N/A

This PR replaces #1173

Rationale for this change

Changes:

• Fix: Report write time accurately (this is now just the time for writing to disk but it previously included some of the IPC encoding time)
• New: Number of input batches
• New: Encoding and compression time
• New: Total time spent in native shuffle code

Before

After

What changes are included in this PR?

How are these changes tested?

[andygrove]

@mbutrovich @parthchandra fyi

[codecov-commenter]

Codecov Report

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

Project coverage is 34.32%. Comparing base (95727aa) to head (6adb04c).
Report is 14 commits behind head on main.

Files with missing lines	Patch %	Lines
...apache/spark/sql/comet/CometCollectLimitExec.scala	50.00%	0 Missing and 1 partial ⚠️
...ark/sql/comet/CometTakeOrderedAndProjectExec.scala	50.00%	0 Missing and 1 partial ⚠️
...t/execution/shuffle/CometShuffleExchangeExec.scala	80.00%	0 Missing and 1 partial ⚠️

Additional details and impacted files

@@             Coverage Diff              @@
##               main    #1175      +/-   ##
============================================
- Coverage     34.32%   34.32%   -0.01%     
  Complexity      899      899              
============================================
  Files           115      115              
  Lines         43500    43506       +6     
  Branches       9496     9498       +2     
============================================
+ Hits          14931    14932       +1     
- Misses        25659    25661       +2     
- Partials       2910     2913       +3     

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

[parthchandra]

Just for clarification - what's the relation between shuffle write time, encoding and compression total time, and native shuffle total time?
I would think shuffle write time + encoding and compression total time = native shuffle total time, but that does not seem to be the case?

[andygrove]

Just for clarification - what's the relation between shuffle write time, encoding and compression total time, and native shuffle total time? I would think shuffle write time + encoding and compression total time = native shuffle total time, but that does not seem to be the case?

There is also evaluating the partition expressions (typically very fast if they are just column references) and then the time to actually split the batches into partitions.

[parthchandra]

Just for clarification - what's the relation between shuffle write time, encoding and compression total time, and native shuffle total time? I would think shuffle write time + encoding and compression total time = native shuffle total time, but that does not seem to be the case?

There is also evaluating the partition expressions (typically very fast if they are just column references) and then the time to actually split the batches into partitions.

From the above screenshot, shuffle write time + encoding and compression total time = 18.4s and native shuffle total time=28.8s, so there is a difference of 10.4s which is substantial. Wondering if we are missing something.

Nonetheless, the PR certainly improves on the current.

[andygrove]

There is also interaction with the memory pool, which makes JNI calls into synchronized code in the JVM.

I will see if I can make the metrics more complete in this PR.

[andygrove]

@parthchandra The numbers almost add up now.

[andygrove]

Here is Gluten's equivalent for comparison:

[parthchandra]

@parthchandra The numbers almost add up now.

Brilliant!