szehon-ho commented on code in PR #58153: URL: https://github.com/apache/spark/pull/58153#discussion_r4140126618
########## sql/core/src/test/scala/org/apache/spark/sql/connector/CreateTableWriteOrderSuite.scala: ########## @@ -0,0 +1,780 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.connector + +import java.util + +import scala.collection.mutable.ArrayBuffer +import scala.jdk.CollectionConverters._ + +import org.apache.spark.sql.{AnalysisException, QueryTest, Row} +import org.apache.spark.sql.catalyst.analysis.TableAlreadyExistsException +import org.apache.spark.sql.catalyst.parser.ParseException +import org.apache.spark.sql.catalyst.plans.logical.{CreateTable, CreateTableAsSelect, LogicalPlan, ReplaceTable, ReplaceTableAsSelect, V2CreateTablePlan} +import org.apache.spark.sql.connector.catalog.{Column, DelegatingTable, Identifier, InMemoryTable, InMemoryTableCatalog, StagedTable, StagingInMemoryTableCatalog, Table, TableCatalogCapability, TableInfo} +import org.apache.spark.sql.connector.catalog.CatalogV2Implicits._ +import org.apache.spark.sql.connector.expressions.{ClusterByTransform, FieldReference, IdentityTransform, LogicalExpressions, NullOrdering, SortDirection, SortOrder} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.test.SharedSparkSession +import org.apache.spark.sql.types.IntegerType + +/** + * Tests the create-time write distribution and ordering clauses: CREATE/REPLACE TABLE ... + * DISTRIBUTED BY PARTITION ... [LOCALLY] ORDERED BY ... | UNORDERED. Covers what the parser + * produces, what reaches the catalog, and how a catalog without the capability is rejected. + */ +class CreateTableWriteOrderSuite extends QueryTest with SharedSparkSession { + + // The catalog manager caches catalogs for the session; reset it so tests don't share tables. + override def afterEach(): Unit = { + spark.sessionState.catalogManager.reset() + super.afterEach() + } + + private def parse(sql: String): LogicalPlan = spark.sessionState.sqlParser.parsePlan(sql) + + private def analyze(sql: String): LogicalPlan = + spark.sessionState.executePlan(parse(sql)).analyzed + + private def orderingOf(plan: LogicalPlan): Seq[String] = plan match { + case create: V2CreateTablePlan => WriteSpecCall.render(create.writeOrdering) + case other => fail(s"unexpected plan: $other") + } + + private def calls(catalogName: String): Seq[WriteSpecCall] = + spark.sessionState.catalogManager.catalog(catalogName) + .asInstanceOf[RecordsWriteSpecs].recordedCalls + + test("parse DISTRIBUTED BY PARTITION / ORDERED BY on CREATE TABLE") { + parse("CREATE TABLE t (id INT, c STRING) USING foo PARTITIONED BY (c) " + + "DISTRIBUTED BY PARTITION ORDERED BY id ASC NULLS FIRST") match { + case c: CreateTable => + assert(c.writeDistributionMode === "hash") + assert(c.writeOrdering.length === 1) + val order = c.writeOrdering.head + assert(order.expression().isInstanceOf[IdentityTransform]) + assert(order.direction() === SortDirection.ASCENDING) + assert(order.nullOrdering() === NullOrdering.NULLS_FIRST) + case other => fail(s"unexpected plan: $other") + } + } + + test("bare ORDERED BY implies range, LOCALLY implies none, UNORDERED implies none") { + Seq( + ("ORDERED BY id", "range", 1), + ("LOCALLY ORDERED BY id", "none", 1), + ("UNORDERED", "none", 0) + ).foreach { case (clause, expectedMode, expectedOrderingSize) => + parse(s"CREATE TABLE t (id INT) USING foo $clause") match { + case c: CreateTable => + assert(c.writeDistributionMode === expectedMode, s"for $clause") + assert(c.writeOrdering.length === expectedOrderingSize, s"for $clause") + case other => fail(s"unexpected plan for $clause: $other") + } + } + } + + test("an explicit DISTRIBUTED BY PARTITION decides the distribution on its own") { + Seq( + ("ORDERED BY id", 1), + ("LOCALLY ORDERED BY id", 1), + ("UNORDERED", 0) + ).foreach { case (clause, expectedOrderingSize) => + parse(s"CREATE TABLE t (id INT, c STRING) USING foo PARTITIONED BY (c) " + + s"DISTRIBUTED BY PARTITION $clause") match { + case c: CreateTable => + assert(c.writeDistributionMode === "hash", s"for $clause") + assert(c.writeOrdering.length === expectedOrderingSize, s"for $clause") + case other => fail(s"unexpected plan for $clause: $other") + } + } + } + + test("no clause leaves the mode unset (null), distinct from an explicit none") { + parse("CREATE TABLE t (id INT) USING foo") match { + case c: CreateTable => + assert(c.writeDistributionMode === null) + assert(c.writeOrdering.isEmpty) + case other => fail(s"unexpected plan: $other") + } + } + + test("CTAS / REPLACE / RTAS carry the clauses too") { + parse("CREATE TABLE t USING foo ORDERED BY id AS SELECT 1 AS id") match { + case c: CreateTableAsSelect => + assert(c.writeDistributionMode === "range" && c.writeOrdering.length == 1) + case other => fail(s"unexpected plan: $other") + } + parse("REPLACE TABLE t (id INT) USING foo ORDERED BY id DESC NULLS LAST") match { + case r: ReplaceTable => + assert(r.writeDistributionMode === "range") + assert(r.writeOrdering.head.direction() === SortDirection.DESCENDING) + assert(r.writeOrdering.head.nullOrdering() === NullOrdering.NULLS_LAST) + case other => fail(s"unexpected plan: $other") + } + parse("REPLACE TABLE t USING foo UNORDERED AS SELECT 1 AS id") match { + case r: ReplaceTableAsSelect => + assert(r.writeDistributionMode === "none" && r.writeOrdering.isEmpty) + case other => fail(s"unexpected plan: $other") + } + } + + test("transforms are allowed in the ordering") { + val stmt = "CREATE TABLE t (id INT, ts TIMESTAMP) USING foo " + + "ORDERED BY days(ts), bucket(16, id)" + parse(stmt) match { + case c: CreateTable => + assert(c.writeOrdering.length === 2) + assert(c.writeOrdering.map(_.expression().describe()) === + Seq("days(ts)", "bucket(16, id)")) + case other => fail(s"unexpected plan: $other") + } + } + + test("the parenthesised ORDERED BY form parses the same") { + Seq("ORDERED BY id DESC, c", "ORDERED BY (id DESC, c)").foreach { clause => + parse(s"CREATE TABLE t (id INT, c STRING) USING foo $clause") match { + case c: CreateTable => + assert(WriteSpecCall.render(c.writeOrdering) === + Seq("id DESC NULLS LAST", "c ASC NULLS FIRST"), s"for $clause") + case other => fail(s"unexpected plan for $clause: $other") + } + } + } + + test("the write clauses may appear in any position among the other create-table clauses") { + Seq( + "ORDERED BY id PARTITIONED BY (c) DISTRIBUTED BY PARTITION", + "DISTRIBUTED BY PARTITION PARTITIONED BY (c) LOCALLY ORDERED BY id", + "PARTITIONED BY (c) ORDERED BY id COMMENT 'a table' DISTRIBUTED BY PARTITION" + ).foreach { clauses => + parse(s"CREATE TABLE t (id INT, c STRING) USING foo $clauses") match { + case c: CreateTable => + assert(c.writeDistributionMode === "hash", s"for $clauses") + assert(c.writeOrdering.length === 1, s"for $clauses") + case other => fail(s"unexpected plan for $clauses: $other") + } + } + } + + test("a repeated write clause is rejected") { + Seq( + ("DISTRIBUTED BY PARTITION DISTRIBUTED BY PARTITION", "DISTRIBUTED BY PARTITION"), + ("ORDERED BY id UNORDERED", "ORDERED BY/UNORDERED"), + ("UNORDERED LOCALLY ORDERED BY id", "ORDERED BY/UNORDERED") + ).foreach { case (clauses, clauseName) => + val e = intercept[ParseException] { + parse(s"CREATE TABLE t (id INT, c STRING) USING foo PARTITIONED BY (c) $clauses") + } + assert(e.getCondition === "DUPLICATE_CLAUSES", s"for $clauses") + assert(e.getMessageParameters.get("clauseName") === clauseName, s"for $clauses") + } + } + + test("DISTRIBUTED BY PARTITION requires a partitioned table") { + Seq( + "CREATE TABLE t (id INT) USING foo DISTRIBUTED BY PARTITION", + "CREATE TABLE t USING foo DISTRIBUTED BY PARTITION AS SELECT 1 AS id", + "REPLACE TABLE t (id INT) USING foo DISTRIBUTED BY PARTITION", + "REPLACE TABLE t USING foo DISTRIBUTED BY PARTITION AS SELECT 1 AS id" + ).foreach { stmt => + val e = intercept[ParseException](parse(stmt)) + assert(e.getCondition === + "SPECIFY_DISTRIBUTED_BY_PARTITION_WITHOUT_PARTITIONING_IS_NOT_ALLOWED", s"for $stmt") + } + } + + test("ordering references are normalized to the schema's spelling") { + withSQLConf("spark.sql.catalog.testcat" -> classOf[RecordingInMemoryTableCatalog].getName) { + assert(orderingOf(analyze("CREATE TABLE testcat.t (ID INT, TS TIMESTAMP) USING foo " + + "ORDERED BY id, days(ts) DESC")) === Seq("ID ASC NULLS FIRST", "days(TS) DESC NULLS LAST")) + + assert(orderingOf(analyze( + "CREATE TABLE testcat.t USING foo ORDERED BY id AS SELECT 1 AS ID")) === + Seq("ID ASC NULLS FIRST")) + + assert(orderingOf(analyze("CREATE TABLE testcat.t (ID INT) USING foo " + + "ORDERED BY truncate(4, ID)")) === Seq("truncate(4, ID) ASC NULLS FIRST")) + val e = intercept[AnalysisException] { + analyze("CREATE TABLE testcat.t (ID INT) USING foo ORDERED BY truncate(4, id)") + } + assert(e.getCondition === "WRITE_ORDERING_WITH_UNKNOWN_COLUMN") + assert(e.getMessageParameters.get("cols") === "`id`") + + // references are normalized independently, so only the unresolvable one is reported + val multi = intercept[AnalysisException] { + analyze("CREATE TABLE testcat.t (ID INT) USING foo ORDERED BY bucket(4, id, nope)") + } + assert(multi.getMessageParameters.get("cols") === "`nope`") + } + } + + test("REPLACE TABLE and RTAS normalize the ordering too") { + withSQLConf("spark.sql.catalog.testcat" -> classOf[RecordingInMemoryTableCatalog].getName) { + assert(orderingOf(analyze("REPLACE TABLE testcat.t (ID INT) USING foo ORDERED BY id")) === + Seq("ID ASC NULLS FIRST")) + assert(orderingOf(analyze( + "REPLACE TABLE testcat.t USING foo ORDERED BY id AS SELECT 1 AS ID")) === + Seq("ID ASC NULLS FIRST")) + } + } + + test("a case-sensitive session resolves the ordering case-sensitively") { + withSQLConf( + "spark.sql.catalog.testcat" -> classOf[RecordingInMemoryTableCatalog].getName, + SQLConf.CASE_SENSITIVE.key -> "true") { + assert(orderingOf(analyze("CREATE TABLE testcat.t (ID INT) USING foo ORDERED BY ID")) === + Seq("ID ASC NULLS FIRST")) + val e = intercept[AnalysisException] { + analyze("CREATE TABLE testcat.t (ID INT) USING foo ORDERED BY id") + } + assert(e.getCondition === "WRITE_ORDERING_WITH_UNKNOWN_COLUMN") + assert(e.getMessageParameters.get("cols") === "`id`") + } + } + + test("a case-insensitive session does not normalize an ApplyTransform's references") { + withSQLConf("spark.sql.catalog.testcat" -> classOf[RecordingInMemoryTableCatalog].getName) { + assert(orderingOf(analyze("CREATE TABLE testcat.t (id INT) USING foo ORDERED BY ID")) === + Seq("id ASC NULLS FIRST")) + + val e = intercept[AnalysisException] { + analyze("CREATE TABLE testcat.t (id INT) USING foo ORDERED BY truncate(4, ID)") + } + assert(e.getCondition === "WRITE_ORDERING_WITH_UNKNOWN_COLUMN") + assert(e.getMessageParameters.get("cols") === "`ID`") + + val p = intercept[AnalysisException] { + analyze("CREATE TABLE testcat.t (id INT) USING foo PARTITIONED BY (truncate(4, ID))") + } + assert(p.getCondition === "UNSUPPORTED_FEATURE.PARTITION_WITH_NESTED_COLUMN_IS_UNSUPPORTED") + } + } + + test("ORDERED BY needs a schema to resolve against") { + withSQLConf("spark.sql.catalog.testcat" -> classOf[RecordingInMemoryTableCatalog].getName) { + Seq("ORDERED BY id", "LOCALLY ORDERED BY id").foreach { clause => + val e = intercept[AnalysisException] { + analyze(s"CREATE TABLE testcat.t USING foo $clause") + } + assert(e.getCondition === "SPECIFY_WRITE_ORDERING_IS_NOT_ALLOWED", s"for $clause") + } + + // UNORDERED asks for no ordering at all, so it has nothing to resolve and stays allowed + analyze("CREATE TABLE testcat.t USING foo UNORDERED") + } + } + + test("an ordering on an unknown column is rejected during analysis") { + withSQLConf("spark.sql.catalog.testcat" -> classOf[RecordingInMemoryTableCatalog].getName) { + Seq( + "CREATE TABLE testcat.t (id INT) USING foo ORDERED BY nope", + "CREATE TABLE testcat.t (id INT) USING foo ORDERED BY days(nope)", + "CREATE TABLE testcat.t (id INT) USING foo ORDERED BY truncate(4, nope)", + "CREATE TABLE testcat.t USING foo ORDERED BY nope AS SELECT 1 AS id" + ).foreach { stmt => + val e = intercept[AnalysisException](analyze(stmt)) + assert(e.getCondition === "WRITE_ORDERING_WITH_UNKNOWN_COLUMN", stmt) + assert(e.getMessageParameters.get("cols") === "`nope`", stmt) + } + + // a nested struct field resolves + analyze("CREATE TABLE testcat.t (p STRUCT<x: INT>) USING foo ORDERED BY p.x") + } + } + + test("a repeated ordering column is accepted, unlike a repeated partition column") { + withSQLConf("spark.sql.catalog.testcat" -> classOf[RecordingInMemoryTableCatalog].getName) { + assert(orderingOf(analyze("CREATE TABLE testcat.t (id INT) USING foo " + + "ORDERED BY bucket(4, id), bucket(8, id)")) === + Seq("bucket(4, id) ASC NULLS FIRST", "bucket(8, id) ASC NULLS FIRST")) + analyze("CREATE TABLE testcat.t (id INT) USING foo ORDERED BY id, id DESC") + analyze("CREATE TABLE testcat.t (ts TIMESTAMP) USING foo ORDERED BY days(ts), hours(ts)") + } + } + + test("DISTRIBUTED BY PARTITION is not satisfied by CLUSTER BY") { + val e = intercept[ParseException] { + parse("CREATE TABLE t (id INT, c STRING) USING foo CLUSTER BY (c) DISTRIBUTED BY PARTITION") + } + assert(e.getCondition === + "SPECIFY_DISTRIBUTED_BY_PARTITION_WITHOUT_PARTITIONING_IS_NOT_ALLOWED") + + parse("CREATE TABLE t (id INT, c STRING) USING foo CLUSTERED BY (c) INTO 4 BUCKETS " + + "DISTRIBUTED BY PARTITION") match { + case c: CreateTable => assert(c.writeDistributionMode === "hash") + case other => fail(s"unexpected plan: $other") + } + } + + test("CLUSTER BY and the write clauses both reach the catalog") { Review Comment: Could we complement this successful pass-through test with a catalog that advertises `SUPPORTS_CREATE_TABLE_WITH_WRITE_DISTRIBUTION_AND_ORDERING` but deliberately rejects one combination, for example `CLUSTER BY (a) UNORDERED`? This would exercise the connector-rejection contract separately from the existing tests where the catalog lacks the entire capability. For CREATE, assert that the catalog's error reaches the caller and no table is published. For staged REPLACE/RTAS, start with a populated table and assert that its rows, schema, and clustering metadata survive the rejection. The unchanged-table assertion should be scoped to staging: non-staging `ReplaceTableExec` drops the original before calling `createTable`, so late catalog rejection there retains the existing non-atomic replacement limitation. It would also be useful to include `CLUSTER BY (a) LOCALLY ORDERED BY (b)` in this positive test, asserting mode `none` with the ordering retained. ########## docs/sql-ref-syntax-ddl-create-table-datasource.md: ########## @@ -98,6 +106,64 @@ as any order. For example, you can write COMMENT table_comment after TBLPROPERTI A list of key-value pairs that is used to tag the table definition. +* **DISTRIBUTED BY PARTITION** + + Requests that every write to the table be clustered by the table's partitioning, so each + partition is written by a single task rather than by every task that holds rows for it. + + Requires the table to actually be partitioned, with `PARTITIONED BY` or with + `CLUSTERED BY ... INTO ... BUCKETS`. Note that `CLUSTER BY` -- a different clause from + `CLUSTERED BY ... INTO ... BUCKETS` -- does **not** qualify: it lists clustering columns for the + data source to interpret rather than defining a partitioning, and it cannot be combined with + `PARTITIONED BY` or `CLUSTERED BY ... INTO ... BUCKETS`, so a table using it has no partitioning + to distribute by. + +* **ORDERED BY** + + Requests a sort order for every write to the table, recorded on the table so that later writes + honor it too. `UNORDERED` asks for no ordering at all, which is different from omitting the + clause -- omitting it leaves the choice to the catalog. The parentheses are optional: + `ORDERED BY (a, b)` and `ORDERED BY a, b` are the same. The sort keys must resolve against the + table's columns, so a `CREATE TABLE` with neither a column list nor `AS SELECT` cannot use this + clause. + + The distribution decides how far the order reaches, and this clause picks one when + `DISTRIBUTED BY PARTITION` is absent: a bare `ORDERED BY` range-partitions each write, so the + order holds across the tasks of a write, not only within one, while `LOCALLY ORDERED BY` asks + for it to hold within each write task only, without a shuffle. `UNORDERED` on its own asks for + no distribution either. + + When `DISTRIBUTED BY PARTITION` is given it decides the distribution instead, and the order then + holds within each write task. `LOCALLY` therefore adds nothing beside it, and `UNORDERED` beside + it contributes only "no sort keys": + + ```sql + -- range-partition each write by id, so the order holds across that write's tasks + CREATE TABLE t (id INT, c STRING) USING iceberg PARTITIONED BY (c) ORDERED BY (id); + + -- cluster each write by partition instead, and sort by id within each task + CREATE TABLE t (id INT, c STRING) USING iceberg PARTITIONED BY (c) + DISTRIBUTED BY PARTITION ORDERED BY (id); + + -- cluster each write by partition, with no sort order; UNORDERED here is the same as + -- omitting it, since DISTRIBUTED BY PARTITION already fixed the distribution + CREATE TABLE t (id INT, c STRING) USING iceberg PARTITIONED BY (c) + DISTRIBUTED BY PARTITION UNORDERED; + ``` + + Both clauses are passed to the catalog, which has to support them: a catalog that does not + advertise support for a write distribution and ordering rejects the statement rather than + creating a table that silently lacks the requested layout. The built-in catalogs do not + support them. Both may also be combined with `CLUSTER BY`. Spark passes the clustering columns Review Comment: `Both may also be combined with CLUSTER BY` appears to include `DISTRIBUTED BY PARTITION`, but that combination is rejected by the parser and by the test at `CreateTableWriteOrderSuite:311`. Could we name the allowed forms explicitly here: `ORDERED BY`, `LOCALLY ORDERED BY`, and `UNORDERED` may be combined with `CLUSTER BY`, subject to the catalog accepting the combination? ########## sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ShowCreateTableExec.scala: ########## @@ -125,6 +127,53 @@ case class ShowCreateTableExec( } } + /** + * True for a sort key expression the `transformArgument` grammar rule can represent: a plain + * column reference, or a transform whose name is an identifier and whose own arguments are + * references or constants. The name check approximates `identifier`: a reserved word such as + * `select` passes it but fails to parse under ANSI. + */ + private def isSpellable(e: V2Expression): Boolean = e match { + case _: NamedReference => true + case t: Transform => + t.name().matches("[a-zA-Z_][a-zA-Z0-9_]*") && t.arguments().nonEmpty && + t.arguments().forall(a => a.isInstanceOf[NamedReference] || a.isInstanceOf[Literal[_]]) + case _ => false + } + + /** + * Emits the table's declared write distribution and ordering as clauses. Pairs with no clause + * form, such as `hash` without partitioning or an unspellable sort key, are omitted. + */ + private def showTableWriteDistributionAndOrdering( + table: Table, + builder: StringBuilder): Unit = { + if (table.writeOrdering().forall(o => isSpellable(o.expression()))) { + val orderBy = if (table.writeOrdering().nonEmpty) { + Some(table.writeOrdering() + .map(WriteDistributionAndOrdering.describeSortOrder) + .mkString("ORDERED BY (", ", ", ")")) + } else { + None + } + // Bucketing counts as partitioning here; CLUSTER BY does not. + val hasPartitioning = table.partitioning.exists(!_.isInstanceOf[ClusterByTransform]) Review Comment: [P2] Could we use the `ClusterByTransform(_)` extractor here instead of an implementation-class check? A connector can return `Expressions.apply("cluster_by", Expressions.column("a"))` (or its own `Transform` implementation). The existing extractor recognizes that as clustering, but `isInstanceOf[ClusterByTransform]` is false, so this guard treats it as actual partitioning. If the table reports declared mode `hash`, SHOW CREATE TABLE then emits `DISTRIBUTED BY PARTITION` for a clustering-only table, although this method intends to omit that pair. A focused regression test could stub `Table.partitioning()` with this generic transform and report mode `hash`, with both empty and non-empty write ordering. Assert that neither `DISTRIBUTED BY PARTITION` nor a standalone `ORDERED BY` is emitted for that unrepresentable pair. Keep an actual partition or bucket transform as a positive control. `TransformExtractorSuite` already covers recognition of a connector-defined `cluster_by` transform. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
