[multistage] Add Support for Broker Server/Segment Pruning #15959

ankitsultana · 2025-06-01T00:55:44Z

Summary

Adds support for enabling broker pruning (for servers and segments). This can especially be beneficial for MSE Lite Mode that aims to mimic the scatter gather pattern of the V1 Engine to enable high-qps use-cases.

Test Plan

Added some unit-tests and tested on Colocated Join Quickstart. I had to add the routing config to enable partition pruning in the broker to enable this.

SET useMultistageEngine=true;
SET usePhysicalOptimizer=true;
-- SET useBrokerPruning=true;

explain implementation plan for select count(*)
from userAttributes
where userUUID = 'user-1'

[0]@10.24.6.236:33955|[0] MAIL_RECEIVE(SINGLETON)
└── [1]@10.24.6.236:33321|[0] MAIL_SEND(SINGLETON)->{[0]@10.24.6.236:33955|[0]}
    └── [1]@10.24.6.236:33321|[0] AGGREGATE_FINAL
        └── [1]@10.24.6.236:33321|[0] MAIL_RECEIVE(SINGLETON)
            ├── [2]@10.24.6.236:38493|[1] MAIL_SEND(SINGLETON)->{[1]@10.24.6.236:35465|[0]} (Subtree Omitted)
            └── [2]@10.24.6.236:33321|[0] MAIL_SEND(SINGLETON)->{[1]@10.24.6.236:35465|[0]}
                └── [2]@10.24.6.236:33321|[0] AGGREGATE_LEAF
                    └── [2]@10.24.6.236:33321|[0] FILTER
                        └── [2]@10.24.6.236:33321|[0] TABLE SCAN (userAttributes) null

SET useMultistageEngine=true;
SET usePhysicalOptimizer=true;
SET useBrokerPruning=true;

explain implementation plan for select count(*)
from userAttributes
where userUUID = 'user-1'

[0]@10.24.6.236:33955|[0] MAIL_RECEIVE(SINGLETON)
└── [1]@10.24.6.236:44355|[0] MAIL_SEND(SINGLETON)->{[0]@10.24.6.236:33955|[0]}
    └── [1]@10.24.6.236:44355|[0] AGGREGATE_DIRECT
        └── [1]@10.24.6.236:44355|[0] FILTER
            └── [1]@10.24.6.236:44355|[0] TABLE SCAN (userAttributes) null

codecov-commenter · 2025-06-01T01:37:35Z

Codecov Report

Attention: Patch coverage is 89.28571% with 9 lines in your changes missing coverage. Please review.

Project coverage is 63.41%. Comparing base (1a476de) to head (0ad1251).
Report is 416 commits behind head on master.

Files with missing lines	Patch %	Lines
...t/query/planner/logical/LeafStageToPinotQuery.java	89.18%	0 Missing and 4 partials ⚠️
...he/pinot/query/context/PhysicalPlannerContext.java	50.00%	1 Missing and 1 partial ⚠️
...pinot/query/parser/CalciteRexExpressionParser.java	84.61%	2 Missing ⚠️
...al/v2/opt/rules/LeafStageWorkerAssignmentRule.java	96.00%	1 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff              @@
##             master   #15959      +/-   ##
============================================
+ Coverage     62.90%   63.41%   +0.50%     
+ Complexity     1386     1355      -31     
============================================
  Files          2867     2911      +44     
  Lines        163354   166938    +3584     
  Branches      24952    25533     +581     
============================================
+ Hits         102755   105859    +3104     
- Misses        52847    53057     +210     
- Partials       7752     8022     +270

Flag	Coverage Δ
custom-integration1	`100.00% <ø> (ø)`
integration	`100.00% <ø> (ø)`
integration1	`100.00% <ø> (ø)`
integration2	`0.00% <ø> (ø)`
java-11	`63.37% <89.28%> (+0.50%)`	⬆️
java-21	`63.38% <89.28%> (+0.56%)`	⬆️
skip-bytebuffers-false	`?`
skip-bytebuffers-true	`?`
temurin	`63.41% <89.28%> (+0.50%)`	⬆️
unittests	`63.40% <89.28%> (+0.50%)`	⬆️
unittests1	`56.47% <89.28%> (+0.65%)`	⬆️
unittests2	`33.45% <0.00%> (-0.12%)`	⬇️

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.

ankitsultana · 2025-06-03T04:04:51Z

pinot-query-planner/src/main/java/org/apache/pinot/query/parser/CalciteRexExpressionParser.java

  // --------------------------------------------------------------------------

-  public static List<Expression> convertRexNodes(List<RexExpression> rexNodes, PinotQuery pinotQuery) {
+  public static List<Expression> convertRexNodes(List<RexExpression> rexNodes, List<Expression> selectList) {


simplifying this code. It was taking in a wider object than required. Only access to the selectList was required

ankitsultana · 2025-06-03T16:36:12Z

...java/org/apache/pinot/query/planner/physical/v2/opt/rules/LeafStageWorkerAssignmentRule.java

  }

+  @VisibleForTesting
+  static PinotQuery createPinotQuery(String tableName, TableScan tableScan, Deque<PRelNode> parents,


note for reviewers: This logic to create a PinotQuery out of a leaf stage is similar to ServerPlanRequestVisitor which is run in the servers. However, that visitor deals with PlanNode types, whereas here we deal with PRelNode.

Jackie-Jiang

This is very useful for MSE. Is it possible to extract the code into a common util class? Maybe move the code from SSE to a util?

Jackie-Jiang · 2025-06-03T22:14:48Z

pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java

+        // TODO(mse-physical): Consider removing this query option and making this the default, since there's already
+        //   a table config to enable broker pruning (it is disabled by default).


Which table config are you referring to?

I meant the routing. segmentPrunerTypes key in the table config. Based on my tests we don't prune segments unless this is set (maybe I am missing something)

ankitsultana · 2025-06-04T04:37:52Z

This is very useful for MSE. Is it possible to extract the code into a common util class? Maybe move the code from SSE to a util?

Moved the logic to a util class.

...uery-planner/src/main/java/org/apache/pinot/query/planner/logical/LeafStageToPinotQuery.java

wirybeaver

LGTM. IIUC, after applying the BrokerPruner, the Rel Trait of the Data distribution of the TableScan will be singleton so that the following Aggregation Pushdown Rule would apply the AggregateDirect.

ankitsultana added multi-stage Related to the multi-stage query engine mse-physical-optimizer labels Jun 1, 2025

ankitsultana added 2 commits June 3, 2025 02:16

[multistage] Add Support for Broker Server/Segment Pruning

8e870e4

minor cleanup

f718813

ankitsultana force-pushed the mse-physical-p8 branch from 391f9bf to f718813 Compare June 3, 2025 04:00

ankitsultana commented Jun 3, 2025

View reviewed changes

add doc comments for query option

aeab4b4

ankitsultana marked this pull request as ready for review June 3, 2025 04:09

ankitsultana commented Jun 3, 2025

View reviewed changes

Jackie-Jiang reviewed Jun 3, 2025

View reviewed changes

ankitsultana added 3 commits June 4, 2025 04:11

make it re-usable

d9bfca9

fix bug

24dc178

add precondition

c005b5a

ankitsultana requested a review from Jackie-Jiang June 4, 2025 04:29

invert boolean

080870f

shauryachats approved these changes Jun 4, 2025

View reviewed changes

ankitsultana commented Jun 4, 2025

View reviewed changes

...uery-planner/src/main/java/org/apache/pinot/query/planner/logical/LeafStageToPinotQuery.java Outdated Show resolved Hide resolved

self-review: rename method

0ad1251

shauryachats approved these changes Jun 4, 2025

View reviewed changes

wirybeaver reviewed Jun 4, 2025

View reviewed changes

wirybeaver approved these changes Jun 4, 2025

View reviewed changes

itschrispeck approved these changes Jun 4, 2025

View reviewed changes

ankitsultana merged commit fb2e4df into apache:master Jun 4, 2025
18 checks passed

ankitsultana mentioned this pull request Jun 9, 2025

Pinot fanout constant queries to all servers which could cause unexpected DNS query spikes #16024

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[multistage] Add Support for Broker Server/Segment Pruning #15959

[multistage] Add Support for Broker Server/Segment Pruning #15959

Uh oh!

ankitsultana commented Jun 1, 2025 •

edited

Loading

Uh oh!

codecov-commenter commented Jun 1, 2025 •

edited

Loading

Uh oh!

ankitsultana Jun 3, 2025

Uh oh!

ankitsultana Jun 3, 2025

Uh oh!

Jackie-Jiang left a comment

Uh oh!

Jackie-Jiang Jun 3, 2025

Uh oh!

ankitsultana Jun 3, 2025 •

edited

Loading

Uh oh!

ankitsultana commented Jun 4, 2025

Uh oh!

Uh oh!

wirybeaver left a comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

		// TODO(mse-physical): Consider removing this query option and making this the default, since there's already
		// a table config to enable broker pruning (it is disabled by default).

[multistage] Add Support for Broker Server/Segment Pruning #15959

[multistage] Add Support for Broker Server/Segment Pruning #15959

Uh oh!

Conversation

ankitsultana commented Jun 1, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Test Plan

Uh oh!

codecov-commenter commented Jun 1, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

ankitsultana Jun 3, 2025

Choose a reason for hiding this comment

Uh oh!

ankitsultana Jun 3, 2025

Choose a reason for hiding this comment

Uh oh!

Jackie-Jiang left a comment

Choose a reason for hiding this comment

Uh oh!

Jackie-Jiang Jun 3, 2025

Choose a reason for hiding this comment

Uh oh!

ankitsultana Jun 3, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ankitsultana commented Jun 4, 2025

Uh oh!

Uh oh!

wirybeaver left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

ankitsultana commented Jun 1, 2025 •

edited

Loading

codecov-commenter commented Jun 1, 2025 •

edited

Loading

ankitsultana Jun 3, 2025 •

edited

Loading