[multistage] Add Support for Inferring Invalid Segment Partition Id

pinot-common/src/main/java/org/apache/pinot/common/metrics/BrokerMeter.java

@@ -204,6 +204,14 @@ public class BrokerMeter implements AbstractMetrics.Meter {
    */
   public static final BrokerMeter WINDOW_COUNT = create("WINDOW_COUNT", "queries", true);
 
+  /**
+   * How many MSE queries have encountered segments with invalid partitions.
+   * <p>
+   * This is only emitted for when usePhysicalOptimizer is set to true.
+   */
+  public static final BrokerMeter INVALID_SEGMENT_PARTITION_IN_QUERY = create("INVALID_SEGMENT_PARTITION_IN_QUERY",
+      "queries", false);
+
   /**
    * Number of queries executed with cursors. This count includes queries that use SSE and MSE
    */

pinot-common/src/main/java/org/apache/pinot/common/utils/config/QueryOptionsUtils.java

@@ -379,6 +379,10 @@ public static Boolean isUseMSEToFillEmptySchema(Map<String, String> queryOptions
     return useMSEToFillEmptySchema != null ? Boolean.parseBoolean(useMSEToFillEmptySchema) : defaultValue;
   }
 
+  public static boolean isInferInvalidSegmentPartition(Map<String, String> queryOptions) {
+    return Boolean.parseBoolean(queryOptions.getOrDefault(QueryOptionKey.INFER_INVALID_SEGMENT_PARTITION, "false"));
+  }
+
   @Nullable
   private static Integer uncheckedParseInt(String optionName, @Nullable String optionValue) {
     if (optionValue == null) {

pinot-query-planner/src/main/java/org/apache/pinot/query/QueryEnvironment.java

@@ -79,6 +79,8 @@
 import org.apache.pinot.query.planner.logical.TransformationTracker;
 import org.apache.pinot.query.planner.physical.DispatchableSubPlan;
 import org.apache.pinot.query.planner.physical.PinotDispatchPlanner;
+import org.apache.pinot.query.planner.physical.v2.PRelNode;
+import org.apache.pinot.query.planner.physical.v2.PRelNodeTreeValidator;
 import org.apache.pinot.query.planner.physical.v2.PlanFragmentAndMailboxAssignment;
 import org.apache.pinot.query.planner.physical.v2.RelToPRelConverter;
 import org.apache.pinot.query.planner.plannode.PlanNode;
@@ -175,7 +177,7 @@ private PlannerContext getPlannerContext(SqlNodeAndOptions sqlNodeAndOptions) {
       workerManager = _envConfig.getWorkerManager();
       physicalPlannerContext = new PhysicalPlannerContext(workerManager.getRoutingManager(),
           workerManager.getHostName(), workerManager.getPort(), _envConfig.getRequestId(),
-          workerManager.getInstanceId());
+          workerManager.getInstanceId(), sqlNodeAndOptions.getOptions());
     }
     return new PlannerContext(_config, _catalogReader, _typeFactory, _optProgram, traitProgram,
         sqlNodeAndOptions.getOptions(), _envConfig, format, physicalPlannerContext);
@@ -333,6 +335,7 @@ private RelRoot compileQuery(SqlNode sqlNode, PlannerContext plannerContext) {
       Preconditions.checkNotNull(plannerContext.getPhysicalPlannerContext(), "Physical planner context is null");
       optimized = RelToPRelConverter.toPRelNode(optimized, plannerContext.getPhysicalPlannerContext(),
           _envConfig.getTableCache()).unwrap();
+      PRelNodeTreeValidator.validate((PRelNode) optimized);
     }
     return relation.withRel(optimized);
   }

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

@@ -56,6 +56,7 @@ public Integer get() {
    * Instance ID of the instance corresponding to this process.
    */
   private final String _instanceId;
+  private final Map<String, String> _queryOptions;
 
   /**
    * Used by controller when it needs to extract table names from the query.
@@ -67,15 +68,17 @@ public PhysicalPlannerContext() {
     _port = 0;
     _requestId = 0;
     _instanceId = "";
+    _queryOptions = Map.of();
   }
 
   public PhysicalPlannerContext(RoutingManager routingManager, String hostName, int port, long requestId,
-      String instanceId) {
+      String instanceId, Map<String, String> queryOptions) {
     _routingManager = routingManager;
     _hostName = hostName;
     _port = port;
     _requestId = requestId;
     _instanceId = instanceId;
+    _queryOptions = queryOptions == null ? Map.of() : queryOptions;
   }
 
   public Supplier<Integer> getNodeIdGenerator() {
@@ -107,6 +110,10 @@ public String getInstanceId() {
     return _instanceId;
   }
 
+  public Map<String, String> getQueryOptions() {
+    return _queryOptions;
+  }
+
   public static boolean isUsePhysicalOptimizer(@Nullable Map<String, String> queryOptions) {
     if (queryOptions == null) {
       return false;

pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/PinotLogicalQueryPlanner.java

@@ -40,6 +40,7 @@
 import org.apache.pinot.query.planner.SubPlan;
 import org.apache.pinot.query.planner.SubPlanMetadata;
 import org.apache.pinot.query.planner.physical.v2.PRelNode;
+import org.apache.pinot.query.planner.physical.v2.PRelNodeTreeValidator;
 import org.apache.pinot.query.planner.physical.v2.PlanFragmentAndMailboxAssignment;
 import org.apache.pinot.query.planner.plannode.BasePlanNode;
 import org.apache.pinot.query.planner.plannode.ExchangeNode;
@@ -95,6 +96,8 @@ public static SubPlan makePlan(RelRoot relRoot,
   public static Pair<SubPlan, PlanFragmentAndMailboxAssignment.Result> makePlanV2(RelRoot relRoot,
       PhysicalPlannerContext physicalPlannerContext) {
     PRelNode pRelNode = (PRelNode) relRoot.rel;
+    // TODO(mse-physical): Don't emit metrics for explain statements.
+    PRelNodeTreeValidator.emitMetrics(pRelNode);
     PlanFragmentAndMailboxAssignment planFragmentAndMailboxAssignment = new PlanFragmentAndMailboxAssignment();
     PlanFragmentAndMailboxAssignment.Result result =
         planFragmentAndMailboxAssignment.compute(pRelNode, physicalPlannerContext);

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

@@ -0,0 +1,76 @@
+/**
+ * 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;
+
+import org.apache.calcite.rel.core.Join;
+import org.apache.calcite.rel.core.Window;
+import org.apache.pinot.common.metrics.BrokerMeter;
+import org.apache.pinot.common.metrics.BrokerMetrics;
+
+
+/**
+ * Centralizes validations for the optimized PRelNode tree.
+ */
+public class PRelNodeTreeValidator {
+  private static final BrokerMetrics BROKER_METRICS = BrokerMetrics.get();
+
+  private PRelNodeTreeValidator() {
+  }
+
+  /**
+   * Validate the tree rooted at the given PRelNode. Ideally all issues with the plan should be caught even with an
+   * EXPLAIN, hence this method should be called as part of query compilation itself.
+   */
+  public static void validate(PRelNode rootNode) {
+    // TODO(mse-physical): Add plan validations here.
+  }
+
+  /**
+   * Emit metrics about the given plan tree. This should be avoided for Explain statements since metrics are not really
+   * helpful there and can be misleading.
+   */
+  public static void emitMetrics(PRelNode pRelNode) {
+    Context context = new Context();
+    walk(pRelNode, context);
+    if (context._joinCount > 0) {
+      BROKER_METRICS.addMeteredGlobalValue(BrokerMeter.JOIN_COUNT, context._joinCount);
+      BROKER_METRICS.addMeteredGlobalValue(BrokerMeter.QUERIES_WITH_JOINS, 1);
+    }
+    if (context._windowCount > 0) {
+      BROKER_METRICS.addMeteredGlobalValue(BrokerMeter.WINDOW_COUNT, context._windowCount);
+      BROKER_METRICS.addMeteredGlobalValue(BrokerMeter.QUERIES_WITH_WINDOW, 1);
+    }
+  }
+
+  private static void walk(PRelNode pRelNode, Context context) {
+    if (pRelNode.unwrap() instanceof Join) {
+      context._joinCount++;
+    } else if (pRelNode.unwrap() instanceof Window) {
+      context._windowCount++;
+    }
+    for (PRelNode input : pRelNode.getPRelInputs()) {
+      walk(input, context);
+    }
+  }
+
+  private static class Context {
+    int _joinCount = 0;
+    int _windowCount = 0;
+  }
+}

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

@@ -22,7 +22,6 @@
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Set;
-import javax.annotation.Nullable;
 import org.apache.calcite.plan.RelOptTable;
 import org.apache.calcite.rel.RelDistribution;
 import org.apache.calcite.rel.RelNode;
@@ -44,14 +43,12 @@
 import org.apache.calcite.sql.SqlLiteral;
 import org.apache.calcite.sql.type.SqlTypeName;
 import org.apache.pinot.calcite.rel.hint.PinotHintOptions;
-import org.apache.pinot.common.metrics.BrokerMetrics;
 import org.apache.pinot.common.utils.DataSchema;
 import org.apache.pinot.common.utils.DataSchema.ColumnDataType;
 import org.apache.pinot.common.utils.DatabaseUtils;
 import org.apache.pinot.common.utils.request.RequestUtils;
 import org.apache.pinot.query.planner.logical.RexExpression;
 import org.apache.pinot.query.planner.logical.RexExpressionUtils;
-import org.apache.pinot.query.planner.logical.TransformationTracker;
 import org.apache.pinot.query.planner.physical.v2.nodes.PhysicalAggregate;
 import org.apache.pinot.query.planner.physical.v2.nodes.PhysicalExchange;
 import org.apache.pinot.query.planner.physical.v2.nodes.PhysicalJoin;
@@ -75,14 +72,7 @@ public class PRelToPlanNodeConverter {
   private static final Logger LOGGER = LoggerFactory.getLogger(PRelToPlanNodeConverter.class);
   private static final int DEFAULT_STAGE_ID = -1;
 
-  private final BrokerMetrics _brokerMetrics = BrokerMetrics.get();
-  private boolean _joinFound;
-  private boolean _windowFunctionFound;
-  @Nullable
-  private final TransformationTracker.Builder<PlanNode, RelNode> _tracker;
-
-  public PRelToPlanNodeConverter(@Nullable TransformationTracker.Builder<PlanNode, RelNode> tracker) {
-    _tracker = tracker;
+  private PRelToPlanNodeConverter() {
   }
 
   /**
@@ -105,18 +95,8 @@ public static PlanNode toPlanNode(PRelNode pRelNode, int stageId) {
     } else if (node instanceof Exchange) {
       result = convertPhysicalExchange((PhysicalExchange) node);
     } else if (node instanceof PhysicalJoin) {
-      /* _brokerMetrics.addMeteredGlobalValue(BrokerMeter.JOIN_COUNT, 1);
-      if (!_joinFound) {
-        _brokerMetrics.addMeteredGlobalValue(BrokerMeter.QUERIES_WITH_JOINS, 1);
-        _joinFound = true;
-      } */
       result = convertJoin((PhysicalJoin) node);
     } else if (node instanceof Window) {
-      /* _brokerMetrics.addMeteredGlobalValue(BrokerMeter.WINDOW_COUNT, 1);
-      if (!_windowFunctionFound) {
-        _brokerMetrics.addMeteredGlobalValue(BrokerMeter.QUERIES_WITH_WINDOW, 1);
-        _windowFunctionFound = true;
-      } */
       result = convertWindow((Window) node);
     } else if (node instanceof Values) {
       result = convertValues((Values) node);
@@ -126,9 +106,6 @@ public static PlanNode toPlanNode(PRelNode pRelNode, int stageId) {
       throw new IllegalStateException("Unsupported RelNode: " + node);
     }
     result.setStageId(stageId);
-    /* if (_tracker != null) {
-      _tracker.trackCreation(node, result);
-    } */
     return result;
   }