[multistage] refactor traversals of stage nodes into visitor pattern #9560

agavra · 2022-10-08T03:57:42Z

This PR improves the traversal code around StageNode. It sets up a common pattern for visiting nodes, collecting information, rewriting and making other changes. This PR is setup for one that will help us implement a global sort stage for LIMIT/OFFSET queries and support sort push down.

There are five main parts to look at:

I added the StageNodeVisitor interface and implemented visit in all of the Stage Node implementations
I refactored the partitionKey optimization (that removes a shuffle if not necessary) into a Visitor (ShuffleRewriter)
I refactored constructing the QueryPlan metadata into a visitor (this is in preparation for the next PR) (QueryPlanGenerator)
I refactored constructing the Operator into a visitor (PhyscialPlanBuilder)
I refactored QueryPlan#explain into a visitor, and also improved the functionality (see new plan explain below)

Lastly, I added some quality of life improvements in debug-ability and I identified a "bug" in nested loop joins - though I'll fix that one in a future PR (see FIXME comment)

Example of the improved explain (it now properly recognizes which nodes are executing what code):

[0]@localhost:51925 MAIL_RECEIVE(RANDOM_DISTRIBUTED)
├── [1]@localhost:51923 MAIL_SEND(RANDOM_DISTRIBUTED)->{[0]@localhost:51925} (Subtree Omitted)
└── [1]@localhost:51924 MAIL_SEND(RANDOM_DISTRIBUTED)->{[0]@localhost:51925}
   └── [1]@localhost:51924 JOIN
      ├── [1]@localhost:51924 MAIL_RECEIVE(HASH_DISTRIBUTED)
      │   ├── [2]@localhost:51924 MAIL_SEND(HASH_DISTRIBUTED)->{[1]@localhost:51923,[1]@localhost:51924}
      │   │   └── [2]@localhost:51924 TABLE SCAN (a) {REALTIME=[a3]}
      │   └── [2]@localhost:51923 MAIL_SEND(HASH_DISTRIBUTED)->{[1]@localhost:51923,[1]@localhost:51924}
      │      └── [2]@localhost:51923 TABLE SCAN (a) {REALTIME=[a1, a2]}
      └── [1]@localhost:51924 MAIL_RECEIVE(HASH_DISTRIBUTED)
         └── [3]@localhost:51923 MAIL_SEND(HASH_DISTRIBUTED)->{[1]@localhost:51923,[1]@localhost:51924}
            └── [3]@localhost:51923 TABLE SCAN (b) {REALTIME=[b1]}

This PR improves the traversal code around StageNode. It sets up a common pattern for visiting nodes, collecting information, rewriting and making other changes. This PR is setup for one that will help us implement a global sort stage for LIMIT/OFFSET queries.

codecov-commenter · 2022-10-11T23:41:58Z

Codecov Report

Merging #9560 (f390492) into master (8cdae92) will increase coverage by 40.25%.
The diff coverage is 63.67%.

@@              Coverage Diff              @@
##             master    #9560       +/-   ##
=============================================
+ Coverage     28.28%   68.54%   +40.25%     
- Complexity       53     4920     +4867     
=============================================
  Files          1917     1938       +21     
  Lines        102594   103383      +789     
  Branches      15586    15683       +97     
=============================================
+ Hits          29022    70867    +41845     
+ Misses        70735    27497    -43238     
- Partials       2837     5019     +2182

Flag	Coverage Δ
integration1	`?`
integration2	`24.71% <0.00%> (+0.13%)`	⬆️
unittests1	`67.34% <63.67%> (?)`
unittests2	`15.60% <63.67%> (?)`

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

Impacted Files	Coverage Δ
...e/pinot/query/planner/ExplainPlanStageVisitor.java	`0.00% <0.00%> (ø)`
...java/org/apache/pinot/query/planner/QueryPlan.java	`88.88% <0.00%> (+88.88%)`	⬆️
.../org/apache/pinot/query/planner/StageMetadata.java	`95.00% <0.00%> (+95.00%)`	⬆️
...e/pinot/query/planner/stage/AbstractStageNode.java	`100.00% <ø> (+100.00%)`	⬆️
...apache/pinot/query/mailbox/GrpcMailboxService.java	`94.11% <0.00%> (+94.11%)`	⬆️
.../pinot/query/runtime/blocks/TransferableBlock.java	`69.56% <0.00%> (+69.56%)`	⬆️
...inot/query/runtime/operator/AggregateOperator.java	`84.21% <ø> (+84.21%)`	⬆️
...pinot/query/runtime/operator/HashJoinOperator.java	`82.60% <ø> (+82.60%)`	⬆️
...query/runtime/operator/MailboxReceiveOperator.java	`79.31% <ø> (+79.31%)`	⬆️
...ot/query/runtime/operator/MailboxSendOperator.java	`90.42% <ø> (+90.42%)`	⬆️
... and 1423 more

📣 We’re building smart automated test selection to slash your CI/CD build times. Learn more

walterddr

partial review on the planner side. will do another pass on the runtime a bit later

pinot-query-planner/src/main/java/org/apache/pinot/query/planner/stage/MailboxReceiveNode.java

pinot-query-planner/src/main/java/org/apache/pinot/query/planner/ExplainPlanStageVisitor.java

pinot-query-planner/src/main/java/org/apache/pinot/query/planner/stage/StageNodeVisitor.java

walterddr · 2022-10-12T15:19:55Z

pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/QueryRunner.java

@@ -87,7 +87,7 @@ public void init(PinotConfiguration config, InstanceDataManager instanceDataMana
      _serverExecutor = new ServerQueryExecutorV1Impl();
      _serverExecutor.init(config, instanceDataManager, serverMetrics);
      _workerExecutor = new WorkerQueryExecutor();
-      _workerExecutor.init(config, serverMetrics, _mailboxService, _hostname, _port);
+      _workerExecutor.init(_mailboxService, _hostname, _port);


why reducing the init method signature?

these weren't used - so just some cleanup. I can remove it if I'm missing something

please do not remove them for now. config is useful for overrides and metrics we will be using in the future for sure

this isn't fixed

oops! sorry I thought I had done that but looks like I only fixed the other code path 😬 shame on me

could you revert this change?

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

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

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

walterddr

looks good to me overall. please see the comments and otherwise good to go

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

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

pinot-query-planner/src/main/java/org/apache/pinot/query/planner/QueryPlan.java

pinot-query-planner/src/main/java/org/apache/pinot/query/planner/stage/MailboxReceiveNode.java

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

walterddr · 2022-10-13T20:59:18Z

pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/QueryRunner.java

@@ -87,7 +87,7 @@ public void init(PinotConfiguration config, InstanceDataManager instanceDataMana
      _serverExecutor = new ServerQueryExecutorV1Impl();
      _serverExecutor.init(config, instanceDataManager, serverMetrics);
      _workerExecutor = new WorkerQueryExecutor();
-      _workerExecutor.init(config, serverMetrics, _mailboxService, _hostname, _port);
+      _workerExecutor.init(_mailboxService, _hostname, _port);


this isn't fixed

...query-runtime/src/main/java/org/apache/pinot/query/runtime/executor/PhysicalPlanBuilder.java

...query-runtime/src/main/java/org/apache/pinot/query/runtime/executor/WorkerQueryExecutor.java

walterddr · 2022-10-13T21:01:23Z

...ry-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MailboxReceiveOperator.java

+      // FIXME: there's a bug where singletonInstance may be null in the case of a JOIN where
+      // one side is BROADCAST and the other is SINGLETON (this is the case with nested loop
+      // joins for inequality conditions). This causes NPEs in the logs, but actually works
+      // because the side that hits the NPE doesn't expect to get any data anyway (that's the
+      // side that gets the broadcast from one side but nothing from the SINGLETON)


can we create an issue and link it here?

#9592 - added to comment as well

walterddr

lgtm.

walterddr · 2022-10-17T13:51:08Z

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

+ * inspecting whether all data required by a specific subtree already resides on
+ * a single host. It gathers the information recursively by checking which partitioned


Suggested change

* inspecting whether all data required by a specific subtree already resides on

* a single host. It gathers the information recursively by checking which partitioned

* inspecting whether all data required by a specific subtree are already colocated.

* It gathers the information recursively by checking which partitioned

walterddr · 2022-10-17T13:52:50Z

...query-runtime/src/main/java/org/apache/pinot/query/runtime/executor/PhysicalPlanVisitor.java

+import org.apache.pinot.query.runtime.operator.TransformOperator;
+
+
+public class PhysicalPlanVisitor implements StageNodeVisitor<Operator<TransferableBlock>, Void> {


agavra added 2 commits October 7, 2022 20:53

also refactor QueryPlan#explan

d12d5ff

agavra force-pushed the stage_refactor branch from 0d6aa1f to d12d5ff Compare October 10, 2022 00:31

agavra added 2 commits October 9, 2022 18:45

fix SerDe to account for StageNode as field

388139a

checkstyle

e9e1a3e

walterddr reviewed Oct 12, 2022

View reviewed changes

agavra added 3 commits October 12, 2022 10:39

address feedback

3e0c583

minor javadoc addition

08e6177

checkstyle import again

3178f03

walterddr reviewed Oct 13, 2022

View reviewed changes

address feedback 2

f390492

walterddr approved these changes Oct 17, 2022

View reviewed changes

add javadoc and make PhysicalPlanBuilder static usage

fcfb1e3

walterddr merged commit 4935326 into apache:master Oct 17, 2022

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[multistage] refactor traversals of stage nodes into visitor pattern #9560

[multistage] refactor traversals of stage nodes into visitor pattern #9560

agavra commented Oct 8, 2022 •

edited

Loading

codecov-commenter commented Oct 11, 2022 •

edited

Loading

walterddr left a comment

walterddr Oct 12, 2022

agavra Oct 12, 2022

walterddr Oct 12, 2022

walterddr Oct 13, 2022

agavra Oct 13, 2022

walterddr Oct 17, 2022

walterddr left a comment

walterddr Oct 13, 2022

walterddr Oct 13, 2022

agavra Oct 13, 2022

walterddr left a comment

walterddr Oct 17, 2022

walterddr Oct 17, 2022

		* inspecting whether all data required by a specific subtree already resides on
		* a single host. It gathers the information recursively by checking which partitioned

		import org.apache.pinot.query.runtime.operator.TransformOperator;


		public class PhysicalPlanVisitor implements StageNodeVisitor<Operator<TransferableBlock>, Void> {

[multistage] refactor traversals of stage nodes into visitor pattern #9560

[multistage] refactor traversals of stage nodes into visitor pattern #9560

Conversation

agavra commented Oct 8, 2022 • edited Loading

codecov-commenter commented Oct 11, 2022 • edited Loading

Codecov Report

walterddr left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

walterddr left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

walterddr left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

agavra commented Oct 8, 2022 •

edited

Loading

codecov-commenter commented Oct 11, 2022 •

edited

Loading