dongjoon-hyun commented on code in PR #58153: URL: https://github.com/apache/spark/pull/58153#discussion_r4235543840
########## sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/WriteDistributionAndOrdering.scala: ########## @@ -0,0 +1,256 @@ +/* + * 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.{Expression, GeneralScalarExpression, GetArrayItem, IdentityTransform, Literal, NamedReference, SortOrder, Transform, VariantGet} +import org.apache.spark.sql.connector.expressions.filter.PartitionPredicate +import org.apache.spark.sql.errors.QueryCompilationErrors +import org.apache.spark.sql.internal.connector.{ExpressionWithToString, ToStringSQLBuilder} +import org.apache.spark.sql.types.{FloatType, StructType, TimestampLTZNanosType, TimestampType, TimeType} +import org.apache.spark.unsafe.types.TimestampNanosVal +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 + * [[replayStatement]] for the clauses and returns the pair it declares, does not give back the + * same pair with every referenced column in `schema`. 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], + schema: StructType, + 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) && + referencesExist(schema, ordering) + } + } + } + + /** + * A minimal statement whose parse depends only on `clauses`. The partitioning is there only so + * that the parser accepts `DISTRIBUTED BY PARTITION`. + */ + def replayStatement(clauses: String): String = + s"CREATE TABLE t USING foo PARTITIONED BY (p) $clauses" + + // Checks the parsed ordering exactly, which is at least as strict as CheckAnalysis on the + // replayed statement. + private 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. + private 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(sortKey(l.expression())) == keyOf(sortKey(r.expression())) + } + } + + // A plain column is a key's reference itself or, only at the top of a key, `identity(col)`. + private def sortKey(e: Expression): Expression = e match { + case t: IdentityTransform => t.ref + case other => other + } + + 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: Transform => (t.name, t.arguments().toSeq.map(keyOf)) + case other => other + } + + private def sortOrderToSQL(sortOrder: SortOrder, quote: String => String): String = { + val key = toSQL(sortKey(sortOrder.expression()), quote) + s"$key ${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, with the exceptions in `literalToSQL`. A literal + // Catalyst cannot represent falls back to `describe`. 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 null => "null" + 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 o: SortOrder => sortOrderToSQL(o, quote) + case t: Transform => + t.arguments().map(toSQL(_, quote)).mkString(s"${quote(t.name)}(", ", ", ")") + case _: ExpressionWithToString => + Try(new DescribingSQLBuilder(quote).build(e)).getOrElse(childrenToSQL(e, quote)) + case other => Try(other.describe).getOrElse(childrenToSQL(other, quote)) + } + + private def childrenToSQL(e: Expression, quote: String => String): String = { + val name = e match { + case g: GeneralScalarExpression => g.name + case _ => e.getClass.getSimpleName + } + e.children().map(toSQL(_, quote)).mkString(s"$name(", ", ", ")") + } + + // `ToStringSQLBuilder` renders a few expressions through their children's `describe`, and for an Review Comment: **[40] The recursion this comment describes is already fixed on master.** SPARK-59983 (#59241, 5926a10389e) changed `ToStringSQLBuilder.visitUnexpectedExpr` to render a `GeneralScalarExpression` with an unknown name as a function call. It went into master after this branch's base, so "recurses without end for a `GeneralScalarExpression` name it does not know" no longer holds after a rebase. SPARK-60075, which you filed for this in your last reply, looks like a duplicate of SPARK-59983 and can probably be closed. `DescribingSQLBuilder` itself is still useful, because it routes literals and references through `toSQL` and renders other unknown expressions from their children. Could we rebase and reword this comment to give only those reasons? ########## sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DescribeTableExec.scala: ########## @@ -226,6 +227,27 @@ case class DescribeTableExec( } } + /** + * Reports the table's declared write distribution and ordering, whether or not SHOW CREATE TABLE + * can reproduce them. + */ + private def addTableWriteDistributionAndOrdering(rows: ArrayBuffer[InternalRow]): Unit = { Review Comment: **[42] Each accessor is called up to three times.** `table.writeDistributionMode()` and `table.writeOrdering()` are each called once for `isRequested`, once for the null/empty check and once for the row. A connector may build these on every call, for example by converting its own sort order into a new `SortOrder[]` from refreshed metadata. That repeats the work, and the rows can come from different snapshots than the `isRequested` check. Reading both into local vals once at the top would avoid both. `ShowCreateTableExec.showTableWriteDistributionAndOrdering` already does this for `writeOrdering`. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ShowCreateTableExec.scala: ########## @@ -125,6 +131,57 @@ case class ShowCreateTableExec( } } + /** + * Emits the table's declared write distribution and ordering as clauses when parsing them in + * this session declares the same pair. Otherwise the pair is omitted: it has no clause form, or + * the catalog does not accept the clauses, or the statement would create a v1 table, or a sort + * key references a column the table does not have or does not parse back to the same key. + */ + private def showTableWriteDistributionAndOrdering( + resolvedTable: ResolvedTable, + builder: StringBuilder): Unit = { + import org.apache.spark.sql.connector.catalog.CatalogV2Implicits._ + val table = resolvedTable.table + val writeOrdering = table.writeOrdering().toImmutableArraySeq + val accepted = resolvedTable.catalog.capabilities().contains( + TableCatalogCapability.SUPPORTS_CREATE_TABLE_WITH_WRITE_DISTRIBUTION_AND_ORDERING) + if (accepted && !createsV1Table(resolvedTable)) { + WriteDistributionAndOrdering.writeClausesSQL( + table.writeDistributionMode(), + writeOrdering, + table.partitioning.toImmutableArraySeq, + table.columns.asSchema, + replay + ).foreach(clauses => builder ++= s"$clauses\n") + } + } + + // In the session catalog, a CREATE TABLE whose provider is not a v2 source creates a v1 table, + // which cannot record the clauses, so the replay would be rejected. + private def createsV1Table(resolvedTable: ResolvedTable): Boolean = { Review Comment: **[41] `createsV1Table` copies `ResolveSessionCatalog`'s v1-fallback rule inline.** `v1Capable` is `ResolveSessionCatalog.supportsV1Command` (L1030), and the provider check is `isV2Provider` (L950). Both are private there, so this is a second copy that has to change whenever the rule changes, such as the provider/serde resolution in `getStorageFormatAndProvider` (L765) or the v1 source list. If the copies drift, SHOW CREATE TABLE either emits clauses that the replayed CREATE TABLE rejects with `UNSUPPORTED_FEATURE.TABLE_OPERATION`, or drops clauses that would have worked, and the parse-back check cannot catch this because it only parses. Could we move these two predicates into a shared helper, e.g. next to `DataSourceV2Utils.getTableProvider`, and call it from both places? -- 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]
