[multistage] Fix null handling for several aggregation functions

[dang-stripe]

This is a complementary change to #17750 to address errors with Cannot invoke "java.util.Set.size()" because "intermediateResult" is null that we are running internally. We don't have #17750 pulled down yet, but figured it's worth upstreaming this anyway since it's an extra layer of defense against the shape of error. We did notice a bug with the count aggregation returning null instead of 0 so we also fixed that here too.

This happened to us after rolling out MSE broker pruning and when all segments were pruned in the query since there's no MSE broker logic yet to short circuit the query in the broker. This adds a comment to implement support.

cc @yashmayya @Jackie-Jiang @timothy-e

[codecov-commenter]

Codecov Report

❌ Patch coverage is 0% with 5 lines in your changes missing coverage. Please review.
✅ Project coverage is 63.67%. Comparing base (5fd3183) to head (d74384d).
⚠️ Report is 3 commits behind head on master.

Files with missing lines	Patch %	Lines
...gation/function/CovarianceAggregationFunction.java	0.00%	1 Missing and 1 partial ⚠️
...aggregation/function/IdSetAggregationFunction.java	0.00%	1 Missing and 1 partial ⚠️
...aggregation/function/CountAggregationFunction.java	0.00%	0 Missing and 1 partial ⚠️

Additional details and impacted files

@@             Coverage Diff              @@
##             master   #18471      +/-   ##
============================================
- Coverage     63.68%   63.67%   -0.01%     
- Complexity     1682     1684       +2     
============================================
  Files          3262     3262              
  Lines        199826   199836      +10     
  Branches      31031    31033       +2     
============================================
  Hits         127255   127255              
+ Misses        62421    62420       -1     
- Partials      10150    10161      +11     

Flag	Coverage Δ
custom-integration1	100.00% <ø> (ø)
integration	100.00% <ø> (ø)
integration1	100.00% <ø> (ø)
integration2	0.00% <ø> (ø)
java-21	63.67% <0.00%> (-0.01%)	⬇️
temurin	63.67% <0.00%> (-0.01%)	⬇️
unittests	63.67% <0.00%> (-0.01%)	⬇️
unittests1	55.76% <0.00%> (-0.02%)	⬇️
unittests2	34.95% <0.00%> (+<0.01%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

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

🚀 New features to boost your workflow:

• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[Jackie-Jiang]

Annotate it @Nullable

[dang-stripe]

added

[Jackie-Jiang]

Why do we need this?

[dang-stripe]

I originally added this in our fork before I saw the fix in #17750. It was intended to be a different way of addressing the same bug.

When all segments are pruned by the broker, the final aggregate stage receives no input blocks so _mergeResultHolder[i] will be null and calling getResult() will trigger an NPE. This felt like a more comprehensive change so we don't need to update each aggregation function for null handling.

But then I realized returning identity for all aggregation types might not be right here (depends on the aggregation) so I've reverted this change.

[Jackie-Jiang]

(minor) Also annotate return as @Nullable, same for other places

[dang-stripe]

@yashmayya @Jackie-Jiang wanted to discuss if there's a better fix that doesn't require updating all aggregation functions for null handling. i checked postgres and all functions return null except count-based functions.

https://www.postgresql.org/docs/current/functions-aggregate.html

It should be noted that except for count, these functions return a null value when no rows are selected. In particular, sum of no rows returns null, not zero as one might expect, and array_agg returns null rather than an empty array when there are no input rows. `

so we could do something like this in getResult(), but someone writing a new count-based agg function will need to include their function in isCountFamily. what do you think?

  case FINAL:
      Object intermediate = _mergeResultHolder[i];
      if (intermediate == null) {
          value = isCountFamily(aggFunction.getType()) ? 0L : null;
      } else {
          value = aggFunction.extractFinalResult(intermediate);
      }