Spark 4.1: Add scoped replacement writes (INSERT INTO ... REPLACE USING)

docs/docs/spark-writes.md

@@ -29,18 +29,20 @@ Iceberg uses Apache Spark's DataSourceV2 API for data source and catalog impleme
 | Feature support                                  | Spark | Notes                                                                       |
 |--------------------------------------------------|---------|-----------------------------------------------------------------------------|
 | [SQL insert into](#insert-into)                  | ✔️      | ⚠ Requires `spark.sql.storeAssignmentPolicy=ANSI` (default since Spark 3.0) |
+| [SQL scoped replace](#insert-into--replace-using) | ✔️      | ⚠ Requires Iceberg Spark extensions and Spark 4.1 or higher                 |
 | [SQL merge into](#merge-into)                    | ✔️      | ⚠ Requires Iceberg Spark extensions                                         |
 | [SQL insert overwrite](#insert-overwrite)        | ✔️      | ⚠ Requires `spark.sql.storeAssignmentPolicy=ANSI` (default since Spark 3.0) |
 | [SQL delete from](#delete-from)                  | ✔️      | ⚠ Row-level delete requires Iceberg Spark extensions                        |
 | [SQL update](#update)                            | ✔️      | ⚠ Requires Iceberg Spark extensions                                         |
 | [DataFrame append](#appending-data)              | ✔️      |                                                                             |
 | [DataFrame overwrite](#overwriting-data)         | ✔️      |                                                                             |
+| [DataFrame scoped replace](#scoped-replacement)  | ✔️      | ⚠ Requires Iceberg Spark extensions and Spark 4.1 or higher                 |
 | [DataFrame CTAS and RTAS](#creating-tables)      | ✔️      | ⚠ Requires DSv2 API                                                         |
 | [DataFrame merge into](#merging-data)            | ✔️      | ⚠ Requires DSv2 API (Spark 4.0 and later)                                   |
 
 ## Writing with SQL
 
-Spark supports SQL `INSERT INTO`, `MERGE INTO`, and `INSERT OVERWRITE`, as well as the new `DataFrameWriterV2` API.
+Spark supports SQL `INSERT INTO`, `INSERT INTO ... REPLACE USING`, `MERGE INTO`, and `INSERT OVERWRITE`, as well as the new `DataFrameWriterV2` API.
 
 ### `INSERT INTO`
 
@@ -53,6 +55,31 @@ INSERT INTO prod.db.table VALUES (1, 'a'), (2, 'b')
 INSERT INTO prod.db.table SELECT ...
 ```
 
+### `INSERT INTO ... REPLACE USING`
+
+Iceberg supports scoped replacement writes with Iceberg Spark extensions in Spark 4.1 and higher. A scoped replace deletes existing target rows whose replacement scope appears in the source query, then inserts all rows from the source query in the same commit.
+
+```sql
+INSERT INTO prod.db.table
+REPLACE USING (scope_col_1, scope_col_2)
+SELECT ...
+```
+
+The columns listed in `REPLACE USING` define the replacement scope. For each distinct tuple of scope values produced by the source query, matching target rows are removed using null-safe equality, and the full source query output is appended.
+
+For example, this query replaces all existing rows for categories present in `prod.db.staged_rows` and keeps rows for other categories:
+
+```sql
+INSERT INTO prod.db.sample
+REPLACE USING (category)
+SELECT id, data, category, ts
+FROM prod.db.staged_rows
+```
+
+`REPLACE USING` is useful when incoming data contains complete replacement slices for one or more logical groups, such as tenants, departments, dates, or regions. Unlike `INSERT OVERWRITE`, the replacement scope is based on table columns in the source data and does not depend on the table's partition spec.
+
+The source query must produce the same number of columns as the target table, using the same assignment rules as `INSERT INTO`. Each `REPLACE USING` column must exist in both the target table and the source query output.
+
 ### `MERGE INTO`
 
 Spark supports `MERGE INTO` queries that can express row-level updates.
@@ -313,6 +340,34 @@ To explicitly overwrite partitions, use `overwrite` to supply a filter:
 data.writeTo("prod.db.table").overwrite($"level" === "INFO")
 ```
 
+### Scoped replacement
+
+Iceberg Spark extensions in Spark 4.1 and higher support scoped replacement writes through the
+`DataFrameWriterV2` API. A scoped replace deletes existing target rows whose replacement scope
+appears in the source `DataFrame`, then inserts all rows from the source in the same commit.
+
+```scala
+val staged: DataFrame = ...
+
+staged.writeTo("prod.db.sample")
+    .option("replace-using", "category")
+    .overwrite(lit(true))
+```
+
+Use the `replace-using` write option to list the target columns that define the replacement scope.
+Multiple columns are comma-separated:
+
+```scala
+staged.writeTo("prod.db.sample")
+    .option("replace-using", "tenant_id,business_date")
+    .option("target-file-size-bytes", "134217728")
+    .overwrite(lit(true))
+```
+
+The `overwrite(lit(true))` call is required so Spark applies overwrite privileges and Iceberg can
+replace matching target rows. Other Iceberg write options, such as target file size, compression,
+distribution, snapshot properties, and branch selection, can be supplied with `option` as usual.
+
 ### Creating tables
 
 To run a CTAS or RTAS, use `create`, `replace`, or `createOrReplace` operations:

spark/v4.1/spark-extensions/src/main/antlr/org.apache.spark.sql.catalyst.parser.extensions/IcebergSqlExtensions.g4

@@ -131,6 +131,13 @@ singleOrder
     : order EOF
     ;
 
+// Parses only the command head of `INSERT INTO t REPLACE USING (cols) <query>`.
+// The query tail remains Spark SQL and is delegated to Spark's parser until this syntax
+// can be represented directly in Spark's `insertInto` grammar.
+singleScopedReplaceHead
+    : INSERT INTO TABLE? multipartIdentifier REPLACE USING '(' fieldList ')' EOF
+    ;
+
 order
     : fields+=orderField (',' fields+=orderField)*
     | '(' fields+=orderField (',' fields+=orderField)* ')'
@@ -211,6 +218,7 @@ nonReserved
     | DISTRIBUTED | LOCALLY | MINUTES | MONTHS | UNORDERED | REPLACE | RETAIN | VERSION | WITH | IDENTIFIER_KW | FIELDS | SET | SNAPSHOT | SNAPSHOTS
     | TAG | TRUE | FALSE
     | MAP
+    | INSERT | INTO | USING
     ;
 
 snapshotId
@@ -243,6 +251,8 @@ FIELDS: 'FIELDS';
 FIRST: 'FIRST';
 HOURS: 'HOURS';
 IF : 'IF';
+INSERT: 'INSERT';
+INTO: 'INTO';
 LAST: 'LAST';
 LOCALLY: 'LOCALLY';
 MINUTES: 'MINUTES';
@@ -264,6 +274,7 @@ SNAPSHOTS: 'SNAPSHOTS';
 TABLE: 'TABLE';
 TAG: 'TAG';
 UNORDERED: 'UNORDERED';
+USING: 'USING';
 VERSION: 'VERSION';
 WITH: 'WITH';
 WRITE: 'WRITE';

spark/v4.1/spark-extensions/src/main/scala/org/apache/iceberg/spark/extensions/IcebergSparkSessionExtensions.scala

@@ -22,6 +22,8 @@ import org.apache.spark.sql.SparkSessionExtensions
 import org.apache.spark.sql.catalyst.analysis.CheckViews
 import org.apache.spark.sql.catalyst.analysis.ResolveBranch
 import org.apache.spark.sql.catalyst.analysis.ResolveViews
+import org.apache.spark.sql.catalyst.analysis.RewriteDataFrameScopedReplace
+import org.apache.spark.sql.catalyst.analysis.RewriteScopedReplace
 import org.apache.spark.sql.catalyst.optimizer.ReplaceStaticInvoke
 import org.apache.spark.sql.catalyst.parser.extensions.IcebergSparkSqlExtensionsParser
 import org.apache.spark.sql.execution.datasources.v2.ExtendedDataSourceV2Strategy
@@ -35,6 +37,8 @@ class IcebergSparkSessionExtensions extends (SparkSessionExtensions => Unit) {
     // analyzer extensions
     extensions.injectResolutionRule { spark => ResolveViews(spark) }
     extensions.injectPostHocResolutionRule { spark => ResolveBranch(spark) }
+    extensions.injectPostHocResolutionRule { spark => RewriteDataFrameScopedReplace(spark) }
+    extensions.injectPostHocResolutionRule { _ => RewriteScopedReplace }
     extensions.injectCheckRule(_ => CheckViews)
 
     // optimizer extensions

spark/v4.1/spark-extensions/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteDataFrameScopedReplace.scala

@@ -0,0 +1,142 @@
+/*
+ * 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 org.apache.iceberg.relocated.com.google.common.base.Splitter
+import org.apache.iceberg.spark.SparkTableUtil
+import org.apache.iceberg.spark.source.SparkTable
+import org.apache.spark.sql.AnalysisException
+import org.apache.spark.sql.SparkSession
+import org.apache.spark.sql.catalyst.expressions.Literal.TrueLiteral
+import org.apache.spark.sql.catalyst.plans.logical.AppendData
+import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
+import org.apache.spark.sql.catalyst.plans.logical.OverwriteByExpression
+import org.apache.spark.sql.catalyst.plans.logical.OverwritePartitionsDynamic
+import org.apache.spark.sql.catalyst.plans.logical.ReplaceScopedData
+import org.apache.spark.sql.catalyst.rules.Rule
+import org.apache.spark.sql.execution.datasources.v2.DataSourceV2Relation
+import org.apache.spark.sql.util.CaseInsensitiveStringMap
+import scala.jdk.CollectionConverters._
+
+/**
+ * Rewrites the DataFrameWriterV2 scoped-replace API into the shared logical command.
+ *
+ * The API is carried by a write option on `overwrite(true)`, so Spark resolves the query using the
+ * standard DataFrameWriterV2 by-name write path before this rule runs. Normal Iceberg write options
+ * are moved onto the relation because Spark's row-level write planning reads options from the
+ * relation after [[ReplaceScopedData]] is lowered to [[org.apache.spark.sql.catalyst.plans.logical.ReplaceData]]
+ * or [[org.apache.spark.sql.catalyst.plans.logical.WriteDelta]].
+ */
+case class RewriteDataFrameScopedReplace(spark: SparkSession) extends Rule[LogicalPlan] {
+
+  import RewriteDataFrameScopedReplace._
+
+  override def apply(plan: LogicalPlan): LogicalPlan = plan resolveOperators {
+    case overwrite @ OverwriteByExpression(
+          r @ DataSourceV2Relation(table: SparkTable, _, _, _, relationOptions, _),
+          deleteExpr,
+          query,
+          writeOptions,
+          _,
+          None,
+          _) if overwrite.resolved && replaceUsing(writeOptions).isDefined =>
+
+      if (!deleteExpr.semanticEquals(TrueLiteral)) {
+        throw analysisError(
+          s"Scoped replace must be invoked with overwrite(true), found: ${deleteExpr.sql}")
+      }
+
+      if (query.isStreaming) {
+        throw analysisError("Scoped replace is not supported for streaming DataFrames")
+      }
+
+      val mergedOptions = writeOptionsWithoutReplaceUsing(writeOptions, relationOptions)
+      val targetBranch = SparkTableUtil.determineWriteBranch(spark, table, mergedOptions)
+      val targetTable =
+        if (table.branch() == targetBranch) table else table.copyWithBranch(targetBranch)
+      val scopeColumns = parseScopeColumns(replaceUsing(writeOptions).get)
+
+      ReplaceScopedData(r.copy(table = targetTable, options = mergedOptions), scopeColumns, query)
+
+    case append @ AppendData(
+          DataSourceV2Relation(_: SparkTable, _, _, _, _, _),
+          _,
+          writeOptions,
+          _,
+          _,
+          _) if append.resolved && replaceUsing(writeOptions).isDefined =>
+      throw analysisError("Scoped replace must be invoked with overwrite(true), not append()")
+
+    case overwritePartitions @ OverwritePartitionsDynamic(
+          DataSourceV2Relation(_: SparkTable, _, _, _, _, _),
+          _,
+          writeOptions,
+          _,
+          _) if overwritePartitions.resolved && replaceUsing(writeOptions).isDefined =>
+      throw analysisError(
+        "Scoped replace must be invoked with overwrite(true), not overwritePartitions()")
+  }
+
+  private def writeOptionsWithoutReplaceUsing(
+      writeOptions: Map[String, String],
+      relationOptions: CaseInsensitiveStringMap): CaseInsensitiveStringMap = {
+    val relationMap = relationOptions.asCaseSensitiveMap.asScala.toMap
+    val merged = removeReplaceUsing(relationMap) ++ removeReplaceUsing(writeOptions)
+    new CaseInsensitiveStringMap(merged.asJava)
+  }
+
+  private def removeReplaceUsing(options: Map[String, String]): Map[String, String] = {
+    options.filterNot { case (key, _) => key.equalsIgnoreCase(ReplaceUsingOption) }
+  }
+
+  private def replaceUsing(options: Map[String, String]): Option[String] = {
+    Option(new CaseInsensitiveStringMap(options.asJava).get(ReplaceUsingOption))
+  }
+
+  // `replace-using` follows the standard Iceberg comma-separated option convention, so the list is
+  // split on commas and each token is parsed as a multi-part identifier by the session parser (the
+  // same parser the SQL `REPLACE USING` path delegates to). A consequence of splitting first is that
+  // a column whose quoted name itself contains a comma cannot be expressed through this option; such
+  // names are not representable in a comma-delimited option and must use the SQL `REPLACE USING` form.
+  private def parseScopeColumns(columns: String): Seq[Seq[String]] = {
+    Splitter
+      .on(',')
+      .trimResults()
+      .splitToList(columns)
+      .asScala
+      .map { column =>
+        spark.sessionState.sqlParser.parseMultipartIdentifier(column)
+      }
+      .toSeq
+  }
+
+  private def analysisError(message: String): AnalysisException = {
+    new AnalysisException(
+      errorClass = "_LEGACY_ERROR_TEMP_ICEBERG_SCOPED_REPLACE",
+      sqlState = null,
+      messageTemplate = message,
+      messageParameters = Map.empty[String, String],
+      cause = None,
+      message = Some(message))
+  }
+}
+
+object RewriteDataFrameScopedReplace {
+  private val ReplaceUsingOption = "replace-using"
+}