Extend the PARTITIONED BY clause to accept string literals for the time partitioning

docs/multi-stage-query/reference.md

@@ -199,6 +199,33 @@ The following ISO 8601 periods are supported for `TIME_FLOOR` and the string con
 - P3M
 - P1Y
 
+The string constant can also include any of the keywords mentioned above:
+
+- `HOUR` - Same as `'PT1H'`
+- `DAY` - Same as `'P1D'`
+- `MONTH` - Same as `'P1M'`
+- `YEAR` - Same as `'P1Y'`
+- `ALL TIME`
+- `ALL` - Alias for `ALL TIME`
+
+The `WEEK` granularity is deprecated and not supported in MSQ.
+
+Examples:
+
+```SQL
+-- Keyword
+PARTITIONED BY HOUR
+
+-- String constant
+PARTITIONED BY 'HOUR'
+
+-- Or
+PARTITIOND BY 'PT1H'
+
+-- TIME_FLOOR function
+PARTITIONED BY TIME_FLOOR(__time, 'PT1H')
+```
+
 For more information about partitioning, see [Partitioning](concepts.md#partitioning-by-time). <br /><br />
 *Avoid  partitioning by week, `P1W`, because weeks don't align neatly with months and years, making it difficult to partition by coarser granularities later.
 

sql/src/main/codegen/config.fmpp

@@ -52,6 +52,7 @@ data: {
       "org.apache.calcite.sql.SqlNode"
       "org.apache.calcite.sql.SqlInsert"
       "org.apache.druid.java.util.common.granularity.Granularity"
+      "org.apache.druid.java.util.common.granularity.GranularityType"
       "org.apache.druid.java.util.common.granularity.Granularities"
       "org.apache.druid.sql.calcite.parser.DruidSqlInsert"
       "org.apache.druid.sql.calcite.parser.DruidSqlParserUtils"

sql/src/main/codegen/includes/common.ftl

@@ -17,59 +17,53 @@
  * under the License.
  */
 
-// Using fully qualified name for Pair class, since Calcite also has a same class name being used in the Parser.jj
-org.apache.druid.java.util.common.Pair<Granularity, String> PartitionGranularity() :
+SqlNode PartitionGranularity() :
 {
   SqlNode e;
-  Granularity granularity;
-  String unparseString;
+  SqlNode result;
 }
 {
   (
     <HOUR>
     {
-      granularity = Granularities.HOUR;
-      unparseString = "HOUR";
+      result = SqlLiteral.createSymbol(GranularityType.HOUR, getPos());
     }
   |
     <DAY>
     {
-      granularity = Granularities.DAY;
-      unparseString = "DAY";
+      result = SqlLiteral.createSymbol(GranularityType.DAY, getPos());
     }
   |
     <MONTH>
     {
-      granularity = Granularities.MONTH;
-      unparseString = "MONTH";
+      result = SqlLiteral.createSymbol(GranularityType.MONTH, getPos());
     }
   |
     <YEAR>
     {
-      granularity = Granularities.YEAR;
-      unparseString = "YEAR";
+      result = SqlLiteral.createSymbol(GranularityType.YEAR, getPos());
     }
   |
     <ALL>
     {
-      granularity = Granularities.ALL;
-      unparseString = "ALL";
+      result = SqlLiteral.createSymbol(GranularityType.ALL, getPos());
     }
     [
       <TIME>
       {
-        unparseString += " TIME";
+        result = SqlLiteral.createSymbol(GranularityType.ALL, getPos());
       }
     ]
   |
     e = Expression(ExprContext.ACCEPT_SUB_QUERY)
     {
-      granularity = DruidSqlParserUtils.convertSqlNodeToGranularityThrowingParseExceptions(e);
-      unparseString = e.toString();
+      // validate
+      DruidSqlParserUtils.convertSqlNodeToGranularityThrowingParseExceptions(e);
+      result = e;
     }
   )
   {
-    return new org.apache.druid.java.util.common.Pair(granularity, unparseString);
+    return result;
   }
 }
 

sql/src/main/codegen/includes/insert.ftl

@@ -78,13 +78,14 @@ SqlNode DruidSqlInsert() :
 SqlNode DruidSqlInsertEof() :
 {
   SqlNode insertNode;
-  org.apache.druid.java.util.common.Pair<Granularity, String> partitionedBy = new org.apache.druid.java.util.common.Pair(null, null);
+  SqlNode partitionedBy = null;
   SqlNodeList clusteredBy = null;
 }
 {
   insertNode = DruidSqlInsert()
-  // PARTITIONED BY is necessary, but is kept optional in the grammar. It is asserted that it is not missing in the
-  // DruidSqlInsert constructor so that we can return a custom error message.
+  // PARTITIONED BY is necessary. It can be provided either in this statement or in the catalog.
+  // As a result, it is optional in the grammar. It is asserted that it is not missing in the
+  // insert analysis step.
   [
     <PARTITIONED> <BY>
     partitionedBy = PartitionGranularity()
@@ -93,7 +94,7 @@ SqlNode DruidSqlInsertEof() :
     clusteredBy = ClusteredBy()
   ]
   {
-      if (clusteredBy != null && partitionedBy.lhs == null) {
+      if (clusteredBy != null && partitionedBy == null) {
         throw org.apache.druid.sql.calcite.parser.DruidSqlParserUtils.problemParsing(
           "CLUSTERED BY found before PARTITIONED BY, CLUSTERED BY must come after the PARTITIONED BY clause"
         );
@@ -111,6 +112,6 @@ SqlNode DruidSqlInsertEof() :
       return insertNode;
     }
     SqlInsert sqlInsert = (SqlInsert) insertNode;
-    return new DruidSqlInsert(sqlInsert, partitionedBy.lhs, partitionedBy.rhs, clusteredBy);
+    return DruidSqlInsert.create(sqlInsert, partitionedBy, clusteredBy);
   }
 }

sql/src/main/codegen/includes/replace.ftl

@@ -25,8 +25,7 @@ SqlNode DruidSqlReplaceEof() :
     SqlNodeList columnList = null;
     final Span s;
     SqlInsert sqlInsert;
-    // Using fully qualified name for Pair class, since Calcite also has a same class name being used in the Parser.jj
-    org.apache.druid.java.util.common.Pair<Granularity, String> partitionedBy = new org.apache.druid.java.util.common.Pair(null, null);
+    SqlNode partitionedBy = null;
     SqlNodeList clusteredBy = null;
     final Pair<SqlNodeList, SqlNodeList> p;
     SqlNode replaceTimeQuery = null;
@@ -59,7 +58,7 @@ SqlNode DruidSqlReplaceEof() :
       clusteredBy = ClusteredBy()
     ]
     {
-        if (clusteredBy != null && partitionedBy.lhs == null) {
+        if (clusteredBy != null && partitionedBy == null) {
           throw org.apache.druid.sql.calcite.parser.DruidSqlParserUtils.problemParsing(
             "CLUSTERED BY found before PARTITIONED BY, CLUSTERED BY must come after the PARTITIONED BY clause"
           );
@@ -72,7 +71,7 @@ SqlNode DruidSqlReplaceEof() :
     <EOF>
     {
         sqlInsert = new SqlInsert(s.end(source), SqlNodeList.EMPTY, table, source, columnList);
-        return new DruidSqlReplace(sqlInsert, partitionedBy.lhs, partitionedBy.rhs, clusteredBy, replaceTimeQuery);
+        return DruidSqlReplace.create(sqlInsert, partitionedBy, clusteredBy, replaceTimeQuery);
     }
 }
 

sql/src/main/java/org/apache/druid/sql/calcite/parser/DruidSqlIngest.java

@@ -23,9 +23,10 @@
 import org.apache.calcite.sql.SqlNode;
 import org.apache.calcite.sql.SqlNodeList;
 import org.apache.calcite.sql.parser.SqlParserPos;
-import org.apache.druid.java.util.common.granularity.Granularity;
+import org.apache.calcite.util.ImmutableNullableList;
 
 import javax.annotation.Nullable;
+import java.util.List;
 
 /**
  * Common base class to the two Druid "ingest" statements: INSERT and REPLACE.
@@ -34,10 +35,8 @@
  */
 public abstract class DruidSqlIngest extends SqlInsert
 {
-  protected final Granularity partitionedBy;
-
-  // Used in the unparse function to generate the original query since we convert the string to an enum
-  protected final String partitionedByStringForUnparse;
+  @Nullable
+  protected final SqlNode partitionedBy;
 
   @Nullable
   protected final SqlNodeList clusteredBy;
@@ -48,19 +47,18 @@ public DruidSqlIngest(
       SqlNode targetTable,
       SqlNode source,
       SqlNodeList columnList,
-      @Nullable Granularity partitionedBy,
-      @Nullable String partitionedByStringForUnparse,
+      @Nullable SqlNode partitionedBy,
       @Nullable SqlNodeList clusteredBy
   )
   {
     super(pos, keywords, targetTable, source, columnList);
 
-    this.partitionedByStringForUnparse = partitionedByStringForUnparse;
     this.partitionedBy = partitionedBy;
     this.clusteredBy = clusteredBy;
   }
 
-  public Granularity getPartitionedBy()
+  @Nullable
+  public SqlNode getPartitionedBy()
   {
     return partitionedBy;
   }
@@ -70,4 +68,14 @@ public SqlNodeList getClusteredBy()
   {
     return clusteredBy;
   }
+
+  @Override
+  public List<SqlNode> getOperandList()
+  {
+    return ImmutableNullableList.<SqlNode>builder()
+        .addAll(super.getOperandList())
+        .add(partitionedBy)
+        .add(clusteredBy)
+        .build();
+  }
 }

sql/src/main/java/org/apache/druid/sql/calcite/parser/DruidSqlInsert.java

@@ -24,7 +24,7 @@
 import org.apache.calcite.sql.SqlNodeList;
 import org.apache.calcite.sql.SqlOperator;
 import org.apache.calcite.sql.SqlWriter;
-import org.apache.druid.java.util.common.granularity.Granularity;
+import org.apache.calcite.sql.parser.SqlParserPos;
 
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
@@ -41,27 +41,40 @@ public class DruidSqlInsert extends DruidSqlIngest
   // This allows reusing super.unparse
   public static final SqlOperator OPERATOR = SqlInsert.OPERATOR;
 
-  /**
-   * While partitionedBy and partitionedByStringForUnparse can be null as arguments to the constructor, this is
-   * disallowed (semantically) and the constructor performs checks to ensure that. This helps in producing friendly
-   * errors when the PARTITIONED BY custom clause is not present, and keeps its error separate from JavaCC/Calcite's
-   * custom errors which can be cryptic when someone accidentally forgets to explicitly specify the PARTITIONED BY clause
-   */
-  public DruidSqlInsert(
+  public static DruidSqlInsert create(
       @Nonnull SqlInsert insertNode,
-      @Nullable Granularity partitionedBy,
-      @Nullable String partitionedByStringForUnparse,
+      @Nullable SqlNode partitionedBy,
       @Nullable SqlNodeList clusteredBy
   )
   {
-    super(
+    return new DruidSqlInsert(
         insertNode.getParserPosition(),
         (SqlNodeList) insertNode.getOperandList().get(0), // No better getter to extract this
         insertNode.getTargetTable(),
         insertNode.getSource(),
         insertNode.getTargetColumnList(),
         partitionedBy,
-        partitionedByStringForUnparse,
+        clusteredBy
+    );
+  }
+
+  public DruidSqlInsert(
+      SqlParserPos pos,
+      SqlNodeList keywords,
+      SqlNode targetTable,
+      SqlNode source,
+      SqlNodeList columnList,
+      @Nullable SqlNode partitionedBy,
+      @Nullable SqlNodeList clusteredBy
+  )
+  {
+    super(
+        pos,
+        keywords,
+        targetTable,
+        source,
+        columnList,
+        partitionedBy,
         clusteredBy
     );
   }
@@ -77,8 +90,10 @@ public SqlOperator getOperator()
   public void unparse(SqlWriter writer, int leftPrec, int rightPrec)
   {
     super.unparse(writer, leftPrec, rightPrec);
-    writer.keyword("PARTITIONED BY");
-    writer.keyword(partitionedByStringForUnparse);
+    if (getPartitionedBy() != null) {
+      writer.keyword("PARTITIONED BY");
+      writer.keyword(partitionedBy.toString());
+    }
     if (getClusteredBy() != null) {
       writer.keyword("CLUSTERED BY");
       SqlWriter.Frame frame = writer.startList("", "");