feat(sql): blob v2 copy-through for MERGE/UPDATE row-level commands (#634)

Author: geruh

docs/src/operations/dml/insert-into.md

@@ -209,12 +209,23 @@ Blob v2 columns [read as descriptor structs](../../config.md#blob-v2-reads). Cop
 INSERT INTO documents_copy SELECT id, content FROM documents;
 ```
 
-`writeTo().append()` and `.overwrite()` work the same way. Joins and filters are fine if the blob column stays a direct select or alias.
-
-**Limitations**
-
-- Do not transform the blob column (`CASE`, `content.size`, `GROUP BY`, `UNION`, `DISTINCT` fail the write).
-- `MERGE`, `UPDATE`, and dynamic partition overwrite are not supported.
+Only direct column references (or simple aliases) from a single Lance source table are
+copied this way; `ORDER BY` and `LIMIT` on the SELECT are supported. `writeTo().append()`
+and `.overwrite()` work the same way. Transformed values (e.g. `content.size`),
+`DISTINCT`, `GROUP BY`, `UNION`, ordering by the blob column itself, and joins of multiple
+blob sources keep their descriptor semantics and cannot be inserted into a blob column. See
+[CREATE TABLE](../ddl/create-table.md#copying-blob-v2-columns-with-ctas) for creating a new
+table from a blob v2 query.
+
+On Spark 3.5+, `MERGE INTO` and `UPDATE` deep-copy blob v2 columns the same way: a blob
+assigned from a source column (`SET t.content = s.content`) is copied from the source table,
+and blobs the command does not touch are carried forward from the target. Additional
+limitations apply:
+
+- Assigning a binary literal to a blob v2 column through `UPDATE` or `MERGE` is rejected at
+  analysis; use `INSERT` for direct binary writes.
+- Dynamic partition overwrite is not supported.
+- On Spark 3.4, `MERGE` and `UPDATE` are not supported for blob v2 tables.
 
 ## Writing Large String Data
 

lance-spark-3.4_2.12/src/test/java/org/lance/spark/BlobV2RowLevelUnsupportedTest.java

@@ -0,0 +1,69 @@
+/*
+ * Licensed 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.lance.spark;
+
+import org.junit.jupiter.api.Test;
+
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+public class BlobV2RowLevelUnsupportedTest extends AbstractBlobV2CopyTest {
+
+  @Test
+  public void testMergeIntoBlobV2TableFailsLoudly() throws Exception {
+    String src = "v2_rl34_merge_src_" + System.currentTimeMillis();
+    String tgt = "v2_rl34_merge_tgt_" + System.currentTimeMillis();
+    String fqSrc = fq(src);
+    String fqTgt = fq(tgt);
+    createV2BlobSource(fqSrc, row(1, deterministicBlob(141, 16)));
+    createV2BlobSource(fqTgt, row(1, deterministicBlob(142, 16)));
+    try {
+      Exception ex =
+          assertThrows(
+              Exception.class,
+              () ->
+                  spark.sql(
+                      "MERGE INTO "
+                          + fqTgt
+                          + " t USING "
+                          + fqSrc
+                          + " s ON t.id = s.id"
+                          + " WHEN MATCHED THEN UPDATE SET t.data = s.data"));
+      assertTrue(
+          ex.getMessage() != null && !ex.getMessage().isEmpty(),
+          "MERGE must fail with a message, got: " + ex);
+    } finally {
+      spark.sql("DROP TABLE IF EXISTS " + fqSrc);
+      spark.sql("DROP TABLE IF EXISTS " + fqTgt);
+    }
+  }
+
+  @Test
+  public void testUpdateBlobV2TableFailsLoudly() throws Exception {
+    String tgt = "v2_rl34_update_tgt_" + System.currentTimeMillis();
+    String fqTgt = fq(tgt);
+    createV2BlobSource(fqTgt, row(1, deterministicBlob(143, 16)));
+    try {
+      Exception ex =
+          assertThrows(
+              Exception.class,
+              () -> spark.sql("UPDATE " + fqTgt + " SET data = X'00' WHERE id = 1"));
+      assertTrue(
+          ex.getMessage() != null && !ex.getMessage().isEmpty(),
+          "UPDATE must fail with a message, got: " + ex);
+    } finally {
+      spark.sql("DROP TABLE IF EXISTS " + fqTgt);
+    }
+  }
+}

lance-spark-3.5_2.12/pom.xml

@@ -97,6 +97,7 @@
                                 <source>src/main/java</source>
                                 <source>src/main/java11</source>
                                 <source>src/main/scala</source>
+                                <source>src/main/scala-rowlevel</source>
                                 <source>${project.build.directory}/generated-sources/antlr4</source>
                             </sources>
                         </configuration>

lance-spark-3.5_2.12/src/main/java/org/lance/spark/LancePositionDeltaOperation.java

@@ -14,6 +14,8 @@
 package org.lance.spark;
 
 import org.lance.spark.read.LanceScanBuilder;
+import org.lance.spark.utils.BlobSourceContext;
+import org.lance.spark.utils.BlobUtils;
 import org.lance.spark.write.SparkPositionDeltaWriteBuilder;
 
 import org.apache.spark.sql.connector.expressions.Expressions;
@@ -23,9 +25,11 @@
 import org.apache.spark.sql.connector.write.LogicalWriteInfo;
 import org.apache.spark.sql.connector.write.RowLevelOperation;
 import org.apache.spark.sql.connector.write.SupportsDelta;
+import org.apache.spark.sql.types.StructField;
 import org.apache.spark.sql.types.StructType;
 import org.apache.spark.sql.util.CaseInsensitiveStringMap;
 
+import java.util.Collections;
 import java.util.Map;
 
 public class LancePositionDeltaOperation implements RowLevelOperation, SupportsDelta {
@@ -50,6 +54,8 @@ public class LancePositionDeltaOperation implements RowLevelOperation, SupportsD
 
   private final Map<String, String> tableProperties;
 
+  private final Map<String, BlobSourceContext> blobSourceContexts;
+
   public LancePositionDeltaOperation(
       Command command,
       StructType sparkSchema,
@@ -60,6 +66,30 @@ public LancePositionDeltaOperation(
       boolean managedVersioning,
       String fileFormatVersion,
       Map<String, String> tableProperties) {
+    this(
+        command,
+        sparkSchema,
+        readOptions,
+        initialStorageOptions,
+        namespaceImpl,
+        namespaceProperties,
+        managedVersioning,
+        fileFormatVersion,
+        tableProperties,
+        Collections.emptyMap());
+  }
+
+  private LancePositionDeltaOperation(
+      Command command,
+      StructType sparkSchema,
+      LanceSparkReadOptions readOptions,
+      Map<String, String> initialStorageOptions,
+      String namespaceImpl,
+      Map<String, String> namespaceProperties,
+      boolean managedVersioning,
+      String fileFormatVersion,
+      Map<String, String> tableProperties,
+      Map<String, BlobSourceContext> blobSourceContexts) {
     this.command = command;
     this.sparkSchema = sparkSchema;
     this.readOptions = readOptions;
@@ -69,6 +99,22 @@ public LancePositionDeltaOperation(
     this.managedVersioning = managedVersioning;
     this.fileFormatVersion = fileFormatVersion;
     this.tableProperties = tableProperties;
+    this.blobSourceContexts = blobSourceContexts;
+  }
+
+  public LancePositionDeltaOperation withBlobSourceContexts(
+      Map<String, BlobSourceContext> contexts) {
+    return new LancePositionDeltaOperation(
+        command,
+        sparkSchema,
+        readOptions,
+        initialStorageOptions,
+        namespaceImpl,
+        namespaceProperties,
+        managedVersioning,
+        fileFormatVersion,
+        tableProperties,
+        contexts);
   }
 
   @Override
@@ -84,7 +130,7 @@ public ScanBuilder newScanBuilder(CaseInsensitiveStringMap caseInsensitiveString
 
   @Override
   public DeltaWriteBuilder newWriteBuilder(LogicalWriteInfo logicalWriteInfo) {
-    // Create write options from read options for delta operations
+    rejectBlobV2Descriptors(logicalWriteInfo.schema());
     LanceSparkWriteOptions.Builder writeOptionsBuilder =
         LanceSparkWriteOptions.builder()
             .datasetUri(readOptions.getDatasetUri())
@@ -102,14 +148,44 @@ public DeltaWriteBuilder newWriteBuilder(LogicalWriteInfo logicalWriteInfo) {
       }
     }
     LanceSparkWriteOptions writeOptions = writeOptionsBuilder.build();
+    Map<String, BlobSourceContext> contexts =
+        blobSourceContexts.isEmpty()
+            ? BlobSourceContext.decodeFromWriteOptions(logicalWriteInfo.options())
+            : blobSourceContexts;
     return new SparkPositionDeltaWriteBuilder(
         sparkSchema,
         writeOptions,
         initialStorageOptions,
         namespaceImpl,
         namespaceProperties,
         managedVersioning,
-        readOptions.getTableId());
+        readOptions.getTableId(),
+        contexts);
+  }
+
+  /**
+   * A blob v2 column accepts only BINARY on write. If a descriptor struct reaches the delta writer
+   * it would serialize as opaque bytes and silently corrupt the column, so any row-level plan the
+   * copy-through rewrite could not convert to copy tokens must fail here.
+   */
+  private void rejectBlobV2Descriptors(StructType writeSchema) {
+    for (StructField tableField : sparkSchema.fields()) {
+      if (!BlobUtils.isBlobV2SparkField(tableField)) {
+        continue;
+      }
+      for (StructField writeField : writeSchema.fields()) {
+        if (writeField.name().equals(tableField.name())
+            && writeField.dataType() instanceof StructType) {
+          throw new IllegalArgumentException(
+              String.format(
+                  "Cannot run this MERGE/UPDATE on Lance table %s: blob v2 column '%s' would "
+                      + "receive descriptor structs instead of binary data. Assign the column "
+                      + "directly from a blob v2 source column (deep copy) or omit it from the "
+                      + "command",
+                  readOptions.getDatasetUri(), tableField.name()));
+        }
+      }
+    }
   }
 
   @Override
@@ -129,4 +205,8 @@ public NamedReference[] requiredMetadataAttributes() {
   public boolean representUpdateAsDeleteAndInsert() {
     return command != Command.UPDATE;
   }
+
+  Map<String, BlobSourceContext> blobSourceContexts() {
+    return blobSourceContexts;
+  }
 }

lance-spark-3.5_2.12/src/main/java/org/lance/spark/write/SparkPositionDeltaWrite.java

@@ -27,6 +27,8 @@
 import org.lance.spark.LanceSparkCatalogConfig;
 import org.lance.spark.LanceSparkWriteOptions;
 import org.lance.spark.function.LanceFragmentIdWithDefaultFunction;
+import org.lance.spark.utils.BlobReferenceResolver;
+import org.lance.spark.utils.BlobSourceContext;
 import org.lance.spark.utils.Utils;
 
 import com.google.common.collect.ImmutableList;
@@ -94,14 +96,18 @@ public class SparkPositionDeltaWrite implements DeltaWrite, RequiresDistribution
   private final boolean managedVersioning;
   private final boolean hasStableRowIds;
 
+  /** Blob source contexts for resolving copy tokens, keyed by source dataset URI. */
+  private final Map<String, BlobSourceContext> blobSourceContexts;
+
   public SparkPositionDeltaWrite(
       StructType sparkSchema,
       LanceSparkWriteOptions writeOptions,
       Map<String, String> initialStorageOptions,
       String namespaceImpl,
       Map<String, String> namespaceProperties,
       boolean managedVersioning,
-      List<String> tableId) {
+      List<String> tableId,
+      Map<String, BlobSourceContext> blobSourceContexts) {
     this.sparkSchema = sparkSchema;
     try (Dataset ds = Utils.openDatasetBuilder(writeOptions).build()) {
       this.writeOptions = writeOptions.withVersion(ds.version());
@@ -116,6 +122,7 @@ public SparkPositionDeltaWrite(
     this.namespaceProperties = namespaceProperties;
     this.tableId = tableId;
     this.managedVersioning = managedVersioning;
+    this.blobSourceContexts = blobSourceContexts;
   }
 
   @Override
@@ -151,7 +158,8 @@ public DeltaWriterFactory createBatchWriterFactory(PhysicalWriteInfo info) {
           namespaceImpl,
           namespaceProperties,
           tableId,
-          hasStableRowIds);
+          hasStableRowIds,
+          blobSourceContexts);
     }
 
     @Override
@@ -251,6 +259,7 @@ private static class PositionDeltaWriteFactory implements DeltaWriterFactory {
     private final Map<String, String> namespaceProperties;
     private final List<String> tableId;
     private final boolean hasStableRowIds;
+    private final Map<String, BlobSourceContext> blobSourceContexts;
 
     PositionDeltaWriteFactory(
         StructType sparkSchema,
@@ -259,34 +268,44 @@ private static class PositionDeltaWriteFactory implements DeltaWriterFactory {
         String namespaceImpl,
         Map<String, String> namespaceProperties,
         List<String> tableId,
-        boolean hasStableRowIds) {
+        boolean hasStableRowIds,
+        Map<String, BlobSourceContext> blobSourceContexts) {
       this.sparkSchema = sparkSchema;
       this.writeOptions = writeOptions;
       this.initialStorageOptions = initialStorageOptions;
       this.namespaceImpl = namespaceImpl;
       this.namespaceProperties = namespaceProperties;
       this.tableId = tableId;
       this.hasStableRowIds = hasStableRowIds;
+      this.blobSourceContexts = blobSourceContexts;
     }
 
     @Override
     public DeltaWriter<InternalRow> createWriter(int partitionId, long taskId) {
       int batchSize = writeOptions.getBatchSize();
       boolean useQueuedBuffer = writeOptions.isUseQueuedWriteBuffer();
       boolean useLargeVarTypes = writeOptions.isUseLargeVarTypes();
+      long maxBatchBytes = writeOptions.getMaxBatchBytes();
 
       LanceSparkWriteOptions fragmentWriteOptions =
           writeOptions.toBuilder().enableStableRowIds(false).build();
       WriteParams params = fragmentWriteOptions.toWriteParams(initialStorageOptions);
 
+      // One resolver per write task, closed via LanceDataWriter when the task finishes. Always
+      // created so copy tokens resolve even without captured contexts (falls back to open-by-URI).
+      BlobReferenceResolver blobResolver = new BlobReferenceResolver(blobSourceContexts);
+
       // Select buffer type based on configuration
       ArrowBatchWriteBuffer writeBuffer;
       if (useQueuedBuffer) {
         int queueDepth = writeOptions.getQueueDepth();
         writeBuffer =
-            new QueuedArrowBatchWriteBuffer(sparkSchema, batchSize, queueDepth, useLargeVarTypes);
+            new QueuedArrowBatchWriteBuffer(
+                sparkSchema, batchSize, queueDepth, useLargeVarTypes, maxBatchBytes, blobResolver);
       } else {
-        writeBuffer = new SemaphoreArrowBatchWriteBuffer(sparkSchema, batchSize, useLargeVarTypes);
+        writeBuffer =
+            new SemaphoreArrowBatchWriteBuffer(
+                sparkSchema, batchSize, useLargeVarTypes, maxBatchBytes, blobResolver);
       }
 
       // Create fragment in background thread
@@ -305,7 +324,8 @@ public DeltaWriter<InternalRow> createWriter(int partitionId, long taskId) {
 
       return new LanceDeltaWriter(
           writeOptions,
-          new LanceDataWriter(writeBuffer, fragmentCreationTask, fragmentCreationThread),
+          new LanceDataWriter(
+              writeBuffer, fragmentCreationTask, fragmentCreationThread, null, null, blobResolver),
           initialStorageOptions,
           hasStableRowIds);
     }