如何在 Spark UDF 中设置小数返回类型的精度和小数位数？

Question

这是我的示例代码。我期望小数（16,4）作为 UDF 的返回类型，但它是小数（38,18）。

有没有更好的解决方案？

我不期待答案“cast(price as decimal(16,4))”，因为我的 UDF 中除了强制转换之外还有其他一些业务逻辑。

提前致谢。

import scala.util.Try
import org.apache.spark.sql.functions.udf
import org.apache.spark.sql.types.Decimal
val spark = SparkSession.builder().master("local[*]").appName("Test").getOrCreate()
import spark.implicits._

val stringToDecimal = udf((s:String, precision:Int, scale: Int) => {
  Try(Decimal(BigDecimal(s), precision, scale)).toOption
})

spark.udf.register("stringToDecimal", stringToDecimal)

val inDf = Seq(
  ("1", "864.412"),
  ("2", "1.600"),
  ("3", "2,56")).toDF("id", "price")

val outDf = inDf.selectExpr("id", "stringToDecimal(price, 16, 4) as price")
outDf.printSchema()
outDf.show()

------------------output----------------
root
  |-- id: string (nullable = true)
  |-- price: decimal(38,18) (nullable = true)

+---+--------------------+
| id|               price|
+---+--------------------+
|  1|864.4120000000000...|
|  2|1.600000000000000000|
|  3|                null|
+---+--------------------+

Answer 1

对于 Spark 3.0 及更低版本，您无法设置 Spark 用户定义函数 (UDF) 返回的十进制精度和小数位数，因为在创建 UDF 时会删除精度和小数位数。

说明

要创建 UDF，无论是通过使用 lambda/function 作为参数调用函数 udf，还是使用 sparkSession.udf.register 方法直接将 lambda/function 注册为 UDF，Spark 都需要转换参数类型并返回lambda/函数到Spark's DataType

为此，Spark 将使用 ScalaReflection 类中的方法 schemaFor 将 scala 类型映射到 Spark 的 DataType。

对于BigDecimal或Decimal类型，映射如下：

case t if isSubtype(t, localTypeOf[BigDecimal]) =>
  Schema(DecimalType.SYSTEM_DEFAULT, nullable = true)
case t if isSubtype(t, localTypeOf[java.math.BigDecimal]) =>
  Schema(DecimalType.SYSTEM_DEFAULT, nullable = true)
case t if isSubtype(t, localTypeOf[Decimal]) =>
  Schema(DecimalType.SYSTEM_DEFAULT, nullable = true)

意味着当您的 lambda/函数返回 BigDecimal 或 Decimal 时，UDF 的返回类型将为 DecimalType.SYSTEM_DEFAULT。 DecimalType.SYSTEM_DEFAULT 类型是 Decimal，精度为 38，比例为 18：

val MAX_PRECISION = 38
...
val SYSTEM_DEFAULT: DecimalType = DecimalType(MAX_PRECISION, 18)

结论

因此，每次您将 lambda 或返回 Decimal 或 BigDecimal 的函数转换为 Spark 的 UDF 时，都会以默认精度 38 和小数位数 18 擦除精度和小数位数。

所以你唯一的方法是关注previous answer 并在调用它时转换 UDF 的返回值

Answer 2

Spark 将 Decimal 与 decimal(38, 18) 关联。你需要一个明确的演员表

$"price".cast(DataTypes.createDecimalType(32,2))