Fixing aliased columns bug in GAPFILL based queries (apache#15052)

xiangfu0 · Feb 18, 2025 · dd37e64 · dd37e64
1 parent d76fe22
commit dd37e64
Show file tree

Hide file tree

Showing 3 changed files with 53 additions and 3 deletions.
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/BaseGapfillProcessor.java b/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/BaseGapfillProcessor.java
@@ -146,6 +146,7 @@ protected void replaceColumnNameWithAlias(DataSchema dataSchema) {
    */
   public void process(BrokerResponseNative brokerResponseNative) {
     DataSchema dataSchema = brokerResponseNative.getResultTable().getDataSchema();
+    replaceColumnNameWithAlias(dataSchema);
     DataSchema resultTableSchema = getResultTableDataSchema(dataSchema);
     if (brokerResponseNative.getResultTable().getRows().isEmpty()) {
       brokerResponseNative.setResultTable(new ResultTable(resultTableSchema, Collections.emptyList()));
@@ -174,7 +175,6 @@ public void process(BrokerResponseNative brokerResponseNative) {
     }
 
     List<Object[]> rows = brokerResponseNative.getResultTable().getRows();
-    replaceColumnNameWithAlias(dataSchema);
     List<Object[]> resultRows = gapFillAndAggregate(rows, dataSchema, resultTableSchema);
     brokerResponseNative.setResultTable(new ResultTable(resultTableSchema, resultRows));
   }

diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/GapfillProcessor.java b/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/GapfillProcessor.java
@@ -69,6 +69,7 @@ public class GapfillProcessor extends BaseGapfillProcessor {
    */
   public void process(BrokerResponseNative brokerResponseNative) {
     DataSchema dataSchema = brokerResponseNative.getResultTable().getDataSchema();
+    replaceColumnNameWithAlias(dataSchema);
     DataSchema resultTableSchema = getResultTableDataSchema(dataSchema);
     if (brokerResponseNative.getResultTable().getRows().isEmpty()) {
       brokerResponseNative.setResultTable(new ResultTable(resultTableSchema, Collections.emptyList()));
@@ -98,8 +99,6 @@ public void process(BrokerResponseNative brokerResponseNative) {
 
     List<Object[]>[] timeBucketedRawRows = putRawRowsIntoTimeBucket(brokerResponseNative.getResultTable().getRows());
 
-    replaceColumnNameWithAlias(dataSchema);
-
     if (_queryContext.getAggregationFunctions() == null) {
       Map<String, Integer> sourceColumnsIndexes = new HashMap<>();
       for (int i = 0; i < dataSchema.getColumnNames().length; i++) {

diff --git a/pinot-core/src/test/java/org/apache/pinot/queries/GapfillQueriesTest.java b/pinot-core/src/test/java/org/apache/pinot/queries/GapfillQueriesTest.java
@@ -783,6 +783,57 @@ public void datetimeconvertGapfillTestAggregateAggregateWithHavingClause() {
     }
   }
 
+  @Test
+  public void datetimeconvertGapfillTestAliasedColumnsInTimeserieson() {
+    String gapfillQuery1 = "SELECT "
+        + "time_col, alias_levelId, SUM(occupied) as occupied_slots_count, time_col "
+        + "FROM ("
+        + "  SELECT GapFill(time_col, "
+        + "    '1:MILLISECONDS:SIMPLE_DATE_FORMAT:yyyy-MM-dd HH:mm:ss.SSS', "
+        + "    '2021-11-07 4:00:00.000',  '2021-11-07 12:00:00.000', '1:HOURS',"
+        + "     FILL(occupied, 'FILL_PREVIOUS_VALUE'), TIMESERIESON(alias_levelId, alias_lotId)),"
+        + "     occupied, alias_lotId, alias_levelId"
+        + "  FROM ("
+        + "    SELECT DATETIMECONVERT(eventTime, '1:MILLISECONDS:EPOCH', "
+        + "      '1:MILLISECONDS:SIMPLE_DATE_FORMAT:yyyy-MM-dd HH:mm:ss.SSS', '1:HOURS') AS time_col,"
+        + "       lastWithTime(isOccupied, eventTime, 'INT') as occupied, "
+        + "       cast(lotId as varchar) as alias_lotId, "
+        + "       cast(levelId as varchar) as alias_levelId "
+        + "    FROM parkingData "
+        + "    WHERE eventTime >= 1636257600000 AND eventTime <= 1636286400000 "
+        + "    GROUP BY time_col, alias_levelId, alias_lotId "
+        + "    LIMIT 200 "
+        + "  ) "
+        + "  LIMIT 200 "
+        + ") "
+        + " GROUP BY time_col, alias_levelId "
+        + " HAVING occupied_slots_count > 0"
+        + " LIMIT 200 ";
+
+    BrokerResponseNative gapfillBrokerResponse1 = getBrokerResponse(gapfillQuery1);
+
+    double[] expectedOccupiedSlotsCounts1 = new double[]{1, 2, 3, 4, 3, 2, 1};
+    ResultTable gapFillResultTable1 = gapfillBrokerResponse1.getResultTable();
+    List<Object[]> gapFillRows1 = gapFillResultTable1.getRows();
+    Assert.assertEquals(gapFillRows1.size(), expectedOccupiedSlotsCounts1.length * 2);
+    DateTimeFormatSpec dateTimeFormatter =
+        new DateTimeFormatSpec("1:MILLISECONDS:SIMPLE_DATE_FORMAT:yyyy-MM-dd HH:mm:ss.SSS");
+    DateTimeGranularitySpec dateTimeGranularity = new DateTimeGranularitySpec("1:HOURS");
+
+    long start = dateTimeFormatter.fromFormatToMillis("2021-11-07 04:00:00.000");
+    for (int i = 0; i < expectedOccupiedSlotsCounts1.length; i += 2) {
+      String firstTimeCol = (String) gapFillRows1.get(i)[0];
+      long timeStamp = dateTimeFormatter.fromFormatToMillis(firstTimeCol);
+      Assert.assertEquals(timeStamp, start);
+      Assert.assertEquals(expectedOccupiedSlotsCounts1[i / 2], gapFillRows1.get(i)[2]);
+      firstTimeCol = (String) gapFillRows1.get(i + 1)[0];
+      timeStamp = dateTimeFormatter.fromFormatToMillis(firstTimeCol);
+      Assert.assertEquals(timeStamp, start);
+      Assert.assertEquals(expectedOccupiedSlotsCounts1[i / 2], gapFillRows1.get(i)[2]);
+      start += dateTimeGranularity.granularityToMillis();
+    }
+  }
+
   @Test
   public void toEpochHoursGapfillTestSelectSelect() {
     DateTimeFormatSpec dateTimeFormatter = new DateTimeFormatSpec("1:HOURS:EPOCH");