Addressing PR comments

filodb · Apr 25, 2024 · d3b1391 · d3b1391
1 parent 0a8e85b
commit d3b1391
Show file tree

Hide file tree

Showing 2 changed files with 43 additions and 27 deletions.
diff --git a/coordinator/src/test/scala/filodb.coordinator/queryplanner/LogicalPlanUtilsSpec.scala b/coordinator/src/test/scala/filodb.coordinator/queryplanner/LogicalPlanUtilsSpec.scala
@@ -6,7 +6,7 @@ import filodb.core.query.ColumnFilter
 import filodb.core.query.Filter.Equals
 import filodb.prometheus.ast.TimeStepParams
 import filodb.prometheus.parse.Parser
-import filodb.query.{LogicalPlan, ScalarFixedDoublePlan}
+import filodb.query.LogicalPlan
 import org.scalatest.funspec.AnyFunSpec
 import org.scalatest.matchers.should.Matchers
 
@@ -132,14 +132,27 @@ class LogicalPlanUtilsSpec   extends AnyFunSpec with Matchers {
   it ("getColumnFilterGroupWithLogical plans should return result as expected") {
     val timeParamsSec = TimeStepParams(1000, 10, 10000)
     val query1 = """sum(count_over_time((test_metric{_ws_="test-ws", _ns_="test-ns", usecase="test"} > 20000)[21600s:])) or vector(0)"""
-    val lp = Parser.queryRangeToLogicalPlan(query1, timeParamsSec)
-    val columnGroupWithPlan = LogicalPlan.getColumnFilterGroupWithLogicalPlan(lp)
-    val scalarPlansCount = columnGroupWithPlan.count(x => x.logicalPlan == ScalarFixedDoublePlan)
-    scalarPlansCount shouldEqual 2
-    val query2 = """rate(test_metric{_ns_="test-ns", _ws_="test-ns", cluster="test1"}[5m]) * 1000"""
-    val lp2 = Parser.queryRangeToLogicalPlan(query2, timeParamsSec)
-    val columnGroupWithPlan2 = LogicalPlan.getColumnFilterGroupWithLogicalPlan(lp2)
-    val scalarPlansCount2 = columnGroupWithPlan2.count(x => x.logicalPlan == ScalarFixedDoublePlan)
-    scalarPlansCount2 shouldEqual 1
+    val query2 = """sum(count_over_time((test_metric{_ws_="test-ws", _ns_="test-ns", usecase="test"} > 20000)[21600s:])) or vector(0) or sum(count_over_time((test_metric{_ws_="wrong_ws", _ns_="wrong-ns", usecase="test"} > 20000)[21600s:]))"""
+    val query3 = """rate(test_metric{_ns_="test-ns", _ws_="test-ns", cluster="test1"}[5m]) * 1000"""
+
+    def getColumnFilterWithAndWithoutScalarCount(query: String) : (Int, Int) = {
+      val lp = Parser.queryRangeToLogicalPlan(query, timeParamsSec)
+      val columnGroups = LogicalPlan.getColumnFilterGroup(lp)
+      val columnGroupsWithoutScalar = LogicalPlan.getColumnFilterGroupFilteringLeafScalarPlans(lp)
+      val scalarPlansCount = columnGroups.count(x => x.isEmpty)
+      val scalarPlansWithoutCount = columnGroupsWithoutScalar.count(x => x.isEmpty)
+      (scalarPlansCount, scalarPlansWithoutCount)
+    }
+    var counts = getColumnFilterWithAndWithoutScalarCount(query1)
+    counts._1 shouldEqual 2
+    counts._2 shouldEqual 0
+
+    counts = getColumnFilterWithAndWithoutScalarCount(query2)
+    counts._1 shouldEqual 3
+    counts._2 shouldEqual 0
+
+    counts = getColumnFilterWithAndWithoutScalarCount(query3)
+    counts._1 shouldEqual 1
+    counts._2 shouldEqual 0
   }
 }
diff --git a/query/src/main/scala/filodb/query/LogicalPlan.scala b/query/src/main/scala/filodb/query/LogicalPlan.scala
@@ -671,8 +671,6 @@ case class ApplyLimitFunction(vectors: PeriodicSeriesPlan,
               vectors = vectors.replacePeriodicSeriesFilters(filters))
 }
 
-case class ColumnFilterListWithLogicalPlan(columnFilterSet: Set[ColumnFilter], logicalPlan: Any)
-
 object LogicalPlan {
   /**
     * Get leaf Logical Plans
@@ -764,30 +762,35 @@ object LogicalPlan {
   }
 
   /**
-   * Given a LogicalPlan, the function finds a Seq of all Child nodes, and returns a Set of ColumnFilters and the
-   * associated logical plan for each of the Leaf node
+   * Given a LogicalPlan, the function finds a Seq of all Child nodes,and returns a Set of ColumnFilters.
+   * It filters out the column filter group for all the scalar leaf plans
    *
    * @param logicalPlan the root LogicalPlan
    * @return Seq[ColumnFilterListWithLogicalPlan], Seq has size same as the number of leaf nodes
    */
-  def getColumnFilterGroupWithLogicalPlan(logicalPlan: LogicalPlan): Seq[ColumnFilterListWithLogicalPlan] = {
+  def getColumnFilterGroupFilteringLeafScalarPlans(logicalPlan: LogicalPlan): Seq[Set[ColumnFilter]] = {
     LogicalPlan.findLeafLogicalPlans(logicalPlan) map { lp =>
       lp match {
-        case lp: LabelValues => ColumnFilterListWithLogicalPlan(lp.filters toSet, LabelValues)
-        case lp: LabelNames => ColumnFilterListWithLogicalPlan(lp.filters toSet, LabelNames)
-        case lp: RawSeries => ColumnFilterListWithLogicalPlan(lp.filters toSet, RawSeries)
-        case lp: RawChunkMeta => ColumnFilterListWithLogicalPlan(lp.filters toSet, RawChunkMeta)
-        case lp: SeriesKeysByFilters => ColumnFilterListWithLogicalPlan(lp.filters toSet, SeriesKeysByFilters)
-        case lp: LabelCardinality => ColumnFilterListWithLogicalPlan(lp.filters.toSet, LabelCardinality)
-        case lp: TsCardinalities => ColumnFilterListWithLogicalPlan(lp.filters.toSet, TsCardinalities)
-        case _: ScalarTimeBasedPlan => ColumnFilterListWithLogicalPlan(Set.empty[ColumnFilter], ScalarTimeBasedPlan)
-        case _: ScalarFixedDoublePlan => ColumnFilterListWithLogicalPlan(Set.empty[ColumnFilter], ScalarFixedDoublePlan)
-        case _: ScalarBinaryOperation => ColumnFilterListWithLogicalPlan(Set.empty[ColumnFilter], ScalarBinaryOperation)
+        case lp: LabelValues => (lp.filters toSet, true)
+        case lp: LabelNames => (lp.filters toSet, true)
+        case lp: RawSeries => (lp.filters toSet, true)
+        case lp: RawChunkMeta => (lp.filters toSet, true)
+        case lp: SeriesKeysByFilters => (lp.filters toSet, true)
+        case lp: LabelCardinality => (lp.filters.toSet, true)
+        case lp: TsCardinalities => (lp.filters.toSet, true)
+        case _: ScalarTimeBasedPlan => (Set.empty[ColumnFilter], false)
+        case _: ScalarFixedDoublePlan => (Set.empty[ColumnFilter], false)
+        case _: ScalarBinaryOperation => (Set.empty[ColumnFilter], false)
         case _ => throw new BadQueryException(s"Invalid logical plan $logicalPlan")
       }
     } match {
-      case groupSeq: Seq[ColumnFilterListWithLogicalPlan] =>
-        if (groupSeq.isEmpty || groupSeq.forall( x => x.columnFilterSet.isEmpty)) Seq.empty else groupSeq
+      case groupSeq: Seq[(Set[ColumnFilter], Boolean)] =>
+        if (groupSeq.isEmpty || groupSeq.forall( x => x._1.isEmpty)) Seq.empty else {
+          // filter out the scalar plans with empty column filter group
+          // and only include the column filter of raw leaf plans
+          groupSeq.filter(y => y._2)
+            .map(x => x._1)
+        }
       case _ => Seq.empty
     }
   }