anuragmantri commented on code in PR #58153: URL: https://github.com/apache/spark/pull/58153#discussion_r4213528940
########## sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/WriteDistributionAndOrdering.scala: ########## @@ -0,0 +1,190 @@ +/* + * 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.catalyst.util + +import java.time.ZoneOffset + +import scala.util.Try +import scala.util.control.NonFatal + +import org.apache.spark.sql.catalyst.expressions.{Literal => CatalystLiteral} +import org.apache.spark.sql.connector.catalog.{Identifier, TableCatalog, TableCatalogCapability, WriteDistributionMode} +import org.apache.spark.sql.connector.expressions.{Cast, Expression, GeneralScalarExpression, IdentityTransform, Literal, NamedReference, SortOrder, Transform} +import org.apache.spark.sql.errors.QueryCompilationErrors +import org.apache.spark.sql.types.{FloatType, StructType, TimestampType} +import org.apache.spark.util.ArrayImplicits._ + +/** + * Utility methods for the create-time write distribution and ordering requested by + * `CREATE`/`REPLACE TABLE ... DISTRIBUTED BY PARTITION / [LOCALLY] ORDERED BY`. + */ +object WriteDistributionAndOrdering { + + /** Names the write clauses in an error's operation, as in `CREATE TABLE ... <CLAUSES>`. */ + val CLAUSES = "DISTRIBUTED BY/ORDERED BY/UNORDERED" + + /** + * True when a CREATE/REPLACE TABLE statement asked for a write distribution or ordering. + * `UNORDERED` counts as a request. + */ + def isRequested( + writeDistributionMode: WriteDistributionMode, + writeOrdering: Seq[SortOrder]): Boolean = { + writeDistributionMode != null || writeOrdering.nonEmpty + } + + /** + * True when the transforms define a partitioning that `DISTRIBUTED BY PARTITION` can distribute + * by. Bucketing counts; a `cluster_by` transform, from CLUSTER BY or reported by a connector, + * does not. Matches by name only, since a connector's `cluster_by` may take any arguments. + */ + def hasPartitioning(partitioning: Seq[Transform]): Boolean = { + partitioning.exists(_.name != "cluster_by") + } + + /** + * Rejects a create-time write distribution or ordering that the catalog has not advertised + * support for, before anything is created or dropped. + */ + def validateCatalogForWriteDistributionAndOrdering( + catalog: TableCatalog, + ident: Identifier, + operation: String, + writeDistributionMode: WriteDistributionMode, + writeOrdering: Seq[SortOrder]): Unit = { + if (isRequested(writeDistributionMode, writeOrdering) && + !catalog.capabilities().contains( + TableCatalogCapability.SUPPORTS_CREATE_TABLE_WITH_WRITE_DISTRIBUTION_AND_ORDERING)) { + throw QueryCompilationErrors.unsupportedTableOperationError( + catalog, ident, s"$operation ... $CLAUSES") + } + } + + /** Renders a requested sort key as SQL, for DESCRIBE. */ + def describeSortOrder(sortOrder: SortOrder): String = sortOrderToSQL(sortOrder, quoteIfNeeded) + + /** + * Renders the declared write distribution and ordering as `DISTRIBUTED BY PARTITION`, + * `[LOCALLY] ORDERED BY` or `UNORDERED` clauses for SHOW CREATE TABLE. Returns None when the + * pair has no clause form, such as `HASH` without partitioning, or when `replay`, which parses a + * statement carrying the clauses and returns the pair it declares, does not give back the same + * pair. Names are quoted only where needed, and on a second try all of them, which reserved + * keywords under `spark.sql.ansi.enforceReservedKeywords` need. + */ + def writeClausesSQL( + writeDistributionMode: WriteDistributionMode, + writeOrdering: Seq[SortOrder], + partitioning: Seq[Transform], + replay: String => Option[(WriteDistributionMode, Seq[SortOrder])]): Option[String] = { + lazy val partitioned = hasPartitioning(partitioning) + Seq[String => String](quoteIfNeeded, quoteIdentifier).iterator.flatMap { quote => + val orderBy = if (writeOrdering.nonEmpty) { + Some(writeOrdering.map(sortOrderToSQL(_, quote)).mkString("ORDERED BY (", ", ", ")")) + } else { + None + } + (writeDistributionMode, orderBy) match { + case (WriteDistributionMode.HASH, Some(o)) if partitioned => + Some(s"DISTRIBUTED BY PARTITION $o") + case (WriteDistributionMode.HASH, None) if partitioned => Some("DISTRIBUTED BY PARTITION") + case (WriteDistributionMode.RANGE, Some(o)) => Some(o) + case (WriteDistributionMode.NONE, Some(o)) => Some(s"LOCALLY $o") + case (WriteDistributionMode.NONE, None) => Some("UNORDERED") + case _ => None + } + }.find { clauses => + replay(clauses).exists { case (mode, ordering) => + mode == writeDistributionMode && sameSortOrders(ordering, writeOrdering) + } + } + } + + /** + * True when every column a sort key references is in `schema`, matched as `CheckAnalysis` + * matches the ordering of a CREATE/REPLACE TABLE statement. + */ + def referencesExist(schema: StructType, writeOrdering: Seq[SortOrder]): Boolean = { + writeOrdering.flatMap(_.expression().references()).forall { ref => Review Comment: Good catch. The column check now runs on the replayed ordering, which is what CheckAnalysis sees, so `f(identity(missing))` is omitted. A nested `identity(...)` is also no longer compared as a plain column: only a key's top level counts as `identity(col)`, so `f(identity(id))` is omitted too. Both are in the unit suite and in `ORDERING_OVERRIDE` with omit-and-replay rows. ########## sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/WriteDistributionAndOrdering.scala: ########## @@ -0,0 +1,190 @@ +/* + * 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.catalyst.util + +import java.time.ZoneOffset + +import scala.util.Try +import scala.util.control.NonFatal + +import org.apache.spark.sql.catalyst.expressions.{Literal => CatalystLiteral} +import org.apache.spark.sql.connector.catalog.{Identifier, TableCatalog, TableCatalogCapability, WriteDistributionMode} +import org.apache.spark.sql.connector.expressions.{Cast, Expression, GeneralScalarExpression, IdentityTransform, Literal, NamedReference, SortOrder, Transform} +import org.apache.spark.sql.errors.QueryCompilationErrors +import org.apache.spark.sql.types.{FloatType, StructType, TimestampType} +import org.apache.spark.util.ArrayImplicits._ + +/** + * Utility methods for the create-time write distribution and ordering requested by + * `CREATE`/`REPLACE TABLE ... DISTRIBUTED BY PARTITION / [LOCALLY] ORDERED BY`. + */ +object WriteDistributionAndOrdering { + + /** Names the write clauses in an error's operation, as in `CREATE TABLE ... <CLAUSES>`. */ + val CLAUSES = "DISTRIBUTED BY/ORDERED BY/UNORDERED" + + /** + * True when a CREATE/REPLACE TABLE statement asked for a write distribution or ordering. + * `UNORDERED` counts as a request. + */ + def isRequested( + writeDistributionMode: WriteDistributionMode, + writeOrdering: Seq[SortOrder]): Boolean = { + writeDistributionMode != null || writeOrdering.nonEmpty + } + + /** + * True when the transforms define a partitioning that `DISTRIBUTED BY PARTITION` can distribute + * by. Bucketing counts; a `cluster_by` transform, from CLUSTER BY or reported by a connector, + * does not. Matches by name only, since a connector's `cluster_by` may take any arguments. + */ + def hasPartitioning(partitioning: Seq[Transform]): Boolean = { + partitioning.exists(_.name != "cluster_by") + } + + /** + * Rejects a create-time write distribution or ordering that the catalog has not advertised + * support for, before anything is created or dropped. + */ + def validateCatalogForWriteDistributionAndOrdering( + catalog: TableCatalog, + ident: Identifier, + operation: String, + writeDistributionMode: WriteDistributionMode, + writeOrdering: Seq[SortOrder]): Unit = { + if (isRequested(writeDistributionMode, writeOrdering) && + !catalog.capabilities().contains( + TableCatalogCapability.SUPPORTS_CREATE_TABLE_WITH_WRITE_DISTRIBUTION_AND_ORDERING)) { + throw QueryCompilationErrors.unsupportedTableOperationError( + catalog, ident, s"$operation ... $CLAUSES") + } + } + + /** Renders a requested sort key as SQL, for DESCRIBE. */ + def describeSortOrder(sortOrder: SortOrder): String = sortOrderToSQL(sortOrder, quoteIfNeeded) + + /** + * Renders the declared write distribution and ordering as `DISTRIBUTED BY PARTITION`, + * `[LOCALLY] ORDERED BY` or `UNORDERED` clauses for SHOW CREATE TABLE. Returns None when the + * pair has no clause form, such as `HASH` without partitioning, or when `replay`, which parses a + * statement carrying the clauses and returns the pair it declares, does not give back the same + * pair. Names are quoted only where needed, and on a second try all of them, which reserved + * keywords under `spark.sql.ansi.enforceReservedKeywords` need. + */ + def writeClausesSQL( + writeDistributionMode: WriteDistributionMode, + writeOrdering: Seq[SortOrder], + partitioning: Seq[Transform], + replay: String => Option[(WriteDistributionMode, Seq[SortOrder])]): Option[String] = { + lazy val partitioned = hasPartitioning(partitioning) + Seq[String => String](quoteIfNeeded, quoteIdentifier).iterator.flatMap { quote => + val orderBy = if (writeOrdering.nonEmpty) { + Some(writeOrdering.map(sortOrderToSQL(_, quote)).mkString("ORDERED BY (", ", ", ")")) + } else { + None + } + (writeDistributionMode, orderBy) match { + case (WriteDistributionMode.HASH, Some(o)) if partitioned => + Some(s"DISTRIBUTED BY PARTITION $o") + case (WriteDistributionMode.HASH, None) if partitioned => Some("DISTRIBUTED BY PARTITION") + case (WriteDistributionMode.RANGE, Some(o)) => Some(o) + case (WriteDistributionMode.NONE, Some(o)) => Some(s"LOCALLY $o") + case (WriteDistributionMode.NONE, None) => Some("UNORDERED") + case _ => None + } + }.find { clauses => + replay(clauses).exists { case (mode, ordering) => + mode == writeDistributionMode && sameSortOrders(ordering, writeOrdering) + } + } + } + + /** + * True when every column a sort key references is in `schema`, matched as `CheckAnalysis` + * matches the ordering of a CREATE/REPLACE TABLE statement. + */ + def referencesExist(schema: StructType, writeOrdering: Seq[SortOrder]): Boolean = { + writeOrdering.flatMap(_.expression().references()).forall { ref => + ref.fieldNames().nonEmpty && + Try(schema.findNestedField(ref.fieldNames().toImmutableArraySeq)).toOption.flatten.isDefined + } + } + + /** + * True when two orderings declare the same keys: the same column names, transform names and + * argument structure, literal values and types, directions and null orderings. + */ + def sameSortOrders(left: Seq[SortOrder], right: Seq[SortOrder]): Boolean = { + left.length == right.length && left.zip(right).forall { case (l, r) => + l.direction() == r.direction() && l.nullOrdering() == r.nullOrdering() && + keyOf(l.expression()) == keyOf(r.expression()) + } + } + + private def keyOf(e: Expression): Any = e match { + case r: NamedReference => r.fieldNames().toSeq + case l: Literal[_] => Try(CatalystLiteral.create(l.value, l.dataType)).toOption + case t: IdentityTransform => keyOf(t.ref) + case t: Transform => (t.name, t.arguments().toSeq.map(keyOf)) + case other => other + } + + private def sortOrderToSQL(sortOrder: SortOrder, quote: String => String): String = { + s"${toSQL(sortOrder.expression(), quote)} ${sortOrder.direction()} ${sortOrder.nullOrdering()}" + } + + // `describe` prints a literal's internal value, e.g. `0` for DATE '1970-01-01', so literals are + // rendered through Catalyst to keep their type. Catalyst renders a FLOAT as a CAST, which the + // parser does not accept here, so a finite FLOAT is rendered as `<v>F` instead, and a TIMESTAMP + // in UTC with an explicit offset so that it does not depend on the session time zone or + // timestamp type. A literal Catalyst cannot represent falls back to `describe`. Other connector + // expressions are rendered from their children: `describe` on one that + // `V2ExpressionSQLBuilder` does not know recurses without end. PARTITIONED BY in SHOW CREATE + // TABLE and the `Part N` rows of DESCRIBE render transforms with `describe`, so the same + // transform can print differently there. + private def toSQL(e: Expression, quote: String => String): String = e match { + case r: NamedReference => r.fieldNames.map(quote).mkString(".") + case l: Literal[_] => + try literalToSQL(CatalystLiteral.create(l.value, l.dataType)) catch { + case NonFatal(_) => describeLiteral(l) + } + case t: IdentityTransform => toSQL(t.ref, quote) + case t: Transform => + t.arguments().map(toSQL(_, quote)).mkString(s"${quote(t.name)}(", ", ", ")") + case c: Cast => s"CAST(${toSQL(c.expression(), quote)} AS ${c.dataType().sql})" + case g: GeneralScalarExpression => + g.children().map(toSQL(_, quote)).mkString(s"${g.name}(", ", ", ")") + case other => + other.children().map(toSQL(_, quote)).mkString(s"${other.getClass.getSimpleName}(", ", ", ")") + } + + // `LiteralValue.describe` also rejects a value that does not match its type. + private def describeLiteral(l: Literal[_]): String = { + try l.describe catch { + case NonFatal(_) => String.valueOf(l.value) + } + } + + private def literalToSQL(l: CatalystLiteral): String = (l.value, l.dataType) match { + case (f: Float, FloatType) if java.lang.Float.isFinite(f) => s"${f}F" + case (micros: Long, TimestampType) => Review Comment: Done. A `TimestampLTZNanosType` literal now renders in UTC with a `Z` suffix, padded to its precision. A test checks `TIMESTAMP_LTZ '2020-11-01 01:30:00.123456789-08:00'` in America/Los_Angeles with `spark.sql.timestampNanosTypes.enabled`, and a nanosecond TIMESTAMP_NTZ row was added too. ########## sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/TableInfo.java: ########## @@ -65,9 +72,45 @@ public Transform[] partitions() { public Constraint[] constraints() { return constraints; } + /** + * The requested write distribution, or null when the statement did not ask for one, which + * leaves the choice to the catalog. + * <p> + * Only catalogs that report + * {@link TableCatalogCapability#SUPPORTS_CREATE_TABLE_WITH_WRITE_DISTRIBUTION_AND_ORDERING} see a + * request; for any other catalog, Spark rejects the statement. + * <p> + * Spark passes the request only to the {@code TableInfo} overloads of + * {@link TableCatalog#createTable(Identifier, TableInfo)} and of the {@link StagingTableCatalog} + * methods {@code stageCreate}, {@code stageReplace} and {@code stageCreateOrReplace}. Their + * default implementations drop it, so a catalog that reports the capability must override each + * one it can be reached through. {@link TableCatalog#createTableLike} never gets a request: Spark + * does not copy the source table's write distribution and ordering into its {@code TableInfo}. + * + * @since 4.4.0 + */ + public WriteDistributionMode writeDistributionMode() { return writeDistributionMode; } + + /** + * The requested write ordering; never null, and empty when none was requested. Gated on the + * same capability and delivered the same way as {@link #writeDistributionMode()}. + * <p> + * A plain column is a {@link org.apache.spark.sql.connector.expressions.NamedReference}; any + * other key is a {@link Transform}, such as {@code bucket(16, id)}. Spark checks that each + * referenced column exists in the table schema, and the parser checks the arguments of Review Comment: Thanks. It now says that the parser checks the argument kinds of the lowercase `bucket`, `years`, `months`, `days` and `hours`, and that Spark does not check orderability, a positive bucket count, a bucket's column list, or the arguments of any other transform. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/rules.scala: ########## @@ -323,10 +323,17 @@ case class PreprocessTableCreation(catalog: SessionCatalog) extends Rule[Logical SchemaUtils.checkTransformDuplication( partitioning, "in the partitioning", isCaseSensitive) + // The write ordering has no duplicate check: a repeated sort key is redundant, and + // `bucket(4, id), bucket(8, id)` are distinct keys. + if (schema.isEmpty) { if (partitioning.nonEmpty) { throw QueryCompilationErrors.specifyPartitionNotAllowedWhenTableSchemaNotDefinedError() } + if (create.writeOrdering.nonEmpty) { + throw QueryCompilationErrors Review Comment: Done. The error now takes `create.origin`, and the test pins the statement as its context. ########## sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/WriteDistributionAndOrdering.scala: ########## @@ -0,0 +1,190 @@ +/* + * 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.catalyst.util + +import java.time.ZoneOffset + +import scala.util.Try +import scala.util.control.NonFatal + +import org.apache.spark.sql.catalyst.expressions.{Literal => CatalystLiteral} +import org.apache.spark.sql.connector.catalog.{Identifier, TableCatalog, TableCatalogCapability, WriteDistributionMode} +import org.apache.spark.sql.connector.expressions.{Cast, Expression, GeneralScalarExpression, IdentityTransform, Literal, NamedReference, SortOrder, Transform} +import org.apache.spark.sql.errors.QueryCompilationErrors +import org.apache.spark.sql.types.{FloatType, StructType, TimestampType} +import org.apache.spark.util.ArrayImplicits._ + +/** + * Utility methods for the create-time write distribution and ordering requested by + * `CREATE`/`REPLACE TABLE ... DISTRIBUTED BY PARTITION / [LOCALLY] ORDERED BY`. + */ +object WriteDistributionAndOrdering { + + /** Names the write clauses in an error's operation, as in `CREATE TABLE ... <CLAUSES>`. */ + val CLAUSES = "DISTRIBUTED BY/ORDERED BY/UNORDERED" + + /** + * True when a CREATE/REPLACE TABLE statement asked for a write distribution or ordering. + * `UNORDERED` counts as a request. + */ + def isRequested( + writeDistributionMode: WriteDistributionMode, + writeOrdering: Seq[SortOrder]): Boolean = { + writeDistributionMode != null || writeOrdering.nonEmpty + } + + /** + * True when the transforms define a partitioning that `DISTRIBUTED BY PARTITION` can distribute + * by. Bucketing counts; a `cluster_by` transform, from CLUSTER BY or reported by a connector, + * does not. Matches by name only, since a connector's `cluster_by` may take any arguments. + */ + def hasPartitioning(partitioning: Seq[Transform]): Boolean = { + partitioning.exists(_.name != "cluster_by") + } + + /** + * Rejects a create-time write distribution or ordering that the catalog has not advertised + * support for, before anything is created or dropped. + */ + def validateCatalogForWriteDistributionAndOrdering( + catalog: TableCatalog, + ident: Identifier, + operation: String, + writeDistributionMode: WriteDistributionMode, + writeOrdering: Seq[SortOrder]): Unit = { + if (isRequested(writeDistributionMode, writeOrdering) && + !catalog.capabilities().contains( + TableCatalogCapability.SUPPORTS_CREATE_TABLE_WITH_WRITE_DISTRIBUTION_AND_ORDERING)) { + throw QueryCompilationErrors.unsupportedTableOperationError( + catalog, ident, s"$operation ... $CLAUSES") + } + } + + /** Renders a requested sort key as SQL, for DESCRIBE. */ + def describeSortOrder(sortOrder: SortOrder): String = sortOrderToSQL(sortOrder, quoteIfNeeded) + + /** + * Renders the declared write distribution and ordering as `DISTRIBUTED BY PARTITION`, + * `[LOCALLY] ORDERED BY` or `UNORDERED` clauses for SHOW CREATE TABLE. Returns None when the + * pair has no clause form, such as `HASH` without partitioning, or when `replay`, which parses a + * statement carrying the clauses and returns the pair it declares, does not give back the same + * pair. Names are quoted only where needed, and on a second try all of them, which reserved + * keywords under `spark.sql.ansi.enforceReservedKeywords` need. + */ + def writeClausesSQL( + writeDistributionMode: WriteDistributionMode, + writeOrdering: Seq[SortOrder], + partitioning: Seq[Transform], + replay: String => Option[(WriteDistributionMode, Seq[SortOrder])]): Option[String] = { + lazy val partitioned = hasPartitioning(partitioning) + Seq[String => String](quoteIfNeeded, quoteIdentifier).iterator.flatMap { quote => + val orderBy = if (writeOrdering.nonEmpty) { + Some(writeOrdering.map(sortOrderToSQL(_, quote)).mkString("ORDERED BY (", ", ", ")")) + } else { + None + } + (writeDistributionMode, orderBy) match { + case (WriteDistributionMode.HASH, Some(o)) if partitioned => + Some(s"DISTRIBUTED BY PARTITION $o") + case (WriteDistributionMode.HASH, None) if partitioned => Some("DISTRIBUTED BY PARTITION") + case (WriteDistributionMode.RANGE, Some(o)) => Some(o) + case (WriteDistributionMode.NONE, Some(o)) => Some(s"LOCALLY $o") + case (WriteDistributionMode.NONE, None) => Some("UNORDERED") + case _ => None + } + }.find { clauses => + replay(clauses).exists { case (mode, ordering) => + mode == writeDistributionMode && sameSortOrders(ordering, writeOrdering) + } + } + } + + /** + * True when every column a sort key references is in `schema`, matched as `CheckAnalysis` + * matches the ordering of a CREATE/REPLACE TABLE statement. + */ + def referencesExist(schema: StructType, writeOrdering: Seq[SortOrder]): Boolean = { + writeOrdering.flatMap(_.expression().references()).forall { ref => + ref.fieldNames().nonEmpty && + Try(schema.findNestedField(ref.fieldNames().toImmutableArraySeq)).toOption.flatten.isDefined + } + } + + /** + * True when two orderings declare the same keys: the same column names, transform names and + * argument structure, literal values and types, directions and null orderings. + */ + def sameSortOrders(left: Seq[SortOrder], right: Seq[SortOrder]): Boolean = { + left.length == right.length && left.zip(right).forall { case (l, r) => + l.direction() == r.direction() && l.nullOrdering() == r.nullOrdering() && + keyOf(l.expression()) == keyOf(r.expression()) + } + } + + private def keyOf(e: Expression): Any = e match { + case r: NamedReference => r.fieldNames().toSeq + case l: Literal[_] => Try(CatalystLiteral.create(l.value, l.dataType)).toOption + case t: IdentityTransform => keyOf(t.ref) + case t: Transform => (t.name, t.arguments().toSeq.map(keyOf)) + case other => other + } + + private def sortOrderToSQL(sortOrder: SortOrder, quote: String => String): String = { + s"${toSQL(sortOrder.expression(), quote)} ${sortOrder.direction()} ${sortOrder.nullOrdering()}" + } + + // `describe` prints a literal's internal value, e.g. `0` for DATE '1970-01-01', so literals are + // rendered through Catalyst to keep their type. Catalyst renders a FLOAT as a CAST, which the + // parser does not accept here, so a finite FLOAT is rendered as `<v>F` instead, and a TIMESTAMP + // in UTC with an explicit offset so that it does not depend on the session time zone or + // timestamp type. A literal Catalyst cannot represent falls back to `describe`. Other connector + // expressions are rendered from their children: `describe` on one that + // `V2ExpressionSQLBuilder` does not know recurses without end. PARTITIONED BY in SHOW CREATE + // TABLE and the `Part N` rows of DESCRIBE render transforms with `describe`, so the same + // transform can print differently there. + private def toSQL(e: Expression, quote: String => String): String = e match { + case r: NamedReference => r.fieldNames.map(quote).mkString(".") + case l: Literal[_] => + try literalToSQL(CatalystLiteral.create(l.value, l.dataType)) catch { + case NonFatal(_) => describeLiteral(l) + } + case t: IdentityTransform => toSQL(t.ref, quote) + case t: Transform => + t.arguments().map(toSQL(_, quote)).mkString(s"${quote(t.name)}(", ", ", ")") + case c: Cast => s"CAST(${toSQL(c.expression(), quote)} AS ${c.dataType().sql})" + case g: GeneralScalarExpression => + g.children().map(toSQL(_, quote)).mkString(s"${g.name}(", ", ", ")") + case other => + other.children().map(toSQL(_, quote)).mkString(s"${other.getClass.getSimpleName}(", ", ", ")") Review Comment: Thanks. DESCRIBE renders connector expressions with `ToStringSQLBuilder` again, so `EXTRACT(YEAR FROM ts)`, `my_udf(id)` and `id + 1` print as `describe` printed them. The builder is a subclass that renders literals and references with their types and quoting, and renders `GetArrayItem` and `VariantGet` through their children rather than through `toString`. It renders an expression it does not know as `name(children)` instead of throwing, so a known parent keeps its name, e.g. `my_udf(DATE_TRUNC(id))`. Tests cover these cases, plus a sort order and a null child nested in a key. ########## sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/WriteDistributionAndOrdering.scala: ########## @@ -0,0 +1,190 @@ +/* + * 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.catalyst.util + +import java.time.ZoneOffset + +import scala.util.Try +import scala.util.control.NonFatal + +import org.apache.spark.sql.catalyst.expressions.{Literal => CatalystLiteral} +import org.apache.spark.sql.connector.catalog.{Identifier, TableCatalog, TableCatalogCapability, WriteDistributionMode} +import org.apache.spark.sql.connector.expressions.{Cast, Expression, GeneralScalarExpression, IdentityTransform, Literal, NamedReference, SortOrder, Transform} +import org.apache.spark.sql.errors.QueryCompilationErrors +import org.apache.spark.sql.types.{FloatType, StructType, TimestampType} +import org.apache.spark.util.ArrayImplicits._ + +/** + * Utility methods for the create-time write distribution and ordering requested by + * `CREATE`/`REPLACE TABLE ... DISTRIBUTED BY PARTITION / [LOCALLY] ORDERED BY`. + */ +object WriteDistributionAndOrdering { + + /** Names the write clauses in an error's operation, as in `CREATE TABLE ... <CLAUSES>`. */ + val CLAUSES = "DISTRIBUTED BY/ORDERED BY/UNORDERED" + + /** + * True when a CREATE/REPLACE TABLE statement asked for a write distribution or ordering. + * `UNORDERED` counts as a request. + */ + def isRequested( + writeDistributionMode: WriteDistributionMode, + writeOrdering: Seq[SortOrder]): Boolean = { + writeDistributionMode != null || writeOrdering.nonEmpty + } + + /** + * True when the transforms define a partitioning that `DISTRIBUTED BY PARTITION` can distribute + * by. Bucketing counts; a `cluster_by` transform, from CLUSTER BY or reported by a connector, + * does not. Matches by name only, since a connector's `cluster_by` may take any arguments. + */ + def hasPartitioning(partitioning: Seq[Transform]): Boolean = { + partitioning.exists(_.name != "cluster_by") + } + + /** + * Rejects a create-time write distribution or ordering that the catalog has not advertised + * support for, before anything is created or dropped. + */ + def validateCatalogForWriteDistributionAndOrdering( + catalog: TableCatalog, + ident: Identifier, + operation: String, + writeDistributionMode: WriteDistributionMode, + writeOrdering: Seq[SortOrder]): Unit = { + if (isRequested(writeDistributionMode, writeOrdering) && + !catalog.capabilities().contains( + TableCatalogCapability.SUPPORTS_CREATE_TABLE_WITH_WRITE_DISTRIBUTION_AND_ORDERING)) { + throw QueryCompilationErrors.unsupportedTableOperationError( + catalog, ident, s"$operation ... $CLAUSES") + } + } + + /** Renders a requested sort key as SQL, for DESCRIBE. */ + def describeSortOrder(sortOrder: SortOrder): String = sortOrderToSQL(sortOrder, quoteIfNeeded) + + /** + * Renders the declared write distribution and ordering as `DISTRIBUTED BY PARTITION`, + * `[LOCALLY] ORDERED BY` or `UNORDERED` clauses for SHOW CREATE TABLE. Returns None when the + * pair has no clause form, such as `HASH` without partitioning, or when `replay`, which parses a + * statement carrying the clauses and returns the pair it declares, does not give back the same + * pair. Names are quoted only where needed, and on a second try all of them, which reserved + * keywords under `spark.sql.ansi.enforceReservedKeywords` need. + */ + def writeClausesSQL( + writeDistributionMode: WriteDistributionMode, + writeOrdering: Seq[SortOrder], + partitioning: Seq[Transform], + replay: String => Option[(WriteDistributionMode, Seq[SortOrder])]): Option[String] = { + lazy val partitioned = hasPartitioning(partitioning) + Seq[String => String](quoteIfNeeded, quoteIdentifier).iterator.flatMap { quote => + val orderBy = if (writeOrdering.nonEmpty) { + Some(writeOrdering.map(sortOrderToSQL(_, quote)).mkString("ORDERED BY (", ", ", ")")) + } else { + None + } + (writeDistributionMode, orderBy) match { + case (WriteDistributionMode.HASH, Some(o)) if partitioned => + Some(s"DISTRIBUTED BY PARTITION $o") + case (WriteDistributionMode.HASH, None) if partitioned => Some("DISTRIBUTED BY PARTITION") + case (WriteDistributionMode.RANGE, Some(o)) => Some(o) + case (WriteDistributionMode.NONE, Some(o)) => Some(s"LOCALLY $o") + case (WriteDistributionMode.NONE, None) => Some("UNORDERED") + case _ => None + } + }.find { clauses => + replay(clauses).exists { case (mode, ordering) => + mode == writeDistributionMode && sameSortOrders(ordering, writeOrdering) + } + } + } + + /** + * True when every column a sort key references is in `schema`, matched as `CheckAnalysis` + * matches the ordering of a CREATE/REPLACE TABLE statement. + */ + def referencesExist(schema: StructType, writeOrdering: Seq[SortOrder]): Boolean = { + writeOrdering.flatMap(_.expression().references()).forall { ref => + ref.fieldNames().nonEmpty && + Try(schema.findNestedField(ref.fieldNames().toImmutableArraySeq)).toOption.flatten.isDefined + } + } + + /** + * True when two orderings declare the same keys: the same column names, transform names and + * argument structure, literal values and types, directions and null orderings. + */ + def sameSortOrders(left: Seq[SortOrder], right: Seq[SortOrder]): Boolean = { + left.length == right.length && left.zip(right).forall { case (l, r) => + l.direction() == r.direction() && l.nullOrdering() == r.nullOrdering() && + keyOf(l.expression()) == keyOf(r.expression()) + } + } + + private def keyOf(e: Expression): Any = e match { + case r: NamedReference => r.fieldNames().toSeq + case l: Literal[_] => Try(CatalystLiteral.create(l.value, l.dataType)).toOption + case t: IdentityTransform => keyOf(t.ref) + case t: Transform => (t.name, t.arguments().toSeq.map(keyOf)) + case other => other + } + + private def sortOrderToSQL(sortOrder: SortOrder, quote: String => String): String = { + s"${toSQL(sortOrder.expression(), quote)} ${sortOrder.direction()} ${sortOrder.nullOrdering()}" + } + + // `describe` prints a literal's internal value, e.g. `0` for DATE '1970-01-01', so literals are + // rendered through Catalyst to keep their type. Catalyst renders a FLOAT as a CAST, which the + // parser does not accept here, so a finite FLOAT is rendered as `<v>F` instead, and a TIMESTAMP + // in UTC with an explicit offset so that it does not depend on the session time zone or + // timestamp type. A literal Catalyst cannot represent falls back to `describe`. Other connector + // expressions are rendered from their children: `describe` on one that + // `V2ExpressionSQLBuilder` does not know recurses without end. PARTITIONED BY in SHOW CREATE + // TABLE and the `Part N` rows of DESCRIBE render transforms with `describe`, so the same + // transform can print differently there. + private def toSQL(e: Expression, quote: String => String): String = e match { + case r: NamedReference => r.fieldNames.map(quote).mkString(".") + case l: Literal[_] => + try literalToSQL(CatalystLiteral.create(l.value, l.dataType)) catch { + case NonFatal(_) => describeLiteral(l) + } + case t: IdentityTransform => toSQL(t.ref, quote) + case t: Transform => + t.arguments().map(toSQL(_, quote)).mkString(s"${quote(t.name)}(", ", ", ")") + case c: Cast => s"CAST(${toSQL(c.expression(), quote)} AS ${c.dataType().sql})" + case g: GeneralScalarExpression => + g.children().map(toSQL(_, quote)).mkString(s"${g.name}(", ", ", ")") + case other => + other.children().map(toSQL(_, quote)).mkString(s"${other.getClass.getSimpleName}(", ", ", ")") + } + + // `LiteralValue.describe` also rejects a value that does not match its type. + private def describeLiteral(l: Literal[_]): String = { + try l.describe catch { + case NonFatal(_) => String.valueOf(l.value) + } + } + + private def literalToSQL(l: CatalystLiteral): String = (l.value, l.dataType) match { + case (f: Float, FloatType) if java.lang.Float.isFinite(f) => s"${f}F" + case (micros: Long, TimestampType) => + val utc = TimestampFormatter.getFractionFormatter(ZoneOffset.UTC).format(micros) + s"TIMESTAMP_LTZ '${utc}Z'" + case _ => l.sql Review Comment: Done. `literalToSQL` pads a TIME with more than microsecond precision to its precision, so `TIME '12:00:00.100000000'` round-trips as TIME(9). The unit suite has TIME(7) and TIME(9) rows with and without trailing zeros. ########## docs/sql-ref-syntax-ddl-create-table-datasource.md: ########## @@ -37,9 +37,17 @@ CREATE [ EXTERNAL ] TABLE [ IF NOT EXISTS ] table_identifier [ LOCATION path ] [ COMMENT table_comment ] [ TBLPROPERTIES ( key1=val1, key2=val2, ... ) ] + [ DISTRIBUTED BY PARTITION ] + [ [ LOCALLY ] ORDERED BY ( write_order_field [ , ... ] ) | UNORDERED ] Review Comment: Done: `[ LOCALLY ] ORDERED BY { ( write_order_field [ , ... ] ) | write_order_field [ , ... ] }`. ########## sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/util/WriteDistributionAndOrderingUtilsSuite.scala: ########## @@ -0,0 +1,213 @@ +/* + * 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.catalyst.util + +import java.math.BigDecimal +import java.nio.ByteBuffer +import java.time.{Duration, Instant, LocalDate, LocalDateTime, LocalTime, Period} + +import scala.util.Try + +import org.apache.spark.SparkFunSuite +import org.apache.spark.sql.catalyst.parser.CatalystSqlParser +import org.apache.spark.sql.catalyst.plans.SQLHelper +import org.apache.spark.sql.catalyst.plans.logical.CreateTable +import org.apache.spark.sql.connector.catalog.WriteDistributionMode +import org.apache.spark.sql.connector.expressions.{Expression, Expressions, FieldReference, GeneralScalarExpression, IdentityTransform, NamedReference, NullOrdering, SortDirection, SortOrder, Transform} +import org.apache.spark.sql.connector.expressions.LogicalExpressions.{apply => transform, literal, sort} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types._ + +class WriteDistributionAndOrderingUtilsSuite extends SparkFunSuite with SQLHelper { + import WriteDistributionAndOrdering._ + + private class ConnectorReference(names: String*) extends NamedReference { + override def fieldNames(): Array[String] = names.toArray + } + + private val id = FieldReference("id") + + private def f(arg: Expression): Transform = transform("f", id, arg) + + private def key(e: Expression): SortOrder = + sort(e, SortDirection.ASCENDING, NullOrdering.NULLS_FIRST) + + private def replay(clauses: String): Option[(WriteDistributionMode, Seq[SortOrder])] = { + Try(CatalystSqlParser.parsePlan(s"CREATE TABLE t (id INT) USING foo $clauses")).toOption + .collect { case c: CreateTable => (c.writeDistributionMode, c.writeOrdering) } + } + + private def emitted(e: Expression): Option[String] = { + writeClausesSQL(WriteDistributionMode.RANGE, Seq(key(e)), Seq.empty, replay) + } + + private val spellable: Seq[(String, Expression)] = Seq( + "reference" -> id, + "connector reference" -> transform("f", new ConnectorReference("order-id")), + "identity" -> IdentityTransform(id), + "byte" -> f(literal(1.toByte)), + "short" -> f(literal(1.toShort)), + "int" -> f(literal(1)), + "long" -> f(literal(1L)), + "float" -> f(literal(1.5F)), + "double" -> f(literal(1.5D)), + "decimal" -> f(literal(new BigDecimal("1.5"))), + "string" -> f(literal("x")), + "string with escapes" -> f(literal("it's a \\")), + "binary" -> f(literal(Array[Byte](1, 2))), + "date" -> f(literal(LocalDate.of(1970, 1, 1))), + "timestamp" -> f(literal(Instant.parse("2020-11-01T09:30:00Z"))), + "timestamp_ntz" -> f(literal(LocalDateTime.of(2020, 1, 1, 10, 0))), + "time" -> f(literal(LocalTime.of(12, 0))), + "day-time interval" -> f(literal(Duration.ofHours(1))), + "year-month interval" -> f(literal(Period.ofMonths(14))), + "bucket" -> transform("bucket", literal(16), id), + "bucket without columns" -> transform("bucket", literal(16)), + "Spark's bucket" -> Expressions.bucket(16, "id"), + "days" -> transform("days", id), + "unparsed name" -> transform("Days", id, literal("UTC"))) + + private val unspellable: Seq[(String, Expression)] = Seq( + "nested transform" -> transform("f", transform("g", id)), + "true" -> f(literal(true)), + "false" -> f(literal(false)), + "null" -> f(literal(null)), + "typed null" -> f(literal(null, IntegerType)), + "decimal with trailing zeros" -> + f(literal(new BigDecimal("1.50"), DecimalType(10, 2))), + "decimal in exponent form" -> + f(literal(new BigDecimal("0.0000001"), DecimalType(10, 7))), + "char" -> f(literal("x", CharType(3))), + "varchar" -> f(literal("x", VarcharType(3))), + "collated string" -> f(literal("x", StringType("UTF8_LCASE"))), + "NaN" -> f(literal(Float.NaN)), + "infinity" -> f(literal(Double.PositiveInfinity)), + "maximal float" -> f(literal(Float.MaxValue)), + "time with nanoseconds" -> f(literal(LocalTime.of(12, 0, 0, 100000000), TimeType(7))), + "time(6) holding a nanosecond" -> f(literal(LocalTime.ofNanoOfDay(1))), + "year interval not in whole years" -> + f(literal(13, YearMonthIntervalType(YearMonthIntervalType.YEAR))), + "hour interval not in whole hours" -> + f(literal(Duration.ofMinutes(90).toNanos / 1000, DayTimeIntervalType( + DayTimeIntervalType.HOUR))), + "negative-scale decimal" -> f(Expressions.literal(new BigDecimal("1E+3"))), + "int value with long type" -> f(literal(5, LongType)), + "byte buffer value" -> f(literal(ByteBuffer.wrap(Array[Byte](1, 2)), BinaryType)), + "bucket with columns first" -> transform("bucket", id, literal(16)), + "bucket with long count" -> transform("bucket", literal(16L), id), + "bucket with literal column" -> transform("bucket", literal(16), literal(1)), + "days with two arguments" -> transform("days", id, literal("UTC")), + "hours of a literal" -> transform("hours", literal(1))) + + test("a sort key is emitted iff its rendering replays as the same key") { + assert(spellable.filter { case (_, e) => + !emitted(e).contains(s"ORDERED BY (${describeSortOrder(key(e))})") + }.map(_._1) === Seq.empty) + assert(unspellable.filter { case (_, e) => emitted(e).isDefined }.map(_._1) === Seq.empty) + } + + test("a sort key is emitted iff it replays as the same key under parser confs") { + val timestamp = f(literal(Instant.parse("2020-11-01T09:30:00Z"))) + val wrong = Seq( + (SQLConf.SESSION_LOCAL_TIMEZONE.key -> "America/Los_Angeles", timestamp, true), + (SQLConf.TIMESTAMP_TYPE.key -> "TIMESTAMP_NTZ", timestamp, true), + (SQLConf.LEGACY_INTERVAL_ENABLED.key -> "true", f(literal(Duration.ofHours(1))), false), + (SQLConf.LEGACY_INTERVAL_ENABLED.key -> "true", f(literal(Period.ofMonths(14))), false), + (SQLConf.ESCAPED_STRING_LITERALS.key -> "true", f(literal("it's")), false), + (SQLConf.ESCAPED_STRING_LITERALS.key -> "true", f(literal("a\\b")), false), + (SQLConf.ESCAPED_STRING_LITERALS.key -> "true", f(literal("x")), true), + (SQLConf.LEGACY_TIME_PARSER_POLICY.key -> "LEGACY", + f(literal(LocalDate.of(-1, 1, 1))), false), + (SQLConf.LEGACY_FROM_DAYTIME_STRING.key -> "true", + f(literal(Duration.ofHours(26), DayTimeIntervalType( + DayTimeIntervalType.DAY, DayTimeIntervalType.HOUR))), false) + ).filter { case (conf, e, spelled) => withSQLConf(conf)(emitted(e).isDefined != spelled) } + assert(wrong.map(_._2.describe) === Seq.empty) + } + + test("a name that is a reserved keyword is emitted quoted only where the parser needs it") { + val order = FieldReference("order") + assert(emitted(order) === Some("ORDERED BY (order ASC NULLS FIRST)")) + withSQLConf( + SQLConf.ANSI_ENABLED.key -> "true", + SQLConf.ENFORCE_RESERVED_KEYWORDS.key -> "true") { + assert(emitted(order) === Some("ORDERED BY (`order` ASC NULLS FIRST)")) + } + } + + test("the pairs with no clause form are not emitted") { + val keys = Seq(key(id)) + Seq( + (WriteDistributionMode.HASH, keys, Seq.empty[Transform]), Review Comment: Thanks. The suite now replays the same statement as SHOW CREATE TABLE (`replayStatement`, shared with `ShowCreateTableExec`). The no-clause-form test uses a replay that accepts any clauses, so only the guards decide, and it covers a null mode. A new test emits each mode's clause form, including `DISTRIBUTED BY PARTITION ORDERED BY (id ASC NULLS FIRST)` and `LOCALLY ORDERED BY ...`. Another omits an unspellable key under HASH, RANGE and NONE. -- 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]
