[multistage] Push Collation from Exchange to Lite Mode Sort

[ankitsultana]

Summary

Minor improvement to pushdown collation from the enclosing Exchange to the Lite Mode generated Sort.

This is good for queries like the following. Without this patch the generated Sort only has a fetch and no collation.

SET useMultistageEngine = true;
SET usePhysicalOptimizer = true;
SET useLiteMode = true;

EXPLAIN PLAN FOR WITH ordered_events AS (
  SELECT 
    cityName,
    tripAmount,
    ROW_NUMBER() OVER (
      ORDER BY ts DESC
    ) as row_num
  FROM userFactEvents
),
filtered_events AS (
  SELECT 
    *
  FROM ordered_events
  WHERE row_num < 1000
)
SELECT 
  cityName,
  SUM(tripAmount) as cityTotal
FROM filtered_events
GROUP BY cityName

The plan now for this query is:

Execution Plan
PhysicalAggregate(group=[{0}], agg#0=[$SUM0($1)], aggType=[DIRECT])
  PhysicalFilter(condition=[<($3, 1000)])
    PhysicalWindow(window#0=[window(order by [2 DESC] rows between UNBOUNDED PRECEDING and CURRENT ROW aggs [ROW_NUMBER()])])
      PhysicalExchange(exchangeStrategy=[SINGLETON_EXCHANGE], collation=[[2 DESC]])
        PhysicalSort(fetch=[100000])
          PhysicalProject(cityName=[$3], tripAmount=[$7], ts=[$9])
            PhysicalTableScan(table=[[default, userFactEvents]])

Test Plan

Added a UT and also verified on Quickstart.

[codecov-commenter]

Codecov Report

❌ Patch coverage is 75.00000% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 63.25%. Comparing base (9f74fc7) to head (a74ef0d).
⚠️ Report is 6 commits behind head on master.

Files with missing lines	Patch %	Lines
.../physical/v2/opt/rules/LiteModeSortInsertRule.java	75.00%	0 Missing and 2 partials ⚠️

Additional details and impacted files

@@            Coverage Diff            @@
##             master   #16551   +/-   ##
=========================================
  Coverage     63.25%   63.25%           
- Complexity     1362     1363    +1     
=========================================
  Files          3012     3012           
  Lines        174294   174369   +75     
  Branches      26684    26701   +17     
=========================================
+ Hits         110244   110302   +58     
- Misses        55644    55646    +2     
- Partials       8406     8421   +15     

Flag	Coverage Δ
custom-integration1	100.00% <ø> (ø)
integration	100.00% <ø> (ø)
integration1	100.00% <ø> (ø)
integration2	0.00% <ø> (ø)
java-11	63.23% <75.00%> (+0.04%)	⬆️
java-21	63.22% <75.00%> (-0.01%)	⬇️
temurin	63.25% <75.00%> (+<0.01%)	⬆️
unittests	63.25% <75.00%> (+<0.01%)	⬆️
unittests1	56.36% <75.00%> (-0.05%)	⬇️
unittests2	33.32% <0.00%> (+0.03%)	⬆️

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.

[wirybeaver]

I was wondering which phase generate the PhysicalSort(fetch=[10000]) without the collation before this PR. IIUC, the Calcite will generate the LogicalWindow Node with collocation and doesn't create a LogicalSort Node along the way. If so, it means the pinot has glitches when translate from LogicalWindow to Physical Window by wrongly inserting the PhysicalSort Node without collation.

Moreover, what's the behavior of the general MSE plan? Shall we apply the same patch for non lite mode?

[wirybeaver]

LGTM with a minor comment above.

[wirybeaver]

I was wondering which phase generate the PhysicalSort(fetch=[10000]) without the collation before this PR. IIUC, the Calcite will generate the LogicalWindow Node with collocation and doesn't create a LogicalSort Node along the way. If so, it means the pinot has glitches when translate from LogicalWindow to Physical Window by wrongly inserting the PhysicalSort Node without collation.

Moreover, what's the behavior of the general MSE plan? Shall we apply the same patch for non lite mode?

Ankit provides the answers in slack. I pasted the explanation for future reference:

The default MSE generates a SortExchange under the LogicalWindow WindowExchangeNodeInsertRule if a collation is applicable.
Physical optimizer uses TraitAssignment and adds a trait under the Window for the collation. See
For Lite Mode, we add a sort in the leaf stage. Currently the sort node was blindly added with a row limit. But, it could be that the Exchange just above leaf stage had a collation. That would happen when the node above Exchange had pushed down a Collation trait. So this bug was kinda specific to lite mode