[SPARK-58311][SQL] Gate generated column values on write with a table capability

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

@@ -106,5 +106,25 @@ public enum TableCapability {
    * write modes, like {@link #TRUNCATE}, and {@link #OVERWRITE_BY_FILTER}, but cannot support
    * {@link #OVERWRITE_DYNAMIC}.
    */
-  V1_BATCH_WRITE
+  V1_BATCH_WRITE,
+
+  /**
+   * Signals that the table wants Spark to auto-fill generated column values and enforce generated
+   * column constraints during writes.
+   * <p>
+   * When this capability is present, Spark will:
+   * <ul>
+   *   <li>Auto-compute missing generated column values using the generation expression for
+   *       by-name writes. Ordinary by-position writes still require the input to provide a
+   *       value for every table column.</li>
+   *   <li>Validate explicitly-provided generated column values against the generation
+   *       expression.</li>
+   * </ul>
+   * <p>
+   * Without this capability, the connector is responsible for handling generated column values
+   * during writes.
+   *
+   * @since 4.3.0
+   */
+  GENERATE_COLUMN_VALUES_ON_WRITE
 }

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.
-   * <p>
-   * When this capability is present, Spark will:
-   * <ul>
-   *   <li>Auto-compute missing generated column values using the generation expression.</li>
-   *   <li>Validate explicitly-provided generated column values against the generation
-   *       expression.</li>
-   * </ul>
-   * <p>
-   * 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
 }

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
         }

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
     }
 

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")

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)
   }
 

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 = {

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
   }
 

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")
         }

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")
+    }
   }
 }