Introduce OR Predicate Execution On Star Tree Index

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

@@ -22,17 +22,18 @@
 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.CompositePredicateEvaluator;
 import org.apache.pinot.core.startree.StarTreeUtils;
 import org.apache.pinot.core.startree.plan.StarTreeTransformPlanNode;
 import org.apache.pinot.segment.spi.IndexSegment;
 import org.apache.pinot.segment.spi.index.startree.AggregationFunctionColumnPair;
 import org.apache.pinot.segment.spi.index.startree.StarTreeV2;
 
+
 /**
  * The <code>AggregationGroupByOrderByPlanNode</code> class provides the execution plan for aggregation group-by order-by query on a
  * single segment.
@@ -67,7 +68,7 @@ public AggregationGroupByOrderByPlanNode(IndexSegment indexSegment, QueryContext
       AggregationFunctionColumnPair[] aggregationFunctionColumnPairs =
           StarTreeUtils.extractAggregationFunctionPairs(_aggregationFunctions);
       if (aggregationFunctionColumnPairs != null) {
-        Map<String, List<PredicateEvaluator>> predicateEvaluatorsMap =
+        Map<String, List<CompositePredicateEvaluator>> 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.CompositePredicateEvaluator;
 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<CompositePredicateEvaluator>> 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.CompositePredicateEvaluator;
 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<CompositePredicateEvaluator>> 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/CompositePredicateEvaluator.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.List;
+import org.apache.pinot.core.operator.filter.predicate.PredicateEvaluator;
+
+
+/**
+ * Represents a composite predicate.
+ *
+ * A composite predicate evaluator represents a single predicate evaluator or multiple predicate evaluators conjoined
+ * with OR.
+ * 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 CompositePredicateEvaluator {
+  private final List<PredicateEvaluator> _predicateEvaluators;
+
+  public CompositePredicateEvaluator(List<PredicateEvaluator> predicateEvaluators) {
+    assert !predicateEvaluators.isEmpty();
+    _predicateEvaluators = predicateEvaluators;
+  }
+
+  public List<PredicateEvaluator> getPredicateEvaluators() {
+    return _predicateEvaluators;
+  }
+
+  /**
+   * Applies a dictionary id to the composite predicate evaluator. Returns {@code true} if the dictionary id matches any
+   * predicate evaluator, {@code false} otherwise.
+   */
+  public boolean apply(int dictId) {
+    for (PredicateEvaluator predicateEvaluator : _predicateEvaluators) {
+      if (predicateEvaluator.applySV(dictId)) {
+        return true;
+      }
+    }
+    return false;
+  }
+}

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,20 @@ 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<CompositePredicateEvaluator>> extractPredicateEvaluatorsMap(IndexSegment indexSegment,
       @Nullable FilterContext filter) {
     if (filter == null) {
       return Collections.emptyMap();
     }
 
-    Map<String, List<PredicateEvaluator>> predicateEvaluatorsMap = new HashMap<>();
+    Map<String, List<CompositePredicateEvaluator>> predicateEvaluatorsMap = new HashMap<>();
     Queue<FilterContext> queue = new LinkedList<>();
     queue.add(filter);
     FilterContext filterNode;
@@ -102,40 +108,26 @@ public static Map<String, List<PredicateEvaluator>> extractPredicateEvaluatorsMa
           queue.addAll(filterNode.getChildren());
           break;
         case OR:
-          // Star-tree does not support OR filter
-          return null;
-        case PREDICATE:
-          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
+          Pair<String, List<PredicateEvaluator>> pair = isOrClauseValidForStarTree(indexSegment, filterNode);
+          if (pair == null) {
             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;
+          List<PredicateEvaluator> predicateEvaluators = pair.getRight();
+          // NOTE: Empty list means always true
+          if (!predicateEvaluators.isEmpty()) {
+            predicateEvaluatorsMap.computeIfAbsent(pair.getLeft(), k -> new ArrayList<>())
+                .add(new CompositePredicateEvaluator(predicateEvaluators));
           }
-          PredicateEvaluator predicateEvaluator = PredicateEvaluatorProvider
-              .getPredicateEvaluator(predicate, dictionary, dataSource.getDataSourceMetadata().getDataType());
-          if (predicateEvaluator.isAlwaysFalse()) {
-            // Do not use star-tree if there is no matching record
+          break;
+        case PREDICATE:
+          Predicate predicate = filterNode.getPredicate();
+          PredicateEvaluator predicateEvaluator = getPredicateEvaluatorForPredicate(indexSegment, predicate);
+          if (predicateEvaluator == null) {
             return null;
           }
           if (!predicateEvaluator.isAlwaysTrue()) {
-            predicateEvaluatorsMap.computeIfAbsent(column, k -> new ArrayList<>()).add(predicateEvaluator);
+            predicateEvaluatorsMap.computeIfAbsent(predicate.getLhs().getIdentifier(), k -> new ArrayList<>())
+                .add(new CompositePredicateEvaluator(Collections.singletonList(predicateEvaluator)));
           }
           break;
         default:
@@ -179,4 +171,103 @@ public static boolean isFitForStarTree(StarTreeV2Metadata starTreeV2Metadata,
     // Check predicate columns
     return starTreeDimensions.containsAll(predicateColumns);
   }
+
+  /**
+   * Evaluates whether the given OR clause is valid for StarTree processing.
+   * StarTree supports OR predicates on a single dimension only (d1 < 10 OR d1 > 50).
+   *
+   * @return The pair of single identifier and predicate evaluators applied to it if true; {@code null} if the OR clause
+   *         cannot be solved with star-tree; empty predicate evaluator list if the OR clause always evaluates to true.
+   */
+  @Nullable
+  private static Pair<String, List<PredicateEvaluator>> isOrClauseValidForStarTree(IndexSegment indexSegment,
+      FilterContext filter) {
+    assert filter.getType() == FilterContext.Type.OR;
+
+    List<Predicate> predicates = new ArrayList<>();
+    extractOrClausePredicates(filter, predicates);
+
+    String identifier = null;
+    List<PredicateEvaluator> predicateEvaluators = new ArrayList<>();
+    for (Predicate predicate : predicates) {
+      PredicateEvaluator predicateEvaluator = getPredicateEvaluatorForPredicate(indexSegment, predicate);
+      if (predicateEvaluator == null) {
+        // The predicate cannot be solved with star-tree
+        return null;
+      }
+      if (predicateEvaluator.isAlwaysTrue()) {
+        // Use empty predicate evaluators to represent always true
+        return Pair.of(null, Collections.emptyList());
+      }
+      if (!predicateEvaluator.isAlwaysFalse()) {
+        String predicateIdentifier = predicate.getLhs().getIdentifier();
+        if (identifier == null) {
+          identifier = predicateIdentifier;
+        } else {
+          if (!identifier.equals(predicateIdentifier)) {
+            // The predicates are applied to multiple columns
+            return null;
+          }
+        }
+        predicateEvaluators.add(predicateEvaluator);
+      }
+    }
+    return Pair.of(identifier, predicateEvaluators);
+  }
+
+  /**
+   * Extracts the predicates under the given OR clause, returns {@code false} if there is nested AND under OR clause.
+   */
+  private static boolean extractOrClausePredicates(FilterContext filter, List<Predicate> predicates) {
+    assert filter.getType() == FilterContext.Type.OR;
+
+    for (FilterContext child : filter.getChildren()) {
+      switch (child.getType()) {
+        case AND:
+          return false;
+        case OR:
+          if (!extractOrClausePredicates(child, predicates)) {
+            return false;
+          }
+        case PREDICATE:
+          predicates.add(child.getPredicate());
+          break;
+        default:
+          throw new IllegalStateException();
+      }
+    }
+    return true;
+  }
+
+  /**
+   * Returns the predicate evaluator for the given predicate, or {@code null} if the predicate cannot be solved with
+   * star-tree.
+   */
+  @Nullable
+  private static PredicateEvaluator getPredicateEvaluatorForPredicate(IndexSegment indexSegment, Predicate predicate) {
+    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;
+    }
+    return PredicateEvaluatorProvider
+        .getPredicateEvaluator(predicate, dictionary, dataSource.getDataSourceMetadata().getDataType());
+  }
 }