rdar://72811621 Control distribution and ordering during table creation (apache#939)

### What changes were proposed in this pull request?

This PR adds support for defining distribution and ordering during table creation.

### Why are the changes needed?

This change is needed for feature parity with ADT in Spark 2.

### Does this PR introduce any user-facing change?

This PR adds new optional table creation clauses but it should only affect Iceberg users.

### How was this patch tested?

This PR comes with new tests. More tests are in Iceberg.

Author: aokolnychyi

pom.xml

@@ -194,7 +194,7 @@
     <selenium.version>3.141.59</selenium.version>
     <htmlunit.version>2.40.0</htmlunit.version>
     <acispark.callhome.version>0.2.0</acispark.callhome.version>
-    <iceberg.version>0.12.0-apple-preview</iceberg.version>
+    <iceberg.version>0.12.0-apple-preview-2-SNAPSHOT</iceberg.version>
     <!--
     Managed up from older version from Avro; sync with jackson-module-paranamer dependency version
     -->
@@ -294,6 +294,10 @@
       <id>apple</id>
       <url>https://artifacts.apple.com/libs-release</url>
     </repository>
+    <repository>
+      <id>apple-snapshot</id>
+      <url>https://artifacts.apple.com/libs-snapshot</url>
+    </repository>
   </repositories>
   <pluginRepositories>
     <pluginRepository>

project/SparkBuild.scala

@@ -270,6 +270,7 @@ object SparkBuild extends PomBuild {
       "gcs-maven-central-mirror" at "https://maven-central.storage-download.googleapis.com/maven2/",
       DefaultMavenRepository,
       "apple" at "https://artifacts.apple.com/libs-release",
+      "apple-snapshot" at "https://artifacts.apple.com/libs-snapshot",
       Resolver.mavenLocal,
       Resolver.file("ivyLocal", file(Path.userHome.absolutePath + "/.ivy2/local"))(Resolver.ivyStylePatterns)
     ),

sql/catalyst/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBase.g4

@@ -396,6 +396,7 @@ createTableClauses
      (PARTITIONED BY partitioning=partitionFieldList) |
      skewSpec |
      bucketSpec |
+     writeSpec |
      rowFormat |
      createFileFormat |
      locationSpec |

sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/StagingTableCatalog.java

@@ -19,7 +19,10 @@
 
 import java.util.Map;
 
+import com.google.common.base.Preconditions;
+
 import org.apache.spark.annotation.Evolving;
+import org.apache.spark.sql.connector.expressions.SortOrder;
 import org.apache.spark.sql.connector.expressions.Transform;
 import org.apache.spark.sql.connector.write.LogicalWriteInfo;
 import org.apache.spark.sql.catalyst.analysis.NoSuchNamespaceException;
@@ -78,6 +81,30 @@ StagedTable stageCreate(
       Transform[] partitions,
       Map<String, String> properties) throws TableAlreadyExistsException, NoSuchNamespaceException;
 
+  /**
+   * Stage the creation of a table with a specific distribution and ordering.
+   */
+  default StagedTable stageCreate(
+      Identifier ident,
+      StructType schema,
+      Transform[] partitions,
+      Map<String, String> properties,
+      String distributionMode,
+      SortOrder[] ordering) throws TableAlreadyExistsException, NoSuchNamespaceException {
+
+    Preconditions.checkArgument(
+        distributionMode.equals("none"),
+        "%s does not support tables with a specific distribution",
+        this.getClass().getName());
+
+    Preconditions.checkArgument(
+        ordering.length == 0,
+        "%s does not support tables with a specific ordering",
+        this.getClass().getName());
+
+    return stageCreate(ident, schema, partitions, properties);
+  }
+
   /**
    * Stage the replacement of a table, preparing it to be committed into the metastore when the
    * returned table's {@link StagedTable#commitStagedChanges()} is called.
@@ -111,6 +138,30 @@ StagedTable stageReplace(
       Transform[] partitions,
       Map<String, String> properties) throws NoSuchNamespaceException, NoSuchTableException;
 
+  /**
+   * Stage the replacement of a table with a specific distribution and ordering.
+   */
+  default StagedTable stageReplace(
+      Identifier ident,
+      StructType schema,
+      Transform[] partitions,
+      Map<String, String> properties,
+      String distributionMode,
+      SortOrder[] ordering) throws NoSuchNamespaceException, NoSuchTableException {
+
+    Preconditions.checkArgument(
+        distributionMode.equals("none"),
+        "%s does not support tables with a specific distribution",
+        this.getClass().getName());
+
+    Preconditions.checkArgument(
+        ordering.length == 0,
+        "%s does not support tables with a specific ordering",
+        this.getClass().getName());
+
+    return stageReplace(ident, schema, partitions, properties);
+  }
+
   /**
    * Stage the creation or replacement of a table, preparing it to be committed into the metastore
    * when the returned table's {@link StagedTable#commitStagedChanges()} is called.
@@ -141,4 +192,28 @@ StagedTable stageCreateOrReplace(
       StructType schema,
       Transform[] partitions,
       Map<String, String> properties) throws NoSuchNamespaceException;
+
+  /**
+   * Stage the creation or replacement of a table with a specific distribution and ordering.
+   */
+  default StagedTable stageCreateOrReplace(
+      Identifier ident,
+      StructType schema,
+      Transform[] partitions,
+      Map<String, String> properties,
+      String distributionMode,
+      SortOrder[] ordering) throws NoSuchNamespaceException {
+
+    Preconditions.checkArgument(
+        distributionMode.equals("none"),
+        "%s does not support tables with a specific distribution",
+        this.getClass().getName());
+
+    Preconditions.checkArgument(
+        ordering.length == 0,
+        "%s does not support tables with a specific ordering",
+        this.getClass().getName());
+
+    return stageCreateOrReplace(ident, schema, partitions, properties);
+  }
 }

sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/TableCatalog.java

@@ -18,6 +18,7 @@
 package org.apache.spark.sql.connector.catalog;
 
 import org.apache.spark.annotation.Evolving;
+import org.apache.spark.sql.connector.expressions.SortOrder;
 import org.apache.spark.sql.connector.expressions.Transform;
 import org.apache.spark.sql.catalyst.analysis.NoSuchNamespaceException;
 import org.apache.spark.sql.catalyst.analysis.NoSuchTableException;
@@ -26,6 +27,8 @@
 
 import java.util.Map;
 
+import com.google.common.base.Preconditions;
+
 /**
  * Catalog methods for working with Tables.
  * <p>
@@ -140,6 +143,30 @@ Table createTable(
       Transform[] partitions,
       Map<String, String> properties) throws TableAlreadyExistsException, NoSuchNamespaceException;
 
+  /**
+   * Create a table with a specific distribution and ordering.
+   */
+  default Table createTable(
+      Identifier ident,
+      StructType schema,
+      Transform[] partitions,
+      Map<String, String> properties,
+      String distributionMode,
+      SortOrder[] ordering) throws TableAlreadyExistsException, NoSuchNamespaceException {
+
+    Preconditions.checkArgument(
+        distributionMode.equals("none"),
+        "%s does not support tables with a specific distribution",
+        this.getClass().getName());
+
+    Preconditions.checkArgument(
+        ordering.length == 0,
+        "%s does not support tables with a specific ordering",
+        this.getClass().getName());
+
+    return createTable(ident, schema, partitions, properties);
+  }
+
   /**
    * Apply a set of {@link TableChange changes} to a table in the catalog.
    * <p>

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveCatalogs.scala

@@ -155,7 +155,7 @@ class ResolveCatalogs(val catalogManager: CatalogManager)
       RenameTable(catalog.asTableCatalog, oldName.asIdentifier, newNameParts.asIdentifier)
 
     case c @ CreateTableStatement(
-         NonSessionCatalogAndTable(catalog, tbl), _, _, _, _, _, _, _, _, _, _, _) =>
+         NonSessionCatalogAndTable(catalog, tbl), _, _, _, _, _, _, _, _, _, _, _, _, _) =>
       assertNoNullTypeInSchema(c.tableSchema)
       CreateV2Table(
         catalog.asTableCatalog,
@@ -164,10 +164,12 @@ class ResolveCatalogs(val catalogManager: CatalogManager)
         // convert the bucket spec and add it as a transform
         c.partitioning ++ c.bucketSpec.map(_.asTransform),
         convertTableProperties(c),
-        ignoreIfExists = c.ifNotExists)
+        ignoreIfExists = c.ifNotExists,
+        c.distributionMode,
+        c.ordering)
 
     case c @ CreateTableAsSelectStatement(
-         NonSessionCatalogAndTable(catalog, tbl), _, _, _, _, _, _, _, _, _, _, _, _) =>
+         NonSessionCatalogAndTable(catalog, tbl), _, _, _, _, _, _, _, _, _, _, _, _, _, _) =>
       if (c.asSelect.resolved) {
         assertNoNullTypeInSchema(c.asSelect.schema)
       }
@@ -179,7 +181,9 @@ class ResolveCatalogs(val catalogManager: CatalogManager)
         c.asSelect,
         convertTableProperties(c),
         writeOptions = c.writeOptions,
-        ignoreIfExists = c.ifNotExists)
+        ignoreIfExists = c.ifNotExists,
+        c.distributionMode,
+        c.ordering)
 
     case m @ MigrateTableStatement(NonSessionCatalogAndTable(catalog, tbl), _, _) =>
       if (!catalog.isInstanceOf[SupportsMigrate]) {
@@ -207,7 +211,7 @@ class ResolveCatalogs(val catalogManager: CatalogManager)
         )
 
     case c @ ReplaceTableStatement(
-         NonSessionCatalogAndTable(catalog, tbl), _, _, _, _, _, _, _, _, _, _) =>
+         NonSessionCatalogAndTable(catalog, tbl), _, _, _, _, _, _, _, _, _, _, _, _) =>
       assertNoNullTypeInSchema(c.tableSchema)
       ReplaceTable(
         catalog.asTableCatalog,
@@ -216,10 +220,12 @@ class ResolveCatalogs(val catalogManager: CatalogManager)
         // convert the bucket spec and add it as a transform
         c.partitioning ++ c.bucketSpec.map(_.asTransform),
         convertTableProperties(c),
-        orCreate = c.orCreate)
+        orCreate = c.orCreate,
+        distributionMode = c.distributionMode,
+        ordering = c.ordering)
 
     case c @ ReplaceTableAsSelectStatement(
-         NonSessionCatalogAndTable(catalog, tbl), _, _, _, _, _, _, _, _, _, _, _) =>
+         NonSessionCatalogAndTable(catalog, tbl), _, _, _, _, _, _, _, _, _, _, _, _, _) =>
       if (c.asSelect.resolved) {
         assertNoNullTypeInSchema(c.asSelect.schema)
       }
@@ -231,7 +237,9 @@ class ResolveCatalogs(val catalogManager: CatalogManager)
         c.asSelect,
         convertTableProperties(c),
         writeOptions = c.writeOptions,
-        orCreate = c.orCreate)
+        orCreate = c.orCreate,
+        distributionMode = c.distributionMode,
+        ordering = c.ordering)
 
     case DropViewStatement(NonSessionCatalogAndTable(catalog, viewName), _) =>
       throw new AnalysisException(

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/parser/AstBuilder.scala

@@ -2517,6 +2517,13 @@ class AstBuilder extends SqlBaseBaseVisitor[AnyRef] with SQLConfHelper with Logg
       Seq[Transform], Seq[StructField], Option[BucketSpec], Map[String, String],
       Map[String, String], Option[String], Option[String], Option[SerdeInfo])
 
+  /**
+   * Type to keep track of custom table clauses:
+   * - distribution mode
+   * - ordering
+   */
+  type CustomTableClauses = (String, Seq[V2SortOrder])
+
   /**
    * Validate a create table statement and return the [[TableIdentifier]].
    */
@@ -2968,6 +2975,21 @@ class AstBuilder extends SqlBaseBaseVisitor[AnyRef] with SQLConfHelper with Logg
     }
   }
 
+  private def visitCustomCreateTableClauses(ctx: CreateTableClausesContext): CustomTableClauses = {
+    checkDuplicateClauses(ctx.writeSpec, "DISTRIBUTED/ORDERED BY", ctx)
+
+    ctx.writeSpec.asScala.headOption match {
+      case Some(writeSpec) =>
+        val (distributionSpec, orderingSpec) = toDistributionAndOrderingSpec(writeSpec)
+        val distributionMode = toDistributionMode(distributionSpec, orderingSpec)
+        val ordering = toOrdering(orderingSpec)
+        (distributionMode, ordering)
+
+      case None =>
+        ("none", Array.empty[V2SortOrder])
+    }
+  }
+
   override def visitCreateTableClauses(ctx: CreateTableClausesContext): TableClauses = {
     checkDuplicateClauses(ctx.TBLPROPERTIES, "TBLPROPERTIES", ctx)
     checkDuplicateClauses(ctx.OPTIONS, "OPTIONS", ctx)
@@ -3080,6 +3102,11 @@ class AstBuilder extends SqlBaseBaseVisitor[AnyRef] with SQLConfHelper with Logg
     }
 
     val partitioning = partitionExpressions(partTransforms, partCols, ctx)
+    val (distributionMode, ordering) = visitCustomCreateTableClauses(ctx.createTableClauses)
+
+    if (distributionMode == "hash" && partitioning.isEmpty) {
+      operationNotAllowed("DISTRIBUTED BY PARTITION is supported only for partitioned tables", ctx)
+    }
 
     Option(ctx.query).map(plan) match {
       case Some(_) if columns.nonEmpty =>
@@ -3096,14 +3123,16 @@ class AstBuilder extends SqlBaseBaseVisitor[AnyRef] with SQLConfHelper with Logg
       case Some(query) =>
         CreateTableAsSelectStatement(
           table, query, partitioning, bucketSpec, properties, provider, options, location, comment,
-          writeOptions = Map.empty, serdeInfo, external = external, ifNotExists = ifNotExists)
+          writeOptions = Map.empty, serdeInfo, external = external, ifNotExists = ifNotExists,
+          distributionMode, ordering)
 
       case _ =>
         // Note: table schema includes both the table columns list and the partition columns
         // with data type.
         val schema = StructType(columns ++ partCols)
         CreateTableStatement(table, schema, partitioning, bucketSpec, properties, provider,
-          options, location, comment, serdeInfo, external = external, ifNotExists = ifNotExists)
+          options, location, comment, serdeInfo, external = external, ifNotExists = ifNotExists,
+          distributionMode, ordering)
     }
   }
 
@@ -3160,6 +3189,11 @@ class AstBuilder extends SqlBaseBaseVisitor[AnyRef] with SQLConfHelper with Logg
     }
 
     val partitioning = partitionExpressions(partTransforms, partCols, ctx)
+    val (distributionMode, ordering) = visitCustomCreateTableClauses(ctx.createTableClauses)
+
+    if (distributionMode == "hash" && partitioning.isEmpty) {
+      operationNotAllowed("DISTRIBUTED BY PARTITION is supported only for partitioned tables", ctx)
+    }
 
     Option(ctx.query).map(plan) match {
       case Some(_) if columns.nonEmpty =>
@@ -3176,14 +3210,14 @@ class AstBuilder extends SqlBaseBaseVisitor[AnyRef] with SQLConfHelper with Logg
       case Some(query) =>
         ReplaceTableAsSelectStatement(table, query, partitioning, bucketSpec, properties,
           provider, options, location, comment, writeOptions = Map.empty, serdeInfo,
-          orCreate = orCreate)
+          orCreate = orCreate, distributionMode, ordering)
 
       case _ =>
         // Note: table schema includes both the table columns list and the partition columns
         // with data type.
         val schema = StructType(columns ++ partCols)
         ReplaceTableStatement(table, schema, partitioning, bucketSpec, properties, provider,
-          options, location, comment, serdeInfo, orCreate = orCreate)
+          options, location, comment, serdeInfo, orCreate = orCreate, distributionMode, ordering)
     }
   }
 

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/statements.scala

@@ -145,7 +145,9 @@ case class CreateTableStatement(
     comment: Option[String],
     serde: Option[SerdeInfo],
     external: Boolean,
-    ifNotExists: Boolean) extends ParsedStatement
+    ifNotExists: Boolean,
+    distributionMode: String = "none",
+    ordering: Seq[SortOrder] = Seq.empty) extends ParsedStatement
 
 /**
  * A CREATE TABLE AS SELECT command, as parsed from SQL.
@@ -163,7 +165,9 @@ case class CreateTableAsSelectStatement(
     writeOptions: Map[String, String],
     serde: Option[SerdeInfo],
     external: Boolean,
-    ifNotExists: Boolean) extends ParsedStatement {
+    ifNotExists: Boolean,
+    distributionMode: String = "none",
+    ordering: Seq[SortOrder] = Seq.empty) extends ParsedStatement {
 
   override def children: Seq[LogicalPlan] = Seq(asSelect)
 }
@@ -199,7 +203,9 @@ case class ReplaceTableStatement(
     location: Option[String],
     comment: Option[String],
     serde: Option[SerdeInfo],
-    orCreate: Boolean) extends ParsedStatement
+    orCreate: Boolean,
+    distributionMode: String = "none",
+    ordering: Seq[SortOrder] = Seq.empty) extends ParsedStatement
 
 /**
  * A REPLACE TABLE AS SELECT command, as parsed from SQL.
@@ -216,7 +222,9 @@ case class ReplaceTableAsSelectStatement(
     comment: Option[String],
     writeOptions: Map[String, String],
     serde: Option[SerdeInfo],
-    orCreate: Boolean) extends ParsedStatement {
+    orCreate: Boolean,
+    distributionMode: String = "none",
+    ordering: Seq[SortOrder] = Seq.empty) extends ParsedStatement {
 
   override def children: Seq[LogicalPlan] = Seq(asSelect)
 }