From e0cae0f24600c403accf130c059bfdf67e12a618 Mon Sep 17 00:00:00 2001
From: Parag Jain <pjain1@apache.org>
Date: Sun, 1 Sep 2024 00:25:55 +0530
Subject: [PATCH] fix extraction of timeseries results from result level cache
 (#16895)

* fix extraction of timeseries results from result level cache

* remove unneded import

* add test
---
 .../TimeseriesQueryQueryToolChest.java        |  7 ++++
 .../TimeseriesQueryQueryToolChestTest.java    |  4 +-
 .../sql/calcite/CalciteSelectQueryTest.java   | 41 +++++++++++++++++++
 3 files changed, 51 insertions(+), 1 deletion(-)

diff --git a/processing/src/main/java/org/apache/druid/query/timeseries/TimeseriesQueryQueryToolChest.java b/processing/src/main/java/org/apache/druid/query/timeseries/TimeseriesQueryQueryToolChest.java
index 05b25d05c219..18bcc713dece 100644
--- a/processing/src/main/java/org/apache/druid/query/timeseries/TimeseriesQueryQueryToolChest.java
+++ b/processing/src/main/java/org/apache/druid/query/timeseries/TimeseriesQueryQueryToolChest.java
@@ -383,6 +383,13 @@ public Result<TimeseriesResultValue> apply(Object input)
               timestamp = granularity.toDateTime(Preconditions.checkNotNull(timestampNumber, "timestamp").longValue());
             }
 
+            // If "timestampResultField" is set, we must include a copy of the timestamp in the result.
+            // This is used by the SQL layer when it generates a Timeseries query for a group-by-time-floor SQL query.
+            // The SQL layer expects the result of the time-floor to have a specific name that is not going to be "__time".
+            if (StringUtils.isNotEmpty(query.getTimestampResultField()) && timestamp != null) {
+              retVal.put(query.getTimestampResultField(), timestamp.getMillis());
+            }
+
             CacheStrategy.fetchAggregatorsFromCache(
                 aggs,
                 resultIter,
diff --git a/processing/src/test/java/org/apache/druid/query/timeseries/TimeseriesQueryQueryToolChestTest.java b/processing/src/test/java/org/apache/druid/query/timeseries/TimeseriesQueryQueryToolChestTest.java
index f1afabe9c75e..8d9166d27467 100644
--- a/processing/src/test/java/org/apache/druid/query/timeseries/TimeseriesQueryQueryToolChestTest.java
+++ b/processing/src/test/java/org/apache/druid/query/timeseries/TimeseriesQueryQueryToolChestTest.java
@@ -97,7 +97,7 @@ public void testCacheStrategy() throws Exception
                 ),
                 ImmutableList.of(new ConstantPostAggregator("post", 10)),
                 0,
-                null
+                ImmutableMap.of(TimeseriesQuery.CTX_TIMESTAMP_RESULT_FIELD, "ts_field")
             )
         );
 
@@ -106,6 +106,7 @@ public void testCacheStrategy() throws Exception
         DateTimes.utc(123L),
         new TimeseriesResultValue(
             ImmutableMap.of(
+                "ts_field", 123L,
                 "metric1", 2,
                 "metric0", 3,
                 "complexMetric", new SerializablePairLongString(123L, "val1")
@@ -130,6 +131,7 @@ public void testCacheStrategy() throws Exception
         DateTimes.utc(123L),
         new TimeseriesResultValue(
             ImmutableMap.of(
+                "ts_field", 123L,
                 "metric1", 2,
                 "metric0", 3,
                 "complexMetric", "val1",
diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/CalciteSelectQueryTest.java b/sql/src/test/java/org/apache/druid/sql/calcite/CalciteSelectQueryTest.java
index 6e06c5d4b5d2..973b1eb4df1e 100644
--- a/sql/src/test/java/org/apache/druid/sql/calcite/CalciteSelectQueryTest.java
+++ b/sql/src/test/java/org/apache/druid/sql/calcite/CalciteSelectQueryTest.java
@@ -2118,6 +2118,47 @@ public void testCacheKeyConsistency()
         .run();
   }
 
+  @SqlTestFrameworkConfig.ResultCache(ResultCacheMode.ENABLED)
+  @Test
+  public void testTimseriesResultCachePullConsistency()
+  {
+    skipVectorize();
+
+    String query = "SELECT \n"
+                   + "  (t1.\"__time\") AS \"__time\", \n"
+                   + "  (ANY_VALUE(t1.\"added\")) AS \"added\" \n"
+                   + "FROM \n"
+                   + "  (\n"
+                   + "    SELECT \n"
+                   + "      (time_floor(\"__time\", 'PT1H')) AS \"__time\", \n"
+                   + "      (SUM(added)) AS \"added\" \n"
+                   + "    FROM \"wikipedia\" \n"
+                   + "    GROUP BY 1 \n"
+                   + "    ORDER BY \"__time\" \n"
+                   + "    LIMIT 1\n"
+                   + "  ) t1 \n"
+                   + "GROUP BY 1 \n"
+                   + "ORDER BY \"__time\"";
+
+    testBuilder()
+        .sql(query)
+        .expectedResults(
+            ImmutableList.of(
+                new Object[]{1442016000000L, 32251L}
+            )
+        )
+        .run();
+
+    testBuilder()
+        .sql(query)
+        .expectedResults(
+            ImmutableList.of(
+                new Object[]{1442016000000L, 32251L}
+            )
+        )
+        .run();
+  }
+
   @Test
   public void testSqlToRelInConversion()
   {