[multistage] Multistage Engine Lite Mode (prototype)

pinot-query-planner/src/main/java/org/apache/pinot/query/context/PhysicalPlannerContext.java

@@ -57,6 +57,7 @@ public Integer get() {
    */
   private final String _instanceId;
   private final Map<String, String> _queryOptions;
+  private final boolean _useLiteMode;
 
   /**
    * Used by controller when it needs to extract table names from the query.
@@ -69,6 +70,7 @@ public PhysicalPlannerContext() {
     _requestId = 0;
     _instanceId = "";
     _queryOptions = Map.of();
+    _useLiteMode = false;
   }
 
   public PhysicalPlannerContext(RoutingManager routingManager, String hostName, int port, long requestId,
@@ -79,6 +81,7 @@ public PhysicalPlannerContext(RoutingManager routingManager, String hostName, in
     _requestId = requestId;
     _instanceId = instanceId;
     _queryOptions = queryOptions == null ? Map.of() : queryOptions;
+    _useLiteMode = PhysicalPlannerContext.useLiteMode(queryOptions);
   }
 
   public Supplier<Integer> getNodeIdGenerator() {
@@ -114,10 +117,21 @@ public Map<String, String> getQueryOptions() {
     return _queryOptions;
   }
 
+  public boolean isUseLiteMode() {
+    return _useLiteMode;
+  }
+
   public static boolean isUsePhysicalOptimizer(@Nullable Map<String, String> queryOptions) {
     if (queryOptions == null) {
       return false;
     }
     return Boolean.parseBoolean(queryOptions.getOrDefault(QueryOptionKey.USE_PHYSICAL_OPTIMIZER, "false"));
   }
+
+  private static boolean useLiteMode(@Nullable Map<String, String> queryOptions) {
+    if (queryOptions == null) {
+      return false;
+    }
+    return Boolean.parseBoolean(queryOptions.getOrDefault(QueryOptionKey.USE_LITE_MODE, "false"));
+  }
 }

pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/nodes/PhysicalAggregate.java

@@ -129,4 +129,10 @@ public List<RelFieldCollation> getCollations() {
   public int getLimit() {
     return _limit;
   }
+
+  public PhysicalAggregate withLimit(int newLimit) {
+    return new PhysicalAggregate(getCluster(), getTraitSet(), getHints(), groupSet, groupSets, aggCalls, _nodeId,
+        _pRelInputs.get(0), _pinotDataDistribution, _leafStage, _aggType, _leafReturnFinalResult, _collations,
+        newLimit);
+  }
 }

pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/nodes/PhysicalSort.java

@@ -96,4 +96,9 @@ public PRelNode asLeafStage() {
     return new PhysicalSort(getCluster(), getTraitSet(), getHints(), getCollation(), offset, fetch, _pRelInputs.get(0),
         _nodeId, _pinotDataDistribution, true);
   }
+
+  public PhysicalSort withFetch(RexNode newFetch) {
+    return new PhysicalSort(getCluster(), getTraitSet(), getHints(), getCollation(), offset, newFetch,
+        _pRelInputs.get(0), _nodeId, _pinotDataDistribution, _leafStage);
+  }
 }

pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/opt/PRelOptRule.java

@@ -49,4 +49,9 @@ public PRelNode onDone(PRelNode currentNode) {
   public PRelNode getParentNode(PRelOptRuleCall call) {
     return call._parents.isEmpty() ? null : call._parents.getLast();
   }
+
+  public boolean isLeafBoundary(PRelOptRuleCall call) {
+    PRelNode parentNode = getParentNode(call);
+    return call._currentNode.isLeafStage() && (parentNode == null || !parentNode.isLeafStage());
+  }
 }

pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/opt/PhysicalOptRuleSet.java

@@ -26,6 +26,8 @@
 import org.apache.pinot.query.planner.physical.v2.opt.rules.LeafStageAggregateRule;
 import org.apache.pinot.query.planner.physical.v2.opt.rules.LeafStageBoundaryRule;
 import org.apache.pinot.query.planner.physical.v2.opt.rules.LeafStageWorkerAssignmentRule;
+import org.apache.pinot.query.planner.physical.v2.opt.rules.LiteModeSortInsertRule;
+import org.apache.pinot.query.planner.physical.v2.opt.rules.LiteModeWorkerAssignmentRule;
 import org.apache.pinot.query.planner.physical.v2.opt.rules.SortPushdownRule;
 import org.apache.pinot.query.planner.physical.v2.opt.rules.WorkerExchangeAssignmentRule;
 
@@ -40,13 +42,21 @@ public static List<PRelNodeTransformer> create(PhysicalPlannerContext context, T
     transformers.add(create(new LeafStageWorkerAssignmentRule(context, tableCache), RuleExecutors.Type.POST_ORDER,
         context));
     transformers.add(create(new LeafStageAggregateRule(context), RuleExecutors.Type.POST_ORDER, context));
-    transformers.add(new WorkerExchangeAssignmentRule(context));
+    transformers.add(createWorkerAssignmentRule(context));
     transformers.add(create(new AggregatePushdownRule(context), RuleExecutors.Type.POST_ORDER, context));
     transformers.add(create(new SortPushdownRule(context), RuleExecutors.Type.POST_ORDER, context));
+    if (context.isUseLiteMode()) {
+      transformers.add(create(new LiteModeSortInsertRule(context), RuleExecutors.Type.POST_ORDER, context));
+    }
     return transformers;
   }
 
   private static PRelNodeTransformer create(PRelOptRule rule, RuleExecutors.Type type, PhysicalPlannerContext context) {
     return (pRelNode) -> RuleExecutors.create(type, rule, context).execute(pRelNode);
   }
+
+  private static PRelNodeTransformer createWorkerAssignmentRule(PhysicalPlannerContext context) {
+    return context.isUseLiteMode() ? new LiteModeWorkerAssignmentRule(context)
+        : new WorkerExchangeAssignmentRule(context);
+  }
 }

pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/opt/rules/LiteModeSortInsertRule.java

@@ -0,0 +1,100 @@
+/**
+ * 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.query.planner.physical.v2.opt.rules;
+
+import com.google.common.base.Preconditions;
+import java.util.List;
+import org.apache.calcite.plan.RelTraitSet;
+import org.apache.calcite.rel.RelCollations;
+import org.apache.calcite.rex.RexBuilder;
+import org.apache.calcite.rex.RexNode;
+import org.apache.calcite.sql.type.SqlTypeName;
+import org.apache.pinot.query.context.PhysicalPlannerContext;
+import org.apache.pinot.query.planner.logical.RexExpressionUtils;
+import org.apache.pinot.query.planner.physical.v2.PRelNode;
+import org.apache.pinot.query.planner.physical.v2.nodes.PhysicalAggregate;
+import org.apache.pinot.query.planner.physical.v2.nodes.PhysicalSort;
+import org.apache.pinot.query.planner.physical.v2.opt.PRelOptRule;
+import org.apache.pinot.query.planner.physical.v2.opt.PRelOptRuleCall;
+import org.apache.pinot.query.type.TypeFactory;
+
+
+/**
+ * Lite mode sets a hard limit on the number of rows that the leaf stage is allowed to return. This rule ensures the
+ * same. This is done by adding a Sort in the leaf stage if one doesn't exist already.
+ * <p>
+ *  When the leaf stage has an aggregation and no Sort, then we can add the limit to the aggregate itself to enable
+ *  server-level group trimming, which achieves the same purpose as adding a Sort. If the aggregation limit is higher
+ *  than the hard-limit, then this Rule will throw an error.
+ * </p>
+ * <p>
+ *   When the leaf stage has a Sort already, we verify that the limit doesn't exceed the configured hard limit.
+ * </p>
+ */
+public class LiteModeSortInsertRule extends PRelOptRule {
+  private static final TypeFactory TYPE_FACTORY = new TypeFactory();
+  private static final RexBuilder REX_BUILDER = new RexBuilder(TYPE_FACTORY);
+  // TODO: This should be configurable at broker and via SET statements.
+  private static final int DEFAULT_SERVER_STAGE_LIMIT = 100_000;
+  private final PhysicalPlannerContext _context;
+
+  public LiteModeSortInsertRule(PhysicalPlannerContext context) {
+    _context = context;
+  }
+
+  @Override
+  public boolean matches(PRelOptRuleCall call) {
+    return isLeafBoundary(call);
+  }
+
+  @Override
+  public PRelNode onMatch(PRelOptRuleCall call) {
+    RexNode newFetch = REX_BUILDER.makeLiteral(DEFAULT_SERVER_STAGE_LIMIT, TYPE_FACTORY.createSqlType(
+        SqlTypeName.INTEGER));
+    if (call._currentNode instanceof PhysicalSort) {
+      // When current node is a Sort, if it has a fetch already, verify it is less than the hard limit. Otherwise,
+      // set the configured hard limit within the same Sort.
+      PhysicalSort sort = (PhysicalSort) call._currentNode;
+      if (sort.fetch != null) {
+        int currentFetch = RexExpressionUtils.getValueAsInt(sort.fetch);
+        Preconditions.checkState(currentFetch <= DEFAULT_SERVER_STAGE_LIMIT,
+            "Attempted to stream %s records from server which exceed limit %s", currentFetch,
+            DEFAULT_SERVER_STAGE_LIMIT);
+        return sort;
+      }
+      return sort.withFetch(newFetch);
+    }
+    if (call._currentNode instanceof PhysicalAggregate) {
+      // When current node is aggregate, add the limit to the Aggregate itself and skip adding the Sort.
+      PhysicalAggregate aggregate = (PhysicalAggregate) call._currentNode;
+      Preconditions.checkState(aggregate.getLimit() <= DEFAULT_SERVER_STAGE_LIMIT,
+          "Group trim limit={} exceeds server stage limit={}", aggregate.getLimit(), DEFAULT_SERVER_STAGE_LIMIT);
+      // TODO(mse-physical): This resets the limit to server stage limit. Should we stick with group-trim limit?
+      return aggregate.withLimit(DEFAULT_SERVER_STAGE_LIMIT);
+    }
+    PRelNode input = call._currentNode;
+    return new PhysicalSort(input.unwrap().getCluster(), RelTraitSet.createEmpty(), List.of(),
+        RelCollations.EMPTY, null /* offset */, newFetch, input, nodeId(), input.getPinotDataDistributionOrThrow(),
+        true);
+  }
+
+  private int nodeId() {
+    return _context.getNodeIdGenerator().get();
+  }
+}

pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/opt/rules/LiteModeWorkerAssignmentRule.java

@@ -0,0 +1,112 @@
+/**
+ * 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.query.planner.physical.v2.opt.rules;
+
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Preconditions;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Random;
+import java.util.Set;
+import java.util.stream.Collectors;
+import javax.annotation.Nullable;
+import org.apache.calcite.rel.RelDistribution;
+import org.apache.pinot.calcite.rel.traits.PinotExecStrategyTrait;
+import org.apache.pinot.query.context.PhysicalPlannerContext;
+import org.apache.pinot.query.planner.physical.v2.ExchangeStrategy;
+import org.apache.pinot.query.planner.physical.v2.PRelNode;
+import org.apache.pinot.query.planner.physical.v2.PinotDataDistribution;
+import org.apache.pinot.query.planner.physical.v2.nodes.PhysicalExchange;
+import org.apache.pinot.query.planner.physical.v2.opt.PRelNodeTransformer;
+
+
+/**
+ * Lite mode uses a single worker for all stages except the leaf stage. Since leaf stage assignment is done before this
+ * rule is called, we simply need to sample a random worker from the leaf stage, and assign it to all the non-leaf
+ * plan nodes.
+ */
+public class LiteModeWorkerAssignmentRule implements PRelNodeTransformer {
+  private static final Random RANDOM = new Random();
+  private final PhysicalPlannerContext _context;
+
+  public LiteModeWorkerAssignmentRule(PhysicalPlannerContext context) {
+    _context = context;
+  }
+
+  @Override
+  public PRelNode execute(PRelNode currentNode) {
+    Set<String> workerSet = new HashSet<>();
+    accumulateWorkers(currentNode, workerSet);
+    List<String> workers = List.of(sampleWorker(new ArrayList<>(workerSet)));
+    PinotDataDistribution pdd = new PinotDataDistribution(RelDistribution.Type.SINGLETON, workers, workers.hashCode(),
+        null, null);
+    return addExchangeAndWorkers(currentNode, null, pdd);
+  }
+
+  public PRelNode addExchangeAndWorkers(PRelNode currentNode, @Nullable PRelNode parent, PinotDataDistribution pdd) {
+    if (currentNode.isLeafStage()) {
+      if (parent == null) {
+        return currentNode;
+      }
+      return new PhysicalExchange(nodeId(), currentNode, pdd, Collections.emptyList(),
+          ExchangeStrategy.SINGLETON_EXCHANGE, null, PinotExecStrategyTrait.getDefaultExecStrategy());
+    }
+    List<PRelNode> newInputs = new ArrayList<>();
+    for (PRelNode input : currentNode.getPRelInputs()) {
+      newInputs.add(addExchangeAndWorkers(input, currentNode, pdd));
+    }
+    return currentNode.with(newInputs, pdd);
+  }
+
+  /**
+   * Stores workers assigned to the leaf stage nodes into the provided Set. Note that each worker has an integer prefix
+   * which denotes the "workerId". We remove that prefix before storing them in the set.
+   */
+  @VisibleForTesting
+  static void accumulateWorkers(PRelNode currentNode, Set<String> workerSink) {
+    if (currentNode.isLeafStage()) {
+      workerSink.addAll(currentNode.getPinotDataDistributionOrThrow().getWorkers().stream()
+          .map(LiteModeWorkerAssignmentRule::stripIdPrefixFromWorker).collect(Collectors.toList()));
+      return;
+    }
+    for (PRelNode input : currentNode.getPRelInputs()) {
+      accumulateWorkers(input, workerSink);
+    }
+  }
+
+  /**
+   * Samples a worker from the given list.
+   */
+  @VisibleForTesting
+  static String sampleWorker(List<String> instanceIds) {
+    Preconditions.checkState(!instanceIds.isEmpty(), "No workers in leaf stage");
+    return String.format("0@%s", instanceIds.get(RANDOM.nextInt(instanceIds.size())));
+  }
+
+  @VisibleForTesting
+  static String stripIdPrefixFromWorker(String worker) {
+    return worker.split("@")[1];
+  }
+
+  private int nodeId() {
+    return _context.getNodeIdGenerator().get();
+  }
+}

pinot-query-planner/src/test/java/org/apache/pinot/query/planner/physical/v2/opt/rules/LiteModeWorkerAssignmentRuleTest.java

@@ -0,0 +1,66 @@
+/**
+ * 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.query.planner.physical.v2.opt.rules;
+
+import java.util.HashSet;
+import java.util.List;
+import java.util.Set;
+import org.apache.pinot.query.planner.physical.v2.PRelNode;
+import org.apache.pinot.query.planner.physical.v2.PinotDataDistribution;
+import org.mockito.Mockito;
+import org.testng.annotations.Test;
+
+import static org.mockito.Mockito.doReturn;
+import static org.testng.Assert.*;
+
+
+public class LiteModeWorkerAssignmentRuleTest {
+  @Test
+  public void testAccumulateWorkers() {
+    PRelNode leafOne = create(List.of(), true, List.of("0@server-1", "1@server-2"));
+    PRelNode leafTwo = create(List.of(), true, List.of("0@server-2", "1@server-1"));
+    PRelNode intermediateNode = create(List.of(leafOne, leafTwo), false, List.of("0@server-3", "1@server-4"));
+    Set<String> workers = new HashSet<>();
+    LiteModeWorkerAssignmentRule.accumulateWorkers(intermediateNode, workers);
+    assertEquals(workers, Set.of("server-1", "server-2"));
+  }
+
+  @Test
+  public void testSampleWorker() {
+    List<String> workers = List.of("worker-0", "worker-1", "worker-2");
+    Set<String> selectionCandidates = Set.of("0@worker-0", "0@worker-1", "0@worker-2");
+    Set<String> selectedWorkers = new HashSet<>();
+    for (int iteration = 0; iteration < 1000; iteration++) {
+      selectedWorkers.add(LiteModeWorkerAssignmentRule.sampleWorker(workers));
+    }
+    assertEquals(selectedWorkers, selectionCandidates);
+  }
+
+  private PRelNode create(List<PRelNode> inputs, boolean isLeafStage, List<String> workers) {
+    // Setup mock pinot data distribution.
+    PinotDataDistribution mockPDD = Mockito.mock(PinotDataDistribution.class);
+    doReturn(workers).when(mockPDD).getWorkers();
+    // Setup mock PRelNode.
+    PRelNode current = Mockito.mock(PRelNode.class);
+    doReturn(isLeafStage).when(current).isLeafStage();
+    doReturn(mockPDD).when(current).getPinotDataDistributionOrThrow();
+    doReturn(inputs).when(current).getPRelInputs();
+    return current;
+  }
+}