misc(query): Adding helper functions for hqe (#1883)

filodb · Nov 9, 2024 · b7455b0 · b7455b0
1 parent 29506e6
commit b7455b0
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 0 deletions.
diff --git a/coordinator/src/main/scala/filodb.coordinator/queryplanner/LogicalPlanUtils.scala b/coordinator/src/main/scala/filodb.coordinator/queryplanner/LogicalPlanUtils.scala
@@ -247,6 +247,14 @@ object LogicalPlanUtils extends StrictLogging {
     }
   }
 
+  /**
+   * @param logicalPlan LogicalPlan
+   * @return maximum offset in milliseconds
+   */
+  def getMaxOffetInMillis(logicalPlan: LogicalPlan): Long = {
+    getOffsetMillis(logicalPlan).max
+  }
+
   def getLookBackMillis(logicalPlan: LogicalPlan): Seq[Long] = {
     // getLookBackMillis is used primarily for LongTimeRangePlanner,
     // SubqueryWtihWindowing returns lookback but TopLevelSubquery does not. The conceptual meaning of the lookback
@@ -260,6 +268,14 @@ object LogicalPlanUtils extends StrictLogging {
     }
   }
 
+  /**
+   * @param logicalPlan LogicalPlan
+   * @return maximum lookback in milliseconds
+   */
+  def getMaxLookBackInMillis(logicalPlan: LogicalPlan): Long = {
+    getLookBackMillis(logicalPlan).max
+  }
+
   def getMetricName(logicalPlan: LogicalPlan, datasetMetricColumn: String): Set[String] = {
     val columnFilterGroup = LogicalPlan.getColumnFilterGroup(logicalPlan)
     val metricName = LogicalPlan.getColumnValues(columnFilterGroup, PromMetricLabel)

diff --git a/coordinator/src/test/scala/filodb.coordinator/queryplanner/LogicalPlanUtilsSpec.scala b/coordinator/src/test/scala/filodb.coordinator/queryplanner/LogicalPlanUtilsSpec.scala
@@ -155,4 +155,30 @@ class LogicalPlanUtilsSpec   extends AnyFunSpec with Matchers {
     counts._1 shouldEqual 1
     counts._2 shouldEqual 0
   }
+
+  it ("getMaxOffetInMillis should return result as expected") {
+    val timeParamsSec = TimeStepParams(1000, 10, 10000)
+    val query1 = """rate(test_metric{_ns_="test-ns", _ws_="test-ns", cluster="test1"}[5m]) * 1000"""
+    val query2 = """rate(test_metric{_ns_="test-ns", _ws_="test-ns", cluster="test1"}[5m] offset 1h) + rate(test_metric{_ns_="test-ns", _ws_="test-ns", cluster="test1"}[5m] offset 2h)"""
+
+    def getMaxOffsetInMillis(query: String) : Long = {
+      val lp = Parser.queryRangeToLogicalPlan(query, timeParamsSec)
+      LogicalPlanUtils.getMaxOffetInMillis(lp)
+    }
+    getMaxOffsetInMillis(query1) shouldEqual 0
+    getMaxOffsetInMillis(query2) shouldEqual 7200000
+  }
+
+  it ("getMaxLookbackInMillis should return result as expected") {
+    val timeParamsSec = TimeStepParams(1000, 10, 10000)
+    val query1 = """rate(test_metric{_ns_="test-ns", _ws_="test-ns", cluster="test1"}[15m]) + rate(test_metric{_ns_="test-ns", _ws_="test-ns", cluster="test1"}[10m])"""
+    val query2 = """test_metric{_ns_="test-ns", _ws_="test-ns", cluster="test1"} offset 1d * 1000"""
+
+    def getMaxLookbackInMillis(query: String) : Long = {
+      val lp = Parser.queryRangeToLogicalPlan(query, timeParamsSec)
+      LogicalPlanUtils.getMaxLookBackInMillis(lp)
+    }
+    getMaxLookbackInMillis(query1) shouldEqual 900000 // max of 15 and 10m
+    getMaxLookbackInMillis(query2) shouldEqual 300000 // default lookback is 5m
+  }
 }