[spool] Define equivalence between stages #14296

gortiz · 2024-10-24T14:26:04Z

This PR includes the code that defines when two stages are equivalent and it is the first step to implement #14196.

Instead of implementing the ability to reuse common expressions in a single and very large PR, I decided to create this first one where:

The new code is not being called by production code (yet)
Classes are well tested and documented

My hope is that this will be easier to review.

gortiz · 2024-10-24T14:26:24Z

cc @bziobrowski @vrajat

codecov-commenter · 2024-10-24T15:20:28Z

Codecov Report

Attention: Patch coverage is 44.60784% with 113 lines in your changes missing coverage. Please review.

Project coverage is 63.75%. Comparing base (59551e4) to head (7758934).
Report is 1307 commits behind head on master.

Files with missing lines	Patch %	Lines
.../query/planner/logical/EquivalentStagesFinder.java	40.17%	52 Missing and 18 partials ⚠️
.../pinot/query/planner/plannode/PlanNodeVisitor.java	40.00%	20 Missing and 1 partial ⚠️
...che/pinot/query/planner/logical/GroupedStages.java	75.00%	7 Missing and 3 partials ⚠️
...anner/logical/ParentToChildrenStageCalculator.java	0.00%	8 Missing ⚠️
.../apache/pinot/query/planner/plannode/PlanNode.java	0.00%	3 Missing ⚠️
.../pinot/query/planner/plannode/MailboxSendNode.java	0.00%	1 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff              @@
##             master   #14296      +/-   ##
============================================
+ Coverage     61.75%   63.75%   +2.00%     
- Complexity      207     1569    +1362     
============================================
  Files          2436     2665     +229     
  Lines        133233   146223   +12990     
  Branches      20636    22432    +1796     
============================================
+ Hits          82274    93221   +10947     
- Misses        44911    46087    +1176     
- Partials       6048     6915     +867

Flag	Coverage Δ
custom-integration1	`100.00% <ø> (+99.99%)`	⬆️
integration	`100.00% <ø> (+99.99%)`	⬆️
integration1	`100.00% <ø> (+99.99%)`	⬆️
integration2	`0.00% <ø> (ø)`
java-11	`63.71% <44.60%> (+2.00%)`	⬆️
java-21	`63.64% <44.60%> (+2.01%)`	⬆️
skip-bytebuffers-false	`63.74% <44.60%> (+1.99%)`	⬆️
skip-bytebuffers-true	`63.59% <44.60%> (+35.87%)`	⬆️
temurin	`63.75% <44.60%> (+2.00%)`	⬆️
unittests	`63.74% <44.60%> (+2.00%)`	⬆️
unittests1	`55.51% <44.60%> (+8.62%)`	⬆️
unittests2	`34.09% <0.00%> (+6.36%)`	⬆️

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.

vrajat · 2024-10-24T17:07:39Z

...planner/src/test/java/org/apache/pinot/query/planner/logical/EquivalentStagesFinderTest.java

+import static org.testng.Assert.*;
+
+
+public class EquivalentStagesFinderTest extends StagesTestBase {


Confused by the these tests. Are the tests dependent on each other ? Every test adds another stage to _stages ?

It is explained in StagesTestBase. Basically each test starts from an clean environment where these _stageRoots is empty. That way each test is self-contained and easier to read.

Suggestion: We can rename plan() method to when(), which is a more common name in Given-When-Then tests. Do you think that would make it easier to read?

Do you think plan() is better? It is worth the change?

I get the logic now. when() is fine. What confused me initially is that I missed _stageRoots.clear() call after every test.

vrajat · 2024-10-24T17:17:01Z

pinot-query-planner/src/main/java/org/apache/pinot/query/planner/plannode/PlanNodeVisitor.java

@@ -63,4 +63,150 @@ public interface PlanNodeVisitor<T, C> {
  T visitExchange(ExchangeNode exchangeNode, C context);

  T visitExplained(ExplainedNode node, C context);
+
+  /**


Why did you choose to add this class to the same file ?

Because it is closely related to this visitor. Also I don't need to call it PlanNodeDepthFirstVisitor given it is already contained inside PlanNodeVisitor. But I'm open to move as a regular class if it is preferred.

vrajat · 2024-10-24T19:38:56Z

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

+     * same nodes multiple times. The side effect of that is that the second argument for {@link #areEquivalent} must be
+     * a node that was already visited.
+     */
+    private class NodeEquivalence implements PlanNodeVisitor<Boolean, PlanNode> {


I need to step through patch to understand. Where is the return value (Boolean) checked ?
Also when are methods like visitAggregate called as I cannot find visitChildren calls in the NodeEquivalence class.

Where is the return value (Boolean) checked ?

The entry point here is areEquivalent. That method returns the Boolean returned by each visit method. And areEquivalent is being called by Visitor in its visitMailboxSend.

Also when are methods like visitAggregate called as I cannot find visitChildren calls in the NodeEquivalence class.

There is no visitChildren in this method. NodeEquivalence is a normal visitor, not a DepthFirstVisitor. Recursion is done in baseNode, which is called by all visit specific methods.

yashmayya

Posting a partial review, I'm yet to take a look at the tests.

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

pinot-query-planner/src/main/java/org/apache/pinot/query/planner/plannode/MailboxSendNode.java

...planner/src/main/java/org/apache/pinot/query/planner/logical/ParentToChildrenCalculator.java

pinot-query-planner/src/main/java/org/apache/pinot/query/planner/plannode/PlanNodeVisitor.java

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

yashmayya · 2024-10-25T11:49:32Z

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

+      return group;
+    }
+
+    public Mutable removeStage(MailboxSendNode stage) {


When would we need to remove a stage from the set of grouped stages?

In the prune phase I was talking about before, although I'm not sure if that will be done in a copy fashion or actually modifying the object.

Basically in the next phase we are going to need to transforma a GroupedStages into a simplified expression tree where for each equivalent group G we create a new synthetic stage that send data to each stage g that belongs to G. But imagine a case like:

flowchart LR root --> S1 S1 --> S2 S1 --> S5 S2 --> S3[S3 Read A] S2 --> S4[S4 Read B] S5 --> S6[S6 Read A] S5 --> S7[S7 Read B]

Loading

Where:

Stage 3 is equivalent to Stage 6

Stage 4 is equivalent to Stage 7

Stage 2 is equivalent to Stage 5

We wouldn't want to create synthetic classes for each of these equivalence groups because in the end all (S3, S6) and (S4, S7) could be pruned once (S2, S5) creates its own equivalence.

The final plan should be something like:

flowchart LR root --> S1 S1 --> S2 S1 --> S5 S2 --> N1 S5 --> N1 N1 --> N2[N2 Read A] N1 --> N3[N2 Read B]

Loading

Where Ns are the new synthetic stages.

One way to implement this is to remove stages that will be pruned it its parent is equivalent. In our case S3, S4, S6 and S7 would be pruned.

I'm not 100% sure if this method will be called in the final, but at least in my first draft it is useful

If S2 and S5 are equivalent, why will both of them be retained in the final plan?

Good question. We need to keep them, but we also modify them. It is explained in the design document, but basically there are cases where we need to keep them. For example imagine that the sender node in S2 (the one that sends to S1) uses hash distribution while the sender of S5 sends using random. We could consider them not equivalent but it is better to just keep both stages, schedule them in a way they are running in the same servers than N1 and then apply each distribution.

Wouldn't they still be replaced by new stages? The new ones will retain the sender nodes from the original but the rest of the stage will be replaced by a single receive node which gets data from the synthetic stage that contains all the other nodes (apart from the sender) from the equivalent stages right?

Yes, S2 and S5 will change internally so we could consider them new as well. It helps me to think that they are not new but they have been modified, but this is just a concept level, not programatically (they will be different objects)

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

yashmayya

Thanks for the new test framework for query plan stages @gortiz! The tests using the framework are so pleasant to read 😄

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

yashmayya · 2024-10-29T10:46:43Z

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

+      return group;
+    }
+
+    public Mutable removeStage(MailboxSendNode stage) {


If S2 and S5 are equivalent, why will both of them be retained in the final plan?

pinot-query-planner/src/test/java/org/apache/pinot/query/planner/logical/StagesTestBase.java

yashmayya · 2024-10-29T11:55:11Z

pinot-query-planner/src/test/java/org/apache/pinot/query/planner/logical/GroupedStagesTest.java

+    when(
+        join(
+            exchange(1, tableScan("T1")),
+            exchange(2, tableScan("T2"))
+        )
+    );
+    GroupedStages.Mutable mutable = new GroupedStages.Mutable()
+        .addNewGroup(stage(0))
+        .addNewGroup(stage(1))
+        .addToGroup(stage(0), stage(2));
+    assertEquals(mutable.toString(), "[[0, 2], [1]]");


The tests in these classes probably don't need to use the StagesTestBase framework right? For instance, here it's pretty confusing because stage 0 and stage 2 are not actually equivalent?

I don't follow you. StagesTestBase is just an API to create these nodes. It doesn't care about correctness. Same here. GroupedStagesTest doesn't care about what being equivalent mean. It is just a matemathical partition of a set of nodes. How nodes are partition is defined in EquivalentStagesFinder.

I understand that but my point was that the test looks misleading because we're asserting that the send node above the join node is equivalent to the send node above one of the join node's input table scan nodes. Since we're only testing GroupedStages here and we don't care about the definition of equivalence like you said, wouldn't it be clearer to use synthetic MailboxSendNodes here (mocks or otherwise)?

I think we have a different perspective. This tests is not testing the semantic of the equivalence. It is just testing the API of GroupedStages. In here you can do crazy things that doesn't make any sense. We are just testing that the API works in the same way we can create a test on a java.util.list that verifies that if we have a list of [1, 2] and we call add(1, 3) the result will be [1, 3, 2]. Whether that list is later being used to store ints order by their value or not is not important for the list and the same happens in this GroupedStages class.

pinot-query-planner/src/test/java/org/apache/pinot/query/planner/logical/StagesTestBase.java

yashmayya · 2024-10-29T13:37:03Z

pinot-query-planner/src/test/java/org/apache/pinot/query/planner/logical/StagesTestBase.java

+   * {@link #tableScan(String)} and others to chain instances of this class.
+   */
+  public interface ChildBuilder<P extends PlanNode> {
+    P build(int stageId, @Nullable DataSchema dataSchema, @Nullable PlanNode.NodeHint hints);


This isn't currently being used anywhere directly? All external uses seem to be through SimpleChildBuilder?

No, it is not used externally and probably should never be. But we need it to create the machinary that runs withDataSchema and withHints

yashmayya · 2024-10-29T13:43:32Z

I think this branch needs to be rebased / merged with master, all the builds are currently failing (at least one reason seems to be conflict with #14294).

...ava/org/apache/pinot/segment/local/segment/index/readerwriter/ValueReaderComparisonTest.java

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

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

yashmayya · 2024-10-30T06:20:56Z

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

+      return group;
+    }
+
+    public Mutable removeStage(MailboxSendNode stage) {


Wouldn't they still be replaced by new stages? The new ones will retain the sender nodes from the original but the rest of the stage will be replaced by a single receive node which gets data from the synthetic stage that contains all the other nodes (apart from the sender) from the equivalent stages right?

pinot-query-planner/src/test/java/org/apache/pinot/query/planner/logical/StagesTestBase.java

yashmayya

LGTM, I just have one optional suggestion on GroupedStagesTest.

yashmayya · 2024-11-07T12:21:51Z

pinot-query-planner/src/test/java/org/apache/pinot/query/planner/logical/GroupedStagesTest.java

+    when(
+        join(
+            exchange(1, tableScan("T1")),
+            exchange(2, tableScan("T2"))
+        )
+    );
+    GroupedStages.Mutable mutable = new GroupedStages.Mutable()
+        .addNewGroup(stage(0))
+        .addNewGroup(stage(1))
+        .addToGroup(stage(0), stage(2));
+    assertEquals(mutable.toString(), "[[0, 2], [1]]");


I understand that but my point was that the test looks misleading because we're asserting that the send node above the join node is equivalent to the send node above one of the join node's input table scan nodes. Since we're only testing GroupedStages here and we don't care about the definition of equivalence like you said, wouldn't it be clearer to use synthetic MailboxSendNodes here (mocks or otherwise)?

vrajat · 2024-11-07T12:50:53Z

...planner/src/test/java/org/apache/pinot/query/planner/logical/EquivalentStagesFinderTest.java

+        )
+    );
+    GroupedStages result = EquivalentStagesFinder.findEquivalentStages(stage(0));
+    assertEquals(result.toString(), "[[0], [1, 2]]");


Does this notation mean that 1, 2 are equivalent and therefore grouped together ? In the previous test, [[0], [1], [2]] implies that there no equivalent stages ?

Exactly (to both questions). And each group is sorted by their first element (that is why [0] is shown before [1, 2], because 0 < 1).

I hope tests are easy to read once that is clear.

…ulator

Add some classes to calculate stage equivalence

gortiz requested review from Jackie-Jiang and yashmayya October 24, 2024 14:26

gortiz force-pushed the spool branch from f5c8319 to 6e4b9a4 Compare October 24, 2024 14:43

gortiz mentioned this pull request Oct 24, 2024

[proposal] Reuse common expressions in a query #14196

Closed

8 tasks

vrajat reviewed Oct 24, 2024

View reviewed changes

yashmayya reviewed Oct 28, 2024

View reviewed changes

yashmayya reviewed Oct 29, 2024

View reviewed changes

gortiz force-pushed the spool branch from 87ddc87 to 9be76e4 Compare October 29, 2024 15:12

yashmayya reviewed Oct 30, 2024

View reviewed changes

yashmayya approved these changes Nov 7, 2024

View reviewed changes

vrajat reviewed Nov 7, 2024

View reviewed changes

vrajat approved these changes Nov 8, 2024

View reviewed changes

gortiz added 14 commits November 11, 2024 09:46

spool: Add some classes to calculate stage equivalence

2a857b9

spool: Add javadoc to DepthFirstVisitor

5bd9fe6

spool: Fix cleanup annotation

574bfa0

spool: Rename the builder method from plan to when

b2ded33

spool: Apply suggestions proposed during review

11f9a03

spool: Simplify NodeEquivalence.visitMailboxSend

b9014b4

spool: Rename ParentToChildrenCalculator as ParentToChildrenStageCalc…

6d7fa99

…ulator

spool: Rename a hook method in StagesTestBase

27dec68

spool: Add some @FunctionalInterface annotations

6a31424

spool: Fix an error in exchange SimpleChildBuilder

43216da

spool: Fix method duplicated during merge

0edf0f2

spool: Fix another incorrectly merged line

d2b4864

spool: Fix another incorrectly merged line

0682023

Remove unintended change

7758934

gortiz force-pushed the spool branch from 0b5746c to 7758934 Compare November 11, 2024 08:46

Jackie-Jiang added feature multi-stage Related to the multi-stage query engine labels Nov 16, 2024

Jackie-Jiang approved these changes Nov 16, 2024

View reviewed changes

gortiz merged commit f61f47b into apache:master Nov 18, 2024
21 checks passed

This was referenced Nov 19, 2024

[Spool] Introduce stage replacer and change send nodes to be able to send to more than one stage #14495

Merged

[Spool] Actual implementation #14507

Merged

davecromberge pushed a commit to davecromberge/pinot that referenced this pull request Nov 22, 2024

[spool] Define equivalence between stages (apache#14296)

e5a3076

Add some classes to calculate stage equivalence

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[spool] Define equivalence between stages #14296

[spool] Define equivalence between stages #14296

gortiz commented Oct 24, 2024

gortiz commented Oct 24, 2024

codecov-commenter commented Oct 24, 2024 •

edited

Loading

vrajat Oct 24, 2024

gortiz Oct 25, 2024

gortiz Nov 7, 2024

vrajat Nov 7, 2024

vrajat Oct 24, 2024

gortiz Oct 25, 2024

vrajat Oct 24, 2024

gortiz Oct 25, 2024

yashmayya left a comment

yashmayya Oct 25, 2024

gortiz Oct 28, 2024

yashmayya Oct 29, 2024

gortiz Oct 29, 2024

yashmayya Oct 30, 2024

gortiz Nov 7, 2024

yashmayya left a comment

yashmayya Oct 29, 2024

yashmayya Oct 29, 2024

gortiz Oct 29, 2024

yashmayya Nov 7, 2024

gortiz Nov 7, 2024

yashmayya Oct 29, 2024

gortiz Oct 29, 2024

yashmayya commented Oct 29, 2024

yashmayya Oct 30, 2024

yashmayya left a comment

yashmayya Nov 7, 2024

vrajat Nov 7, 2024

gortiz Nov 7, 2024

		import static org.testng.Assert.*;


		public class EquivalentStagesFinderTest extends StagesTestBase {

[spool] Define equivalence between stages #14296

[spool] Define equivalence between stages #14296

Conversation

gortiz commented Oct 24, 2024

gortiz commented Oct 24, 2024

codecov-commenter commented Oct 24, 2024 • edited Loading

Codecov Report

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

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

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

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

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

yashmayya commented Oct 29, 2024

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

Choose a reason for hiding this comment

Choose a reason for hiding this comment

codecov-commenter commented Oct 24, 2024 •

edited

Loading