populate queryOption down to leaf

Rong Rong · Rong Rong · commit 7b545ea6318b · 2023-04-17T10:32:33.000-07:00
diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/ServerRequestPlanVisitor.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/ServerRequestPlanVisitor.java
@@ -19,9 +19,9 @@
 package org.apache.pinot.query.runtime.plan;
 
 import com.google.common.collect.ImmutableList;
-import com.google.common.collect.ImmutableMap;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.stream.Collectors;
@@ -126,8 +126,7 @@ public static ServerPlanRequestContext build(MailboxService<TransferableBlock> m
     QUERY_OPTIMIZER.optimize(pinotQuery, tableConfig, schema);
 
     // 2. set pinot query options according to requestMetadataMap
-    pinotQuery.setQueryOptions(
-        ImmutableMap.of(CommonConstants.Broker.Request.QueryOptionKey.TIMEOUT_MS, String.valueOf(timeoutMs)));
+    updateQueryOptions(pinotQuery, requestMetadataMap, timeoutMs, traceEnabled);
 
     // 3. wrapped around in broker request
     BrokerRequest brokerRequest = new BrokerRequest();
@@ -151,6 +150,19 @@ public static ServerPlanRequestContext build(MailboxService<TransferableBlock> m
     return context;
   }
 
+  private static void updateQueryOptions(PinotQuery pinotQuery, Map<String, String> requestMetadataMap, long timeoutMs,
+      boolean traceEnabled) {
+    Map<String, String> queryOptions = new HashMap<>();
+    // put default timeout and trace options
+    queryOptions.put(CommonConstants.Broker.Request.QueryOptionKey.TIMEOUT_MS, String.valueOf(timeoutMs));
+    if (traceEnabled) {
+      queryOptions.put(CommonConstants.Broker.Request.TRACE, "true");
+    }
+    // overwrite with requestMetadataMap to carry query options from request:
+    queryOptions.putAll(requestMetadataMap);
+    pinotQuery.setQueryOptions(queryOptions);
+  }
+
   private static void walkStageNode(StageNode node, ServerPlanRequestContext context) {
     node.visit(INSTANCE, context);
   }
diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/QueryRunnerTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/QueryRunnerTest.java
@@ -54,10 +54,8 @@
 
 
 /**
- * all legacy tests.
- *
- * @deprecated do not add to this test set. this class will be broken down and clean up.
- * add your test to appropraite files in {@link org.apache.pinot.query.runtime.queries} instead.
+ * all special tests that doesn't fit into {@link org.apache.pinot.query.runtime.queries.ResourceBasedQueriesTest}
+ * pattern goes here.
  */
 public class QueryRunnerTest extends QueryRunnerTestBase {
   public static final Object[][] ROWS = new Object[][]{
@@ -156,13 +154,23 @@ public void tearDown() {
     _mailboxService.shutdown();
   }
 
-  @Test(dataProvider = "testDataWithSqlToFinalRowCount")
+
+  /**
+   * Test compares with expected row count only.
+   */
+   @Test(dataProvider = "testDataWithSqlToFinalRowCount")
   public void testSqlWithFinalRowCountChecker(String sql, int expectedRows)
       throws Exception {
     List<Object[]> resultRows = queryRunner(sql, null);
     Assert.assertEquals(resultRows.size(), expectedRows);
   }
 
+  /**
+   * Test automatically compares against H2.
+   *
+   * @deprecated do not add to this test set. this class will be broken down and clean up.
+   *   add your test to the appropriate files in {@link org.apache.pinot.query.runtime.queries} instead.
+   */
   @Test(dataProvider = "testSql")
   public void testSqlWithH2Checker(String sql)
       throws Exception {
@@ -172,6 +180,9 @@ public void testSqlWithH2Checker(String sql)
     compareRowEquals(resultRows, expectedRows);
   }
 
+  /**
+   * Test compares against its desired exceptions.
+   */
   @Test(dataProvider = "testDataWithSqlExecutionExceptions")
   public void testSqlWithExceptionMsgChecker(String sql, String exceptionMsg) {
     long requestId = RANDOM_REQUEST_ID_GEN.nextLong();
@@ -235,6 +246,10 @@ private Object[][] provideTestSqlAndRowCount() {
         new Object[]{"SELECT round_decimal(col3) FROM a", 15},
         new Object[]{"SELECT col1, roundDecimal(COUNT(*)) FROM a GROUP BY col1", 5},
         new Object[]{"SELECT col1, round_decimal(COUNT(*)) FROM a GROUP BY col1", 5},
+
+        // test queries with special query options attached
+        //   - when leaf limit is set, each server returns multiStageLeafLimit number of rows only.
+        new Object[]{"SET multiStageLeafLimit = 1; SELECT * FROM a", 2},
     };
   }
 
diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/QueryRunnerTestBase.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/QueryRunnerTestBase.java
@@ -61,6 +61,8 @@
 import org.apache.pinot.spi.utils.BytesUtils;
 import org.apache.pinot.spi.utils.CommonConstants;
 import org.apache.pinot.spi.utils.StringUtil;
+import org.apache.pinot.sql.parsers.CalciteSqlParser;
+import org.apache.pinot.sql.parsers.SqlNodeAndOptions;
 import org.h2.jdbc.JdbcArray;
 import org.testng.Assert;
 
@@ -84,14 +86,22 @@ public abstract class QueryRunnerTestBase extends QueryTestSet {
   // --------------------------------------------------------------------------
   // QUERY UTILS
   // --------------------------------------------------------------------------
+
+  /**
+   * Dispatch query to each pinot-server. The logic should mimic
+   * {@link QueryDispatcher#submit(long, QueryPlan, long, Map)} but does not actually make ser/de dispatches.
+   */
   protected List<Object[]> queryRunner(String sql, Map<Integer, ExecutionStatsAggregator> executionStatsAggregatorMap) {
-    QueryPlan queryPlan = _queryEnvironment.planQuery(sql);
     long requestId = RANDOM_REQUEST_ID_GEN.nextLong();
+    SqlNodeAndOptions sqlNodeAndOptions = CalciteSqlParser.compileToSqlNodeAndOptions(sql);
+    QueryPlan queryPlan = _queryEnvironment.planQuery(sql, sqlNodeAndOptions, requestId).getQueryPlan();
     Map<String, String> requestMetadataMap = new HashMap<>();
     requestMetadataMap.put(QueryConfig.KEY_OF_BROKER_REQUEST_ID, String.valueOf(requestId));
     requestMetadataMap.put(QueryConfig.KEY_OF_BROKER_REQUEST_TIMEOUT_MS,
         String.valueOf(CommonConstants.Broker.DEFAULT_BROKER_TIMEOUT_MS));
+    requestMetadataMap.putAll(sqlNodeAndOptions.getOptions());
 
+    // Putting trace testing here as extra options as it doesn't go along with the rest of the items.
     if (executionStatsAggregatorMap != null) {
       requestMetadataMap.put(CommonConstants.Broker.Request.TRACE, "true");
     }
