[Spool] Actual implementation #14507

gortiz · 2024-11-20T13:21:30Z

This PR is a continuation of #14495 and the next step on #14196.

While previous PRs were focused the machinery to find the common stages and how to replace them in the plan, this one is focused on sending these plans to the servers and how servers deal with them.

The main change is that now send operator may need to send to multiple stages. That implies a change in the protocol, which I made backward compatible by deprecating the old attribute and adding a new one that is repeated. A deprecated attribute can still be used and reserves the id in protobuf. Servers still recognize the old attribute (required in case a new server talks with an old broker) but they first try to read the new attribute (which will be null/empty if sent by an old broker) and in case it is empty, use the old attribute.

To understand the change in the operators is important to understand the topology used by BlockExchange and Mailbox. Mailboxes receive blocks and send them to exactly one worker. They could send the block through memory or GRPC. BlockExchanges are the ones that decide how to distribute the block between different mailboxes. For example BroadcastExchange send the block to all Mailboxes while RandomExchange sends the block to a random mailbox.

With the multi-sender operators we need to send to broadcast the blocks to multiple stages, which they may need to apply another distribution for each worker. What this PR does is to create a new layer of BlockExchanges when a sender is multi-sender. The root of this layer is always a standard BroadcastExchange, but it sends to a new type of SendingMailbox: A BlockExchangeSendingMailbox. This new mailbox is just an adaptor pattern that wraps a BlockExchange and shows it as a SendingMailbox. The BlockExchanges used in this second layer is the one normally used for each stage. Remember that this distribution must the same for all stages in the same multi-sender.

As an example of how it works, you can run the following in ColocatedJoinEngineQuickStart

SET useSpools = true;
explain implementation plan for
select * from userAttributes as a1
join userGroups as a2
on a1.userUUID = a2.userUUID
join userAttributes as a3
on a1.userUUID = a3.userUUID
limit 10

which returns:

[0]@192.168.1.42:46155|[0] MAIL_RECEIVE(BROADCAST_DISTRIBUTED)
├── [1]@192.168.1.42:44265|[0] MAIL_SEND(BROADCAST_DISTRIBUTED)->{[0]@192.168.1.42:46155|[0]} (Subtree Omitted)
├── [1]@192.168.1.42:34477|[1] MAIL_SEND(BROADCAST_DISTRIBUTED)->{[0]@192.168.1.42:46155|[0]} (Subtree Omitted)
├── [1]@192.168.1.42:44185|[2] MAIL_SEND(BROADCAST_DISTRIBUTED)->{[0]@192.168.1.42:46155|[0]} (Subtree Omitted)
└── [1]@192.168.1.42:34303|[3] MAIL_SEND(BROADCAST_DISTRIBUTED)->{[0]@192.168.1.42:46155|[0]}
    └── [1]@192.168.1.42:34303|[3] SORT LIMIT 10
        └── [1]@192.168.1.42:34303|[3] MAIL_RECEIVE(HASH_DISTRIBUTED)
            ├── [2]@192.168.1.42:44265|[0] MAIL_SEND(HASH_DISTRIBUTED)->{[1]@192.168.1.42:37717|[3],[1]@192.168.1.42:43409|[1],[1]@192.168.1.42:43851|[2],[1]@192.168.1.42:46503|[0]} (Subtree Omitted)
            ├── [2]@192.168.1.42:34477|[1] MAIL_SEND(HASH_DISTRIBUTED)->{[1]@192.168.1.42:37717|[3],[1]@192.168.1.42:43409|[1],[1]@192.168.1.42:43851|[2],[1]@192.168.1.42:46503|[0]} (Subtree Omitted)
            ├── [2]@192.168.1.42:44185|[2] MAIL_SEND(HASH_DISTRIBUTED)->{[1]@192.168.1.42:37717|[3],[1]@192.168.1.42:43409|[1],[1]@192.168.1.42:43851|[2],[1]@192.168.1.42:46503|[0]} (Subtree Omitted)
            └── [2]@192.168.1.42:34303|[3] MAIL_SEND(HASH_DISTRIBUTED)->{[1]@192.168.1.42:37717|[3],[1]@192.168.1.42:43409|[1],[1]@192.168.1.42:43851|[2],[1]@192.168.1.42:46503|[0]}
                └── [2]@192.168.1.42:34303|[3] SORT LIMIT 10
                    └── [2]@192.168.1.42:34303|[3] JOIN
                        ├── [2]@192.168.1.42:34303|[3] MAIL_RECEIVE(HASH_DISTRIBUTED)
                        │   ├── [3]@192.168.1.42:44265|[0] MAIL_SEND(HASH_DISTRIBUTED)[PARTITIONED]->{[2]@192.168.1.42:46503|[0]} (Subtree Omitted)
                        │   ├── [3]@192.168.1.42:34477|[1] MAIL_SEND(HASH_DISTRIBUTED)[PARTITIONED]->{[2]@192.168.1.42:43409|[1]} (Subtree Omitted)
                        │   ├── [3]@192.168.1.42:44185|[2] MAIL_SEND(HASH_DISTRIBUTED)[PARTITIONED]->{[2]@192.168.1.42:43851|[2]} (Subtree Omitted)
                        │   └── [3]@192.168.1.42:34303|[3] MAIL_SEND(HASH_DISTRIBUTED)[PARTITIONED]->{[2]@192.168.1.42:37717|[3]}
                        │       └── [3]@192.168.1.42:34303|[3] JOIN
                        │           ├── [3]@192.168.1.42:34303|[3] MAIL_RECEIVE(HASH_DISTRIBUTED)
spool here ->           │           │   ├── [4]@192.168.1.42:44185|[1] MAIL_SEND(HASH_DISTRIBUTED)->{[2]@192.168.1.42:37717|[3],[2]@192.168.1.42:43409|[1],[2]@192.168.1.42:43851|[2],[2]@192.168.1.42:46503|[0],[3]@192.168.1.42:37717|[3],[3]@192.168.1.42:43409|[1],[3]@192.168.1.42:43851|[2],[3]@192.168.1.42:46503|[0]} (Subtree Omitted)
                        │           │   └── [4]@192.168.1.42:44265|[0] MAIL_SEND(HASH_DISTRIBUTED)->{[2]@192.168.1.42:37717|[3],[2]@192.168.1.42:43409|[1],[2]@192.168.1.42:43851|[2],[2]@192.168.1.42:46503|[0],[3]@192.168.1.42:37717|[3],[3]@192.168.1.42:43409|[1],[3]@192.168.1.42:43851|[2],[3]@192.168.1.42:46503|[0]}
                        │           │       └── [4]@192.168.1.42:44265|[0] PROJECT
                        │           │           └── [4]@192.168.1.42:44265|[0] TABLE SCAN (userAttributes) null
                        │           └── [3]@192.168.1.42:34303|[3] MAIL_RECEIVE(HASH_DISTRIBUTED)
                        │               ├── [5]@192.168.1.42:44185|[1] MAIL_SEND(HASH_DISTRIBUTED)->{[3]@192.168.1.42:37717|[3],[3]@192.168.1.42:43409|[1],[3]@192.168.1.42:43851|[2],[3]@192.168.1.42:46503|[0]} (Subtree Omitted)
                        │               └── [5]@192.168.1.42:44265|[0] MAIL_SEND(HASH_DISTRIBUTED)->{[3]@192.168.1.42:37717|[3],[3]@192.168.1.42:43409|[1],[3]@192.168.1.42:43851|[2],[3]@192.168.1.42:46503|[0]}
                        │                   └── [5]@192.168.1.42:44265|[0] PROJECT
                        │                       └── [5]@192.168.1.42:44265|[0] TABLE SCAN (userGroups) null
                        └── [2]@192.168.1.42:34303|[3] MAIL_RECEIVE(HASH_DISTRIBUTED)
spool here ->               ├── [4]@192.168.1.42:44185|[1] MAIL_SEND(HASH_DISTRIBUTED)->{[2]@192.168.1.42:37717|[3],[2]@192.168.1.42:43409|[1],[2]@192.168.1.42:43851|[2],[2]@192.168.1.42:46503|[0],[3]@192.168.1.42:37717|[3],[3]@192.168.1.42:43409|[1],[3]@192.168.1.42:43851|[2],[3]@192.168.1.42:46503|[0]} (Subtree Omitted)
                            └── [4]@192.168.1.42:44265|[0] MAIL_SEND(HASH_DISTRIBUTED)->{[2]@192.168.1.42:37717|[3],[2]@192.168.1.42:43409|[1],[2]@192.168.1.42:43851|[2],[2]@192.168.1.42:46503|[0],[3]@192.168.1.42:37717|[3],[3]@192.168.1.42:43409|[1],[3]@192.168.1.42:43851|[2],[3]@192.168.1.42:46503|[0]}
                                └── [4]@192.168.1.42:44265|[0] PROJECT
                                    └── [4]@192.168.1.42:44265|[0] TABLE SCAN (userAttributes) null

TODO:

Fix queries (right now this query return an empty result, which is incorrect).
Automatic tests

codecov-commenter · 2024-11-20T18:00:51Z

Codecov Report

Attention: Patch coverage is 76.19048% with 55 lines in your changes missing coverage. Please review.

Project coverage is 63.84%. Comparing base (59551e4) to head (cfeb8c4).
Report is 1526 commits behind head on master.

Files with missing lines	Patch %	Lines
...query/runtime/operator/exchange/BlockExchange.java	48.38%	10 Missing and 6 partials ⚠️
...hysical/colocated/GreedyShuffleRewriteVisitor.java	0.00%	15 Missing ⚠️
...me/operator/utils/BlockingMultiStreamConsumer.java	0.00%	9 Missing ⚠️
...ery/planner/physical/MailboxAssignmentVisitor.java	90.62%	4 Missing and 2 partials ⚠️
...requesthandler/MultiStageBrokerRequestHandler.java	0.00%	2 Missing ⚠️
...inot/query/planner/serde/PlanNodeDeserializer.java	66.66%	1 Missing and 1 partial ⚠️
...apache/pinot/query/mailbox/GrpcSendingMailbox.java	50.00%	2 Missing ⚠️
...he/pinot/query/planner/logical/PlanFragmenter.java	85.71%	0 Missing and 1 partial ⚠️
.../pinot/query/planner/serde/PlanNodeSerializer.java	83.33%	0 Missing and 1 partial ⚠️
...he/pinot/query/mailbox/InMemorySendingMailbox.java	0.00%	1 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff              @@
##             master   #14507      +/-   ##
============================================
+ Coverage     61.75%   63.84%   +2.08%     
- Complexity      207     1608    +1401     
============================================
  Files          2436     2704     +268     
  Lines        133233   150849   +17616     
  Branches      20636    23309    +2673     
============================================
+ Hits          82274    96302   +14028     
- Misses        44911    47335    +2424     
- Partials       6048     7212    +1164

Flag	Coverage Δ
custom-integration1	`100.00% <ø> (+99.99%)`	⬆️
integration	`100.00% <ø> (+99.99%)`	⬆️
integration1	`100.00% <ø> (+99.99%)`	⬆️
integration2	`0.00% <ø> (ø)`
java-11	`63.79% <76.19%> (+2.08%)`	⬆️
java-21	`63.73% <76.19%> (+2.11%)`	⬆️
skip-bytebuffers-false	`63.81% <76.19%> (+2.06%)`	⬆️
skip-bytebuffers-true	`63.71% <76.19%> (+35.98%)`	⬆️
temurin	`63.84% <76.19%> (+2.08%)`	⬆️
unittests	`63.83% <76.19%> (+2.08%)`	⬆️
unittests1	`56.28% <75.54%> (+9.39%)`	⬆️
unittests2	`34.13% <3.46%> (+6.40%)`	⬆️

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.

rename some test methods Co-authored-by: Yash Mayya <[email protected]>

Explain implementation is working, but not actual queries.

gortiz · 2024-12-17T10:24:50Z

I think this PR is ready to review

gortiz · 2024-12-17T10:28:46Z

...-planner/src/main/java/org/apache/pinot/query/planner/physical/MailboxAssignmentVisitor.java

-      Map<Integer, QueryServerInstance> receiverServerMap = receiverMetadata.getWorkerIdToServerInstanceMap();
-      Map<Integer, Map<Integer, MailboxInfos>> senderMailboxesMap = senderMetadata.getWorkerIdToMailboxesMap();
-      Map<Integer, Map<Integer, MailboxInfos>> receiverMailboxesMap = receiverMetadata.getWorkerIdToMailboxesMap();
+      for (Integer receiverStageId : sendNode.getReceiverStageIds()) {


I recommend to review this class with the hide whitespace option in GH

gortiz · 2024-12-17T10:31:05Z

...t-query-planner/src/main/java/org/apache/pinot/query/planner/serde/PlanNodeDeserializer.java

+    if (protoReceiverIds == null || protoReceiverIds.isEmpty()) {
+      // This should only happen if a not updated broker sends the request
+      receiverIds = List.of(protoMailboxSendNode.getReceiverStageId());


This (and the change in proto) is what keeps backward compatibility

gortiz · 2024-12-17T11:51:26Z

pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/GrpcSendingMailbox.java

+
+  @Override
+  public String toString() {
+    return "g" + _id;


I've added some toString to mailboxes to make it easier to understand them in debug/trace logs

gortiz · 2024-12-17T11:54:13Z

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

+        if (LOGGER.isTraceEnabled()) {
+          LOGGER.trace("Sending EOS metadata. Only mailbox #{} will get stats", mailboxIdToSendMetadata);
+        }


Remember that this verbose way to log (first check then log) is used to avoid the unnecessary varargs array allocation in LOGGER.trace(string, object...) when the log is not going to be printed. Given this is the hotpath, we are extra careful trying to not include extra allocations.

gortiz · 2024-12-17T11:56:10Z

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

   */
  public boolean send(TransferableBlock block)
-      throws Exception {
+      throws IOException, TimeoutException {


this is more precise and IIRC we need to limit the scope of the exception in order to create the BlockExchangeSendingMailbox

gortiz · 2024-12-17T12:10:56Z

Improve PR description

yashmayya

I think it might be very useful to have a couple of integration tests for some basic spool use cases (verifying both the execution output as well as the plan potentially)?

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

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

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

...planner/src/main/java/org/apache/pinot/query/planner/explain/PhysicalExplainPlanVisitor.java

...y-planner/src/main/java/org/apache/pinot/query/planner/physical/DispatchablePlanVisitor.java

...t-query-planner/src/main/java/org/apache/pinot/query/planner/serde/PlanNodeDeserializer.java

pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/InMemorySendingMailbox.java

yashmayya · 2024-12-23T05:37:35Z

...main/java/org/apache/pinot/query/planner/physical/colocated/GreedyShuffleRewriteVisitor.java

+      if (context.isJoinStage(receiverStageId)) {
+        sendsToJoin = true;
+        break;
+      }


We're only checking if at least one receiver stage is a join stage here?

I've just changed this code to make it closer to the original code. I think it will be easier to read this way.

...src/main/java/org/apache/pinot/query/runtime/operator/utils/BlockingMultiStreamConsumer.java

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

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

gortiz · 2024-12-23T14:07:09Z

We have some tests that verify that, but in different files. For example:

pinot-query-planner/src/test/resources/queries/ExplainPhysicalPlans.json verifies the plan.
pinot-query-runtime/src/test/resources/queries/Spool.json verifies the result.

TBH there is no test verifying with the plan that spools are used AND the result is correct (ie comparing with h2), although the query in these two cases is the same.

Semantic should be the same, but the PR should be easier to read.

gortiz · 2025-01-02T14:15:15Z

...main/java/org/apache/pinot/query/planner/physical/colocated/GreedyShuffleRewriteVisitor.java

      Set<ColocationKey> colocationKeys;
-      if (canSkipShuffleBasic && areServersSuperset(node.getReceiverStageId(), node.getStageId())) {
+      if (canSkipShuffleBasic && allAreSuperSet(receiverStageIds, node)) {


Otherwise we cannot apply the shuffle optimization.

This means that if we find two stages that are equivalent but one can be optimized with colocated join while the other cannot, we need to decide whether we want to apply spool or colocated.

Which one is better? I'm not sure. Probably we will need data to understand the difference. In theory if we don't apply spooling, we are going to end up executing the sender stage twice. In one of them we are going to skip the shuffle, but in the other we are going to shuffle anyway. Therefore the asymptotic cost will be the same. If we apply spooling, the same amount of data will be shuffled but we would end up doing less work because the sender stage would be executed only once.

Thanks for the explanation! Might be useful to add this as a comment in the code as well.

yashmayya

LGTM, apologies for the delay!

yashmayya · 2025-01-07T10:18:51Z

...main/java/org/apache/pinot/query/planner/physical/colocated/GreedyShuffleRewriteVisitor.java

      Set<ColocationKey> colocationKeys;
-      if (canSkipShuffleBasic && areServersSuperset(node.getReceiverStageId(), node.getStageId())) {
+      if (canSkipShuffleBasic && allAreSuperSet(receiverStageIds, node)) {


Thanks for the explanation! Might be useful to add this as a comment in the code as well.

gortiz added 4 commits November 19, 2024 16:29

Support stage replacement

982ac13

Verify the case where a spool tries to send twice to the same stage

d02b5cb

Add tests that verifies that different distribution breaks equivalence

05b894e

Fix root mailbox send node build using null receiver stage

1597f5b

gortiz changed the title ~~Spool3~~ [Spool] Actual implementation Nov 20, 2024

gortiz requested review from yashmayya, vrajat and Jackie-Jiang November 20, 2024 13:40

gortiz added 3 commits November 20, 2024 17:05

Remove unused variable

91b1b55

Fix MailboxReceiveNode.equals, broken in previous spool commits

37168a0

Hide the receiver stage bit set in MailboxSendNode

fe22305

gortiz force-pushed the spool3 branch from 6d50b24 to 6658753 Compare November 20, 2024 17:23

gortiz mentioned this pull request Nov 20, 2024

[proposal] Reuse common expressions in a query #14196

Closed

8 tasks

gortiz and others added 4 commits November 21, 2024 09:08

Apply suggestions from code review

a72157f

rename some test methods Co-authored-by: Yash Mayya <[email protected]>

Make _senderStageId mutable to simplify getSenderStageId() logic

edf8542

Make StagesTestBase.Spool static

5bff873

Rename Spool as SpoolBuilder and add javadoc explaining the class

8263de3

gortiz force-pushed the spool3 branch from 6658753 to 6b56d2e Compare November 21, 2024 08:55

gortiz added 4 commits November 21, 2024 13:12

Fix error added in a previous commit

3b252cf

Support spool in runtime. Still not complete.

32fa712

Explain implementation is working, but not actual queries.

Fix design issues in multi-sender MailboxSendNode

f3bffb3

Add some explain physical plans with spool

f06138a

gortiz force-pushed the spool3 branch from 6b56d2e to f06138a Compare November 21, 2024 15:40

gortiz added 5 commits December 13, 2024 15:25

Merge branch 'master' into spool3

0b45a95

Add some logs in trace mode

cb7ff1f

Rename spool resource tests

7d63fbf

Fix error detected (and only affecting) tests. Improve logs

0479a2b

Merge remote-tracking branch 'origin/master' into spool3

53cebc8

gortiz marked this pull request as ready for review December 17, 2024 07:55

gortiz mentioned this pull request Dec 17, 2024

Spool4 #14672

Open

gortiz commented Dec 17, 2024

View reviewed changes

Merge branch 'master' into spool3

6907740

yashmayya reviewed Dec 21, 2024

View reviewed changes

yashmayya reviewed Dec 23, 2024

View reviewed changes

yashmayya added multi-stage Related to the multi-stage query engine feature release-notes Referenced by PRs that need attention when compiling the next release notes labels Dec 23, 2024

Change a log as suggested in PR

5ce60e5

gortiz added 2 commits December 23, 2024 15:41

Merge remote-tracking branch 'origin/master' into spool3

c4f5823

Change GreedyShuffleRewriteVisitor to be closer to the original one

cfeb8c4

Semantic should be the same, but the PR should be easier to read.

gortiz commented Jan 2, 2025

View reviewed changes

yashmayya approved these changes Jan 7, 2025

View reviewed changes

gortiz merged commit 3adcedb into apache:master Jan 7, 2025
21 checks passed

zeronerdzerogeekzerocool pushed a commit to zeronerdzerogeekzerocool/pinot that referenced this pull request Feb 20, 2025

[Spool] Actual implementation (apache#14507)

b77807a

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Spool] Actual implementation #14507

[Spool] Actual implementation #14507

gortiz commented Nov 20, 2024 •

edited

Loading

codecov-commenter commented Nov 20, 2024 •

edited

Loading

gortiz commented Dec 17, 2024

gortiz Dec 17, 2024

gortiz Dec 17, 2024

gortiz Dec 17, 2024

gortiz Dec 17, 2024

gortiz Dec 17, 2024

gortiz commented Dec 17, 2024

yashmayya left a comment

yashmayya Dec 23, 2024

gortiz Jan 2, 2025

gortiz commented Dec 23, 2024

gortiz Jan 2, 2025

yashmayya Jan 7, 2025

yashmayya left a comment

yashmayya Jan 7, 2025

[Spool] Actual implementation #14507

[Spool] Actual implementation #14507

Conversation

gortiz commented Nov 20, 2024 • edited Loading

codecov-commenter commented Nov 20, 2024 • edited Loading

Codecov Report

gortiz commented Dec 17, 2024

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

gortiz commented Dec 17, 2024

yashmayya left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

gortiz commented Dec 23, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

yashmayya left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

gortiz commented Nov 20, 2024 •

edited

Loading

codecov-commenter commented Nov 20, 2024 •

edited

Loading