Support timestamp cast literals in table-model time planning

integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBComplexQueryIT.java

@@ -19,19 +19,29 @@
 
 package org.apache.iotdb.relational.it.query.recent;
 
+import org.apache.iotdb.isession.SessionConfig;
 import org.apache.iotdb.it.env.EnvFactory;
 import org.apache.iotdb.it.framework.IoTDBTestRunner;
 import org.apache.iotdb.itbase.category.TableClusterIT;
 import org.apache.iotdb.itbase.category.TableLocalStandaloneIT;
+import org.apache.iotdb.itbase.env.BaseEnv;
 
+import com.google.gson.JsonElement;
+import com.google.gson.JsonObject;
+import com.google.gson.JsonParser;
 import org.junit.AfterClass;
 import org.junit.BeforeClass;
 import org.junit.Test;
 import org.junit.experimental.categories.Category;
 import org.junit.runner.RunWith;
 
+import java.sql.Connection;
+import java.sql.ResultSet;
+import java.sql.Statement;
+
 import static org.apache.iotdb.db.it.utils.TestUtils.prepareTableData;
 import static org.apache.iotdb.db.it.utils.TestUtils.tableResultSetEqualTest;
+import static org.junit.Assert.assertTrue;
 
 @RunWith(IoTDBTestRunner.class)
 @Category({TableLocalStandaloneIT.class, TableClusterIT.class})
@@ -91,6 +101,105 @@ public void testTableLessQueryWithOrderBy() {
         DATABASE_NAME);
   }
 
+  @Test
+  public void testTimeFilterWithTimestampCast() throws Exception {
+    String[] expectedHeader = new String[] {"time", "name"};
+    String rangeQuery =
+        "SELECT time, name FROM employees"
+            + " WHERE time >= CAST(2 AS TIMESTAMP) AND time < CAST(5 AS TIMESTAMP)"
+            + " ORDER BY time";
+    tableResultSetEqualTest(
+        rangeQuery,
+        expectedHeader,
+        new String[] {
+          "1970-01-01T00:00:00.002Z,John,",
+          "1970-01-01T00:00:00.003Z,Nancy,",
+          "1970-01-01T00:00:00.004Z,Jack,"
+        },
+        DATABASE_NAME);
+    // Time predicates must reach the table scan so partition pruning remains effective.
+    assertHasTimePredicate(rangeQuery);
+
+    String betweenQuery =
+        "SELECT time, name FROM employees"
+            + " WHERE time BETWEEN CAST(2 AS TIMESTAMP) AND CAST(4 AS TIMESTAMP)"
+            + " ORDER BY time";
+    tableResultSetEqualTest(
+        betweenQuery,
+        expectedHeader,
+        new String[] {
+          "1970-01-01T00:00:00.002Z,John,",
+          "1970-01-01T00:00:00.003Z,Nancy,",
+          "1970-01-01T00:00:00.004Z,Jack,"
+        },
+        DATABASE_NAME);
+    assertHasTimePredicate(betweenQuery);
+
+    tableResultSetEqualTest(
+        "SELECT time, name FROM employees"
+            + " WHERE time IN (CAST(2 AS TIMESTAMP), CAST(4 AS TIMESTAMP)) ORDER BY time",
+        expectedHeader,
+        new String[] {"1970-01-01T00:00:00.002Z,John,", "1970-01-01T00:00:00.004Z,Jack,"},
+        DATABASE_NAME);
+
+    String singleValueInQuery =
+        "SELECT time, name FROM employees WHERE time IN (CAST(2 AS TIMESTAMP)) ORDER BY time";
+    tableResultSetEqualTest(
+        singleValueInQuery,
+        expectedHeader,
+        new String[] {"1970-01-01T00:00:00.002Z,John,"},
+        DATABASE_NAME);
+    assertHasTimePredicate(singleValueInQuery);
+  }
+
+  private static void assertHasTimePredicate(String query) throws Exception {
+    // EXPLAIN coordinates and node IDs may differ by DataNode, so verify the plan on one node
+    // instead of comparing all cluster plans as regular query results.
+    try (Connection connection =
+            EnvFactory.getEnv()
+                .getConnection(
+                    EnvFactory.getEnv().getDataNodeWrapper(0),
+                    SessionConfig.DEFAULT_USER,
+                    SessionConfig.DEFAULT_PASSWORD,
+                    BaseEnv.TABLE_SQL_DIALECT);
+        Statement statement = connection.createStatement()) {
+      statement.execute("USE " + DATABASE_NAME);
+      try (ResultSet resultSet = statement.executeQuery("EXPLAIN (FORMAT JSON) " + query)) {
+        assertTrue(resultSet.next());
+        assertTrue(
+            "Expected DeviceTableScanNode to contain a TimePredicate",
+            containsTimePredicate(JsonParser.parseString(resultSet.getString(1))));
+      }
+    }
+  }
+
+  private static boolean containsTimePredicate(JsonElement element) {
+    if (element == null || element.isJsonNull()) {
+      return false;
+    }
+    if (element.isJsonObject()) {
+      JsonObject object = element.getAsJsonObject();
+      JsonElement properties = object.get("properties");
+      if (properties != null
+          && properties.isJsonObject()
+          && properties.getAsJsonObject().has("TimePredicate")) {
+        return true;
+      }
+      for (String key : object.keySet()) {
+        if (containsTimePredicate(object.get(key))) {
+          return true;
+        }
+      }
+    } else if (element.isJsonArray()) {
+      for (JsonElement child : element.getAsJsonArray()) {
+        if (containsTimePredicate(child)) {
+          return true;
+        }
+      }
+    }
+    return false;
+  }
+
   @Test
   public void testTableLessQuery() {
     String[] expectedHeader;

integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBGapFillTableIT.java

@@ -31,6 +31,8 @@
 import org.junit.experimental.categories.Category;
 import org.junit.runner.RunWith;
 
+import java.time.Instant;
+
 import static org.apache.iotdb.db.it.utils.TestUtils.prepareTableData;
 import static org.apache.iotdb.db.it.utils.TestUtils.tableAssertTestFail;
 import static org.apache.iotdb.db.it.utils.TestUtils.tableResultSetEqualTest;
@@ -264,6 +266,44 @@ public void normalGapFillTest() {
         DATABASE_NAME);
   }
 
+  @Test
+  public void gapFillWithTimestampCastBoundsTest() {
+    long startTime = Instant.parse("2024-09-24T04:00:00Z").toEpochMilli();
+    long endTime = Instant.parse("2024-09-24T12:00:00Z").toEpochMilli();
+    String[] expectedHeader = new String[] {"hour_time", "avg_s1"};
+    String[] retArray =
+        new String[] {
+          "2024-09-24T04:00:00.000Z,null,",
+          "2024-09-24T05:00:00.000Z,null,",
+          "2024-09-24T06:00:00.000Z,null,",
+          "2024-09-24T07:00:00.000Z,27.2,",
+          "2024-09-24T08:00:00.000Z,27.3,",
+          "2024-09-24T09:00:00.000Z,null,",
+          "2024-09-24T10:00:00.000Z,null,",
+          "2024-09-24T11:00:00.000Z,29.3,"
+        };
+    tableResultSetEqualTest(
+        "select date_bin_gapfill(1h, time) as hour_time, avg(s1) as avg_s1 from table1"
+            + " where time >= CAST("
+            + startTime
+            + " AS TIMESTAMP) and time < CAST("
+            + endTime
+            + " AS TIMESTAMP) and device_id = 'd1' group by 1 order by 1",
+        expectedHeader,
+        retArray,
+        DATABASE_NAME);
+    tableResultSetEqualTest(
+        "select date_bin_gapfill(1h, time) as hour_time, avg(s1) as avg_s1 from table1"
+            + " where time between CAST("
+            + startTime
+            + " AS TIMESTAMP) and CAST("
+            + (endTime - 1)
+            + " AS TIMESTAMP) and device_id = 'd1' group by 1 order by 1",
+        expectedHeader,
+        retArray,
+        DATABASE_NAME);
+  }
+
   @Test
   public void gapFillWithFillClauseTest() {
 

iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/plan/relational/planner/ir/GlobalTimePredicateExtractVisitor.java

@@ -27,6 +27,7 @@
 import org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.Expression;
 import org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.Extract;
 import org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.FunctionCall;
+import org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.GenericLiteral;
 import org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.IfExpression;
 import org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.InPredicate;
 import org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.IsNotNullPredicate;
@@ -286,13 +287,22 @@ private static boolean checkIsTimeFilter(
       Expression timeExpression, String timeColumnName, Expression valueExpression) {
     return (isTimeColumn(timeExpression, timeColumnName)
             || isExtractTimeColumn(timeExpression, timeColumnName))
-        && valueExpression instanceof LongLiteral;
+        && isTimeValueLiteral(valueExpression);
   }
 
   private static boolean checkBetweenConstantSatisfy(Expression e1, Expression e2) {
-    return e1 instanceof LongLiteral
-        && e2 instanceof LongLiteral
-        && ((LongLiteral) e1).getParsedValue() <= ((LongLiteral) e2).getParsedValue();
+    return isTimeValueLiteral(e1) && isTimeValueLiteral(e2) && getLongValue(e1) <= getLongValue(e2);
+  }
+
+  private static boolean isTimeValueLiteral(Expression expression) {
+    return expression instanceof LongLiteral || expression instanceof GenericLiteral;
+  }
+
+  private static long getLongValue(Expression expression) {
+    if (expression instanceof GenericLiteral) {
+      return Long.parseLong(((GenericLiteral) expression).getValue());
+    }
+    return ((LongLiteral) expression).getParsedValue();
   }
 
   public static class Context {

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/analyzer/predicate/ConvertPredicateToTimeFilterVisitor.java

@@ -23,6 +23,7 @@
 import org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.ComparisonExpression;
 import org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.Expression;
 import org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.Extract;
+import org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.GenericLiteral;
 import org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.IfExpression;
 import org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.InListExpression;
 import org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.InPredicate;
@@ -70,14 +71,14 @@ public Filter visitInPredicate(InPredicate node, Void context) {
     checkArgument(valueList instanceof InListExpression);
     List<Expression> values = ((InListExpression) valueList).getValues();
     for (Expression value : values) {
-      checkArgument(value instanceof LongLiteral);
+      checkArgument(isLongValue(value));
     }
     if (values.size() == 1) {
       return TimeFilterApi.eq(getLongValue(values.get(0)));
     }
     Set<Long> longValues = new HashSet<>();
     for (Expression value : values) {
-      longValues.add(((LongLiteral) value).getParsedValue());
+      longValues.add(getLongValue(value));
     }
     return TimeFilterApi.in(longValues);
   }
@@ -317,6 +318,13 @@ public Filter visitBetweenPredicate(BetweenPredicate node, Void context) {
   }
 
   public static long getLongValue(Expression expression) {
+    if (expression instanceof GenericLiteral) {
+      return Long.parseLong(((GenericLiteral) expression).getValue());
+    }
     return ((LongLiteral) expression).getParsedValue();
   }
+
+  private static boolean isLongValue(Expression expression) {
+    return expression instanceof LongLiteral || expression instanceof GenericLiteral;
+  }
 }

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/QueryPlanner.java

@@ -1356,7 +1356,12 @@ private PlanBuilder gapFill(
     long origin = ((LongLiteral) gapFillColumn.getChildren().get(3)).getParsedValue();
     long[] startAndEndTime =
         getStartTimeAndEndTimeOfGapFill(
-            timeColumn, wherePredicate, origin, monthDuration, nonMonthDuration);
+            timeColumn,
+            wherePredicate,
+            origin,
+            monthDuration,
+            nonMonthDuration,
+            subPlan.getTranslations().getPlannerContext());
     return subPlan.withNewRoot(
         new GapFillNode(
             queryIdAllocator.genPlanNodeId(),
@@ -1375,11 +1380,12 @@ private long[] getStartTimeAndEndTimeOfGapFill(
       Expression wherePredicate,
       long origin,
       int monthDuration,
-      long nonMonthDuration) {
+      long nonMonthDuration,
+      PlannerContext plannerContext) {
     GapFillStartAndEndTimeExtractVisitor.Context context =
         new GapFillStartAndEndTimeExtractVisitor.Context();
     GapFillStartAndEndTimeExtractVisitor visitor =
-        new GapFillStartAndEndTimeExtractVisitor(timeColumn);
+        new GapFillStartAndEndTimeExtractVisitor(timeColumn, plannerContext, session);
     if (!Boolean.TRUE.equals(wherePredicate.accept(visitor, context))) {
       throw new SemanticException(CAN_NOT_INFER_TIME_RANGE);
     } else {

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/TranslationMap.java

@@ -73,7 +73,6 @@ public class TranslationMap {
   private final Scope scope;
   private final Analysis analysis;
   private final Optional<TranslationMap> outerContext;
-  // TODO unused now
   private final PlannerContext plannerContext;
 
   // current mappings of underlying field -> symbol for translating direct field references
@@ -202,6 +201,10 @@ public Analysis getAnalysis() {
     return analysis;
   }
 
+  public PlannerContext getPlannerContext() {
+    return plannerContext;
+  }
+
   public boolean canTranslate(Expression expression) {
     verifyAstExpression(expression);