perf(query): reduce regex query fanout

coordinator/src/main/scala/filodb.coordinator/queryplanner/LogicalPlanUtils.scala

@@ -6,7 +6,8 @@ import scala.collection.mutable.ArrayBuffer
 import com.typesafe.scalalogging.StrictLogging
 
 import filodb.core.TargetSchemaProvider
-import filodb.core.query.{ColumnFilter, QueryContext, RangeParams}
+import filodb.core.query.{ColumnFilter, QueryContext, QueryUtils, RangeParams}
+import filodb.core.query.Filter.{Equals, EqualsRegex}
 import filodb.prometheus.ast.SubqueryUtils
 import filodb.prometheus.ast.Vectors.PromMetricLabel
 import filodb.prometheus.ast.WindowConstants
@@ -28,17 +29,17 @@ object LogicalPlanUtils extends StrictLogging {
   }
 
   /**
-   * Given a LogicalPlan check if any descendent of plan is an aggregate operation
+   * Given a LogicalPlan check if any descendent of plan is either an aggregate or join operation
    * @param lp the LogicalPlan instance
-   * @return true if a descendent is an aggregate else false
+   * @return true if a descendent is an aggregate or join else false
    */
-  def hasDescendantAggregate(lp: LogicalPlan): Boolean = lp match {
+  def hasDescendantAggregateOrJoin(lp: LogicalPlan): Boolean = lp match {
     case _: Aggregate                 => true
     // consider this BinaryJoin example foo + on(h) + bar.
     // partition1 has foo{h=1}, bar{h1=2} and partition2 has foo{h=2}, bar{h1=1}
     // the binary join cannot happen on a partition locally. InProcessPlanDispatcher is required.
     case _: BinaryJoin               => true
-    case nonLeaf: NonLeafLogicalPlan  => nonLeaf.children.exists(hasDescendantAggregate(_))
+    case nonLeaf: NonLeafLogicalPlan  => nonLeaf.children.exists(hasDescendantAggregateOrJoin(_))
     case _                            => false
   }
   /**
@@ -420,9 +421,22 @@ object LogicalPlanUtils extends StrictLogging {
     val rsTschemaOpts = LogicalPlan.findLeafLogicalPlans(plan)
       .filter(_.isInstanceOf[RawSeries])
       .map(_.asInstanceOf[RawSeries]).flatMap{ rs =>
-      val shardKeyFilters = getShardKeyFilters(rs)
-      val interval = LogicalPlanUtils.getSpanningIntervalSelector(rs)
-      shardKeyFilters.map{ shardKey =>
+        val interval = LogicalPlanUtils.getSpanningIntervalSelector(rs)
+        val rawShardKeyFilters = getShardKeyFilters(rs)
+        // The filters might contain pipe-concatenated EqualsRegex values.
+        // Convert these into sets of single-valued Equals filters.
+        val resolvedShardKeyFilters = rawShardKeyFilters.flatMap{ filters =>
+          val equalsFilters: Seq[Seq[ColumnFilter]] = filters.map { filter =>
+            filter.filter match {
+              case EqualsRegex(values: String) if QueryUtils.containsPipeOnlyRegex(values) =>
+                QueryUtils.splitAtUnescapedPipes(values).map(value => ColumnFilter(filter.column, Equals(value)))
+              case _ => Seq(filter)
+            }}
+          // E.g. foo{key1=~"baz|bat",key2=~"bar|bak"} would give the following combos:
+          // [[baz,bar], [baz,bak], [bat,bar], [bat,bak]]
+          QueryUtils.combinations(equalsFilters)
+      }
+      resolvedShardKeyFilters.map{ shardKey =>
         val filters = LogicalPlanUtils.upsertFilters(rs.filters, shardKey)
         LogicalPlanUtils.getTargetSchemaIfUnchanging(targetSchemaProvider, filters, interval)
       }

coordinator/src/main/scala/filodb.coordinator/queryplanner/PartitionLocationPlanner.scala

@@ -0,0 +1,94 @@
+package filodb.coordinator.queryplanner
+
+import filodb.coordinator.queryplanner.LogicalPlanUtils.getLookBackMillis
+import filodb.core.metadata.Dataset
+import filodb.core.query.{PromQlQueryParams, QueryUtils}
+import filodb.core.query.Filter.{Equals, EqualsRegex}
+import filodb.query.LogicalPlan
+import filodb.query.exec.{ExecPlan, PromQlRemoteExec}
+
+/**
+ * Abstract class for planners that need getPartitions functionality.
+ *
+ * FIXME: the ShardKeyRegexPlanner and MultiPartitionPlanner share purpose/responsibility
+ *   and should eventually be merged. Currently, the SKRP needs getPartitions to group
+ *   resolved shard-keys by partition before it individually materializes each of these
+ *   groups with the MPP. The MPP will again find the corresponding partition for each group
+ *   and materialize accordingly.
+ */
+abstract class PartitionLocationPlanner(dataset: Dataset,
+                                        partitionLocationProvider: PartitionLocationProvider)
+  extends QueryPlanner with DefaultPlanner {
+
+  // scalastyle:off method.length
+  /**
+   * Gets the partition Assignment for the given plan
+   */
+  protected def getPartitions(logicalPlan: LogicalPlan,
+                              queryParams: PromQlQueryParams,
+                              infiniteTimeRange: Boolean = false) : Seq[PartitionAssignment] = {
+
+    //1.  Get a Seq of all Leaf node filters
+    val leafFilters = LogicalPlan.getColumnFilterGroup(logicalPlan)
+    val nonMetricColumnSet = dataset.options.nonMetricShardColumns.toSet
+    //2. Filter from each leaf node filters to keep only nonMetricShardKeyColumns and convert them to key value map
+    val routingKeyMap: Seq[Map[String, String]] = leafFilters
+      .filter(_.nonEmpty)
+      .map(_.filter(col => nonMetricColumnSet.contains(col.column)))
+      .map{ filters =>
+        filters.map { filter =>
+          val values = filter.filter match {
+            case Equals(value) => Seq(value.toString)
+            case EqualsRegex(value: String) if QueryUtils.containsPipeOnlyRegex(value) =>
+              QueryUtils.splitAtUnescapedPipes(value)
+            case _ => throw new IllegalArgumentException(
+              s"""shard keys must be filtered by equality or "|"-only regex. filter=${filter}""")
+          }
+          (filter.column, values)
+        }
+      }
+      .flatMap(keyToVals => QueryUtils.makeAllKeyValueCombos(keyToVals.toMap))
+
+    // 3. Determine the query time range
+    val queryTimeRange = if (infiniteTimeRange) {
+      TimeRange(0, Long.MaxValue)
+    } else {
+      // 3a. Get the start and end time is ms based on the lookback, offset and the user provided start and end time
+      val (maxOffsetMs, minOffsetMs) = LogicalPlanUtils.getOffsetMillis(logicalPlan)
+        .foldLeft((Long.MinValue, Long.MaxValue)) {
+          case ((accMax, accMin), currValue) => (accMax.max(currValue), accMin.min(currValue))
+        }
+
+      val periodicSeriesTimeWithOffset = TimeRange((queryParams.startSecs * 1000) - maxOffsetMs,
+        (queryParams.endSecs * 1000) - minOffsetMs)
+      val lookBackMs = getLookBackMillis(logicalPlan).max
+
+      //3b Get the Query time range based on user provided range, offsets in previous steps and lookback
+      TimeRange(periodicSeriesTimeWithOffset.startMs - lookBackMs,
+        periodicSeriesTimeWithOffset.endMs)
+    }
+
+    //4. Based on the map in 2 and time range in 5, get the partitions to query
+    routingKeyMap.flatMap(metricMap =>
+      partitionLocationProvider.getPartitions(metricMap, queryTimeRange))
+  }
+  // scalastyle:on method.length
+
+  /**
+   * Checks if all the PartitionAssignments belong to same partition
+   */
+  protected def isSinglePartition(partitions: Seq[PartitionAssignment]) : Boolean = {
+    if (partitions.isEmpty)
+      true
+    else {
+      val partName = partitions.head.partitionName
+      partitions.forall(_.partitionName.equals(partName))
+    }
+  }
+
+  protected def canSupportMultiPartitionCalls(execPlans: Seq[ExecPlan]): Boolean =
+    execPlans.forall {
+      case _: PromQlRemoteExec => false
+      case _ => true
+    }
+}