[SPARK-41982][SQL] Partitions of type string should not be treated as numeric types

docs/sql-migration-guide.md

@@ -36,6 +36,7 @@ license: |
   - Since Spark 3.4, Spark throws only `PartitionsAlreadyExistException` when it creates partitions but some of them exist already. In Spark 3.3 or earlier, Spark can throw either `PartitionsAlreadyExistException` or `PartitionAlreadyExistsException`.
   - Since Spark 3.4, Spark will do validation for partition spec in ALTER PARTITION to follow the behavior of `spark.sql.storeAssignmentPolicy` which may cause an exception if type conversion fails, e.g. `ALTER TABLE .. ADD PARTITION(p='a')` if column `p` is int type. To restore the legacy behavior, set `spark.sql.legacy.skipTypeValidationOnAlterPartition` to `true`.
   - Since Spark 3.4, vectorized readers are enabled by default for the nested data types (array, map and struct). To restore the legacy behavior, set `spark.sql.orc.enableNestedColumnVectorizedReader` and `spark.sql.parquet.enableNestedColumnVectorizedReader` to `false`.
+  - Since Spark 3.4, When inserting(adding/renaming/dropping) a partition, the partition field is of `String` and the partition value is not quoted, it will be treated as a field of `Numeric` by default. If you want to keep the partition field of `String`, we can set `spark.sql.legacy.keepPartitionSpecAsStringLiteral` to `true`.
 
 ## Upgrading from Spark SQL 3.2 to 3.3
 

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/parser/AstBuilder.scala

@@ -507,14 +507,19 @@ class AstBuilder extends SqlBaseParserBaseVisitor[AnyRef] with SQLConfHelper wit
       ctx: PartitionSpecContext): Map[String, Option[String]] = withOrigin(ctx) {
     val legacyNullAsString =
       conf.getConf(SQLConf.LEGACY_PARSE_NULL_PARTITION_SPEC_AS_STRING_LITERAL)
+    val keepPartitionSpecAsString =
+      conf.getConf(SQLConf.LEGACY_KEEP_PARTITION_SPEC_AS_STRING_LITERAL)
+
     val parts = ctx.partitionVal.asScala.map { pVal =>
       // Check if the query attempted to refer to a DEFAULT column value within the PARTITION clause
       // and return a specific error to help guide the user, since this is not allowed.
       if (pVal.DEFAULT != null) {
         throw QueryParsingErrors.defaultColumnReferencesNotAllowedInPartitionSpec(ctx)
       }
       val name = pVal.identifier.getText
-      val value = Option(pVal.constant).map(v => visitStringConstant(v, legacyNullAsString))
+      val value = Option(pVal.constant).map(v => {
+        visitStringConstant(v, legacyNullAsString, keepPartitionSpecAsString)
+      })
       name -> value
     }
     // Before calling `toMap`, we check duplicated keys to avoid silently ignore partition values
@@ -546,14 +551,19 @@ class AstBuilder extends SqlBaseParserBaseVisitor[AnyRef] with SQLConfHelper wit
    */
   protected def visitStringConstant(
       ctx: ConstantContext,
-      legacyNullAsString: Boolean): String = withOrigin(ctx) {
+      legacyNullAsString: Boolean = false,
+      keepPartitionSpecAsString: Boolean = false): String = withOrigin(ctx) {
     expression(ctx) match {
       case Literal(null, _) if !legacyNullAsString => null
       case l @ Literal(null, _) => l.toString
       case l: Literal =>
-        // TODO For v2 commands, we will cast the string back to its actual value,
-        //  which is a waste and can be improved in the future.
-        Cast(l, StringType, Some(conf.sessionLocalTimeZone)).eval().toString
+        if (keepPartitionSpecAsString && !ctx.isInstanceOf[StringLiteralContext]) {
+          ctx.getText
+        } else {
+          // TODO For v2 commands, we will cast the string back to its actual value,
+          //  which is a waste and can be improved in the future.
+          Cast(l, StringType, Some(conf.sessionLocalTimeZone)).eval().toString
+        }
       case other =>
         throw new IllegalArgumentException(s"Only literals are allowed in the " +
           s"partition spec, but got ${other.sql}")

sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala

@@ -3288,6 +3288,15 @@ object SQLConf {
       .booleanConf
       .createWithDefault(false)
 
+  val LEGACY_KEEP_PARTITION_SPEC_AS_STRING_LITERAL =
+    buildConf("spark.sql.legacy.keepPartitionSpecAsStringLiteral")
+      .internal()
+      .doc("If it is set to true and the type of the partition is string, " +
+        "partition value will be treated as string value but not numeric value")
+      .version("3.4.0")
+      .booleanConf
+      .createWithDefault(false)
+
   val LEGACY_REPLACE_DATABRICKS_SPARK_AVRO_ENABLED =
     buildConf("spark.sql.legacy.replaceDatabricksSparkAvro.enabled")
       .internal()

sql/core/src/main/scala/org/apache/spark/sql/execution/SparkSqlParser.scala

@@ -363,7 +363,7 @@ class SparkSqlAstBuilder extends AstBuilder {
    * Convert a constants list into a String sequence.
    */
   override def visitConstantList(ctx: ConstantListContext): Seq[String] = withOrigin(ctx) {
-    ctx.constant.asScala.map(v => visitStringConstant(v, legacyNullAsString = false)).toSeq
+    ctx.constant.asScala.map(v => visitStringConstant(v)).toSeq
   }
 
   /**

sql/core/src/test/scala/org/apache/spark/sql/SQLInsertTestSuite.scala

@@ -351,6 +351,35 @@ trait SQLInsertTestSuite extends QueryTest with SQLTestUtils {
       }
     }
   }
+
+  test("SPARK-41982: " +
+    "treat the partition field as a string type when keepPartitionSpecAsStringLiteral is enabled") {
+    withSQLConf(SQLConf.LEGACY_KEEP_PARTITION_SPEC_AS_STRING_LITERAL.key -> "true") {
+      withTable("t") {
+        sql("create table t(i string, j int) using orc partitioned by (dt string)")
+        sql("insert into t partition(dt=08) values('a', 10)")
+        Seq(
+          "select * from t where dt='08'",
+          "select * from t where dt=08"
+        ).foreach { query =>
+          checkAnswer(sql(query), Seq(Row("a", 10, "08")))
+        }
+        val e = intercept[AnalysisException](sql("alter table t drop partition(dt='8')"))
+        assert(e.getMessage.contains("PARTITIONS_NOT_FOUND"))
+      }
+    }
+
+    withSQLConf(SQLConf.LEGACY_KEEP_PARTITION_SPEC_AS_STRING_LITERAL.key -> "false") {
+      withTable("t") {
+        sql("create table t(i string, j int) using orc partitioned by (dt string)")
+        sql("insert into t partition(dt=08) values('a', 10)")
+        checkAnswer(sql("select * from t where dt='08'"), sql("select * from t where dt='07'"))
+        checkAnswer(sql("select * from t where dt=08"), Seq(Row("a", 10, "8")))
+        val e = intercept[AnalysisException](sql("alter table t drop partition(dt='08')"))
+        assert(e.getMessage.contains("PARTITIONS_NOT_FOUND"))
+      }
+    }
+  }
 }
 
 class FileSourceSQLInsertTestSuite extends SQLInsertTestSuite with SharedSparkSession {

sql/core/src/test/scala/org/apache/spark/sql/execution/command/AlterTableAddPartitionSuiteBase.scala

@@ -257,4 +257,26 @@ trait AlterTableAddPartitionSuiteBase extends QueryTest with DDLCommandTestUtils
       }
     }
   }
+
+  test("SPARK-41982: add partition when keepPartitionSpecAsString set `true`") {
+    withSQLConf(SQLConf.LEGACY_KEEP_PARTITION_SPEC_AS_STRING_LITERAL.key -> "true") {
+      withNamespaceAndTable("ns", "tbl") { t =>
+        sql(s"CREATE TABLE $t(name STRING, age INT) USING PARQUET PARTITIONED BY (dt STRING)")
+        sql(s"ALTER TABLE $t ADD PARTITION(dt = 08)")
+        checkPartitions(t, Map("dt" -> "08"))
+        sql(s"ALTER TABLE $t ADD PARTITION(dt = '09')")
+        checkPartitions(t, Map("dt" -> "09"), Map("dt" -> "08"))
+      }
+    }
+
+    withSQLConf(SQLConf.LEGACY_KEEP_PARTITION_SPEC_AS_STRING_LITERAL.key -> "false") {
+      withNamespaceAndTable("ns", "tb2") { t =>
+        sql(s"CREATE TABLE $t(name STRING, age INT) USING PARQUET PARTITIONED BY (dt STRING)")
+        sql(s"ALTER TABLE $t ADD PARTITION(dt = 08)")
+        checkPartitions(t, Map("dt" -> "8"))
+        sql(s"ALTER TABLE $t ADD PARTITION(dt = '09')")
+        checkPartitions(t, Map("dt" -> "09"), Map("dt" -> "8"))
+      }
+    }
+  }
 }

sql/core/src/test/scala/org/apache/spark/sql/execution/command/AlterTableDropPartitionSuiteBase.scala

@@ -254,4 +254,34 @@ trait AlterTableDropPartitionSuiteBase extends QueryTest with DDLCommandTestUtil
       checkPartitions(t)
     }
   }
+
+  test("SPARK-41982: drop partition when keepPartitionSpecAsString set `true`") {
+    withSQLConf(SQLConf.LEGACY_KEEP_PARTITION_SPEC_AS_STRING_LITERAL.key -> "true") {
+      withNamespaceAndTable("ns", "tbl") { t =>
+        sql(s"CREATE TABLE $t(name STRING, age INT) using orc PARTITIONED BY (dt STRING)")
+        sql(s"ALTER TABLE $t ADD PARTITION(dt = 08)")
+        checkPartitions(t, Map("dt" -> "08"))
+        sql(s"ALTER TABLE $t DROP PARTITION (dt = 08)")
+        checkPartitions(t)
+        sql(s"ALTER TABLE $t ADD PARTITION(dt = '08')")
+        checkPartitions(t, Map("dt" -> "08"))
+        sql(s"ALTER TABLE $t DROP PARTITION (dt = '08')")
+        checkPartitions(t)
+      }
+    }
+
+    withSQLConf(SQLConf.LEGACY_KEEP_PARTITION_SPEC_AS_STRING_LITERAL.key -> "false") {
+      withNamespaceAndTable("ns", "tb2") { t =>
+        sql(s"CREATE TABLE $t(name STRING, age INT) using orc PARTITIONED BY (dt STRING)")
+        sql(s"ALTER TABLE $t ADD PARTITION(dt = 08)")
+        checkPartitions(t, Map("dt" -> "8"))
+        sql(s"ALTER TABLE $t DROP PARTITION (dt = 08)")
+        checkPartitions(t)
+        sql(s"ALTER TABLE $t ADD PARTITION(dt = 08)")
+        checkPartitions(t, Map("dt" -> "8"))
+        sql(s"ALTER TABLE $t DROP PARTITION (dt = 8)")
+        checkPartitions(t)
+      }
+    }
+  }
 }

sql/core/src/test/scala/org/apache/spark/sql/execution/command/AlterTableRenamePartitionSuiteBase.scala

@@ -242,4 +242,40 @@ trait AlterTableRenamePartitionSuiteBase extends QueryTest with DDLCommandTestUt
       checkPartitions(t, Map("part" -> "2020-01-02"))
     }
   }
+
+  test("SPARK-41982: rename partition when keepPartitionSpecAsString set `true`") {
+    withSQLConf(SQLConf.LEGACY_KEEP_PARTITION_SPEC_AS_STRING_LITERAL.key -> "true") {
+      withNamespaceAndTable("ns", "tbl") { t =>
+        sql(s"CREATE TABLE $t(name STRING, age INT) USING PARQUET PARTITIONED BY (dt STRING)")
+        sql(s"ALTER TABLE $t ADD PARTITION(dt = 08)")
+        checkPartitions(t, Map("dt" -> "08"))
+        sql(s"ALTER TABLE $t PARTITION (dt = 08)" +
+          s" RENAME TO PARTITION (dt = 09)")
+        checkPartitions(t, Map("dt" -> "09"))
+        sql(s"ALTER TABLE $t PARTITION (dt = 09)" +
+          s" RENAME TO PARTITION (dt = '08')")
+        checkPartitions(t, Map("dt" -> "08"))
+        sql(s"ALTER TABLE $t PARTITION (dt = '08')" +
+          s" RENAME TO PARTITION (dt = '09')")
+        checkPartitions(t, Map("dt" -> "09"))
+      }
+    }
+
+    withSQLConf(SQLConf.LEGACY_KEEP_PARTITION_SPEC_AS_STRING_LITERAL.key -> "false") {
+      withNamespaceAndTable("ns", "tb2") { t =>
+        sql(s"CREATE TABLE $t(name STRING, age INT) USING PARQUET PARTITIONED BY (dt STRING)")
+        sql(s"ALTER TABLE $t ADD PARTITION(dt = 08)")
+        checkPartitions(t, Map("dt" -> "8"))
+        sql(s"ALTER TABLE $t PARTITION (dt = 08)" +
+          s" RENAME TO PARTITION (dt = 09)")
+        checkPartitions(t, Map("dt" -> "9"))
+        sql(s"ALTER TABLE $t PARTITION (dt = 09)" +
+          s" RENAME TO PARTITION (dt = '08')")
+        checkPartitions(t, Map("dt" -> "08"))
+        sql(s"ALTER TABLE $t PARTITION (dt = '08')" +
+          s" RENAME TO PARTITION (dt = '09')")
+        checkPartitions(t, Map("dt" -> "09"))
+      }
+    }
+  }
 }