spark overwrite功能：支持带schema的表、分区子表、包含特殊字符的表名；spark支持repartition

Author: zgdgod

hologres-connector-spark-2.x/README.md

@@ -8,7 +8,7 @@
     <parent>
         <groupId>com.alibaba.hologres</groupId>
         <artifactId>hologres-connector-parent</artifactId>
-        <version>1.4.0-SNAPSHOT</version>
+        <version>1.4.1-SNAPSHOT</version>
     </parent>
 
     <artifactId>hologres-connector-spark-2.x</artifactId>

hologres-connector-spark-2.x/pom.xml

@@ -26,15 +26,21 @@ class HoloWriter(
   logger.info("HoloWriter begin: " + LocalDateTime.now())
   val hologresConfigs: HologresConfigs = new HologresConfigs(holoOptions)
   var is_overwrite: Boolean = mode == SaveMode.Overwrite
+  private var partitionInfo: (String, String) = _
+
   if (is_overwrite) {
-    hologresConfigs.tempTableForOverwrite = hologresConfigs.table + new Random().nextInt(Int.MaxValue) + "_temp"
+    hologresConfigs.tempTableForOverwrite = JDBCUtil.generateTempTableNameForOverwrite(hologresConfigs)
     JDBCUtil.createTempTableForOverWrite(hologresConfigs)
   }
 
   override def commit(messages: Array[WriterCommitMessage]): Unit = {
     logger.info("HoloWriter commit: " + LocalDateTime.now())
     if (is_overwrite) {
-      JDBCUtil.renameTempTableForOverWrite(hologresConfigs)
+      if (partitionInfo.eq(null)) {
+        JDBCUtil.renameTempTableForOverWrite(hologresConfigs)
+      } else {
+        JDBCUtil.renameTempTableForOverWrite(hologresConfigs, partitionInfo._1, partitionInfo._2)
+      }
     }
   }
 
@@ -56,12 +62,15 @@ class HoloWriter(
   override def createWriterFactory(): DataWriterFactory[InternalRow] = {
     val holoClient: HoloClient = new HoloClient(hologresConfigs.holoConfig)
     try {
-      var holoSchema: TableSchema = null
+      var holoSchema = holoClient.getTableSchema(TableName.valueOf(hologresConfigs.table))
+      if (holoSchema.isPartitionParentTable && is_overwrite) {
+        throw new IOException("Partition parent table can not be insert overwrite now.")
+      }
+      partitionInfo = JDBCUtil.getChildTablePartitionInfo(hologresConfigs)
+
       if (is_overwrite) {
         // insert overwrite 会先写在一张临时表中，写入成功时替换原表。
         holoSchema = holoClient.getTableSchema(TableName.valueOf(hologresConfigs.tempTableForOverwrite))
-      } else {
-        holoSchema = holoClient.getTableSchema(TableName.valueOf(hologresConfigs.table))
       }
 
       var holoVersion: HoloVersion = null

hologres-connector-spark-2.x/src/main/scala/com/alibaba/hologres/spark2/sink/HoloWriter.scala

@@ -234,6 +234,100 @@ class SparkHoloReadWriteSuite extends SparkHoloSuiteBase {
   }
 
 
+  test("SaveMode = overwrite child table with schema") {
+    val parentTable = "test.\"Table-Parent\""
+    val partitionValue = "20240527"
+    val table = "test.\"Table-Child_20240527\""
+
+    testUtils.dropTable(parentTable)
+    testUtils.createSchema("test")
+    testUtils.createPartitionTable(defaultCreateHoloParentTableDDL, parentTable, table, partitionValue)
+
+    val byteA = Array(4.toByte, 5.toByte, 6.toByte, 'q'.toByte, 'e'.toByte)
+    val intA = Array(4, 5, 6)
+    val doubleA = Array(2.333, 3.444, 4.555)
+    val date = Date.valueOf("2024-05-27")
+
+    val data1 = Seq(
+      Row(0L, -7L, 20, "phone1", 6.7F, Timestamp.valueOf("2021-03-29 00:00:00"), byteA, intA, doubleA, date),
+      Row(1L, 6L, -30, "phone2", 7.8F, Timestamp.valueOf("2021-04-01 12:00:00"), byteA, intA, doubleA, date)
+    )
+
+    val data2 = Seq(
+      Row(0L, -7L, 20, "phone1", 6.7F, Timestamp.valueOf("2021-03-29 00:00:00"), byteA, intA, doubleA, date),
+      Row(1L, -7L, 20, "phone1", 6.7F, Timestamp.valueOf("2021-03-29 00:00:00"), byteA, intA, doubleA, date),
+      Row(2L, 6L, -30, "phone2", 7.8F, Timestamp.valueOf("2021-04-01 12:00:00"), byteA, intA, doubleA, date),
+      Row(3L, 6L, -30, "phone2", 7.8F, Timestamp.valueOf("2021-04-01 12:00:00"), byteA, intA, doubleA, date)
+    )
+
+    val newSchema = StructType(Array(
+      StructField("pk", LongType),
+      StructField("id", LongType),
+      StructField("count", IntegerType),
+      StructField("name", StringType),
+      StructField("thick", FloatType),
+      StructField("time", TimestampType),
+      StructField("by", BinaryType),
+      StructField("inta", ArrayType(IntegerType)),
+      StructField("doublea", ArrayType(DoubleType)),
+      StructField("dt", DateType)
+    ))
+
+    var df = spark.createDataFrame(
+      spark.sparkContext.parallelize(data1),
+      newSchema
+    ).orderBy("pk").cache()
+
+    df.write
+      .format("hologres")
+      .option(SourceProvider.USERNAME, testUtils.username)
+      .option(SourceProvider.PASSWORD, testUtils.password)
+      .option(SourceProvider.JDBCURL, testUtils.jdbcUrl)
+      .option(SourceProvider.TABLE, table)
+      .option(SourceProvider.WRITE_MODE, "insertOrUpdate")
+      .option(SourceProvider.COPY_WRITE_MODE, "true")
+      .option(SourceProvider.ENABLE_TARGET_SHARDS, "true")
+      .option(SourceProvider.COPY_WRITE_DIRTY_DATA_CHECK, "true")
+      .mode(SaveMode.Overwrite)
+      .save()
+
+    df = spark.createDataFrame(
+      spark.sparkContext.parallelize(data2),
+      newSchema
+    ).orderBy("pk").cache()
+
+    df.write
+      .format("hologres")
+      .option(SourceProvider.USERNAME, testUtils.username)
+      .option(SourceProvider.PASSWORD, testUtils.password)
+      .option(SourceProvider.JDBCURL, testUtils.jdbcUrl)
+      .option(SourceProvider.TABLE, table)
+      .option(SourceProvider.WRITE_MODE, "insertOrUpdate")
+      .option(SourceProvider.COPY_WRITE_MODE, "true")
+      .option(SourceProvider.BULK_LOAD, "true")
+      .option(SourceProvider.ENABLE_TARGET_SHARDS, "true")
+      .option(SourceProvider.COPY_WRITE_DIRTY_DATA_CHECK, "true")
+      .mode(SaveMode.Overwrite)
+      .save()
+
+    val readDf = spark.read
+      .format("hologres")
+      .schema(newSchema) // 指定读取哪些字段
+      .option(SourceProvider.USERNAME, testUtils.username)
+      .option(SourceProvider.PASSWORD, testUtils.password)
+      .option(SourceProvider.JDBCURL, testUtils.jdbcUrl)
+      .option(SourceProvider.TABLE, table)
+      .load().orderBy("pk").cache()
+
+    assert(df.count() == 4)
+    // compare read and write
+    if (df.except(readDf).count() > 0) {
+      df.show()
+      readDf.show()
+      throw new Exception("The data read is inconsistent with the data written！！！")
+    }
+  }
+
   test("write or read not exists columns.") {
     val table = "table_for_holo_test_1"
     val data = Seq(

hologres-connector-spark-2.x/src/test/scala/com/alibaba/hologres/spark2/SparkHoloReadWriteSuite.scala

@@ -78,4 +78,32 @@ abstract class SparkHoloSuiteBase extends QueryTest with SharedSparkSession {
     "    json_column json," +
     "    jsonb_column jsonb," +
     "    rb_column roaringbitmap);"
+
+  val defaultCreateHoloParentTableDDL =
+    "create table PARENT_TABLE_NAME (" +
+      "    pk bigint," +
+      "    st smallint," +
+      "    id bigint," +
+      "    count int," +
+      "    name text," +
+      "    price numeric(38, 12)," +
+      "    out_of_stock bool," +
+      "    weight double precision," +
+      "    thick float4," +
+      "    time timestamptz," +
+      "    dt date," +
+      "    by bytea," +
+      "    inta int4[]," +
+      "    longa int8[]," +
+      "    floata float4[]," +
+      "    doublea float8[]," +
+      "    boola boolean[]," +
+      "    stringa text[]," +
+      "    json_column json," +
+      "    jsonb_column jsonb," +
+      "    rb_column roaringbitmap" +
+      // "    primary key(pk, dt)" +
+      ") PARTITION BY LIST(dt);\n" +
+      "CREATE TABLE TABLE_NAME PARTITION OF PARENT_TABLE_NAME FOR VALUES IN ('PARTITION_VALUE');"
+
 }