You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@druid.apache.org by "cryptoe (via GitHub)" <gi...@apache.org> on 2023/09/15 07:00:05 UTC

[PR] Limit pages size to a configurable limit (druid)

cryptoe opened a new pull request, #14994:
URL: https://github.com/apache/druid/pull/14994

   Adding the ability to limit the pages sizes of select queries. 
   
   * We piggyback on the same machinery that is used to control the `numRowsPerSegment`. 
   * This patch introduces a new context parameter `rowsPerPage` for which the default value is set to `100000` rows. 
   * This patch also optimizes adding the last `selectResults` stage only when the previous stages have sorted outputs. Currently for each select query with `selectDestination=durableStorage`, we used to add this extra `selectResults` stage. 
   
   
   
   
   This PR has:
   
   - [x] been self-reviewed.
      - [ ] using the [concurrency checklist](https://github.com/apache/druid/blob/master/dev/code-review/concurrency.md) (Remove this item if the PR doesn't have any relation to concurrency.)
   - [x] added documentation for new or modified features or behaviors.
   - [] a release note entry in the PR description.
   - [x] added Javadocs for most classes and all non-trivial methods. Linked related entities via Javadoc links.
   - [ ] added or updated version, license, or notice information in [licenses.yaml](https://github.com/apache/druid/blob/master/dev/license.md)
   - [x] added comments explaining the "why" and the intent of the code wherever would not be obvious for an unfamiliar reader.
   - [x] added unit tests or modified existing tests to cover new code paths, ensuring the threshold for [code coverage](https://github.com/apache/druid/blob/master/dev/code-review/code-coverage.md) is met.
   - [ ] added integration tests.
   - [x] been tested in a test Druid cluster.
   


-- 
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-unsubscribe@druid.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org


Re: [PR] Limit pages size to a configurable limit (druid)

Posted by "LakshSingla (via GitHub)" <gi...@apache.org>.
LakshSingla commented on PR #14994:
URL: https://github.com/apache/druid/pull/14994#issuecomment-1960295243

   @marzi312 It is happening because you have a LIMIT in front of your query. It is a bug in the MSQ engine. I'll raise a patch for the same.


-- 
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-unsubscribe@druid.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org


Re: [PR] Limit pages size to a configurable limit (druid)

Posted by "LakshSingla (via GitHub)" <gi...@apache.org>.
LakshSingla commented on PR #14994:
URL: https://github.com/apache/druid/pull/14994#issuecomment-1758123496

   Okay I looked at `SqlStatementResourceHelper#populatePageList`, and we are using counters during the execution time to determine the page list. Shouldn't that be risky, considering counters are a best-effort user-facing statistics for the job and not something we can rely upon. 


-- 
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-unsubscribe@druid.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org


Re: [PR] Limit pages size to a configurable limit (druid)

Posted by "cryptoe (via GitHub)" <gi...@apache.org>.
cryptoe merged PR #14994:
URL: https://github.com/apache/druid/pull/14994


-- 
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-unsubscribe@druid.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org


Re: [PR] Limit pages size to a configurable limit (druid)

Posted by "cryptoe (via GitHub)" <gi...@apache.org>.
cryptoe commented on PR #14994:
URL: https://github.com/apache/druid/pull/14994#issuecomment-1759062612

   > > @LakshSingla that test case is already added if you see SqlStatementResourceHelperTest#testDistinctPartitionsOnEachWorker() .
   > 
   > Also, if possible please add a MSQ select test for the regression. Using counters is risky as is and we'd wanna make sure that we are tackling the issue we are seeing E2E.
   
   Added another test `SqlMSQStatementResourcePostTest#testMultipleWorkersWithPageSizeLimiting` which covers that case e2e. 


-- 
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-unsubscribe@druid.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org


Re: [PR] Limit pages size to a configurable limit (druid)

Posted by "LakshSingla (via GitHub)" <gi...@apache.org>.
LakshSingla commented on code in PR #14994:
URL: https://github.com/apache/druid/pull/14994#discussion_r1331928033


##########
extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/ControllerImpl.java:
##########
@@ -1556,8 +1556,13 @@ private static QueryDefinition makeQueryDefinition(
       shuffleSpecFactory = ShuffleSpecFactories.singlePartition();
       queryToPlan = querySpec.getQuery();
     } else if (querySpec.getDestination() instanceof DurableStorageMSQDestination) {
-      // we add a final stage which generates one partition per worker.
-      shuffleSpecFactory = ShuffleSpecFactories.globalSortWithMaxPartitionCount(tuningConfig.getMaxNumWorkers());
+
+      shuffleSpecFactory = (clusterBy, aggregate) ->
+          new GlobalSortTargetSizeShuffleSpec(
+              clusterBy,
+              MultiStageQueryContext.getRowsPerPage(querySpec.getQuery().context()),
+              aggregate
+          );

Review Comment:
   Let's refactor this lambda to ShuffleSpecFactories#globalSortWithTargetSize and change the other instance here https://github.com/apache/druid/blob/79f882f48c15eef77d627f1d31212d6c09cae315/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/ControllerImpl.java#L1585 to use that factory as well. 



##########
docs/multi-stage-query/reference.md:
##########
@@ -246,6 +246,7 @@ The following table lists the context parameters for the MSQ task engine:
 | `durableShuffleStorage` | SELECT, INSERT, REPLACE <br /><br />Whether to use durable storage for shuffle mesh. To use this feature, configure the durable storage at the server level using `druid.msq.intermediate.storage.enable=true`). If these properties are not configured, any query with the context variable `durableShuffleStorage=true` fails with a configuration error. <br /><br />                                                                                                                                                                                                                                                                                                                                                                                                                   | `false` |
 | `faultTolerance` | SELECT, INSERT, REPLACE<br /><br /> Whether to turn on fault tolerance mode or not. Failed workers are retried based on [Limits](#limits). Cannot be used when `durableShuffleStorage` is explicitly set to false.                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                         | `false` |
 | `selectDestination` | SELECT<br /><br /> Controls where the final result of the select query is written. <br />Use `taskReport`(the default) to write select results to the task report. <b> This is not scalable since task reports size explodes for large results </b> <br/>Use `durableStorage` to write results to durable storage location. <b>For large results sets, its recommended to use `durableStorage` </b>. To configure durable storage see [`this`](#durable-storage) section.                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                          | `taskRep
 ort` |
+| `rowsPerPage` | SELECT<br /><br />The number of rows per page to target. The actual number of rows per page may be somewhat higher or lower than this number. In most cases, use the default.<br /> This property comes into effect only when `selectDestination` is set to `durableStorage`                                                                                                                                                                                                                                                                                                                                                                                                                                                    | 100000 |

Review Comment:
   Does `page` mean something for the end user reading the documentation? If not, I'd suggest
   * Use some commonly used cloud terminology in the explanation of the parameter - Maximum number of rows that can be present in a single object stored in the durable storage
   * We should mention when this should be set by the user - and in our case, one good use case would be if the file/object size exceeds the maximum allowable limit. 
   * Perhaps change the context parameter to something else, though I am fine with it, since this shouldn't be used by the user very frequently.



-- 
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-unsubscribe@druid.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org


Re: [PR] Limit pages size to a configurable limit (druid)

Posted by "cryptoe (via GitHub)" <gi...@apache.org>.
cryptoe commented on code in PR #14994:
URL: https://github.com/apache/druid/pull/14994#discussion_r1354616728


##########
docs/multi-stage-query/reference.md:
##########
@@ -246,6 +246,7 @@ The following table lists the context parameters for the MSQ task engine:
 | `durableShuffleStorage` | SELECT, INSERT, REPLACE <br /><br />Whether to use durable storage for shuffle mesh. To use this feature, configure the durable storage at the server level using `druid.msq.intermediate.storage.enable=true`). If these properties are not configured, any query with the context variable `durableShuffleStorage=true` fails with a configuration error. <br /><br />                                                                                                                                                                                                                                                                                                                                                                                                                   | `false` |
 | `faultTolerance` | SELECT, INSERT, REPLACE<br /><br /> Whether to turn on fault tolerance mode or not. Failed workers are retried based on [Limits](#limits). Cannot be used when `durableShuffleStorage` is explicitly set to false.                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                         | `false` |
 | `selectDestination` | SELECT<br /><br /> Controls where the final result of the select query is written. <br />Use `taskReport`(the default) to write select results to the task report. <b> This is not scalable since task reports size explodes for large results </b> <br/>Use `durableStorage` to write results to durable storage location. <b>For large results sets, its recommended to use `durableStorage` </b>. To configure durable storage see [`this`](#durable-storage) section.                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                          | `taskRep
 ort` |
+| `rowsPerPage` | SELECT<br /><br />The number of rows per page to target. The actual number of rows per page may be somewhat higher or lower than this number. In most cases, use the default.<br /> This property comes into effect only when `selectDestination` is set to `durableStorage`                                                                                                                                                                                                                                                                                                                                                                                                                                                    | 100000 |

Review Comment:
   page is a concept explained here : https://druid.apache.org/docs/latest/querying/query-deep-storage/#get-query-status



-- 
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-unsubscribe@druid.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org


Re: [PR] Limit pages size to a configurable limit (druid)

Posted by "cryptoe (via GitHub)" <gi...@apache.org>.
cryptoe commented on code in PR #14994:
URL: https://github.com/apache/druid/pull/14994#discussion_r1356112683


##########
extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/ControllerImpl.java:
##########
@@ -1693,8 +1687,9 @@ private static QueryDefinition makeQueryDefinition(
       shuffleSpecFactory = ShuffleSpecFactories.singlePartition();
       queryToPlan = querySpec.getQuery();
     } else if (querySpec.getDestination() instanceof DurableStorageMSQDestination) {
-      // we add a final stage which generates one partition per worker.
-      shuffleSpecFactory = ShuffleSpecFactories.globalSortWithMaxPartitionCount(tuningConfig.getMaxNumWorkers());
+      shuffleSpecFactory = ShuffleSpecFactories.getGlobalSortWithTargetSize(
+          MultiStageQueryContext.getRowsPerPage(querySpec.getQuery().context())

Review Comment:
   GlobalSortTargetSizeShuffleSpec enforces a limit on partition size globally yes. 
   
   So there can be 2 cases: 
   1. If the last stage is group by post shuffle, then we know that each partition will only be present on distinct worker only. Hence the page size will control the number of rows in that partition. 
   
   2. If the last stage is scanStage, then we add a new `QueryResultFrameProcessor` since data needs to be sorted on the boost column. The queryResultFrameProcessor will merge the result in the same partition and write out a single partition. Since the partition cuts on sizes are done globally, in the controller, we would have the final partition equal to the page size configured. 
   
   



-- 
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-unsubscribe@druid.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org


Re: [PR] Limit pages size to a configurable limit (druid)

Posted by "cryptoe (via GitHub)" <gi...@apache.org>.
cryptoe commented on PR #14994:
URL: https://github.com/apache/druid/pull/14994#issuecomment-1758021185

   @LakshSingla that test case is already added if you see SqlStatementResourceHelperTest#testDistinctPartitionsOnEachWorker() . 
   
   


-- 
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-unsubscribe@druid.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org


Re: [PR] Limit pages size to a configurable limit (druid)

Posted by "marzi312 (via GitHub)" <gi...@apache.org>.
marzi312 commented on PR #14994:
URL: https://github.com/apache/druid/pull/14994#issuecomment-1959955932

   hey @cryptoe 
   I've tested this feature with a curl command, POST to ``/v2/sql/statements`` but it completely ignores `rowsPerPage ` parameter. 
   Example: 
   `curl -XPOST -H "Content-Type: application/json" -d '{"context": {"includeSegmentSource":"REALTIME","selectDestination":"DURABLESTORAGE", "executionMode":"ASYNC", "durableShuffleStorage": "true", "rowsPerPage":"1000"}, "query":"SELECT * FROM my_table limit 200000"}' https://<url>/druid/v2/sql/statements
   `
   
   `rowsPerPage` default value is 100000 but it seems it even ignores the default value for me. 
    I tested this with Druid V29 
    
    and here's the results from get query status endpoint
   `...."pages":[{"id":0,"numRows":200000,"sizeInBytes":123026540}]}`
   


-- 
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-unsubscribe@druid.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org


Re: [PR] Limit pages size to a configurable limit (druid)

Posted by "adarshsanjeev (via GitHub)" <gi...@apache.org>.
adarshsanjeev commented on code in PR #14994:
URL: https://github.com/apache/druid/pull/14994#discussion_r1355075260


##########
extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/ControllerImpl.java:
##########
@@ -1772,27 +1767,32 @@ private static QueryDefinition makeQueryDefinition(
       return queryDef;
     } else if (querySpec.getDestination() instanceof DurableStorageMSQDestination) {
 
-      // attaching new query results stage always.
+      // attaching new query results stage if the final stage does sort during shuffle so that results are ordered.
       StageDefinition finalShuffleStageDef = queryDef.getFinalStageDefinition();
-      final QueryDefinitionBuilder builder = QueryDefinition.builder();
-      for (final StageDefinition stageDef : queryDef.getStageDefinitions()) {
-        builder.add(StageDefinition.builder(stageDef));
-      }
-
-      builder.add(StageDefinition.builder(queryDef.getNextStageNumber())
-                                 .inputs(new StageInputSpec(queryDef.getFinalStageDefinition().getStageNumber()))
-                                 .maxWorkerCount(tuningConfig.getMaxNumWorkers())
-                                 .signature(finalShuffleStageDef.getSignature())
-                                 .shuffleSpec(null)
-                                 .processorFactory(new QueryResultFrameProcessorFactory())
-      );
+      if (finalShuffleStageDef.doesSortDuringShuffle()) {
+        final QueryDefinitionBuilder builder = QueryDefinition.builder();
+        for (final StageDefinition stageDef : queryDef.getStageDefinitions()) {

Review Comment:
   nit: could use org.apache.druid.msq.kernel.QueryDefinitionBuilder#addAll(org.apache.druid.msq.kernel.QueryDefinition) here



##########
extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/ControllerImpl.java:
##########
@@ -1693,8 +1687,9 @@ private static QueryDefinition makeQueryDefinition(
       shuffleSpecFactory = ShuffleSpecFactories.singlePartition();
       queryToPlan = querySpec.getQuery();
     } else if (querySpec.getDestination() instanceof DurableStorageMSQDestination) {
-      // we add a final stage which generates one partition per worker.
-      shuffleSpecFactory = ShuffleSpecFactories.globalSortWithMaxPartitionCount(tuningConfig.getMaxNumWorkers());
+      shuffleSpecFactory = ShuffleSpecFactories.getGlobalSortWithTargetSize(
+          MultiStageQueryContext.getRowsPerPage(querySpec.getQuery().context())

Review Comment:
   I might be missing something in somewhere else in this PR, but doesn't GlobalSortTargetSizeShuffleSpec enforce the limit on the total partition size summed across all workers? Since we create a new page for each worker parition combination, would the limit be enforced?



-- 
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-unsubscribe@druid.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org


Re: [PR] Limit pages size to a configurable limit (druid)

Posted by "LakshSingla (via GitHub)" <gi...@apache.org>.
LakshSingla commented on PR #14994:
URL: https://github.com/apache/druid/pull/14994#issuecomment-1758984326

   > @LakshSingla that test case is already added if you see SqlStatementResourceHelperTest#testDistinctPartitionsOnEachWorker() .
   
   Also, if possible please add a MSQ select test for the regression. Using counters is risky as is and we'd wanna make sure that we are tackling the issue we are seeing E2E. 


-- 
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-unsubscribe@druid.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org


Re: [PR] Limit pages size to a configurable limit (druid)

Posted by "LakshSingla (via GitHub)" <gi...@apache.org>.
LakshSingla commented on code in PR #14994:
URL: https://github.com/apache/druid/pull/14994#discussion_r1332632293


##########
extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQSelectTest.java:
##########
@@ -215,16 +216,6 @@ public void testSelectOnFoo()
                 .with().totalFiles(1),
             0, 0, "input0"
         )
-        .setExpectedCountersForStageWorkerChannel(
-            CounterSnapshotMatcher
-                .with().rows(6).frames(1),
-            0, 0, "output"
-        )
-        .setExpectedCountersForStageWorkerChannel(
-            CounterSnapshotMatcher
-                .with().rows(6).frames(1),
-            0, 0, "shuffle"
-        )

Review Comment:
   Why are we removing the assertion?



##########
processing/src/main/java/org/apache/druid/frame/Frame.java:
##########
@@ -105,7 +105,7 @@ public class Frame
   private final int numRegions;
   private final boolean permuted;
 
-  private Frame(Memory memory, FrameType frameType, long numBytes, int numRows, int numRegions, boolean permuted)
+  protected Frame(Memory memory, FrameType frameType, long numBytes, int numRows, int numRegions, boolean permuted)

Review Comment:
   ```suggestion
     protected Frame(Memory memory, FrameType frameType, long numBytes, int numRows, int numRegions, boolean permuted)
   ```
   
   Please revert this change. We shouldn't change the access level in the main class to allow testing. At worst, we can mock the class in the tests, but if we do want to assert something.
   
   We are basically annotating the constructor here with @VisibleForTesting, without mentioning it
   https://github.com/apache/druid/pull/11848#discussion_r739022142



##########
extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestBase.java:
##########
@@ -1329,24 +1332,36 @@ public Pair<MSQSpec, Pair<List<MSQResultsReport.ColumnAndType>, List<Object[]>>>
           } else {
             StageDefinition finalStage = Objects.requireNonNull(SqlStatementResourceHelper.getFinalStage(
                 payload)).getStageDefinition();
-            Closer closer = Closer.create();
-            InputChannelFactory inputChannelFactory = DurableStorageInputChannelFactory.createStandardImplementation(
-                controllerId,
-                localFileStorageConnector,
-                closer,
-                true
+
+            Optional<List<PageInformation>> pages = SqlStatementResourceHelper.populatePageList(
+                payload,
+                spec.getDestination()
             );
-            rows = new FrameChannelSequence(inputChannelFactory.openChannel(
-                finalStage.getId(),
-                0,
-                0
-            )).flatMap(frame -> SqlStatementResourceHelper.getResultSequence(
-                msqControllerTask,
-                finalStage,
-                frame,
-                objectMapper
-            )).withBaggage(closer).toList();
 
+            if (!pages.isPresent()) {
+              throw new ISE("Query no results found");

Review Comment:
   Grammar



##########
extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/util/SqlStatementResourceHelperTest.java:
##########
@@ -0,0 +1,334 @@
+/*
+ * 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.druid.msq.util;
+
+import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.Maps;
+import org.apache.datasketches.memory.Memory;
+import org.apache.druid.frame.Frame;
+import org.apache.druid.frame.FrameType;
+import org.apache.druid.indexer.TaskState;
+import org.apache.druid.java.util.common.Pair;
+import org.apache.druid.java.util.common.logger.Logger;
+import org.apache.druid.msq.counters.ChannelCounters;
+import org.apache.druid.msq.counters.CounterSnapshots;
+import org.apache.druid.msq.counters.CounterSnapshotsTree;
+import org.apache.druid.msq.indexing.destination.DurableStorageMSQDestination;
+import org.apache.druid.msq.indexing.report.MSQStagesReport;
+import org.apache.druid.msq.indexing.report.MSQStatusReport;
+import org.apache.druid.msq.indexing.report.MSQTaskReportPayload;
+import org.apache.druid.msq.indexing.report.MSQTaskReportTest;
+import org.apache.druid.msq.sql.entity.PageInformation;
+import org.junit.Assert;
+import org.junit.Test;
+
+import java.util.ArrayDeque;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.TreeMap;
+
+public class SqlStatementResourceHelperTest
+{
+
+  private static final Logger log = new Logger(SqlStatementResourceHelperTest.class);
+
+  @Test
+  public void testDistinctPartitionsOnEachWorker()
+  {
+    CounterSnapshotsTree counterSnapshots = new CounterSnapshotsTree();
+    ChannelCounters worker0 = createChannelCounters(new int[]{0, 3, 6});
+    ChannelCounters worker1 = createChannelCounters(new int[]{1, 4, 4, 7, 9, 10, 13});
+    ChannelCounters worker2 = createChannelCounters(new int[]{2, 5, 8, 11, 14});
+
+    counterSnapshots.put(0, 0, new CounterSnapshots(ImmutableMap.of("output", worker0.snapshot())));
+    counterSnapshots.put(0, 1, new CounterSnapshots(ImmutableMap.of("output", worker1.snapshot())));
+    counterSnapshots.put(0, 2, new CounterSnapshots(ImmutableMap.of("output", worker2.snapshot())));
+
+    MSQTaskReportPayload payload = new MSQTaskReportPayload(new MSQStatusReport(
+        TaskState.SUCCESS,
+        null,
+        new ArrayDeque<>(),
+        null,
+        0,
+        1,
+        2
+    ), MSQStagesReport.create(
+        MSQTaskReportTest.QUERY_DEFINITION,
+        ImmutableMap.of(),
+        ImmutableMap.of(),
+        ImmutableMap.of(0, 3),
+        ImmutableMap.of(0, 15)
+    ), counterSnapshots, null);
+
+    Optional<List<PageInformation>> pages = SqlStatementResourceHelper.populatePageList(
+        payload,
+        DurableStorageMSQDestination.instance()
+    );
+    validatePages(pages.get(), createValidationMap(worker0, worker1, worker2));
+  }
+
+  @Test
+  public void testOnePartitionOnEachWorker()
+  {
+    CounterSnapshotsTree counterSnapshots = new CounterSnapshotsTree();
+    ChannelCounters worker0 = createChannelCounters(new int[]{0});
+    ChannelCounters worker1 = createChannelCounters(new int[]{1});
+    ChannelCounters worker2 = createChannelCounters(new int[]{2});
+    ChannelCounters worker3 = createChannelCounters(new int[]{4});
+
+    counterSnapshots.put(0, 0, new CounterSnapshots(ImmutableMap.of("output", worker0.snapshot())));
+    counterSnapshots.put(0, 1, new CounterSnapshots(ImmutableMap.of("output", worker1.snapshot())));
+    counterSnapshots.put(0, 2, new CounterSnapshots(ImmutableMap.of("output", worker2.snapshot())));
+    counterSnapshots.put(0, 3, new CounterSnapshots(ImmutableMap.of("output", worker3.snapshot())));
+
+    MSQTaskReportPayload payload = new MSQTaskReportPayload(new MSQStatusReport(
+        TaskState.SUCCESS,
+        null,
+        new ArrayDeque<>(),
+        null,
+        0,
+        1,
+        2
+    ), MSQStagesReport.create(
+        MSQTaskReportTest.QUERY_DEFINITION,
+        ImmutableMap.of(),
+        ImmutableMap.of(),
+        ImmutableMap.of(0, 4),
+        ImmutableMap.of(0, 4)
+    ), counterSnapshots, null);
+
+    Optional<List<PageInformation>> pages = SqlStatementResourceHelper.populatePageList(
+        payload,
+        DurableStorageMSQDestination.instance()
+    );
+    validatePages(pages.get(), createValidationMap(worker0, worker1, worker2));
+  }
+
+
+  @Test
+  public void testCommonPartitionsOnEachWorker()
+  {
+    CounterSnapshotsTree counterSnapshots = new CounterSnapshotsTree();
+    ChannelCounters worker0 = createChannelCounters(new int[]{0, 1, 2, 3, 8, 9});
+    ChannelCounters worker1 = createChannelCounters(new int[]{1, 4, 12});
+    ChannelCounters worker2 = createChannelCounters(new int[]{20});
+    ChannelCounters worker3 = createChannelCounters(new int[]{2, 2, 5, 6, 7, 9, 15});
+
+    counterSnapshots.put(0, 0, new CounterSnapshots(ImmutableMap.of("output", worker0.snapshot())));
+    counterSnapshots.put(0, 1, new CounterSnapshots(ImmutableMap.of("output", worker1.snapshot())));
+    counterSnapshots.put(0, 2, new CounterSnapshots(ImmutableMap.of("output", worker2.snapshot())));
+    counterSnapshots.put(0, 3, new CounterSnapshots(ImmutableMap.of("output", worker3.snapshot())));
+
+    MSQTaskReportPayload payload = new MSQTaskReportPayload(new MSQStatusReport(
+        TaskState.SUCCESS,
+        null,
+        new ArrayDeque<>(),
+        null,
+        0,
+        1,
+        2
+    ), MSQStagesReport.create(
+        MSQTaskReportTest.QUERY_DEFINITION,
+        ImmutableMap.of(),
+        ImmutableMap.of(),
+        ImmutableMap.of(0, 4),
+        ImmutableMap.of(0, 21)
+    ), counterSnapshots, null);
+
+    Optional<List<PageInformation>> pages =
+        SqlStatementResourceHelper.populatePageList(payload, DurableStorageMSQDestination.instance());
+    validatePages(pages.get(), createValidationMap(worker0, worker1, worker2, worker3));
+  }
+
+
+  @Test
+  public void testNullChannelCounters()
+  {
+    CounterSnapshotsTree counterSnapshots = new CounterSnapshotsTree();
+    ChannelCounters worker0 = createChannelCounters(new int[0]);
+    ChannelCounters worker1 = createChannelCounters(new int[]{1, 4, 12});
+    ChannelCounters worker2 = createChannelCounters(new int[]{20});
+    ChannelCounters worker3 = createChannelCounters(new int[]{2, 2, 5, 6, 7, 9, 15});
+
+    counterSnapshots.put(0, 0, new CounterSnapshots(Maps.newHashMap()));
+    counterSnapshots.put(0, 1, new CounterSnapshots(ImmutableMap.of("output", worker1.snapshot())));
+    counterSnapshots.put(0, 2, new CounterSnapshots(ImmutableMap.of("output", worker2.snapshot())));
+    counterSnapshots.put(0, 3, new CounterSnapshots(ImmutableMap.of("output", worker3.snapshot())));
+
+    MSQTaskReportPayload payload = new MSQTaskReportPayload(new MSQStatusReport(
+        TaskState.SUCCESS,
+        null,
+        new ArrayDeque<>(),
+        null,
+        0,
+        1,
+        2
+    ), MSQStagesReport.create(
+        MSQTaskReportTest.QUERY_DEFINITION,
+        ImmutableMap.of(),
+        ImmutableMap.of(),
+        ImmutableMap.of(0, 4),
+        ImmutableMap.of(0, 21)
+    ), counterSnapshots, null);
+
+    Optional<List<PageInformation>> pages = SqlStatementResourceHelper.populatePageList(
+        payload,
+        DurableStorageMSQDestination.instance()
+    );
+    validatePages(pages.get(), createValidationMap(worker0, worker1, worker2, worker3));
+  }
+
+
+  @Test
+  public void testConsecutivePartitionsOnEachWorker()
+  {
+    CounterSnapshotsTree counterSnapshots = new CounterSnapshotsTree();
+    ChannelCounters worker0 = createChannelCounters(new int[]{0, 1, 2});
+    ChannelCounters worker1 = createChannelCounters(new int[]{3, 4, 5});
+    ChannelCounters worker2 = createChannelCounters(new int[]{6, 7, 8});
+    ChannelCounters worker3 = createChannelCounters(new int[]{9, 10, 11, 12});
+
+    counterSnapshots.put(0, 0, new CounterSnapshots(ImmutableMap.of("output", worker0.snapshot())));
+    counterSnapshots.put(0, 1, new CounterSnapshots(ImmutableMap.of("output", worker1.snapshot())));
+    counterSnapshots.put(0, 2, new CounterSnapshots(ImmutableMap.of("output", worker2.snapshot())));
+    counterSnapshots.put(0, 3, new CounterSnapshots(ImmutableMap.of("output", worker3.snapshot())));
+
+    MSQTaskReportPayload payload = new MSQTaskReportPayload(new MSQStatusReport(
+        TaskState.SUCCESS,
+        null,
+        new ArrayDeque<>(),
+        null,
+        0,
+        1,
+        2
+    ), MSQStagesReport.create(
+        MSQTaskReportTest.QUERY_DEFINITION,
+        ImmutableMap.of(),
+        ImmutableMap.of(),
+        ImmutableMap.of(0, 4),
+        ImmutableMap.of(0, 13)
+    ), counterSnapshots, null);
+
+    Optional<List<PageInformation>> pages = SqlStatementResourceHelper.populatePageList(
+        payload,
+        DurableStorageMSQDestination.instance()
+    );
+    validatePages(pages.get(), createValidationMap(worker0, worker1, worker2, worker3));
+  }
+
+
+  private void validatePages(
+      List<PageInformation> pageList,
+      Map<Integer, Map<Integer, Pair<Long, Long>>> partitionToWorkerToRowsBytes
+  )
+  {
+    int currentPage = 0;
+    for (Map.Entry<Integer, Map<Integer, Pair<Long, Long>>> partitionWorker : partitionToWorkerToRowsBytes.entrySet()) {
+      for (Map.Entry<Integer, Pair<Long, Long>> workerRowsBytes : partitionWorker.getValue().entrySet()) {
+        PageInformation pageInformation = pageList.get(currentPage);
+        Assert.assertEquals(currentPage, pageInformation.getId());
+        Assert.assertEquals(workerRowsBytes.getValue().lhs, pageInformation.getNumRows());
+        Assert.assertEquals(workerRowsBytes.getValue().rhs, pageInformation.getSizeInBytes());
+        Assert.assertEquals(partitionWorker.getKey(), pageInformation.getPartition());
+        Assert.assertEquals(workerRowsBytes.getKey(), pageInformation.getWorker());
+        log.debug(pageInformation.toString());
+        currentPage++;
+      }
+    }
+    Assert.assertEquals(currentPage, pageList.size());
+  }
+
+  private Map<Integer, Map<Integer, Pair<Long, Long>>> createValidationMap(
+      ChannelCounters... workers
+  )
+  {
+    if (workers == null || workers.length == 0) {
+      return Maps.newHashMap();
+    } else {
+      Map<Integer, Map<Integer, Pair<Long, Long>>> partitionToWorkerToRowsBytes = new TreeMap<>();
+      for (int worker = 0; worker < workers.length; worker++) {
+        ChannelCounters.Snapshot workerCounter = workers[worker].snapshot();
+        for (int partition = 0; workerCounter != null && partition < workerCounter.getRows().length; partition++) {
+          Map<Integer, Pair<Long, Long>> workerMap = partitionToWorkerToRowsBytes.computeIfAbsent(
+              partition,
+              k -> new TreeMap<>()
+          );
+
+          if (workerCounter.getRows()[partition] != 0) {
+            workerMap.put(
+                worker,
+                new Pair<>(
+                    workerCounter.getRows()[partition],
+                    workerCounter.getBytes()[partition]
+                )
+            );
+          }
+
+        }
+      }
+      return partitionToWorkerToRowsBytes;
+    }
+  }
+
+
+  private ChannelCounters createChannelCounters(int[] partitions)
+  {
+    if (partitions == null || partitions.length == 0) {
+      return new ChannelCounters();
+    }
+    ChannelCounters channelCounters = new ChannelCounters();
+    int prev = -1;
+    for (int current : partitions) {
+      if (prev > current) {
+        throw new IllegalArgumentException("Channel numbers should be in increasing order");
+      }
+      channelCounters.addFrame(current, new TestFrame(current * 10 + 1, 100));
+      prev = current;
+    }
+    return channelCounters;
+  }
+
+  public static class TestFrame extends Frame

Review Comment:
   ```suggestion
     private static class TestFrame extends Frame
   ```
   Let's not create a separate class for a single use 



-- 
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-unsubscribe@druid.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org


Re: [PR] Limit pages size to a configurable limit (druid)

Posted by "LakshSingla (via GitHub)" <gi...@apache.org>.
LakshSingla commented on code in PR #14994:
URL: https://github.com/apache/druid/pull/14994#discussion_r1332641725


##########
processing/src/main/java/org/apache/druid/frame/Frame.java:
##########
@@ -105,7 +105,7 @@ public class Frame
   private final int numRegions;
   private final boolean permuted;
 
-  private Frame(Memory memory, FrameType frameType, long numBytes, int numRows, int numRegions, boolean permuted)
+  protected Frame(Memory memory, FrameType frameType, long numBytes, int numRows, int numRegions, boolean permuted)

Review Comment:
   ```suggestion
     private Frame(Memory memory, FrameType frameType, long numBytes, int numRows, int numRegions, boolean permuted)
   ```
   
   Please revert this change. We shouldn't change the access level in the main class to allow testing. At worst, we can mock the class in the tests, but if we do want to assert something that won't ever be possible and we are setting up a defensive check for the same.
   
   We are basically annotating the constructor here with @VisibleForTesting, without mentioning it
   https://github.com/apache/druid/pull/11848#discussion_r739022142



-- 
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-unsubscribe@druid.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org


Re: [PR] Limit pages size to a configurable limit (druid)

Posted by "cryptoe (via GitHub)" <gi...@apache.org>.
cryptoe commented on code in PR #14994:
URL: https://github.com/apache/druid/pull/14994#discussion_r1354906061


##########
extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQSelectTest.java:
##########
@@ -215,16 +216,6 @@ public void testSelectOnFoo()
                 .with().totalFiles(1),
             0, 0, "input0"
         )
-        .setExpectedCountersForStageWorkerChannel(
-            CounterSnapshotMatcher
-                .with().rows(6).frames(1),
-            0, 0, "output"
-        )
-        .setExpectedCountersForStageWorkerChannel(
-            CounterSnapshotMatcher
-                .with().rows(6).frames(1),
-            0, 0, "shuffle"
-        )

Review Comment:
   Its was more of an influx thing. Fixed it 



-- 
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-unsubscribe@druid.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org


Re: [PR] Limit pages size to a configurable limit (druid)

Posted by "cryptoe (via GitHub)" <gi...@apache.org>.
cryptoe commented on PR #14994:
URL: https://github.com/apache/druid/pull/14994#issuecomment-1759169471

   Thanks @adarshsanjeev @LakshSingla for the review. 


-- 
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-unsubscribe@druid.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org


Re: [PR] Limit pages size to a configurable limit (druid)

Posted by "cryptoe (via GitHub)" <gi...@apache.org>.
cryptoe commented on PR #14994:
URL: https://github.com/apache/druid/pull/14994#issuecomment-1758999786

   > Why do we need indexes on frames for this?
   That will help us in supporting usecases in the result api with s`startRowOffset` and `range`. Then we can hide the page concept totally from the end user.  
   


-- 
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-unsubscribe@druid.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org


Re: [PR] Limit pages size to a configurable limit (druid)

Posted by "cryptoe (via GitHub)" <gi...@apache.org>.
cryptoe commented on PR #14994:
URL: https://github.com/apache/druid/pull/14994#issuecomment-1758974725

   @LakshSingla 
   In theory I agree with you that counters are a little risky in the long run and we should have something else in the task report but there are no good alternatives yet. Once we go down on generating frame indexes, then we can think of changing the approach of this. 


-- 
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-unsubscribe@druid.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org


Re: [PR] Limit pages size to a configurable limit (druid)

Posted by "cryptoe (via GitHub)" <gi...@apache.org>.
cryptoe commented on code in PR #14994:
URL: https://github.com/apache/druid/pull/14994#discussion_r1354908031


##########
extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/util/SqlStatementResourceHelperTest.java:
##########
@@ -0,0 +1,334 @@
+/*
+ * 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.druid.msq.util;
+
+import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.Maps;
+import org.apache.datasketches.memory.Memory;
+import org.apache.druid.frame.Frame;
+import org.apache.druid.frame.FrameType;
+import org.apache.druid.indexer.TaskState;
+import org.apache.druid.java.util.common.Pair;
+import org.apache.druid.java.util.common.logger.Logger;
+import org.apache.druid.msq.counters.ChannelCounters;
+import org.apache.druid.msq.counters.CounterSnapshots;
+import org.apache.druid.msq.counters.CounterSnapshotsTree;
+import org.apache.druid.msq.indexing.destination.DurableStorageMSQDestination;
+import org.apache.druid.msq.indexing.report.MSQStagesReport;
+import org.apache.druid.msq.indexing.report.MSQStatusReport;
+import org.apache.druid.msq.indexing.report.MSQTaskReportPayload;
+import org.apache.druid.msq.indexing.report.MSQTaskReportTest;
+import org.apache.druid.msq.sql.entity.PageInformation;
+import org.junit.Assert;
+import org.junit.Test;
+
+import java.util.ArrayDeque;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.TreeMap;
+
+public class SqlStatementResourceHelperTest
+{
+
+  private static final Logger log = new Logger(SqlStatementResourceHelperTest.class);
+
+  @Test
+  public void testDistinctPartitionsOnEachWorker()
+  {
+    CounterSnapshotsTree counterSnapshots = new CounterSnapshotsTree();
+    ChannelCounters worker0 = createChannelCounters(new int[]{0, 3, 6});
+    ChannelCounters worker1 = createChannelCounters(new int[]{1, 4, 4, 7, 9, 10, 13});
+    ChannelCounters worker2 = createChannelCounters(new int[]{2, 5, 8, 11, 14});
+
+    counterSnapshots.put(0, 0, new CounterSnapshots(ImmutableMap.of("output", worker0.snapshot())));
+    counterSnapshots.put(0, 1, new CounterSnapshots(ImmutableMap.of("output", worker1.snapshot())));
+    counterSnapshots.put(0, 2, new CounterSnapshots(ImmutableMap.of("output", worker2.snapshot())));
+
+    MSQTaskReportPayload payload = new MSQTaskReportPayload(new MSQStatusReport(
+        TaskState.SUCCESS,
+        null,
+        new ArrayDeque<>(),
+        null,
+        0,
+        1,
+        2
+    ), MSQStagesReport.create(
+        MSQTaskReportTest.QUERY_DEFINITION,
+        ImmutableMap.of(),
+        ImmutableMap.of(),
+        ImmutableMap.of(0, 3),
+        ImmutableMap.of(0, 15)
+    ), counterSnapshots, null);
+
+    Optional<List<PageInformation>> pages = SqlStatementResourceHelper.populatePageList(
+        payload,
+        DurableStorageMSQDestination.instance()
+    );
+    validatePages(pages.get(), createValidationMap(worker0, worker1, worker2));
+  }
+
+  @Test
+  public void testOnePartitionOnEachWorker()
+  {
+    CounterSnapshotsTree counterSnapshots = new CounterSnapshotsTree();
+    ChannelCounters worker0 = createChannelCounters(new int[]{0});
+    ChannelCounters worker1 = createChannelCounters(new int[]{1});
+    ChannelCounters worker2 = createChannelCounters(new int[]{2});
+    ChannelCounters worker3 = createChannelCounters(new int[]{4});
+
+    counterSnapshots.put(0, 0, new CounterSnapshots(ImmutableMap.of("output", worker0.snapshot())));
+    counterSnapshots.put(0, 1, new CounterSnapshots(ImmutableMap.of("output", worker1.snapshot())));
+    counterSnapshots.put(0, 2, new CounterSnapshots(ImmutableMap.of("output", worker2.snapshot())));
+    counterSnapshots.put(0, 3, new CounterSnapshots(ImmutableMap.of("output", worker3.snapshot())));
+
+    MSQTaskReportPayload payload = new MSQTaskReportPayload(new MSQStatusReport(
+        TaskState.SUCCESS,
+        null,
+        new ArrayDeque<>(),
+        null,
+        0,
+        1,
+        2
+    ), MSQStagesReport.create(
+        MSQTaskReportTest.QUERY_DEFINITION,
+        ImmutableMap.of(),
+        ImmutableMap.of(),
+        ImmutableMap.of(0, 4),
+        ImmutableMap.of(0, 4)
+    ), counterSnapshots, null);
+
+    Optional<List<PageInformation>> pages = SqlStatementResourceHelper.populatePageList(
+        payload,
+        DurableStorageMSQDestination.instance()
+    );
+    validatePages(pages.get(), createValidationMap(worker0, worker1, worker2));
+  }
+
+
+  @Test
+  public void testCommonPartitionsOnEachWorker()
+  {
+    CounterSnapshotsTree counterSnapshots = new CounterSnapshotsTree();
+    ChannelCounters worker0 = createChannelCounters(new int[]{0, 1, 2, 3, 8, 9});
+    ChannelCounters worker1 = createChannelCounters(new int[]{1, 4, 12});
+    ChannelCounters worker2 = createChannelCounters(new int[]{20});
+    ChannelCounters worker3 = createChannelCounters(new int[]{2, 2, 5, 6, 7, 9, 15});
+
+    counterSnapshots.put(0, 0, new CounterSnapshots(ImmutableMap.of("output", worker0.snapshot())));
+    counterSnapshots.put(0, 1, new CounterSnapshots(ImmutableMap.of("output", worker1.snapshot())));
+    counterSnapshots.put(0, 2, new CounterSnapshots(ImmutableMap.of("output", worker2.snapshot())));
+    counterSnapshots.put(0, 3, new CounterSnapshots(ImmutableMap.of("output", worker3.snapshot())));
+
+    MSQTaskReportPayload payload = new MSQTaskReportPayload(new MSQStatusReport(
+        TaskState.SUCCESS,
+        null,
+        new ArrayDeque<>(),
+        null,
+        0,
+        1,
+        2
+    ), MSQStagesReport.create(
+        MSQTaskReportTest.QUERY_DEFINITION,
+        ImmutableMap.of(),
+        ImmutableMap.of(),
+        ImmutableMap.of(0, 4),
+        ImmutableMap.of(0, 21)
+    ), counterSnapshots, null);
+
+    Optional<List<PageInformation>> pages =
+        SqlStatementResourceHelper.populatePageList(payload, DurableStorageMSQDestination.instance());
+    validatePages(pages.get(), createValidationMap(worker0, worker1, worker2, worker3));
+  }
+
+
+  @Test
+  public void testNullChannelCounters()
+  {
+    CounterSnapshotsTree counterSnapshots = new CounterSnapshotsTree();
+    ChannelCounters worker0 = createChannelCounters(new int[0]);
+    ChannelCounters worker1 = createChannelCounters(new int[]{1, 4, 12});
+    ChannelCounters worker2 = createChannelCounters(new int[]{20});
+    ChannelCounters worker3 = createChannelCounters(new int[]{2, 2, 5, 6, 7, 9, 15});
+
+    counterSnapshots.put(0, 0, new CounterSnapshots(Maps.newHashMap()));
+    counterSnapshots.put(0, 1, new CounterSnapshots(ImmutableMap.of("output", worker1.snapshot())));
+    counterSnapshots.put(0, 2, new CounterSnapshots(ImmutableMap.of("output", worker2.snapshot())));
+    counterSnapshots.put(0, 3, new CounterSnapshots(ImmutableMap.of("output", worker3.snapshot())));
+
+    MSQTaskReportPayload payload = new MSQTaskReportPayload(new MSQStatusReport(
+        TaskState.SUCCESS,
+        null,
+        new ArrayDeque<>(),
+        null,
+        0,
+        1,
+        2
+    ), MSQStagesReport.create(
+        MSQTaskReportTest.QUERY_DEFINITION,
+        ImmutableMap.of(),
+        ImmutableMap.of(),
+        ImmutableMap.of(0, 4),
+        ImmutableMap.of(0, 21)
+    ), counterSnapshots, null);
+
+    Optional<List<PageInformation>> pages = SqlStatementResourceHelper.populatePageList(
+        payload,
+        DurableStorageMSQDestination.instance()
+    );
+    validatePages(pages.get(), createValidationMap(worker0, worker1, worker2, worker3));
+  }
+
+
+  @Test
+  public void testConsecutivePartitionsOnEachWorker()
+  {
+    CounterSnapshotsTree counterSnapshots = new CounterSnapshotsTree();
+    ChannelCounters worker0 = createChannelCounters(new int[]{0, 1, 2});
+    ChannelCounters worker1 = createChannelCounters(new int[]{3, 4, 5});
+    ChannelCounters worker2 = createChannelCounters(new int[]{6, 7, 8});
+    ChannelCounters worker3 = createChannelCounters(new int[]{9, 10, 11, 12});
+
+    counterSnapshots.put(0, 0, new CounterSnapshots(ImmutableMap.of("output", worker0.snapshot())));
+    counterSnapshots.put(0, 1, new CounterSnapshots(ImmutableMap.of("output", worker1.snapshot())));
+    counterSnapshots.put(0, 2, new CounterSnapshots(ImmutableMap.of("output", worker2.snapshot())));
+    counterSnapshots.put(0, 3, new CounterSnapshots(ImmutableMap.of("output", worker3.snapshot())));
+
+    MSQTaskReportPayload payload = new MSQTaskReportPayload(new MSQStatusReport(
+        TaskState.SUCCESS,
+        null,
+        new ArrayDeque<>(),
+        null,
+        0,
+        1,
+        2
+    ), MSQStagesReport.create(
+        MSQTaskReportTest.QUERY_DEFINITION,
+        ImmutableMap.of(),
+        ImmutableMap.of(),
+        ImmutableMap.of(0, 4),
+        ImmutableMap.of(0, 13)
+    ), counterSnapshots, null);
+
+    Optional<List<PageInformation>> pages = SqlStatementResourceHelper.populatePageList(
+        payload,
+        DurableStorageMSQDestination.instance()
+    );
+    validatePages(pages.get(), createValidationMap(worker0, worker1, worker2, worker3));
+  }
+
+
+  private void validatePages(
+      List<PageInformation> pageList,
+      Map<Integer, Map<Integer, Pair<Long, Long>>> partitionToWorkerToRowsBytes
+  )
+  {
+    int currentPage = 0;
+    for (Map.Entry<Integer, Map<Integer, Pair<Long, Long>>> partitionWorker : partitionToWorkerToRowsBytes.entrySet()) {
+      for (Map.Entry<Integer, Pair<Long, Long>> workerRowsBytes : partitionWorker.getValue().entrySet()) {
+        PageInformation pageInformation = pageList.get(currentPage);
+        Assert.assertEquals(currentPage, pageInformation.getId());
+        Assert.assertEquals(workerRowsBytes.getValue().lhs, pageInformation.getNumRows());
+        Assert.assertEquals(workerRowsBytes.getValue().rhs, pageInformation.getSizeInBytes());
+        Assert.assertEquals(partitionWorker.getKey(), pageInformation.getPartition());
+        Assert.assertEquals(workerRowsBytes.getKey(), pageInformation.getWorker());
+        log.debug(pageInformation.toString());
+        currentPage++;
+      }
+    }
+    Assert.assertEquals(currentPage, pageList.size());
+  }
+
+  private Map<Integer, Map<Integer, Pair<Long, Long>>> createValidationMap(
+      ChannelCounters... workers
+  )
+  {
+    if (workers == null || workers.length == 0) {
+      return Maps.newHashMap();
+    } else {
+      Map<Integer, Map<Integer, Pair<Long, Long>>> partitionToWorkerToRowsBytes = new TreeMap<>();
+      for (int worker = 0; worker < workers.length; worker++) {
+        ChannelCounters.Snapshot workerCounter = workers[worker].snapshot();
+        for (int partition = 0; workerCounter != null && partition < workerCounter.getRows().length; partition++) {
+          Map<Integer, Pair<Long, Long>> workerMap = partitionToWorkerToRowsBytes.computeIfAbsent(
+              partition,
+              k -> new TreeMap<>()
+          );
+
+          if (workerCounter.getRows()[partition] != 0) {
+            workerMap.put(
+                worker,
+                new Pair<>(
+                    workerCounter.getRows()[partition],
+                    workerCounter.getBytes()[partition]
+                )
+            );
+          }
+
+        }
+      }
+      return partitionToWorkerToRowsBytes;
+    }
+  }
+
+
+  private ChannelCounters createChannelCounters(int[] partitions)
+  {
+    if (partitions == null || partitions.length == 0) {
+      return new ChannelCounters();
+    }
+    ChannelCounters channelCounters = new ChannelCounters();
+    int prev = -1;
+    for (int current : partitions) {
+      if (prev > current) {
+        throw new IllegalArgumentException("Channel numbers should be in increasing order");
+      }
+      channelCounters.addFrame(current, new TestFrame(current * 10 + 1, 100));
+      prev = current;
+    }
+    return channelCounters;
+  }
+
+  public static class TestFrame extends Frame

Review Comment:
   I just removed this class and went with mocks. 



-- 
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-unsubscribe@druid.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org


Re: [PR] Limit pages size to a configurable limit (druid)

Posted by "cryptoe (via GitHub)" <gi...@apache.org>.
cryptoe commented on code in PR #14994:
URL: https://github.com/apache/druid/pull/14994#discussion_r1356115325


##########
extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/ControllerImpl.java:
##########
@@ -1693,8 +1687,9 @@ private static QueryDefinition makeQueryDefinition(
       shuffleSpecFactory = ShuffleSpecFactories.singlePartition();
       queryToPlan = querySpec.getQuery();
     } else if (querySpec.getDestination() instanceof DurableStorageMSQDestination) {
-      // we add a final stage which generates one partition per worker.
-      shuffleSpecFactory = ShuffleSpecFactories.globalSortWithMaxPartitionCount(tuningConfig.getMaxNumWorkers());
+      shuffleSpecFactory = ShuffleSpecFactories.getGlobalSortWithTargetSize(
+          MultiStageQueryContext.getRowsPerPage(querySpec.getQuery().context())

Review Comment:
   I have added a new testcase `testExternSelectWithMultipleWorkers`. You can look at the counter checks to see whats happening with a scan query. 



-- 
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-unsubscribe@druid.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org


Re: [PR] Limit pages size to a configurable limit (druid)

Posted by "LakshSingla (via GitHub)" <gi...@apache.org>.
LakshSingla commented on PR #14994:
URL: https://github.com/apache/druid/pull/14994#issuecomment-1758979412

   > Once we go down on generating frame indexes, then we can think of changing the approach of this.
   Why do we need indexes on frames for this?


-- 
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-unsubscribe@druid.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org