From a914940b378ae5a6b7b489c570d5c0ba73100730 Mon Sep 17 00:00:00 2001 From: Yash Mayya Date: Wed, 20 Nov 2024 15:03:38 +0700 Subject: [PATCH] Fix timestamp literal handling in the multi-stage query engine (#14502) --- .../query/planner/logical/RexExpressionUtils.java | 10 +++++++++- .../apache/pinot/query/QueryEnvironmentTestBase.java | 3 +++ 2 files changed, 12 insertions(+), 1 deletion(-) diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/RexExpressionUtils.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/RexExpressionUtils.java index bf92e28f8e2b..a1b707585057 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/RexExpressionUtils.java +++ b/pinot-query-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 index 3b0b811928af..685252bbbe43 100644 --- 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)"} }; }