Add support for Spark Overwrite save mode

spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerErrorCode.java

@@ -24,6 +24,7 @@ public enum SpannerErrorCode {
   DECIMAL_OUT_OF_RANGE(6),
   INVALID_ARGUMENT(7),
   SCHEMA_VALIDATION_ERROR(8),
+  DDL_EXCEPTION(9),
   // Should be last
   UNSUPPORTED(9998),
   UNKNOWN(9999);

spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerSchemaConverter.java

@@ -0,0 +1,127 @@
+// Copyright 2026 Google LLC
+//
+// 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 com.google.cloud.spark.spanner;
+
+import com.google.cloud.spanner.Dialect;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.stream.Collectors;
+import org.apache.spark.sql.types.DataType;
+import org.apache.spark.sql.types.DataTypes;
+import org.apache.spark.sql.types.DecimalType;
+import org.apache.spark.sql.types.StructField;
+import org.apache.spark.sql.types.StructType;
+
+public class SpannerSchemaConverter {
+
+  private final Dialect dialect;
+
+  public SpannerSchemaConverter(Dialect dialect) {
+    this.dialect = dialect;
+  }
+
+  public String sparkSchemaToSpannerDDL(StructType schema, String tableName) {
+    List<String> colDefs = new ArrayList<>();
+    List<String> pkCols = new ArrayList<>();
+    String quote = getQuote();
+
+    for (StructField field : schema.fields()) {
+      String colName = field.name();
+      String spannerType = sparkTypeToSpannerType(field.dataType());
+      String suffix = "";
+      if (field.metadata().contains("pk") && field.metadata().getBoolean("pk")) {
+        pkCols.add(colName);
+      }
+      if (!field.nullable()) {
+        suffix = " NOT NULL";
+      }
+      colDefs.add(quote + colName + quote + " " + spannerType + suffix);
+    }
+
+    String pkDef =
+        "PRIMARY KEY ("
+            + pkCols.stream().map(c -> quote + c + quote).collect(Collectors.joining(", "))
+            + ")";
+
+    return "CREATE TABLE "
+        + quote
+        + tableName
+        + quote
+        + " ("
+        + String.join(", ", colDefs)
+        + ") "
+        + pkDef;
+  }
+
+  private String getQuote() {
+    if (this.dialect == Dialect.POSTGRESQL) {
+      return "\"";
+    }
+    return "`";
+  }
+
+  public String sparkTypeToSpannerType(DataType sparkType) {
+    if (sparkType instanceof DecimalType) {
+      return "NUMERIC";
+    }
+    if (dialect == Dialect.POSTGRESQL) {
+      if (sparkType.equals(DataTypes.LongType)) {
+        return "int8";
+      }
+      if (sparkType.equals(DataTypes.StringType)) {
+        return "varchar";
+      }
+      if (sparkType.equals(DataTypes.BooleanType)) {
+        return "bool";
+      }
+      if (sparkType.equals(DataTypes.DoubleType)) {
+        return "float8";
+      }
+      if (sparkType.equals(DataTypes.BinaryType)) {
+        return "bytea";
+      }
+      if (sparkType.equals(DataTypes.TimestampType)) {
+        return "timestamptz";
+      }
+      if (sparkType.equals(DataTypes.DateType)) {
+        return "date";
+      }
+    }
+    // Default to Google Standard SQL
+    if (sparkType.equals(DataTypes.LongType)) {
+      return "INT64";
+    }
+    if (sparkType.equals(DataTypes.StringType)) {
+      return "STRING(MAX)";
+    }
+    if (sparkType.equals(DataTypes.BooleanType)) {
+      return "BOOL";
+    }
+    if (sparkType.equals(DataTypes.DoubleType)) {
+      return "FLOAT64";
+    }
+    if (sparkType.equals(DataTypes.BinaryType)) {
+      return "BYTES(MAX)";
+    }
+    if (sparkType.equals(DataTypes.TimestampType)) {
+      return "TIMESTAMP";
+    }
+    if (sparkType.equals(DataTypes.DateType)) {
+      return "DATE";
+    }
+    // Fallback for unknown types.
+    return "STRING(MAX)";
+  }
+}

spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerTable.java

@@ -51,9 +51,9 @@ public class SpannerTable implements Table, SupportsRead, SupportsWrite {
   private final SpannerTableSchema dbSchema;
   private final StructType sparkSchema;
   private static final ImmutableSet<TableCapability> tableCapabilities =
-      ImmutableSet.of(TableCapability.BATCH_READ, TableCapability.BATCH_WRITE);
+      ImmutableSet.of(
+          TableCapability.BATCH_READ, TableCapability.BATCH_WRITE, TableCapability.TRUNCATE);
   private final Map<String, String> properties;
-
   private static final Logger log = LoggerFactory.getLogger(SpannerTable.class);
 
   public SpannerTable(Map<String, String> properties) {

spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerUtils.java

@@ -64,6 +64,7 @@
 import org.apache.spark.sql.types.Decimal;
 import org.apache.spark.sql.types.StructField;
 import org.apache.spark.sql.types.StructType;
+import org.apache.spark.sql.util.CaseInsensitiveStringMap;
 import org.apache.spark.unsafe.types.UTF8String;
 import org.threeten.bp.Duration;
 
@@ -190,6 +191,21 @@ public static BatchClientWithCloser batchClientFromProperties(Map<String, String
 
   public static BatchClientWithCloser batchClientFromProperties(
       Map<String, String> properties, SessionPoolOptions sessionPoolOptions) {
+    SpannerOptions options = buildSpannerOptions(properties, sessionPoolOptions);
+    Spanner spanner = options.getService();
+    DatabaseId databaseId =
+        DatabaseId.of(
+            options.getProjectId(), properties.get("instanceId"), properties.get("databaseId"));
+    return new BatchClientWithCloser(
+        spanner, spanner.getBatchClient(databaseId), spanner.getDatabaseClient(databaseId));
+  }
+
+  public static SpannerOptions buildSpannerOptions(CaseInsensitiveStringMap properties) {
+    return buildSpannerOptions(properties, defaultSessionPoolOptions);
+  }
+
+  public static SpannerOptions buildSpannerOptions(
+      Map<String, String> properties, SessionPoolOptions sessionPoolOptions) {
     SpannerOptions.Builder builder =
         SpannerOptions.newBuilder()
             .setSessionPoolOption(sessionPoolOptions)
@@ -222,13 +238,7 @@ public static BatchClientWithCloser batchClientFromProperties(
     }
     System.setProperty("com.google.cloud.spanner.watchdogTimeoutSeconds", "7200");
 
-    SpannerOptions options = builder.build();
-    Spanner spanner = options.getService();
-    DatabaseId databaseId =
-        DatabaseId.of(
-            options.getProjectId(), properties.get("instanceId"), properties.get("databaseId"));
-    return new BatchClientWithCloser(
-        spanner, spanner.getBatchClient(databaseId), spanner.getDatabaseClient(databaseId));
+    return builder.build();
   }
 
   public static void toSparkDecimal(GenericInternalRow dest, java.math.BigDecimal v, int at) {
@@ -503,7 +513,7 @@ public static StructType pruneSchema(
     return prunedSchema;
   }
 
-  static String getRequiredOption(Map<String, String> properties, String option) {
+  public static String getRequiredOption(Map<String, String> properties, String option) {
     String tableName = properties.get(option);
     if (tableName == null) {
       throw new SpannerConnectorException(

spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerWriteBuilder.java

@@ -14,19 +14,121 @@
 
 package com.google.cloud.spark.spanner;
 
+import com.google.cloud.spanner.DatabaseAdminClient;
+import com.google.cloud.spanner.DatabaseClient;
+import com.google.cloud.spanner.DatabaseId;
+import com.google.cloud.spanner.Dialect;
+import com.google.cloud.spanner.Spanner;
+import com.google.cloud.spanner.SpannerException;
+import com.google.cloud.spanner.Statement;
+import com.google.cloud.spanner.connection.Connection;
+import java.util.Arrays;
+import java.util.concurrent.ExecutionException;
 import org.apache.spark.sql.connector.write.BatchWrite;
 import org.apache.spark.sql.connector.write.LogicalWriteInfo;
+import org.apache.spark.sql.connector.write.SupportsTruncate;
 import org.apache.spark.sql.connector.write.WriteBuilder;
+import org.apache.spark.sql.types.StructType;
+import org.apache.spark.sql.util.CaseInsensitiveStringMap;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
-public class SpannerWriteBuilder implements WriteBuilder {
+public class SpannerWriteBuilder implements WriteBuilder, SupportsTruncate {
+
+  private static final Logger log = LoggerFactory.getLogger(SpannerWriteBuilder.class);
   private final LogicalWriteInfo info;
+  private final StructType schema;
 
   public SpannerWriteBuilder(LogicalWriteInfo info) {
     this.info = info;
+    this.schema = info.schema();
   }
 
   @Override
   public BatchWrite buildForBatch() {
     return new SpannerBatchWrite(info);
   }
+
+  @Override
+  public WriteBuilder truncate() {
+    CaseInsensitiveStringMap opts = new CaseInsensitiveStringMap(this.info.options());
+    String overwriteMode = opts.getOrDefault("overwriteMode", "truncate");
+
+    if (overwriteMode.equalsIgnoreCase("recreate")) {
+      recreateTable(opts);
+    } else {
+      truncateTable(opts);
+    }
+
+    return this;
+  }
+
+  private void recreateTable(CaseInsensitiveStringMap opts) {
+    String instanceId = SpannerUtils.getRequiredOption(opts, "instanceId");
+    String databaseId = SpannerUtils.getRequiredOption(opts, "databaseId");
+    String tableName = SpannerUtils.getRequiredOption(opts, "table");
+
+    try (Spanner spanner = SpannerUtils.buildSpannerOptions(opts).getService()) {
+      DatabaseAdminClient dbAdminClient = spanner.getDatabaseAdminClient();
+      Dialect dialect;
+      try (Connection conn = SpannerUtils.connectionFromProperties(opts.asCaseSensitiveMap())) {
+        dialect = conn.getDialect();
+      }
+      // TODO Re-use drop table and create table code from SpannerCatalog
+      // Drop the table.
+      dbAdminClient
+          .updateDatabaseDdl(
+              instanceId, databaseId, Arrays.asList("DROP TABLE `" + tableName + "`"), null)
+          .get();
+
+      // Create the table.
+      SpannerSchemaConverter converter = new SpannerSchemaConverter(dialect);
+      String createTableDdl = converter.sparkSchemaToSpannerDDL(this.schema, tableName);
+      dbAdminClient
+          .updateDatabaseDdl(instanceId, databaseId, Arrays.asList(createTableDdl), null)
+          .get();
+
+    } catch (InterruptedException | ExecutionException e) {
+      throw new SpannerConnectorException(
+          SpannerErrorCode.DDL_EXCEPTION, "Error recreating table " + tableName, e);
+    }
+  }
+
+  private void truncateTable(CaseInsensitiveStringMap opts) {
+    String projectId = SpannerUtils.getRequiredOption(opts, "projectId");
+    String instanceId = SpannerUtils.getRequiredOption(opts, "instanceId");
+    String databaseId = SpannerUtils.getRequiredOption(opts, "databaseId");
+    String tableName = SpannerUtils.getRequiredOption(opts, "table");
+
+    try (Spanner spanner = SpannerUtils.buildSpannerOptions(opts).getService()) {
+      DatabaseClient dbClient =
+          spanner.getDatabaseClient(DatabaseId.of(projectId, instanceId, databaseId));
+      truncateTable(dbClient, tableName);
+    } catch (Exception e) {
+      throw new SpannerConnectorException(
+          SpannerErrorCode.DDL_EXCEPTION, "Error truncating table " + tableName, e);
+    }
+  }
+
+  private long truncateTable(DatabaseClient dbClient, String tableName) {
+
+    // 1. Construct the DML Statement
+    // Spanner requires a WHERE clause for PDML, even if you are deleting everything.
+    String sql = "DELETE FROM `" + tableName.replace("`", "``") + "` WHERE true";
+    Statement statement = Statement.of(sql);
+
+    try {
+      // 2. Execute the Partitioned Update
+      // This is a blocking call. The Spanner client will divide the table into
+      // partitions and run concurrent background transactions to delete the data.
+      long deletedRowCount = dbClient.executePartitionedUpdate(statement);
+      log.info("Successfully deleted " + deletedRowCount + " rows.");
+      return deletedRowCount;
+
+    } catch (SpannerException e) {
+      // SpannerExceptions wrap underlying gRPC errors (e.g., DEADLINE_EXCEEDED, PERMISSION_DENIED)
+      log.error("Failed to execute Partitioned DML on table: " + tableName, e);
+      throw e;
+    }
+  }
 }