diff --git a/src/main/scala/za/co/absa/standardization/udf/UDFBuilder.scala b/src/main/scala/za/co/absa/standardization/udf/UDFBuilder.scala index fa06da2..1de9cd6 100644 --- a/src/main/scala/za/co/absa/standardization/udf/UDFBuilder.scala +++ b/src/main/scala/za/co/absa/standardization/udf/UDFBuilder.scala @@ -19,7 +19,7 @@ package za.co.absa.standardization.udf import org.apache.spark.sql.expressions.UserDefinedFunction import org.apache.spark.sql.functions.udf import org.apache.spark.sql.types.DataType -import za.co.absa.standardization.config.StandardizationConfig +import za.co.absa.standardization.config.{ErrorCodesConfig, StandardizationConfig} import za.co.absa.standardization.types.parsers.NumericParser import za.co.absa.standardization.types.parsers.NumericParser.NumericParserException @@ -39,9 +39,20 @@ object UDFBuilder { val vColumnNameForError = columnNameForError val vDefaultValue = defaultValue val vColumnNullable = columnNullable - val vStdConfig = stdConfig + val vErrorCodes = stdConfig.errorCodes - udf[UDFResult[T], String](numericParserToTyped(_, sourceDataType, targetDataType, vParser, vColumnNullable, vColumnNameForError, vStdConfig, vDefaultValue)) + udf[UDFResult[T], String]( + numericParserToTyped( + _, + sourceDataType, + targetDataType, + vParser, + vColumnNullable, + vColumnNameForError, + vErrorCodes, + vDefaultValue + ) + ) } private def numericParserToTyped[T](input: String, @@ -50,14 +61,14 @@ object UDFBuilder { parser: NumericParser[T], columnNullable: Boolean, columnNameForError: String, - stdConfig: StandardizationConfig, + errorCodes: ErrorCodesConfig, defaultValue: Option[T]): UDFResult[T] = { val result = Option(input) match { case Some(string) => parser.parse(string).map(Some(_)) case None if columnNullable => Success(None) case None => Failure(nullException) } - UDFResult.fromTry(result, columnNameForError, input, sourceDataType.typeName, targetDataType.typeName, None, stdConfig, defaultValue) + UDFResult.fromTry(result, columnNameForError, input, sourceDataType.typeName, targetDataType.typeName, None, errorCodes, defaultValue) } private val nullException = new NumericParserException("Null value on input for non-nullable field") diff --git a/src/main/scala/za/co/absa/standardization/udf/UDFResult.scala b/src/main/scala/za/co/absa/standardization/udf/UDFResult.scala index 6329726..cd8f229 100644 --- a/src/main/scala/za/co/absa/standardization/udf/UDFResult.scala +++ b/src/main/scala/za/co/absa/standardization/udf/UDFResult.scala @@ -18,7 +18,7 @@ package za.co.absa.standardization.udf import za.co.absa.standardization.ErrorMessage import za.co.absa.standardization.StandardizationErrorMessage -import za.co.absa.standardization.config.StandardizationConfig +import za.co.absa.standardization.config.{ErrorCodesConfig, StandardizationConfig} import scala.util.{Failure, Success, Try} @@ -38,11 +38,22 @@ object UDFResult { pattern: Option[String], stdConfig: StandardizationConfig, defaultValue: Option[T] = None): UDFResult[T] = { + fromTry(result, columnName, rawValue, sourceType, targetType, pattern, stdConfig.errorCodes, defaultValue) + } + + def fromTry[T](result: Try[Option[T]], + columnName: String, + rawValue: String, + sourceType: String, + targetType: String, + pattern: Option[String], + errorCodes: ErrorCodesConfig, + defaultValue: Option[T]): UDFResult[T] = { result match { case Success(success) => UDFResult.success(success) - case Failure(_) if Option(rawValue).isEmpty => UDFResult(defaultValue, Seq(StandardizationErrorMessage.stdNullErr(columnName)(stdConfig.errorCodes))) + case Failure(_) if Option(rawValue).isEmpty => UDFResult(defaultValue, Seq(StandardizationErrorMessage.stdNullErr(columnName)(errorCodes))) case Failure(_) => - UDFResult(defaultValue, Seq(StandardizationErrorMessage.stdCastErr(columnName, rawValue, sourceType, targetType, pattern)(stdConfig.errorCodes))) + UDFResult(defaultValue, Seq(StandardizationErrorMessage.stdCastErr(columnName, rawValue, sourceType, targetType, pattern)(errorCodes))) } } } diff --git a/src/test/scala/za/co/absa/standardization/udf/UDFBuilderSuite.scala b/src/test/scala/za/co/absa/standardization/udf/UDFBuilderSuite.scala index 383f892..2539c32 100644 --- a/src/test/scala/za/co/absa/standardization/udf/UDFBuilderSuite.scala +++ b/src/test/scala/za/co/absa/standardization/udf/UDFBuilderSuite.scala @@ -18,17 +18,27 @@ package za.co.absa.standardization.udf import java.io.{ByteArrayInputStream, ByteArrayOutputStream, ObjectInputStream, ObjectOutputStream, ObjectStreamClass} import org.apache.spark.sql.expressions.UserDefinedFunction +import org.apache.spark.sql.functions.col import org.apache.spark.sql.types._ import org.scalatest.funsuite.AnyFunSuite +import za.co.absa.spark.commons.test.SparkTestBase import za.co.absa.standardization.RecordIdGeneration.IdType.NoId -import za.co.absa.standardization.config.{BasicMetadataColumnsConfig, BasicStandardizationConfig, StandardizationConfig} +import za.co.absa.standardization.config.{ + BasicErrorCodesConfig, + BasicMetadataColumnsConfig, + BasicStandardizationConfig, + DefaultStandardizationConfig, + StandardizationConfig +} import za.co.absa.standardization.schema.MetadataKeys import za.co.absa.standardization.types.TypedStructField._ import za.co.absa.standardization.types.parsers.IntegralParser.{PatternIntegralParser, RadixIntegralParser} import za.co.absa.standardization.types.parsers.{DecimalParser, FractionalParser} import za.co.absa.standardization.types.{TypeDefaults, CommonTypeDefaults, TypedStructField} -class UDFBuilderSuite extends AnyFunSuite { +import scala.util.{Failure, Success} + +class UDFBuilderSuite extends AnyFunSuite with SparkTestBase { private implicit val defaults: TypeDefaults = CommonTypeDefaults private val stdConfig = BasicStandardizationConfig .fromDefault() @@ -144,4 +154,97 @@ class UDFBuilderSuite extends AnyFunSuite { ois.readObject().asInstanceOf[UserDefinedFunction] } + test("Serialization and deserialization of stringUdfViaNumericParser with default config") { + val fieldName = "test" + val field: StructField = StructField(fieldName, IntegerType, nullable = true, new MetadataBuilder() + .putString(MetadataKeys.Pattern, "000000") + .build) + val typedField = TypedStructField(field) + + val numericTypeField = typedField.asInstanceOf[NumericTypeStructField[Int]] + val defaultValue: Option[Int] = typedField.defaultValueWithGlobal.get.map(_.asInstanceOf[Int]) + val parser = numericTypeField.parser.get.asInstanceOf[PatternIntegralParser[Int]] + val udfFnc = UDFBuilder.stringUdfViaNumericParser( + StringType, + field.dataType, + parser, + numericTypeField.nullable, + fieldName, + DefaultStandardizationConfig, + defaultValue + ) + //write + val baos = new ByteArrayOutputStream + val oos = new ObjectOutputStream(baos) + oos.writeObject(udfFnc) + oos.flush() + val serialized = baos.toByteArray + assert(serialized.nonEmpty) + //read + val ois = new ObjectInputStream(new ByteArrayInputStream(serialized)) { + override def resolveClass(desc: ObjectStreamClass): Class[_] = + Class.forName(desc.getName, false, loader) + } + ois.readObject().asInstanceOf[UserDefinedFunction] + import spark.implicits._ + + val rows = Seq("000123", "bad").toDF("input").select(udfFnc(col("input")).as("result")).collect() + assert(rows.length === 2) + } + + test("UDFResult.fromTry uses provided error codes config") { + val errorCodes = BasicErrorCodesConfig("cast-code", "null-code", "type-code", "schema-code") + + val successResult = UDFResult.fromTry[Int]( + Success(Some(2)), + "field", + "2", + "string", + "integer", + None, + errorCodes, + None + ) + val castResult = UDFResult.fromTry[Int]( + Failure(new RuntimeException("boom")), + "field", + "bad", + "string", + "integer", + None, + errorCodes, + Some(0) + ) + val nullResult = UDFResult.fromTry[Int]( + Failure(new RuntimeException("boom")), + "field", + null, + "string", + "integer", + None, + errorCodes, + Some(1) + ) + + assert(successResult === UDFResult.success(Some(2))) + assert(castResult.result === Some(0)) + assert(castResult.error.map(_.errCode) === Seq("cast-code")) + assert(nullResult.result === Some(1)) + assert(nullResult.error.map(_.errCode) === Seq("null-code")) + } + + test("UDFResult.fromTry keeps StandardizationConfig overload") { + val result = UDFResult.fromTry[Int]( + Failure(new RuntimeException("boom")), + "field", + "bad", + "string", + "integer", + None, + stdConfig + ) + + assert(result.error.map(_.errCode) === Seq(stdConfig.errorCodes.castError)) + } + }