[SPARK-51830] Exception handling for partition datatype conversion call

Author: Madhukar525722

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/PartitioningUtils.scala

@@ -38,6 +38,7 @@ import org.apache.spark.sql.catalyst.expressions.{Attribute, Cast, Literal}
 import org.apache.spark.sql.catalyst.types.DataTypeUtils
 import org.apache.spark.sql.catalyst.util.{CaseInsensitiveMap, DateFormatter, DateTimeUtils, TimeFormatter, TimestampFormatter}
 import org.apache.spark.sql.errors.{QueryCompilationErrors, QueryExecutionErrors}
+import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.types._
 import org.apache.spark.sql.util.SchemaUtils
 import org.apache.spark.unsafe.types.UTF8String
@@ -361,12 +362,21 @@ object PartitioningUtils extends SQLConfHelper {
     }.mkString("/")
   }
 
-  def removeLeadingZerosFromNumberTypePartition(value: String, dataType: DataType): String =
-    dataType match {
-      case ByteType | ShortType | IntegerType | LongType | FloatType | DoubleType =>
-        Option(castPartValueToDesiredType(dataType, value, null)).map(_.toString).orNull
-      case _ => value
+  def removeLeadingZerosFromNumberTypePartition(value: String, dataType: DataType): String = {
+    try {
+      dataType match {
+        case ByteType | ShortType | IntegerType | LongType | FloatType | DoubleType =>
+          Option(castPartValueToDesiredType(dataType, value, null)).map(_.toString).orNull
+        case _ => value
+      }
+    } catch {
+      case NonFatal(_) =>
+        if (SQLConf.get.validatePartitionColumns) {
+          throw QueryExecutionErrors.failedToCastValueToDataTypeForPartitionColumnError(
+              value, dataType, null)
+        } else value
     }
+  }
 
   def getPathFragment(spec: TablePartitionSpec, partitionColumns: Seq[Attribute]): String = {
     getPathFragment(spec, DataTypeUtils.fromAttributes(partitionColumns))

sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetPartitionDiscoverySuite.scala

@@ -34,6 +34,7 @@ import org.apache.spark.sql.catalyst.catalog.ExternalCatalogUtils
 import org.apache.spark.sql.catalyst.types.DataTypeUtils
 import org.apache.spark.sql.catalyst.util.{DateFormatter, DateTimeUtils, TimeFormatter, TimestampFormatter}
 import org.apache.spark.sql.catalyst.util.DateTimeUtils.localDateTimeToMicros
+import org.apache.spark.sql.errors.QueryExecutionErrors
 import org.apache.spark.sql.execution.datasources._
 import org.apache.spark.sql.execution.datasources.{PartitionPath => Partition}
 import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2Relation, FileTable}
@@ -1405,6 +1406,28 @@ class ParquetV2PartitionDiscoverySuite extends ParquetPartitionDiscoverySuite {
     assert("p_int=10/p_float=1.0" === path)
   }
 
+  test("SPARK-51830 handle exception and validate partition column as true") {
+    SQLConf.get.setConf(SQLConf.VALIDATE_PARTITION_COLUMNS, true)
+    val spec = Map("p_int" -> "not_a_number")
+    val schema = new StructType().add("p_int", "int")
+    val exception = intercept[Exception] {
+      PartitioningUtils.getPathFragment(spec, schema)
+    }
+    assert(
+        exception.isInstanceOf[
+          QueryExecutionErrors.FailedToCastValueToDataTypeForPartitionColumnError
+        ]
+      )
+  }
+
+  test("SPARK-51830 handle exception and validate partition column as false") {
+    SQLConf.get.setConf(SQLConf.VALIDATE_PARTITION_COLUMNS, false)
+    val spec = Map("p_int" -> "not_a_number")
+    val schema = new StructType().add("p_int", "int")
+    val path = PartitioningUtils.getPathFragment(spec, schema)
+    assert("p_int=not_a_number" === path)
+  }
+
   test("SPARK-39417: Null partition value") {
     // null partition value is replaced by DEFAULT_PARTITION_NAME before hitting getPathFragment.
     val spec = Map("p_int"-> ExternalCatalogUtils.DEFAULT_PARTITION_NAME)