[IcebergIO] extend table partitioning to Beam SQL (#35268)

* add sql partitioning

* trigger ITs and add to CHANGES

* simplify test

* include emoty namespace fix

* add error when using partitioning for unsupported tables

* fix test

Author: ahmedabu98

.github/trigger_files/IO_Iceberg_Integration_Tests.json

@@ -1,4 +1,4 @@
 {
     "comment": "Modify this file in a trivial way to cause this test suite to run.",
-    "modification": 6
+    "modification": 3
 }

.github/trigger_files/beam_PostCommit_SQL.json

@@ -1,4 +1,4 @@
 {
   "comment": "Modify this file in a trivial way to cause this test suite to run",
-  "modification": 2
+  "modification": 3
 }

CHANGES.md

@@ -78,7 +78,7 @@
 * [IcebergIO] Now available with Beam SQL! ([#34799](https://github.com/apache/beam/pull/34799))
 * [IcebergIO] Support reading with column pruning ([#34856](https://github.com/apache/beam/pull/34856))
 * [IcebergIO] Support reading with pushdown filtering ([#34827](https://github.com/apache/beam/pull/34827))
-* [IcebergIO] Create tables with a specified partition spec ([#34966](https://github.com/apache/beam/pull/34966))
+* [IcebergIO] Create tables with a specified partition spec ([#34966](https://github.com/apache/beam/pull/34966), [#35268](https://github.com/apache/beam/pull/35268))
 * [IcebergIO] Dynamically create namespaces if needed ([#35228](https://github.com/apache/beam/pull/35228))
 
 ## New Features / Improvements

sdks/java/extensions/sql/build.gradle

@@ -114,6 +114,9 @@ dependencies {
   provided library.java.hadoop_client
   permitUnusedDeclared library.java.hadoop_client
   provided library.java.kafka_clients
+
+  testImplementation "org.apache.iceberg:iceberg-api:1.6.1"
+  testImplementation "org.apache.iceberg:iceberg-core:1.6.1"
   testImplementation library.java.vendored_calcite_1_28_0
   testImplementation library.java.vendored_guava_32_1_2_jre
   testImplementation library.java.junit

sdks/java/extensions/sql/src/main/codegen/config.fmpp

@@ -45,6 +45,7 @@ data: {
         "LOCATION"
         "TBLPROPERTIES"
         "PROPERTIES"
+        "PARTITIONED"
       ]
 
       # List of keywords from "keywords" section that are not reserved.

sdks/java/extensions/sql/src/main/codegen/includes/parserImpls.ftl

@@ -255,12 +255,28 @@ SqlDrop SqlDropCatalog(Span s, boolean replace) :
     }
 }
 
+SqlNodeList PartitionFieldList() :
+{
+    final List<SqlNode> list = new ArrayList<SqlNode>();
+    SqlNode field;
+}
+{
+    field = StringLiteral() { list.add(field); }
+    (
+        <COMMA> field = StringLiteral() { list.add(field); }
+    )*
+    {
+        return new SqlNodeList(list, getPos());
+    }
+}
+
 /**
  * Note: This example is probably out of sync with the code.
  *
- * CREATE TABLE ( IF NOT EXISTS )?
+ * CREATE EXTERNAL TABLE ( IF NOT EXISTS )?
  *   ( database_name '.' )? table_name '(' column_def ( ',' column_def )* ')'
  *   TYPE type_name
+ *   ( PARTITIONED BY '(' partition_field ( ',' partition_field )* ')' )?
  *   ( COMMENT comment_string )?
  *   ( LOCATION location_string )?
  *   ( TBLPROPERTIES tbl_properties )?
@@ -271,6 +287,7 @@ SqlCreate SqlCreateExternalTable(Span s, boolean replace) :
     final SqlIdentifier id;
     List<Schema.Field> fieldList = null;
     final SqlNode type;
+    SqlNodeList partitionFields = null;
     SqlNode comment = null;
     SqlNode location = null;
     SqlNode tblProperties = null;
@@ -290,6 +307,7 @@ SqlCreate SqlCreateExternalTable(Span s, boolean replace) :
     |
         type = SimpleIdentifier()
     )
+    [ <PARTITIONED> <BY> <LPAREN> partitionFields = PartitionFieldList() <RPAREN> ]
     [ <COMMENT> comment = StringLiteral() ]
     [ <LOCATION> location = StringLiteral() ]
     [ <TBLPROPERTIES> tblProperties = StringLiteral() ]
@@ -302,6 +320,7 @@ SqlCreate SqlCreateExternalTable(Span s, boolean replace) :
                 id,
                 fieldList,
                 type,
+                partitionFields,
                 comment,
                 location,
                 tblProperties);

sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/parser/SqlCreateExternalTable.java

@@ -19,9 +19,11 @@
 
 import static org.apache.beam.sdk.schemas.Schema.toSchema;
 import static org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.util.Static.RESOURCE;
+import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument;
 import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkNotNull;
 
 import java.util.List;
+import java.util.stream.Collectors;
 import org.apache.beam.sdk.extensions.sql.TableUtils;
 import org.apache.beam.sdk.extensions.sql.impl.BeamCalciteSchema;
 import org.apache.beam.sdk.extensions.sql.impl.utils.CalciteUtils;
@@ -33,12 +35,14 @@
 import org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.sql.SqlIdentifier;
 import org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.sql.SqlKind;
 import org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.sql.SqlNode;
+import org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.sql.SqlNodeList;
 import org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.sql.SqlOperator;
 import org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.sql.SqlSpecialOperator;
 import org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.sql.SqlUtil;
 import org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.sql.SqlWriter;
 import org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.sql.parser.SqlParserPos;
 import org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.util.Pair;
+import org.checkerframework.checker.nullness.qual.Nullable;
 
 /** Parse tree for {@code CREATE EXTERNAL TABLE} statement. */
 @SuppressWarnings({
@@ -51,6 +55,7 @@ public class SqlCreateExternalTable extends SqlCreate implements BeamSqlParser.E
   private final SqlNode comment;
   private final SqlNode location;
   private final SqlNode tblProperties;
+  private final @Nullable SqlNodeList partitionFields;
 
   private static final SqlOperator OPERATOR =
       new SqlSpecialOperator("CREATE EXTERNAL TABLE", SqlKind.OTHER_DDL);
@@ -63,13 +68,15 @@ public SqlCreateExternalTable(
       SqlIdentifier name,
       List<Schema.Field> columnList,
       SqlNode type,
+      SqlNodeList partitionFields,
       SqlNode comment,
       SqlNode location,
       SqlNode tblProperties) {
     super(OPERATOR, pos, replace, ifNotExists);
     this.name = checkNotNull(name);
     this.columnList = columnList; // may be null
     this.type = checkNotNull(type);
+    this.partitionFields = partitionFields;
     this.comment = comment; // may be null
     this.location = location; // may be null
     this.tblProperties = tblProperties; // may be null
@@ -98,6 +105,19 @@ public void unparse(SqlWriter writer, int leftPrec, int rightPrec) {
     }
     writer.keyword("TYPE");
     type.unparse(writer, 0, 0);
+    if (partitionFields != null) {
+      writer.keyword("PARTITIONED");
+      writer.keyword("BY");
+      writer.sep("(");
+      for (int i = 0; i < partitionFields.size(); i++) {
+        if (i > 0) {
+          writer.sep(",");
+        }
+        SqlNode field = partitionFields.get(i);
+        field.unparse(writer, 0, 0);
+      }
+      writer.sep(")");
+    }
     if (comment != null) {
       writer.keyword("COMMENT");
       comment.unparse(writer, 0, 0);
@@ -130,7 +150,17 @@ public void execute(CalcitePrepare.Context context) {
           name.getParserPosition(),
           RESOURCE.internal("Schema is not instanceof BeamCalciteSchema"));
     }
+
     BeamCalciteSchema schema = (BeamCalciteSchema) pair.left.schema;
+    if (partitionFields != null) {
+      checkArgument(
+          schema.resolveMetastore().supportsPartitioning(),
+          "Invalid use of 'PARTITIONED BY()': Table '%s' of type '%s' "
+              + "does not support partitioning.",
+          SqlDdlNodes.name(name),
+          SqlDdlNodes.getString(type));
+    }
+
     schema.resolveMetastore().createTable(toTable());
   }
 
@@ -149,11 +179,19 @@ private void unparseColumn(SqlWriter writer, Schema.Field column) {
     }
   }
 
+  private @Nullable List<String> parsePartitionFields() {
+    if (partitionFields == null) {
+      return null;
+    }
+    return partitionFields.stream().map(SqlDdlNodes::getString).collect(Collectors.toList());
+  }
+
   private Table toTable() {
     return Table.builder()
         .type(SqlDdlNodes.getString(type))
         .name(SqlDdlNodes.name(name))
         .schema(columnList.stream().collect(toSchema()))
+        .partitionFields(parsePartitionFields())
         .comment(SqlDdlNodes.getString(comment))
         .location(SqlDdlNodes.getString(location))
         .properties(

sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/meta/Table.java

@@ -20,6 +20,7 @@
 import com.fasterxml.jackson.databind.node.ObjectNode;
 import com.google.auto.value.AutoValue;
 import java.io.Serializable;
+import java.util.List;
 import org.apache.beam.sdk.extensions.sql.TableUtils;
 import org.apache.beam.sdk.schemas.Schema;
 import org.checkerframework.checker.nullness.qual.Nullable;
@@ -34,6 +35,8 @@ public abstract class Table implements Serializable {
 
   public abstract Schema getSchema();
 
+  public abstract @Nullable List<String> getPartitionFields();
+
   public abstract @Nullable String getComment();
 
   public abstract @Nullable String getLocation();
@@ -55,6 +58,8 @@ public abstract static class Builder {
 
     public abstract Builder schema(Schema getSchema);
 
+    public abstract Builder partitionFields(List<String> fields);
+
     public abstract Builder comment(String name);
 
     public abstract Builder location(String location);

sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/meta/catalog/InMemoryCatalog.java

@@ -25,7 +25,7 @@
 public class InMemoryCatalog implements Catalog {
   private final String name;
   private final Map<String, String> properties;
-  protected final InMemoryMetaStore metaStore = new InMemoryMetaStore();
+  private final InMemoryMetaStore metaStore = new InMemoryMetaStore();
 
   public InMemoryCatalog(String name, Map<String, String> properties) {
     this.name = name;

sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/meta/provider/TableProvider.java

@@ -79,4 +79,8 @@ default Set<String> getSubProviders() {
   default TableProvider getSubProvider(String name) {
     return null;
   }
+
+  default boolean supportsPartitioning() {
+    return false;
+  }
 }