/*
 * 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.hudi.functional

import org.apache.hudi.HoodieSparkUtils
import org.apache.hudi.common.config.TypedProperties
import org.apache.hudi.common.util.StringUtils
import org.apache.hudi.exception.HoodieException
import org.apache.hudi.functional.TestSparkSqlWithCustomKeyGenerator._
import org.apache.hudi.testutils.HoodieClientTestUtils.createMetaClient
import org.apache.hudi.util.SparkKeyGenUtils

import org.apache.spark.sql.SaveMode
import org.apache.spark.sql.hudi.common.HoodieSparkSqlTestBase
import org.joda.time.DateTime
import org.joda.time.format.DateTimeFormat
import org.junit.jupiter.api.Assertions.{assertEquals, assertFalse, assertTrue}
import org.slf4j.LoggerFactory

/**
 * Tests Spark SQL DML with custom key generator and write configs.
 */
class TestSparkSqlWithCustomKeyGenerator extends HoodieSparkSqlTestBase {
  private val LOG = LoggerFactory.getLogger(getClass)

  test("Test Spark SQL DML with custom key generator") {
    withTempDir { tmp =>
      Seq(
        Seq("COPY_ON_WRITE", "ts:timestamp,segment:simple",
          "(ts=202401, segment='cat2')", "202401/cat2",
          Seq("202312/cat2", "202312/cat4", "202401/cat1", "202401/cat3", "202402/cat1", "202402/cat3", "202402/cat5"),
          TS_FORMATTER_FUNC,
          (ts: Integer, segment: String) => TS_FORMATTER_FUNC.apply(ts) + "/" + segment),
        Seq("MERGE_ON_READ", "segment:simple",
          "(segment='cat3')", "cat3",
          Seq("cat1", "cat2", "cat4", "cat5"),
          TS_TO_STRING_FUNC,
          (_: Integer, segment: String) => segment),
        Seq("MERGE_ON_READ", "ts:timestamp",
          "(ts=202312)", "202312",
          Seq("202401", "202402"),
          TS_TO_STRING_FUNC,
          (ts: Integer, _: String) => TS_FORMATTER_FUNC.apply(ts)),
        Seq("MERGE_ON_READ", "ts:timestamp,segment:simple",
          "(ts=202401, segment='cat2')", "202401/cat2",
          Seq("202312/cat2", "202312/cat4", "202401/cat1", "202401/cat3", "202402/cat1", "202402/cat3", "202402/cat5"),
          TS_TO_STRING_FUNC,
          (ts: Integer, segment: String) => TS_FORMATTER_FUNC.apply(ts) + "/" + segment)
      ).foreach { testParams =>
        withTable(generateTableName) { tableName =>
          LOG.warn("Testing with parameters: " + testParams)
          val tableType = testParams(0).asInstanceOf[String]
          val writePartitionFields = testParams(1).asInstanceOf[String]
          val dropPartitionStatement = testParams(2).asInstanceOf[String]
          val droppedPartition = testParams(3).asInstanceOf[String]
          val expectedPartitions = testParams(4).asInstanceOf[Seq[String]]
          val tsGenFunc = testParams(5).asInstanceOf[Integer => String]
          val partitionGenFunc = testParams(6).asInstanceOf[(Integer, String) => String]
          val tablePath = tmp.getCanonicalPath + "/" + tableName
          val timestampKeyGeneratorConfig = if (writePartitionFields.contains("timestamp")) {
            TS_KEY_GEN_CONFIGS
          } else {
            Map[String, String]()
          }
          val timestampKeyGenProps = if (timestampKeyGeneratorConfig.nonEmpty) {
            ", " + timestampKeyGeneratorConfig.map(e => e._1 + " = '" + e._2 + "'").mkString(", ")
          } else {
            ""
          }

          prepareTableWithKeyGenerator(
            tableName, tablePath, tableType,
            CUSTOM_KEY_GEN_CLASS_NAME, writePartitionFields, timestampKeyGeneratorConfig)

          // SQL CTAS with table properties containing key generator write configs
          createTableWithSql(tableName, tablePath,
            s"hoodie.datasource.write.partitionpath.field = '$writePartitionFields'" + timestampKeyGenProps)

          // Prepare source and test SQL INSERT INTO
          val sourceTableName = tableName + "_source"
          prepareParquetSource(sourceTableName, Seq(
            "(7, 'a7', 1399.0, 1706800227, 'cat1')",
            "(8, 'a8', 26.9, 1706800227, 'cat3')",
            "(9, 'a9', 299.0, 1701443427, 'cat4')"))
          spark.sql(
            s"""
               | INSERT INTO $tableName
               | SELECT * from ${tableName}_source
               | """.stripMargin)
          validateResults(
            tableName,
            s"SELECT id, name, cast(price as string), cast(ts as string), segment from $tableName",
            tsGenFunc,
            partitionGenFunc,
            Seq(),
            Seq(1, "a1", "1.6", 1704121827, "cat1"),
            Seq(2, "a2", "10.8", 1704121827, "cat1"),
            Seq(3, "a3", "30.0", 1706800227, "cat1"),
            Seq(4, "a4", "103.4", 1701443427, "cat2"),
            Seq(5, "a5", "1999.0", 1704121827, "cat2"),
            Seq(6, "a6", "80.0", 1704121827, "cat3"),
            Seq(7, "a7", "1399.0", 1706800227, "cat1"),
            Seq(8, "a8", "26.9", 1706800227, "cat3"),
            Seq(9, "a9", "299.0", 1701443427, "cat4")
          )

          // Test SQL UPDATE
          spark.sql(
            s"""
               | UPDATE $tableName
               | SET price = price + 10.0
               | WHERE id between 4 and 7
               | """.stripMargin)
          validateResults(
            tableName,
            s"SELECT id, name, cast(price as string), cast(ts as string), segment from $tableName",
            tsGenFunc,
            partitionGenFunc,
            Seq(),
            Seq(1, "a1", "1.6", 1704121827, "cat1"),
            Seq(2, "a2", "10.8", 1704121827, "cat1"),
            Seq(3, "a3", "30.0", 1706800227, "cat1"),
            Seq(4, "a4", "113.4", 1701443427, "cat2"),
            Seq(5, "a5", "2009.0", 1704121827, "cat2"),
            Seq(6, "a6", "90.0", 1704121827, "cat3"),
            Seq(7, "a7", "1409.0", 1706800227, "cat1"),
            Seq(8, "a8", "26.9", 1706800227, "cat3"),
            Seq(9, "a9", "299.0", 1701443427, "cat4")
          )

          // Test SQL MERGE INTO
          spark.sql(
            s"""
               | MERGE INTO $tableName as target
               | USING (
               |   SELECT 1 as id, 'a1' as name, 1.6 as price, 1704121827 as ts, 'cat1' as segment, 'delete' as flag
               |   UNION
               |   SELECT 2 as id, 'a2' as name, 11.9 as price, 1704121827 as ts, 'cat1' as segment, '' as flag
               |   UNION
               |   SELECT 6 as id, 'a6' as name, 99.0 as price, 1704121827 as ts, 'cat3' as segment, '' as flag
               |   UNION
               |   SELECT 8 as id, 'a8' as name, 24.9 as price, 1706800227 as ts, 'cat3' as segment, '' as flag
               |   UNION
               |   SELECT 10 as id, 'a10' as name, 888.8 as price, 1706800227 as ts, 'cat5' as segment, '' as flag
               | ) source
               | on target.id = source.id
               | WHEN MATCHED AND flag != 'delete' THEN UPDATE SET
               |   id = source.id, name = source.name, price = source.price, ts = source.ts, segment = source.segment
               | WHEN MATCHED AND flag = 'delete' THEN DELETE
               | WHEN NOT MATCHED THEN INSERT (id, name, price, ts, segment)
               |   values (source.id, source.name, source.price, source.ts, source.segment)
               | """.stripMargin)
          validateResults(
            tableName,
            s"SELECT id, name, cast(price as string), cast(ts as string), segment from $tableName",
            tsGenFunc,
            partitionGenFunc,
            Seq(),
            Seq(2, "a2", "11.9", 1704121827, "cat1"),
            Seq(3, "a3", "30.0", 1706800227, "cat1"),
            Seq(4, "a4", "113.4", 1701443427, "cat2"),
            Seq(5, "a5", "2009.0", 1704121827, "cat2"),
            Seq(6, "a6", "99.0", 1704121827, "cat3"),
            Seq(7, "a7", "1409.0", 1706800227, "cat1"),
            Seq(8, "a8", "24.9", 1706800227, "cat3"),
            Seq(9, "a9", "299.0", 1701443427, "cat4"),
            Seq(10, "a10", "888.8", 1706800227, "cat5")
          )

          // Test SQL DELETE
          spark.sql(
            s"""
               | DELETE FROM $tableName
               | WHERE id = 7
               | """.stripMargin)
          validateResults(
            tableName,
            s"SELECT id, name, cast(price as string), cast(ts as string), segment from $tableName",
            tsGenFunc,
            partitionGenFunc,
            Seq(),
            Seq(2, "a2", "11.9", 1704121827, "cat1"),
            Seq(3, "a3", "30.0", 1706800227, "cat1"),
            Seq(4, "a4", "113.4", 1701443427, "cat2"),
            Seq(5, "a5", "2009.0", 1704121827, "cat2"),
            Seq(6, "a6", "99.0", 1704121827, "cat3"),
            Seq(8, "a8", "24.9", 1706800227, "cat3"),
            Seq(9, "a9", "299.0", 1701443427, "cat4"),
            Seq(10, "a10", "888.8", 1706800227, "cat5")
          )

          // Test DROP PARTITION
          assertTrue(getSortedTablePartitions(tableName).contains(droppedPartition))
          spark.sql(
            s"""
               | ALTER TABLE $tableName DROP PARTITION $dropPartitionStatement
               |""".stripMargin)
          validatePartitions(tableName, Seq(droppedPartition), expectedPartitions)

          if (HoodieSparkUtils.isSpark3) {
            // Test INSERT OVERWRITE, only supported in Spark 3.x
            spark.sql(
              s"""
                 | INSERT OVERWRITE $tableName
                 | SELECT 100 as id, 'a100' as name, 299.0 as price, 1706800227 as ts, 'cat10' as segment
                 | """.stripMargin)
            validateResults(
              tableName,
              s"SELECT id, name, cast(price as string), cast(ts as string), segment from $tableName",
              tsGenFunc,
              partitionGenFunc,
              Seq(),
              Seq(100, "a100", "299.0", 1706800227, "cat10")
            )
          }
        }
      }
    }
  }

  test("Test table property isolation for partition path field config "
    + "with custom key generator for Spark 3.1 and above") {
    // Only testing Spark 3.1 and above as lower Spark versions do not support
    // ALTER TABLE .. SET TBLPROPERTIES .. to store table-level properties in Hudi Catalog
    if (HoodieSparkUtils.gteqSpark3_1) {
      withTempDir { tmp => {
        val tableNameNonPartitioned = generateTableName
        val tableNameSimpleKey = generateTableName
        val tableNameCustom1 = generateTableName
        val tableNameCustom2 = generateTableName

        val tablePathNonPartitioned = tmp.getCanonicalPath + "/" + tableNameNonPartitioned
        val tablePathSimpleKey = tmp.getCanonicalPath + "/" + tableNameSimpleKey
        val tablePathCustom1 = tmp.getCanonicalPath + "/" + tableNameCustom1
        val tablePathCustom2 = tmp.getCanonicalPath + "/" + tableNameCustom2

        val tableType = "MERGE_ON_READ"
        val writePartitionFields1 = "segment:simple"
        val writePartitionFields2 = "ts:timestamp,segment:simple"

        prepareTableWithKeyGenerator(
          tableNameNonPartitioned, tablePathNonPartitioned, tableType,
          NONPARTITIONED_KEY_GEN_CLASS_NAME, "", Map())
        prepareTableWithKeyGenerator(
          tableNameSimpleKey, tablePathSimpleKey, tableType,
          SIMPLE_KEY_GEN_CLASS_NAME, "segment", Map())
        prepareTableWithKeyGenerator(
          tableNameCustom1, tablePathCustom1, tableType,
          CUSTOM_KEY_GEN_CLASS_NAME, writePartitionFields1, Map())
        prepareTableWithKeyGenerator(
          tableNameCustom2, tablePathCustom2, tableType,
          CUSTOM_KEY_GEN_CLASS_NAME, writePartitionFields2, TS_KEY_GEN_CONFIGS)

        // Non-partitioned table does not require additional partition path field write config
        createTableWithSql(tableNameNonPartitioned, tablePathNonPartitioned, "")
        // Partitioned table with simple key generator does not require additional partition path field write config
        createTableWithSql(tableNameSimpleKey, tablePathSimpleKey, "")
        // Partitioned table with custom key generator requires additional partition path field write config
        // Without that, right now the SQL DML fails
        createTableWithSql(tableNameCustom1, tablePathCustom1, "")
        createTableWithSql(tableNameCustom2, tablePathCustom2,
          s"hoodie.datasource.write.partitionpath.field = '$writePartitionFields2', "
            + TS_KEY_GEN_CONFIGS.map(e => e._1 + " = '" + e._2 + "'").mkString(", "))

        val segmentPartitionFunc = (_: Integer, segment: String) => segment
        val customPartitionFunc = (ts: Integer, segment: String) => TS_FORMATTER_FUNC.apply(ts) + "/" + segment

        testFirstRoundInserts(tableNameNonPartitioned, TS_TO_STRING_FUNC, (_, _) => "")
        testFirstRoundInserts(tableNameSimpleKey, TS_TO_STRING_FUNC, segmentPartitionFunc)
        // INSERT INTO should fail for tableNameCustom1
        val sourceTableName = tableNameCustom1 + "_source"
        prepareParquetSource(sourceTableName, Seq("(7, 'a7', 1399.0, 1706800227, 'cat1')"))
        assertThrows[HoodieException] {
          spark.sql(
            s"""
               | INSERT INTO $tableNameCustom1
               | SELECT * from $sourceTableName
               | """.stripMargin)
        }
        testFirstRoundInserts(tableNameCustom2, TS_TO_STRING_FUNC, customPartitionFunc)

        // Now add the missing partition path field write config for tableNameCustom1
        spark.sql(
          s"""ALTER TABLE $tableNameCustom1
             | SET TBLPROPERTIES (hoodie.datasource.write.partitionpath.field = '$writePartitionFields1')
             | """.stripMargin)

        // All tables should be able to do INSERT INTO without any problem,
        // since the scope of the added write config is at the catalog table level
        testSecondRoundInserts(tableNameNonPartitioned, TS_TO_STRING_FUNC, (_, _) => "")
        testSecondRoundInserts(tableNameSimpleKey, TS_TO_STRING_FUNC, segmentPartitionFunc)
        testFirstRoundInserts(tableNameCustom1, TS_TO_STRING_FUNC, segmentPartitionFunc)
        testSecondRoundInserts(tableNameCustom2, TS_TO_STRING_FUNC, customPartitionFunc)
      }
      }
    }
  }

  test("Test wrong partition path field write config with custom key generator") {
    withTempDir { tmp => {
      val tableName = generateTableName
      val tablePath = tmp.getCanonicalPath + "/" + tableName
      val tableType = "MERGE_ON_READ"
      val writePartitionFields = "segment:simple,ts:timestamp"
      val wrongWritePartitionFields = "segment:simple"
      val customPartitionFunc = (ts: Integer, segment: String) => segment + "/" + TS_FORMATTER_FUNC.apply(ts)

      prepareTableWithKeyGenerator(
        tableName, tablePath, "MERGE_ON_READ",
        CUSTOM_KEY_GEN_CLASS_NAME, writePartitionFields, TS_KEY_GEN_CONFIGS)

      // CREATE TABLE should fail due to config conflict
      assertThrows[HoodieException] {
        createTableWithSql(tableName, tablePath,
          s"hoodie.datasource.write.partitionpath.field = '$wrongWritePartitionFields', "
            + TS_KEY_GEN_CONFIGS.map(e => e._1 + " = '" + e._2 + "'").mkString(", "))
      }

      createTableWithSql(tableName, tablePath,
        s"hoodie.datasource.write.partitionpath.field = '$writePartitionFields', "
          + TS_KEY_GEN_CONFIGS.map(e => e._1 + " = '" + e._2 + "'").mkString(", "))
      // Set wrong write config
      spark.sql(
        s"""ALTER TABLE $tableName
           | SET TBLPROPERTIES (hoodie.datasource.write.partitionpath.field = '$wrongWritePartitionFields')
           | """.stripMargin)

      // INSERT INTO should fail due to conflict between write and table config of partition path fields
      val sourceTableName = tableName + "_source"
      prepareParquetSource(sourceTableName, Seq("(7, 'a7', 1399.0, 1706800227, 'cat1')"))
      assertThrows[HoodieException] {
        spark.sql(
          s"""
             | INSERT INTO $tableName
             | SELECT * from $sourceTableName
             | """.stripMargin)
      }

      // Only testing Spark 3.1 and above as lower Spark versions do not support
      // ALTER TABLE .. SET TBLPROPERTIES .. to store table-level properties in Hudi Catalog
      if (HoodieSparkUtils.gteqSpark3_1) {
        // Now fix the partition path field write config for tableName
        spark.sql(
          s"""ALTER TABLE $tableName
             | SET TBLPROPERTIES (hoodie.datasource.write.partitionpath.field = '$writePartitionFields')
             | """.stripMargin)

        // INSERT INTO should succeed now
        testFirstRoundInserts(tableName, TS_TO_STRING_FUNC, customPartitionFunc)
      }
    }
    }
  }

  private def testFirstRoundInserts(tableName: String,
                                    tsGenFunc: Integer => String,
                                    partitionGenFunc: (Integer, String) => String): Unit = {
    val sourceTableName = tableName + "_source1"
    prepareParquetSource(sourceTableName, Seq("(7, 'a7', 1399.0, 1706800227, 'cat1')"))
    spark.sql(
      s"""
         | INSERT INTO $tableName
         | SELECT * from $sourceTableName
         | """.stripMargin)
    validateResults(
      tableName,
      s"SELECT id, name, cast(price as string), cast(ts as string), segment from $tableName",
      tsGenFunc,
      partitionGenFunc,
      Seq(),
      Seq(1, "a1", "1.6", 1704121827, "cat1"),
      Seq(2, "a2", "10.8", 1704121827, "cat1"),
      Seq(3, "a3", "30.0", 1706800227, "cat1"),
      Seq(4, "a4", "103.4", 1701443427, "cat2"),
      Seq(5, "a5", "1999.0", 1704121827, "cat2"),
      Seq(6, "a6", "80.0", 1704121827, "cat3"),
      Seq(7, "a7", "1399.0", 1706800227, "cat1")
    )
  }

  private def testSecondRoundInserts(tableName: String,
                                     tsGenFunc: Integer => String,
                                     partitionGenFunc: (Integer, String) => String): Unit = {
    val sourceTableName = tableName + "_source2"
    prepareParquetSource(sourceTableName, Seq("(8, 'a8', 26.9, 1706800227, 'cat3')"))
    spark.sql(
      s"""
         | INSERT INTO $tableName
         | SELECT * from $sourceTableName
         | """.stripMargin)
    validateResults(
      tableName,
      s"SELECT id, name, cast(price as string), cast(ts as string), segment from $tableName",
      tsGenFunc,
      partitionGenFunc,
      Seq(),
      Seq(1, "a1", "1.6", 1704121827, "cat1"),
      Seq(2, "a2", "10.8", 1704121827, "cat1"),
      Seq(3, "a3", "30.0", 1706800227, "cat1"),
      Seq(4, "a4", "103.4", 1701443427, "cat2"),
      Seq(5, "a5", "1999.0", 1704121827, "cat2"),
      Seq(6, "a6", "80.0", 1704121827, "cat3"),
      Seq(7, "a7", "1399.0", 1706800227, "cat1"),
      Seq(8, "a8", "26.9", 1706800227, "cat3")
    )
  }

  private def prepareTableWithKeyGenerator(tableName: String,
                                           tablePath: String,
                                           tableType: String,
                                           keyGenClassName: String,
                                           writePartitionFields: String,
                                           timestampKeyGeneratorConfig: Map[String, String]): Unit = {
    val df = spark.sql(
      s"""SELECT 1 as id, 'a1' as name, 1.6 as price, 1704121827 as ts, 'cat1' as segment
         | UNION
         | SELECT 2 as id, 'a2' as name, 10.8 as price, 1704121827 as ts, 'cat1' as segment
         | UNION
         | SELECT 3 as id, 'a3' as name, 30.0 as price, 1706800227 as ts, 'cat1' as segment
         | UNION
         | SELECT 4 as id, 'a4' as name, 103.4 as price, 1701443427 as ts, 'cat2' as segment
         | UNION
         | SELECT 5 as id, 'a5' as name, 1999.0 as price, 1704121827 as ts, 'cat2' as segment
         | UNION
         | SELECT 6 as id, 'a6' as name, 80.0 as price, 1704121827 as ts, 'cat3' as segment
         |""".stripMargin)

    df.write.format("hudi")
      .option("hoodie.datasource.write.table.type", tableType)
      .option("hoodie.datasource.write.keygenerator.class", keyGenClassName)
      .option("hoodie.datasource.write.partitionpath.field", writePartitionFields)
      .option("hoodie.datasource.write.recordkey.field", "id")
      .option("hoodie.datasource.write.precombine.field", "name")
      .option("hoodie.table.name", tableName)
      .option("hoodie.insert.shuffle.parallelism", "1")
      .option("hoodie.upsert.shuffle.parallelism", "1")
      .option("hoodie.bulkinsert.shuffle.parallelism", "1")
      .options(timestampKeyGeneratorConfig)
      .mode(SaveMode.Overwrite)
      .save(tablePath)

    // Validate that the generated table has expected table configs of key generator and partition path fields
    val metaClient = createMetaClient(spark, tablePath)
    assertEquals(keyGenClassName, metaClient.getTableConfig.getKeyGeneratorClassName)
    // Validate that that partition path fields in the table config should always
    // contain the field names only (no key generator type like "segment:simple")
    if (CUSTOM_KEY_GEN_CLASS_NAME.equals(keyGenClassName)) {
      val props = new TypedProperties()
      props.put("hoodie.datasource.write.partitionpath.field", writePartitionFields)
      timestampKeyGeneratorConfig.foreach(e => {
        props.put(e._1, e._2)
      })
      // For custom key generator, the "hoodie.datasource.write.partitionpath.field"
      // contains the key generator type, like "ts:timestamp,segment:simple",
      // whereas the partition path fields in table config is "ts,segment"
      assertEquals(
        SparkKeyGenUtils.getPartitionColumns(Option(CUSTOM_KEY_GEN_CLASS_NAME), props),
        metaClient.getTableConfig.getPartitionFieldProp)
    } else {
      assertEquals(writePartitionFields, metaClient.getTableConfig.getPartitionFieldProp)
    }
  }

  private def createTableWithSql(tableName: String,
                                 tablePath: String,
                                 tblProps: String): Unit = {
    val tblPropsStatement = if (StringUtils.isNullOrEmpty(tblProps)) {
      ""
    } else {
      "TBLPROPERTIES (\n" + tblProps + "\n)"
    }
    spark.sql(
      s"""
         | CREATE TABLE $tableName USING HUDI
         | location '$tablePath'
         | $tblPropsStatement
         | """.stripMargin)
  }

  private def prepareParquetSource(sourceTableName: String,
                                   rows: Seq[String]): Unit = {
    spark.sql(
      s"""CREATE TABLE $sourceTableName
         | (id int, name string, price decimal(5, 1), ts int, segment string)
         | USING PARQUET
         |""".stripMargin)
    spark.sql(
      s"""
         | INSERT INTO $sourceTableName values
         | ${rows.mkString(", ")}
         | """.stripMargin)
  }

  private def validateResults(tableName: String,
                              sql: String,
                              tsGenFunc: Integer => String,
                              partitionGenFunc: (Integer, String) => String,
                              droppedPartitions: Seq[String],
                              expects: Seq[Any]*): Unit = {
    checkAnswer(sql)(
      expects.map(e => Seq(e(0), e(1), e(2), tsGenFunc.apply(e(3).asInstanceOf[Integer]), e(4))): _*
    )
    val expectedPartitions: Seq[String] = expects
      .map(e => partitionGenFunc.apply(e(3).asInstanceOf[Integer], e(4).asInstanceOf[String]))
      .distinct.sorted
    validatePartitions(tableName, droppedPartitions, expectedPartitions)
  }

  private def getSortedTablePartitions(tableName: String): Seq[String] = {
    spark.sql(s"SHOW PARTITIONS $tableName").collect()
      .map(row => row.getString(0))
      .sorted.toSeq
  }

  private def validatePartitions(tableName: String,
                                 droppedPartitions: Seq[String],
                                 expectedPartitions: Seq[String]): Unit = {
    val actualPartitions: Seq[String] = getSortedTablePartitions(tableName)
    if (expectedPartitions.size == 1 && expectedPartitions.head.isEmpty) {
      assertTrue(actualPartitions.isEmpty)
    } else {
      assertEquals(expectedPartitions, actualPartitions)
    }
    droppedPartitions.foreach(dropped => assertFalse(actualPartitions.contains(dropped)))
  }
}

object TestSparkSqlWithCustomKeyGenerator {
  val SIMPLE_KEY_GEN_CLASS_NAME = "org.apache.hudi.keygen.SimpleKeyGenerator"
  val NONPARTITIONED_KEY_GEN_CLASS_NAME = "org.apache.hudi.keygen.NonpartitionedKeyGenerator"
  val CUSTOM_KEY_GEN_CLASS_NAME = "org.apache.hudi.keygen.CustomKeyGenerator"
  val DATE_FORMAT_PATTERN = "yyyyMM"
  val TS_KEY_GEN_CONFIGS = Map(
    "hoodie.keygen.timebased.timestamp.type" -> "SCALAR",
    "hoodie.keygen.timebased.output.dateformat" -> DATE_FORMAT_PATTERN,
    "hoodie.keygen.timebased.timestamp.scalar.time.unit" -> "seconds"
  )
  val TS_TO_STRING_FUNC = (tsSeconds: Integer) => tsSeconds.toString
  val TS_FORMATTER_FUNC = (tsSeconds: Integer) => {
    new DateTime(tsSeconds * 1000L).toString(DateTimeFormat.forPattern(DATE_FORMAT_PATTERN))
  }

  def getTimestampKeyGenConfigs: Map[String, String] = {
    Map(
      "hoodie.keygen.timebased.timestamp.type" -> "SCALAR",
      "hoodie.keygen.timebased.output.dateformat" -> DATE_FORMAT_PATTERN,
      "hoodie.keygen.timebased.timestamp.scalar.time.unit" -> "seconds"
    )
  }
}
