* InfinitySupport changed from mix-in train to class (#71)

benedeki · web-flow · commit 4807526d1947 · 2025-08-29T15:51:43.000+02:00
* created `InfinitySupportIso` for support of Iso pattern fallback in date and timestamp parsing
* Integrated the new class into `TypeParser` instead of using it as trait
* Removed some compilation warnings
diff --git a/src/main/scala/za/co/absa/standardization/SchemaValidator.scala b/src/main/scala/za/co/absa/standardization/SchemaValidator.scala
@@ -18,7 +18,6 @@ package za.co.absa.standardization
 
 import org.apache.spark.sql.SparkSession
 import org.apache.spark.sql.types._
-import za.co.absa.standardization.ErrorMessage
 import za.co.absa.standardization.types.{TypeDefaults, TypedStructField}
 import za.co.absa.standardization.validation.field.FieldValidationIssue
 
@@ -116,7 +115,7 @@ object SchemaValidator {
             fields += prefixedField
         }
       }
-      fields.toSeq
+      fields
     }
 
     def flattenArray(field: StructField, arr: ArrayType, structPath: String): Seq[FlatField] = {
@@ -128,7 +127,7 @@ object SchemaValidator {
           val prefixedField = FlatField(structPath, field)
           arrayFields += prefixedField
       }
-      arrayFields.toSeq
+      arrayFields
     }
 
     flattenStruct(schema, "")
diff --git a/src/main/scala/za/co/absa/standardization/stages/InfinitySupport.scala b/src/main/scala/za/co/absa/standardization/stages/InfinitySupport.scala
@@ -20,26 +20,51 @@ import org.apache.spark.sql.Column
 import org.apache.spark.sql.functions.{lit, when}
 import org.apache.spark.sql.types.DataType
 
-trait InfinitySupport {
-  protected def infMinusSymbol: Option[String]
-  protected def infMinusValue: Option[String]
-  protected def infPlusSymbol: Option[String]
-  protected def infPlusValue: Option[String]
-  protected def canParseInfValue(value: String): Boolean
-  protected val origType: DataType
-
-  def replaceInfinitySymbols(column: Column): Column = {
-    val columnWithNegativeInf: Column = infMinusSymbol.flatMap { minusSymbol =>
-      infMinusValue.map { minusValue =>
-        when(column === lit(minusSymbol), lit(minusValue)).otherwise(column)
+class InfinitySupport(
+                       val infMinusSymbol: Option[String],
+                       val infMinusValue: Option[String],
+                       val infPlusSymbol: Option[String],
+                       val infPlusValue: Option[String],
+                       val origType: DataType
+                     ) {
+  private val hasInfinityDefined: Boolean = (infMinusSymbol.isDefined && infMinusValue.isDefined) ||
+                                              (infPlusSymbol.isDefined && infPlusValue.isDefined)
+
+  protected def replaceSymbol(column: Column,
+                              possibleSymbol: Option[String],
+                              possibleValue: Option[String],
+                              valueToColumn: String => Column
+                             ): Column = {
+    possibleSymbol.flatMap { symbol =>
+      possibleValue.map { value =>
+        when(column === lit(symbol).cast(origType), valueToColumn(value))
       }
     }.getOrElse(column)
+  }
 
-    infPlusSymbol.flatMap { plusSymbol =>
-      infPlusValue.map { plusValue =>
-        when(columnWithNegativeInf === lit(plusSymbol), lit(plusValue))
-          .otherwise(columnWithNegativeInf)
-      }
-    }.getOrElse(columnWithNegativeInf)
+  protected def defaultInfinityValueInjection(value: String): Column = lit(value).cast(origType)
+
+  protected def executeReplacement(column: Column, conversion: Column => Column): Column = {
+    val columnWithNegativeInf = replaceSymbol(column, infMinusSymbol, infMinusValue, defaultInfinityValueInjection)
+    val columnWithPositiveInf = replaceSymbol(columnWithNegativeInf, infPlusSymbol, infPlusValue, defaultInfinityValueInjection)
+    conversion(columnWithPositiveInf.otherwise(column))
+  }
+
+  def replaceInfinitySymbols(column: Column, conversion: Column => Column = c => c): Column = {
+    if (hasInfinityDefined) {
+      executeReplacement(column, conversion)
+    } else {
+      conversion(column)
+    }
+  }
+}
+
+object InfinitySupport {
+  def apply(infMinusSymbol: Option[String],
+            infMinusValue: Option[String],
+            infPlusSymbol: Option[String],
+            infPlusValue: Option[String],
+            origType: DataType): InfinitySupport = {
+    new InfinitySupport(infMinusSymbol, infMinusValue, infPlusSymbol, infPlusValue, origType)
   }
 }
diff --git a/src/main/scala/za/co/absa/standardization/stages/InfinitySupportIso.scala b/src/main/scala/za/co/absa/standardization/stages/InfinitySupportIso.scala
@@ -0,0 +1,116 @@
+/*
+ * Copyright 2025 ABSA Group Limited
+ *
+ * Licensed 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 za.co.absa.standardization.stages
+
+import org.apache.spark.sql.Column
+import org.apache.spark.sql.functions._
+import org.apache.spark.sql.types.{DataType, DateType, TimestampType}
+
+import java.text.SimpleDateFormat
+import java.time.OffsetDateTime
+import java.time.format.DateTimeFormatter.{ISO_DATE, ISO_DATE_TIME}
+import scala.util.Try
+
+abstract class InfinitySupportIso(
+                                   infMinusSymbol: Option[String],
+                                   infMinusValue: Option[String],
+                                   infPlusSymbol: Option[String],
+                                   infPlusValue: Option[String],
+                                   origType: DataType
+                                 ) extends InfinitySupport(infMinusSymbol, infMinusValue, infPlusSymbol, infPlusValue, origType) {
+
+  def isoCast(value: String): Column
+
+  def useIsoForInfMinus: Boolean
+  def useIsoForInfPlus: Boolean
+
+  def chooseInjectionFunction(isIso: Boolean, conversion: Column => Column, value: String): Column = {
+    if (isIso) {
+      isoCast(value)
+    } else {
+      conversion(defaultInfinityValueInjection(value))
+    }
+  }
+
+  override def executeReplacement(column: Column, conversion: Column => Column): Column = {
+    if (useIsoForInfMinus || useIsoForInfPlus) {
+      val columnWithNegativeInf = replaceSymbol(column, infMinusSymbol, infMinusValue, chooseInjectionFunction(useIsoForInfMinus, conversion, _))
+      val columnWithPositiveInf = replaceSymbol(columnWithNegativeInf, infPlusSymbol, infPlusValue, chooseInjectionFunction(useIsoForInfPlus, conversion, _))
+      columnWithPositiveInf.otherwise(conversion(column))
+    } else {
+      super.executeReplacement(column, conversion)
+    }
+  }
+}
+
+object InfinitySupportIso {
+  def apply(
+             infMinusSymbol: Option[String],
+             infMinusValue: Option[String],
+             infPlusSymbol: Option[String],
+             infPlusValue: Option[String],
+             origType: DataType,
+             targetType: DataType
+           ): InfinitySupportIso = {
+    targetType match {
+      case DateType => new InfinitySupportIsoDate(infMinusSymbol, infMinusValue, infPlusSymbol, infPlusValue, origType)
+      case TimestampType => new InfinitySupportIsoTimestamp(infMinusSymbol, infMinusValue, infPlusSymbol, infPlusValue, origType)
+      case _ => throw new IllegalArgumentException(s"InfinitySupportIso does not support target type $targetType")
+    }
+  }
+
+  def isOfISODateFormat(dateValue: Option[String]): Boolean = {
+    dateValue.exists(value =>
+      Try {
+        ISO_DATE.parse(value)
+      }.isSuccess)
+  }
+
+  def isOfISOTimestampFormat(timestampValue: Option[String]): Boolean = {
+    timestampValue.exists(value =>
+      Try {
+        OffsetDateTime.parse(value)
+      }.isSuccess)
+  }
+
+  class InfinitySupportIsoDate(
+                                infMinusSymbol: Option[String],
+                                infMinusValue: Option[String],
+                                infPlusSymbol: Option[String],
+                                infPlusValue: Option[String],
+                                origType: DataType
+                              ) extends InfinitySupportIso(infMinusSymbol, infMinusValue, infPlusSymbol, infPlusValue, origType) {
+    override def isoCast(value: String): Column = to_date(lit(value))
+    val useIsoForInfMinus: Boolean = isOfISODateFormat(infMinusValue)
+    val useIsoForInfPlus: Boolean = isOfISODateFormat(infPlusValue)
+  }
+
+  class InfinitySupportIsoTimestamp(
+                                     infMinusSymbol: Option[String],
+                                     infMinusValue: Option[String],
+                                     infPlusSymbol: Option[String],
+                                     infPlusValue: Option[String],
+                                     origType: DataType
+                                   ) extends InfinitySupportIso(infMinusSymbol, infMinusValue, infPlusSymbol, infPlusValue, origType) {
+    override def isoCast(value: String): Column = to_timestamp(lit(value))
+    val useIsoForInfMinus: Boolean = isOfISOTimestampFormat(infMinusValue)
+    val useIsoForInfPlus: Boolean = isOfISOTimestampFormat(infPlusValue)
+  }
+
+}
+
+
diff --git a/src/main/scala/za/co/absa/standardization/stages/TypeParser.scala b/src/main/scala/za/co/absa/standardization/stages/TypeParser.scala
@@ -34,7 +34,7 @@ import za.co.absa.standardization.schema.{MetadataKeys, MetadataValues, StdSchem
 import za.co.absa.standardization.time.DateTimePattern
 import za.co.absa.standardization.typeClasses.{DoubleLike, LongLike}
 import za.co.absa.standardization.types.TypedStructField._
-import za.co.absa.standardization.types.parsers.{DateTimeParser => DateTimeParserImpl}
+import za.co.absa.standardization.types.parsers.DateTimeParser
 import za.co.absa.standardization.types.{ParseOutput, TypeDefaults, TypedStructField}
 import za.co.absa.standardization.udf.{UDFBuilder, UDFNames}
 
@@ -318,13 +318,13 @@ object TypeParser {
   }
 
   private abstract class NumericParser[N: TypeTag](override val field: NumericTypeStructField[N])
-                                                  (implicit defaults: TypeDefaults) extends ScalarParser[N] with InfinitySupport {
-    override protected val infMinusSymbol: Option[String] = metadata.getOptString(MetadataKeys.MinusInfinitySymbol)
-    override protected val infMinusValue: Option[String] = metadata.getOptString(MetadataKeys.MinusInfinityValue)
-    override protected val infPlusSymbol: Option[String] = metadata.getOptString(MetadataKeys.PlusInfinitySymbol)
-    override protected val infPlusValue: Option[String] = metadata.getOptString(MetadataKeys.PlusInfinityValue)
-    override protected def canParseInfValue(value: String): Boolean = false
-    private val columnWithInfinityReplaced = replaceInfinitySymbols(column).cast(origType)
+                                                  (implicit defaults: TypeDefaults) extends ScalarParser[N] {
+    private val infinitySupport = InfinitySupport(
+      metadata.getOptString(MetadataKeys.MinusInfinitySymbol),
+      metadata.getOptString(MetadataKeys.MinusInfinityValue),
+      metadata.getOptString(MetadataKeys.PlusInfinitySymbol),
+      metadata.getOptString(MetadataKeys.PlusInfinityValue),
+      origType)
 
     override protected def standardizeAfterCheck(stdConfig: StandardizationConfig)(implicit logger: Logger): ParseOutput = {
       if (field.needsUdfParsing) {
@@ -345,9 +345,9 @@ object TypeParser {
         val columnWithProperDecimalSymbols: Column = if (replacements.nonEmpty) {
           val from = replacements.keys.mkString
           val to = replacements.values.mkString
-          translate(columnWithInfinityReplaced, from, to)
+          infinitySupport.replaceInfinitySymbols(column, translate(_, from, to))
         } else {
-          columnWithInfinityReplaced
+          infinitySupport.replaceInfinitySymbols(column)
         }
 
         val columnToCast = if (field.allowInfinity && (decimalSymbols.infinityValue != InfinityStr)) {
@@ -504,28 +504,17 @@ object TypeParser {
     * Date          | O                               | ->to_utc_timestamp->to_date
     * Other         | ->String->to_date               | ->String->to_timestamp->to_utc_timestamp->to_date
     */
-  private abstract class DateTimeParser[T](implicit defaults: TypeDefaults) extends PrimitiveParser[T] with InfinitySupport {
+  private abstract class DateTimeParser[T](implicit defaults: TypeDefaults) extends PrimitiveParser[T] {
     override val field: DateTimeTypeStructField[T]
     protected val pattern: DateTimePattern = field.pattern.get.get
-    override protected val infMinusSymbol: Option[String] = metadata.getOptString(MetadataKeys.MinusInfinitySymbol)
-    override protected val infMinusValue: Option[String] = metadata.getOptString(MetadataKeys.MinusInfinityValue)
-    override protected val infPlusSymbol: Option[String] = metadata.getOptString(MetadataKeys.PlusInfinitySymbol)
-    override protected val infPlusValue: Option[String] = metadata.getOptString(MetadataKeys.PlusInfinityValue)
-
-
-    private val IsoDatePattern = "yyyy-MM-dd"
-    private val IsoTimestampPattern = "yyyy-MM-dd HH:mm:ss"
-
-    private lazy val dateTimeParser: DateTimeParserImpl = field.parser.get
-
-    override protected def canParseInfValue(value: String): Boolean = {
-      Try{
-        field.dataType match{
-          case DateType => dateTimeParser.parseDate(value)
-          case TimestampType => dateTimeParser.parseTimestamp(value)
-        }
-      }.isSuccess
-    }
+    private val infinitySupport = InfinitySupportIso(
+      metadata.getOptString(MetadataKeys.MinusInfinitySymbol),
+      metadata.getOptString(MetadataKeys.MinusInfinityValue),
+      metadata.getOptString(MetadataKeys.PlusInfinitySymbol),
+      metadata.getOptString(MetadataKeys.PlusInfinityValue),
+      origType,
+      field.dataType
+    )
 
     protected val replaceCenturyUDF: UserDefinedFunction = udf((inputDate: String, centuryPattern: String) => {
       val centuryIndex = centuryPattern.indexOf(DateTimePattern.patternCenturyChar)
@@ -538,45 +527,6 @@ object TypeParser {
       pendedInput.substring(0, centuryIndex) + modifiedChar + pendedInput.substring(centuryIndex + 1)
     })
 
-    private val columnWithInfinityReplaced: Column = {
-      val replaced = replaceInfinitySymbols(column)
-
-      val originalCastFunc: Column => Column = if (pattern.isEpoch) {
-        col => (col.cast(decimalType) / pattern.epochFactor).cast(TimestampType)
-      } else {
-        col => castStringColumn(col)
-      }
-
-      val isoPattern = field.dataType match {
-        case DateType => IsoDatePattern
-        case TimestampType => IsoTimestampPattern
-      }
-      val isoCastFunc: Column => Column = col => field.dataType match{
-        case DateType => to_date(col,isoPattern)
-        case TimestampType => to_timestamp(col, isoPattern)
-      }
-
-      infMinusSymbol.flatMap { minusSymbol =>
-        infMinusValue.map { minusValue =>
-          if (canParseInfValue(minusValue)){
-            originalCastFunc(replaced)
-          } else {
-            when(replaced === lit(minusValue), isoCastFunc(lit(minusValue))).otherwise(originalCastFunc(replaced))
-          }
-        }
-      }.getOrElse{
-        infPlusSymbol.flatMap{ plusSymbol =>
-          infPlusValue.map { plusValue=>
-            if (canParseInfValue(plusValue)){
-              originalCastFunc(replaced)
-            } else{
-              when(replaced === lit(plusValue), isoCastFunc(lit(plusValue))).otherwise(originalCastFunc(replaced))
-            }
-          }
-        }.getOrElse(originalCastFunc(replaced))
-      }
-    }
-
     override protected def assemblePrimitiveCastLogic: Column = {
       if (pattern.isEpoch) {
         castEpoch()
@@ -598,20 +548,20 @@ object TypeParser {
       // underlyingType match {
       origType match {
         case _: NullType                  => nullColumn
-        case _: DateType                  => castDateColumn(columnWithInfinityReplaced)
-        case _: TimestampType             => castTimestampColumn(columnWithInfinityReplaced)
-        case _: StringType                => castStringColumn(columnWithInfinityReplaced)
+        case _: DateType                  => infinitySupport.replaceInfinitySymbols(column, castDateColumn)
+        case _: TimestampType             => infinitySupport.replaceInfinitySymbols(column, castTimestampColumn)
+        case _: StringType                => infinitySupport.replaceInfinitySymbols(column, castStringColumn)
         case ot: DoubleType               =>
           // this case covers some IBM date format where it's represented as a double ddmmyyyy.hhmmss
           patternNeeded(ot)
-          castFractionalColumn(columnWithInfinityReplaced, ot)
+          infinitySupport.replaceInfinitySymbols(column, castFractionalColumn(_, ot))
         case ot: FloatType                =>
           // this case covers some IBM date format where it's represented as a double ddmmyyyy.hhmmss
           patternNeeded(ot)
-          castFractionalColumn(columnWithInfinityReplaced, ot)
+          infinitySupport.replaceInfinitySymbols(column, castFractionalColumn(_, ot))
         case ot                           =>
           patternNeeded(ot)
-          castNonStringColumn(columnWithInfinityReplaced, ot)
+          infinitySupport.replaceInfinitySymbols(column, castNonStringColumn(_, ot))
       }
     }
 
@@ -633,7 +583,7 @@ object TypeParser {
     }
 
     protected def castEpoch(): Column = {
-      (columnWithInfinityReplaced.cast(decimalType) / pattern.epochFactor).cast(TimestampType)
+      infinitySupport.replaceInfinitySymbols(column, c => (c.cast(decimalType) / pattern.epochFactor).cast(TimestampType))
     }
 
     protected def castStringColumn(stringColumn: Column): Column
diff --git a/src/test/scala/za/co/absa/standardization/StandardizationCsvSuite.scala b/src/test/scala/za/co/absa/standardization/StandardizationCsvSuite.scala
@@ -19,7 +19,6 @@ package za.co.absa.standardization
 import org.apache.spark.sql.functions._
 import org.apache.spark.sql.types._
 import org.scalatest.funsuite.AnyFunSuite
-import za.co.absa.standardization.ErrorMessage
 import za.co.absa.spark.commons.implicits.DataFrameImplicits.DataFrameEnhancements
 import za.co.absa.spark.commons.test.SparkTestBase
 import za.co.absa.standardization.RecordIdGeneration.IdType.NoId
diff --git a/src/test/scala/za/co/absa/standardization/TestSamples.scala b/src/test/scala/za/co/absa/standardization/TestSamples.scala
@@ -16,7 +16,6 @@
 
 package za.co.absa.standardization
 
-import za.co.absa.standardization.ErrorMessage
 import za.co.absa.standardization.RecordIdGeneration.IdType.NoId
 import za.co.absa.standardization.config.{BasicMetadataColumnsConfig, BasicStandardizationConfig, ErrorCodesConfig}
 
diff --git a/src/test/scala/za/co/absa/standardization/stages/InfinitySupportIsoSuite.scala b/src/test/scala/za/co/absa/standardization/stages/InfinitySupportIsoSuite.scala

Original file line number	Diff line number	Diff line change
`@@ -18,7 +18,6 @@ package za.co.absa.standardization`
`18`	`18`
`19`	`19`	`import org.apache.spark.sql.SparkSession`
`20`	`20`	`import org.apache.spark.sql.types._`
`21`		`-import za.co.absa.standardization.ErrorMessage`
`22`	`21`	`import za.co.absa.standardization.types.{TypeDefaults, TypedStructField}`
`23`	`22`	`import za.co.absa.standardization.validation.field.FieldValidationIssue`
`24`	`23`
`@@ -116,7 +115,7 @@ object SchemaValidator {`
`116`	`115`	`fields += prefixedField`
`117`	`116`	`}`
`118`	`117`	`}`
`119`		`- fields.toSeq`
	`118`	`+ fields`
`120`	`119`	`}`
`121`	`120`
`122`	`121`	`def flattenArray(field: StructField, arr: ArrayType, structPath: String): Seq[FlatField] = {`
`@@ -128,7 +127,7 @@ object SchemaValidator {`
`128`	`127`	`val prefixedField = FlatField(structPath, field)`
`129`	`128`	`arrayFields += prefixedField`
`130`	`129`	`}`
`131`		`- arrayFields.toSeq`
	`130`	`+ arrayFields`
`132`	`131`	`}`
`133`	`132`
`134`	`133`	`flattenStruct(schema, "")`