[SPARK-51902][SQL] Enforce check constraint on table insertion

common/utils/src/main/resources/error/error-conditions.json

@@ -544,6 +544,14 @@
     ],
     "sqlState" : "56000"
   },
+  "CHECK_CONSTRAINT_VIOLATION" : {
+    "message" : [
+      "CHECK constraint <constraintName> <expression> violated by row with values:",
+      "<values>",
+      ""
+    ],
+    "sqlState" : "23001"
+  },
   "CIRCULAR_CLASS_REFERENCE" : {
     "message" : [
       "Cannot have circular references in class, but got the circular reference of class <t>."

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala

@@ -400,6 +400,7 @@ class Analyzer(override val catalogManager: CatalogManager) extends RuleExecutor
       ResolveWindowFrame ::
       ResolveNaturalAndUsingJoin ::
       ResolveOutputRelation ::
+      new ResolveTableConstraints(catalogManager) ::
       new ResolveDataFrameDropColumns(catalogManager) ::
       new ResolveSetVariable(catalogManager) ::
       ExtractWindowExpressions ::

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveTableConstraints.scala

@@ -0,0 +1,84 @@
+/*
+ * 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.analysis
+
+import scala.collection.mutable
+
+import org.apache.spark.sql.catalyst.expressions.{And, CheckInvariant, Expression, V2ExpressionUtils}
+import org.apache.spark.sql.catalyst.plans.logical.{Filter, LogicalPlan, V2WriteCommand}
+import org.apache.spark.sql.catalyst.rules.Rule
+import org.apache.spark.sql.catalyst.trees.TreePattern.COMMAND
+import org.apache.spark.sql.connector.catalog.CatalogManager
+import org.apache.spark.sql.connector.catalog.constraints.Check
+import org.apache.spark.sql.execution.datasources.v2.DataSourceV2Relation
+
+class ResolveTableConstraints(val catalogManager: CatalogManager) extends Rule[LogicalPlan] {
+
+  override def apply(plan: LogicalPlan): LogicalPlan = plan.resolveOperatorsWithPruning(
+    _.containsPattern(COMMAND), ruleId) {
+    case v2Write: V2WriteCommand
+      if v2Write.table.resolved && v2Write.query.resolved &&
+        !containsCheckInvariant(v2Write.query) && v2Write.outputResolved =>
+      v2Write.table match {
+        case r: DataSourceV2Relation
+          if r.table.constraints != null && r.table.constraints.nonEmpty =>
+          // Check constraint is the only enforced constraint for DSV2 tables.
+          val checkInvariants = r.table.constraints.collect {
+            case c: Check =>
+              val unresolvedExpr = buildCatalystExpression(c)
+              val columnExtractors = mutable.Map[String, Expression]()
+              buildColumnExtractors(unresolvedExpr, columnExtractors)
+              CheckInvariant(unresolvedExpr, columnExtractors.toSeq, c.name, c.predicateSql)
+          }
+          // Combine the check invariants into a single expression using conjunctive AND.
+          checkInvariants.reduceOption(And).fold(v2Write)(
+            condition => v2Write.withNewQuery(Filter(condition, v2Write.query)))
+        case _ =>
+          v2Write
+      }
+  }
+
+  private def containsCheckInvariant(plan: LogicalPlan): Boolean = {
+    plan match {
+      case Filter(condition, _) =>
+        condition.exists(_.isInstanceOf[CheckInvariant])
+
+      case _ => false
+    }
+  }
+
+  private def buildCatalystExpression(c: Check): Expression = {
+    Option(c.predicate())
+      .flatMap(V2ExpressionUtils.toCatalyst)
+      .getOrElse(catalogManager.v1SessionCatalog.parser.parseExpression(c.predicateSql()))
+  }
+
+  private def buildColumnExtractors(
+      expr: Expression,
+      columnExtractors: mutable.Map[String, Expression]): Unit = {
+    expr match {
+      case u: UnresolvedExtractValue =>
+        // When extracting a value from a Map or Array type, we display only the specific extracted
+        // value rather than the entire Map or Array structure for clarity and readability.
+        columnExtractors(u.sql) = u
+      case u: UnresolvedAttribute =>
+        columnExtractors(u.name) = u
+      case other =>
+        other.children.foreach(buildColumnExtractors(_, columnExtractors))
+    }
+  }
+}

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/constraints.scala

@@ -18,13 +18,17 @@ package org.apache.spark.sql.catalyst.expressions
 
 import java.util.UUID
 
-import org.apache.spark.sql.catalyst.analysis.UnresolvedException
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.analysis.{TypeCheckResult, UnresolvedException}
+import org.apache.spark.sql.catalyst.expressions.codegen.{Block, CodegenContext, ExprCode, FalseLiteral, TrueLiteral}
+import org.apache.spark.sql.catalyst.expressions.codegen.Block.BlockHelper
 import org.apache.spark.sql.catalyst.parser.ParseException
 import org.apache.spark.sql.catalyst.trees.CurrentOrigin
 import org.apache.spark.sql.catalyst.util.V2ExpressionBuilder
 import org.apache.spark.sql.connector.catalog.constraints.Constraint
 import org.apache.spark.sql.connector.expressions.FieldReference
-import org.apache.spark.sql.types.DataType
+import org.apache.spark.sql.errors.QueryExecutionErrors
+import org.apache.spark.sql.types.{BooleanType, DataType}
 
 trait TableConstraint extends Expression with Unevaluable {
   /** Convert to a data source v2 constraint */
@@ -259,3 +263,105 @@ case class ForeignKeyConstraint(
     copy(userProvidedCharacteristic = c)
   }
 }
+
+/**
+ * An expression that validates a check constraint on a column.
+ * If the evaluation result is false, it throws a [[SparkRuntimeException]] indicating constraint
+ * violation. Otherwise, it returns true to indicate that the constraint is satisfied, even if the
+ * expression is null.
+ *
+ * @param child The fully resolved expression to be evaluated to check the constraint.
+ * @param columnExtractors Extractors for each referenced column. Used to generate readable errors.
+ * @param constraintName The name of the constraint.
+ * @param predicateSql The SQL representation of the constraint.
+ */
+case class CheckInvariant(
+    child: Expression,
+    columnExtractors: Seq[(String, Expression)],
+    constraintName: String,
+    predicateSql: String)
+  extends Expression with NonSQLExpression {
+
+  override def children: Seq[Expression] = child +: columnExtractors.map(_._2)
+  override def dataType: DataType = BooleanType
+  override def nullable: Boolean = false
+
+  override def checkInputDataTypes(): TypeCheckResult = {
+    if (child.dataType != BooleanType) {
+      TypeCheckResult.TypeCheckFailure(
+        s"CHECK constraint must evaluate to a boolean expression, but got ${child.dataType}")
+    } else {
+      TypeCheckResult.TypeCheckSuccess
+    }
+  }
+
+  override def eval(input: InternalRow): Any = {
+    val result = child.eval(input)
+    if (result == false) {
+      val values = columnExtractors.map {
+        case (column, extractor) => column -> extractor.eval(input)
+      }.toMap
+      throw QueryExecutionErrors.checkViolation(constraintName, predicateSql, values)
+    }
+    true
+  }
+
+  /**
+   * Generate the code to extract values for the columns referenced in a violated CHECK constraint.
+   * We build parallel lists of full column names and their extracted values in the row which
+   * violates the constraint, to be passed to the constraint violation error
+   * in [[generateExpressionValidationCode()]].
+   *
+   * Note that this code is a bit expensive, so it shouldn't be run until we already
+   * know the constraint has been violated.
+   */
+  private def generateColumnValuesCode(
+    colList: String, valList: String, ctx: CodegenContext): Block = {
+    val start =
+      code"""
+            |java.util.List<String> $colList = new java.util.ArrayList<String>();
+            |java.util.List<Object> $valList = new java.util.ArrayList<Object>();
+            |""".stripMargin
+    columnExtractors.map {
+      case (name, extractor) =>
+        val colValue = extractor.genCode(ctx)
+        code"""
+              |$colList.add("$name");
+              |${colValue.code}
+              |if (${colValue.isNull}) {
+              |  $valList.add(null);
+              |} else {
+              |  $valList.add(${colValue.value});
+              |}
+              |""".stripMargin
+    }.fold(start)(_ + _)
+  }
+
+  private def generateExpressionValidationCode(ctx: CodegenContext): Block = {
+    val elementValue = child.genCode(ctx)
+    val colListName = ctx.freshName("colList")
+    val valListName = ctx.freshName("valList")
+    val ret = code"""${elementValue.code}
+          |
+          |if (!${elementValue.isNull} && ${elementValue.value} == false) {
+          |  ${generateColumnValuesCode(colListName, valListName, ctx)}
+          |  throw org.apache.spark.sql.errors.QueryExecutionErrors.checkViolationJava(
+          |     "$constraintName", "$predicateSql", $colListName, $valListName);
+          |}
+     """.stripMargin
+    ret
+  }
+
+  override protected def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
+    val code = generateExpressionValidationCode(ctx)
+    ev.copy(code = code, isNull = FalseLiteral, value = TrueLiteral)
+  }
+
+  override protected def withNewChildrenInternal(
+    newChildren: IndexedSeq[Expression]): Expression = {
+    copy(
+      child = newChildren.head,
+      columnExtractors = columnExtractors.map(_._1).zip(newChildren.tail)
+    )
+  }
+}

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/rules/RuleIdCollection.scala

@@ -112,6 +112,7 @@ object RuleIdCollection {
       "org.apache.spark.sql.catalyst.expressions.EliminatePipeOperators" ::
       "org.apache.spark.sql.catalyst.expressions.ValidateAndStripPipeExpressions" ::
       "org.apache.spark.sql.catalyst.analysis.ResolveUnresolvedHaving" ::
+      "org.apache.spark.sql.catalyst.analysis.ResolveTableConstraints" ::
       // Catalyst Optimizer rules
       "org.apache.spark.sql.catalyst.optimizer.BooleanSimplification" ::
       "org.apache.spark.sql.catalyst.optimizer.CollapseProject" ::

sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryExecutionErrors.scala

@@ -24,6 +24,8 @@ import java.time.DateTimeException
 import java.util.Locale
 import java.util.concurrent.TimeoutException
 
+import scala.jdk.CollectionConverters._
+
 import com.fasterxml.jackson.core.{JsonParser, JsonToken}
 import org.apache.hadoop.fs.{FileAlreadyExistsException, FileStatus, Path}
 import org.apache.hadoop.fs.permission.FsPermission
@@ -2993,4 +2995,35 @@ private[sql] object QueryExecutionErrors extends QueryErrorsBase with ExecutionE
   def notAbsolutePathError(path: Path): SparkException = {
     SparkException.internalError(s"$path is not absolute path.")
   }
+
+  // Throws a SparkRuntimeException when a CHECK constraint is violated, including details of the
+  // violation.
+  def checkViolation(
+      constraintName: String,
+      sqlStr: String,
+      values: Map[String, Any]): SparkRuntimeException = {
+    // Sort by the column name to generate consistent error messages in Scala 2.12 and 2.13.
+    val valueLines = values.toSeq.sortBy(_._1).map {
+      case (column, value) =>
+        s" - $column : $value"
+    }.mkString("\n")
+    new SparkRuntimeException(
+      errorClass = "CHECK_CONSTRAINT_VIOLATION",
+      messageParameters = Map(
+        "constraintName" -> constraintName,
+        "expression" -> sqlStr,
+        "values" -> valueLines
+      )
+    )
+  }
+
+  // Throws a SparkRuntimeException when a CHECK constraint is violated, including details of the
+  // violation. This is a Java-friendly version of the above method.
+  def checkViolationJava(
+      constraintName: String,
+      sqlStr: String,
+      columns: java.util.List[String],
+      values: java.util.List[Any]): SparkRuntimeException = {
+    checkViolation(constraintName, sqlStr, columns.asScala.zip(values.asScala).toMap)
+  }
 }

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/ConstraintExpressionSuite.scala

@@ -0,0 +1,45 @@
+/*
+ * 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.expressions
+
+import org.apache.spark.{SparkFunSuite, SparkRuntimeException}
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.types.IntegerType
+
+class ConstraintExpressionSuite extends SparkFunSuite with ExpressionEvalHelper {
+  private val boundRef = BoundReference(0, IntegerType, nullable = true)
+  private val expr =
+    CheckInvariant(GreaterThan(boundRef, Literal(0)), Seq(("a", boundRef)), "c1", "a > 0")
+
+  def expectedMessage(value: String): String =
+    s"""|[CHECK_CONSTRAINT_VIOLATION] CHECK constraint c1 a > 0 violated by row with values:
+        | - a : $value
+        | SQLSTATE: 23001""".stripMargin
+
+  test("CheckInvariant: returns true if column 'a' > 0") {
+    checkEvaluation(expr, true, InternalRow(1))
+  }
+
+  test("CheckInvariant: return true if column 'a' is null") {
+    checkEvaluation(expr, true, InternalRow(null))
+  }
+
+  test("CheckInvariant: throws exception if column 'a' <= 0") {
+    checkExceptionInExpression[SparkRuntimeException](
+      expr, InternalRow(-1), expectedMessage("-1"))
+  }
+}