Fix filters on TIMESTAMP_TO_MILLIS over TIME_FLOOR. (#17871)

* Fix filters on TIMESTAMP_TO_MILLIS over TIME_FLOOR.

It is possible for toQueryGranularity to return nonnull for BIGINT rhs,
in the case where the time floor expression is wrapped in
TIMESTAMP_TO_MILLIS. Prior to this patch, that would throw an (incorrect)
error about expecting a literal.

This patch fixes the defensive error to have the correct message
(in Calcites) and also fixes the filter translation code to not hit that
defensive check (in Expressions).

* Fix error message.

Author: gianm

sql/src/main/java/org/apache/druid/sql/calcite/expression/Expressions.java

@@ -648,12 +648,19 @@ private static DimFilter toSimpleLeafFilter(
         return null;
       }
 
-      // Special handling for filters on FLOOR(__time TO granularity).
+      // Special handling for filters like FLOOR(__time TO granularity).
       final Granularity queryGranularity =
           toQueryGranularity(lhsExpression, plannerContext.getExpressionParser());
-      if (queryGranularity != null) {
-        // lhs is FLOOR(__time TO granularity); rhs must be a timestamp
-        final long rhsMillis = Calcites.calciteDateTimeLiteralToJoda(rhs, plannerContext.getTimeZone()).getMillis();
+      if (queryGranularity != null && !RexLiteral.isNullLiteral(rhs)) {
+        // lhs is a time-floor expression; rhs must be a timestamp or millis
+        final long rhsMillis;
+
+        if (rhs.getType().getSqlTypeName() == SqlTypeName.BIGINT) {
+          rhsMillis = ((Number) RexLiteral.value(rhs)).longValue();
+        } else {
+          rhsMillis = Calcites.calciteDateTimeLiteralToJoda(rhs, plannerContext.getTimeZone()).getMillis();
+        }
+
         return buildTimeFloorFilter(
             ColumnHolder.TIME_COLUMN_NAME,
             queryGranularity,

sql/src/main/java/org/apache/druid/sql/calcite/planner/Calcites.java

@@ -40,8 +40,8 @@
 import org.apache.calcite.util.DateString;
 import org.apache.calcite.util.TimeString;
 import org.apache.calcite.util.TimestampString;
+import org.apache.druid.error.DruidException;
 import org.apache.druid.java.util.common.DateTimes;
-import org.apache.druid.java.util.common.IAE;
 import org.apache.druid.java.util.common.StringUtils;
 import org.apache.druid.math.expr.ExpressionProcessing;
 import org.apache.druid.math.expr.ExpressionProcessingConfig;
@@ -449,8 +449,8 @@ public static DateString jodaToCalciteDateString(final DateTime dateTime, final
   public static DateTime calciteDateTimeLiteralToJoda(final RexNode literal, final DateTimeZone timeZone)
   {
     final SqlTypeName typeName = literal.getType().getSqlTypeName();
-    if (literal.getKind() != SqlKind.LITERAL || (typeName != SqlTypeName.TIMESTAMP && typeName != SqlTypeName.DATE)) {
-      throw new IAE("Expected literal but got[%s]", literal.getKind());
+    if (literal.getKind() != SqlKind.LITERAL) {
+      throw DruidException.defensive("Expected literal but got[%s]", literal.getKind());
     }
 
     if (typeName == SqlTypeName.TIMESTAMP) {
@@ -460,7 +460,7 @@ public static DateTime calciteDateTimeLiteralToJoda(final RexNode literal, final
       final DateString dateString = (DateString) RexLiteral.value(literal);
       return CALCITE_DATE_PARSER.parse(dateString.toString()).withZoneRetainFields(timeZone);
     } else {
-      throw new IAE("Expected TIMESTAMP or DATE but got[%s]", typeName);
+      throw DruidException.defensive("Expected TIMESTAMP or DATE but got[%s]", typeName);
     }
   }
 

sql/src/test/java/org/apache/druid/sql/calcite/CalciteQueryTest.java

@@ -6106,6 +6106,28 @@ public void testCountStarWithFloorTimeFilter()
     );
   }
 
+  @Test
+  public void testCountStarWithFloorTimeFilterUsingMilliseconds()
+  {
+    testQuery(
+        "SELECT COUNT(*) FROM druid.foo "
+        + "WHERE TIMESTAMP_TO_MILLIS(FLOOR(__time TO DAY)) >= 946684800000 AND "
+        + "TIMESTAMP_TO_MILLIS(FLOOR(__time TO DAY)) < 978307200000",
+        ImmutableList.of(
+            Druids.newTimeseriesQueryBuilder()
+                  .dataSource(CalciteTests.DATASOURCE1)
+                  .intervals(querySegmentSpec(Intervals.of("2000-01-01/2001-01-01")))
+                  .granularity(Granularities.ALL)
+                  .aggregators(aggregators(new CountAggregatorFactory("a0")))
+                  .context(QUERY_CONTEXT_DEFAULT)
+                  .build()
+        ),
+        ImmutableList.of(
+            new Object[]{3L}
+        )
+    );
+  }
+
   @Test
   public void testCountStarWithMisalignedFloorTimeFilter()
   {