[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 75.00000% with 59 lines in your changes missing coverage. Please review.

Project coverage is 64.02%. Comparing base (59551e4) to head (6907740).
Report is 1485 commits behind head on master.

Files with missing lines	Patch %	Lines
...hysical/colocated/GreedyShuffleRewriteVisitor.java	0.00%	19 Missing ⚠️
...query/runtime/operator/exchange/BlockExchange.java	48.38%	10 Missing and 6 partials ⚠️
...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	60.00%	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	87.50%	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%   64.02%   +2.27%     
- Complexity      207     1602    +1395     
============================================
  Files          2436     2707     +271     
  Lines        133233   149187   +15954     
  Branches      20636    22860    +2224     
============================================
+ Hits          82274    95519   +13245     
- Misses        44911    46656    +1745     
- Partials       6048     7012     +964

Flag	Coverage Δ
custom-integration1	`100.00% <ø> (+99.99%)`	⬆️
integration	`100.00% <ø> (+99.99%)`	⬆️
integration1	`100.00% <ø> (+99.99%)`	⬆️
integration2	`0.00% <ø> (ø)`
java-11	`64.00% <75.00%> (+2.29%)`	⬆️
java-21	`63.89% <75.00%> (+2.26%)`	⬆️
skip-bytebuffers-false	`64.01% <75.00%> (+2.27%)`	⬆️
skip-bytebuffers-true	`63.87% <75.00%> (+36.14%)`	⬆️
temurin	`64.02% <75.00%> (+2.27%)`	⬆️
unittests	`64.02% <75.00%> (+2.27%)`	⬆️
unittests1	`56.36% <74.35%> (+9.47%)`	⬆️
unittests2	`34.42% <3.38%> (+6.69%)`	⬆️

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)?

yashmayya · 2024-12-20T17:43:13Z

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

@@ -207,13 +207,17 @@ protected TransferableBlock getNextBlock()
      buildBroadcastHashTable();
    }
    if (_upstreamErrorBlock != null) {
+      LOGGER.trace("Returning upstream error block for join operator");


Are all the changes in this file unrelated to spools?

Yes, all the changes are just new logs I've added while debuging some blocking issues. These logs are disabled by default and should be easily removed by either the jit given they are not allocating. I recommend to keep them in case we need to debug something in the future.

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?

yashmayya · 2024-12-23T06:11:57Z

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

        LOGGER.debug("==[RECEIVE]== EOS received : " + _id + " in mailbox: " + removed.getId()
-            + " (" + _mailboxes.size() + " mailboxes alive)");
+            + " (" + ids + " mailboxes alive)");


Suggested change

+ " (" + ids + " mailboxes alive)");

+ " (Mailboxes alive: " + ids + ")");

...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.

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

Open

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

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 20, 2024

gortiz Dec 23, 2024

yashmayya Dec 23, 2024

yashmayya Dec 23, 2024

gortiz commented Dec 23, 2024

	+ " (" + ids + " mailboxes alive)");
	+ " (Mailboxes alive: " + ids + ")");

[Spool] Actual implementation #14507

Are you sure you want to change the base?

[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

Choose a reason for hiding this comment

Choose a reason for hiding this comment

gortiz commented Dec 23, 2024

gortiz commented Nov 20, 2024 •

edited

Loading

codecov-commenter commented Nov 20, 2024 •

edited

Loading