gortiz commented on code in PR #15024:
URL: https://github.com/apache/pinot/pull/15024#discussion_r1950933276
##########
pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MailboxSendOperator.java:
##########
@@ -110,8 +110,40 @@ private static BlockExchange
getBlockExchange(OpChainExecutionContext ctx, Mailb
getBlockExchange(ctx, receiverStageId, node.getDistributionType(),
node.getKeys(), statMap, innerSplitter);
perStageSendingMailboxes.add(blockExchange.asSendingMailbox(Integer.toString(receiverStageId)));
}
+
+ Function<List<SendingMailbox>, Integer> statsIndexChooser =
getStatsIndexChooser(ctx, node);
return BlockExchange.getExchange(perStageSendingMailboxes,
RelDistribution.Type.BROADCAST_DISTRIBUTED,
- Collections.emptyList(), mainSplitter);
+ Collections.emptyList(), mainSplitter, statsIndexChooser);
+ }
+
+ private static Function<List<SendingMailbox>, Integer>
getStatsIndexChooser(OpChainExecutionContext ctx,
+ MailboxSendNode node) {
Review Comment:
I'm afraid mailboxes don't maintain the stage. Remember that initially
(until we added spools), all mailboxes of the same sender were sent to the same
stage.
Also, I don't see the problem here. It is fair to assume that the iterable
order will be repeatable. We can even specify that in the javadoc of
getReceiverStageIds.
We can also calculate the index in the same loop we are using to create the
mailboxes, but I think that would make the code more difficult to read.
The block I'm talking about starts on line 107 and contains the following
code:
```java
BlockSplitter innerSplitter = BlockSplitter.NO_OP;
for (int receiverStageId : node.getReceiverStageIds()) {
BlockExchange blockExchange =
getBlockExchange(ctx, receiverStageId, node.getDistributionType(),
node.getKeys(), statMap, innerSplitter);
perStageSendingMailboxes.add(blockExchange.asSendingMailbox(Integer.toString(receiverStageId)));
}
```
Alternatively, we can change the function to return the mailbox to which we
need to send stats instead of the index. However, in that case, we must change
BlockExchange to compare by equality instead of index.
##########
pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/SpoolIntegrationTest.java:
##########
@@ -0,0 +1,143 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.integration.tests;
+
+import com.fasterxml.jackson.databind.JsonNode;
+import com.jayway.jsonpath.DocumentContext;
+import com.jayway.jsonpath.JsonPath;
+import java.io.File;
+import java.util.List;
+import java.util.Map;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.data.Schema;
+import org.apache.pinot.spi.env.PinotConfiguration;
+import org.apache.pinot.spi.utils.CommonConstants;
+import org.apache.pinot.util.TestUtils;
+import org.testcontainers.shaded.org.apache.commons.io.FileUtils;
+import org.testng.Assert;
+import org.testng.annotations.AfterClass;
+import org.testng.annotations.BeforeClass;
+import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.Test;
+
+
+public class SpoolIntegrationTest extends BaseClusterIntegrationTest
+ implements ExplainIntegrationTestTrait {
+
+ @BeforeClass
+ public void setUp()
+ throws Exception {
+ TestUtils.ensureDirectoriesExistAndEmpty(_tempDir, _segmentDir, _tarDir);
+
+ // Start the Pinot cluster
+ startZk();
+ startController();
+ startBroker();
+ startServers(2);
+
+ // Create and upload the schema and table config
+ Schema schema = createSchema();
+ addSchema(schema);
+ TableConfig tableConfig = createOfflineTableConfig();
+ addTableConfig(tableConfig);
+
+ // Unpack the Avro files
+ List<File> avroFiles = unpackAvroData(_tempDir);
+
+ // Create and upload segments
+ ClusterIntegrationTestUtils.buildSegmentsFromAvro(avroFiles, tableConfig,
schema, 0, _segmentDir, _tarDir);
+ uploadSegments(getTableName(), _tarDir);
+
+ // Wait for all documents loaded
+ waitForAllDocsLoaded(600_000L);
+ }
+
+ protected void overrideBrokerConf(PinotConfiguration brokerConf) {
+ String property =
CommonConstants.MultiStageQueryRunner.KEY_OF_MULTISTAGE_EXPLAIN_INCLUDE_SEGMENT_PLAN;
+ brokerConf.setProperty(property, "true");
+ }
+
+ @BeforeMethod
+ public void resetMultiStage() {
+ setUseMultiStageQueryEngine(true);
+ }
+
+ // Test that intermediate stages can be spooled.
+ // In this case Stage 4 is an intermediate stage whose single child is stage
5.
+ // Stage 4 is spooled and sends data to stages 3 and 7
+ @Test
+ public void intermediateSpool()
+ throws Exception {
+ JsonNode jsonNode = postQuery("SET useSpools = true;\n"
+ + "WITH group_and_sum AS (\n"
+ + " SELECT ArrTimeBlk,\n"
+ + " Dest,\n"
+ + " SUM(ArrTime) AS ArrTime\n"
+ + " FROM mytable\n"
+ + " GROUP BY ArrTimeBlk,\n"
+ + " Dest\n"
+ + " limit 1000\n"
+ + "),\n"
+ + "aggregated_data AS (\n"
+ + " SELECT\n"
+ + " Dest,\n"
+ + " SUM(ArrTime) AS ArrTime\n"
+ + " FROM group_and_sum\n"
+ + " GROUP BY\n"
+ + " Dest\n"
+ + "),\n"
+ + "joined AS (\n"
+ + " SELECT\n"
+ + " s.Dest,\n"
+ + " s.ArrTime,\n"
+ + " (o.ArrTime) AS ArrTime2\n"
+ + " FROM group_and_sum s\n"
+ + " JOIN aggregated_data o\n"
+ + " ON s.Dest = o.Dest\n"
+ + ")\n"
+ + "SELECT *\n"
+ + "FROM joined\n"
+ + "LIMIT 1");
+ JsonNode stats = jsonNode.get("stageStats");
+ Assert.assertNotNull(stats, "Stage stats should be present. Please verify
the query didn't fail");
Review Comment:
done
##########
pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/SpoolIntegrationTest.java:
##########
@@ -0,0 +1,143 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.integration.tests;
+
+import com.fasterxml.jackson.databind.JsonNode;
+import com.jayway.jsonpath.DocumentContext;
+import com.jayway.jsonpath.JsonPath;
+import java.io.File;
+import java.util.List;
+import java.util.Map;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.data.Schema;
+import org.apache.pinot.spi.env.PinotConfiguration;
+import org.apache.pinot.spi.utils.CommonConstants;
+import org.apache.pinot.util.TestUtils;
+import org.testcontainers.shaded.org.apache.commons.io.FileUtils;
+import org.testng.Assert;
+import org.testng.annotations.AfterClass;
+import org.testng.annotations.BeforeClass;
+import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.Test;
+
+
+public class SpoolIntegrationTest extends BaseClusterIntegrationTest
+ implements ExplainIntegrationTestTrait {
+
+ @BeforeClass
+ public void setUp()
+ throws Exception {
+ TestUtils.ensureDirectoriesExistAndEmpty(_tempDir, _segmentDir, _tarDir);
+
+ // Start the Pinot cluster
+ startZk();
+ startController();
+ startBroker();
+ startServers(2);
+
+ // Create and upload the schema and table config
+ Schema schema = createSchema();
+ addSchema(schema);
+ TableConfig tableConfig = createOfflineTableConfig();
+ addTableConfig(tableConfig);
+
+ // Unpack the Avro files
+ List<File> avroFiles = unpackAvroData(_tempDir);
+
+ // Create and upload segments
+ ClusterIntegrationTestUtils.buildSegmentsFromAvro(avroFiles, tableConfig,
schema, 0, _segmentDir, _tarDir);
+ uploadSegments(getTableName(), _tarDir);
+
+ // Wait for all documents loaded
+ waitForAllDocsLoaded(600_000L);
+ }
+
+ protected void overrideBrokerConf(PinotConfiguration brokerConf) {
+ String property =
CommonConstants.MultiStageQueryRunner.KEY_OF_MULTISTAGE_EXPLAIN_INCLUDE_SEGMENT_PLAN;
+ brokerConf.setProperty(property, "true");
+ }
+
+ @BeforeMethod
+ public void resetMultiStage() {
+ setUseMultiStageQueryEngine(true);
+ }
+
+ // Test that intermediate stages can be spooled.
+ // In this case Stage 4 is an intermediate stage whose single child is stage
5.
+ // Stage 4 is spooled and sends data to stages 3 and 7
+ @Test
+ public void intermediateSpool()
+ throws Exception {
+ JsonNode jsonNode = postQuery("SET useSpools = true;\n"
+ + "WITH group_and_sum AS (\n"
+ + " SELECT ArrTimeBlk,\n"
+ + " Dest,\n"
+ + " SUM(ArrTime) AS ArrTime\n"
+ + " FROM mytable\n"
+ + " GROUP BY ArrTimeBlk,\n"
+ + " Dest\n"
+ + " limit 1000\n"
+ + "),\n"
+ + "aggregated_data AS (\n"
+ + " SELECT\n"
+ + " Dest,\n"
+ + " SUM(ArrTime) AS ArrTime\n"
+ + " FROM group_and_sum\n"
+ + " GROUP BY\n"
+ + " Dest\n"
+ + "),\n"
+ + "joined AS (\n"
+ + " SELECT\n"
+ + " s.Dest,\n"
+ + " s.ArrTime,\n"
+ + " (o.ArrTime) AS ArrTime2\n"
+ + " FROM group_and_sum s\n"
+ + " JOIN aggregated_data o\n"
+ + " ON s.Dest = o.Dest\n"
+ + ")\n"
+ + "SELECT *\n"
+ + "FROM joined\n"
+ + "LIMIT 1");
+ JsonNode stats = jsonNode.get("stageStats");
+ Assert.assertNotNull(stats, "Stage stats should be present. Please verify
the query didn't fail");
+ DocumentContext parsed = JsonPath.parse(stats.toString());
+ List<Map<String, Object>> stage4On3 = parsed.read("$..[?(@.stage ==
3)]..[?(@.stage == 4)]");
+ Assert.assertNotNull(stage4On3, "Stage 4 should be a descendant of stage
3");
Review Comment:
done
##########
pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/PlanFragmenter.java:
##########
@@ -89,12 +89,29 @@ private PlanNode process(PlanNode node, Context context) {
@Override
public void onSubstitution(int receiver, int oldSender, int newSender) {
+ // Change the sender of the receiver to the new sender
IntList senders = _childPlanFragmentIdsMap.get(receiver);
senders.rem(oldSender);
if (!senders.contains(newSender)) {
senders.add(newSender);
}
+
+ // Remove the old sender and its children from the plan fragment map
_planFragmentMap.remove(oldSender);
+
+ IntList pending = new IntArrayList();
Review Comment:
I've also changed the code to include the first iteration in the while loop
##########
pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/PlanFragmenter.java:
##########
@@ -89,12 +89,29 @@ private PlanNode process(PlanNode node, Context context) {
@Override
public void onSubstitution(int receiver, int oldSender, int newSender) {
+ // Change the sender of the receiver to the new sender
IntList senders = _childPlanFragmentIdsMap.get(receiver);
senders.rem(oldSender);
if (!senders.contains(newSender)) {
senders.add(newSender);
}
+
+ // Remove the old sender and its children from the plan fragment map
_planFragmentMap.remove(oldSender);
+
+ IntList pending = new IntArrayList();
Review Comment:
done
##########
pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MailboxSendOperator.java:
##########
@@ -110,8 +110,40 @@ private static BlockExchange
getBlockExchange(OpChainExecutionContext ctx, Mailb
getBlockExchange(ctx, receiverStageId, node.getDistributionType(),
node.getKeys(), statMap, innerSplitter);
perStageSendingMailboxes.add(blockExchange.asSendingMailbox(Integer.toString(receiverStageId)));
}
+
+ Function<List<SendingMailbox>, Integer> statsIndexChooser =
getStatsIndexChooser(ctx, node);
return BlockExchange.getExchange(perStageSendingMailboxes,
RelDistribution.Type.BROADCAST_DISTRIBUTED,
- Collections.emptyList(), mainSplitter);
+ Collections.emptyList(), mainSplitter, statsIndexChooser);
+ }
+
+ private static Function<List<SendingMailbox>, Integer>
getStatsIndexChooser(OpChainExecutionContext ctx,
+ MailboxSendNode node) {
Review Comment:
It is not needed for this implementation but for the other implementation
used by default, which returns a random index from the list. I think it makes
sense that the strategy takes the list of mailboxes as an argument.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscr...@pinot.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscr...@pinot.apache.org
For additional commands, e-mail: commits-h...@pinot.apache.org
