yashmayya commented on code in PR #15024:
URL: https://github.com/apache/pinot/pull/15024#discussion_r1950707569
##########
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:
Also I'm not sure I follow why the return function's input needs to be a
`List<SendingMailbox>` - we're not actually using the list contents to
determine the one to send to right? And the precondition check to verify that
the min stage index is less than the size of the list of sending mailboxes
doesn't seem to buy us much in terms of safety or user experience IMO.
##########
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:
nit: `pending` doesn't seem to descriptive enough IMO - what about something
like `stagesToRemove` or `fragmentsToRemove` instead?
##########
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:
nit: Can add an `assertNoError(jsonNode);` above this
##########
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:
I think the parser creates an empty list and not `null` in this case, so we
can probably remove this check and rely on only the next one (`stage4On3.size()
== 1`).
##########
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");
Review Comment:
Kudos for creating this query 😄
##########
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:
This feels a little convoluted - we're choosing the stage to send to via
information in the `MailboxSendNode` but then using the _index_ of this stage
in the receiver stage IDs list (which is actually maintained as a `BitSet` in
the `MailboxSendNode`) to pick the `SendingMailbox` from the list maintained in
the `BlockExchange`? In this case, I think it is working because we create that
list using the same receiver stage IDs iterable for spools where each
`SendingMailbox` is actually a `BlockExchange` itself. But this seems like a
fragile dependency that isn't cleanly enforced. Can we use the stage ID itself
instead of the index or is that not possible because the sending mailboxes
don't maintain stage metadata?
--
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
