Fix timestamp literal handling in the multi-stage query engine (#14502)

apache · Nov 20, 2024 · a914940 · a914940
1 parent 89622c0
commit a914940
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 1 deletion.
diff --git a/...uery-planner/src/main/java/org/apache/pinot/query/planner/logical/RexExpressionUtils.java b/...uery-planner/src/main/java/org/apache/pinot/query/planner/logical/RexExpressionUtils.java
@@ -255,7 +255,15 @@ private static RexExpression.Literal fromRexLiteralValue(ColumnDataType dataType
         value = Boolean.TRUE.equals(value) ? BooleanUtils.INTERNAL_TRUE : BooleanUtils.INTERNAL_FALSE;
         break;
       case TIMESTAMP:
-        value = ((Calendar) value).getTimeInMillis();
+        if (value instanceof Calendar) {
+          value = ((Calendar) value).getTimeInMillis();
+        } else if (value instanceof TimestampString) {
+          value = ((TimestampString) value).getMillisSinceEpoch();
+        } else if (value instanceof Long) {
+          // Already in millis
+        } else {
+          throw new IllegalStateException("Unsupported value type for TIMESTAMP: " + value.getClass().getName());
+        }
         break;
       case STRING:
         value = ((NlsString) value).getValue();

diff --git a/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryEnvironmentTestBase.java b/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryEnvironmentTestBase.java
@@ -250,6 +250,9 @@ protected Object[][] provideQueries() {
         new Object[]{"SELECT JSON_EXTRACT_SCALAR(col1, '$.foo', 'FLOAT_ARRAY') FROM a"},
         new Object[]{"SELECT JSON_EXTRACT_SCALAR(col1, '$.foo', 'DOUBLE_ARRAY') FROM a"},
         new Object[]{"SELECT JSON_EXTRACT_SCALAR(col1, '$.foo', 'STRING_ARRAY') FROM a"},
+        new Object[]{"SELECT ts_timestamp FROM a WHERE ts_timestamp BETWEEN TIMESTAMP '2016-01-01 00:00:00' AND "
+            + "TIMESTAMP '2016-01-01 10:00:00'"},
+        new Object[]{"SELECT ts_timestamp FROM a WHERE ts_timestamp >= CAST(1454284798000 AS TIMESTAMP)"}
     };
   }