+ * When this capability is present, Spark will: + *
+ * Without this capability, the connector is responsible for handling generated column values + * during writes. + * + * @since 4.3.0 + */ + GENERATE_COLUMN_VALUES_ON_WRITE } diff --git a/sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/TableCatalogCapability.java b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/TableCatalogCapability.java index ac829395b9bcd..a60c827d5ace1 100644 --- a/sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/TableCatalogCapability.java +++ b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/TableCatalogCapability.java @@ -92,23 +92,5 @@ public enum TableCatalogCapability { * {@link TableCatalog#createTable}. * See {@link Column#identityColumnSpec()}. */ - SUPPORTS_CREATE_TABLE_WITH_IDENTITY_COLUMNS, - - /** - * Signals that the TableCatalog supports Spark auto-filling generated column values and - * enforcing generated column constraints during writes. - *
- * When this capability is present, Spark will: - *
- * Without this capability, the connector is responsible for handling generated column values - * during writes. - * - * @since 4.3.0 - */ - SUPPORT_GENERATED_COLUMN_ON_WRITE + SUPPORTS_CREATE_TABLE_WITH_IDENTITY_COLUMNS } diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala index 7c00522887737..981e92ea69ef9 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala @@ -3930,11 +3930,11 @@ class Analyzer( val defaultValueFillMode = if (conf.coerceInsertNestedTypes && v2Write.schemaEvolutionEnabled) RECURSE else FILL - // Only let TableOutputResolver see generation expression metadata if the catalog + // Only let TableOutputResolver see generation expression metadata if the table // supports auto-filling generated columns on write. val expected = v2Write.table match { case r: DataSourceV2Relation - if !GeneratedColumn.supportsGeneratedColumnsOnWrite(r.catalog) => + if !GeneratedColumn.supportsGeneratedColumnsOnWrite(r.table) => r.output.map(GeneratedColumn.removeGenerationExpressionMetadata) case _ => v2Write.table.output } diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveTableConstraints.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveTableConstraints.scala index de4490143a83d..2a4b5693c705d 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveTableConstraints.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveTableConstraints.scala @@ -88,7 +88,7 @@ class ResolveTableConstraints(val catalogManager: CatalogManager) extends Rule[L private def buildGeneratedColumnConstraints( r: DataSourceV2Relation, v2Write: V2WriteCommand): Seq[Expression] = { - if (!GeneratedColumn.supportsGeneratedColumnsOnWrite(r.catalog, r.table.columns())) { + if (!GeneratedColumn.supportsGeneratedColumnsOnWrite(r.table, r.table.columns())) { return Seq.empty } diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala index 60143ae46c3c9..a64bdb6cbf58f 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala @@ -62,7 +62,7 @@ trait RewriteRowLevelCommand extends Rule[LogicalPlan] { relation: DataSourceV2Relation, command: Command): Unit = { if (GeneratedColumn.supportsGeneratedColumnsOnWrite( - relation.catalog, relation.table.columns())) { + relation.table, relation.table.columns())) { throw QueryCompilationErrors.unsupportedTableOperationError( relation.catalog.get, relation.identifier.get, s"${command.toString} with generated columns") diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/GeneratedColumn.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/GeneratedColumn.scala index 0469cfc7678e8..f8ca60c3a40af 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/GeneratedColumn.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/GeneratedColumn.scala @@ -19,8 +19,8 @@ package org.apache.spark.sql.catalyst.util import org.apache.spark.sql.catalyst.expressions.Attribute import org.apache.spark.sql.catalyst.plans.logical.ColumnDefinition -import org.apache.spark.sql.connector.catalog.{CatalogPlugin, Column, Identifier, TableCatalog, - TableCatalogCapability} +import org.apache.spark.sql.connector.catalog.{Column, Identifier, Table, TableCapability, + TableCatalog, TableCatalogCapability} import org.apache.spark.sql.errors.QueryCompilationErrors import org.apache.spark.sql.types.{Metadata, MetadataBuilder, StructField, StructType} @@ -92,28 +92,24 @@ object GeneratedColumn { } /** - * Whether the table catalog supports Spark auto-filling generated column values and enforcing - * generated column constraints during writes (the - * [[TableCatalogCapability.SUPPORT_GENERATED_COLUMN_ON_WRITE]] capability). Without it, the - * connector is responsible for handling generated column values. + * Whether the table wants Spark to auto-fill generated column values and enforce generated + * column constraints during writes (the + * [[TableCapability.GENERATE_COLUMN_VALUES_ON_WRITE]] capability). Without it, the connector is + * responsible for handling generated column values. */ - def supportsGeneratedColumnsOnWrite(catalog: Option[CatalogPlugin]): Boolean = { - catalog.exists { - case tc: TableCatalog => - tc.capabilities().contains(TableCatalogCapability.SUPPORT_GENERATED_COLUMN_ON_WRITE) - case _ => false - } + def supportsGeneratedColumnsOnWrite(table: Table): Boolean = { + table.capabilities().contains(TableCapability.GENERATE_COLUMN_VALUES_ON_WRITE) } /** - * Whether the catalog supports generated columns on write (see + * Whether the table supports generated columns on write (see * [[supportsGeneratedColumnsOnWrite]]) and the given columns include at least one generated * column. */ def supportsGeneratedColumnsOnWrite( - catalog: Option[CatalogPlugin], + table: Table, columns: Array[Column]): Boolean = { - supportsGeneratedColumnsOnWrite(catalog) && + supportsGeneratedColumnsOnWrite(table) && columns.exists(_.columnGenerationExpression() != null) } diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryBaseTable.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryBaseTable.scala index 78775d80cfa6d..fba80eb3d4cbe 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryBaseTable.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryBaseTable.scala @@ -195,6 +195,8 @@ abstract class InMemoryBaseTable( private val acceptAnySchema = properties.getOrDefault("accept-any-schema", "false").toBoolean private val autoSchemaEvolution = properties.getOrDefault("auto-schema-evolution", "true") .toBoolean + private val generateColumnValuesOnWrite = + properties.getOrDefault("generate-column-values-on-write", "true").toBoolean partitioning.foreach { case _: IdentityTransform => @@ -488,7 +490,10 @@ abstract class InMemoryBaseTable( override def capabilities(): util.Set[TableCapability] = (baseCapabiilities ++ (if (acceptAnySchema) Seq(TableCapability.ACCEPT_ANY_SCHEMA) else Seq.empty) ++ - (if (autoSchemaEvolution) Seq(TableCapability.AUTOMATIC_SCHEMA_EVOLUTION) else Seq.empty)) + (if (autoSchemaEvolution) Seq(TableCapability.AUTOMATIC_SCHEMA_EVOLUTION) else Seq.empty) ++ + (if (generateColumnValuesOnWrite) { + Seq(TableCapability.GENERATE_COLUMN_VALUES_ON_WRITE) + } else Seq.empty)) .asJava override def newScanBuilder(options: CaseInsensitiveStringMap): ScanBuilder = { diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTableCatalog.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTableCatalog.scala index 6466e3b3da300..bb137ba4830df 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTableCatalog.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTableCatalog.scala @@ -286,8 +286,7 @@ class InMemoryTableCatalog extends BasicInMemoryTableCatalog with SupportsNamesp TableCatalogCapability.SUPPORT_COLUMN_DEFAULT_VALUE, TableCatalogCapability.SUPPORT_TABLE_CONSTRAINT, TableCatalogCapability.SUPPORTS_CREATE_TABLE_WITH_GENERATED_COLUMNS, - TableCatalogCapability.SUPPORTS_CREATE_TABLE_WITH_IDENTITY_COLUMNS, - TableCatalogCapability.SUPPORT_GENERATED_COLUMN_ON_WRITE + TableCatalogCapability.SUPPORTS_CREATE_TABLE_WITH_IDENTITY_COLUMNS ).asJava } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/ResolveWriteToStream.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/ResolveWriteToStream.scala index 2afb2ad2f4eb8..4167f12885b0e 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/ResolveWriteToStream.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/ResolveWriteToStream.scala @@ -67,7 +67,7 @@ object ResolveWriteToStream extends Rule[LogicalPlan] { // Streaming writes with generated columns are not yet supported. s.catalogAndIdent.foreach { case (catalog, ident) => - if (GeneratedColumn.supportsGeneratedColumnsOnWrite(Some(catalog), s.sink.columns())) { + if (GeneratedColumn.supportsGeneratedColumnsOnWrite(s.sink, s.sink.columns())) { throw QueryCompilationErrors.unsupportedTableOperationError( catalog, ident, "streaming write with generated columns") } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/GeneratedColumnWriteSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/GeneratedColumnWriteSuite.scala index 54875c54a1032..81b2ea59cad9e 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/connector/GeneratedColumnWriteSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/GeneratedColumnWriteSuite.scala @@ -21,8 +21,7 @@ import org.apache.spark.SparkRuntimeException import org.apache.spark.sql.{AnalysisException, QueryTest, Row} import org.apache.spark.sql.catalyst.QueryPlanningTracker import org.apache.spark.sql.catalyst.expressions.CheckInvariant -import org.apache.spark.sql.connector.catalog.{InMemoryCatalog, InMemoryRowLevelOperationTableCatalog, - TableCatalogCapability} +import org.apache.spark.sql.connector.catalog.InMemoryRowLevelOperationTableCatalog import org.apache.spark.sql.execution.streaming.runtime.MemoryStream import org.apache.spark.sql.internal.SQLConf @@ -40,16 +39,10 @@ class GeneratedColumnWriteSuite extends QueryTest with DatasourceV2SQLBase { } } - // A catalog that supports creating generated columns but does NOT declare - // SUPPORT_GENERATED_COLUMN_ON_WRITE, so Spark must not auto-fill or enforce them. - private val noWriteCapCat = "nowritecapcat" - - private def withNoWriteCapCatalog(f: => Unit): Unit = { - withSQLConf(s"spark.sql.catalog.$noWriteCapCat" -> - classOf[InMemoryNoGenColWriteCatalog].getName) { - f - } - } + // A table property that makes the in-memory test table omit the + // GENERATE_COLUMN_VALUES_ON_WRITE capability, so Spark must not auto-fill or enforce + // generated columns for it. + private val noWriteCapProp = "generate-column-values-on-write" private def hasCheckInvariant(sqlText: String): Boolean = { val parsed = spark.sessionState.sqlParser.parsePlan(sqlText) @@ -1154,44 +1147,29 @@ class GeneratedColumnWriteSuite extends QueryTest with DatasourceV2SQLBase { } } - test("catalog without write capability does not auto-fill or enforce generated columns") { - withNoWriteCapCatalog { - val tblName = "my_tab" - withTable(s"$noWriteCapCat.$tblName") { - sql(s"""CREATE TABLE $noWriteCapCat.$tblName( - | a INT, - | b INT GENERATED ALWAYS AS (a + 1) - |) USING foo""".stripMargin) + test("table without write capability does not auto-fill or enforce generated columns") { + val tblName = "my_tab" + withTable(s"testcat.$tblName") { + sql(s"""CREATE TABLE testcat.$tblName( + | a INT, + | b INT GENERATED ALWAYS AS (a + 1) + |) USING foo TBLPROPERTIES ('$noWriteCapProp' = 'false')""".stripMargin) - // A user-provided value that does NOT match the generation expression is written - // as-is: no constraint is enforced because the catalog does not opt in. - sql(s"INSERT INTO $noWriteCapCat.$tblName(a, b) VALUES (5, 999)") - checkAnswer(spark.table(s"$noWriteCapCat.$tblName"), Row(5, 999)) - - // Omitting the generated column does not auto-fill; it is treated as a regular - // nullable column and filled with null. - sql(s"INSERT INTO $noWriteCapCat.$tblName(a) VALUES (7)") - checkAnswer( - spark.table(s"$noWriteCapCat.$tblName"), - Row(5, 999) :: Row(7, null) :: Nil) - - // No CheckInvariant is added to the plan for the generated column. - assert(!hasCheckInvariant(s"INSERT INTO $noWriteCapCat.$tblName(a, b) VALUES (5, 6)"), - "No CheckInvariant should be added when the catalog lacks the write capability") - } - } - } -} + // A user-provided value that does NOT match the generation expression is written + // as-is: no constraint is enforced because the table does not opt in. + sql(s"INSERT INTO testcat.$tblName(a, b) VALUES (5, 999)") + checkAnswer(spark.table(s"testcat.$tblName"), Row(5, 999)) -/** - * A catalog that supports creating tables with generated columns but does NOT declare - * [[TableCatalogCapability.SUPPORT_GENERATED_COLUMN_ON_WRITE]], so Spark leaves generated - * column handling to the connector (no auto-fill, no constraint enforcement on write). - */ -class InMemoryNoGenColWriteCatalog extends InMemoryCatalog { - override def capabilities: java.util.Set[TableCatalogCapability] = { - val caps = new java.util.HashSet[TableCatalogCapability](super.capabilities) - caps.remove(TableCatalogCapability.SUPPORT_GENERATED_COLUMN_ON_WRITE) - caps + // Omitting the generated column does not auto-fill; it is treated as a regular + // nullable column and filled with null. + sql(s"INSERT INTO testcat.$tblName(a) VALUES (7)") + checkAnswer( + spark.table(s"testcat.$tblName"), + Row(5, 999) :: Row(7, null) :: Nil) + + // No CheckInvariant is added to the plan for the generated column. + assert(!hasCheckInvariant(s"INSERT INTO testcat.$tblName(a, b) VALUES (5, 6)"), + "No CheckInvariant should be added when the table lacks the write capability") + } } }