Introduce OR Predicate Execution On Star Tree Index

pinot-core/src/main/java/org/apache/pinot/core/plan/AggregationGroupByOrderByPlanNode.java

@@ -22,11 +22,11 @@
 import java.util.Map;
 import java.util.Set;
 import org.apache.pinot.common.request.context.ExpressionContext;
-import org.apache.pinot.core.operator.filter.predicate.PredicateEvaluator;
 import org.apache.pinot.core.operator.query.AggregationGroupByOrderByOperator;
 import org.apache.pinot.core.query.aggregation.function.AggregationFunction;
 import org.apache.pinot.core.query.aggregation.function.AggregationFunctionUtils;
 import org.apache.pinot.core.query.request.context.QueryContext;
+import org.apache.pinot.core.startree.CompositePredicate;
 import org.apache.pinot.core.startree.StarTreeUtils;
 import org.apache.pinot.core.startree.plan.StarTreeTransformPlanNode;
 import org.apache.pinot.segment.spi.IndexSegment;
@@ -67,7 +67,7 @@ public AggregationGroupByOrderByPlanNode(IndexSegment indexSegment, QueryContext
       AggregationFunctionColumnPair[] aggregationFunctionColumnPairs =
           StarTreeUtils.extractAggregationFunctionPairs(_aggregationFunctions);
       if (aggregationFunctionColumnPairs != null) {
-        Map<String, List<PredicateEvaluator>> predicateEvaluatorsMap =
+        Map<String, List<CompositePredicate>> predicateEvaluatorsMap =
             StarTreeUtils.extractPredicateEvaluatorsMap(indexSegment, queryContext.getFilter());
         if (predicateEvaluatorsMap != null) {
           for (StarTreeV2 starTreeV2 : starTrees) {

pinot-core/src/main/java/org/apache/pinot/core/plan/AggregationGroupByPlanNode.java

@@ -22,11 +22,11 @@
 import java.util.Map;
 import java.util.Set;
 import org.apache.pinot.common.request.context.ExpressionContext;
-import org.apache.pinot.core.operator.filter.predicate.PredicateEvaluator;
 import org.apache.pinot.core.operator.query.AggregationGroupByOperator;
 import org.apache.pinot.core.query.aggregation.function.AggregationFunction;
 import org.apache.pinot.core.query.aggregation.function.AggregationFunctionUtils;
 import org.apache.pinot.core.query.request.context.QueryContext;
+import org.apache.pinot.core.startree.CompositePredicate;
 import org.apache.pinot.core.startree.StarTreeUtils;
 import org.apache.pinot.core.startree.plan.StarTreeTransformPlanNode;
 import org.apache.pinot.segment.spi.IndexSegment;
@@ -64,7 +64,7 @@ public AggregationGroupByPlanNode(IndexSegment indexSegment, QueryContext queryC
       AggregationFunctionColumnPair[] aggregationFunctionColumnPairs =
           StarTreeUtils.extractAggregationFunctionPairs(_aggregationFunctions);
       if (aggregationFunctionColumnPairs != null) {
-        Map<String, List<PredicateEvaluator>> predicateEvaluatorsMap =
+        Map<String, List<CompositePredicate>> predicateEvaluatorsMap =
             StarTreeUtils.extractPredicateEvaluatorsMap(indexSegment, queryContext.getFilter());
         if (predicateEvaluatorsMap != null) {
           for (StarTreeV2 starTreeV2 : starTrees) {

pinot-core/src/main/java/org/apache/pinot/core/plan/AggregationPlanNode.java

@@ -22,11 +22,11 @@
 import java.util.Map;
 import java.util.Set;
 import org.apache.pinot.common.request.context.ExpressionContext;
-import org.apache.pinot.core.operator.filter.predicate.PredicateEvaluator;
 import org.apache.pinot.core.operator.query.AggregationOperator;
 import org.apache.pinot.core.query.aggregation.function.AggregationFunction;
 import org.apache.pinot.core.query.aggregation.function.AggregationFunctionUtils;
 import org.apache.pinot.core.query.request.context.QueryContext;
+import org.apache.pinot.core.startree.CompositePredicate;
 import org.apache.pinot.core.startree.StarTreeUtils;
 import org.apache.pinot.core.startree.plan.StarTreeTransformPlanNode;
 import org.apache.pinot.segment.spi.IndexSegment;
@@ -55,13 +55,14 @@ public AggregationPlanNode(IndexSegment indexSegment, QueryContext queryContext)
       AggregationFunctionColumnPair[] aggregationFunctionColumnPairs =
           StarTreeUtils.extractAggregationFunctionPairs(_aggregationFunctions);
       if (aggregationFunctionColumnPairs != null) {
-        Map<String, List<PredicateEvaluator>> predicateEvaluatorsMap =
+        Map<String, List<CompositePredicate>> predicateEvaluatorsMap =
             StarTreeUtils.extractPredicateEvaluatorsMap(indexSegment, queryContext.getFilter());
         if (predicateEvaluatorsMap != null) {
           for (StarTreeV2 starTreeV2 : starTrees) {
             if (StarTreeUtils.isFitForStarTree(starTreeV2.getMetadata(), aggregationFunctionColumnPairs, null,
                 predicateEvaluatorsMap.keySet())) {
               _transformPlanNode = null;
+
               _starTreeTransformPlanNode =
                   new StarTreeTransformPlanNode(starTreeV2, aggregationFunctionColumnPairs, null,
                       predicateEvaluatorsMap, queryContext.getDebugOptions());

pinot-core/src/main/java/org/apache/pinot/core/startree/CompositePredicate.java

@@ -0,0 +1,57 @@
+/**
+ * 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.core.startree;
+
+import java.util.ArrayList;
+import java.util.List;
+import org.apache.pinot.common.request.context.FilterContext;
+import org.apache.pinot.core.operator.filter.predicate.PredicateEvaluator;
+
+
+/**
+ * Represents a composite predicate.
+ *
+ * A composite predicate represents a set of predicates conjoined by a given relation.
+ * Consider the given predicate: (d1 > 10 OR d1 < 50). A composite predicate will represent
+ * two predicates -- d1 > 10 and d1 < 50 and represent that they are related by the operator OR.
+ */
+public class CompositePredicate {
+  private final FilterContext.Type _filterContextType;
+  private final List<PredicateEvaluator> _predicateEvaluators;
+
+  public CompositePredicate(FilterContext.Type filterContextType) {
+    _filterContextType = filterContextType;
+
+    _predicateEvaluators = new ArrayList<>();
+  }
+
+  public void addPredicateEvaluator(PredicateEvaluator predicateEvaluator) {
+    assert predicateEvaluator != null;
+
+    _predicateEvaluators.add(predicateEvaluator);
+  }
+
+  public List<PredicateEvaluator> getPredicateEvaluators() {
+    return _predicateEvaluators;
+  }
+
+  public FilterContext.Type getFilterContextType() {
+    return _filterContextType;
+  }
+}

pinot-core/src/main/java/org/apache/pinot/core/startree/StarTreeUtils.java

@@ -28,6 +28,7 @@
 import java.util.Queue;
 import java.util.Set;
 import javax.annotation.Nullable;
+import org.apache.commons.lang3.tuple.Pair;
 import org.apache.pinot.common.request.context.ExpressionContext;
 import org.apache.pinot.common.request.context.FilterContext;
 import org.apache.pinot.common.request.context.predicate.Predicate;
@@ -84,15 +85,19 @@ public static AggregationFunctionColumnPair[] extractAggregationFunctionPairs(
   /**
    * Extracts a map from the column to a list of {@link PredicateEvaluator}s for it. Returns {@code null} if the filter
    * cannot be solved by the star-tree.
+   *
+   * A predicate can be simple (d1 > 10) or composite (d1 > 10 AND d2 < 50) or multi levelled (d1 > 50 AND (d2 > 10 OR d2 < 35)
+   * This method represents a list of CompositePredicates per dimension. For each dimension, all CompositePredicates in the
+   * list are implicitly ANDed together. Any OR predicates are nested within a CompositePredicate.
    */
   @Nullable
-  public static Map<String, List<PredicateEvaluator>> extractPredicateEvaluatorsMap(IndexSegment indexSegment,
+  public static Map<String, List<CompositePredicate>> extractPredicateEvaluatorsMap(IndexSegment indexSegment,
       @Nullable FilterContext filter) {
     if (filter == null) {
       return Collections.emptyMap();
     }
 
-    Map<String, List<PredicateEvaluator>> predicateEvaluatorsMap = new HashMap<>();
+    Map<String, List<CompositePredicate>> predicateEvaluatorsMap = new HashMap<>();
     Queue<FilterContext> queue = new LinkedList<>();
     queue.add(filter);
     FilterContext filterNode;
@@ -102,8 +107,23 @@ public static Map<String, List<PredicateEvaluator>> extractPredicateEvaluatorsMa
           queue.addAll(filterNode.getChildren());
           break;
         case OR:
-          // Star-tree does not support OR filter
-          return null;
+          Pair<Boolean, String> result = isOrClauseValidForStarTree(filterNode);
+
+          if (result.getLeft()) {
+            assert result.getRight() != null;
+            String column = result.getRight();
+
+            CompositePredicate compositePredicate = new CompositePredicate(
+                FilterContext.Type.OR);
+
+            filterNode.getChildren().forEach(k -> compositePredicate
+                .addPredicateEvaluator(getPredicateEvaluatorForPredicate(indexSegment, k)));
+
+            predicateEvaluatorsMap.computeIfAbsent(column, k -> new ArrayList<>()).add(compositePredicate);
+            break;
+          } else {
+            return null;
+          }
         case PREDICATE:
           Predicate predicate = filterNode.getPredicate();
           ExpressionContext lhs = predicate.getLhs();
@@ -112,30 +132,15 @@ public static Map<String, List<PredicateEvaluator>> extractPredicateEvaluatorsMa
             return null;
           }
           String column = lhs.getIdentifier();
-          DataSource dataSource = indexSegment.getDataSource(column);
-          Dictionary dictionary = dataSource.getDictionary();
-          if (dictionary == null) {
-            // Star-tree does not support non-dictionary encoded dimension
-            return null;
-          }
-          switch (predicate.getType()) {
-            // Do not use star-tree for the following predicates because:
-            //   - REGEXP_LIKE: Need to scan the whole dictionary to gather the matching dictionary ids
-            //   - TEXT_MATCH/IS_NULL/IS_NOT_NULL: No way to gather the matching dictionary ids
-            case REGEXP_LIKE:
-            case TEXT_MATCH:
-            case IS_NULL:
-            case IS_NOT_NULL:
-              return null;
-          }
-          PredicateEvaluator predicateEvaluator = PredicateEvaluatorProvider
-              .getPredicateEvaluator(predicate, dictionary, dataSource.getDataSourceMetadata().getDataType());
-          if (predicateEvaluator.isAlwaysFalse()) {
-            // Do not use star-tree if there is no matching record
-            return null;
-          }
-          if (!predicateEvaluator.isAlwaysTrue()) {
-            predicateEvaluatorsMap.computeIfAbsent(column, k -> new ArrayList<>()).add(predicateEvaluator);
+
+          PredicateEvaluator predicateEvaluator = getPredicateEvaluatorForPredicate(indexSegment, filterNode);
+
+          if (predicateEvaluator != null && !predicateEvaluator.isAlwaysTrue()) {
+            CompositePredicate compositePredicate = new CompositePredicate(
+                FilterContext.Type.PREDICATE);
+
+            compositePredicate.addPredicateEvaluator(predicateEvaluator);
+            predicateEvaluatorsMap.computeIfAbsent(column, k -> new ArrayList<>()).add(compositePredicate);
           }
           break;
         default:
@@ -179,4 +184,96 @@ public static boolean isFitForStarTree(StarTreeV2Metadata starTreeV2Metadata,
     // Check predicate columns
     return starTreeDimensions.containsAll(predicateColumns);
   }
+
+  /**
+   * Evaluates a given filter to ascertain if the OR clause is valid for StarTree processing.
+   *
+   * StarTree supports OR predicates on a single dimension only (d1 < 10 OR d1 > 50).
+   *
+   * @return Pair of the result and the single literal on which the predicate is based. If the
+   * predicate is not valid for execution on StarTree, literal value will be null
+   */
+  private static Pair<Boolean, String> isOrClauseValidForStarTree(FilterContext filterContext) {
+    assert filterContext != null;
+
+    Set<String> seenLiterals = new HashSet<>();
+
+    if (!isOrClauseValidForStarTreeInternal(filterContext, seenLiterals)) {
+      return Pair.of(false, null);
+    }
+
+    boolean result = seenLiterals.size() == 1;
+
+    return Pair.of(result, result ? seenLiterals.iterator().next() : null);
+  }
+
+  /** Internal processor for the above evaluator */
+  private static boolean isOrClauseValidForStarTreeInternal(FilterContext filterContext, Set<String> seenLiterals) {
+    assert filterContext != null;
+
+    if (filterContext.getType() == FilterContext.Type.OR) {
+      List<FilterContext> childFilterContexts = filterContext.getChildren();
+
+      for (FilterContext childFilterContext : childFilterContexts) {
+        isOrClauseValidForStarTreeInternal(childFilterContext, seenLiterals);
+      }
+    } else if (filterContext.getType() == FilterContext.Type.PREDICATE) {
+      String literalValue = validateExpressionAndExtractIdentifier(filterContext.getPredicate());
+
+      if (literalValue != null) {
+        seenLiterals.add(literalValue);
+      } else {
+        return false;
+      }
+    }
+
+    return true;
+  }
+
+  /** Checks if the given predicate has an expression which is of type identifier and returns it */
+  private static String validateExpressionAndExtractIdentifier(Predicate predicate) {
+    assert predicate != null;
+
+    ExpressionContext expressionContext = predicate.getLhs();
+
+   if (expressionContext.getType() == ExpressionContext.Type.IDENTIFIER) {
+      return expressionContext.getIdentifier();
+   }
+
+    return null;
+  }
+
+  private static PredicateEvaluator getPredicateEvaluatorForPredicate(IndexSegment indexSegment, FilterContext filterNode) {
+    Predicate predicate = filterNode.getPredicate();
+    ExpressionContext lhs = predicate.getLhs();
+    if (lhs.getType() != ExpressionContext.Type.IDENTIFIER) {
+      // Star-tree does not support non-identifier expression
+      return null;
+    }
+    String column = lhs.getIdentifier();
+    DataSource dataSource = indexSegment.getDataSource(column);
+    Dictionary dictionary = dataSource.getDictionary();
+    if (dictionary == null) {
+      // Star-tree does not support non-dictionary encoded dimension
+      return null;
+    }
+    switch (predicate.getType()) {
+      // Do not use star-tree for the following predicates because:
+      //   - REGEXP_LIKE: Need to scan the whole dictionary to gather the matching dictionary ids
+      //   - TEXT_MATCH/IS_NULL/IS_NOT_NULL: No way to gather the matching dictionary ids
+      case REGEXP_LIKE:
+      case TEXT_MATCH:
+      case IS_NULL:
+      case IS_NOT_NULL:
+        return null;
+    }
+    PredicateEvaluator predicateEvaluator = PredicateEvaluatorProvider
+        .getPredicateEvaluator(predicate, dictionary, dataSource.getDataSourceMetadata().getDataType());
+    if (predicateEvaluator.isAlwaysFalse()) {
+      // Do not use star-tree if there is no matching record
+      return null;
+    }
+
+    return predicateEvaluator;
+  }
 }