fix extraction of timeseries results from result level cache (apache#…

…16895) * fix extraction of timeseries results from result level cache * remove unneded import * add test
edgar2020 · Sep 5, 2024 · e0cae0f · e0cae0f
1 parent cbb45ff
commit e0cae0f
Show file tree

Hide file tree

Showing 3 changed files with 51 additions and 1 deletion.
diff --git a/...essing/src/main/java/org/apache/druid/query/timeseries/TimeseriesQueryQueryToolChest.java b/...essing/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/...ng/src/test/java/org/apache/druid/query/timeseries/TimeseriesQueryQueryToolChestTest.java b/...ng/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
@@ -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()
   {