From b7455b0619d7562305eddd267805832aa89e93e0 Mon Sep 17 00:00:00 2001
From: sandeep6189 <sandeep6189@gmail.com>
Date: Fri, 8 Nov 2024 20:22:23 -0800
Subject: [PATCH] misc(query): Adding helper functions for hqe (#1883)

---
 .../queryplanner/LogicalPlanUtils.scala       | 16 ++++++++++++
 .../queryplanner/LogicalPlanUtilsSpec.scala   | 26 +++++++++++++++++++
 2 files changed, 42 insertions(+)

diff --git a/coordinator/src/main/scala/filodb.coordinator/queryplanner/LogicalPlanUtils.scala b/coordinator/src/main/scala/filodb.coordinator/queryplanner/LogicalPlanUtils.scala
index d5ed505bf8..5a8f690f41 100644
--- 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
index f15da53b66..2c75054084 100644
--- 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
+  }
 }