From 7be6f38c6c2ca390e86c41ca36d7684bc1be0d2a Mon Sep 17 00:00:00 2001 From: Max Ksyunz Date: Mon, 2 Mar 2026 21:47:41 -0800 Subject: [PATCH 01/10] Implement overwrite mode as `delete from ... where true` sent as partitioned DML. Signed-off-by: Max Ksyunz --- .../cloud/spark/spanner/SpannerTable.java | 2 +- .../spark/spanner/SpannerWriteBuilder.java | 60 ++++++++++++++++++- 2 files changed, 60 insertions(+), 2 deletions(-) diff --git a/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerTable.java b/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerTable.java index 729daaac..2afe61e0 100644 --- a/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerTable.java +++ b/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerTable.java @@ -54,7 +54,7 @@ public class SpannerTable implements Table, SupportsRead, SupportsWrite { private final SpannerTableSchema dbSchema; private final @Nullable StructType dfSchema; private static final ImmutableSet tableCapabilities = - ImmutableSet.of(TableCapability.BATCH_READ, TableCapability.BATCH_WRITE); + ImmutableSet.of(TableCapability.BATCH_READ, TableCapability.BATCH_WRITE, TableCapability.TRUNCATE); private final CaseInsensitiveStringMap properties; private static final Logger log = LoggerFactory.getLogger(SpannerTable.class); diff --git a/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerWriteBuilder.java b/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerWriteBuilder.java index 92fcc79f..8313a2b5 100644 --- a/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerWriteBuilder.java +++ b/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerWriteBuilder.java @@ -14,11 +14,18 @@ package com.google.cloud.spark.spanner; +import com.google.cloud.spanner.DatabaseClient; +import com.google.cloud.spanner.DatabaseId; +import com.google.cloud.spanner.Spanner; +import com.google.cloud.spanner.SpannerException; +import com.google.cloud.spanner.Statement; 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.util.CaseInsensitiveStringMap; -public class SpannerWriteBuilder implements WriteBuilder { +public class SpannerWriteBuilder implements WriteBuilder, SupportsTruncate { private final LogicalWriteInfo info; public SpannerWriteBuilder(LogicalWriteInfo info) { @@ -29,4 +36,55 @@ public SpannerWriteBuilder(LogicalWriteInfo info) { public BatchWrite buildForBatch() { return new SpannerBatchWrite(info); } + + @Override + public WriteBuilder truncate() { + CaseInsensitiveStringMap opts = new CaseInsensitiveStringMap(this.info.options()); + 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); + } + return this; + } + + /** + * Deletes all rows from a specified Spanner table using Partitioned DML. + * * @param dbClient The initialized Cloud Spanner DatabaseClient. + * + * @param tableName The name of the table to truncate. + * @return The total number of rows deleted. + */ + public static 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 + " WHERE true"; + Statement statement = Statement.of(sql); + + System.out.println("Starting Partitioned DML execution: " + 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); + + System.out.println("Successfully deleted " + deletedRowCount + " rows."); + return deletedRowCount; + + } catch (SpannerException e) { + // SpannerExceptions wrap underlying gRPC errors (e.g., DEADLINE_EXCEEDED, PERMISSION_DENIED) + System.err.println("Failed to execute Partitioned DML on table: " + tableName); + throw e; + } + } } From 78b312ef2d472b08f659ebd8f2d4b9d662fede53 Mon Sep 17 00:00:00 2001 From: Max Ksyunz Date: Tue, 3 Mar 2026 18:23:05 -0800 Subject: [PATCH 02/10] Fix-up formatting Signed-off-by: Max Ksyunz --- .../java/com/google/cloud/spark/spanner/SpannerWriteBuilder.java | 1 + 1 file changed, 1 insertion(+) diff --git a/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerWriteBuilder.java b/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerWriteBuilder.java index 8313a2b5..69f96229 100644 --- a/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerWriteBuilder.java +++ b/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerWriteBuilder.java @@ -57,6 +57,7 @@ public WriteBuilder truncate() { } /** + * * Deletes all rows from a specified Spanner table using Partitioned DML. * * @param dbClient The initialized Cloud Spanner DatabaseClient. * From ff9c8ec56b20c72550e35a8c89b69e0eacb8dc3c Mon Sep 17 00:00:00 2001 From: Max Ksyunz Date: Tue, 3 Mar 2026 18:26:28 -0800 Subject: [PATCH 03/10] Fix-up formatting II Signed-off-by: Max Ksyunz --- .../com/google/cloud/spark/spanner/SpannerWriteBuilder.java | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerWriteBuilder.java b/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerWriteBuilder.java index 69f96229..6ed2e203 100644 --- a/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerWriteBuilder.java +++ b/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerWriteBuilder.java @@ -57,9 +57,9 @@ public WriteBuilder truncate() { } /** - * * Deletes all rows from a specified Spanner table using Partitioned DML. - * * @param dbClient The initialized Cloud Spanner DatabaseClient. + * @param dbClient The + * initialized Cloud Spanner DatabaseClient. * * @param tableName The name of the table to truncate. * @return The total number of rows deleted. From f6063ad52b92f9ae72c4bbbea4a719ef3930e4ba Mon Sep 17 00:00:00 2001 From: Max Ksyunz Date: Tue, 3 Mar 2026 18:31:09 -0800 Subject: [PATCH 04/10] No comments, no problems. spotless:apply was not fixing error that spotless:check was reporting. Signed-off-by: Max Ksyunz --- .../google/cloud/spark/spanner/SpannerWriteBuilder.java | 8 -------- 1 file changed, 8 deletions(-) diff --git a/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerWriteBuilder.java b/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerWriteBuilder.java index 6ed2e203..cdb9b159 100644 --- a/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerWriteBuilder.java +++ b/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerWriteBuilder.java @@ -56,14 +56,6 @@ public WriteBuilder truncate() { return this; } - /** - * Deletes all rows from a specified Spanner table using Partitioned DML. - * @param dbClient The - * initialized Cloud Spanner DatabaseClient. - * - * @param tableName The name of the table to truncate. - * @return The total number of rows deleted. - */ public static long truncateTable(DatabaseClient dbClient, String tableName) { // 1. Construct the DML Statement From 897dc6887bf8f5e63e10fca674bd3b6ff7dd7f52 Mon Sep 17 00:00:00 2001 From: Max Ksyunz Date: Tue, 3 Mar 2026 19:19:10 -0800 Subject: [PATCH 05/10] Update spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerWriteBuilder.java Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- .../com/google/cloud/spark/spanner/SpannerWriteBuilder.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerWriteBuilder.java b/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerWriteBuilder.java index cdb9b159..05d638ee 100644 --- a/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerWriteBuilder.java +++ b/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerWriteBuilder.java @@ -60,7 +60,7 @@ public static 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 + " WHERE true"; + String sql = "DELETE FROM `" + tableName.replace("`", "``") + "` WHERE true"; Statement statement = Statement.of(sql); System.out.println("Starting Partitioned DML execution: " + sql); From 4565ae7afbd0986e2808e7ad99ab4aff828eeec2 Mon Sep 17 00:00:00 2001 From: Max Ksyunz Date: Thu, 5 Mar 2026 12:05:47 -0800 Subject: [PATCH 06/10] Replace print statements with logs. Signed-off-by: Max Ksyunz --- .../cloud/spark/spanner/SpannerWriteBuilder.java | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerWriteBuilder.java b/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerWriteBuilder.java index 05d638ee..8d6c5f10 100644 --- a/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerWriteBuilder.java +++ b/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerWriteBuilder.java @@ -24,8 +24,12 @@ import org.apache.spark.sql.connector.write.SupportsTruncate; import org.apache.spark.sql.connector.write.WriteBuilder; import org.apache.spark.sql.util.CaseInsensitiveStringMap; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; public class SpannerWriteBuilder implements WriteBuilder, SupportsTruncate { + + private static final Logger log = LoggerFactory.getLogger(SpannerWriteBuilder.class); private final LogicalWriteInfo info; public SpannerWriteBuilder(LogicalWriteInfo info) { @@ -56,27 +60,24 @@ public WriteBuilder truncate() { return this; } - public static long truncateTable(DatabaseClient dbClient, String tableName) { + 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); - System.out.println("Starting Partitioned DML execution: " + 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); - - System.out.println("Successfully deleted " + deletedRowCount + " rows."); + log.info("Successfully deleted " + deletedRowCount + " rows."); return deletedRowCount; } catch (SpannerException e) { // SpannerExceptions wrap underlying gRPC errors (e.g., DEADLINE_EXCEEDED, PERMISSION_DENIED) - System.err.println("Failed to execute Partitioned DML on table: " + tableName); + log.error("Failed to execute Partitioned DML on table: " + tableName, e); throw e; } } From 820feee175b685c151ac62543e2fc453e6450cda Mon Sep 17 00:00:00 2001 From: Max Ksyunz Date: Thu, 5 Mar 2026 22:39:47 -0800 Subject: [PATCH 07/10] add support for overwriteMode option for Overwrite save mode, support "recreate" Signed-off-by: Max Ksyunz --- .../spark/spanner/SpannerSchemaConverter.java | 127 ++++++++++++++++++ .../spark/spanner/SpannerWriteBuilder.java | 52 ++++++- .../integration/WriteIntegrationTest.java | 12 ++ 3 files changed, 190 insertions(+), 1 deletion(-) create mode 100644 spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerSchemaConverter.java diff --git a/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerSchemaConverter.java b/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerSchemaConverter.java new file mode 100644 index 00000000..c82c6181 --- /dev/null +++ b/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 colDefs = new ArrayList<>(); + List 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)"; + } +} diff --git a/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerWriteBuilder.java b/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerWriteBuilder.java index 8d6c5f10..c6a75a09 100644 --- a/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerWriteBuilder.java +++ b/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerWriteBuilder.java @@ -14,15 +14,21 @@ 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; @@ -31,9 +37,11 @@ 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 @@ -44,6 +52,49 @@ public BatchWrite buildForBatch() { @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"); @@ -57,7 +108,6 @@ public WriteBuilder truncate() { throw new SpannerConnectorException( SpannerErrorCode.DDL_EXCEPTION, "Error truncating table " + tableName, e); } - return this; } private long truncateTable(DatabaseClient dbClient, String tableName) { diff --git a/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/integration/WriteIntegrationTest.java b/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/integration/WriteIntegrationTest.java index 356e181d..956c1aa2 100644 --- a/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/integration/WriteIntegrationTest.java +++ b/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/integration/WriteIntegrationTest.java @@ -28,6 +28,11 @@ import com.google.cloud.spanner.ErrorCode; import com.google.cloud.spanner.SpannerException; import com.google.cloud.spark.spanner.SpannerCatalog; +import com.google.cloud.spanner.DatabaseAdminClient; +import com.google.cloud.spanner.Dialect; +import com.google.cloud.spanner.Spanner; +import com.google.cloud.spark.spanner.SpannerSchemaConverter; +import com.google.cloud.spark.spanner.SpannerUtils; import com.google.cloud.spark.spanner.TestData; import java.math.BigDecimal; import java.math.MathContext; @@ -44,9 +49,11 @@ import org.apache.spark.sql.RowFactory; import org.apache.spark.sql.SaveMode; import org.apache.spark.sql.types.DataTypes; +import org.apache.spark.sql.types.Metadata; import org.apache.spark.sql.types.StructField; import org.apache.spark.sql.types.StructType; import org.junit.Assert; +import org.apache.spark.sql.util.CaseInsensitiveStringMap; import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.Parameterized; @@ -114,6 +121,11 @@ public WriteIntegrationTest(boolean usePostgresSql) { this.usePostgresSql = usePostgresSql; } + @Test + public void testOverwriteRecreateMode() { + + } + @Override protected boolean getUsePostgreSql() { return usePostgresSql; From 8fba7efbc557e1aa6558737aeb4a00af2903bff6 Mon Sep 17 00:00:00 2001 From: Max Ksyunz Date: Thu, 5 Mar 2026 23:30:24 -0800 Subject: [PATCH 08/10] Update spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerWriteBuilder.java Temporary code. will be merged with catalog branch. Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- .../com/google/cloud/spark/spanner/SpannerWriteBuilder.java | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerWriteBuilder.java b/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerWriteBuilder.java index c6a75a09..056e84a2 100644 --- a/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerWriteBuilder.java +++ b/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerWriteBuilder.java @@ -76,9 +76,11 @@ private void recreateTable(CaseInsensitiveStringMap opts) { } // TODO Re-use drop table and create table code from SpannerCatalog // Drop the table. + // Drop the table. + String quote = dialect == Dialect.POSTGRESQL ? "\"" : "`"; dbAdminClient .updateDatabaseDdl( - instanceId, databaseId, Arrays.asList("DROP TABLE `" + tableName + "`"), null) + instanceId, databaseId, Arrays.asList("DROP TABLE " + quote + tableName + quote), null) .get(); // Create the table. From 511c411e5a658b18f3a6b3efec5875c6862c5a47 Mon Sep 17 00:00:00 2001 From: Max Ksyunz Date: Fri, 6 Mar 2026 10:11:06 -0800 Subject: [PATCH 09/10] Rework recreate mode to use same DDL generation code as SpannerCatalog class. Signed-off-by: Max Ksyunz --- .../cloud/spark/spanner/SpannerCatalog.java | 4 +- .../spanner/SpannerInformationSchema.java | 10 +- .../spark/spanner/SpannerSchemaConverter.java | 127 ------------------ .../cloud/spark/spanner/SpannerTable.java | 3 +- .../spark/spanner/SpannerWriteBuilder.java | 38 +++--- .../spark/spanner/SpannerCatalogTest.java | 4 +- .../integration/WriteIntegrationTest.java | 106 ++++++++++++++- 7 files changed, 132 insertions(+), 160 deletions(-) delete mode 100644 spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerSchemaConverter.java diff --git a/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerCatalog.java b/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerCatalog.java index 30c9817d..8ed8f4d4 100644 --- a/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerCatalog.java +++ b/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerCatalog.java @@ -167,7 +167,7 @@ public Table createTable( DatabaseClient dbClient = getDatabaseClient(); Dialect dialect = dbClient.getDialect(); SpannerInformationSchema schemaInfo = createSchemaInfo(dialect); - String ddl = schemaInfo.toDdl(ident, schema); + String ddl = schemaInfo.createTableDdl(ident, schema); DatabaseAdminClient dbAdminClient = spanner.getDatabaseAdminClient(); OperationFuture op = dbAdminClient.updateDatabaseDdl( @@ -216,7 +216,7 @@ public Table alterTable(Identifier ident, TableChange... changes) { public boolean dropTable(Identifier ident) { DatabaseClient dbClient = getDatabaseClient(); SpannerInformationSchema schemaInfo = createSchemaInfo(dbClient.getDialect()); - String ddl = "DROP TABLE " + schemaInfo.quoteIdentifier(ident.name()); + String ddl = schemaInfo.dropTableDdl(ident.name()); DatabaseAdminClient dbAdminClient = spanner.getDatabaseAdminClient(); OperationFuture op = diff --git a/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerInformationSchema.java b/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerInformationSchema.java index 95641134..6a1446fd 100644 --- a/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerInformationSchema.java +++ b/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerInformationSchema.java @@ -37,7 +37,7 @@ public interface SpannerInformationSchema { String sparkTypeToSpannerType(StructField field); - default String toDdl(Identifier ident, StructType schema) { + default String createTableDdl(Identifier ident, StructType schema) { StringBuilder ddl = new StringBuilder(); ddl.append("CREATE TABLE ").append(quoteIdentifier(ident.name())).append(" ("); for (StructField field : schema.fields()) { @@ -70,6 +70,14 @@ default String toDdl(Identifier ident, StructType schema) { return ddl.toString(); } + default Statement truncateTableDml(String tableName) { + return Statement.of("DELETE FROM " + quoteIdentifier(tableName) + " WHERE true"); + } + + default String dropTableDdl(String tableName) { + return "DROP TABLE " + quoteIdentifier(tableName); + } + static SpannerInformationSchema create(Dialect dialect) { switch (dialect) { case POSTGRESQL: diff --git a/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerSchemaConverter.java b/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerSchemaConverter.java deleted file mode 100644 index c82c6181..00000000 --- a/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerSchemaConverter.java +++ /dev/null @@ -1,127 +0,0 @@ -// 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 colDefs = new ArrayList<>(); - List 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)"; - } -} diff --git a/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerTable.java b/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerTable.java index 2afe61e0..b786f283 100644 --- a/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerTable.java +++ b/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerTable.java @@ -54,7 +54,8 @@ public class SpannerTable implements Table, SupportsRead, SupportsWrite { private final SpannerTableSchema dbSchema; private final @Nullable StructType dfSchema; private static final ImmutableSet tableCapabilities = - ImmutableSet.of(TableCapability.BATCH_READ, TableCapability.BATCH_WRITE, TableCapability.TRUNCATE); + ImmutableSet.of( + TableCapability.BATCH_READ, TableCapability.BATCH_WRITE, TableCapability.TRUNCATE); private final CaseInsensitiveStringMap properties; private static final Logger log = LoggerFactory.getLogger(SpannerTable.class); diff --git a/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerWriteBuilder.java b/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerWriteBuilder.java index 056e84a2..87ddfb4e 100644 --- a/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerWriteBuilder.java +++ b/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerWriteBuilder.java @@ -22,8 +22,9 @@ 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.Collections; import java.util.concurrent.ExecutionException; +import org.apache.spark.sql.connector.catalog.Identifier; import org.apache.spark.sql.connector.write.BatchWrite; import org.apache.spark.sql.connector.write.LogicalWriteInfo; import org.apache.spark.sql.connector.write.SupportsTruncate; @@ -74,20 +75,18 @@ private void recreateTable(CaseInsensitiveStringMap opts) { try (Connection conn = SpannerUtils.connectionFromProperties(opts.asCaseSensitiveMap())) { dialect = conn.getDialect(); } - // TODO Re-use drop table and create table code from SpannerCatalog - // Drop the table. - // Drop the table. - String quote = dialect == Dialect.POSTGRESQL ? "\"" : "`"; + SpannerInformationSchema schemaInfo = SpannerInformationSchema.create(dialect); + + String dropDdl = schemaInfo.dropTableDdl(tableName); dbAdminClient - .updateDatabaseDdl( - instanceId, databaseId, Arrays.asList("DROP TABLE " + quote + tableName + quote), null) + .updateDatabaseDdl(instanceId, databaseId, Collections.singletonList(dropDdl), null) .get(); // Create the table. - SpannerSchemaConverter converter = new SpannerSchemaConverter(dialect); - String createTableDdl = converter.sparkSchemaToSpannerDDL(this.schema, tableName); + Identifier ident = Identifier.of(new String[0], tableName); + String createDdl = schemaInfo.createTableDdl(ident, this.schema); dbAdminClient - .updateDatabaseDdl(instanceId, databaseId, Arrays.asList(createTableDdl), null) + .updateDatabaseDdl(instanceId, databaseId, Collections.singletonList(createDdl), null) .get(); } catch (InterruptedException | ExecutionException e) { @@ -105,30 +104,29 @@ private void truncateTable(CaseInsensitiveStringMap opts) { try (Spanner spanner = SpannerUtils.buildSpannerOptions(opts).getService()) { DatabaseClient dbClient = spanner.getDatabaseClient(DatabaseId.of(projectId, instanceId, databaseId)); - truncateTable(dbClient, tableName); + Dialect dialect = dbClient.getDialect(); + + SpannerInformationSchema informationSchema = SpannerInformationSchema.create(dialect); + + truncateTable(dbClient, tableName, informationSchema); } catch (Exception e) { throw new SpannerConnectorException( SpannerErrorCode.DDL_EXCEPTION, "Error truncating table " + tableName, e); } } - private long truncateTable(DatabaseClient dbClient, String tableName) { + private long truncateTable( + DatabaseClient dbClient, String tableName, SpannerInformationSchema informationSchema) { - // 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); + Statement statement = informationSchema.truncateTableDml(tableName); 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. + // Execute partitioned update. This is a blocking call. 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; } diff --git a/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/SpannerCatalogTest.java b/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/SpannerCatalogTest.java index 72a1992e..7ff6d447 100644 --- a/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/SpannerCatalogTest.java +++ b/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/SpannerCatalogTest.java @@ -207,7 +207,7 @@ public void createTableShouldThrowExceptionOnNoPrimaryKey() { thrown.expectMessage( "No primary key found for table no_pk_table. Please specify at least one primary key column."); - SpannerInformationSchema.create(dialect).toDdl(ident, schema); + SpannerInformationSchema.create(dialect).createTableDdl(ident, schema); } @Test @@ -279,7 +279,7 @@ public void testToDdl() { new StructField("price", DataTypes.createDecimalType(10, 2), true, Metadata.empty()), }); - String ddl = SpannerInformationSchema.create(dialect).toDdl(ident, schema); + String ddl = SpannerInformationSchema.create(dialect).createTableDdl(ident, schema); if (dialect == Dialect.POSTGRESQL) { assertEquals( diff --git a/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/integration/WriteIntegrationTest.java b/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/integration/WriteIntegrationTest.java index 956c1aa2..b18611f5 100644 --- a/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/integration/WriteIntegrationTest.java +++ b/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/integration/WriteIntegrationTest.java @@ -28,11 +28,6 @@ import com.google.cloud.spanner.ErrorCode; import com.google.cloud.spanner.SpannerException; import com.google.cloud.spark.spanner.SpannerCatalog; -import com.google.cloud.spanner.DatabaseAdminClient; -import com.google.cloud.spanner.Dialect; -import com.google.cloud.spanner.Spanner; -import com.google.cloud.spark.spanner.SpannerSchemaConverter; -import com.google.cloud.spark.spanner.SpannerUtils; import com.google.cloud.spark.spanner.TestData; import java.math.BigDecimal; import java.math.MathContext; @@ -49,11 +44,9 @@ import org.apache.spark.sql.RowFactory; import org.apache.spark.sql.SaveMode; import org.apache.spark.sql.types.DataTypes; -import org.apache.spark.sql.types.Metadata; import org.apache.spark.sql.types.StructField; import org.apache.spark.sql.types.StructType; import org.junit.Assert; -import org.apache.spark.sql.util.CaseInsensitiveStringMap; import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.Parameterized; @@ -123,7 +116,106 @@ public WriteIntegrationTest(boolean usePostgresSql) { @Test public void testOverwriteRecreateMode() { + String tableName = TestData.WRITE_TABLE_NAME + "_RECREATE"; + spark.sql("DROP TABLE IF EXISTS spanner." + tableName); + + // 1. Define schema with primary key metadata (needed for table recreation) + StructType schema = + new StructType( + new StructField[] { + DataTypes.createStructField( + "long_col", DataTypes.LongType, false, SpannerCatalog.PRIMARY_KEY_METADATA), + DataTypes.createStructField("string_col", DataTypes.StringType, true), + }); + + Map props = connectionProperties(usePostgresSql); + props.put("table", tableName); + + // 2. Write initial data (creates the table via ErrorIfExists) + List initialRows = + Arrays.asList(RowFactory.create(1L, "initial-one"), RowFactory.create(2L, "initial-two")); + Dataset initialDf = spark.createDataFrame(initialRows, schema); + initialDf.write().format("cloud-spanner").options(props).mode(SaveMode.ErrorIfExists).save(); + + // 3. Verify initial data + Dataset dfAfterInitialWrite = spark.read().format("cloud-spanner").options(props).load(); + assertEquals(2, dfAfterInitialWrite.count()); + + // 4. Overwrite with recreate mode + List newRows = + Arrays.asList( + RowFactory.create(3L, "new-three"), + RowFactory.create(4L, "new-four"), + RowFactory.create(5L, "new-five")); + Dataset newDf = spark.createDataFrame(newRows, schema); + + Map overwriteProps = connectionProperties(usePostgresSql); + overwriteProps.put("table", tableName); + overwriteProps.put("overwriteMode", "recreate"); + + newDf.write().format("cloud-spanner").options(overwriteProps).mode(SaveMode.Overwrite).save(); + + // 5. Verify only new data exists + Dataset finalDf = spark.read().format("cloud-spanner").options(props).load(); + assertEquals(3, finalDf.count()); + + Map finalRows = + finalDf.collectAsList().stream() + .collect(java.util.stream.Collectors.toMap(r -> r.getLong(0), r -> r)); + + assertThat(finalRows.get(3L).getString(1)).isEqualTo("new-three"); + assertThat(finalRows.get(4L).getString(1)).isEqualTo("new-four"); + assertThat(finalRows.get(5L).getString(1)).isEqualTo("new-five"); + } + + @Test + public void testOverwriteTruncateMode() { + String tableName = TestData.WRITE_TABLE_NAME + "_TRUNCATE"; + spark.sql("DROP TABLE IF EXISTS spanner." + tableName); + + // 1. Define schema with primary key metadata (needed for initial table creation) + StructType schema = + new StructType( + new StructField[] { + DataTypes.createStructField( + "long_col", DataTypes.LongType, false, SpannerCatalog.PRIMARY_KEY_METADATA), + DataTypes.createStructField("string_col", DataTypes.StringType, true), + }); + + Map props = connectionProperties(usePostgresSql); + props.put("table", tableName); + + // 2. Write initial data (creates the table via ErrorIfExists) + List initialRows = + Arrays.asList(RowFactory.create(1L, "initial-one"), RowFactory.create(2L, "initial-two")); + Dataset initialDf = spark.createDataFrame(initialRows, schema); + initialDf.write().format("cloud-spanner").options(props).mode(SaveMode.ErrorIfExists).save(); + + // 3. Verify initial data + Dataset dfAfterInitialWrite = spark.read().format("cloud-spanner").options(props).load(); + assertEquals(2, dfAfterInitialWrite.count()); + + // 4. Overwrite with default truncate mode + List newRows = + Arrays.asList( + RowFactory.create(3L, "new-three"), + RowFactory.create(4L, "new-four"), + RowFactory.create(5L, "new-five")); + Dataset newDf = spark.createDataFrame(newRows, schema); + + newDf.write().format("cloud-spanner").options(props).mode(SaveMode.Overwrite).save(); + + // 5. Verify only new data exists + Dataset finalDf = spark.read().format("cloud-spanner").options(props).load(); + assertEquals(3, finalDf.count()); + + Map finalRows = + finalDf.collectAsList().stream() + .collect(java.util.stream.Collectors.toMap(r -> r.getLong(0), r -> r)); + assertThat(finalRows.get(3L).getString(1)).isEqualTo("new-three"); + assertThat(finalRows.get(4L).getString(1)).isEqualTo("new-four"); + assertThat(finalRows.get(5L).getString(1)).isEqualTo("new-five"); } @Override From c87acb371488880cc819b7a96ad3fd941716173f Mon Sep 17 00:00:00 2001 From: Max Ksyunz Date: Fri, 6 Mar 2026 12:22:51 -0800 Subject: [PATCH 10/10] Update spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerWriteBuilder.java Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- .../google/cloud/spark/spanner/SpannerWriteBuilder.java | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerWriteBuilder.java b/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerWriteBuilder.java index 87ddfb4e..23fad8f5 100644 --- a/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerWriteBuilder.java +++ b/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerWriteBuilder.java @@ -57,8 +57,14 @@ public WriteBuilder truncate() { if (overwriteMode.equalsIgnoreCase("recreate")) { recreateTable(opts); - } else { + } else if (overwriteMode.equalsIgnoreCase("truncate")) { truncateTable(opts); + } else { + throw new SpannerConnectorException( + SpannerErrorCode.INVALID_ARGUMENT, + "Unsupported overwriteMode '" + + overwriteMode + + "'. Supported modes are 'recreate' and 'truncate'."); } return this;